Caustics Interface: Functional

caustics is a powerful gravitational lensing simulator that can support users from beginner to highly advanced. In this tutorial we will cover the basics of the caustics functional interface. This one is a bit different from the other two since we are building from the ground up.

Simulating an SIE lens

Here we will demo the very basics of lensing with a classic SIE lens model. We will see what it takes to make an SIE model, lens a background Sersic source, and sample the resulting image using a Simulator. caustics simulators can generalize to very complex scenarios, here we will use a built-in simulator which handles a common use case (lensing a background source). To start, we of course need to import some modules. For the minimal example, this is just matplotlib a common package used for plotting, torch which is a numerical package for GPU/autodiff (much like numpy), and caustics the reason you are here.

In this tutorial, we will guide you through the process of simulating an SIE lens the functional method, which is a bit different than the other two. This tutorial is mirrored in two other tutorials so you can see the yaml, object oriented, and functional interfaces.

First, let’s import the necessary packages:

Import the Necessary Packages

Note: These packages need to be imported for any method

import matplotlib.pyplot as plt
import torch
import caustics

Define sampling grid

To replicate the other mirrored tutorials, here we will sample on the same 100x100 grid. In those examples we do gaussian quadrature sub pixel integration so we will show how to do that here as well. The coordinates specified here are what we will eventually sample on in the simulator.

# First lets define the sampling grid
npix = 100
pixelscale = 0.05
thx, thy = caustics.utils.meshgrid(pixelscale, npix, npix)
gqx, gqy, gqW = caustics.utils.gaussian_quadrature_grid(
    pixelscale, thx, thy, quad_level=3
)

Simulation parameters

Here we use the same parameters as the other tutorials. Note that z_s and z_l are not actually needed, but we keep them anyway for consistency. In the caustics object oriented framework, some lenses need to know their redshift so to keep a uniform interface we have all lenses require a redshift.

x = torch.tensor([
#   z_s  z_l   x0   y0   q    phi     Rein 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, 
#   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

Build a simulator

Here we manually construct our own simulator using caustics.func. This will perform exactly the same operations as the other introduction tutorials, using the same functions. The only difference is that this version is more opaque and harder to debug since we have done everything manually. Note that it is more fragile too, internally we have specified the indices of x that are needed for the lens and light distributions, if we swapped out the sie with and nfw then all the indices would need to change. This is why using the yaml or oop caustics interfaces is easier to play around with different profiles and configurations without everything breaking. Also note that we aren’t doing PSF convolution, that would add another layer of complexity and would be tedious to do here, though it is very easy in the other interfaces.

Still, this is such a simple function it represents the optimal use of compute resources, it will be faster than the other methods. Further, you have complete freedom to change whatever you like! This makes it ideal for exploring new ideas.

def sim(x):
    # Compute deflection angles
    ax, ay = caustics.func.reduced_deflection_angle_sie(*x[2:7], gqx, gqy)
    # Raytrace with lens equation
    bx, by = gqx - ax, gqy - ay
    # Lens background source light
    k = caustics.func.k_sersic(x[11])
    mu_fine = caustics.func.brightness_sersic(*x[7:14], bx, by, k)
    mu = caustics.utils.gaussian_quadrature_integrator(mu_fine, gqW)
    # Add lens light
    k = caustics.func.k_sersic(x[18])
    mu_fine = caustics.func.brightness_sersic(*x[14:], gqx, gqy, k)
    mu += caustics.utils.gaussian_quadrature_integrator(mu_fine, gqW)
    return mu

Plot the Results!

This section is mostly self explanatory. We evaluate the simulator configuration by calling it like a function.

# Substitute sim with sim for yaml method
image = sim(x).detach().cpu().numpy()

plt.imshow(image, origin="lower")
plt.axis("off")
plt.show()

The Simulator Graph

There is no simulator graph for our manually constructed simulator function, not unless we make it!

Sampling with a Simulator

Now let’s see how to use some of the powerful features of the simulator we have created. Note that it is a function, allowing us to take advantage of many PyTorch features. To start, lets see how we can run batches of lens simulations using vmap.

newx = x.repeat(20, 1)
newx += torch.normal(mean=0, std=0.1 * torch.ones_like(newx))

images = torch.vmap(sim)(newx)  # Substitute minisim with sim for the yaml method

fig, axarr = plt.subplots(4, 5, figsize=(20, 16))
for ax, im in zip(axarr.flatten(), images):
    ax.imshow(im, origin="lower")
    ax.axis("off")
plt.tight_layout()
plt.show()

Gradients with autodiff

Batching is useful for fully parallelizing code and maximally using computational resources, but autodiff gradients allow whole new algorithms and techniques to be used in gravitational lensing! Let’s try computing the Jacobian for the lensing configuration that we have been using so far. The result is a grid of images that show how the lensing simulation image would change if we adjusted each parameter individually. Thanks to autodiff, these derivatives have no finite differences approximation error, they are exact up to the machine precision.

J = torch.func.jacfwd(sim)(x)

Show code cell source

Hide code cell source

fig, axarr = plt.subplots(3, 7, figsize=(20, 9))
labels = [
    "z_s",
    "z_l",
    "lens x0",
    "lens y0",
    "lens q",
    "lens phi",
    "lens Rein",
    "source x0",
    "source y0",
    "source q",
    "source phi",
    "source n",
    "source Re",
    "source Ie",
    "lenslight x0",
    "lenslight y0",
    "lenslight q",
    "lenslight phi",
    "lenslight n",
    "lenslight Re",
    "lenslight Ie",
]
for i, ax in enumerate(axarr.flatten()):
    ax.imshow(J[..., i], origin="lower")
    ax.set_title(labels[i])
    ax.axis("off")
plt.tight_layout()
plt.show()

Model Parameters

In order, here is an explanation of the parameters.

• z_s is the redshift of the source.

• z_l is the lens redshift which tells the lens how far away it is from the observer (us).

• The next two parameters x0 and y0 indicate where the lens is relative to the main optical axis, which is the coordinates (0, 0).

• The q parameter gives the axis ratio for the SIE, so it knows how elongated it is.

• phi indicates the position angle (where the ellipse is pointing).

• Rein gives the Einstein radius (in arcsec) of the lens.

• The next x0 and y0 provide the position relative to the main optical axis of the Sersic source, here we offset the source slightly to make for an interesting figure.

• The q parameter defines the axis ratio of the Sersic ellipse.

• phi defines the position angle of the ellipse.

• n is the Sersic index which determines how concentrated the light is; n=0.5 is a Gaussian distribution, n=1 is an exponential, n=4 is a De Vaucouleurs profile.

• Re is the radius (in arcsec) within which half the total light of the profile is enclosed.

• Ie is the brightness at Re.

• The next set of parameters are also Sersic parameters, but this time they are for the lens light model.