feat: support Kaplinghat halos in substructure arrays

Author: Jammy2211

autolens/lens/substructure_util.py

@@ -3,6 +3,60 @@
 import autogalaxy as ag
 
 
+def _halo_parameter_extractor_from(profile_cls):
+    kaplinghat_profile_classes = tuple(
+        cls
+        for cls in (
+            getattr(ag.mp, "KaplinghatCoredNFWSph", None),
+            getattr(ag.mp, "KaplinghatCoredNFWMCRLudlowSph", None),
+        )
+        if cls is not None
+    )
+
+    if profile_cls is ag.mp.cNFWSph:
+
+        def extract(prof):
+            return [
+                prof.centre[0],
+                prof.centre[1],
+                prof.kappa_s,
+                prof.scale_radius,
+                prof.core_radius,
+            ]
+
+        return 5, extract
+
+    if profile_cls in kaplinghat_profile_classes:
+
+        def extract(prof):
+            return [
+                prof.centre[0],
+                prof.centre[1],
+                prof.kappa_s,
+                prof.scale_radius,
+                prof.interaction_radius,
+                prof.central_density,
+                prof.isothermal_radius,
+            ]
+
+        return 7, extract
+
+    if profile_cls is ag.mp.NFWTruncatedSph:
+
+        def extract(prof):
+            return [
+                prof.centre[0],
+                prof.centre[1],
+                prof.kappa_s,
+                prof.scale_radius,
+                prof.truncation_radius,
+            ]
+
+        return 5, extract
+
+    raise ValueError(f"Unsupported halo profile class: {profile_cls}")
+
+
 def precompute_scaling_matrix(plane_redshifts, cosmology=None):
     import jax.numpy as jnp
 
@@ -28,21 +82,7 @@ def galaxies_to_halo_arrays(galaxies, plane_redshifts, max_n, profile_cls):
     import jax.numpy as jnp
 
     n_planes = len(plane_redshifts)
-
-    if profile_cls is ag.mp.cNFWSph:
-        n_params = 5
-        def extract(prof):
-            return [
-                prof.centre[0], prof.centre[1],
-                prof.kappa_s, prof.scale_radius, prof.core_radius,
-            ]
-    else:
-        n_params = 5
-        def extract(prof):
-            return [
-                prof.centre[0], prof.centre[1],
-                prof.kappa_s, prof.scale_radius, prof.truncation_radius,
-            ]
+    n_params, extract = _halo_parameter_extractor_from(profile_cls=profile_cls)
 
     params = np.zeros((n_planes, max_n, n_params))
     mask = np.zeros((n_planes, max_n), dtype=bool)

test_autolens/config/priors/dark_mass_profiles.yaml

@@ -261,8 +261,8 @@ gNFW:
     width_modifier:
       type: Relative
       value: 0.2
-gNFWSph:
-  centre_0:
+gNFWSph:
+  centre_0:
     limits:
       lower: -inf
       upper: inf
@@ -302,13 +302,155 @@ gNFWSph:
     width_modifier:
       type: Relative
       value: 0.2
