{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "76315cb9-11f7-42b2-b393-692511f6d67a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# ces_01\n",
    "%reset -f\n",
    "\n",
    "from math import exp, pi, sin, cos, tan, atan, sqrt, floor, log10, nan, log   # math functions\n",
    "import numpy as np\n",
    "import scipy.linalg as la             # linear algebra functions\n",
    "import scipy.interpolate as interp    # linear algebra functions\n",
    "from scipy.integrate import solve_ivp\n",
    "import matplotlib\n",
    "from matplotlib import pyplot as plt  # general plotting \n",
    "from matplotlib.colors import ListedColormap\n",
    "import time;\n",
    "\n",
    "# function to make a numpy N-by-1 column vector\n",
    "# c=eda_cvec(v1, v2, ...) from a list of several\n",
    "# array-like entities v1, v2, including a number\n",
    "# a list of numbers, a tuple of numbers, an N-by-0 np array\n",
    "# and a N-by-1 np array. The function\n",
    "# also insures that, if v is an np array, that\n",
    "# c is a copy, as contrasted to a view, of it\n",
    "# It promotes integers to floats, and integers\n",
    "# and floats to complex, by context.\n",
    "# This version concatenates many argments,\n",
    "# whereas c=eda_cvec1(v1) takes just one argiment.\n",
    "# I recommend always using eda_cvec(v1, v2, ...)\n",
    "def gda_cvec(*argv):\n",
    "    t = int;\n",
    "    Nt = 0;\n",
    "    for a in argv:\n",
    "        v = gda_cvec1(a);\n",
    "        N,M = np.shape(v);\n",
    "        Nt = Nt + N;\n",
    "        if( N==0 ):\n",
    "            continue; # skip vector of zero length\n",
    "        if (t==int) and isinstance(v[0,0],float):\n",
    "            t=float;\n",
    "        elif isinstance(v[0,0],complex):\n",
    "            t=complex;\n",
    "    w = np.zeros((Nt,1),dtype=t);\n",
    "    Nt = 0;\n",
    "    for a in argv:\n",
    "        v = gda_cvec1(a);\n",
    "        N,M = np.shape(v);\n",
    "        w[Nt:Nt+N,0:1] = v;  # patch 20230418 was #w[Nt:Nt+N,0] = v[0:N,0];\n",
    "        Nt = Nt + N;\n",
    "    return w;\n",
    "\n",
    "# function to make a numpy N-by-1 column vector\n",
    "# c=gda_cvec1(v) from entity v that is array-like,\n",
    "# including a number, a list of numbers, a tuple\n",
    "# of numbers, an N-by-0 np array and a N-by1 np array.\n",
    "# It promotes integers to floats, and integers\n",
    "# and floats to complex, by context. The function\n",
    "# also insures that, if v is an np array, that\n",
    "# c is a copy, as contrasted to a view, of it.\n",
    "# This version takes just one input argment.\n",
    "# whereas c=gda_cvec(v1,v2,...) concatenates\n",
    "# many argiments.\n",
    "def gda_cvec1(v):\n",
    "    if isinstance(v, int) or isinstance(v, np.int32):\n",
    "        w = np.zeros((1,1),dtype=int);\n",
    "        w[0,0] = v;\n",
    "        return w;\n",
    "    elif isinstance(v, float):\n",
    "        w = np.zeros((1,1),dtype=float);\n",
    "        w[0,0] = v;\n",
    "        return w;\n",
    "    elif isinstance(v, complex):\n",
    "        w = np.zeros((1,1),dtype=complex);\n",
    "        w[0,0] = v;\n",
    "        return w;\n",
    "    elif isinstance(v, np.ndarray):\n",
    "        s = np.shape(v);\n",
    "        if len(s) == 1:\n",
    "            return np.copy(np.reshape(v,(s[0],1)));\n",
    "        else:\n",
    "            [r,c]=s;\n",
    "            if( c==1 ):\n",
    "                return(np.copy(v));\n",
    "            elif(r==1):\n",
    "                return(np.copy(v.T));\n",
    "            else:\n",
    "                raise TypeError(\"gda_cvec: %d by %d ndarray not allowed\" % (r, c));\n",
    "    elif isinstance(v, list):\n",
    "        r = len(v);\n",
    "        t = int;\n",
    "        for vi in v:\n",
    "            if isinstance(vi,int) or isinstance(vi, np.int32): #patch v->vi 20230418\n",
    "                pass;\n",
    "            elif isinstance(vi,float):\n",
    "                t=float;\n",
    "            elif isinstance(vi,complex):\n",
    "                t=complex;\n",
    "                break;\n",
    "            else:\n",
    "                raise TypeError(\"gda_cvec: list contains unsupported type %s\" % type(vi));\n",
    "        w = np.zeros((r,1),dtype=t);\n",
    "        w[:,0] = np.array(v); # patch v -> np.array(v)\n",
    "        return w;\n",
    "    elif isinstance(v, tuple):\n",
    "        r = len(v);\n",
    "        t = int;\n",
    "        for vi in v:\n",
    "            if isinstance(vi,int) or isinstance(vi, np.int32): #patch v->vi 20230418\n",
    "                pass;\n",
    "            elif isinstance(vi,float):\n",
    "                t=float;\n",
    "            elif isinstance(vi,complex):\n",
    "                t=complex;\n",
    "                break;\n",
    "            else:\n",
    "                raise TypeError(\"gda_cvec: tuple contains unsupported type %s\" % type(vi));\n",
    "        w = np.zeros((r,1),dtype=t);\n",
    "        w[:,0] = np.array(list(v)); # patch v -> np.array(list(v));\n",
    "        return w;\n",
    "    else:\n",
    "        raise TypeError(\"gda_cvec: %s not supported\" % type(v));\n",
    "\n",
    "        \n",
    "# gda_draw function makes a \"pictorial matrix equation\"\n",
    "# arguments are vectors, matrices and strings\n",
    "# which are plotted in the order that the appear\n",
    "# except that strings starting with 'title ' are plotted\n",
    "# under the subseqeunt matrix or vector\n",
    "# always returns a status of 1\n",
    "def gda_draw(*argv):\n",
    "    DOCOLOR=True;\n",
    "    if( DOCOLOR ):\n",
    "        bwcmap = cm.get_cmap('jet');\n",
    "    else:\n",
    "        bw = np.zeros((256,4));\n",
    "        v = 0.9*(256 - np.linspace( 0, 255, 256 ))/255;\n",
    "        bw[:,0] = v;\n",
    "        bw[:,1] = v;\n",
    "        bw[:,2] = v;\n",
    "        bw[:,3] = np.ones(256);\n",
    "        bwcmap = ListedColormap(bw);\n",
    "    # size of plot\n",
    "    W = 16;\n",
    "    H = 4;\n",
    "    fig1 = plt.figure(1);\n",
    "    # figsize width and height in inches\n",
    "    fig1.set_size_inches(W,H);\n",
    "    ax1 = plt.subplot(1,1,1);\n",
    "    plt.axis([0, W, -H/2, H/2]);\n",
    "    plt.axis('off');\n",
    "    LM = W/6;    # matrix width and heoght\n",
    "    LV = W/40;   # vector width\n",
    "    FS = 0.12;    # character width\n",
    "    TO = 0.4;    # title vertical offset\n",
    "    SP = 0.2;    # space between objects\n",
    "    LS = 0.2;    # leading space\n",
    "    p = LS; # starting x-position\n",
    "    istitle=0; # flags presence of a title\n",
    "    for a in argv:\n",
    "        if isinstance(a,np.ndarray):\n",
    "            sh = np.shape(a);\n",
    "            if len(sh) == 1:  # conversion to nx1 array\n",
    "                n = sh[0];\n",
    "                m = 1;\n",
    "                ap = a;\n",
    "                a = np.zeros((n,1));\n",
    "                a[:,0] = ap;\n",
    "            else:\n",
    "                n = sh[0];\n",
    "                m = sh[1];\n",
    "            if m==1:\n",
    "                pold=p;\n",
    "                left=p;\n",
    "                right=p+LV;\n",
    "                bottom=-LM/2;\n",
    "                top=LM/2;\n",
    "                plt.imshow( a, cmap=bwcmap, vmin=np.min(a), vmax=np.max(a), extent=(left,right,bottom,top) );\n",
    "                p = p+LV;\n",
    "                pm = (p+pold)/2;\n",
    "                if istitle:\n",
    "                    plt.text(pm,-(LM/2)-TO,titlestr,horizontalalignment='center');\n",
    "                    istitle=0;\n",
    "                p = p+SP;\n",
    "            else:\n",
    "                pold=p;\n",
    "                left=p;\n",
    "                right=p+LM;\n",
    "                bottom=-LM/2;\n",
    "                top=LM/2;\n",
    "                plt.imshow( a, cmap=bwcmap, vmin=np.min(a), vmax=np.max(a), extent=(left,right,bottom,top) );\n",
    "                p = p+LM;\n",
    "                pm = (p+pold)/2;\n",
    "                if istitle:\n",
    "                    plt.text(pm,-(LM/2)-TO,titlestr,horizontalalignment='center');\n",
    "                    istitle=0;\n",
    "                p = p+SP;\n",
    "        elif isinstance(a,str):\n",
    "            ns = len(a);\n",
    "            istitle=0;\n",
    "            if( ns>=6 ):\n",
    "                if 'title ' in a[0:6]:\n",
    "                    istitle=1;\n",
    "                    titlestr=a[6:];\n",
    "            if( istitle != 1):\n",
    "                plt.text(p,0,a);\n",
    "                p = p + ns*FS + SP;\n",
    "    plt.show();\n",
    "    return 1;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9c1dc601-cee5-42c0-af73-a754f3e0c2fb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "51077 data read\n",
      "   but truncated to 51076\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 16768 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# initializastion part\n",
    "\n",
    "# read data\n",
    "D = np.genfromtxt(\"MYD.txt\", delimiter=\"\\t\");\n",
    "N, i = np.shape(D);\n",
    "print(\"%d data read\" % (N));\n",
    "if( N%2 ):\n",
    "    N=N-1;\n",
    "    print(\"   but truncated to %d\" % (N));\n",
    "t1 = np.copy(D[0:N,0:1]);\n",
    "t1min=t1[0,0];\n",
    "t1max=t1[N-1,0];\n",
    "d1 = np.copy(D[0:N,1:2])*1e10;\n",
    "d1absmax = np.max(np.abs(d1));\n",
    "t2 = np.copy(D[0:N,2:3]);\n",
    "t2min=t2[0,0];\n",
    "t2max=t2[N-1,0];\n",
    "d2 = np.copy(D[0:N,3:4])*1e10;\n",
    "d2absmax = np.max(np.abs(d2));\n",
    "Dt = t1[1,0]-t1[0,0]; # assummed same for t1 and t2\n",
    "\n",
    "# w axis\n",
    "Nf = int(N/2)+1;\n",
    "Nw = Nf;\n",
    "fmax = 1.0/(2.0*Dt);\n",
    "wmax = (2.0*pi)*fmax;\n",
    "fpos = gda_cvec(np.linspace(0.0,fmax,Nw));\n",
    "wpos = 2.0*pi*fpos;\n",
    "f = np.concatenate( (fpos,-np.flipud(fpos[1:Nf-1,0:1])), axis=0 );\n",
    "w = 2.0*pi*f;\n",
    "Df = f[1,0]-f[0,0];\n",
    "Dw = 2.0*pi*Df;\n",
    "\n",
    "# initialize dispersion curve\n",
    "fposlist = gda_cvec( 0.0, 0.05, 0.2, 25.0, 50.0);\n",
    "sposlist = gda_cvec( 0.2, 0.2,  0.2, 0.2, 0.2 );\n",
    "cs = interp.CubicSpline(fposlist.ravel(), sposlist.ravel(), bc_type='natural', extrapolate=True);\n",
    "sposint = gda_cvec(cs(fpos));\n",
    "fposint = fpos;\n",
    "\n",
    "fplot=0.1;\n",
    "fig1 = plt.figure();\n",
    "ax1 = plt.subplot(1,1,1);\n",
    "plt.axis( [0.0, fplot, 0.0, 0.5 ] );\n",
    "plt.plot(fposint,sposint,'k-');\n",
    "plt.plot(fposlist,sposlist,'k*');\n",
    "plt.xlabel('frequency f (Hz)');\n",
    "plt.ylabel('slowness (s/m)');\n",
    "plt.show();\n",
    "\n",
    "d1a = d1;\n",
    "elast = d1a-d2;\n",
    "Elast = np.matmul(elast.T,elast);\n",
    "d1alast = d1a;\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "97b662ea-06d0-4948-973d-49f7718e5eee",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 143240 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 100736 characters]",
      "text/plain": [
       "<Figure size 800x800 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: 1.9315e+04\n"
     ]
    }
   ],
   "source": [
    "# MAIN PART\n",
    "\n",
    "def advance( d, N, Dt, Drange, fposlist, sposlist):\n",
    "\n",
    "    # w axis\n",
    "    Nf = int(N/2)+1;\n",
    "    Nw = Nf;\n",
    "    fmax = 1.0/(2.0*Dt);\n",
    "    wmax = (2.0*pi)*fmax;\n",
    "    fpos = gda_cvec(np.linspace(0.0,fmax,Nw));\n",
    "    wpos = 2.0*pi*fpos;\n",
    "    f = np.concatenate( (fpos,-np.flipud(fpos[1:Nf-1,0:1])), axis=0 );\n",
    "    w = 2.0*pi*f;\n",
    "    Df = f[1,0]-f[0,0];\n",
    "    Dw = 2.0*pi*Df;\n",
    "    \n",
    "    # interpolate slowness\n",
    "    cs = interp.CubicSpline(fposlist.ravel(), sposlist.ravel(), bc_type='natural', extrapolate=True);\n",
    "    spos = gda_cvec(cs(fpos));\n",
    "    s = np.concatenate( (spos, np.flipud(spos[1:Nf-1,0:1])), axis=0 );\n",
    "    \n",
    "    # advance d1\n",
    "    dt = np.fft.fft( d, axis=0 );\n",
    "    dt = np.multiply( dt , np.exp( -complex(0.0,1.0) * Drange * np.multiply(w,s) ) );\n",
    "    da = np.real( np.fft.ifft(dt,axis=0) );\n",
    "    \n",
    "    return da, fpos, spos;\n",
    "\n",
    "#  interactive graphics from withing the Jupyter Notebook is a bit dicey\n",
    "# if METHOD=1 fails, try METHOD=2 (or vice-versa)\n",
    "METHOD=0;\n",
    "# get current backend\n",
    "if( METHOD==1):\n",
    "    bkend = matplotlib.get_backend();\n",
    "    \n",
    "#  use interactive plotting\n",
    "if( METHOD==1 ):\n",
    "    matplotlib.use('Qt5Agg');\n",
    "else:\n",
    "    %matplotlib qt5\n",
    "       \n",
    "fig1 = plt.figure(1,figsize=(8,8));\n",
    "ax1 = plt.subplot(3,1,1);\n",
    "ax2 = plt.subplot(3,1,2);\n",
    "ax3 = plt.subplot(3,1,3);\n",
    "\n",
    "for ii in range(10):\n",
    "\n",
    "    ax2.cla();\n",
    "    ax2.axis( [t2min, t2max, -1.1*d2absmax, 1.1*d2absmax ] );\n",
    "    ax2.plot(t2,d2,'k-',lw=4);\n",
    "    ax2.plot(t1,d1alast,'r-',lw=2);\n",
    "    ax2.set_xlabel('t (s)');\n",
    "    ax2.set_ylabel('d2(t)');\n",
    "    \n",
    "    ax3.cla();\n",
    "    ax3.axis( [t2min, t2max, -1.1*d2absmax, 1.1*d2absmax ] );\n",
    "    ax3.plot(t2,elast,'r-',lw=2);\n",
    "    ax3.set_xlabel('time t (s)');\n",
    "    ax3.set_ylabel('e(t)');\n",
    "    ax3.set_title(\"Error: %.4e\" % (Elast));\n",
    "\n",
    "    ax1.cla();\n",
    "    ax1.axis( [0.0, fplot, 0.0, 0.5 ] );\n",
    "    ax1.plot(fposint,sposint,'r-');\n",
    "    ax1.plot(fposlist,sposlist,'r*');\n",
    "    ax1.set_xlabel('frequency f (Hz)');\n",
    "    ax1.set_ylabel('slowness (s/m)');\n",
    "    ax1.set_title('Pick >=3 points, click just below top to finish; <3 terminates');\n",
    "    plt.draw();  # must draw to update plot\n",
    "\n",
    "    MAXPTS = 20;\n",
    "    fposnew=np.zeros((MAXPTS,1));\n",
    "    sposnew=np.zeros((MAXPTS,1));\n",
    "    M=0;\n",
    "    stop=0;\n",
    "    for i in range(MAXPTS):\n",
    "        p = fig1.ginput(1);  # returns list of (x,y) tupples\n",
    "        if( len(p)==0 ):\n",
    "            continue;\n",
    "        if( ( p[0][1] > 0.45 ) ):\n",
    "            if( M>=3 ):\n",
    "                stop=0;\n",
    "            else:\n",
    "                stop=1;\n",
    "            break;\n",
    "        fposnew[M,0] = p[0][0];\n",
    "        sposnew[M,0] = p[0][1];\n",
    "        ax1.plot(p[0][0],p[0][1],'ko');\n",
    "        plt.draw();  # must draw to update plot\n",
    "        M=M+1;\n",
    "    if( stop ):\n",
    "        break;\n",
    "    fposnew = fposnew[0:M,0:1];\n",
    "    sposnew = sposnew[0:M,0:1];\n",
    "\n",
    "    k = np.argsort(fposnew.ravel());\n",
    "    fposnew = gda_cvec( fposnew[k,0] );\n",
    "    sposnew = gda_cvec( sposnew[k,0] );\n",
    "\n",
    "    range1 = 2.8731e+03;\n",
    "    range2 = 2.8462e+03;\n",
    "    d1a, fposnewint, sposnewint = advance( d1, N, Dt, range2-range1, fposnew, sposnew);\n",
    "    #ax1.plot(fposnewint,sposnewint,'k-');\n",
    "    #ax1.plot(fposnew,sposnew,'ko');\n",
    "    \n",
    "    #ax2.plot(t1,d1a,'k-',lw=4);\n",
    "\n",
    "    e = d1a-d2;\n",
    "    E = np.matmul(e.T,e);\n",
    "\n",
    "    plt.draw();\n",
    "\n",
    "    # update lists\n",
    "    sposlist = np.copy(sposnew)\n",
    "    fposlist = np.copy(fposnew);\n",
    "    sposint = np.copy(sposnewint);\n",
    "    fposint = np.copy(fposnewint);\n",
    "    elast=e;\n",
    "    Elast=E;\n",
    "    d1alast = d1a;\n",
    "\n",
    "# switch back to standard plotting\n",
    "if( METHOD==1 ):\n",
    "    matplotlib.use(bkend);\n",
    "else:\n",
    "    %matplotlib inline\n",
    "    \n",
    "fig2 = plt.figure(2,figsize=(8,8));\n",
    "ax1 = plt.subplot(3,1,1);\n",
    "ax2 = plt.subplot(3,1,2);\n",
    "ax3 = plt.subplot(3,1,3);\n",
    "\n",
    "ax2.axis( [t2min, t2max, -1.1*d2absmax, 1.1*d2absmax ] );\n",
    "ax2.plot(t2,d2,'k-');\n",
    "ax2.plot(t2,d1a,'r-');\n",
    "ax2.set_xlabel('t (s)');\n",
    "ax2.set_ylabel('d2(t)');\n",
    "\n",
    "ax1.axis( [0.0, fplot, 0.0, 0.5 ] );\n",
    "ax1.plot(fposint,sposint,'r-');\n",
    "ax1.plot(fposlist,sposlist,'r*');\n",
    "ax1.set_xlabel('frequency f (Hz)');\n",
    "ax1.set_ylabel('slowness (s/m)');\n",
    "\n",
    "ax3.axis( [t2min, t2max, -1.1*d2absmax, 1.1*d2absmax ] );\n",
    "ax3.plot(t2,elast,'k-',lw=4);\n",
    "ax3.set_xlabel('time t (s)');\n",
    "ax3.set_ylabel('e(t)');\n",
    "    \n",
    "plt.show();\n",
    "print(\"Error: %.4e\" % (E));\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d487f23-14c4-4e53-bba6-5b3d4d606bf4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "958ab3cf-7b45-4844-9c89-0068eefe33fa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
