From 9d63599803b81387965505711346313fc94f9233 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 17 Jul 2026 08:59:14 +0100 Subject: [PATCH] Add PowerLawIntermediate: intermediate-axis (COOLEST) Einstein radius New elliptical power-law whose einstein_radius parameter follows the intermediate-axis convention of COOLEST/lenstronomy/herculens (theta_int = sqrt(q) (2/(1+q))^(1/(gamma-1)) theta_PL), giving an identity COOLEST PEMD mapping. Implemented via a behavior-preserving einstein_radius_major_from(xp) hook in PowerLawCore/PowerLaw (all existing profiles numerically unchanged; full suite green), a thin xp-threaded override, ag.mp export, and a priors-config block. mass_profile_from gains intermediate=True to rebuild PEMDs as PowerLawIntermediate. Part of PyAutoLabs/PyAutoLens#616. Co-Authored-By: Claude Fable 5 --- .../config/priors/mass/total/power_law.yaml | 66 +++++++++ autogalaxy/interop/coolest/mass.py | 48 ++++++- autogalaxy/profiles/mass/__init__.py | 1 + autogalaxy/profiles/mass/total/__init__.py | 2 +- autogalaxy/profiles/mass/total/power_law.py | 57 +++++++- .../profiles/mass/total/power_law_core.py | 13 +- test_autogalaxy/interop/test_coolest_mass.py | 45 ++++++ .../mass/total/test_power_law_intermediate.py | 128 ++++++++++++++++++ 8 files changed, 356 insertions(+), 4 deletions(-) create mode 100644 test_autogalaxy/profiles/mass/total/test_power_law_intermediate.py diff --git a/autogalaxy/config/priors/mass/total/power_law.yaml b/autogalaxy/config/priors/mass/total/power_law.yaml index d0b73133f..4fbca5c11 100644 --- a/autogalaxy/config/priors/mass/total/power_law.yaml +++ b/autogalaxy/config/priors/mass/total/power_law.yaml @@ -104,3 +104,69 @@ PowerLawSph: limits: lower: 1.0 upper: 3.0 + +PowerLawIntermediate: + centre_0: + type: Gaussian + mean: 0.0 + sigma: 0.1 + width_modifier: + type: Absolute + value: 0.05 + limits: + lower: -inf + upper: inf + centre_1: + type: Gaussian + mean: 0.0 + sigma: 0.1 + width_modifier: + type: Absolute + value: 0.05 + limits: + lower: -inf + upper: inf + einstein_radius: + type: Uniform + lower_limit: 0.0 + upper_limit: 8.0 + width_modifier: + type: Relative + value: 0.25 + limits: + lower: 0.0 + upper: inf + ell_comps_0: + type: TruncatedGaussian + mean: 0.0 + sigma: 0.3 + lower_limit: -1.0 + upper_limit: 1.0 + width_modifier: + type: Absolute + value: 0.2 + limits: + lower: -1.0 + upper: 1.0 + ell_comps_1: + type: TruncatedGaussian + mean: 0.0 + sigma: 0.3 + lower_limit: -1.0 + upper_limit: 1.0 + width_modifier: + type: Absolute + value: 0.2 + limits: + lower: -1.0 + upper: 1.0 + slope: + type: Uniform + lower_limit: 1.5 + upper_limit: 3.0 + width_modifier: + type: Absolute + value: 0.2 + limits: + lower: 1.0 + upper: 3.0 diff --git a/autogalaxy/interop/coolest/mass.py b/autogalaxy/interop/coolest/mass.py index 1bc554d00..e08cf6a46 100644 --- a/autogalaxy/interop/coolest/mass.py +++ b/autogalaxy/interop/coolest/mass.py @@ -46,6 +46,7 @@ from autogalaxy.profiles.mass.total.isothermal import Isothermal from autogalaxy.profiles.mass.total.isothermal import IsothermalSph from autogalaxy.profiles.mass.total.power_law import PowerLaw +from autogalaxy.profiles.mass.total.power_law import PowerLawIntermediate from autogalaxy.profiles.mass.total.power_law import PowerLawSph @@ -174,6 +175,40 @@ def _power_law_from(parameters: Dict) -> PowerLaw: ) +def _pemd_dict_from_intermediate(profile, **kwargs) -> Dict: + """ + The ``PowerLawIntermediate`` Einstein radius is already the COOLEST + intermediate-axis theta_E, so the mapping is an identity on all parameters + (only the ellipticity / angle / centre conventions change). + """ + q, phi = conventions.q_phi_from_ell_comps(ell_comps=profile.ell_comps) + center_x, center_y = conventions.center_from_centre(centre=profile.centre) + return { + "type": "PEMD", + "parameters": { + "gamma": float(profile.slope), + "theta_E": float(profile.einstein_radius), + "q": q, + "phi": phi, + "center_x": center_x, + "center_y": center_y, + }, + } + + +def _power_law_intermediate_from(parameters: Dict) -> PowerLawIntermediate: + return PowerLawIntermediate( + centre=conventions.centre_from_center( + center_x=parameters["center_x"], center_y=parameters["center_y"] + ), + ell_comps=conventions.ell_comps_from_q_phi( + q=parameters["q"], phi=parameters["phi"] + ), + einstein_radius=float(parameters["theta_E"]), + slope=float(parameters["gamma"]), + ) + + def _nfw_dict_from(profile, sigma_crit: Optional[float] = None, **kwargs) -> Dict: """ The PyAutoGalaxy NFW convergence is ``kappa(xi) = 2 kappa_s g(xi / r_s)`` @@ -281,6 +316,7 @@ def _mass_sheet_from(parameters: Dict) -> MassSheet: Isothermal: _sie_dict_from, IsothermalSph: _sie_dict_from, PowerLaw: _pemd_dict_from, + PowerLawIntermediate: _pemd_dict_from_intermediate, PowerLawSph: _pemd_dict_from, NFW: _nfw_dict_from, NFWSph: _nfw_dict_from, @@ -325,7 +361,10 @@ def coolest_dict_from_mass(profile, sigma_crit: Optional[float] = None) -> Dict: def mass_profile_from( - profile_type: str, parameters: Dict, sigma_crit: Optional[float] = None + profile_type: str, + parameters: Dict, + sigma_crit: Optional[float] = None, + intermediate: bool = False, ): """ Build a PyAutoGalaxy mass profile from a COOLEST profile type name and its @@ -342,9 +381,16 @@ def mass_profile_from( sigma_crit The critical surface mass density of the lens configuration, required only for NFW profiles. + intermediate + If True, a PEMD profile is built as a ``PowerLawIntermediate`` (its + ``einstein_radius`` is the template's ``theta_E`` unchanged) instead of + the default ``PowerLaw`` (which converts theta_E to the PyAutoGalaxy + major-axis convention). Both give the identical mass distribution. """ if profile_type == "NFW": return _nfw_from(parameters=parameters, sigma_crit=sigma_crit) + if profile_type == "PEMD" and intermediate: + return _power_law_intermediate_from(parameters=parameters) try: from_coolest = _FROM_COOLEST[profile_type] except KeyError: diff --git a/autogalaxy/profiles/mass/__init__.py b/autogalaxy/profiles/mass/__init__.py index 560a4283e..4468fb5a6 100644 --- a/autogalaxy/profiles/mass/__init__.py +++ b/autogalaxy/profiles/mass/__init__.py @@ -17,6 +17,7 @@ IsothermalCore, IsothermalCoreSph, PowerLaw, + PowerLawIntermediate, PowerLawSph, Isothermal, IsothermalSph, diff --git a/autogalaxy/profiles/mass/total/__init__.py b/autogalaxy/profiles/mass/total/__init__.py index 897ea34cb..63ebb8530 100644 --- a/autogalaxy/profiles/mass/total/__init__.py +++ b/autogalaxy/profiles/mass/total/__init__.py @@ -8,7 +8,7 @@ ) from .isothermal import Isothermal, IsothermalSph from .isothermal_core import IsothermalCore, IsothermalCoreSph -from .power_law import PowerLaw, PowerLawSph +from .power_law import PowerLaw, PowerLawIntermediate, PowerLawSph from .power_law_broken import PowerLawBroken, PowerLawBrokenSph from .power_law_core import PowerLawCore, PowerLawCoreSph from .power_law_multipole import PowerLawMultipole diff --git a/autogalaxy/profiles/mass/total/power_law.py b/autogalaxy/profiles/mass/total/power_law.py index 67be63edd..3be5b3346 100644 --- a/autogalaxy/profiles/mass/total/power_law.py +++ b/autogalaxy/profiles/mass/total/power_law.py @@ -103,7 +103,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): slope = self.slope - 1.0 einstein_radius = ( 2.0 / (self.axis_ratio(xp) ** -0.5 + self.axis_ratio(xp) ** 0.5) - ) * self.einstein_radius + ) * self.einstein_radius_major_from(xp) factor = xp.divide(1.0 - self.axis_ratio(xp), 1.0 + self.axis_ratio(xp)) b = xp.multiply(einstein_radius, xp.sqrt(self.axis_ratio(xp))) @@ -168,6 +168,61 @@ def potential_func(u, y, x, axis_ratio, slope, core_radius, xp=np): ) +class PowerLawIntermediate(PowerLaw): + r"""Elliptical power-law mass profile with an intermediate-axis Einstein radius. + + Identical mass distribution to :class:`PowerLaw`, but the ``einstein_radius`` + parameter follows the *intermediate-axis* convention used by the COOLEST + standard, lenstronomy and herculens: radii are measured as + :math:`r = \sqrt{a b} = \sqrt{q} \, a` on the elliptical isodensity + contours, and the convergence is + + .. math:: + + \kappa(r) = \frac{3 - \gamma}{2} + \left(\frac{\theta_{\rm E}}{r}\right)^{\gamma - 1}. + + The relation to the :class:`PowerLaw` ``einstein_radius`` (:math:`\theta_{\rm PL}`) is: + + .. math:: + + \theta_{\rm E} = \sqrt{q} \left(\frac{2}{1 + q}\right)^{1 / (\gamma - 1)} \theta_{\rm PL}, + + which for the isothermal case (:math:`\gamma = 2`) reduces to + :math:`2 \sqrt{q} / (1 + q) \, \theta_{\rm PL}`. + + Use this profile when parameter values must line up directly with other lens + modeling codes — its COOLEST ``PEMD`` mapping is an identity + (``autogalaxy.interop.coolest``) — at the cost of departing from the + Einstein-radius convention of the rest of **PyAutoGalaxy**. + + Parameters + ---------- + centre : (float, float) + (y, x) arc-second coordinates of the profile centre. + ell_comps : (float, float) + Ellipticity components (e1, e2) of the elliptical coordinate system. + einstein_radius : float + Einstein radius in arcseconds, in the intermediate-axis (COOLEST) convention. + slope : float + Logarithmic density slope :math:`\gamma`. + """ + + def einstein_radius_major_from(self, xp=np): + """ + Convert the stored intermediate-axis ``einstein_radius`` to the major-axis + convention the power-law formulae are written in (the inverse of the + relation in the class docstring). Threaded through ``xp`` so the profile + is JAX-traceable like its parent. + """ + axis_ratio = self.axis_ratio(xp) + return ( + self.einstein_radius + / xp.sqrt(axis_ratio) + * ((1.0 + axis_ratio) / 2.0) ** (1.0 / (self.slope - 1.0)) + ) + + class PowerLawSph(PowerLaw): r"""Spherical power-law mass profile. diff --git a/autogalaxy/profiles/mass/total/power_law_core.py b/autogalaxy/profiles/mass/total/power_law_core.py index c12be621c..5db083003 100644 --- a/autogalaxy/profiles/mass/total/power_law_core.py +++ b/autogalaxy/profiles/mass/total/power_law_core.py @@ -71,6 +71,17 @@ def __init__( self.slope = slope self.core_radius = core_radius + def einstein_radius_major_from(self, xp=np): + """ + The Einstein radius in the convention the profile's convergence and deflection formulae are written in, + where the elliptical radius is the major-axis radius of the isodensity contour. + + For this class and its standard subclasses this is the ``einstein_radius`` parameter itself; subclasses + which parameterise the Einstein radius in a different convention (e.g. ``PowerLawIntermediate``) override + this with the conversion, which every calculation then inherits. + """ + return self.einstein_radius + def einstein_radius_rescaled(self, xp=np): """ Rescale the einstein radius by slope and axis_ratio, to reduce its degeneracy with other mass-profiles @@ -78,7 +89,7 @@ def einstein_radius_rescaled(self, xp=np): """ return ( (3 - self.slope) / (1 + self.axis_ratio(xp)) - ) * self.einstein_radius ** (self.slope - 1) + ) * self.einstein_radius_major_from(xp) ** (self.slope - 1) @aa.over_sample @aa.decorators.to_array diff --git a/test_autogalaxy/interop/test_coolest_mass.py b/test_autogalaxy/interop/test_coolest_mass.py index 416abb0fb..f04d1170a 100644 --- a/test_autogalaxy/interop/test_coolest_mass.py +++ b/test_autogalaxy/interop/test_coolest_mass.py @@ -112,6 +112,51 @@ def test__power_law__round_trip_is_numerically_identical(axis_ratio, angle, slop ) +def test__power_law_intermediate__coolest_mapping_is_identity(): + profile = ag.mp.PowerLawIntermediate( + centre=(0.1, -0.2), + ell_comps=ag.convert.ell_comps_from(axis_ratio=0.7, angle=45.0), + einstein_radius=1.2, + slope=2.1, + ) + + profile_dict = coolest.coolest_dict_from_mass(profile=profile) + + assert profile_dict["type"] == "PEMD" + # No sqrt(q) factor — theta_E is the parameter itself. + assert profile_dict["parameters"]["theta_E"] == 1.2 + assert profile_dict["parameters"]["gamma"] == pytest.approx(2.1) + assert profile_dict["parameters"]["q"] == pytest.approx(0.7) + + # And the analytic COOLEST convergence with these parameters matches. + convergence = profile.convergence_2d_from(grid=aa.Grid2DIrregular(grid_yx())) + convergence_coolest = coolest_power_law_convergence( + parameters=profile_dict["parameters"], grid=grid_yx() + ) + assert np.asarray(convergence) == pytest.approx(convergence_coolest, rel=1e-8) + + profile_back = coolest.mass_profile_from( + profile_type="PEMD", + parameters=profile_dict["parameters"], + intermediate=True, + ) + + assert isinstance(profile_back, ag.mp.PowerLawIntermediate) + assert profile_back.einstein_radius == pytest.approx(1.2, rel=1e-12) + + # Default import still builds the PyAutoGalaxy-convention PowerLaw, with + # the identical mass distribution. + profile_back_default = coolest.mass_profile_from( + profile_type="PEMD", parameters=profile_dict["parameters"] + ) + assert type(profile_back_default) is ag.mp.PowerLaw + + grid = aa.Grid2DIrregular(grid_yx()) + assert np.asarray( + profile_back_default.convergence_2d_from(grid=grid) + ) == pytest.approx(np.asarray(profile.convergence_2d_from(grid=grid)), rel=1e-8) + + def test__isothermal_sph__round_trip(): profile = ag.mp.IsothermalSph(centre=(0.1, 0.2), einstein_radius=0.8) diff --git a/test_autogalaxy/profiles/mass/total/test_power_law_intermediate.py b/test_autogalaxy/profiles/mass/total/test_power_law_intermediate.py new file mode 100644 index 000000000..9f2141d69 --- /dev/null +++ b/test_autogalaxy/profiles/mass/total/test_power_law_intermediate.py @@ -0,0 +1,128 @@ +import numpy as np +import pytest + +import autoarray as aa +import autogalaxy as ag + + +def grid_yx(): + return np.array( + [ + [1.0, 0.5], + [-0.3, 0.7], + [0.2, -1.1], + [0.8, 0.9], + [-1.2, -0.4], + ] + ) + + +def theta_intermediate_from(einstein_radius, axis_ratio, slope): + return ( + np.sqrt(axis_ratio) + * (2.0 / (1.0 + axis_ratio)) ** (1.0 / (slope - 1.0)) + * einstein_radius + ) + + +@pytest.mark.parametrize( + "axis_ratio,angle,slope", + [(0.7, 45.0, 2.0), (0.6, 40.0, 2.3), (0.9, -20.0, 1.7)], +) +def test__equivalent_to_power_law_under_theta_rescale(axis_ratio, angle, slope): + """ + A `PowerLawIntermediate` whose einstein_radius is the intermediate-axis + conversion of a `PowerLaw`'s einstein_radius is the identical mass + distribution. + """ + ell_comps = ag.convert.ell_comps_from(axis_ratio=axis_ratio, angle=angle) + + power_law = ag.mp.PowerLaw( + centre=(0.1, -0.2), + ell_comps=ell_comps, + einstein_radius=1.4, + slope=slope, + ) + intermediate = ag.mp.PowerLawIntermediate( + centre=(0.1, -0.2), + ell_comps=ell_comps, + einstein_radius=theta_intermediate_from( + einstein_radius=1.4, axis_ratio=axis_ratio, slope=slope + ), + slope=slope, + ) + + grid = aa.Grid2DIrregular(grid_yx()) + + assert np.asarray(intermediate.convergence_2d_from(grid=grid)) == pytest.approx( + np.asarray(power_law.convergence_2d_from(grid=grid)), rel=1e-10 + ) + assert np.asarray( + intermediate.deflections_yx_2d_from(grid=grid) + ) == pytest.approx( + np.asarray(power_law.deflections_yx_2d_from(grid=grid)), rel=1e-8 + ) + assert np.asarray(intermediate.potential_2d_from(grid=grid)) == pytest.approx( + np.asarray(power_law.potential_2d_from(grid=grid)), rel=1e-8 + ) + + +def test__convergence_matches_coolest_form_with_identity_theta(): + """ + The COOLEST / lenstronomy PEMD convergence evaluated directly with the + profile's own einstein_radius (no conversion factor) matches the profile — + the defining property of the intermediate-axis convention. + """ + axis_ratio, angle, slope = 0.6, 40.0, 2.3 + + profile = ag.mp.PowerLawIntermediate( + centre=(0.0, 0.0), + ell_comps=ag.convert.ell_comps_from(axis_ratio=axis_ratio, angle=angle), + einstein_radius=1.2, + slope=slope, + ) + + grid = grid_yx() + x = grid[:, 1] + y = grid[:, 0] + + alpha = np.radians(angle) + x_prime = np.cos(alpha) * x + np.sin(alpha) * y + y_prime = -np.sin(alpha) * x + np.cos(alpha) * y + + r = np.sqrt(axis_ratio * x_prime**2 + y_prime**2 / axis_ratio) + convergence_coolest = (3.0 - slope) / 2.0 * (1.2 / r) ** (slope - 1.0) + + convergence = profile.convergence_2d_from(grid=aa.Grid2DIrregular(grid)) + + assert np.asarray(convergence) == pytest.approx(convergence_coolest, rel=1e-8) + + +def test__power_law_regression__unchanged_by_hook_refactor(): + """ + Pin PowerLaw numerics: the einstein_radius_major_from hook must be an + exact no-op for the existing profile. + """ + profile = ag.mp.PowerLaw( + centre=(0.0, 0.0), + ell_comps=(0.05, 0.05), + einstein_radius=1.0, + slope=2.2, + ) + + assert profile.einstein_radius_major_from() == 1.0 + + convergence = profile.convergence_2d_from( + grid=aa.Grid2DIrregular(np.array([[0.5, 0.5]])) + ) + + # (3 - 2.2) / (1 + q) * (1 / xi)^(1.2) evaluated independently. + q, phi = profile.axis_ratio(), profile.angle() + alpha = np.radians(phi) + x_p = np.cos(alpha) * 0.5 + np.sin(alpha) * 0.5 + y_p = -np.sin(alpha) * 0.5 + np.cos(alpha) * 0.5 + xi = np.sqrt(x_p**2 + (y_p / q) ** 2) + + assert np.asarray(convergence)[0] == pytest.approx( + (3 - 2.2) / (1 + q) * xi ** (-1.2), rel=1e-8 + )