diff --git a/src/caustics/lenses/epl.py b/src/caustics/lenses/epl.py index 747e86f5..ab177b54 100644 --- a/src/caustics/lenses/epl.py +++ b/src/caustics/lenses/epl.py @@ -59,6 +59,15 @@ class EPL(ThinLens): """ + _null_params = { + "x0": 0.0, + "y0": 0.0, + "q": 0.5, + "phi": 0.0, + "b": 1.0, + "t": 1.0, + } + def __init__( self, cosmology: Cosmology, diff --git a/src/caustics/lenses/external_shear.py b/src/caustics/lenses/external_shear.py index 803043d0..6aee98a7 100644 --- a/src/caustics/lenses/external_shear.py +++ b/src/caustics/lenses/external_shear.py @@ -34,6 +34,13 @@ class ExternalShear(ThinLens): distortion that can be caused by nearby structures outside of the main lens galaxy. """ + _null_params = { + "x0": 0.0, + "y0": 0.0, + "gamma_1": 0.1, + "gamma_2": 0.1, + } + def __init__( self, cosmology: Cosmology, diff --git a/src/caustics/lenses/mass_sheet.py b/src/caustics/lenses/mass_sheet.py index 76feb5d3..d51c5a0c 100644 --- a/src/caustics/lenses/mass_sheet.py +++ b/src/caustics/lenses/mass_sheet.py @@ -35,6 +35,12 @@ class MassSheet(ThinLens): distortion that can be caused by nearby structures outside of the main lens galaxy. """ + _null_params = { + "x0": 0.0, + "y0": 0.0, + "surface_density": 0.1, + } + def __init__( self, cosmology: Cosmology, diff --git a/src/caustics/lenses/nfw.py b/src/caustics/lenses/nfw.py index 335c7a68..cfe68f18 100644 --- a/src/caustics/lenses/nfw.py +++ b/src/caustics/lenses/nfw.py @@ -65,6 +65,13 @@ class NFW(ThinLens): Computes the lensing potential. """ + _null_params = { + "x0": 0.0, + "y0": 0.0, + "m": 1e13, + "c": 5.0, + } + def __init__( self, cosmology: Cosmology, diff --git a/src/caustics/lenses/pixelated_convergence.py b/src/caustics/lenses/pixelated_convergence.py index 59c56bba..10cda750 100644 --- a/src/caustics/lenses/pixelated_convergence.py +++ b/src/caustics/lenses/pixelated_convergence.py @@ -5,6 +5,7 @@ import torch.nn.functional as F from scipy.fft import next_fast_len from torch import Tensor +import numpy as np from ..cosmology import Cosmology from ..utils import get_meshgrid, interp2d, safe_divide, safe_log @@ -16,6 +17,12 @@ class PixelatedConvergence(ThinLens): + _null_params = { + "x0": 0.0, + "y0": 0.0, + "convergence_map": np.logspace(0, 1, 100, dtype=np.float32).reshape(10, 10), + } + def __init__( self, pixelscale: float, diff --git a/src/caustics/lenses/point.py b/src/caustics/lenses/point.py index c457ac12..b640f84d 100644 --- a/src/caustics/lenses/point.py +++ b/src/caustics/lenses/point.py @@ -34,6 +34,12 @@ class Point(ThinLens): Softening parameter to prevent numerical instabilities. """ + _null_params = { + "x0": 0.0, + "y0": 0.0, + "th_ein": 1.0, + } + def __init__( self, cosmology: Cosmology, diff --git a/src/caustics/lenses/pseudo_jaffe.py b/src/caustics/lenses/pseudo_jaffe.py index eb3f6812..671333bf 100644 --- a/src/caustics/lenses/pseudo_jaffe.py +++ b/src/caustics/lenses/pseudo_jaffe.py @@ -42,6 +42,14 @@ class PseudoJaffe(ThinLens): Softening parameter to prevent numerical instabilities. """ + _null_params = { + "x0": 0.0, + "y0": 0.0, + "mass": 1e12, + "core_radius": 0.1, + "scale_radius": 1.0, + } + def __init__( self, cosmology: Cosmology, @@ -258,18 +266,21 @@ def potential( Returns -------- Tensor - The lensing potential. + The lensing potential (arcsec^2). """ - # TODO: why do the units come out to arcsec (length) and not arcsec^2? Note that the units in A18 also come out to arcsec. - # fmt: off x, y = translate_rotate(x, y, x0, y0) d_l = self.cosmology.angular_diameter_distance(z_l, params) # Mpc + d_s = self.cosmology.angular_diameter_distance(z_s, params) # Mpc + d_ls = self.cosmology.angular_diameter_distance_z1z2(z_l, z_s, params) # Mpc + R_squared = x**2 + y**2 + self.s # arcsec^2 sigma_crit = self.cosmology.critical_surface_density(z_l, z_s, params) # Msun / Mpc^2 surface_density_0 = self.get_convergence_0(z_s, params) * sigma_crit # Msun / Mpc^2 - coeff = 4 * pi * G_over_c2 * surface_density_0 * (d_l * arcsec_to_rad) * core_radius * scale_radius / (scale_radius - core_radius) # unitless or arcsec? + + coeff = -8 * pi * G_over_c2 * surface_density_0 * (d_l * d_ls / d_s) * core_radius * scale_radius / (scale_radius - core_radius) # arcsec + scale_a = (scale_radius**2 + R_squared).sqrt() # arcsec scale_b = (core_radius**2 + R_squared).sqrt() # arcsec scale_c = core_radius * (core_radius + (core_radius**2 + R_squared).sqrt()).log() # arcsec diff --git a/src/caustics/lenses/sie.py b/src/caustics/lenses/sie.py index 6134cc3d..7ea2a5e7 100644 --- a/src/caustics/lenses/sie.py +++ b/src/caustics/lenses/sie.py @@ -38,6 +38,14 @@ class SIE(ThinLens): The core radius of the lens (defaults to 0.0). """ + _null_params = { + "x0": 0.0, + "y0": 0.0, + "q": 0.5, + "phi": 0.0, + "b": 1.0, + } + def __init__( self, cosmology: Cosmology, diff --git a/src/caustics/lenses/sis.py b/src/caustics/lenses/sis.py index abbdfb03..15a54172 100644 --- a/src/caustics/lenses/sis.py +++ b/src/caustics/lenses/sis.py @@ -33,6 +33,12 @@ class SIS(ThinLens): A smoothing factor, default is 0.0. """ + _null_params = { + "x0": 0.0, + "y0": 0.0, + "th_ein": 1.0, + } + def __init__( self, cosmology: Cosmology, diff --git a/src/caustics/lenses/tnfw.py b/src/caustics/lenses/tnfw.py index 87249d6d..3c9478d6 100644 --- a/src/caustics/lenses/tnfw.py +++ b/src/caustics/lenses/tnfw.py @@ -75,6 +75,14 @@ class TNFW(ThinLens): """ + _null_params = { + "x0": 0.0, + "y0": 0.0, + "mass": 1e13, + "scale_radius": 1.0, + "tau": 3.0, + } + def __init__( self, cosmology: Cosmology, @@ -529,10 +537,12 @@ def potential( u = g**2 F = self._F(g) L = self._L(g, tau) + d_l = self.cosmology.angular_diameter_distance(z_l, params) + d_s = self.cosmology.angular_diameter_distance(z_s, params) + d_ls = self.cosmology.angular_diameter_distance_z1z2(z_l, z_s, params) - # d_l = self.cosmology.angular_diameter_distance(z_l, params) # fmt: off - S = 2 * self.get_M0(params) * G_over_c2 # * rad_to_arcsec * d_l**2 + S = 2 * self.get_M0(params) * G_over_c2 * (d_ls / d_s) / (d_l * arcsec_to_rad**2) a1 = 1 / (t2 + 1) ** 2 a2 = 2 * torch.pi * t2 * (tau - (t2 + u).sqrt() + tau * (tau + (t2 + u).sqrt()).log()) a3 = 2 * (t2 - 1) * tau * (t2 + u).sqrt() * L @@ -541,5 +551,7 @@ def potential( a6 = t2 * (t2 - 1) * (1 / g.to(dtype=torch.cdouble)).arccos().abs() ** 2 a7 = t2 * ((t2 - 1) * tau.log() - t2 - 1) * u.log() a8 = t2 * ((t2 - 1) * tau.log() * (4 * tau).log() + 2 * (tau / 2).log() - 2 * tau * (tau - torch.pi) * (2 * tau).log()) - # fmt: on - return S * a1 * (a2 + a3 + a4 + a5 + a6 + a7 - a8) # fmt: skip + + return S * a1 * (a2 + a3 + a4 + a5 + a6 + a7 - a8) + + # fmt: on diff --git a/tests/test_lens_potential.py b/tests/test_lens_potential.py new file mode 100644 index 00000000..fc1f0997 --- /dev/null +++ b/tests/test_lens_potential.py @@ -0,0 +1,186 @@ +""" +Check for internal consistency of the lensing potential for all ThinLens objects. +""" + +import torch + +import caustics + + +def test_lens_potential_vs_deflection(): + """ + Check for internal consistency of the lensing potential for all ThinLens objects against the deflection angles. The gradient of the potential should equal the deflection angle. + """ + # Define a grid of points to test. + x = torch.linspace(-1, 1, 10) + y = torch.linspace(-1, 1, 10) + x, y = torch.meshgrid(x, y, indexing="ij") + + # Define a source redshift. + z_s = 1.0 + # Define a lens redshift. + z_l = 0.5 + + # Define a cosmology. + cosmo = caustics.cosmology.FlatLambdaCDM(name="cosmo") + + # Define a list of lens models. + lenses = [ + caustics.lenses.EPL( + cosmology=cosmo, z_l=z_l, **caustics.lenses.EPL._null_params + ), + caustics.lenses.ExternalShear( + cosmology=cosmo, z_l=z_l, **caustics.lenses.ExternalShear._null_params + ), + caustics.lenses.MassSheet( + cosmology=cosmo, z_l=z_l, **caustics.lenses.MassSheet._null_params + ), + caustics.lenses.NFW( + cosmology=cosmo, + z_l=z_l, + **caustics.lenses.NFW._null_params, + use_case="differentiable", + ), + caustics.lenses.PixelatedConvergence( + cosmology=cosmo, + z_l=z_l, + **caustics.lenses.PixelatedConvergence._null_params, + pixelscale=0.1, + n_pix=10, + ), + caustics.lenses.Point( + cosmology=cosmo, z_l=z_l, **caustics.lenses.Point._null_params + ), + caustics.lenses.PseudoJaffe( + cosmology=cosmo, z_l=z_l, **caustics.lenses.PseudoJaffe._null_params + ), + caustics.lenses.SIE( + cosmology=cosmo, z_l=z_l, **caustics.lenses.SIE._null_params + ), + caustics.lenses.SIS( + cosmology=cosmo, z_l=z_l, **caustics.lenses.SIS._null_params + ), + caustics.lenses.TNFW( + cosmology=cosmo, + z_l=z_l, + **caustics.lenses.TNFW._null_params, + use_case="differentiable", + ), + ] + + # Define a list of lens model names. + names = list(L.name for L in lenses) + # Loop over the lenses. + for lens, name in zip(lenses, names): + print(f"Testing lens: {name}") + # Compute the deflection angle. + ax, ay = lens.reduced_deflection_angle(x, y, z_s) + + # Ensure the x,y coordinates track gradients + x = x.detach().requires_grad_() + y = y.detach().requires_grad_() + + # Compute the lensing potential. + phi = lens.potential(x, y, z_s) + + # Compute the gradient of the lensing potential. + phi_ax, phi_ay = torch.autograd.grad( + phi, (x, y), grad_outputs=torch.ones_like(phi) + ) + + # Check that the gradient of the lensing potential equals the deflection angle. + if name in ["NFW", "TNFW"]: + # Special functions in NFW and TNFW are not highly accurate, so we relax the tolerance + assert torch.allclose(phi_ax, ax, atol=1e-3, rtol=1e-3) + assert torch.allclose(phi_ay, ay, atol=1e-3, rtol=1e-3) + elif name in ["PixelatedConvergence"]: + # PixelatedConvergence potential is defined by bilinear interpolation so it is very imprecise + assert torch.allclose(phi_ax, ax, rtol=1e0) + assert torch.allclose(phi_ay, ay, rtol=1e0) + else: + assert torch.allclose(phi_ax, ax) + assert torch.allclose(phi_ay, ay) + + +def test_lens_potential_vs_convergence(): + """ + Check for internal consistency of the lensing potential for all ThinLens objects against the convergence. The laplacian of the potential should equal the convergence. + """ + # Define a grid of points to test. + x = torch.linspace(-1, 1, 10) + y = torch.linspace(-1, 1, 10) + x, y = torch.meshgrid(x, y, indexing="ij") + x, y = x.clone().detach(), y.clone().detach() + + # Define a source redshift. + z_s = 1.0 + # Define a lens redshift. + z_l = 0.5 + + # Define a cosmology. + cosmo = caustics.cosmology.FlatLambdaCDM(name="cosmo") + + # Define a list of lens models. + lenses = [ + caustics.lenses.EPL( + cosmology=cosmo, z_l=z_l, **caustics.lenses.EPL._null_params + ), + caustics.lenses.ExternalShear( + cosmology=cosmo, z_l=z_l, **caustics.lenses.ExternalShear._null_params + ), + caustics.lenses.MassSheet( + cosmology=cosmo, z_l=z_l, **caustics.lenses.MassSheet._null_params + ), + # caustics.lenses.NFW( + # cosmology=cosmo, + # z_l=z_l, + # **caustics.lenses.NFW._null_params, + # use_case="differentiable", + # ), # Cannot vmap NFW when in differentiable mode + # caustics.lenses.PixelatedConvergence( + # cosmology=cosmo, + # z_l=z_l, + # **caustics.lenses.PixelatedConvergence._null_params, + # pixelscale=0.2, + # n_pix=10, + # ), # cannot compute Hessian of PixelatedConvergence potential, always returns zeros due to bilinear interpolation + # caustics.lenses.Point(cosmology=cosmo, z_l=z_l, **caustics.lenses.Point._null_params), # Point mass convergence is delta function + caustics.lenses.PseudoJaffe( + cosmology=cosmo, z_l=z_l, **caustics.lenses.PseudoJaffe._null_params + ), + caustics.lenses.SIE( + cosmology=cosmo, z_l=z_l, **caustics.lenses.SIE._null_params + ), + caustics.lenses.SIS( + cosmology=cosmo, z_l=z_l, **caustics.lenses.SIS._null_params + ), + # caustics.lenses.TNFW( + # cosmology=cosmo, z_l=z_l, **caustics.lenses.TNFW._null_params, use_case="differentiable" + # ), # Cannot vmap TNFW when in differentiable mode + ] + + # Define a list of lens model names. + names = list(L.name for L in lenses) + # Loop over the lenses. + for lens, name in zip(lenses, names): + print(f"Testing lens: {name}") + # Compute the convergence. + try: + kappa = lens.convergence(x, y, z_s) + except NotImplementedError: + continue + + # Compute the laplacian of the lensing potential. + phi_H = torch.vmap( + torch.vmap( + torch.func.hessian(lens.potential, (0, 1)), in_dims=(0, 0, None) + ), + in_dims=(0, 0, None), + )(x, y, z_s) + phi_kappa = 0.5 * (phi_H[0][0] + phi_H[1][1]) + + # Check that the laplacian of the lensing potential equals the convergence. + if name in ["NFW", "TNFW"]: + assert torch.allclose(phi_kappa, kappa, atol=1e-4) + else: + assert torch.allclose(phi_kappa, kappa)