fill out start fix multiplane

docs/source/gravlensingintro.ipynb

@@ -8,6 +8,13 @@
     "\n",
     "Here we provide a primer course on gravitational lensing. It includes interactive examples using `caustics` to demo various effects.\n",
     "\n",
+    "This gravitational lensing primer will give you a good coverage of lensing concepts and intuition. The visual examples will help drive home the mathematical concepts. However, we do not take it upon ourselves in this tutorial to derive and explore the full depths of all the lensing formalism. There are plenty of excellent resources if you would like a more formal background for these concepts. We recommend the following resources for a good start:\n",
+    "- Introduction to Gravitational Lensing: With Python Examples by Massimo Meneghetti\n",
+    "- Bartelmann and Schneider 2001\n",
+    "- Fleury et al. 2022\n",
+    "- Petkova et al. 2014\n",
+    "- Keeton 2001\n",
+    "\n",
     "Quick note: we have hidden all the messy plotting commands for the sake of a clean tutorial. The code on display is the code needed to generate the data in the visualizations. If you want to see how the plots are made, just expand the bars above each figure which include the various matplotlib commands."
    ]
   },
@@ -1040,10 +1047,10 @@
    "outputs": [],
    "source": [
     "F = 100\n",
-    "source = caustics.Sersic(x0=0.05, y0=0.1, q=1.0, phi=0.0, n=2, Re=0.05)\n",
+    "source = caustics.Sersic(x0=0.1, y0=0.1, q=1.0, phi=0.0, n=2, Re=0.05)\n",
     "T = caustics.Param(\"T\")  # time variable\n",
     "source.Ie = (\n",
-    "    lambda p: torch.sin(5 * 2 * np.pi * (p[\"T\"].value % (2 * np.pi))) ** 2\n",
+    "    lambda p: torch.sin(2 * 2 * np.pi * (p[\"T\"].value % (2 * np.pi))) ** 2\n",
     ")  # periodic brightness\n",
     "source.Ie.link(T)\n",
     "cosmology = caustics.FlatLambdaCDM()\n",
@@ -1146,14 +1153,16 @@
     "\n",
     "where a subscript $s$ indicates the source plane quantity, and $\\mathcal{D}$ is an angular diameter distance. The challenge with this equation is that $\\vec{\\theta}_i$ is an argument of the deflection angle, making the expression cumbersome to evaluate. We now make the transition to comoving distances $D_i = (1+z_i)\\mathcal{D}_i$ and work in comoving coordinates. We can write an expression for the comoving coordinates of the next plane using only the planes before it like so:\n",
     "\n",
-    "$$\\vec{X}_{i+1} = \\vec{X}_i - D_{i+1,i}\\left(\\vec{\\theta}_0 + \\sum_j^{i-1}\\vec{\\hat{\\alpha}}_j(\\vec{X}_j)\\right)$$\n",
+    "$$\\vec{X}_{i+1} = \\vec{X}_i + D_{i+1,i}\\left(\\vec{\\theta}_0 - \\sum_j^{i-1}\\vec{\\hat{\\alpha}}_j(\\vec{X}_j)\\right)$$\n",
     "\n",
     "Which we can further break down into two recursion functions to be applied at each plane.\n",
     "\n",
-    "$$\\vec{\\theta}_{i+1} = \\vec{\\theta}_i + \\vec{\\hat{\\alpha}}_i(\\vec{X}_i / D_i)$$\n",
-    "$$\\vec{X}_{i+1} = \\vec{X}_i + D_{i,i+1}\\vec{\\theta}_{i+1}$$\n",
+    "$$\\vec{\\theta}_{i+1} = \\vec{\\theta}_i - \\vec{\\hat{\\alpha}}_i(\\vec{X}_i / D_i)$$\n",
+    "$$\\vec{X}_{i+1} = \\vec{X}_i + D_{i,i+1}\\vec{\\theta}_{i}$$\n",
+    "\n",
+    "where we see the first equation essentially breaks down the sum of the angles at each plane (only computing them as we need them, as we figure out the path one plane at a time), the second equation gives the recursion relation but with the angles already accumulated.\n",
     "\n",
-    "Where the relations must be computed in order since the $\\vec{X}_i$ one depends on $\\vec{\\theta}_{i+1}$."
+    "Note that the indexing for $\\vec{X}_i$ and $\\vec{\\theta}_i$ is slightly different. It is best for $\\vec{X}_i$ to think of each plane (including the image plane and source plane) as being the index for $i$, so $i=0$ is the image plane, $i=1$ is the first lens plane, until $i=s$ is the source plane. For $\\vec{\\theta}_i$ it is best to think of is as indexing between planes. So for $\\vec{\\theta}_i$ the $i=0$ index is for angles between the image plane and the first lens plane, $i=1$ would be between the first lens plane and the second, and so on. This means that the indexing for $\\vec{X}_i$ ultimately goes one higher than for $\\vec{\\theta}_i$."
    ]
   },
   {
@@ -1180,9 +1189,9 @@
     "\n",
     "theta_x = [torch.linspace(-fov / 2, fov / 2, F, dtype=torch.float32)]\n",
     "theta_y = [torch.zeros(F, dtype=torch.float32)]\n",
-    "D0 = cosmology.transverse_comoving_distance(z_ls[0]) * arcsec_to_rad * 1000\n",
-    "X = [torch.zeros_like(theta_x[0]), theta_x[0] * D0]\n",
-    "Y = [torch.zeros_like(theta_y[0]), theta_y[0] * D0]\n",
+    "D01 = cosmology.transverse_comoving_distance(z_ls[0]) * arcsec_to_rad * 1000\n",
+    "X = [torch.zeros_like(theta_x[0]), theta_x[0] * D01]\n",
+    "Y = [torch.zeros_like(theta_y[0]), theta_y[0] * D01]\n",
     "\n",
     "for i, lens in enumerate(lenses):\n",
     "    # Next plane\n",
@@ -1195,7 +1204,7 @@
     "        * 1000\n",
     "    )\n",
     "    # Deflection\n",
-    "    alpha_x, alpha_y = lens.physical_deflection_angle(X[-1] / Di, Y[-1] / Di, z_next)\n",
+    "    alpha_x, alpha_y = lens.physical_deflection_angle(X[-1] / Di, Y[-1] / Di, z_s)\n",
     "    # Update angles\n",
     "    theta_x.append(theta_x[i] - alpha_x)\n",
     "    theta_y.append(theta_y[i] - alpha_y)\n",