Multi-Source Modelling#
Sometimes a lensing system will have aligned with multiple sources of light. These will be at different source redshifts and so some consideration must be taken to model the image of both objects. Here we will demonstrate a case of a single gravitational lens projecting two sources.
%load_ext autoreload
%autoreload 2
%matplotlib inline
import matplotlib.pyplot as plt
from numpy import pi
import torch
from torch.nn.functional import avg_pool2d
import caustics
from caustics import Module, forward, Param
Lets define the objects we will need in the lensing system. The cosmology and lens are like normal, but now we have two sources. Also, the Einstein radius of the lens is usually set to a single value, but that value is encoding information about the lens configuration and is related to the redshift of the source, this only makes sense for a single plane system. We can’t conveniently pack that away anymore, so we set the SIE to have a physical parametrization (velocity dispersion) and then caustics will work out what the Einstein radius should be relative to each source.
cosmology = caustics.FlatLambdaCDM()
# Set z_s to nan so we know if it gets accidentally used somewhere
# Set the lens to a physical parametrization so multi-source works properly
lens = caustics.SIE(
cosmology=cosmology,
name="lens",
z_l=0.5,
z_s=torch.nan,
x0=0.0,
y0=0.0,
q=0.5,
phi=0.0,
sigma_v=250,
parametrization="velocity_dispersion",
)
src1 = caustics.Sersic(
name="source1", x0=0.0, y0=0.0, q=0.6, phi=pi / 3, n=4.0, Re=0.4, Ie=0.5
)
src2 = caustics.Sersic(
name="source2", x0=0.1, y0=0.2, q=0.3, phi=-pi / 4, n=2.0, Re=1.5, Ie=1.5
)
# Define the pixel grid for imaging
n_pix = 100
res = 0.05
upsample_factor = 3
Below we define a new LensSource Module, except that now it takes two sources, so it is LensTwoSources. During runtime the source redshift for the lens is overridden with the two values. For each source the lensing proceeds like normal.
class LensTwoSources(Module):
def __init__(self, lens, src1, src2, z_s1, z_s2):
super().__init__()
self.lens = lens
self.src1 = src1
self.src2 = src2
self.z_s1 = Param("z_s1", z_s1)
self.z_s2 = Param("z_s2", z_s2)
theta_x, theta_y = caustics.utils.meshgrid(
res / upsample_factor,
upsample_factor * n_pix,
dtype=torch.float32,
)
self.theta_x = theta_x
self.theta_y = theta_y
@forward
def sample(self, z_s1, z_s2, source1=True, source2=True):
mu = torch.zeros_like(self.theta_x)
if source1:
# Note we override the lens z_s parameter just for this one raytrace call
bx, by = self.lens.raytrace(self.theta_x, self.theta_y, z_s=z_s1)
mu += self.src1.brightness(bx, by)
if source2:
# Note here we override z_s with a different value
bx, by = self.lens.raytrace(self.theta_x, self.theta_y, z_s=z_s2)
mu += self.src2.brightness(bx, by)
return avg_pool2d(mu[None][None], upsample_factor).squeeze()
lts_sim = LensTwoSources(lens, src1, src2, z_s1=0.8, z_s2=1.5)
# Sample the image
img = lts_sim.sample()
