refactor: move quad_level to init (#245)

* refactor: move quad_level to init * update unit tests for LensSource * fiw more unit tests * yaml needs null not None * yaml sim cant accept none * refactor: lens source gridding system * fix no source mu shape * fix mu with no source * actually fix it this time * fix example usage docstring * Add psf as parameter * sim test quad level set * sim should work * get outputs right, shape and normalization * add check conv2d vs fft * fix fft roll in lens source * update docs with new quad_level argument fmt * remove unnecessary 'off' option for psf_mode
Ciela-Institute · Jul 27, 2024 · f72d8ad · f72d8ad
1 parent 9806a84
commit f72d8ad
Show file tree

Hide file tree

Showing 9 changed files with 244 additions and 130 deletions.
diff --git a/docs/source/examples/Example_ImageFit_LM.ipynb b/docs/source/examples/Example_ImageFit_LM.ipynb
@@ -56,6 +56,7 @@
     "cosmology.to(dtype=torch.float32)\n",
     "\n",
     "upsample_factor = 1\n",
+    "quad_level = 3\n",
     "thx, thy = caustics.utils.meshgrid(\n",
     "    pixelscale / upsample_factor,\n",
     "    upsample_factor * numPix,\n",
@@ -117,6 +118,7 @@
     "    pixels_x=numPix,\n",
     "    pixelscale=pixelscale,\n",
     "    upsample_factor=upsample_factor,\n",
+    "    quad_level=quad_level,\n",
     "    z_s=2.0,\n",
     ")"
    ]
