Caustics Interface: YAML

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 yaml interface.

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 using the .yaml file method. 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

Using the YAML File Method

The YAML File Method for building a simulator greatly simplifies the process by providing a clear, human-readable, and structured format to define all the components of the simulator. It allows for easy reuse of defined components, such as cosmology or lenses, and makes the configuration of the simulator easily adjustable without touching the actual code. This method enhances maintainability, readability, and scalability of the simulator building process.

# Build the simulator
sim = caustics.build_simulator("example.yml")

# Take a look at the yaml file
!cat example.yml

cosmology: &cosmo
  name: cosmo
  kind: FlatLambdaCDM

lens: &lens
  name: lens
  kind: SIE
  init_kwargs:
    cosmology: *cosmo

src: &src
  name: source
  kind: Sersic

lnslt: &lnslt
  name: lenslight
  kind: Sersic

simulator:
  name: minisim
  kind: LensSource
  init_kwargs:
    # Single lense
    lens: *lens
    source: *src
    lens_light: *lnslt
    pixelscale: 0.05
    pixels_x: 100
    quad_level: 3

# Print out the order of model parameters
# Note that parameters with values are "static" so they don't need to be provided by you
print(sim)

simulator|LensSource
    psf|static: 1
    x0|static: 0
    y0|static: 0
    lens|SIE
        cosmology|FlatLambdaCDM
            h0|static: 0.677
            critical_density_0|static: 1.27e+11
            Om0|static: 0.31
        z_s|dynamic
        z_l|dynamic
        x0|dynamic
        y0|dynamic
        q|dynamic
        phi|dynamic
        Rein|dynamic
    src|Sersic
        x0|dynamic
        y0|dynamic
        q|dynamic
        phi|dynamic
        n|dynamic
        Re|dynamic
        Ie|dynamic
    lnslt|Sersic
        x0|dynamic
        y0|dynamic
        q|dynamic
        phi|dynamic
        n|dynamic
        Re|dynamic
        Ie|dynamic

# Here we build a tensor with the parameters of the model
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

Plot the Results!

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

# Here we sample an image
image = sim(x).detach().cpu().numpy()

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

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 behaves essentially like 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.

# Now lets sample 20 images by batching the operation using vmap
newx = x.repeat(20, 1)
newx += torch.normal(mean=0, std=0.1 * torch.ones_like(newx))

images = torch.vmap(sim)(newx)

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.

# Now lets compute the jacobian of the simulator wrt each parameter
J = torch.func.jacfwd(sim)(x)
# The shape of J is (npixels y, npixels x, nparameters)

Show code cell source

Hide code cell source

fig, axarr = plt.subplots(3, 7, figsize=(20, 9))
labels = tuple(p.name for p in sim.dynamic_params)
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()

The Simulator Graph

Here we take a quick look at the simulator graph for the image we have produced. You will learn much more about what this means in the Simulators tutorial notebook, but let’s cover the basics here. First, note that this is a Directed Acyclic Graph (DAG), this is how all simulator parameters are represented in caustics. At the top of the graph is the LensSource object, you can see in brackets it has a name sim which is used as the identifier for it’s node in the graph. At the next level is the z_s parameter for the redshift of the source. Next are the SIE lens, Sersic source, and Sersic lenslight objects which themselves each hold parameters. You will notice that all the parameters are in white boxes right now, this is because they are dynamic parameters which need values to be passed, grey boxes are used for parameters with fixed values.

# The simulator is internally represented as a directed acyclic graph of operations
sim.graphviz()

Model Parameters

Each of the lens, source, and lenslight models have their own parameters that are needed to sample a given lensing configuration. There are a number of ways to pass these parameters to a caustics simulator, but the most straightforward for most purposes is as a Pytorch Tensor. In this cell, we build the tensor manually for our configuration. In general one would have a script generate the configuration, or an optimizer/sampler would randomly choose parameter values.

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.

Breaking down the YAML file

Here we will go over the parts of the yaml file.

Cosmology

In the YAML file, the cosmology is defined under the cosmology key. This section is used to specify the cosmological model that will be used in the simulator.

Here’s the corresponding section from the YAML file:

cosmology: &cosmo
    name: cosmo
    kind: FlatLambdaCDM

In this section:

• cosmology: is the key that starts the definition of the cosmology.

• &cosmo is a YAML anchor that allows us to give a name to this section for later reference.

• name: cosmo sets the name of the cosmology to cosmo.

• kind: FlatLambdaCDM sets the kind of the cosmology to FlatLambdaCDM, which is a cosmological model assuming a flat Universe dominated by a cosmological constant (Lambda) and cold dark matter (CDM).

This cosmology definition can be referenced elsewhere in the YAML file using the *cosmo alias. This allows us to reuse the same cosmology definition for multiple lenses or other components without having to redefine it each time.

Lens Mass Distribution

