callback implemented and tested

Author: Jammy2211

autogalaxy/profiles/mass/dark/mcr_util.py

@@ -1,3 +1,6 @@
+import jax
+import jax.numpy as jnp
+from jax import ShapeDtypeStruct
 import numpy as np
 import warnings
 
@@ -56,63 +59,147 @@ def kappa_s_and_scale_radius_for_duffy(mass_at_200, redshift_object, redshift_so
     return kappa_s, scale_radius, radius_at_200
 
 
-def kappa_s_and_scale_radius_for_ludlow(
-    mass_at_200, scatter_sigma, redshift_object, redshift_source
+def _ludlow16_cosmology_callback(
+    mass_at_200,
+    redshift_object,
+    redshift_source,
 ):
     """
-    Computes the AutoGalaxy NFW parameters (kappa_s, scale_radius) for an NFW halo of the given
-    mass, enforcing the Ludlow '16 mass-concentration relation.
-
-    Interprets mass as *`M_{200c}`*, not `M_{200m}`.
+    Pure NumPy / Python function.
+    Must NEVER see JAX tracers.
     """
+
+    import numpy as np
     from astropy import units
     from colossus.cosmology import cosmology as col_cosmology
     from colossus.halo.concentration import concentration as col_concentration
-
-    warnings.filterwarnings("ignore")
-
     from autogalaxy.cosmology.wrap import Planck15
 
-    cosmology = Planck15()
-
+    # -----------------------
+    # Colossus cosmology
+    # -----------------------
     col_cosmo = col_cosmology.setCosmology("planck15")
+
     m_input = mass_at_200 * col_cosmo.h
     concentration = col_concentration(
-        m_input, "200c", redshift_object, model="ludlow16"
+        m_input,
+        "200c",
+        redshift_object,
+        model="ludlow16",
     )
 
-    concentration = 10.0 ** (np.log10(concentration) + scatter_sigma * 0.15)
+    # -----------------------
+    # Astropy cosmology
+    # -----------------------
+    cosmology = Planck15()
 
     cosmic_average_density = (
-        cosmology.critical_density(redshift_object).to(units.solMass / units.kpc**3)
-    ).value
+        cosmology.critical_density(redshift_object)
+        .to(units.solMass / units.kpc**3)
+        .value
+    )
 
     critical_surface_density = (
         cosmology.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from(
-            redshift_0=redshift_object, redshift_1=redshift_source
+            redshift_0=redshift_object,
+            redshift_1=redshift_source,
         )
     )
 
     kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object)
 
+    return (
+        np.float64(concentration),
+        np.float64(cosmic_average_density),
+        np.float64(critical_surface_density),
+        np.float64(kpc_per_arcsec),
+    )
+
+
+def ludlow16_cosmology_jax(
+    mass_at_200,
+    redshift_object,
+    redshift_source,
+):
+    """
+    JAX-safe wrapper around Colossus + Astropy cosmology.
+    """
+
+    return jax.pure_callback(
+        _ludlow16_cosmology_callback,
+        (
+            ShapeDtypeStruct((), jnp.float64),  # concentration
+            ShapeDtypeStruct((), jnp.float64),  # rho_crit(z)
+            ShapeDtypeStruct((), jnp.float64),  # Sigma_crit
+            ShapeDtypeStruct((), jnp.float64),  # kpc/arcsec
+        ),
+        mass_at_200,
+        redshift_object,
+        redshift_source,
+    )
+
+
+def kappa_s_and_scale_radius_for_ludlow(
+    mass_at_200,
+    scatter_sigma,
+    redshift_object,
+    redshift_source,
+):
+
+    if isinstance(mass_at_200, (float, np.ndarray, np.float64)):
+        xp = np
+    else:
+        xp = jnp
+
+    # ------------------------------------
+    # Cosmology + concentration (callback)
+    # ------------------------------------
+
+    if xp is np:
+        (
+            concentration,
+            cosmic_average_density,
+            critical_surface_density,
+            kpc_per_arcsec,
+        ) = _ludlow16_cosmology_callback(
+            mass_at_200,
+            redshift_object,
+            redshift_source,
+        )
+    else:
+        (
+            concentration,
+            cosmic_average_density,
+            critical_surface_density,
+            kpc_per_arcsec,
+        ) = ludlow16_cosmology_jax(
+            mass_at_200,
+            redshift_object,
+            redshift_source,
+        )
+
+    # Apply scatter (JAX-safe)
+    concentration = 10.0 ** (xp.log10(concentration) + scatter_sigma * 0.15)
+
+    # ------------------------------------
+    # JAX-native algebra
+    # ------------------------------------
     radius_at_200 = (
-        mass_at_200 / (200.0 * cosmic_average_density * (4.0 * np.pi / 3.0))
-    ) ** (
-        1.0 / 3.0
-    )  # r200
+        mass_at_200 / (200.0 * cosmic_average_density * (4.0 * xp.pi / 3.0))
+    ) ** (1.0 / 3.0)
 
     de_c = (
         200.0
         / 3.0
         * (
             concentration**3
-            / (np.log(1.0 + concentration) - concentration / (1.0 + concentration))
+            / (xp.log(1.0 + concentration) - concentration / (1.0 + concentration))
         )
-    )  # rho_c
+    )
 
-    scale_radius_kpc = radius_at_200 / concentration  # scale radius in kpc
-    rho_s = cosmic_average_density * de_c  # rho_s
-    kappa_s = rho_s * scale_radius_kpc / critical_surface_density  # kappa_s
-    scale_radius = scale_radius_kpc / kpc_per_arcsec  # scale radius in arcsec
+    scale_radius_kpc = radius_at_200 / concentration
+    rho_s = cosmic_average_density * de_c
+    kappa_s = rho_s * scale_radius_kpc / critical_surface_density
+    scale_radius = scale_radius_kpc / kpc_per_arcsec
 
     return kappa_s, scale_radius, radius_at_200