tests: Add consistency test between lensing potential and deflection …

…angle, convergence (#138)

* feat: add _null_params for basic lens evaluation

* test: add test for consistency between lens potential and deflection angle

* test: add potential to convergence test

* fix: lower tolerance for NFW test

* TNFW potential now correct

* PseudoJaffe now passes tests

* PixelatedConvergence now kind of works with low tolerance

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,

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,

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,

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,

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,

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,

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

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,

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,

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

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)