{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Checkout www.pygimli.org for more examples"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Refraction in 3D\n\nThis example shows refracted ray paths in a three-dimensional vertical gradient\nmedium.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>This is a placeholder/proof-of-concept. The code should be refactored\n    partly to `tt.showRayPaths()`</p></div>\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport pygimli as pg\nimport pygimli.meshtools as mt\nfrom pygimli.physics import traveltime as tt\nfrom pygimli.viewer.pv import drawSensors\nimport pyvista"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Build mesh.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "depth = 15\nwidth = 30\nplc = mt.createCube(size=[width, width, depth], pos=[0, 0, -depth/2], area=5)\n\nn_sensors = 8\nsensors = np.zeros((n_sensors, 3))\nsensors[0, 0] = 15\nsensors[0, 1] = -10\nsensors[1:, 0] = -15\nsensors[1:, 1] = np.linspace(-15, 15, n_sensors - 1)\n\nfor pos in sensors:\n    plc.createNode(pos)\nmesh = mt.createMesh(plc)\nmesh.createSecondaryNodes(1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Create vertical gradient model.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "vel = 300 + -pg.z(mesh.cellCenters()) * 100\npg.show(mesh, vel, label=pg.utils.unit(\"vel\"), showMesh=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Set-up data container.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "data = tt.createRAData(sensors)\ndata.markInvalid(data(\"s\") > 1)\ndata.set(\"t\", np.zeros(data.size()))\ndata.removeInvalid()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Do raytracing.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# fop = pg.core.TravelTimeDijkstraModelling(mesh, data)\nfop = tt.TravelTimeModelling()\nfop.setData(data)\nfop.setMesh(mesh)\nprint(fop.mesh())\n\n# This is to show single raypaths.\ngraph = fop.createGraph(1 / vel)\ndij = tt.Dijkstra(graph)\ndij.setStartNode(mesh.findNearestNode([15, -10, 0]))\n\nrays = []\nfor receiver in sensors[1:]:\n    ni = dij.shortestPathTo(mesh.findNearestNode(receiver))\n    pos = mesh.positions(withSecNodes=True)[ni]\n    segs = np.zeros((len(pos), 3))\n    segs[:,0] = pg.x(pos)\n    segs[:,1] = pg.y(pos)\n    segs[:,2] = pg.z(pos)\n    rays.append(segs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Plot final ray paths.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "pl, _ = pg.show(mesh, style='wireframe', line_width=0.1,\n                        hold=True)\ndrawSensors(pl, sensors, diam=0.5, color='yellow')\n\nfor ray in rays:\n    for i in range(len(ray) - 1):\n        start = tuple(ray[i])\n        stop = tuple(ray[i + 1])\n        line = pyvista.Line(start, stop)\n        pl.add_mesh(line, color='green', line_width=2)\n\npl.show()"
      ]
    }
  ],
  "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.11.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}