Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

docs: Demo making new lens model #324

Open
wants to merge 3 commits into
base: dev
Choose a base branch
from
Open

docs: Demo making new lens model #324

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
74eee89
demo making a custom lens
ConnorStoneAstro Feb 10, 2025
c304422
add demo to toc
ConnorStoneAstro Feb 10, 2025
12f8587
hide plotting commands
ConnorStoneAstro Feb 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/source/_toc.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ chapters:
- file: tutorials/VisualizeCaustics
- file: tutorials/MultiplaneDemo
- file: tutorials/InvertLensEquation
- file: tutorials/NewModelDemo
- file: tutorials/MultiSource
- file: tutorials/LensZoo
- file: examples/index
Expand Down
351 changes: 351 additions & 0 deletions docs/source/tutorials/NewModelDemo.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,351 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Making a new lens model\n",
"\n",
"Here we will demo how you can make your own lens model just by defining a potential, `caustics` will take care of the rest."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%load_ext autoreload\n",
"%autoreload 2\n",
"%matplotlib inline\n",
"\n",
"import matplotlib.pyplot as plt\n",
"import torch\n",
"import caustics\n",
"from caustics import forward, Param"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Below we define a class that inherits from `caustics.ThinLens`, this is the abstract class for all single plane lenses in `caustics`. The base class needs a cosmology, lens redshift, and source redshift which are passed via `super()`. After that we define the `Param`s needed by our class (which is just some gaussian parameters). Finally, we define the potential function for our class, which in this case is a gaussian. The potential is convenient because all other lensing quantities (deflection angle and convergence) can be determined from derivatives of the potential. This is why, given only the potential, `caustics` is able to build a full model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class GaussianPotential(caustics.ThinLens):\n",
"\n",
" def __init__(self, cosmology, z_l, z_s, x0, y0, A, sigma):\n",
" super().__init__(cosmology=cosmology, z_l=z_l, z_s=z_s)\n",
"\n",
" self.x0 = Param(\"x0\", x0)\n",
" self.y0 = Param(\"y0\", y0)\n",
" self.A = Param(\"A\", A)\n",
" self.sigma = Param(\"sigma\", sigma)\n",
"\n",
" @forward\n",
" def potential(self, x, y, x0, y0, A, sigma):\n",
" return -A * torch.exp(-((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma**2))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now we can do a really basic simulation just to see everything is working. We take a Sersic source model and use the `LensSource` simulator to make an image of the lensing from our gaussian potential model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"cosmo = caustics.FlatLambdaCDM()\n",
"lens = GaussianPotential(cosmo, z_l=0.5, z_s=1.0, x0=0.0, y0=0.0, A=2.0, sigma=1.0)\n",
"src = caustics.Sersic(x0=0.2, y0=0.2, q=0.6, phi=1.0, Ie=1.0, Re=1.0, n=2.0)\n",
"\n",
"sim = caustics.LensSource(lens, src, pixels_x=100, pixelscale=0.05, upsample_factor=2)\n",
"\n",
"plt.imshow(sim().numpy(), origin=\"lower\")\n",
"plt.axis(\"off\")\n",
"plt.title(\"Sersic lensed with Gaussian potential\")\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that we've tried lensing, lets look at all the basic lensing quantities and map them out for our new lens. The potential is exactly as we specified, the deflection angles are its derivatives, the convergence comes from second derivatives, and so on. We can compute shear, magnification, and the time delay field as well."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"hide-input"
]
},
"outputs": [],
"source": [
"fig, axarr = plt.subplots(2, 4, figsize=(20, 10))\n",
"n_pix = 100\n",
"res = 0.05\n",
"thx, thy = caustics.utils.meshgrid(res, n_pix, dtype=torch.float32)\n",
"axarr[0][0].imshow(lens.potential(thx, thy).numpy(), origin=\"lower\")\n",
"axarr[0][0].set_title(\"Potential\")\n",
"axarr[0][0].axis(\"off\")\n",
"axarr[0][1].imshow(lens.reduced_deflection_angle(thx, thy)[0].numpy(), origin=\"lower\")\n",
"axarr[0][1].set_title(\"Deflection x\")\n",
"axarr[0][1].axis(\"off\")\n",
"axarr[0][2].imshow(lens.reduced_deflection_angle(thx, thy)[1].numpy(), origin=\"lower\")\n",
"axarr[0][2].set_title(\"Deflection y\")\n",
"axarr[0][2].axis(\"off\")\n",
"axarr[0][3].imshow(lens.convergence(thx, thy).numpy(), origin=\"lower\")\n",
"axarr[0][3].set_title(\"Convergence\")\n",
"axarr[0][3].axis(\"off\")\n",
"axarr[1][0].imshow(lens.shear(thx, thy)[0].numpy(), origin=\"lower\")\n",
"axarr[1][0].set_title(\"Shear g1\")\n",
"axarr[1][0].axis(\"off\")\n",
"axarr[1][1].imshow(lens.shear(thx, thy)[1].numpy(), origin=\"lower\")\n",
"axarr[1][1].set_title(\"Shear g2\")\n",
"axarr[1][1].axis(\"off\")\n",
"axarr[1][2].imshow(\n",
" torch.clamp(lens.magnification(thx, thy), -10.0, 20.0).numpy(), origin=\"lower\"\n",
")\n",
"axarr[1][2].set_title(\"Magnification\")\n",
"axarr[1][2].axis(\"off\")\n",
"axarr[1][3].imshow(lens.time_delay(thx, thy).numpy(), origin=\"lower\")\n",
"axarr[1][3].set_title(\"Time delay\")\n",
"axarr[1][3].axis(\"off\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"If you know the analytic form of one of the quantities, you may want to write out the appropriate function and overload the base class method which uses autograd to compute it. This will be faster since you've done some of the work for the code by figuring out the analytic form."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class GaussianPotentialFast(caustics.ThinLens):\n",
"\n",
" def __init__(self, cosmology, z_l, z_s, x0, y0, A, sigma):\n",
" super().__init__(cosmology=cosmology, z_l=z_l, z_s=z_s)\n",
"\n",
" self.x0 = Param(\"x0\", x0)\n",
" self.y0 = Param(\"y0\", y0)\n",
" self.A = Param(\"A\", A)\n",
" self.sigma = Param(\"sigma\", sigma)\n",
"\n",
" @forward\n",
" def potential(self, x, y, x0, y0, A, sigma):\n",
" return -A * torch.exp(-((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma**2))\n",
"\n",
" @forward\n",
" def reduced_deflection_angle(self, x, y, x0, y0, A, sigma):\n",
" ax = -(x - x0) / sigma**2 # derivative of exponent\n",
" ay = -(y - y0) / sigma**2\n",
" p = self.potential(x, y) # exponential stays after derivative\n",
" return ax * p, ay * p\n",
"\n",
" @forward\n",
" def convergence(self, x, y, x0, y0, A, sigma):\n",
" p = self.potential(x, y)\n",
" dx = (x - x0) ** 2 / sigma**4\n",
" dxdx = -1 / sigma**2\n",
" dy = (y - y0) ** 2 / sigma**4\n",
" return 0.5 * (2 * dxdx + dx + dy) * p"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"hide-input"
]
},
"outputs": [],
"source": [
"lens_basic = GaussianPotential(\n",
" cosmo, z_l=0.5, z_s=1.0, x0=0.0, y0=0.0, A=2.0, sigma=1.0\n",
")\n",
"lens_fast = GaussianPotentialFast(\n",
" cosmo, z_l=0.5, z_s=1.0, x0=0.0, y0=0.0, A=2.0, sigma=1.0\n",
")\n",
"\n",
"fig, axarr = plt.subplots(2, 3, figsize=(15, 10))\n",
"axarr[0][0].imshow(\n",
" lens_basic.reduced_deflection_angle(thx, thy)[0].numpy(), origin=\"lower\"\n",
")\n",
"axarr[0][0].set_title(\"Deflection x basic\")\n",
"axarr[0][0].axis(\"off\")\n",
"axarr[0][1].imshow(\n",
" lens_fast.reduced_deflection_angle(thx, thy)[0].numpy(), origin=\"lower\"\n",
")\n",
"axarr[0][1].set_title(\"Deflection x fast\")\n",
"axarr[0][1].axis(\"off\")\n",
"axarr[0][2].imshow(\n",
" lens_basic.reduced_deflection_angle(thx, thy)[0].numpy()\n",
" - lens_fast.reduced_deflection_angle(thx, thy)[0].numpy(),\n",
" origin=\"lower\",\n",
")\n",
"axarr[0][2].set_title(\"Difference\")\n",
"axarr[0][2].axis(\"off\")\n",
"axarr[1][0].imshow(\n",
" lens_basic.reduced_deflection_angle(thx, thy)[1].numpy(), origin=\"lower\"\n",
")\n",
"axarr[1][0].set_title(\"Deflection y basic\")\n",
"axarr[1][0].axis(\"off\")\n",
"axarr[1][1].imshow(\n",
" lens_fast.reduced_deflection_angle(thx, thy)[1].numpy(), origin=\"lower\"\n",
")\n",
"axarr[1][1].set_title(\"Deflection y fast\")\n",
"axarr[1][1].axis(\"off\")\n",
"axarr[1][2].imshow(\n",
" lens_basic.reduced_deflection_angle(thx, thy)[1].numpy()\n",
" - lens_fast.reduced_deflection_angle(thx, thy)[1].numpy(),\n",
" origin=\"lower\",\n",
")\n",
"axarr[1][2].set_title(\"Difference\")\n",
"axarr[1][2].axis(\"off\")\n",
"fig.suptitle(\"Comparison of basic and fast lensing, the two are identical\", fontsize=16)\n",
"plt.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%%timeit\n",
"ax, ay = lens_basic.reduced_deflection_angle(thx, thy)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%%timeit\n",
"ax, ay = lens_fast.reduced_deflection_angle(thx, thy)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Here we see that our new fast version is much faster (almost 10x faster) than the basic one which only uses automatic differentiation from the potential. There are a few reasons for this, in the most straightforward setups it is normal for autograd to be about 2-3x slower than an analytic derivation. Further, because in this case there are many shared calculations between `ax` and `ay`, we were able to save ourselves a bunch of calculations by only doing the shared stuff once.\n",
"\n",
"Next lets look at the convergence, which uses the Hessian of the potential. Below we compare the two ways of computing the convergence, one using autograd and the other using analytic derivatives. We see that the two are nearly identical, the residuals are at the level of `10^-7` which is the precision of floating point operations. Thus the two are identical up to the level that we can tell with our current numerical precision."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"hide-input"
]
},
"outputs": [],
"source": [
"fig, axarr = plt.subplots(1, 3, figsize=(15, 5))\n",
"axarr[0].imshow(lens_basic.convergence(thx, thy).numpy(), origin=\"lower\")\n",
"axarr[0].set_title(\"Convergence basic\")\n",
"axarr[0].axis(\"off\")\n",
"axarr[1].imshow(lens_fast.convergence(thx, thy).numpy(), origin=\"lower\")\n",
"axarr[1].set_title(\"Convergence fast\")\n",
"axarr[1].axis(\"off\")\n",
"im = axarr[2].imshow(\n",
" lens_basic.convergence(thx, thy).numpy() - lens_fast.convergence(thx, thy).numpy(),\n",
" origin=\"lower\",\n",
")\n",
"fig.colorbar(im, ax=axarr[2])\n",
"axarr[2].set_title(\"Difference\")\n",
"axarr[2].axis(\"off\")\n",
"fig.suptitle(\n",
" \"Comparison of basic and fast convergence, the two are identical\", fontsize=16\n",
")\n",
"plt.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%%timeit\n",
"kappa = lens_basic.convergence(thx, thy)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%%timeit\n",
"kappa = lens_fast.convergence(thx, thy)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Since the convergence uses second derivatives, we see an even more dramatic difference between our basic autograd from the potential and an analytic calculation. It is now almost 30x faster, which is another factor of 3 because of the extra autograd operation needed for the basic calculation.\n",
"\n",
"The conclusion here is that using autograd from the potential is easy and reasonably fast, but if performance is a significant value then its worth doing the extra work to get the derivatives yourself. In `caustics` all of the base models have analytic `potential`, `deflection_angle`, and `convergence` so that it is as performant as possible. If you can use a built-in method of `caustics` then it is worth doing so, but if you need to make your own model, you now know how to make it as fast as possible!"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "PY39",
"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.9.5"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Loading