Merge pull request #207 from Ciela-Institute/dev

Dev

.pre-commit-config.yaml

@@ -9,7 +9,7 @@ ci:
 
 repos:
   - repo: https://github.com/psf/black
-    rev: "24.2.0"
+    rev: "24.4.0"
     hooks:
       - id: black-jupyter
 
@@ -20,7 +20,7 @@ repos:
         additional_dependencies: [black==23.7.0]
 
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: "v4.5.0"
+    rev: "v4.6.0"
     hooks:
       - id: check-added-large-files
       - id: check-case-conflict
@@ -50,7 +50,7 @@ repos:
         args: [--prose-wrap=always]
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.3.2"
+    rev: "v0.4.1"
     hooks:
       - id: ruff
         args: ["--fix", "--show-fixes"]

docs/requirements.txt

@@ -1,6 +1,7 @@
 ipywidgets
 jupyter-book
 matplotlib
+pydantic>=2.6.1,<3
 pyro-ppl
 sphinx
 sphinx_rtd_theme

docs/source/tutorials/example.yml

@@ -0,0 +1,28 @@
+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: Lens_Source
+  init_kwargs:
+    # Single lense
+    lens: *lens
+    source: *src
+    lens_light: *lnslt
+    pixelscale: 0.05
+    pixels_x: 100

src/caustics/func.py

@@ -0,0 +1,102 @@
+from .lenses.func import (
+    forward_raytrace,
+    physical_from_reduced_deflection_angle,
+    reduced_from_physical_deflection_angle,
+    reduced_deflection_angle_sie,
+    potential_sie,
+    convergence_sie,
+    reduced_deflection_angle_point,
+    potential_point,
+    convergence_point,
+    reduced_deflection_angle_mass_sheet,
+    potential_mass_sheet,
+    convergence_mass_sheet,
+    reduced_deflection_angle_epl,
+    potential_epl,
+    convergence_epl,
+    reduced_deflection_angle_external_shear,
+    potential_external_shear,
+    physical_deflection_angle_nfw,
+    potential_nfw,
+    convergence_nfw,
+    _f_batchable_nfw,
+    _f_differentiable_nfw,
+    _g_batchable_nfw,
+    _g_differentiable_nfw,
+    _h_batchable_nfw,
+    _h_differentiable_nfw,
+    reduced_deflection_angle_pixelated_convergence,
+    potential_pixelated_convergence,
+    _fft2_padded,
+    build_kernels_pixelated_convergence,
+    convergence_0_pseudo_jaffe,
+    potential_pseudo_jaffe,
+    reduced_deflection_angle_pseudo_jaffe,
+    mass_enclosed_2d_pseudo_jaffe,
+    convergence_pseudo_jaffe,
+    reduced_deflection_angle_sis,
+    potential_sis,
+    convergence_sis,
+    mass_enclosed_2d_tnfw,
+    physical_deflection_angle_tnfw,
+    potential_tnfw,
+    convergence_tnfw,
+    scale_density_tnfw,
+    M0_scalemass_tnfw,
+    M0_totmass_tnfw,
+    concentration_tnfw,
+)
+
+from .light.func import brightness_sersic, k_lenstronomy, k_sersic
+
+__all__ = (
+    "forward_raytrace",
+    "physical_from_reduced_deflection_angle",
+    "reduced_from_physical_deflection_angle",
+    "reduced_deflection_angle_sie",
+    "potential_sie",
+    "convergence_sie",
+    "reduced_deflection_angle_point",
+    "potential_point",
+    "convergence_point",
+    "reduced_deflection_angle_mass_sheet",
+    "potential_mass_sheet",
+    "convergence_mass_sheet",
+    "reduced_deflection_angle_epl",
+    "potential_epl",
+    "convergence_epl",
+    "reduced_deflection_angle_external_shear",
+    "potential_external_shear",
+    "physical_deflection_angle_nfw",
+    "potential_nfw",
+    "convergence_nfw",
+    "_f_batchable_nfw",
+    "_f_differentiable_nfw",
+    "_g_batchable_nfw",
+    "_g_differentiable_nfw",
+    "_h_batchable_nfw",
+    "_h_differentiable_nfw",
+    "reduced_deflection_angle_pixelated_convergence",
+    "potential_pixelated_convergence",
+    "_fft2_padded",
+    "build_kernels_pixelated_convergence",
+    "convergence_0_pseudo_jaffe",
+    "potential_pseudo_jaffe",
+    "reduced_deflection_angle_pseudo_jaffe",
+    "mass_enclosed_2d_pseudo_jaffe",
+    "convergence_pseudo_jaffe",
+    "reduced_deflection_angle_sis",
+    "potential_sis",
+    "convergence_sis",
+    "mass_enclosed_2d_tnfw",
+    "physical_deflection_angle_tnfw",
+    "potential_tnfw",
+    "convergence_tnfw",
+    "scale_density_tnfw",
+    "M0_scalemass_tnfw",
+    "M0_totmass_tnfw",
+    "concentration_tnfw",
+    "brightness_sersic",
+    "k_lenstronomy",
+    "k_sersic",
+)

