{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example for cmesh2d and pygmsh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pygmsh\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from ttcrpy import cmesh2d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define some function for the slowness"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f0(x, y):\n",
    "    return 1./(0.6 + 0.2*y + 0.005*x**2)\n",
    "def f1(x, y):\n",
    "    return 1./(1. + 0.2*y + 0.005*x**2)\n",
    "def f2(x, y):\n",
    "    return 1./(1.3 + 0.2*y - 0.003*x**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a mesh with pygmsh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Info    : Running 'gmsh -3 /var/folders/j6/s0152dxj3q54bsjvw2s42b0m0000gn/T/tmp1zokjhlz.geo -format msh -bin -o /var/folders/j6/s0152dxj3q54bsjvw2s42b0m0000gn/T/tmpl96mo5rb.msh' [Gmsh 4.5.4, 1 node, max. 1 thread]\n",
      "Info    : Started on Sat Mar 21 19:11:20 2020\n",
      "Info    : Reading '/var/folders/j6/s0152dxj3q54bsjvw2s42b0m0000gn/T/tmp1zokjhlz.geo'...\n",
      "Info    : Done reading '/var/folders/j6/s0152dxj3q54bsjvw2s42b0m0000gn/T/tmp1zokjhlz.geo'\n",
      "Info    : Meshing 1D...\n",
      "Info    : [  0 %] Meshing curve 1 (Nurb)\n",
      "Info    : [ 10 %] Meshing curve 2 (Line)\n",
      "Info    : [ 20 %] Meshing curve 3 (Nurb)\n",
      "Info    : [ 30 %] Meshing curve 4 (Line)\n",
      "Info    : [ 40 %] Meshing curve 5 (Line)\n",
      "Info    : [ 50 %] Meshing curve 6 (Line)\n",
      "Info    : [ 60 %] Meshing curve 7 (Line)\n",
      "Info    : [ 70 %] Meshing curve 8 (Line)\n",
      "Info    : [ 80 %] Meshing curve 9 (Line)\n",
      "Info    : [ 90 %] Meshing curve 10 (Line)\n",
      "Info    : [100 %] Meshing curve 11 (Line)\n",
      "Info    : Done meshing 1D (0.009525 s)\n",
      "Info    : Meshing 2D...\n",
      "Info    : [  0 %] Meshing surface 15 (Plane, Frontal)\n",
      "Info    : [ 40 %] Meshing surface 16 (Plane, Frontal)\n",
      "Info    : [ 70 %] Meshing surface 17 (Plane, Frontal)\n",
      "Info    : Done meshing 2D (0.030273 s)\n",
      "Info    : Meshing 3D...\n",
      "Info    : Done meshing 3D (2.2e-05 s)\n",
      "Info    : 652 nodes 1369 elements\n",
      "Info    : Writing '/var/folders/j6/s0152dxj3q54bsjvw2s42b0m0000gn/T/tmpl96mo5rb.msh'...\n",
      "Info    : Done writing '/var/folders/j6/s0152dxj3q54bsjvw2s42b0m0000gn/T/tmpl96mo5rb.msh'\n",
      "Info    : Stopped on Sat Mar 21 19:11:20 2020\n"
     ]
    }
   ],
   "source": [
    "geom = pygmsh.built_in.Geometry()\n",
    "\n",
    "lcar = 0.5\n",
    "\n",
    "pts = []\n",
    "pts.append(geom.add_point((0.0, 0.0, 0.0), lcar))    # 0\n",
    "pts.append(geom.add_point((1.0, 0.0, 0.0), lcar))    # 1\n",
    "pts.append(geom.add_point((5.0, 0.0, -2.0), lcar))   # 2\n",
    "pts.append(geom.add_point((10.0, 0.0, 2.0), lcar))   # 3\n",
    "pts.append(geom.add_point((14.0, 0.0, 0.0), lcar))   # 4\n",
    "pts.append(geom.add_point((15.0, 0.0, 0.0), lcar))   # 5\n",
    "pts.append(geom.add_point((15.0, 0.0, 4.0), lcar))   # 6 \n",
    "pts.append(geom.add_point((10.0, 0.0, 3.0), lcar))   # 7 \n",
    "pts.append(geom.add_point((5.0, 0.0, 5.0), lcar))    # 8 \n",
    "pts.append(geom.add_point((0.0, 0.0, 4.0), lcar))    # 9 \n",
    "pts.append(geom.add_point((15.0, 0.0, 8.0), lcar))   # 10\n",
    "pts.append(geom.add_point((0.0, 0.0, 8.0), lcar))    # 11\n",
    "\n",
    "pts.append(geom.add_point((3.0, 0.0, 5.0), lcar))    # 12\n",
    "pts.append(geom.add_point((9.0, 0.0, 5.0), lcar))    # 13\n",
    "pts.append(geom.add_point((9.0, 0.0, 7.0), lcar))    # 14\n",
    "pts.append(geom.add_point((3.0, 0.0, 7.0), lcar))    # 15\n",
    "\n",
    "\n",
    "lines = []\n",
    "lines.append(geom.add_bspline(pts[:6]))           # 0\n",
    "lines.append(geom.add_line(pts[5], pts[6]))       # 1\n",
    "lines.append(geom.add_bspline(pts[6:10]))         # 2\n",
    "lines.append(geom.add_line(pts[9], pts[0]))       # 3\n",
    "\n",
    "lines.append(geom.add_line(pts[6], pts[10]))      # 4\n",
    "lines.append(geom.add_line(pts[10], pts[11]))     # 5\n",
    "lines.append(geom.add_line(pts[11], pts[9]))      # 6\n",
    "\n",
    "lines.append(geom.add_line(pts[12], pts[13]))     # 7\n",
    "lines.append(geom.add_line(pts[13], pts[14]))     # 8\n",
    "lines.append(geom.add_line(pts[14], pts[15]))     # 9\n",
    "lines.append(geom.add_line(pts[15], pts[12]))     # 10\n",
    "\n",
    "ll0 = geom.add_line_loop(lines[7:])\n",
    "ll1 = geom.add_line_loop(lines[:4])\n",
    "ll2 = geom.add_line_loop((-lines[2], lines[4], lines[5], lines[6]))\n",
    "\n",
    "surf0 = geom.add_plane_surface(ll0)\n",
    "surf1 = geom.add_plane_surface(ll1)\n",
    "surf2 = geom.add_plane_surface(ll2, (ll0,))\n",
    "\n",
    "geom.add_physical(surf0, 'hole')\n",
    "geom.add_physical(surf1, 'layer_1')\n",
    "geom.add_physical(surf2, 'layer_2')\n",
    "\n",
    "mesh = pygmsh.generate_mesh(geom)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```cmesh2d``` want (x,z) coordinates, extract those from the mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "nodes = np.vstack((mesh.points[:,0], mesh.points[:,2])).T\n",
    "\n",
    "triangles = mesh.cells_dict['triangle']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute centroid of triangles and compute slowness at these coordonates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "phys = np.hstack(mesh.cell_data['gmsh:physical'])\n",
    "slowness = np.empty((triangles.shape[0],))\n",
    "for n in range(triangles.shape[0]):\n",
    "    t = triangles[n]\n",
    "    x = (nodes[t[0],0] + nodes[t[1],0] + nodes[t[2],0])/3.0\n",
    "    z = (nodes[t[0],1] + nodes[t[1],1] + nodes[t[2],1])/3.0\n",
    "\n",
    "    if phys[n] == 1:\n",
    "        slowness[n] = f0(np.array(x), np.array(z))\n",
    "    elif phys[n] == 2:\n",
    "        slowness[n] = f1(np.array(x), np.array(z))\n",
    "    else:\n",
    "        slowness[n] = f2(np.array(x), np.array(z))\n",
    "\n",
    "V = 1./slowness   # for plotting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create mesh and define source & receivers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "m1 = cmesh2d.Mesh2D(nodes, triangles, nsecondary=10, nthreads=1)\n",
    "\n",
    "Tx = np.array([[0.15, 0.15]])\n",
    "t0 = np.array([0.0])\n",
    "\n",
    "Rx = np.array([[2.0, -0.4],\n",
    "               [5.0, -0.7],\n",
    "               [7.0, -0.1],\n",
    "               [10.0, 1.0],\n",
    "               [14.0, 0.2],\n",
    "               [14.5, 2.0],\n",
    "               [14.5, 3.0],\n",
    "               [14.5, 4.0],\n",
    "               [14.5, 5.0],\n",
    "               [14.5, 6.0],\n",
    "               [14.5, 7.0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute traveltimes & raypaths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "tt, _, rays = m1.raytrace(slowness, Tx, Rx, t0, nout=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Show raypaths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 125848 characters]\n",
      "text/plain": [
       "<Figure size 720x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(10,4))\n",
    "ax = fig.add_subplot(111, aspect='equal')\n",
    "\n",
    "patches = []\n",
    "for t in triangles:\n",
    "    pts = np.array([[nodes[t[0],0], nodes[t[0],1]],\n",
    "                    [nodes[t[1],0], nodes[t[1],1]],\n",
    "                    [nodes[t[2],0], nodes[t[2],1]]])\n",
    "    p = matplotlib.patches.Polygon(pts, closed=True)\n",
    "    patches.append(p)\n",
    "colors = V\n",
    "\n",
    "cmap = sns.diverging_palette(230, 15, sep=20, as_cmap=True)\n",
    "\n",
    "p = matplotlib.collections.PatchCollection(patches, cmap=cmap)\n",
    "p.set_array(np.array(colors))\n",
    "ax.add_collection(p)\n",
    "\n",
    "plt.colorbar(p, ax=ax)\n",
    "for r in rays:\n",
    "    plt.plot(r[:,0], r[:,1], c=[0.5, 0.5, 0.5])\n",
    "\n",
    "ax.invert_yaxis()\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check traveltime values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 8176 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "plt.plot(tt, 'o')\n",
    "plt.xlabel('Rx number', fontsize=14)\n",
    "plt.ylabel('Traveltime', fontsize=14)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Extract traveltimes at all nodes and export to VTK format for visualization in other software (e.g. Paraview)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "tt_all = m1.get_grid_traveltimes()\n",
    "tt_all.shape\n",
    "fields = {'Velocity': V, 'Traveltime': tt_all}\n",
    "m1.to_vtk(fields, 'example_cmesh2d')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check raytracing options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on built-in function raytrace:\n",
      "\n",
      "raytrace(...) method of ttcrpy.cmesh2d.Mesh2D instance\n",
      "    raytrace(slowness, Tx, Rx, t0, nout, thread_no) -> tt,L,rays\n",
      "    \n",
      "    Compute traveltimes for a single source.  The source can comprise many\n",
      "    points, each having its own otigin time (t0).\n",
      "    \n",
      "    Parameters\n",
      "    ----------\n",
      "    slowness : vector of slowness at cells\n",
      "    Tx : source coordinates, ndarray of size nTx by 2\n",
      "        1st column contains X coordinates,\n",
      "        2nd contains Z coordinates\n",
      "    Rx : receiver coordinates, ndarray of size nRx by 2\n",
      "        1st column contains X coordinates,\n",
      "        2nd contains Z coordinates\n",
      "    t0 : origin times, scalar or ndarray of size nTx\n",
      "    nout : number of output variables (1, 2 or 3)\n",
      "    thread_no : thread/process number on which computation should be run\n",
      "    \n",
      "    Returns\n",
      "    -------\n",
      "    tt : travel times\n",
      "    L : (if nout > 1) data kernel matrix (scipy csr) (tt = L*slowness)\n",
      "    rays : (if nout == 3) coordinates of rays\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(m1.raytrace)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
