feat: add multipole lens profile

src/caustics/__init__.py

@@ -22,6 +22,7 @@
     SinglePlane,
     MassSheet,
     TNFW,
+    Multipole,
 )
 from .light import Source, Pixelated, Sersic  # PROBESDataset conflicts with .data
 from .data import HDF5Dataset, IllustrisKappaDataset, PROBESDataset
@@ -53,6 +54,7 @@
     "SinglePlane",
     "MassSheet",
     "TNFW",
+    "Multipole",
     "Source",
     "Pixelated",
     "Sersic",

src/caustics/func.py

@@ -45,6 +45,9 @@
     M0_scalemass_tnfw,
     M0_totmass_tnfw,
     concentration_tnfw,
+    reduced_deflection_angle_multipole,
+    potential_multipole,
+    convergence_multipole,
 )
 
 from .light.func import brightness_sersic, k_lenstronomy, k_sersic
@@ -96,6 +99,9 @@
     "M0_scalemass_tnfw",
     "M0_totmass_tnfw",
     "concentration_tnfw",
+    "reduced_deflection_angle_multipole",
+    "potential_multipole",
+    "convergence_multipole",
     "brightness_sersic",
     "k_lenstronomy",
     "k_sersic",

src/caustics/lenses/__init__.py

@@ -11,6 +11,7 @@
 from .mass_sheet import MassSheet
 from .tnfw import TNFW
 from .multiplane import Multiplane
+from .multipole import Multipole
 
 
 __all__ = [
@@ -28,4 +29,5 @@
     "SinglePlane",
     "MassSheet",
     "TNFW",
+    "Multipole",
 ]

src/caustics/lenses/func/__init__.py

@@ -53,6 +53,12 @@
     concentration_tnfw,
 )
 
+from .multipole import (
+    reduced_deflection_angle_multipole,
+    potential_multipole,
+    convergence_multipole,
+)
+
 __all__ = (
     "forward_raytrace",
     "physical_from_reduced_deflection_angle",
@@ -102,4 +108,7 @@
     "M0_scalemass_tnfw",
     "M0_totmass_tnfw",
     "concentration_tnfw",
+    "reduced_deflection_angle_multipole",
+    "potential_multipole",
+    "convergence_multipole",
 )

src/caustics/lenses/func/multipole.py

@@ -0,0 +1,124 @@
+import torch
+
+from ...utils import translate_rotate
+
+
+def reduced_deflection_angle_multipole(x0, y0, m, a_m, phi_m, x, y):
+    """
+    Calculates the reduced deflection angle.
+
+    Parameters
+    ----------
+    x: Tensor
+        x-coordinates in the lens plane.
+    y: Tensor
+        y-coordinates in the lens plane.
+    z_s: Tensor
+        Redshifts of the sources.
+    x0: Tensor
+        x-coordinate of the center of the lens.
+    y0: Tensor
+        y-coordinate of the center of the lens.
+    m: Tensor, int
+        The multipole order.
+    a_m: Tensor
+        The multipole amplitude.
+    phi_m: Tensor
+        The multipole orientation.
+
+    Returns
+    -------
+    tuple[Tensor, Tensor]
+        The reduced deflection angles in the x and y directions.
+
+    Equation (B11) and (B12) https://arxiv.org/pdf/1307.4220, Xu et al. 2014
+    """
+    x, y = translate_rotate(x, y, x0, y0)
+
+    phi = torch.arctan2(y, x)
+    ax = torch.cos(phi) * a_m / (1 - m**2) * torch.cos(m * (phi - phi_m)) + torch.sin(
+        phi
+    ) * m * a_m / (1 - m**2) * torch.sin(m * (phi - phi_m))
+    ay = torch.sin(phi) * a_m / (1 - m**2) * torch.cos(m * (phi - phi_m)) - torch.cos(
+        phi
+    ) * m * a_m / (1 - m**2) * torch.sin(m * (phi - phi_m))
+
+    return ax, ay  # derotate(ax, ay, phi)
+
+
+def potential_multipole(x0, y0, m, a_m, phi_m, x, y):
+    """
+    Compute the lensing potential.
+
+    Parameters
+    ----------
+    x: Tensor
+        x-coordinates in the lens plane.
+    y: Tensor
+        y-coordinates in the lens plane.
+    z_s: Tensor
+        Redshifts of the sources.
+    x0: Tensor
+        x-coordinate of the center of the lens.
+    y0: Tensor
+        y-coordinate of the center of the lens.
+    m: Tensor, int
+        The multipole order.
+    a_m: Tensor
+        The multipole amplitude.
+    phi_m: Tensor
+        The multipole orientation.
+
+    Returns
+    -------
+    potential: Tensor
+        Lensing potential.
+
+        *Unit: arcsec^2*
+
+    Equation (B11) and (B3) https://arxiv.org/pdf/1307.4220, Xu et al. 2014
+
+    """
+    x, y = translate_rotate(x, y, x0, y0)
+    r = torch.sqrt(x**2 + y**2)
+    phi = torch.arctan2(y, x)
+    return r * a_m / (1 - m**2) * torch.cos(m * (phi - phi_m))
+
+
+def convergence_multipole(x0, y0, m, a_m, phi_m, x, y):
+    """
+    Compute the lensing convergence.
+
+    Parameters
+    ----------
+    x: Tensor
+        x-coordinates in the lens plane.
+    y: Tensor
+        y-coordinates in the lens plane.
+    z_s: Tensor
+        Redshifts of the sources.
+    x0: Tensor
+        x-coordinate of the center of the lens.
+    y0: Tensor
+        y-coordinate of the center of the lens.
+    m: Tensor, int
+        The multipole order.
+    a_m: Tensor
+        The multipole amplitude.
+    phi_m: Tensor
+        The multipole orientation.
+
+    Returns
+    -------
+    convergence: Tensor
+        Lensing convergence.
+
+        *Unit: unitless*
+
+    Equation (B10) and (B3) https://arxiv.org/pdf/1307.4220, Xu et al. 2014
+
+    """
+    x, y = translate_rotate(x, y, x0, y0)
+    r = torch.sqrt(x**2 + y**2)
+    phi = torch.arctan2(y, x)
+    return 1 / (2 * r) * a_m * torch.cos(m * (phi - phi_m))