examples/acoustic/02-misfit-functions-test/00-Ricker-Test/01_misfit_wavelets-shift.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams[\"font.family\"] = \"Times New Roman\"\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/02-misfit-functions-test/00-Ricker-Test/Figures\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nt = 1000\n",
    "dt = 0.003\n",
    "f0 = 6\n",
    "tlist,ricker = wavelet(nt,dt,f0,amp0=1,t0=0.5)\n",
    "ricker       = np.roll(ricker,330)\n",
    "shift_ricker = np.roll(ricker,-100)\n",
    "\n",
    "plt.figure(figsize=(8,4))\n",
    "plt.plot(tlist,ricker       ,c='r',linestyle=\"-\"    ,label=\"Ricker wavelet\")\n",
    "plt.plot(tlist,shift_ricker ,c='b',linestyle=\"--\"   ,label=\"Shifted ricker wavelet\")\n",
    "plt.legend(fontsize=14)\n",
    "plt.xlabel(\"Times (s)\",fontsize=13)\n",
    "plt.ylabel(\"Amplitude\",fontsize=13)\n",
    "plt.tick_params(labelsize=13)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize(data):\n",
    "    if isinstance(data,list):\n",
    "        data = np.array(data)\n",
    "    res = (data - data.min())/(data.max() - data.min())\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_shift = np.arange(-100,100,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# L1\n",
    "from ADFWI.fwi.misfit import Misfit_waveform_L1\n",
    "loss_fn = Misfit_waveform_L1(dt=1)\n",
    "\n",
    "ricker_tensor = numpy2tensor(ricker)\n",
    "ricker_tensor = ricker_tensor.reshape(1,-1,1)\n",
    "L1_loss = []\n",
    "for i in x_shift:\n",
    "    shift_ricker = np.roll(ricker,i)\n",
    "    shift_ricker_tensor = numpy2tensor(shift_ricker)\n",
    "    shift_ricker_tensor = shift_ricker_tensor.reshape(1,-1,1)\n",
    "    loss = loss_fn.forward(ricker_tensor,shift_ricker_tensor)\n",
    "    L1_loss.append(loss.detach().numpy())\n",
    "\n",
    "\n",
    "plt.figure(figsize = (10,6))\n",
    "plt.plot(x_shift,normalize(L1_loss)         ,c='gray'   ,label=\"L1\")\n",
    "plt.legend(fontsize=14)\n",
    "plt.xlabel(\"Shift Distance\",fontsize=14)\n",
    "plt.ylabel(\"Normalized Misfit\",fontsize=14)\n",
    "plt.tick_params(labelsize=14)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# L2\n",
    "from ADFWI.fwi.misfit import Misfit_waveform_L2\n",
    "loss_fn = Misfit_waveform_L2(dt=1)\n",
    "\n",
    "ricker_tensor = numpy2tensor(ricker)\n",
    "ricker_tensor = ricker_tensor.reshape(1,-1,1)\n",
    "L2_loss = []\n",
    "for i in x_shift:\n",
    "    shift_ricker = np.roll(ricker,i)\n",
    "    shift_ricker_tensor = numpy2tensor(shift_ricker)\n",
    "    shift_ricker_tensor = shift_ricker_tensor.reshape(1,-1,1)\n",
    "    loss = loss_fn.forward(ricker_tensor,shift_ricker_tensor)\n",
    "    L2_loss.append(loss.detach().numpy())\n",
    "\n",
    "plt.figure(figsize = (10,6))\n",
    "plt.plot(x_shift,normalize(L1_loss)         ,c='gray'   ,label=\"L1\")\n",
    "plt.plot(x_shift,normalize(L2_loss)         ,c='silver' ,label=\"L2\")\n",
    "plt.legend(fontsize=14)\n",
    "plt.xlabel(\"Shift Distance\",fontsize=14)\n",
    "plt.ylabel(\"Normalized Misfit\",fontsize=14)\n",
    "plt.tick_params(labelsize=14)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SmoothL1\n",
    "from ADFWI.fwi.misfit import Misfit_waveform_smoothL1\n",
    "loss_fn = Misfit_waveform_smoothL1(dt=1)\n",
    "\n",
    "ricker_tensor = numpy2tensor(ricker)\n",
    "ricker_tensor = ricker_tensor.reshape(1,-1,1)\n",
    "SmoothL1_loss = []\n",
    "for i in x_shift:\n",
    "    shift_ricker = np.roll(ricker,i)\n",
    "    shift_ricker_tensor = numpy2tensor(shift_ricker)\n",
    "    shift_ricker_tensor = shift_ricker_tensor.reshape(1,-1,1)\n",
    "    loss = loss_fn.forward(ricker_tensor,shift_ricker_tensor)\n",
    "    SmoothL1_loss.append(loss.detach().numpy())\n",
    "\n",
    "plt.figure(figsize = (10,6))\n",
    "plt.plot(x_shift,normalize(L1_loss)         ,c='gray'   ,label=\"L1\")\n",
    "plt.plot(x_shift,normalize(L2_loss)         ,c='silver' ,label=\"L2\")\n",
    "plt.plot(x_shift,normalize(SmoothL1_loss)   ,c='k'      ,label=\"SmoothL1\")\n",
    "plt.legend(fontsize=14)\n",
    "plt.xlabel(\"Shift Distance\",fontsize=14)\n",
    "plt.ylabel(\"Normalized Misfit\",fontsize=14)\n",
    "plt.tick_params(labelsize=14)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# StudentT\n",
    "from ADFWI.fwi.misfit import Misfit_waveform_studentT\n",
    "loss_fn = Misfit_waveform_studentT(dt=1)\n",
    "\n",
    "ricker_tensor = numpy2tensor(ricker)\n",
    "ricker_tensor = ricker_tensor.reshape(1,-1,1)\n",
    "StudentT_loss = []\n",
    "for i in x_shift:\n",
    "    shift_ricker = np.roll(ricker,i)\n",
    "    shift_ricker_tensor = numpy2tensor(shift_ricker)\n",
    "    shift_ricker_tensor = shift_ricker_tensor.reshape(1,-1,1)\n",
    "    loss = loss_fn.forward(ricker_tensor,shift_ricker_tensor)\n",
    "    StudentT_loss.append(loss.detach().numpy())\n",
    "\n",
    "plt.figure(figsize = (10,6))\n",
    "plt.plot(x_shift,normalize(L1_loss)         ,c='gray'   ,label=\"L1\")\n",
    "plt.plot(x_shift,normalize(L2_loss)         ,c='silver' ,label=\"L2\")\n",
    "plt.plot(x_shift,normalize(SmoothL1_loss)   ,c='k'      ,label=\"SmoothL1\")\n",
    "plt.plot(x_shift,normalize(StudentT_loss)   ,c='#00BBF9',label=\"StudentT\")\n",
    "plt.legend(fontsize=14)\n",
    "plt.xlabel(\"Shift Distance\",fontsize=14)\n",
    "plt.ylabel(\"Normalized Misfit\",fontsize=14)\n",
    "plt.tick_params(labelsize=14)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Global Correlation\n",
    "from ADFWI.fwi.misfit import Misfit_global_correlation\n",
    "loss_fn = Misfit_global_correlation(dt=1)\n",
    "\n",
    "ricker_tensor = numpy2tensor(ricker)\n",
    "ricker_tensor = ricker_tensor.reshape(1,-1,1)\n",
    "GC_loss = []\n",
    "for i in x_shift:\n",
    "    shift_ricker = np.roll(ricker,i)\n",
    "    shift_ricker_tensor = numpy2tensor(shift_ricker)\n",
    "    shift_ricker_tensor = shift_ricker_tensor.reshape(1,-1,1)\n",
    "    loss = loss_fn.forward(ricker_tensor,shift_ricker_tensor)\n",
    "    GC_loss.append(loss.detach().numpy())\n",
    "\n",
    "plt.figure(figsize = (10,6))\n",
    "plt.plot(x_shift,normalize(L1_loss)         ,c='gray'   ,label=\"L1\")\n",
    "plt.plot(x_shift,normalize(L2_loss)         ,c='silver' ,label=\"L2\")\n",
    "plt.plot(x_shift,normalize(SmoothL1_loss)   ,c='k'      ,label=\"SmoothL1\")\n",
    "plt.plot(x_shift,normalize(StudentT_loss)   ,c='#00BBF9',label=\"StudentT\")\n",
    "plt.plot(x_shift,normalize(GC_loss)         ,c='#22A884',label=\"Global Correlation\")\n",
    "\n",
    "plt.legend(fontsize=14)\n",
    "plt.xlabel(\"Shift Distance\",fontsize=14)\n",
    "plt.ylabel(\"Normalized Misfit\",fontsize=14)\n",
    "plt.tick_params(labelsize=14)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Envelope Correlation\n",
    "from ADFWI.fwi.misfit import Misfit_envelope\n",
    "loss_fn = Misfit_envelope(dt=1,p=1,instaneous_phase=False)\n",
    "\n",
    "ricker_tensor = numpy2tensor(ricker)\n",
    "ricker_tensor = ricker_tensor.reshape(1,-1,1)\n",
    "Envelope_loss = []\n",
    "for i in x_shift:\n",
    "    shift_ricker = np.roll(ricker,i)\n",
    "    shift_ricker_tensor = numpy2tensor(shift_ricker)\n",
    "    shift_ricker_tensor = shift_ricker_tensor.reshape(1,-1,1)\n",
    "    loss = loss_fn.forward(ricker_tensor,shift_ricker_tensor)\n",
    "    Envelope_loss.append(loss.detach().numpy())\n",
    "\n",
    "plt.figure(figsize = (10,6))\n",
    "plt.plot(x_shift,normalize(L1_loss)         ,c='gray'   ,label=\"L1\")\n",
    "plt.plot(x_shift,normalize(L2_loss)         ,c='silver' ,label=\"L2\")\n",
    "plt.plot(x_shift,normalize(SmoothL1_loss)   ,c='k'      ,label=\"SmoothL1\")\n",
    "plt.plot(x_shift,normalize(StudentT_loss)   ,c='#00BBF9',label=\"StudentT\")\n",
    "plt.plot(x_shift,normalize(GC_loss)         ,c='#22A884',label=\"Global Correlation\")\n",
    "plt.plot(x_shift,normalize(Envelope_loss)   ,c='#EB1C4A',label=\"Envelope\")\n",
    "\n",
    "plt.legend(fontsize=14)\n",
    "plt.xlabel(\"Shift Distance\",fontsize=14)\n",
    "plt.ylabel(\"Normalized Misfit\",fontsize=14)\n",
    "plt.tick_params(labelsize=14)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Softdtw\n",
    "from ADFWI.fwi.misfit import Misfit_sdtw\n",
    "loss_fn = Misfit_sdtw(dt=1,gamma=5)\n",
    "\n",
    "ricker_tensor = numpy2tensor(ricker).to(\"cuda:0\")\n",
    "ricker_tensor = ricker_tensor.reshape(1,-1,1)\n",
    "Sdtw_loss = []\n",
    "for i in x_shift:\n",
    "    shift_ricker = np.roll(ricker,i)\n",
    "    shift_ricker_tensor = numpy2tensor(shift_ricker).to(\"cuda:0\")\n",
    "    shift_ricker_tensor = shift_ricker_tensor.reshape(1,-1,1)\n",
    "    loss = loss_fn.forward(ricker_tensor[:,::2,:],shift_ricker_tensor[:,::2,:])\n",
    "    Sdtw_loss.append(loss.cpu().detach().numpy())\n",
    "\n",
    "plt.figure(figsize = (10,6))\n",
    "plt.plot(x_shift,normalize(L1_loss)         ,c='gray'   ,label=\"L1\")\n",
    "plt.plot(x_shift,normalize(L2_loss)         ,c='silver' ,label=\"L2\")\n",
    "plt.plot(x_shift,normalize(SmoothL1_loss)   ,c='k'      ,label=\"SmoothL1\")\n",
    "plt.plot(x_shift,normalize(StudentT_loss)   ,c='#00BBF9',label=\"StudentT\")\n",
    "plt.plot(x_shift,normalize(GC_loss)         ,c='#22A884',label=\"Global Correlation\")\n",
    "plt.plot(x_shift,normalize(Envelope_loss)   ,c='#EB1C4A',label=\"Envelope\")\n",
    "plt.plot(x_shift,normalize(Sdtw_loss)       ,c='#1E90FF',label=\"Softdtw\")\n",
    "\n",
    "plt.legend(fontsize=14)\n",
    "plt.xlabel(\"Shift Distance\",fontsize=14)\n",
    "plt.ylabel(\"Normalized Misfit\",fontsize=14)\n",
    "plt.tick_params(labelsize=14)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wasserstein\n",
    "from ADFWI.fwi.misfit import Misfit_Wasserstein1,Misfit_wasserstein_1d,Misfit_wasserstein_sinkhorn\n",
    "# loss_fn = Misfit_Wasserstein1(trans_type='linear',dt=1,theta=1)\n",
    "# loss_fn = Misfit_wasserstein_1d(dt=dt,p=1)\n",
    "loss_fn = Misfit_wasserstein_sinkhorn(dt=dt,sparse_sampling=1,p=1,blur=1e-2,scaling=0.92)\n",
    "\n",
    "ricker_tensor = numpy2tensor(ricker).to(\"cuda:0\")\n",
    "ricker_tensor = ricker_tensor.reshape(1,-1,1)\n",
    "Wasserstein1_loss = []\n",
    "for i in x_shift:\n",
    "    shift_ricker = np.roll(ricker,i)\n",
    "    shift_ricker_tensor = numpy2tensor(shift_ricker).to(\"cuda:0\")\n",
    "    shift_ricker_tensor = shift_ricker_tensor.reshape(1,-1,1)\n",
    "    loss = loss_fn.forward(ricker_tensor,shift_ricker_tensor)\n",
    "    Wasserstein1_loss.append(loss.cpu().detach().numpy())\n",
    "\n",
    "plt.figure(figsize = (10,6))\n",
    "plt.plot(x_shift,normalize(L1_loss)         ,c='gray'   ,label=\"L1\")\n",
    "plt.plot(x_shift,normalize(L2_loss)         ,c='silver' ,label=\"L2\")\n",
    "plt.plot(x_shift,normalize(SmoothL1_loss)   ,c='k'      ,label=\"SmoothL1\")\n",
    "plt.plot(x_shift,normalize(StudentT_loss)   ,c='#00BBF9',label=\"StudentT\")\n",
    "plt.plot(x_shift,normalize(GC_loss)         ,c='#22A884',label=\"Global Correlation\")\n",
    "plt.plot(x_shift,normalize(Envelope_loss)   ,c='#EB1C4A',label=\"Envelope\")\n",
    "plt.plot(x_shift,normalize(Sdtw_loss)       ,c='#1E90FF',label=\"Softdtw\")\n",
    "plt.plot(x_shift,normalize(Wasserstein1_loss),c='#FF8C00',label=\"Wasserstein1\")\n",
    "\n",
    "plt.legend(fontsize=14)\n",
    "plt.xlabel(\"Shift Distance\",fontsize=14)\n",
    "plt.ylabel(\"Normalized Misfit\",fontsize=14)\n",
    "plt.tick_params(labelsize=14)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Figure For Article"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import gridspec\n",
    "\n",
    "fig = plt.figure(figsize = (12,10))\n",
    "gs  = gridspec.GridSpec(7,2,hspace=0.8)\n",
    "ax0 = fig.add_subplot(gs[0:2,0])\n",
    "ax1 = fig.add_subplot(gs[0:2,1])\n",
    "ax2 = fig.add_subplot(gs[2:,:])\n",
    "\n",
    "################################\n",
    "nt = 1000\n",
    "dt = 0.003\n",
    "f0 = 6\n",
    "tlist,ricker = wavelet(nt,dt,f0,amp0=1,t0=0.5)\n",
    "ricker       = np.roll(ricker,330)\n",
    "shift_ricker = np.roll(ricker,0)\n",
    "\n",
    "ax0.plot(tlist,ricker       ,c='k',linestyle=\"-\"    ,label=\"Ricker\")\n",
    "ax0.plot(tlist,shift_ricker ,c='r',linestyle=\"--\"   ,label=\"Shifted Ricker\")\n",
    "ax0.grid()\n",
    "ax0.legend(fontsize=12,loc='upper right')\n",
    "ax0.set_xlabel(\"Times (s)\",fontsize=13)\n",
    "ax0.set_ylabel(\"Amplitude\",fontsize=13)\n",
    "ax0.tick_params(labelsize=13)\n",
    "\n",
    "################################\n",
    "nt = 1000\n",
    "dt = 0.003\n",
    "f0 = 6\n",
    "tlist,ricker = wavelet(nt,dt,f0,amp0=1,t0=0.5)\n",
    "ricker       = np.roll(ricker,330)\n",
    "shift_ricker = np.roll(ricker,int(0.15//dt))\n",
    "\n",
    "ax1.plot(tlist,ricker       ,c='k',linestyle=\"-\"    ,label=\"Ricker\")\n",
    "ax1.plot(tlist,shift_ricker ,c='r',linestyle=\"--\"   ,label=\"Shifted Ricker\")\n",
    "ax1.grid()\n",
    "ax1.legend(fontsize=12,loc='upper right')\n",
    "ax1.set_xlabel(\"Times (s)\",fontsize=13)\n",
    "ax1.set_ylabel(\"Amplitude\",fontsize=13)\n",
    "ax1.tick_params(labelsize=13)\n",
    "\n",
    "#################################\n",
    "\n",
    "ax2.plot(x_shift*dt,normalize(L1_loss)          ,c='gray'   ,label=\"L1\"                 ,linestyle='-')\n",
    "ax2.plot(x_shift*dt,normalize(L2_loss)          ,c='silver' ,label=\"L2\"                 ,linestyle='-.')\n",
    "ax2.plot(x_shift*dt,normalize(SmoothL1_loss)    ,c='k'      ,label=\"SmoothL1\"           ,linestyle='--')\n",
    "ax2.plot(x_shift*dt,normalize(StudentT_loss)    ,c='#00BBF9',label=\"StudentT\"           ,linestyle='-')\n",
    "ax2.plot(x_shift*dt,normalize(GC_loss)          ,c='#22A884',label=\"Global Correlation\" ,linestyle='-.')\n",
    "ax2.plot(x_shift*dt,normalize(Envelope_loss)    ,c='#EB1C4A',label=\"Envelope\"           ,linestyle='--')\n",
    "ax2.plot(x_shift*dt,normalize(Sdtw_loss)        ,c='#1E90FF',label=\"Softdtw\"            ,linestyle='-')\n",
    "ax2.plot(x_shift*dt,normalize(Wasserstein1_loss),c='#FF8C00',label=\"Wasserstein1\"       ,linestyle='-.')\n",
    "\n",
    "fill_shift = np.arange(35,65)\n",
    "ax2.fill_between(fill_shift*dt, 0, 1,color='green', alpha=0.1,transform=plt.gca().get_xaxis_transform())\n",
    "ax2.fill_between(-fill_shift*dt, 0, 1,color='green', alpha=0.1,transform=plt.gca().get_xaxis_transform())\n",
    "\n",
    "bbox_args = dict(boxstyle=\"round\", fc=\"lightblue\",edgecolor=\"lightblue\")\n",
    "ax2.annotate('Global minimum', \n",
    "             xy=(0, 0),xytext=(0.1, 0.07),\n",
    "             fontsize=15,bbox=bbox_args,\n",
    "             arrowprops=dict(\n",
    "                            facecolor='lightblue', \n",
    "                            alpha=0.6,\n",
    "                            arrowstyle='fancy',\n",
    "                            # arrowstyle='-|>',\n",
    "                            connectionstyle='arc3,rad=-0.4',#有多个参数可选\n",
    "                            color='k',\n",
    "                            ))\n",
    "ax2.annotate('Local minimum', \n",
    "             xy=(0.15, 0.55),xytext=(0.2, 0.3),\n",
    "             fontsize=15,bbox=bbox_args,\n",
    "             arrowprops=dict(\n",
    "                            facecolor='lightblue', \n",
    "                            alpha=0.6,\n",
    "                            arrowstyle='fancy',\n",
    "                            # arrowstyle='-|>',\n",
    "                            connectionstyle='arc3,rad=-0.4',#有多个参数可选\n",
    "                            color='k',\n",
    "                            ))\n",
    "\n",
    "ax2.legend(fontsize=13)\n",
    "ax2.set_xlabel(\"Time shift (s)\",fontsize=16)\n",
    "ax2.set_ylabel(\"Normalized Misfit\",fontsize=16)\n",
    "ax2.tick_params(labelsize=16)\n",
    "plt.savefig(os.path.join(project_path,\"misfit_shift.png\"),dpi=300,bbox_inches=\"tight\")\n",
    "plt.show()"
   ]
  },
  {
   "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
}