@@ -128,9 +130,7 @@
    "source": [
     "## Sample some mock data\n",
     "\n",
-    "Here we write out the true values for all the parameters in the model. In total there are 21 parameters, so this is quite a complex model already! We then plot the data so we can see what it is we re trying to fit.\n",
-    "\n",
-    "Note that when we sample the simulator we call it with `quad_level=7`. This means the simulator will use gaussian quadrature sub-pixel integration to ensure the brightness of each pixel is very accurately computed."
+    "Here we write out the true values for all the parameters in the model. In total there are 21 parameters, so this is quite a complex model already! We then plot the data so we can see what it is we re trying to fit."
    ]
   },
   {
@@ -178,7 +178,7 @@
     "print(true_params)\n",
     "\n",
     "# simulate lens, crop extra evaluation for PSF\n",
-    "true_system = sim(allparams, quad_level=7)  # simulate at high resolution\n",
+    "true_system = sim(allparams)\n",
     "\n",
     "fig, axarr = plt.subplots(1, 2, figsize=(15, 8))\n",
     "axarr[0].imshow(\n",
@@ -231,7 +231,7 @@
     "res = caustics.utils.batch_lm(\n",
     "    batch_inits,\n",
     "    obs_system.reshape(-1).repeat(10, 1),\n",
-    "    lambda x: sim(x, quad_level=3).reshape(-1),\n",
+    "    lambda x: sim(x).reshape(-1),\n",
     "    C=variance.reshape(-1).repeat(10, 1),\n",
     ")\n",
     "best_fit = res[0][np.argmin(res[2].numpy())]\n",
@@ -288,7 +288,7 @@
    "outputs": [],
    "source": [
     "# Compute jacobian\n",
-    "J = torch.func.jacfwd(lambda x: sim(x, quad_level=3))(best_fit)\n",
+    "J = torch.func.jacfwd(lambda x: sim(x))(best_fit)\n",
     "fig, axarr = plt.subplots(3, 7, figsize=(21, 9))\n",
     "for i, ax in enumerate(axarr.flatten()):\n",
     "    ax.imshow(J[..., i], origin=\"lower\")\n",

diff --git a/docs/source/examples/Example_ImageFit_NUTS.ipynb b/docs/source/examples/Example_ImageFit_NUTS.ipynb
@@ -61,6 +61,7 @@
     "cosmology.to(dtype=torch.float32)\n",
     "\n",
     "upsample_factor = 1\n",
+    "quad_level = 3\n",
     "thx, thy = caustics.utils.meshgrid(\n",
     "    pixelscale / upsample_factor,\n",
     "    upsample_factor * numPix,\n",
@@ -133,9 +134,7 @@
    "source": [
     "## Sample some mock data\n",
     "\n",
-    "Here we write out the true values for all the parameters in the model. In total there are 21 parameters, so this is quite a complex model already! We then plot the data so we can see what it is we re trying to fit.\n",
-    "\n",
-    "Note that when we sample the simulator we call it with quad_level=7. This means the simulator will use gaussian quadrature sub-pixel integration to ensure the brightness of each pixel is very accurately computed."
+    "Here we write out the true values for all the parameters in the model. In total there are 21 parameters, so this is quite a complex model already! We then plot the data so we can see what it is we re trying to fit."
    ]
   },
   {
@@ -182,7 +181,7 @@
     "print(true_params)\n",
     "\n",
     "# simulate lens, crop extra evaluation for PSF\n",
-    "true_system = sim(allparams, quad_level=7)  # simulate at high resolution\n",
+    "true_system = sim(allparams)  # simulate at high resolution\n",
     "\n",
     "fig, axarr = plt.subplots(1, 2, figsize=(15, 8))\n",
     "axarr[0].imshow(\n",

diff --git a/docs/source/intro.md b/docs/source/intro.md
@@ -39,10 +39,10 @@ x = torch.tensor([
     5.0, -0.2, 0.0, 0.8, 0.0, 1., 1.0, 10.0
 ])  # fmt: skip
 
-minisim = caustics.LensSource(
-    lens=sie, source=src, lens_light=lnslt, pixelscale=0.05, pixels_x=100
+sim = caustics.LensSource(
+    lens=sie, source=src, lens_light=lnslt, pixelscale=0.05, pixels_x=100, quad_level=3
 )
-plt.imshow(minisim(x, quad_level=3), origin="lower")
+plt.imshow(sim(x), origin="lower")
 plt.show()
 ```
 
@@ -54,7 +54,7 @@ plt.show()
 newx = x.repeat(20, 1)
 newx += torch.normal(mean=0, std=0.1 * torch.ones_like(newx))
 
-images = torch.vmap(minisim)(newx)
+images = torch.vmap(sim)(newx)
 
 fig, axarr = plt.subplots(4, 5, figsize=(20, 16))
 for ax, im in zip(axarr.flatten(), images):
@@ -67,7 +67,7 @@ plt.show()
 ### Automatic Differentiation
 
 ```python
-J = torch.func.jacfwd(minisim)(x)
+J = torch.func.jacfwd(sim)(x)
 
 # Plot the new images
 fig, axarr = plt.subplots(3, 7, figsize=(20, 9))

diff --git a/docs/source/tutorials/InterfaceIntroduction_oop.ipynb b/docs/source/tutorials/InterfaceIntroduction_oop.ipynb
@@ -129,7 +129,7 @@
    "outputs": [],
    "source": [
     "sim = caustics.LensSource(\n",
-    "    lens=sie, source=src, lens_light=lnslt, pixelscale=0.05, pixels_x=100\n",
+    "    lens=sie, source=src, lens_light=lnslt, pixelscale=0.05, pixels_x=100, quad_level=3\n",
     ")"
    ]
   },
@@ -174,7 +174,7 @@
    "outputs": [],
    "source": [
     "# Substitute minisim with sim for yaml method\n",
-    "image = sim(x, quad_level=3).detach().cpu().numpy()\n",
+    "image = sim(x).detach().cpu().numpy()\n",
     "\n",
     "plt.imshow(image, origin=\"lower\")\n",
     "plt.axis(\"off\")\n",

diff --git a/docs/source/tutorials/InterfaceIntroduction_yaml.ipynb b/docs/source/tutorials/InterfaceIntroduction_yaml.ipynb
@@ -100,7 +100,7 @@
    "source": [
     "### Plot the Results!\n",
     "\n",
-    "This section is mostly self explanatory. We evaluate the simulator configuration by calling it like a function. There are several possible arguments to the simulator function, here we use `quad_level` which indicates the depth of quadrature integration to use for each pixel (how accurate to make the flux)."
+    "This section is mostly self explanatory. We evaluate the simulator configuration by calling it like a function. "
    ]
   },
   {
@@ -110,7 +110,7 @@
    "outputs": [],
    "source": [
     "# Here we sample an image\n",
-    "image = sim(x, quad_level=3).detach().cpu().numpy()\n",
+    "image = sim(x).detach().cpu().numpy()\n",
     "\n",
     "plt.imshow(image, origin=\"lower\")\n",
     "plt.axis(\"off\")\n",

diff --git a/docs/source/tutorials/example.yml b/docs/source/tutorials/example.yml
@@ -26,3 +26,4 @@ simulator:
     lens_light: *lnslt
     pixelscale: 0.05
     pixels_x: 100
+    quad_level: 3