Merge pull request #316 from Ciela-Institute/dev

Dev

.pre-commit-config.yaml

@@ -51,13 +51,13 @@ repos:
         args: [--prose-wrap=always]
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.8.3"
+    rev: "v0.9.2"
     hooks:
       - id: ruff
         args: ["--fix", "--show-fixes"]
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: "v1.13.0"
+    rev: "v1.14.1"
     hooks:
       - id: mypy
         files: src

README.md

@@ -36,13 +36,13 @@ import caustics
 import torch
 
 cosmology = caustics.FlatLambdaCDM()
-sie = caustics.SIE(cosmology=cosmology, name="lens")
+sie = caustics.SIE(cosmology=cosmology, name="lens", z_l=0.5, z_s=1.0)
 src = caustics.Sersic(name="source")
 lnslt = caustics.Sersic(name="lenslight")
 
 x = torch.tensor([
-#   z_s  z_l   x0   y0   q    phi     b    x0   y0   q     phi    n    Re
-    1.5, 0.5, -0.2, 0.0, 0.4, 1.5708, 1.7, 0.0, 0.0, 0.5, -0.985, 1.3, 1.0,
+#    x0   y0   q    phi     Rein x0   y0   q     phi    n    Re
+    -0.2, 0.0, 0.4, 1.5708, 1.7, 0.0, 0.0, 0.5, -0.985, 1.3, 1.0,
 #   Ie    x0   y0   q    phi  n   Re   Ie
     5.0, -0.2, 0.0, 0.8, 0.0, 1., 1.0, 10.0
 ])  # fmt: skip

docs/source/_toc.yml

@@ -21,7 +21,7 @@ chapters:
       - file: tutorials/VisualizeCaustics
       - file: tutorials/MultiplaneDemo
       - file: tutorials/InvertLensEquation
-      - file: tutorials/Parameters
+      - file: tutorials/MultiSource
   - file: examples/index
     sections:
       - file: examples/Example_ImageFit_LM

docs/source/examples/Example_ImageFit_LM.ipynb

@@ -86,12 +86,16 @@
     "# Set up the forward model\n",
     "\n",
     "# Lens mass model (SIE + shear)\n",
-    "lens_sie = caustics.SIE(name=\"galaxylens\", cosmology=cosmology, z_l=1.0)\n",
+    "lens_sie = caustics.SIE(name=\"galaxylens\", cosmology=cosmology)\n",
     "lens_shear = caustics.ExternalShear(\n",
-    "    name=\"externalshear\", cosmology=cosmology, x0=0.0, y0=0.0, z_l=1.0\n",
+    "    name=\"externalshear\", cosmology=cosmology, x0=0.0, y0=0.0\n",
     ")\n",
     "lens_mass_model = caustics.SinglePlane(\n",
-    "    name=\"lensmass\", cosmology=cosmology, lenses=[lens_sie, lens_shear], z_l=1.0\n",
+    "    name=\"lensmass\",\n",
+    "    cosmology=cosmology,\n",
+    "    lenses=[lens_sie, lens_shear],\n",
+    "    z_l=z_l,\n",
+    "    z_s=z_s,\n",
     ")\n",
     "\n",
     "# Lens light model (sersic)\n",
@@ -119,7 +123,6 @@
     "    pixelscale=pixelscale,\n",
     "    upsample_factor=upsample_factor,\n",
     "    quad_level=quad_level,\n",
-    "    z_s=2.0,\n",
     ")"
    ]
   },
@@ -151,7 +154,7 @@
     "        \"y0\": 0.0,\n",
     "        \"q\": 0.86,\n",
     "        \"phi\": -0.20,\n",
-    "        \"b\": 0.66,\n",
+    "        \"Rein\": 0.66,\n",
     "    },\n",
     "    \"externalshear\": {\"gamma_1\": 0.0, \"gamma_2\": -0.05},\n",
     "    \"sourcelight\": {\n",

docs/source/examples/Example_ImageFit_MCMC.ipynb

@@ -96,12 +96,16 @@
    "outputs": [],
    "source": [
     "# Lens mass model (SIE + shear)\n",
-    "lens_sie = caustics.SIE(name=\"galaxylens\", cosmology=cosmology, z_l=1.0)\n",
+    "lens_sie = caustics.SIE(name=\"galaxylens\", cosmology=cosmology)\n",
     "lens_shear = caustics.ExternalShear(\n",
-    "    name=\"externalshear\", cosmology=cosmology, x0=0.0, y0=0.0, z_l=1.0\n",
+    "    name=\"externalshear\", cosmology=cosmology, x0=0.0, y0=0.0\n",
     ")\n",
     "lens_mass_model = caustics.SinglePlane(\n",
-    "    name=\"lensmass\", cosmology=cosmology, lenses=[lens_sie, lens_shear], z_l=1.0\n",
+    "    name=\"lensmass\",\n",
+    "    cosmology=cosmology,\n",
+    "    lenses=[lens_sie, lens_shear],\n",
+    "    z_l=z_l,\n",
+    "    z_s=z_s,\n",
     ")\n",
     "\n",
     "# Lens light model (sersic)\n",
@@ -128,7 +132,6 @@
     "    pixels_x=numPix,\n",
     "    pixelscale=pixelscale,\n",
     "    upsample_factor=upsample_factor,\n",
-    "    z_s=2.0,\n",
     ")"
    ]
   },