src/caustics/lenses/base.py

@@ -11,8 +11,8 @@
 from ..cosmology import Cosmology
 from ..parametrized import Parametrized, unpack
 from .utils import get_magnification
-from ..utils import batch_lm
 from ..packed import Packed
+from . import func
 
 __all__ = ("ThinLens", "ThickLens")
 
@@ -185,40 +185,19 @@ def forward_raytrace(
             *Unit: arcsec*
         """
 
-        bxy = torch.stack((bx, by)).repeat(n_init, 1)  # has shape (n_init, Dout:2)
-
         # TODO make FOV more general so that it doesn't have to be centered on zero,zero
         if fov is None:
             raise ValueError("fov must be given to generate initial guesses")
 
-        # Random starting points in image plane
-        guesses = (torch.as_tensor(fov) * (torch.rand(n_init, 2) - 0.5)).to(
-            device=bxy.device
-        )  # Has shape (n_init, Din:2)
-
-        # Optimize guesses in image plane
-        x, l, c = batch_lm(  # noqa: E741 Unused `l` variable
-            guesses,
-            bxy,
-            lambda *a, **k: torch.stack(
-                self.raytrace(a[0][..., 0], a[0][..., 1], *a[1:], **k), dim=-1
-            ),
-            f_args=(z_s, params),
+        return func.forward_raytrace(
+            bx,
+            by,
+            partial(self.raytrace, params=params, z_s=z_s),
+            epsilon,
+            n_init,
+            fov,
         )
 
-        # Clip points that didn't converge
-        x = x[c < 1e-2 * epsilon**2]
-
-        # Cluster results into n-images
-        res = []
-        while len(x) > 0:
-            res.append(x[0])
-            d = torch.linalg.norm(x - x[0], dim=-1)
-            x = x[d > epsilon]
-
-        res = torch.stack(res, dim=0)
-        return res[..., 0], res[..., 1]
-
 
 class ThickLens(Lens):
     """
@@ -782,9 +761,8 @@ def reduced_deflection_angle(
         deflection_angle_x, deflection_angle_y = self.physical_deflection_angle(
             x, y, z_s, params
         )
-        return (
-            (d_ls / d_s) * deflection_angle_x,
-            (d_ls / d_s) * deflection_angle_y,
+        return func.reduced_from_physical_deflection_angle(
+            deflection_angle_x, deflection_angle_y, d_s, d_ls
         )
 
     @unpack
@@ -839,9 +817,8 @@ def physical_deflection_angle(
         deflection_angle_x, deflection_angle_y = self.reduced_deflection_angle(
             x, y, z_s, params
         )
-        return (
-            (d_s / d_ls) * deflection_angle_x,
-            (d_s / d_ls) * deflection_angle_y,
+        return func.physical_from_reduced_deflection_angle(
+            deflection_angle_x, deflection_angle_y, d_s, d_ls
         )
 
     @abstractmethod

src/caustics/lenses/epl.py

@@ -4,10 +4,10 @@
 import torch
 from torch import Tensor
 
-from ..utils import derotate, translate_rotate
 from .base import ThinLens, CosmologyType, NameType, ZLType
 from ..parametrized import unpack
 from ..packed import Packed
+from . import func
 
 __all__ = ("EPL",)
 
@@ -243,19 +243,9 @@ def reduced_deflection_angle(
             *Unit: arcsec*
 
         """
-        x, y = translate_rotate(x, y, x0, y0, phi)
-
-        # follow Tessore et al 2015 (eq. 5)
-        z = q * x + y * 1j
-        r = torch.abs(z)
-
-        # Tessore et al 2015 (eq. 23)
-        r_omega = self._r_omega(z, t, q)
-        alpha_c = 2.0 / (1.0 + q) * (b / r) ** t * r_omega  # fmt: skip
-
-        alpha_real = torch.nan_to_num(alpha_c.real, posinf=10**10, neginf=-(10**10))
-        alpha_imag = torch.nan_to_num(alpha_c.imag, posinf=10**10, neginf=-(10**10))
-        return derotate(alpha_real, alpha_imag, phi)
+        return func.reduced_deflection_angle_epl(
+            x0, y0, q, phi, b, t, x, y, self.n_iter
+        )
 
     def _r_omega(self, z, t, q):
         """
@@ -349,10 +339,7 @@ def potential(
             *Unit: arcsec^2*
 
         """
-        ax, ay = self.reduced_deflection_angle(x, y, z_s, params)
-        ax, ay = derotate(ax, ay, -phi)
-        x, y = translate_rotate(x, y, x0, y0, phi)
-        return (x * ax + y * ay) / (2 - t)  # fmt: skip
+        return func.potential_epl(x0, y0, q, phi, b, t, x, y, self.n_iter)
 
     @unpack
     def convergence(
@@ -402,6 +389,4 @@ def convergence(
             *Unit: unitless*
 
         """
-        x, y = translate_rotate(x, y, x0, y0, phi)
-        psi = (q**2 * (x**2 + self.s**2) + y**2).sqrt()  # fmt: skip
-        return (2 - t) / 2 * (b / psi) ** t  # fmt: skip
+        return func.convergence_epl(x0, y0, q, phi, b, t, x, y, self.s)

src/caustics/lenses/external_shear.py

@@ -2,11 +2,12 @@
 from typing import Optional, Union, Annotated
 
 from torch import Tensor
+import torch
 
-from ..utils import translate_rotate
 from .base import ThinLens, CosmologyType, NameType, ZLType
 from ..parametrized import unpack
 from ..packed import Packed
+from . import func
 
 __all__ = ("ExternalShear",)
 
@@ -135,16 +136,9 @@ def reduced_deflection_angle(
             *Unit: arcsec*
 
         """
-        x, y = translate_rotate(x, y, x0, y0)
-        # Meneghetti eq 3.83
-        # TODO, why is it not:
-        # th = (x**2 + y**2).sqrt() + self.s
-        # a1 = x/th + x * gamma_1 + y * gamma_2
-        # a2 = y/th + x * gamma_2 - y * gamma_1
-        a1 = x * gamma_1 + y * gamma_2
-        a2 = x * gamma_2 - y * gamma_1
-
-        return a1, a2  # I'm not sure but I think no derotation necessary
+        return func.reduced_deflection_angle_external_shear(
+            x0, y0, gamma_1, gamma_2, x, y
+        )
 
     @unpack
     def potential(
@@ -192,9 +186,7 @@ def potential(
             *Unit: arcsec^2*
 
         """
-        ax, ay = self.reduced_deflection_angle(x, y, z_s, params)
-        x, y = translate_rotate(x, y, x0, y0)
-        return 0.5 * (x * ax + y * ay)
+        return func.potential_external_shear(x0, y0, gamma_1, gamma_2, x, y)
 
     @unpack
     def convergence(
@@ -247,4 +239,4 @@ def convergence(
             This method is not implemented as the convergence is not defined
             for an external shear.
         """
-        raise NotImplementedError("convergence undefined for external shear")
+        return torch.zeros_like(x)