docs: Add example with local modelling (#239)

* Add example with local modelling

* fix interface with simulator grid

docs/source/tutorials/InvertLensEquation.ipynb

@@ -115,20 +115,126 @@
     "ax.scatter(x, y, color=\"b\", label=\"forward raytrace\", zorder=10)\n",
     "ax.scatter(bx, by, color=\"r\", marker=\"x\", label=\"source plane\", zorder=9)\n",
     "ax.scatter([sp_x.item()], [sp_y.item()], color=\"g\", label=\"true pos\", zorder=8)\n",
+    "ax.set_axis_off()\n",
     "plt.legend()\n",
     "plt.show()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "0163d5f7",
+   "metadata": {},
+   "source": [
+    "## Lets take a look\n",
+    "\n",
+    "Using the `LensSource` simulator and the forward raytracing coordinates we can focus our calculations on the regions of interest for each image."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "5",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "src = caustics.Sersic(\n",
+    "    x0=0.2, y0=0.2, q=0.9, phi=0.0, n=1.0, Re=0.05, Ie=1.0, name=\"source\"\n",
+    ")\n",
+    "\n",
+    "sim = caustics.LensSource(\n",
+    "    lens=lens, source=src, z_s=z_s, x0=None, y0=None, pixelscale=0.005, pixels_x=100\n",
+    ")\n",
+    "\n",
+    "# Plot the source and lens\n",
+    "fig, ax = plt.subplots()\n",
+    "CS = ax.contour(thx, thy, detA, levels=[0.0], colors=\"b\", zorder=1)\n",
+    "# Get the path from the matplotlib contour plot of the critical line\n",
+    "for path in paths:\n",
+    "    # Collect the path into a discrete set of points\n",
+    "    x1 = torch.tensor(list(float(vs[0]) for vs in path))\n",
+    "    x2 = torch.tensor(list(float(vs[1]) for vs in path))\n",
+    "    # raytrace the points to the source plane\n",
+    "    y1, y2 = lens.raytrace(x1, x2, z_s)\n",
+    "\n",
+    "    # Plot the caustic\n",
+    "    ax.plot(y1, y2, color=\"r\", zorder=1)\n",
+    "for i in range(len(x)):\n",
+    "    ax.imshow(\n",
+    "        sim([x[i], y[i]]),\n",
+    "        extent=(\n",
+    "            -sim.pixelscale * sim.pixels_x / 2 + x[i],\n",
+    "            sim.pixelscale * sim.pixels_x / 2 + x[i],\n",
+    "            -sim.pixelscale * sim.pixels_y / 2 + y[i],\n",
+    "            sim.pixelscale * sim.pixels_y / 2 + y[i],\n",
+    "        ),\n",
+    "        origin=\"lower\",\n",
+    "    )\n",
+    "ax.set_xlim([-1.5, 2])\n",
+    "ax.set_ylim([-1.5, 2])\n",
+    "ax.set_axis_off()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "edb5fa5e",
+   "metadata": {},
+   "source": [
+    "This is much more efficient than evaluating a whole image. Below you can see the same setup but we see how the simulator spends a lot of pixels evaluating low flux areas that don't matter much for modelling."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c0689864",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sim_wide = caustics.LensSource(\n",
+    "    lens=lens, source=src, z_s=z_s, pixelscale=0.005, pixels_x=1000\n",
+    ")\n",
+    "fig, ax = plt.subplots()\n",
+    "CS = ax.contour(thx, thy, detA, levels=[0.0], colors=\"b\", zorder=1)\n",
+    "for path in paths:\n",
+    "    # Collect the path into a discrete set of points\n",
+    "    x1 = torch.tensor(list(float(vs[0]) for vs in path))\n",
+    "    x2 = torch.tensor(list(float(vs[1]) for vs in path))\n",
+    "    # raytrace the points to the source plane\n",
+    "    y1, y2 = lens.raytrace(x1, x2, z_s)\n",
+    "\n",
+    "    # Plot the caustic\n",
+    "    ax.plot(y1, y2, color=\"r\", zorder=1)\n",
+    "ax.imshow(\n",
+    "    sim_wide({}),\n",
+    "    origin=\"lower\",\n",
+    "    extent=(\n",
+    "        -sim_wide.pixelscale * sim_wide.pixels_x / 2,\n",
+    "        sim_wide.pixelscale * sim_wide.pixels_x / 2,\n",
+    "        -sim_wide.pixelscale * sim_wide.pixels_y / 2,\n",
+    "        sim_wide.pixelscale * sim_wide.pixels_y / 2,\n",
+    "    ),\n",
+    ")\n",
+    "ax.set_xlim([-1.5, 2])\n",
+    "ax.set_ylim([-1.5, 2])\n",
+    "ax.set_axis_off()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d2fe0e6f",
+   "metadata": {},
+   "outputs": [],
    "source": []
   }
  ],
  "metadata": {
+  "kernelspec": {
+   "display_name": "PY39",
+   "language": "python",
+   "name": "python3"
+  },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
@@ -138,7 +244,8 @@
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3"
+   "pygments_lexer": "ipython3",
+   "version": "3.9.5"
   }
  },
  "nbformat": 4,