examples/acoustic/01-model-test/03-SEAM-I/02_inversion.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import integrate\n",
    "import sys\n",
    "import os\n",
    "sys.path.append(\"../../../../\")\n",
    "from ADFWI.propagator  import *\n",
    "from ADFWI.model       import *\n",
    "from ADFWI.view        import *\n",
    "from ADFWI.utils       import *\n",
    "from ADFWI.survey      import *\n",
    "from ADFWI.fwi         import *\n",
    "\n",
    "project_path = \"./examples/acoustic/01-model-test/03-SEAM-I/data\"\n",
    "if not os.path.exists(os.path.join(project_path,\"model\")):\n",
    "    os.makedirs(os.path.join(project_path,\"model\"))\n",
    "if not os.path.exists(os.path.join(project_path,\"waveform\")):\n",
    "    os.makedirs(os.path.join(project_path,\"waveform\"))\n",
    "if not os.path.exists(os.path.join(project_path,\"survey\")):\n",
    "    os.makedirs(os.path.join(project_path,\"survey\"))\n",
    "if not os.path.exists(os.path.join(project_path,\"inversion\")):\n",
    "    os.makedirs(os.path.join(project_path,\"inversion\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the basic model parameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda:0\"         # Specify the GPU device\n",
    "dtype  = torch.float32    # Set data type to 32-bit floating point\n",
    "ox, oz = 0, 0             # Origin coordinates for x and z directions\n",
    "nz, nx = 103, 241         # Grid dimensions in z and x directions\n",
    "dx, dz = 50, 50           # Grid spacing in x and z directions\n",
    "nt, dt = 1601, 0.005      # Time steps and time interval\n",
    "nabc = 30                 # Thickness of the absorbing boundary layer\n",
    "f0 = 3                    # Initial frequency in Hz\n",
    "free_surface = True       # Enable free surface boundary condition"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the initial velocity model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ADFWI.propagator.gradient_process import smooth2d\n",
    "\n",
    "model_path =\"./examples/datasets/SEAM-I_source\"\n",
    "\n",
    "vp_true = np.loadtxt(os.path.join(model_path,\"SEAM_I_1201_515_10m.dat\"))[::5,::5].T\n",
    "vp_init = np.copy(vp_true)\n",
    "vp_init = smooth2d(vp_true, span = 10)\n",
    "vp_init[vp_true[:,:]==1490]=1490\n",
    "rho_init    = np.power(vp_init, 0.25) * 310\n",
    "\n",
    "# Create an acoustic model with the specified parameters.\n",
    "model = AcousticModel(\n",
    "    ox, oz, nx, nz, dx, dz,\n",
    "    vp_init, rho_init,\n",
    "    vp_bound=[vp_true.min(), vp_true.max()],\n",
    "    vp_grad=True, \n",
    "    free_surface=free_surface,\n",
    "    abc_type=\"PML\",\n",
    "    abc_jerjan_alpha=0.007,\n",
    "    nabc=nabc, \n",
    "    device=device,\n",
    "    dtype=dtype\n",
    ")\n",
    "\n",
    "# Save the initialized model to a file for later use.\n",
    "model.save(os.path.join(project_path, \"model/init_model.npz\"))\n",
    "\n",
    "# Print the model's representation for verification.\n",
    "print(model.__repr__())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model._plot_vp_rho(figsize=(12,5),wspace=0.15,cbar_pad_fraction=-0.05,cmap='coolwarm',save_path=os.path.join(project_path,\"model/init_vp_rho.png\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the observed systems: Survey = Source + Receiver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define source positions in the model\n",
    "src_z = np.array([90 for i in range(5, nx, 10)])  # Z-coordinates for sources\n",
    "src_x = np.array([i  for i in range(5, nx, 10)])  # X-coordinates for sources\n",
    "\n",
    "# Generate wavelet for the source\n",
    "src_t, src_v = wavelet(nt, dt, f0, amp0=1)  # Create time and wavelet amplitude\n",
    "src_v = integrate.cumtrapz(src_v, axis=-1, initial=0)  # Integrate wavelet to get velocity\n",
    "\n",
    "source = Source(nt=nt, dt=dt, f0=f0)  # Initialize source object\n",
    "\n",
    "# Method 1: Add multiple sources at once (commented out)\n",
    "# source.add_sources(src_x=src_x, src_z=src_z, src_wavelet=src_v, src_type='mt', src_mt=np.array([[1,0,0],[0,1,0],[0,0,1]]))\n",
    "\n",
    "# Method 2: Loop through each source position to add them individually\n",
    "for i in range(len(src_x)):\n",
    "    source.add_source(src_x=src_x[i], src_z=src_z[i], src_wavelet=src_v, src_type=\"mt\", src_mt=np.array([[1,0,0],[0,1,0],[0,0,1]]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define receiver positions in the model\n",
    "rcv_z = np.array([1 for i in range(0, nx, 2)])  # Z-coordinates for receivers\n",
    "rcv_x = np.array([j for j in range(0, nx, 2)])  # X-coordinates for receivers\n",
    "\n",
    "receiver = Receiver(nt=nt, dt=dt)  # Initialize receiver object\n",
    "\n",
    "# Method 1: Add all receivers at once (commented out)\n",
    "# receiver.add_receivers(rcv_x=rcv_x, rcv_z=rcv_z, rcv_type='pr')\n",
    "\n",
    "# Method 2: Loop through each receiver position to add them individually\n",
    "for i in range(len(rcv_x)):\n",
    "    receiver.add_receiver(rcv_x=rcv_x[i], rcv_z=rcv_z[i], rcv_type=\"pr\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a survey object using the defined source and receiver\n",
    "survey = Survey(source=source, receiver=receiver)\n",
    "\n",
    "# Print a representation of the survey object to check its configuration\n",
    "print(survey.__repr__())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the wavelet used in the source\n",
    "source.plot_wavelet(save_path=os.path.join(project_path, \"survey/wavelets.png\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the survey configuration over the velocity model\n",
    "survey.plot(model.vp,cmap='coolwarm',save_path=os.path.join(project_path,\"survey/observed_system_init.png\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the propagator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the wave propagator using the specified model and survey configuration\n",
    "F = AcousticPropagator(model,survey,device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve the damping array from the propagator and plot it to visualize boundary conditions\n",
    "damp = F.damp\n",
    "plot_damp(damp,save_path=os.path.join(project_path,\"model/boundary_condition.png\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load observed datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an instance of SeismicData using the survey object.\n",
    "d_obs = SeismicData(survey)\n",
    "\n",
    "# Load observed waveform data from a specified file.\n",
    "d_obs.load(os.path.join(project_path, \"waveform/obs_data.npz\"))\n",
    "\n",
    "# Print a summary representation of the observed seismic data.\n",
    "print(d_obs.__repr__())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inversion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the L2 misfit function for waveform inversion.\n",
    "from ADFWI.fwi.misfit import Misfit_waveform_L2\n",
    "\n",
    "# Set the number of iterations for the inversion process.\n",
    "iteration = 300\n",
    "\n",
    "# Initialize the optimizer (Adam) for model parameters with a learning rate.\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=10)\n",
    "\n",
    "# Set up a learning rate scheduler to adjust the learning rate over time.\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=200, gamma=0.75, last_epoch=-1)\n",
    "\n",
    "# Configure the misfit function to compute the loss based on the observed data.\n",
    "loss_fn = Misfit_waveform_L2(dt=dt)\n",
    "\n",
    "# Create a gradient mask to selectively process gradients during inversion.\n",
    "gradient_processor = GradProcessor()\n",
    "\n",
    "# Initialize the acoustic full waveform inversion (FWI) object.\n",
    "fwi = AcousticFWI(propagator=F,\n",
    "                  model=model,\n",
    "                  optimizer=optimizer,\n",
    "                  scheduler=scheduler,\n",
    "                  loss_fn=loss_fn,\n",
    "                  obs_data=d_obs,\n",
    "                  gradient_processor=gradient_processor,\n",
    "                  waveform_normalize=True, \n",
    "                  cache_result=True,  \n",
    "                  save_fig_epoch=50,  \n",
    "                  save_fig_path=os.path.join(project_path, \"inversion\"))\n",
    "\n",
    "# Run the forward modeling for the specified number of iterations.\n",
    "fwi.forward(iteration=iteration, batch_size=None, checkpoint_segments=1)\n",
    "\n",
    "# Retrieve the inversion results: updated velocity and loss values.\n",
    "iter_vp = fwi.iter_vp\n",
    "iter_loss = fwi.iter_loss\n",
    "\n",
    "# Save the iteration results to files for later analysis.\n",
    "np.savez(os.path.join(project_path, \"inversion/iter_vp.npz\"), data=np.array(iter_vp))\n",
    "np.savez(os.path.join(project_path, \"inversion/iter_loss.npz\"), data=np.array(iter_loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## visualize the inverted results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the misfit\n",
    "plt.figure(figsize=(8,6))\n",
    "plt.plot(iter_loss,c='k')\n",
    "plt.xlabel(\"Iterations\", fontsize=12)\n",
    "plt.ylabel(\"L2-norm Misfits\", fontsize=12)\n",
    "plt.tick_params(labelsize=12)\n",
    "plt.savefig(os.path.join(project_path,\"inversion/misfit.png\"),bbox_inches='tight',dpi=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the initial model and inverted resutls\n",
    "plt.figure(figsize=(12,8))\n",
    "plt.subplot(121)\n",
    "plt.imshow(vp_init,cmap='jet_r')\n",
    "plt.subplot(122)\n",
    "plt.imshow(iter_vp[-1],cmap='jet_r')\n",
    "plt.savefig(os.path.join(project_path,\"inversion/inverted_res.png\"),bbox_inches='tight',dpi=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the gradient\n",
    "plt.figure()\n",
    "plt.imshow(fwi.iter_vp_grad[0],cmap='coolwarm')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "from IPython.display import HTML\n",
    "\n",
    "# Set up the figure for plotting\n",
    "fig, ax = plt.subplots(figsize=(8, 6))\n",
    "cax = ax.imshow(iter_vp[0], aspect='equal', cmap='jet_r', vmin=vp_true.min(), vmax=vp_true.max())\n",
    "ax.set_title('Inversion Process Visualization')\n",
    "ax.set_xlabel('X Coordinate')\n",
    "ax.set_ylabel('Z Coordinate')\n",
    "\n",
    "# Create a horizontal colorbar\n",
    "cbar = fig.colorbar(cax, ax=ax, orientation='horizontal', fraction=0.046, pad=0.2)\n",
    "cbar.set_label('Velocity (m/s)')\n",
    "\n",
    "# Adjust the layout to minimize white space\n",
    "plt.subplots_adjust(top=0.85, bottom=0.2, left=0.1, right=0.9)\n",
    "\n",
    "# Initialization function\n",
    "def init():\n",
    "    cax.set_array(iter_vp[0])  # Use the 2D array directly\n",
    "    return cax,\n",
    "\n",
    "# Animation function\n",
    "def animate(i):\n",
    "    cax.set_array(iter_vp[i])  # Update with the i-th iteration directly\n",
    "    return cax,\n",
    "\n",
    "# Create the animation\n",
    "ani = animation.FuncAnimation(fig, animate, init_func=init, frames=len(iter_vp), interval=100, blit=True)\n",
    "\n",
    "# Save the animation as a video file (e.g., MP4 format)\n",
    "ani.save(os.path.join(project_path, \"inversion/inversion_process.gif\"), writer='pillow', fps=10)\n",
    "\n",
    "# Display the animation using HTML\n",
    "plt.close(fig)  # Prevents static display of the last frame\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ADinversion",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "1e52d0fb768956d75c9106cc9d1e4cd93f3049c3ef5ccf93d433f27b0d7ebb6b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}