In order for gravitational lensing to occur, we need some mass to bend the light. Here we define a basic Singular Isothermal Ellipsoid (SIE), which is a versatile profile used in many strong gravitational lensing simulations. As the first argument, we pass the cosmology so that the SIE can compute various quantities which make use of redshift information (seen later). Each model must have a unique name so we call this one lens though you can also let caustics automatically pick a unique name.

In the YAML file, the lens mass distribution is defined under the lens key. This section is used to specify the properties of the lens that will be used in the simulator.

Here’s the corresponding section from the YAML file for lens:

lens: &lens
    name: lens
    kind: SIE
    init_kwargs:
        cosmology: *cosmo

In this section:

• lens: is the key that starts the definition of the first lens.

• &lens is a YAML anchor that allows us to give a name to this section for later reference.

• name: lens sets the name of the lens to ‘lens’.

• kind: SIE sets the kind of the lens to ‘SIE’, which stands for Singular Isothermal Ellipsoid, a common model for the mass distribution of a lens in gravitational lensing.

• init_kwargs: is a key for additional parameters required for the lens.

• cosmology: *cosmo sets the cosmology for the lens to the cosmology defined earlier in the YAML file.

Source Light Distribution

If we wish to see anything in our lensing configuration then we need a bright object in the background to produce some light that will pass through (and be bent by) our lens mass distribution. Here we create a Sersic light model which is a common versatile profile for representing galaxies. Note that we don’t need to pass a light model any Cosmology information, since light models essentially just define a function on (x,y) coordinates that gives a brightness, the lens models handle all cosmology related calculations. For the name we very creatively choose source.

In the YAML file, the source light distribution is defined under the src key. This section is used to specify the properties of the source that will be used in the simulator.

Here’s the corresponding section from the YAML file:

src: &src
    name: source
    kind: Sersic

In this section:

• src: is the key that starts the definition of the source.

• &src is a YAML anchor that allows us to give a name to this section for later reference.

• name: source sets the name of the source to ‘source’.

• kind: Sersic sets the kind of the source to ‘Sersic’, which is a common model for the light distribution of a source in astronomical imaging.

This source definition can be referenced elsewhere in the YAML file using the *src alias. This allows us to reuse the same source definition for multiple components without having to redefine it each time.

Lens Light Distribution

The source isn’t the only bright thing in the sky! The lensing galaxy itself will also have bright stars and can be seen as well. Let’s add another Sersic model with the name lenslight.

In the YAML file, the lens light distribution is defined under the lnslt key. This section is used to specify the properties of the lens light that will be used in the simulator.

Here’s the corresponding section from the YAML file:

lnslt: &lnslt
    name: lenslight
    kind: Sersic

In this section:

• lnslt: is the key that starts the definition of the lens light.

• &lnslt is a YAML anchor that allows us to give a name to this section for later reference.

• name: lenslight sets the name of the lens light to ‘lenslight’.

• kind: Sersic sets the kind of the lens light to ‘Sersic’, which is a common model for the light distribution of a lens in astronomical imaging.

This lens light definition can be referenced elsewhere in the YAML file using the *lnslt alias. This allows us to reuse the same lens light definition for multiple components without having to redefine it each time.

LensSource Simulator

Next we pass our configuration to a Simulator in caustics, simulators perform the work of forward modelling various configurations and producing the desired outputs. Here we are interested in a common scenario of producing an image of a background source through a lens distribution. It is possible to make your own simulator to represent all sorts of situations. First, we pass the lens model and the source model defined above. Next we use pixelscale and pixels_x to define the grid of pixels that will be sampled. Finally, we pass the z_s redshift at which the source (Sersic) model should be placed; recall that light models don’t use the cosmology model and so aren’t aware of their placement in space.

In the YAML file, the simulator is defined under the simulator key. This section is used to specify the properties of the simulator that will be used to perform the simulation.

Here’s the corresponding section from the YAML file:

simulator:
    name: minisim
    kind: LensSource
    init_kwargs:
        lens: *lens
        source: *src
        lens_light: *lnslt
        pixelscale: 0.05
        pixels_x: 100

In this section:

• simulator: is the key that starts the definition of the simulator.

• name: minisim sets the name of the simulator to ‘minisim’.

• kind: LensSource sets the kind of the simulator to ‘LensSource’, which indicates that this simulator will simulate a lens and a source.

• init_kwargs: is a key for additional parameters required for the simulator.

• lens: *lens sets the lens for the simulator to the lens defined earlier in the YAML file.

• source: *src sets the source for the simulator to the source defined earlier in the YAML file.

• lens_light: *lnslt sets the lens light for the simulator to the lens light defined earlier in the YAML file.

• pixelscale: 0.05 sets the pixel scale for the simulator to 0.05.

• pixels_x: 100 sets the number of pixels in the x-direction for the simulator to 100.

This section defines the simulator that will be used to perform the simulation. It references the lens, source, and lens light definitions from earlier in the YAML file and sets additional parameters for the simulator.