-  scale_radius:
-    limits:
-      lower: 0.0
-      upper: inf
-    lower_limit: 0.0
-    type: Uniform
-    upper_limit: 30.0
-    width_modifier:
-      type: Relative
-      value: 0.2
+  scale_radius:
+    limits:
+      lower: 0.0
+      upper: inf
+    lower_limit: 0.0
+    type: Uniform
+    upper_limit: 30.0
+    width_modifier:
+      type: Relative
+      value: 0.2
+KaplinghatCoredNFWSph:
+  centre_0:
+    limits:
+      lower: -inf
+      upper: inf
+    mean: 0.0
+    sigma: 0.1
+    type: Gaussian
+    width_modifier:
+      type: Absolute
+      value: 0.05
+  centre_1:
+    limits:
+      lower: -inf
+      upper: inf
+    mean: 0.0
+    sigma: 0.1
+    type: Gaussian
+    width_modifier:
+      type: Absolute
+      value: 0.05
+  kappa_s:
+    limits:
+      lower: 0.0
+      upper: inf
+    lower_limit: 0.0
+    type: Uniform
+    upper_limit: 1.0
+    width_modifier:
+      type: Relative
+      value: 0.2
+  scale_radius:
+    limits:
+      lower: 0.0
+      upper: inf
+    lower_limit: 0.0
+    type: Uniform
+    upper_limit: 30.0
+    width_modifier:
+      type: Relative
+      value: 0.2
+  sigma_over_m:
+    limits:
+      lower: 0.0
+      upper: inf
+    lower_limit: 0.0
+    type: Uniform
+    upper_limit: 10.0
+    width_modifier:
+      type: Relative
+      value: 0.5
+  t_age:
+    limits:
+      lower: 0.0
+      upper: inf
+    lower_limit: 0.0
+    type: Uniform
+    upper_limit: 13.8
+    width_modifier:
+      type: Relative
+      value: 0.2
+  interaction_radius:
+    limits:
+      lower: 0.0
+      upper: inf
+    lower_limit: 0.0
+    type: Uniform
+    upper_limit: 30.0
+    width_modifier:
+      type: Relative
+      value: 0.2
+KaplinghatCoredNFWMCRLudlowSph:
+  centre_0:
+    limits:
+      lower: -inf
+      upper: inf
+    mean: 0.0
+    sigma: 0.1
+    type: Gaussian
+    width_modifier:
+      type: Absolute
+      value: 0.05
+  centre_1:
+    limits:
+      lower: -inf
+      upper: inf
+    mean: 0.0
+    sigma: 0.1
+    type: Gaussian
+    width_modifier:
+      type: Absolute
+      value: 0.05
+  mass_at_200:
+    limits:
+      lower: 0.0
+      upper: inf
+    lower_limit: 100000000.0
+    type: LogUniform
+    upper_limit: 1000000000000000.0
+    width_modifier:
+      type: Relative
+      value: 0.5
+  sigma_over_m:
+    limits:
+      lower: 0.0
+      upper: inf
+    lower_limit: 0.0
+    type: Uniform
+    upper_limit: 10.0
+    width_modifier:
+      type: Relative
+      value: 0.5
+  t_age:
+    limits:
+      lower: 0.0
+      upper: inf
+    lower_limit: 0.0
+    type: Uniform
+    upper_limit: 13.8
+    width_modifier:
+      type: Relative
+      value: 0.2
+  redshift_object:
+    limits:
+      lower: 0.0
+      upper: inf
+    lower_limit: 0.0
+    type: Uniform
+    upper_limit: 1.0
+    width_modifier:
+      type: Relative
+      value: 0.5
+  redshift_source:
+    limits:
+      lower: 0.0
+      upper: inf
+    lower_limit: 0.0
+    type: Uniform
+    upper_limit: 1.0
+    width_modifier:
+      type: Relative
+      value: 0.5

test_autolens/lens/test_substructure_util.py

@@ -0,0 +1,69 @@
+import numpy as np
+import pytest
+
+import autolens as al
+from autolens.lens import substructure_util
+
+
+requires_kaplinghat = pytest.mark.skipif(
+    not hasattr(al.mp, "KaplinghatCoredNFWSph"),
+    reason="Kaplinghat SIDM profiles are provided by the pending PyAutoGalaxy release.",
+)
+
+
+@requires_kaplinghat
+def test__autolens_exposes_kaplinghat_profiles_from_autogalaxy():
+    assert hasattr(al.mp, "KaplinghatCoredNFWSph")
+    assert hasattr(al.mp, "KaplinghatCoredNFWMCRLudlowSph")
+
+
+@requires_kaplinghat
+def test__galaxies_to_halo_arrays__packs_kaplinghat_profile_parameters():
+    profile = al.mp.KaplinghatCoredNFWSph(
+        centre=(0.1, -0.2),
+        kappa_s=0.03,
+        scale_radius=1.7,
+        sigma_over_m=2.0,
+        t_age=8.0,
+        interaction_radius=0.25,
+    )
+    galaxies = [
+        al.Galaxy(redshift=0.5, mass=profile),
+        al.Galaxy(redshift=1.0, mass_sheet=al.mp.MassSheet(kappa=-0.01)),
+    ]
+
+    params, mask, sheet_kappas = substructure_util.galaxies_to_halo_arrays(
+        galaxies=galaxies,
+        plane_redshifts=[0.5, 1.0],
+        max_n=2,
+        profile_cls=al.mp.KaplinghatCoredNFWSph,
+    )
+
+    assert params.shape == (2, 2, 7)
+    assert mask.tolist() == [[True, False], [False, False]]
+    assert sheet_kappas.tolist() == [0.0, -0.01]
+
+    np.testing.assert_allclose(
+        np.asarray(params[0, 0]),
+        np.array(
+            [
+                0.1,
+                -0.2,
+                profile.kappa_s,
+                profile.scale_radius,
+                profile.interaction_radius,
+                profile.central_density,
+                profile.isothermal_radius,
+            ]
+        ),
+    )
+
+
+def test__galaxies_to_halo_arrays__raises_for_unsupported_profile_class():
+    with pytest.raises(ValueError):
+        substructure_util.galaxies_to_halo_arrays(
+            galaxies=[],
+            plane_redshifts=[0.5],
+            max_n=1,
+            profile_cls=al.mp.IsothermalSph,
+        )