{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 28,
   "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",
    "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",
    "import matplotlib.cm as cm\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 = matplotlib.colormaps['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;\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 21972 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 27120 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 27004 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 111796 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 17188 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 17856 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 16156 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# tunable parameters\n",
    "S1 = -0.05;    # the maximum slowness perturbation of the Gaussian anomaly\n",
    "xc = 75;      # x-position of the center of the anomaly\n",
    "yc = 50;      # y-position of the center of the anomaly\n",
    "sxx=5.0;     # length scale of anomaly in x-direction\n",
    "syy=5.0;      # length scale of anomaly in y-direction\n",
    "tilt = 0.0;  # tilt in degrees (0=no tilt)\n",
    "th=10.0;     # the angle of the rays (0=no tilt)\n",
    "\n",
    "# x-axis\n",
    "Nx=101;\n",
    "Dx = 1.0;\n",
    "xmax = Dx*(Nx-1);\n",
    "x = gda_cvec(np.linspace(0.0,xmax,Nx));\n",
    "\n",
    "# y-axis\n",
    "Ny=101;\n",
    "Dy = 1.0;\n",
    "ymax = Dy*(Ny-1);\n",
    "y = gda_cvec(np.linspace(0.0,ymax,Ny));\n",
    "\n",
    "# slowness\n",
    "So = 1.0;\n",
    "S = np.zeros((Ny,Nx));    # slowness\n",
    "gradS = np.zeros((2,Ny,Nx));   # grad ;\n",
    "\n",
    "# rotation\n",
    "s=sin((pi/180.0)*tilt);\n",
    "c=cos((pi/180.0)*tilt);\n",
    "T = np.array( [ [c, -s], [s, c] ] );\n",
    "COV = np.array( [ [sxx**2, 0.0], [0.0, syy**2] ] );\n",
    "C = np.matmul( T, np.matmul(COV,T.T) );\n",
    "CI = la.inv(C);\n",
    "c00 = CI[0,0];\n",
    "c11 = CI[1,1];\n",
    "c01 = CI[0,1];\n",
    "\n",
    "# Gaussian function\n",
    "# exp( -0.5 (c00*Dx**2 + c11*Dy**2 + 2c12*Dx*Dy ) ); \n",
    "def Sfun(x,y):\n",
    "    DX = x-xc;\n",
    "    DY = y-yc;\n",
    "    G = exp(-0.5*(c00*DX**2+c11*DY**2+2.0*c01*DX*DY) );\n",
    "    Gx = -0.5*(2.0*c00*DX+2.0*c01*DY)*G;\n",
    "    Gy = -0.5*(2.0*c11*DY+2.0*c01*DY)*G;\n",
    "    return (So + S1*G), (S1*Gx), (S1*Gy);\n",
    "\n",
    "for iy in range(Ny):\n",
    "    for ix in range(Nx):\n",
    "        myS, mySx, mySy = Sfun( x[ix,0], y[iy,0] );\n",
    "        S[iy,ix] = myS;\n",
    "        gradS[0,iy,ix] = mySx;\n",
    "        gradS[1,iy,ix] = mySy;\n",
    "        \n",
    "# used to check orientation\n",
    "if( 0 ):\n",
    "    S = np.zeros((Ny,Nx));\n",
    "    S[5,5]=1.0;\n",
    "    S[5,Nx-5]=2.0;\n",
    "    S[Ny-5,5]=3.0\n",
    "\n",
    "fig1 = plt.figure(1);\n",
    "ax1 = plt.subplot(1,1,1);\n",
    "mycm = plt.get_cmap('jet');\n",
    "plt.axis( [0.0, xmax, 0.0, ymax ] );\n",
    "plt.imshow(S,cmap=mycm,aspect='auto',extent=(0.0,xmax,ymax,0.0));\n",
    "plt.colorbar();\n",
    "plt.xlabel('x');\n",
    "plt.ylabel('y');\n",
    "plt.title('S');\n",
    "plt.show();\n",
    "\n",
    "fig2 = plt.figure(2);\n",
    "ax1 = plt.subplot(1,1,1);\n",
    "mycm = plt.get_cmap('jet');\n",
    "plt.axis( [0.0, xmax, 0.0, ymax ] );\n",
    "plt.imshow(gradS[0,:,:],cmap=mycm,aspect='auto',extent=(0.0,xmax,ymax,0.0));\n",
    "plt.colorbar();\n",
    "plt.xlabel('x');\n",
    "plt.ylabel('y');\n",
    "plt.title('gradSx');\n",
    "plt.show();\n",
    "\n",
    "fig3 = plt.figure(3);\n",
    "ax1 = plt.subplot(1,1,1);\n",
    "mycm = plt.get_cmap('jet');\n",
    "plt.axis( [0.0, xmax, 0.0, ymax ] );\n",
    "plt.imshow(gradS[1,:,:],cmap=mycm,aspect='auto',extent=(0.0,xmax,ymax,0.0));\n",
    "plt.colorbar();\n",
    "plt.xlabel('x');\n",
    "plt.ylabel('y');\n",
    "plt.title('gradSy');\n",
    "plt.show();\n",
    "\n",
    "# ray equaion (d/ds) S (dx/ds) = grad S (Menke & Abbott, 8.7.11)\n",
    "# now transform to an equation in traveltime t\n",
    "# arclength s, traveltime t, velocity v, slowness S=1/v\n",
    "# s = ct = t/S   so  ds/dt = 1/S  and dt/ds = S\n",
    "# d/ds = dt/ds d/dt = S d/dt\n",
    "# (d/ds) S (d/ds) x = grad S  becomes\n",
    "# S (d/dt) S S (d/dt) x = grad S becomes \n",
    "# (d/dt) S^2 (dx/dt) = (1/S) grad S\n",
    "\n",
    "# with definition\n",
    "# u2 = S^2 (dx/dt)\n",
    "# u3 = S^2 (dy/dt)\n",
    "# we have\n",
    "# (d/dt) U2 = (1/S) gradSx\n",
    "# (d/dt) U3 = (1/S) gradS3\n",
    "# then with the definition\n",
    "# u0 = x\n",
    "# u1 = y\n",
    "# we have\n",
    "# u2 = S^2 (dx/dt) = S^2 du0/dt\n",
    "# u3 = S^2 (dy/dt) = s^2 du1/dt\n",
    "# or\n",
    "# (d/dt) u0 = u2/S^2\n",
    "# (d/dt) u1 = u3/S^2\n",
    "\n",
    "# so\n",
    "# (d/dt) u0 = u2/S^2\n",
    "# (d/dt) u1 = u3/S^2\n",
    "# (d/dt) u2 = (1/S) gradSx\n",
    "# (d/dt) u3 = (1/S) gradSy\n",
    "\n",
    "M=4;\n",
    "\n",
    "def myfun( t, u ):\n",
    "    f = np.zeros((M,));\n",
    "    myS, mySx, mySy = Sfun( u[0], u[1] );\n",
    "    mySr = 1.0/myS;\n",
    "    myS2r = mySr**2;\n",
    "    f[0] = u[2] * myS2r;\n",
    "    f[1] = u[3] * myS2r;\n",
    "    f[2] = mySr * mySx;\n",
    "    f[3] = mySr * mySy;\n",
    "    return f;\n",
    "\n",
    "Dt = 0.05;\n",
    "t0 = 0.0;     \n",
    "tmax = 4.0 * np.max(S) * sqrt(xmax**2 + ymax**2);\n",
    "t = np.arange(t0, t0+tmax, Dt);\n",
    "NN=np.shape(t);\n",
    "N = NN[0];\n",
    "\n",
    "fig4 = plt.figure(4);\n",
    "ax1 = plt.subplot(1,1,1);\n",
    "mycm = plt.get_cmap('jet');\n",
    "plt.axis( [0.0, xmax, 0.0, ymax ] );\n",
    "plt.imshow(S,cmap=mycm,aspect='auto',extent=(0.0,xmax,ymax,0.0));\n",
    "plt.colorbar();\n",
    "plt.xlabel('x');\n",
    "plt.ylabel('y');\n",
    "plt.title('S');\n",
    "\n",
    "T = np.zeros((Ny,Nx));    # travel time\n",
    "# x=xmax traveltimes and positions\n",
    "NTTs = 0;\n",
    "TTend = np.zeros((Ny,1));\n",
    "Yend = np.zeros((Ny,1));\n",
    "for iy0 in range(0,Ny):\n",
    "    # initial conditions\n",
    "    myS, mySx, mySy = Sfun( x[ix,0], y[iy,0] );\n",
    "    u0 = np.zeros((M,));\n",
    "    ix0 = 0;\n",
    "    u0[0] = 0.0;       # x0\n",
    "    u0[1] = y[iy0]+Dx/2.0;  # y0\n",
    "    theta = (pi/180.0)*th;\n",
    "    # tangent = dx/ds = [cos(theta), sin(theta)] then mult by S to get dx/dt\n",
    "    u0[2] =  cos(theta)/myS; # dx/dt\n",
    "    u0[3] =  sin(theta)/myS; # dx/dt\n",
    "\n",
    "    sol=solve_ivp(myfun, [t0, tmax], u0, method = 'DOP853', t_eval=t);\n",
    "\n",
    "    plt.plot( sol.y[0], sol.y[1], 'w-' );\n",
    "    \n",
    "    # traveltime plot\n",
    "    Np, = np.shape(sol.t);\n",
    "    found=0;\n",
    "    for k in range(Np):\n",
    "        tp = sol.t[k];\n",
    "        xp = sol.y[0][k];\n",
    "        yp = sol.y[1][k];\n",
    "        jx = int(xp/Dx);\n",
    "        jy = int(yp/Dy);\n",
    "        if( (jx>=0) & (jx<Nx) & (jy>=0) & (jy<Ny) ):\n",
    "            T[jy,jx] = tp;\n",
    "        if( (xp>=xmax) and (k>0) and (not found) ):\n",
    "            found=1;\n",
    "            if( (yp>=0.0) & (yp<=ymax) ):\n",
    "                # linear interpolation along ray\n",
    "                tp0 = sol.t[k-1];\n",
    "                xp0 = sol.y[0][k-1];\n",
    "                yp0 = sol.y[1][k-1];\n",
    "                dxdT = (xp-xp0)/(tp-tp0); # change in x with traveltime\n",
    "                dydT = (yp-yp0)/(tp-tp0); # change in y with traveltime\n",
    "                myDT = (xmax-xp0)/dxdT;   # DT needed to go to xmax\n",
    "                TTend[NTTs,0] = tp0 + myDT; # TT at xmax\n",
    "                Yend[NTTs,0] = yp0 + dydT * myDT; # y at xmax\n",
    "                NTTs=NTTs+1;\n",
    "plt.show();\n",
    "\n",
    "# truncate to actual length\n",
    "TTend = TTend[0:NTTs,0:1];\n",
    "Yend = Yend[0:NTTs,0:1];\n",
    "   \n",
    "fig5 = plt.figure(5);\n",
    "ax1 = plt.subplot(1,1,1);\n",
    "mycm = plt.get_cmap('jet');\n",
    "plt.axis( [0.0, xmax, 0.0, ymax ] );\n",
    "plt.imshow(T,cmap=mycm,aspect='auto',extent=(0.0,xmax,ymax,0.0));\n",
    "plt.colorbar();\n",
    "plt.xlabel('x');\n",
    "plt.ylabel('y');\n",
    "plt.title('T');\n",
    "plt.show();\n",
    "\n",
    "fig6 = plt.figure(6);\n",
    "ax1 = plt.subplot(1,1,1);\n",
    "TTmin = np.min(TTend);\n",
    "TTmax = np.max(TTend);\n",
    "plt.axis( [0.0, ymax, -10.0, 10.0 ] );\n",
    "plt.xlabel('y');\n",
    "plt.ylabel('T - Tmean');\n",
    "plt.plot(Yend,TTend-np.mean(TTend),'k-', lw=2 );\n",
    "plt.show();\n",
    "\n",
    "figt = plt.figure(7);\n",
    "ax1 = plt.subplot(1,1,1);\n",
    "plt.axis( [0.0, ymax, 0, 2 ] );\n",
    "plt.xlabel('y');\n",
    "plt.ylabel('ray hits');\n",
    "for i in range(NTTs):\n",
    "    plt.plot(gda_cvec(Yend[i,0],Yend[i,0]),gda_cvec(0.0,1.0),'k-', lw=2 );\n",
    "plt.show();\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