@@ -160,7 +163,7 @@
     "        \"y0\": 0.0,\n",
     "        \"q\": 0.86,\n",
     "        \"phi\": -0.20,\n",
-    "        \"b\": 0.66,\n",
+    "        \"Rein\": 0.66,\n",
     "    },\n",
     "    \"externalshear\": {\"gamma_1\": 0.0, \"gamma_2\": -0.05},\n",
     "    \"sourcelight\": {\n",

docs/source/examples/Example_QSOLensFit.ipynb

@@ -52,6 +52,7 @@
     "    cosmology=cosmology,\n",
     "    name=\"sie\",\n",
     "    z_l=z_l,\n",
+    "    z_s=z_s,\n",
     "    s=1e-3,\n",
     ")"
    ]
@@ -75,12 +76,12 @@
     "sp_x = torch.tensor(0.2)\n",
     "sp_y = torch.tensor(0.2)\n",
     "\n",
-    "#  true parameters     x0   y0   q    phi        b\n",
+    "#  true parameters     x0   y0   q    phi        Rein\n",
     "params = torch.tensor([0.0, 0.0, 0.4, np.pi / 5, 1.0])\n",
     "# Points in image plane\n",
-    "x, y = lens.forward_raytrace(sp_x, sp_y, z_s, params)\n",
+    "x, y = lens.forward_raytrace(sp_x, sp_y, params)\n",
     "# get magnifications\n",
-    "mu = lens.magnification(x, y, z_s, params)\n",
+    "mu = lens.magnification(x, y, params)\n",
     "# remove heavily demagnified points\n",
     "x = x[mu > 1e-2]\n",
     "y = y[mu > 1e-2]"
@@ -111,7 +112,7 @@
     "    dtype=torch.float32,\n",
     ")\n",
     "\n",
-    "A = lens.jacobian_lens_equation(thx, thy, z_s, params)\n",
+    "A = lens.jacobian_lens_equation(thx, thy, params)\n",
     "detA = torch.linalg.det(A)"
    ]
   },
@@ -136,7 +137,7 @@
     "    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, params)\n",
+    "    y1, y2 = lens.raytrace(x1, x2, params)\n",
     "\n",
     "    # Plot the caustic\n",
     "    ax.plot(y1, y2, color=\"r\", zorder=1)\n",
@@ -172,7 +173,7 @@
     "# Compute Chi^2 which is just the distance between the four images raytraced back to the source plane.\n",
     "# This returns zero when all QSO images land in the same spot.\n",
     "def loss(P):\n",
-    "    bx, by = lens.raytrace(x, y, z_l, P)\n",
+    "    bx, by = lens.raytrace(x, y, P)\n",
     "    return torch.cat(\n",
     "        (\n",
     "            torch.sum((bx[0] - bx[1:]) ** 2).unsqueeze(-1),\n",
@@ -191,7 +192,7 @@
     "# If you decrease the threshold from 1e-8 to 1e-10 you will get better fits, but fewer of them\n",
     "avg_fit = torch.mean(fit_params[0][fit_params[2] < 1e-8], dim=0)\n",
     "print(avg_fit.numpy())\n",
-    "# Note that the order is: x0, y0, q, phi, b"
+    "# Note that the order is: x0, y0, q, phi, Rein"
    ]
   },
   {
@@ -226,15 +227,15 @@
     "    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, params)\n",
+    "    y1, y2 = lens.raytrace(x1, x2, params)\n",
     "\n",
     "    # Plot the caustic\n",
     "    ax.plot(y1, y2, color=\"r\", label=\"True\", zorder=1)\n",
     "\n",
     "# Plot fitted caustics\n",
     "#################\n",
     "\n",
-    "A = lens.jacobian_lens_equation(thx, thy, z_s, avg_fit)\n",
+    "A = lens.jacobian_lens_equation(thx, thy, avg_fit)\n",
     "detA = torch.linalg.det(A)\n",
     "\n",
     "CS = ax.contour(\n",
@@ -248,7 +249,7 @@
     "    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, params)\n",
+    "    y1, y2 = lens.raytrace(x1, x2, params)\n",
     "\n",
     "    # Plot the caustic\n",
     "    ax.plot(y1, y2, color=\"green\", linestyle=\"--\", label=\"Fit\", zorder=1)\n",