feat: twins stellar-plus-dark transformation and diagnostics (lenstwin 3)

autolens/plot/__init__.py

@@ -109,3 +109,5 @@
     subplot_sensitivity,
     subplot_figures_of_merit_grid as subplot_sensitivity_figures_of_merit,
 )
+
+from autolens.twins.plot.twin_plots import subplot_twin_comparison

autolens/twins/__init__.py

@@ -19,6 +19,11 @@
 Poisson noise is used only when the setup sets ``exposure_time`` and
 ``background_sky_level`` and ``noise="poisson"`` is requested.
 
+``transform.ReplaceTotalMassWithStellarPlusDark`` replaces a galaxy's total mass with
+stellar mass following its MGE light plus a dark halo, normalised inside the parent's
+Einstein radius (``einstein_radius_aperture_from``); ``compare`` and ``compare_to_data``
+report what a transformation changed, and ``aplt.subplot_twin_comparison`` plots it.
+
 NumPy only: nothing here supports JAX.
 """
 
@@ -27,6 +32,8 @@
 from autolens.twins.reference import ReferenceModel
 from autolens.twins.twin import Twin
 from autolens.twins.simulate import SimulatedTwin
+from autolens.twins.aperture import einstein_radius_aperture_from
+from autolens.twins.compare import TwinComparison, compare, compare_to_data
 from autolens.twins import transform
 
 __all__ = [
@@ -35,5 +42,9 @@
     "Twin",
     "SimulatedTwin",
     "TransformationRecord",
+    "TwinComparison",
+    "compare",
+    "compare_to_data",
+    "einstein_radius_aperture_from",
     "transform",
 ]

autolens/twins/_galaxies.py

@@ -0,0 +1,64 @@
+"""
+How ``al.twins`` addresses the galaxies of a tracer.
+
+A tracer loaded from a model-fit's ``files/tracer.json`` holds its galaxies in an
+``af.ModelInstance`` keyed by the model's names (``af.Collection(lens=..., source=...)``
+gives ``"lens"`` and ``"source"``); those keys are the galaxy names. A tracer built
+from a plain list (or a list-backed ``ModelInstance``) names its galaxies by position,
+``"galaxy_0"``, ``"galaxy_1"``, ..., as ``al.interop.coolest`` does. These are the
+names that ``ReferenceModel`` provenance ``roles`` (galaxy name -> role) are keyed by,
+so a transformation also accepts a role, resolved through the reference's ``roles``.
+
+A component path is ``"galaxies.<name>.<attribute>"`` (e.g. ``"galaxies.lens.mass"``),
+and a tracer-level field is ``"fields.<index>"``.
+"""
+
+from typing import Dict
+
+
+def galaxy_name_list_from(tracer) -> list:
+    """The galaxy names of ``tracer``, in the order of ``tracer.galaxies``."""
+    return list(galaxy_dict_from(tracer))
+
+
+def galaxy_dict_from(tracer) -> Dict[str, object]:
+    """
+    The galaxies of ``tracer`` by name: the keys of a dict-backed ``af.ModelInstance``
+    (a tracer from ``tracer.json``), otherwise ``"galaxy_<index>"``.
+    """
+    galaxies = tracer.galaxies
+    items = galaxies.dict.items() if hasattr(galaxies, "dict") else None
+    if items is not None and all(isinstance(key, str) for key, _ in items):
+        return dict(items)
+    return {f"galaxy_{index}": galaxy for index, galaxy in enumerate(galaxies)}
+
+
+def galaxy_name_from(twin, galaxy: str) -> str:
+    """
+    Resolve ``galaxy`` (a galaxy name, or a role from the reference's provenance
+    ``roles``) to the name of one galaxy of ``twin.tracer``.
+
+    Raises
+    ------
+    ValueError
+        If ``galaxy`` is neither a galaxy name nor a role held by exactly one galaxy.
+    """
+    names = galaxy_name_list_from(twin.tracer)
+    if galaxy in names:
+        return galaxy
+
+    roles = twin.reference.provenance.get("roles", {}) or {}
+    matches = [name for name, role in roles.items() if role == galaxy]
+    if len(matches) == 1 and matches[0] in names:
+        return matches[0]
+    if len(matches) > 1:
+        raise ValueError(
+            f"The role {galaxy!r} is held by several galaxies ({matches}); name the "
+            f"galaxy instead."
+        )
+    raise ValueError(
+        f"No galaxy named {galaxy!r} in the twin's tracer, and no galaxy has that role. "
+        f"Galaxy names: {names}; roles: {roles}. A tracer loaded from tracer.json names "
+        f"its galaxies by the model's keys (e.g. 'lens'); a list of galaxies is named "
+        f"'galaxy_0', 'galaxy_1', ..."
+    )

autolens/twins/aperture.py

@@ -0,0 +1,103 @@
+"""
+The Einstein-radius aperture of the mass normalisation convention (ruling R3,
+PyAutoGalaxy ``docs/api/mass_budget.rst``): a circle of radius theta_E,parent, the
+Einstein radius of the whole parent lens model, centred on the mass centre of the
+galaxy being replaced.
+
+theta_E is computed at the tracer level, ``ag.LensCalc.from_tracer(tracer)
+.einstein_radius_from(grid)``: the radius of the circle with the area of the
+tangential critical curve of all mass in the tracer (galaxies, external shear and
+other fields), so it is defined for multi-galaxy and non-isothermal parents.
+"""
+
+from typing import Tuple
+
+import autoarray as aa
+import autogalaxy as ag
+
+
+def einstein_radius_from(
+    tracer, grid: aa.type.Grid2DLike, pixel_scale: float = 0.05
+) -> float:
+    """
+    The tracer-level Einstein radius, in arc-seconds.
+
+    Parameters
+    ----------
+    tracer
+        The lens model (all its mass contributes).
+    grid
+        The grid whose mask sets the region in which the tangential critical curve is
+        searched (``LensCalc``'s evaluation grid zooms to it). ``al.twins`` uses the
+        observation setup's masked grid, ``setup.grid``.
+    pixel_scale
+        The resolution of the evaluation grid on which the critical curve is traced.
+
+    Raises
+    ------
+    ValueError
+        If the tracer has no tangential critical curve inside the grid.
+    """
+    einstein_radius = float(
+        ag.LensCalc.from_tracer(tracer).einstein_radius_from(
+            grid=grid, pixel_scale=pixel_scale
+        )
+    )
+    if not einstein_radius > 0.0:
+        raise ValueError(
+            "The tracer has no tangential critical curve inside the grid, so its "
+            "Einstein radius (and the Einstein-radius aperture) is undefined."
+        )
+    return einstein_radius
+
+
+def einstein_radius_aperture_from(
+    tracer,
+    grid: aa.type.Grid2DLike,
+    centre: Tuple[float, float],
+    pixel_scale: float = 0.05,
+    **integrator_settings,
+) -> ag.CircularAperture:
+    """
+    The circular aperture of radius theta_E,parent (``einstein_radius_from``) centred on
+    ``centre``. The Einstein radius is the returned aperture's ``radius``.
+
+    Parameters
+    ----------
+    tracer
+        The parent lens model.
+    grid
+        The grid on which theta_E is computed (see ``einstein_radius_from``).
+    centre
+        The (y,x) aperture centre in arc-seconds: the mass centre of the replaced galaxy.
+    pixel_scale
+        The resolution of the critical-curve evaluation grid.
+    integrator_settings
+        ``sub_size_interior``, ``sub_size_edge``, ``sub_size_centre`` and
+        ``centre_refine_pixels`` of ``ag.CircularAperture`` (the grid integrator's
+        accuracy). A ``gNFW`` convergence does one ``quad`` per point, so pass coarse
+        sub-sizes for it.
+    """
+    return ag.CircularAperture(
+        radius=einstein_radius_from(tracer=tracer, grid=grid, pixel_scale=pixel_scale),
+        centre=(float(centre[0]), float(centre[1])),
+        **integrator_settings,
+    )
+
+
+def check_aperture_in_grid(aperture: ag.CircularAperture, grid: aa.Grid2D):
+    """
+    Raise ``ValueError`` if the circle extends beyond the extent of the uniform grid the
+    integrator uses (``grid.mask`` shape, pixel scales and origin; the mask is ignored).
+    """
+    mask = grid.mask
+    half_y = mask.shape_native[0] * mask.pixel_scales[0] / 2.0
+    half_x = mask.shape_native[1] * mask.pixel_scales[1] / 2.0
+    dy = abs(aperture.centre[0] - mask.origin[0]) + aperture.radius
+    dx = abs(aperture.centre[1] - mask.origin[1]) + aperture.radius
+    if dy > half_y or dx > half_x:
+        raise ValueError(
+            f'The aperture (radius {aperture.radius:.4f}" at {aperture.centre}) extends '
+            f'beyond the integration grid (half-extent {half_y:.4f}" x {half_x:.4f}"); '
+            f"pass a larger grid."
+        )