From 9541c87bb2b6174839d5adad9faafffc2c5a0809 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 15:35:35 +0000 Subject: [PATCH] feat: twins stellar-plus-dark transformation and diagnostics (lenstwin 3) Public API grew: al.twins.einstein_radius_aperture_from, al.twins.transform.ReplaceTotalMassWithStellarPlusDark, al.twins.compare, al.twins.TwinComparison, al.twins.compare_to_data, aplt.subplot_twin_comparison, TransformationRecord.truth and a flat "truth" list in twin.json. Nothing removed. Issue PyAutoLabs/PyAutoLens#805. Rulings honoured (2026-10-28): R3 (aperture = circle of radius theta_E,parent from the whole tracer, centred on the replaced galaxy; M_ref = the replaced total-mass components only, shear excluded), R4 (NFWSph on the light centre, fixed r_s given in kpc and converted with the setup cosmology) and R6 (library plotter, aplt conventions). The split is solved by ag.MassBudget (PyAutoGalaxy#676). The light basis is re-expressed as lmp.Gaussian in place, so the light is unchanged and not double-counted. Also makes test__simulate__observed_route_leaves_global_random_state order-independent: the first import of jax in a process draws from NumPy's global RNG (jax/_src/clusters/k8s_cluster.py, reached lazily via autoarray AbstractNDArray.__getitem__), so the test warms the route before seeding. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_016nhQ8NtZ6QtuqxKKAnhhdg --- autolens/plot/__init__.py | 2 + autolens/twins/__init__.py | 11 + autolens/twins/_galaxies.py | 64 +++ autolens/twins/aperture.py | 103 +++++ autolens/twins/compare.py | 351 +++++++++++++++++ autolens/twins/plot/__init__.py | 0 autolens/twins/plot/twin_plots.py | 154 ++++++++ autolens/twins/record.py | 10 + autolens/twins/simulate.py | 9 + autolens/twins/transform/__init__.py | 1 + autolens/twins/transform/stellar_dark.py | 479 +++++++++++++++++++++++ docs/api/twins.rst | 87 ++++ test_autolens/twins/conftest.py | 102 +++++ test_autolens/twins/test_compare.py | 108 +++++ test_autolens/twins/test_stellar_dark.py | 392 +++++++++++++++++++ test_autolens/twins/test_twins.py | 8 + 16 files changed, 1881 insertions(+) create mode 100644 autolens/twins/_galaxies.py create mode 100644 autolens/twins/aperture.py create mode 100644 autolens/twins/compare.py create mode 100644 autolens/twins/plot/__init__.py create mode 100644 autolens/twins/plot/twin_plots.py create mode 100644 autolens/twins/transform/stellar_dark.py create mode 100644 test_autolens/twins/conftest.py create mode 100644 test_autolens/twins/test_compare.py create mode 100644 test_autolens/twins/test_stellar_dark.py diff --git a/autolens/plot/__init__.py b/autolens/plot/__init__.py index a9b9495a8..7b31d5146 100644 --- a/autolens/plot/__init__.py +++ b/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 diff --git a/autolens/twins/__init__.py b/autolens/twins/__init__.py index b67f7a6f1..91074aff3 100644 --- a/autolens/twins/__init__.py +++ b/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", ] diff --git a/autolens/twins/_galaxies.py b/autolens/twins/_galaxies.py new file mode 100644 index 000000000..c78ac7fd0 --- /dev/null +++ b/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.."`` (e.g. ``"galaxies.lens.mass"``), +and a tracer-level field is ``"fields."``. +""" + +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_"``. + """ + 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', ..." + ) diff --git a/autolens/twins/aperture.py b/autolens/twins/aperture.py new file mode 100644 index 000000000..a1e755232 --- /dev/null +++ b/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." + ) diff --git a/autolens/twins/compare.py b/autolens/twins/compare.py new file mode 100644 index 000000000..f9ce70068 --- /dev/null +++ b/autolens/twins/compare.py @@ -0,0 +1,351 @@ +""" +Diagnostics of what a transformation does to a twin, before any inference (lens-twin +report section 2.4): + +- the physical model: convergence and deflection difference maps, enclosed-mass + profiles M( np.ndarray: + """ + twin - parent, with entries that are identical (including both infinite, e.g. at + an isothermal centre lying on a pixel centre) set to exactly zero rather than NaN. + """ + twin_values = np.asarray(twin_values) + parent_values = np.asarray(parent_values) + with np.errstate(invalid="ignore"): + return np.where(twin_values == parent_values, 0.0, twin_values - parent_values) + + +class TwinComparison: + def __init__( + self, + grid: aa.Grid2D, + convergence_parent: aa.Array2D, + convergence_twin: aa.Array2D, + deflections_parent, + deflections_twin, + centre: Tuple[float, float], + radii: np.ndarray, + mass_enclosed_parent: np.ndarray, + mass_enclosed_twin: np.ndarray, + tangential_critical_curves_parent: List[aa.Grid2DIrregular], + tangential_critical_curves_twin: List[aa.Grid2DIrregular], + einstein_radius_parent: Optional[float], + einstein_radius_twin: Optional[float], + image_parent: aa.Array2D, + image_twin: aa.Array2D, + noise_map: aa.Array2D, + ): + """ + The result of ``al.twins.compare(parent, twin)``. Differences are twin minus + parent. ``einstein_radius_*`` is ``None`` when that tracer has no tangential + critical curve inside the parent's setup grid. + + Attributes + ---------- + grid + The uniform grid of the convergence and deflection maps. + radii, mass_enclosed_parent, mass_enclosed_twin + M( aa.Array2D: + """kappa_twin - kappa_parent on ``grid``.""" + return aa.Array2D( + values=_difference( + self.convergence_twin.array, self.convergence_parent.array + ), + mask=self.grid.mask, + ) + + @property + def deflections_difference(self) -> aa.Grid2D: + """alpha_twin - alpha_parent, (y, x) on ``grid``.""" + return aa.Grid2D( + values=_difference( + self.deflections_twin.array, self.deflections_parent.array + ), + mask=self.grid.mask, + ) + + @property + def deflections_y_difference(self) -> aa.Array2D: + return aa.Array2D( + values=self.deflections_difference.array[:, 0], mask=self.grid.mask + ) + + @property + def deflections_x_difference(self) -> aa.Array2D: + return aa.Array2D( + values=self.deflections_difference.array[:, 1], mask=self.grid.mask + ) + + @property + def mass_enclosed_difference(self) -> np.ndarray: + return _difference(self.mass_enclosed_twin, self.mass_enclosed_parent) + + @property + def normalized_residual_map(self) -> aa.Array2D: + """(I_twin - I_parent) / sigma on the parent's masked grid.""" + return aa.Array2D( + values=_difference(self.image_twin.array, self.image_parent.array) + / self.noise_map.array, + mask=self.image_parent.mask, + ) + + @property + def delta_chi_squared_image(self) -> float: + """ + Sum over the parent's unmasked pixels of (I_twin - I_parent)^2 / sigma^2, noise + free: how strongly the parent's observation distinguishes the two models. + """ + return float(np.sum(self.normalized_residual_map.array**2)) + + def __repr__(self): + return ( + f"TwinComparison(delta_chi_squared_image={self.delta_chi_squared_image:.6g}, " + f"einstein_radius_parent={self.einstein_radius_parent}, " + f"einstein_radius_twin={self.einstein_radius_twin})" + ) + + +def _uniform_grid_from(mask: aa.Mask2D) -> aa.Grid2D: + return aa.Grid2D.uniform( + shape_native=mask.shape_native, + pixel_scales=mask.pixel_scales, + origin=mask.origin, + over_sample_size=1, + ) + + +def _default_centre_from(tracer) -> Tuple[float, float]: + """The centre of the first non-light, non-shear mass profile of the lowest-redshift galaxy.""" + for galaxy in tracer.galaxies_ascending_redshift: + for profile in galaxy.profile_dict.values(): + if ( + isinstance(profile, MassProfile) + and not isinstance(profile, LightProfile) + and not isinstance(profile, ag.mp.ExternalShear) + ): + return tuple(float(value) for value in profile.centre) + return (0.0, 0.0) + + +def _einstein_radius_or_none(tracer, grid, pixel_scale) -> Optional[float]: + try: + return einstein_radius_from(tracer=tracer, grid=grid, pixel_scale=pixel_scale) + except ValueError: + return None + + +def _on_setup(twin, setup): + from autolens.twins.twin import Twin + + if twin.setup is setup: + return twin + return Twin(reference=twin.reference, tracer=twin.tracer, setup=setup) + + +def compare( + parent, + twin, + grid: Optional[aa.Grid2D] = None, + radii: Optional[Sequence[float]] = None, + centre: Optional[Tuple[float, float]] = None, + pixel_scale: float = 0.05, + aperture_settings: Optional[dict] = None, +) -> TwinComparison: + """ + Compare ``twin`` with ``parent`` (usually ``Twin.baseline(reference)``). + + Parameters + ---------- + parent, twin + The two ``Twin`` objects. Both images are computed on the parent's setup (its + masked grid, PSF and noise map), whatever setup the twin carries. + grid + The uniform grid of the convergence and deflection maps and of the M( float: + """ + Delta chi^2_data = chi^2(twin | data) - chi^2(parent | data): two ``al.FitImaging`` of + the observed ``dataset``, one with each tracer. Positive means the observed data + disfavour the twin relative to the parent. + + This is a different number from ``TwinComparison.delta_chi_squared_image`` (which is + noise free and compares the two models with each other, not with the data). + + Parameters + ---------- + twin + The transformed twin. + dataset + The masked observed ``al.Imaging`` (e.g. ``ObservationSetup. + dataset_from_fit_directory(fit_path)``). + parent + The twin to compare against; defaults to the reference (``twin.reference.tracer``). + + Raises + ------ + NotImplementedError + If either tracer has a pixelized source. + """ + from autolens.imaging.fit_imaging import FitImaging + + tracer_parent = twin.reference.tracer if parent is None else parent.tracer + _check_parametric(tracer_parent, "parent") + _check_parametric(twin.tracer, "twin") + + chi_squared_twin = FitImaging(dataset=dataset, tracer=twin.tracer).chi_squared + chi_squared_parent = FitImaging(dataset=dataset, tracer=tracer_parent).chi_squared + return float(chi_squared_twin - chi_squared_parent) diff --git a/autolens/twins/plot/__init__.py b/autolens/twins/plot/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/autolens/twins/plot/twin_plots.py b/autolens/twins/plot/twin_plots.py new file mode 100644 index 000000000..b149da559 --- /dev/null +++ b/autolens/twins/plot/twin_plots.py @@ -0,0 +1,154 @@ +""" +Plots of ``al.twins.compare`` diagnostics (ruling R6: the plotter lives in the +library and follows the ``aplt`` conventions of ``autolens/lens/plot/tracer_plots.py``). +""" + +from typing import Optional + +import numpy as np + +from autogalaxy.util.plot_utils import plot_array +from autoarray.plot.utils import ( + conf_subplot_figsize, + hide_unused_axes, + numpy_lines as _to_lines, + save_figure, + subplots, + tight_layout, +) + +PARENT_COLOR = "white" +TWIN_COLOR = "red" + + +def _curve_lines(curves): + return _to_lines(list(curves) if curves else []) or [] + + +def subplot_twin_comparison( + comparison, + output_path: Optional[str] = None, + output_filename: str = "twin_comparison", + output_format: Optional[str] = None, + colormap: Optional[str] = None, + title_prefix: Optional[str] = None, +): + """ + A 2x3 subplot of a ``TwinComparison`` (``al.twins.compare(parent, twin)``). + + Panels: + 0: parent noise-free image (parent critical curve white, twin red) + 1: twin noise-free image, on the parent's masked grid + 2: (I_twin - I_parent) / sigma, titled with delta chi^2_image + 3: convergence difference kappa_twin - kappa_parent + 4: enclosed mass M( Dict: + """ + The flat ground-truth parameters the transformation set (``procedure["truth"]``, + e.g. the mass-to-light ratio, ``kappa_s``, scale radius, inner slope and dark + fraction), or an empty dictionary. ``SimulatedTwin.twin_json`` writes them under + ``"truth"`` so inference can be compared with them without re-deriving them. + """ + return copy.deepcopy(dict(self.procedure.get("truth", {}))) + def __repr__(self): return ( f"TransformationRecord(name={self.name!r}, changed={self.changed}, " diff --git a/autolens/twins/simulate.py b/autolens/twins/simulate.py index 6d057e640..b81fdffab 100644 --- a/autolens/twins/simulate.py +++ b/autolens/twins/simulate.py @@ -112,11 +112,20 @@ def twin_json(self) -> dict: the library versions now, the transformation records (dictable dictionaries), the setup fields that differ from the reference setup, and the seed, noise route and simulator settings. + + ``"truth"`` lists, flat and in chain order, the ground-truth parameters each + transformation reports (``TransformationRecord.truth``), each with its + ``"transformation"`` name and ``"record_index"``; records with none are skipped. """ return { "reference": self.twin.reference.provenance, "versions": _util.library_versions(), "records": [to_dict(record) for record in self.twin.records], + "truth": [ + {"transformation": record.name, "record_index": index, **record.truth} + for index, record in enumerate(self.twin.records) + if record.truth + ], "setup_overrides": self.twin.setup_overrides(), "seed": self.seed, "noise": self.noise, diff --git a/autolens/twins/transform/__init__.py b/autolens/twins/transform/__init__.py index 2c1fd7c06..560cfb929 100644 --- a/autolens/twins/transform/__init__.py +++ b/autolens/twins/transform/__init__.py @@ -11,6 +11,7 @@ from typing import Protocol, Tuple, runtime_checkable from autolens.twins.transform.identity import Identity +from autolens.twins.transform.stellar_dark import ReplaceTotalMassWithStellarPlusDark @runtime_checkable diff --git a/autolens/twins/transform/stellar_dark.py b/autolens/twins/transform/stellar_dark.py new file mode 100644 index 000000000..b50a16b28 --- /dev/null +++ b/autolens/twins/transform/stellar_dark.py @@ -0,0 +1,479 @@ +""" +``ReplaceTotalMassWithStellarPlusDark``: replace a galaxy's total-mass profile(s) with +stellar mass that follows its MGE light plus a dark-matter halo, normalised so their +projected mass inside the parent's Einstein radius equals that of the replaced mass. + +The normalisation convention is PyAutoGalaxy's ``docs/api/mass_budget.rst`` (rulings +R3 and R4, 2026-10-28); the solve is ``ag.MassBudget``. +""" + +import copy +from typing import Dict, Optional, Union + +import numpy as np + +import autoarray as aa +import autogalaxy as ag +from autogalaxy.profiles.light.abstract import LightProfile +from autogalaxy.profiles.mass.abstract.abstract import MassProfile +from autonerves.class_path import get_class_path + +from autolens.twins._galaxies import galaxy_dict_from, galaxy_name_from +from autolens.twins.aperture import ( + check_aperture_in_grid, + einstein_radius_aperture_from, + einstein_radius_from, +) +from autolens.twins.record import TransformationRecord + +NAME = "ReplaceTotalMassWithStellarPlusDark" + +CONVENTION_PAGE = "PyAutoGalaxy docs/api/mass_budget.rst (rulings R3, R4)" + +CONVENTION_TEXT = ( + "Aperture (R3): a circle of radius theta_E,parent, the Einstein radius of the whole " + "parent lens model (tracer-level, from the area of its tangential critical curve via " + "LensCalc), centred on the mass centre of the replaced galaxy. Reference mass (R3): " + "M_ref is the projected mass of the replaced mass component(s) inside the aperture, " + "by grid integration of their convergence (mass_angular_within_circle_via_grid_from); " + "external shear and all other galaxies are left untouched and do not enter M_ref. " + "Stellar component: a Basis of lmp.Gaussian copied from the galaxy's light basis " + "(centre, ell_comps, sigma and solved intensity of every Gaussian) with one shared " + "mass-to-light ratio Upsilon = f_* M_ref / M_light(theta_E), M_light being the same " + "grid integral of the basis at Upsilon = 1. Dark component (R4): an NFWSph (or a " + "gNFWSph with a fixed inner_slope) centred on the light with a fixed scale radius, " + "given in kpc and converted to arc-seconds with the setup cosmology at the galaxy " + "redshift; kappa_s = (1 - f_*) M_ref / M(kappa_s = 1), exact because the convergence " + "is linear in kappa_s. All masses are angular (convergence units, arcsec^2). " + f"Source: {CONVENTION_PAGE}." +) + +DEFAULT_APERTURE_SETTINGS = { + "sub_size_interior": 8, + "sub_size_edge": 16, + "sub_size_centre": 64, + "centre_refine_pixels": 2.0, +} + + +def _plain(value): + """ + JSON-able plain data for a record: NumPy scalars become floats / ints / bools and + tuples become lists. ``al.from_dict`` drops ``None`` values inside dictionaries, so + a ``None`` here is a bug and raises. + """ + if value is None: + raise ValueError("A TransformationRecord value must not be None.") + if isinstance(value, dict): + return {str(key): _plain(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_plain(item) for item in value] + if isinstance(value, (bool, np.bool_)): + return bool(value) + if isinstance(value, (int, np.integer)): + return int(value) + if isinstance(value, (float, np.floating)): + return float(value) + if isinstance(value, str): + return value + raise TypeError( + f"A TransformationRecord value must be plain data, got {type(value).__name__}." + ) + + +def _has_mass(profile) -> bool: + if isinstance(profile, ag.lp_basis.Basis): + return len(profile.mass_profile_list) > 0 + return isinstance(profile, MassProfile) + + +def _is_external_shear(profile) -> bool: + return isinstance(profile, ag.mp.ExternalShear) + + +def _grid_geometry(grid: aa.Grid2D) -> Dict: + return { + "shape_native": list(grid.mask.shape_native), + "pixel_scales": [float(scale) for scale in grid.mask.pixel_scales], + "origin": [float(value) for value in grid.mask.origin], + } + + +class ReplaceTotalMassWithStellarPlusDark: + def __init__( + self, + galaxy: str, + light: str, + *, + scale_radius_kpc: float, + f_dark: float = 0.5, + budget: Optional[Union[ag.MassBudget, Dict[str, Dict[str, float]]]] = None, + dark_cls=None, + keep_shear: bool = True, + aperture_settings: Optional[Dict] = None, + einstein_radius_pixel_scale: float = 0.05, + grid: Optional[aa.Grid2D] = None, + dark_name: str = "dark", + **dark_fixed, + ): + """ + Replace the total-mass profile(s) of one galaxy by stellar mass following its MGE + light plus a dark-matter halo, keeping the replaced galaxy's projected mass inside + theta_E,parent (the one deliberate preservation, recorded in ``preserved``). + + ``apply(twin)`` returns ``(tracer, setup, record)`` and mutates nothing: + + - every ``MassProfile`` attribute of the galaxy that is not light and not an + ``ExternalShear`` (e.g. its ``Isothermal``) is removed; + - the light basis attribute ``light`` is replaced by the stellar ``Basis`` of + ``lmp.Gaussian`` that ``ag.MassBudget`` solves. Its members copy the centre, + ``ell_comps``, ``sigma`` and intensity of the light Gaussians, so the galaxy's + light is unchanged and its stellar mass is ``Upsilon`` times that light. (Adding + the stellar basis beside the light basis would count the light twice.) + - the dark profile is added as the attribute ``dark_name``, centred on the light; + - external shear (galaxy attribute or tracer field), the source and all other + galaxies are unchanged (``keep_shear=False`` removes the replaced galaxy's + ``ExternalShear`` attributes and the tracer's ``ExternalShear`` fields). + + The aperture is ``einstein_radius_aperture_from(tracer, setup.grid, centre)``: the + whole parent tracer's Einstein radius, centred on the replaced mass centre. + + Parameters + ---------- + galaxy + The galaxy to transform: its name in the tracer (the key of the fit's + ``af.Collection``, e.g. ``"lens"``, or ``"galaxy_"`` for a tracer built + from a list), or a role from the reference's provenance ``roles``. + light + The name of the galaxy's MGE light attribute: a ``Basis`` of solved-intensity + ``lp.Gaussian`` profiles (linear, operated or non-Gaussian members raise + ``ProfileException``: they have no ``lmp.Gaussian`` counterpart). + scale_radius_kpc + The fixed dark-profile scale radius, in kpc (ruling R4). It is converted to + arc-seconds with ``cosmology.arcsec_per_kpc_from(redshift=galaxy.redshift)``, + the cosmology being the setup's, or the tracer's when the setup has none. + f_dark + The dark fraction of the replaced mass inside theta_E,parent (stellar + fraction ``1 - f_dark``). Ignored when ``budget`` is given. + budget + Alternatively an ``ag.MassBudget`` or its ``fractions`` dictionary, + ``{galaxy: {"stellar": f_*, "dark": f_dark}}``; only its fractions are used + (the aperture is always the Einstein-radius aperture). A budget naming more + than one galaxy raises ``NotImplementedError`` (ruling R7). + dark_cls + The dark profile class; ``al.mp.NFWSph`` by default, or e.g. ``al.mp.gNFWSph`` + with ``inner_slope`` passed in ``dark_fixed``. + keep_shear + Keep the external shear (default). It never enters M_ref. + aperture_settings + ``sub_size_interior``, ``sub_size_edge``, ``sub_size_centre`` and + ``centre_refine_pixels`` of the aperture integrator (defaults: those of + ``ag.CircularAperture``). A ``gNFW`` convergence does one ``quad`` per point, + so pass coarse sub-sizes for it (e.g. 2, 4, 4). + einstein_radius_pixel_scale + The resolution of the grid on which the critical curve, and so theta_E, is + traced. + grid + The uniform grid on which aperture masses are integrated; defaults to the + setup mask's full (unmasked) geometry. The aperture must lie within it. + dark_name + The galaxy attribute the dark profile is stored as. + dark_fixed + Further fixed constructor arguments of ``dark_cls``, e.g. ``inner_slope``. + """ + if budget is not None: + fractions = ( + budget.fractions if isinstance(budget, ag.MassBudget) else dict(budget) + ) + # Validates the fractions; raises NotImplementedError for several galaxies. + ag.MassBudget(aperture=None, fractions=fractions) + self.budget_galaxy = next(iter(fractions)) + self.fractions = { + "stellar": float(fractions[self.budget_galaxy]["stellar"]), + "dark": float(fractions[self.budget_galaxy]["dark"]), + } + else: + if not 0.0 <= f_dark <= 1.0: + raise ValueError(f"f_dark must lie in [0, 1], got {f_dark}.") + self.budget_galaxy = None + self.fractions = {"stellar": 1.0 - float(f_dark), "dark": float(f_dark)} + + if not scale_radius_kpc > 0.0: + raise ValueError( + f"scale_radius_kpc must be positive, got {scale_radius_kpc}." + ) + + self.galaxy = galaxy + self.light = light + self.scale_radius_kpc = float(scale_radius_kpc) + self.dark_cls = ag.mp.NFWSph if dark_cls is None else dark_cls + self.keep_shear = bool(keep_shear) + self.aperture_settings = { + **DEFAULT_APERTURE_SETTINGS, + **dict(aperture_settings or {}), + } + self.einstein_radius_pixel_scale = float(einstein_radius_pixel_scale) + self.grid = grid + self.dark_name = dark_name + self.dark_fixed = dict(dark_fixed) + + @property + def f_dark(self) -> float: + return self.fractions["dark"] + + def requested(self) -> Dict: + """The requested arguments, as plain data (``grid`` as its geometry).""" + return _plain( + { + "galaxy": self.galaxy, + "light": self.light, + "scale_radius_kpc": self.scale_radius_kpc, + "f_dark": self.f_dark, + "fractions": self.fractions, + "dark_cls": get_class_path(self.dark_cls), + "dark_fixed": self.dark_fixed, + "keep_shear": self.keep_shear, + "aperture_settings": self.aperture_settings, + "einstein_radius_pixel_scale": self.einstein_radius_pixel_scale, + "grid": ( + "setup mask geometry" + if self.grid is None + else _grid_geometry(self.grid) + ), + "dark_name": self.dark_name, + } + ) + + def apply(self, twin): + tracer = twin.tracer + setup = twin.setup + + name = galaxy_name_from(twin, self.galaxy) + if self.budget_galaxy is not None and ( + galaxy_name_from(twin, self.budget_galaxy) != name + ): + raise ValueError( + f"The budget names galaxy {self.budget_galaxy!r} but the transformation " + f"replaces {name!r}." + ) + galaxy = galaxy_dict_from(tracer)[name] + path = f"galaxies.{name}" + + light_basis = getattr(galaxy, self.light, None) + if not isinstance(light_basis, ag.lp_basis.Basis): + raise ValueError( + f"{path}.{self.light} must be an MGE light Basis, got " + f"{type(light_basis).__name__}." + ) + # Raises ProfileException for members with no lmp.Gaussian counterpart, before + # any expensive step. + light_basis.mass_basis_from(mass_to_light_ratio=1.0) + + profile_dict = galaxy.profile_dict + stellar_already = [ + key + for key, profile in profile_dict.items() + if isinstance(profile, LightProfile) and _has_mass(profile) + ] + if stellar_already: + raise ValueError( + f"{path} already has light-and-mass components {stellar_already}; " + f"{NAME} replaces a total-mass profile only." + ) + replaced = { + key: profile + for key, profile in profile_dict.items() + if isinstance(profile, MassProfile) + and not isinstance(profile, LightProfile) + and not _is_external_shear(profile) + } + if not replaced: + raise ValueError(f"{path} has no total-mass profile to replace.") + if self.dark_name in profile_dict and self.dark_name not in replaced: + raise ValueError( + f"{path} already has an attribute {self.dark_name!r}; pass another " + f"dark_name." + ) + shear_keys = [ + key for key, profile in profile_dict.items() if _is_external_shear(profile) + ] + shear_field_indexes = [ + index + for index, field in enumerate(getattr(tracer, "fields", [])) + if _is_external_shear(field) + ] + + mass_centre = tuple( + float(value) for value in next(iter(replaced.values())).centre + ) + light_centre = tuple( + float(value) for value in light_basis.profile_list[0].centre + ) + + aperture = einstein_radius_aperture_from( + tracer=tracer, + grid=setup.grid, + centre=mass_centre, + pixel_scale=self.einstein_radius_pixel_scale, + **self.aperture_settings, + ) + einstein_radius_parent = aperture.radius + + grid = self.grid + if grid is None: + grid = aa.Grid2D.uniform( + shape_native=setup.mask.shape_native, + pixel_scales=setup.mask.pixel_scales, + origin=setup.mask.origin, + over_sample_size=1, + ) + check_aperture_in_grid(aperture=aperture, grid=grid) + + cosmology = setup.cosmology if setup.cosmology is not None else tracer.cosmology + arcsec_per_kpc = float(cosmology.arcsec_per_kpc_from(redshift=galaxy.redshift)) + scale_radius_arcsec = self.scale_radius_kpc * arcsec_per_kpc + + result = ag.MassBudget( + aperture=aperture, fractions={name: self.fractions} + ).solve( + reference_mass=ag.Galaxy(redshift=galaxy.redshift, **replaced), + light_basis=light_basis, + dark_cls=self.dark_cls, + scale_radius=scale_radius_arcsec, + centre=light_centre, + grid=grid, + **self.dark_fixed, + ) + + new_tracer = copy.deepcopy(tracer) + new_galaxy = galaxy_dict_from(new_tracer)[name] + + changed = [] + removed = {} + added = {} + for key, profile in replaced.items(): + delattr(new_galaxy, key) + changed.append(f"{path}.{key}") + removed[f"{path}.{key}"] = get_class_path(type(profile)) + if result.stellar is not None: + setattr(new_galaxy, self.light, result.stellar) + changed.append(f"{path}.{self.light}") + added[f"{path}.{self.light}"] = "Basis of lmp.Gaussian" + if result.dark is not None: + setattr(new_galaxy, self.dark_name, result.dark) + if f"{path}.{self.dark_name}" not in changed: + changed.append(f"{path}.{self.dark_name}") + added[f"{path}.{self.dark_name}"] = get_class_path(type(result.dark)) + + side_effects = [] + shear_paths = [f"{path}.{key}" for key in shear_keys] + [ + f"fields.{index}" for index in shear_field_indexes + ] + if not self.keep_shear: + for key in shear_keys: + delattr(new_galaxy, key) + new_tracer.fields = [ + field + for index, field in enumerate(new_tracer.fields) + if index not in shear_field_indexes + ] + changed.extend(shear_paths) + for shear_path in shear_paths: + removed[shear_path] = get_class_path(ag.mp.ExternalShear) + if shear_paths: + side_effects.append(f"external shear removed: {shear_paths}") + elif shear_paths: + side_effects.append( + f"external shear kept unchanged and excluded from M_ref: {shear_paths}" + ) + + try: + einstein_radius_twin = einstein_radius_from( + tracer=new_tracer, + grid=setup.grid, + pixel_scale=self.einstein_radius_pixel_scale, + ) + side_effects.append( + f"Einstein radius of the twin {einstein_radius_twin:.6f} arcsec (parent " + f"{einstein_radius_parent:.6f} arcsec), recomputed on the same grid; " + f"reported, not asserted" + ) + except ValueError: + einstein_radius_twin = None + side_effects.append( + "the twin has no tangential critical curve inside the setup grid; its " + "Einstein radius is undefined (reported, not asserted)" + ) + if result.dark is not None: + side_effects.append( + f"dark halo centred on the light centre {list(light_centre)}; the " + f"replaced mass centre was {list(mass_centre)}" + ) + if result.stellar is not None: + side_effects.append( + f"light basis {path}.{self.light} re-expressed as lmp.Gaussian with " + f"identical centre, ell_comps, intensity and sigma, so the light is " + f"unchanged" + ) + side_effects.append("the source and all other galaxies are unchanged") + + derived = dict(result.derived) + derived["fixed"] = dict(derived.get("fixed") or {}) + dark_profile = result.dark + inner_slope = self.dark_fixed.get( + "inner_slope", getattr(dark_profile, "inner_slope", 1.0) + ) + + truth = { + "galaxy": name, + "redshift": galaxy.redshift, + "mass_to_light_ratio": derived["mass_to_light_ratio"], + "kappa_s": derived["kappa_s"], + "scale_radius_arcsec": scale_radius_arcsec, + "scale_radius_kpc": self.scale_radius_kpc, + "inner_slope": inner_slope, + "dark_cls": get_class_path(self.dark_cls), + "f_dark": derived["f_dark"], + "f_stellar": self.fractions["stellar"], + "einstein_radius_parent": einstein_radius_parent, + "mass_reference": derived["mass_reference"], + "dark_centre": list(light_centre), + } + + procedure = { + "text": CONVENTION_TEXT, + "convention": CONVENTION_PAGE, + "galaxy": name, + "derived": derived, + "einstein_radius_parent": einstein_radius_parent, + "einstein_radius_grid": "setup.grid", + "einstein_radius_pixel_scale": self.einstein_radius_pixel_scale, + "aperture": { + "radius": aperture.radius, + "centre": list(aperture.centre), + **self.aperture_settings, + }, + "aperture_grid": _grid_geometry(grid), + "scale_radius_kpc": self.scale_radius_kpc, + "scale_radius_arcsec": scale_radius_arcsec, + "arcsec_per_kpc": arcsec_per_kpc, + "redshift": galaxy.redshift, + "cosmology": get_class_path(type(cosmology)), + "cosmology_source": "setup" if setup.cosmology is not None else "tracer", + "removed": removed, + "added": added, + "truth": truth, + } + if einstein_radius_twin is not None: + procedure["einstein_radius_twin"] = einstein_radius_twin + + record = TransformationRecord( + name=NAME, + requested=self.requested(), + procedure=_plain(procedure), + changed=changed, + side_effects=side_effects, + preserved=[f"projected mass of {name} inside theta_E,parent"], + ) + + return new_tracer, setup, record diff --git a/docs/api/twins.rst b/docs/api/twins.rst index a2cc34fb6..43717b91e 100644 --- a/docs/api/twins.rst +++ b/docs/api/twins.rst @@ -37,6 +37,77 @@ the setup's ``exposure_time`` and ``background_sky_level``. reloaded = al.twins.SimulatedTwin.from_directory("twins/lens_001/baseline_seed_1") resimulated = reloaded.resimulate() # bit-identical to reloaded.dataset.data +Galaxies are addressed by name: a tracer loaded from a fit's ``files/tracer.json`` +keeps the keys of the model's ``af.Collection`` (e.g. ``"lens"``, ``"source"``), and a +tracer built from a plain list names its galaxies ``"galaxy_0"``, ``"galaxy_1"``, .... +A transformation also accepts a role from the reference's provenance ``roles`` +(``roles={"lens": "main_deflector"}`` lets ``galaxy="main_deflector"`` resolve to +``"lens"``). Component paths in a record read ``"galaxies.lens.mass"``. + +Stellar plus dark matter +------------------------ + +``transform.ReplaceTotalMassWithStellarPlusDark`` replaces a galaxy's total-mass +profile(s) (every mass profile that is not light and not an ``ExternalShear``) with +stellar mass following its MGE light plus a dark-matter halo. The split follows the +mass-normalisation convention of PyAutoGalaxy's Mass Budget page +(``PyAutoGalaxy/docs/api/mass_budget.rst``, rulings R3 and R4): + +- the aperture is a circle of radius theta_E,parent, the Einstein radius of the whole + parent tracer (``LensCalc``, on the setup's masked grid), centred on the replaced + galaxy's mass centre (``einstein_radius_aperture_from``); +- M_ref is the projected mass of the replaced component(s) inside it; external shear + and other galaxies are untouched and excluded; +- ``ag.MassBudget`` solves one mass-to-light ratio for the light ``Basis`` (which is + re-expressed as ``lmp.Gaussian`` members with identical light, so the image of the + lens light is unchanged) and the ``kappa_s`` of an ``NFWSph`` (or ``gNFWSph`` with a + fixed ``inner_slope``) centred on the light, with ``scale_radius_kpc`` converted to + arc-seconds by the setup cosmology (or the tracer's) at the galaxy redshift. + +The record's ``preserved`` names the one deliberate preservation, the galaxy's +projected mass inside theta_E,parent. Its ``procedure`` quotes the convention and the +derived numbers, and its ``truth`` (mass-to-light ratio, ``kappa_s``, scale radius in +arc-seconds and kpc, inner slope, dark fraction) is written flat under ``"truth"`` in +``twin.json``. A ``gNFW`` convergence does one ``quad`` per point: pass coarse +``aperture_settings`` for it. Budgets naming several galaxies raise +``NotImplementedError`` (ruling R7). + +.. code-block:: python + + baseline = al.twins.Twin.baseline(reference) + twin = baseline.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", light="bulge", scale_radius_kpc=20.0, f_dark=0.5 + ) + ) + twin.records[-1].truth # flat ground truth + +Diagnostics +----------- + +``compare(parent, twin)`` returns a ``TwinComparison``: convergence and deflection +difference maps, enclosed-mass profiles M( 0 + + +def test__compare__transformed_twin(witness_baseline, witness_twin): + comparison = al.twins.compare( + witness_baseline, witness_twin, aperture_settings=COARSE_APERTURE + ) + record = witness_twin.records[-1] + + assert comparison.delta_chi_squared_image > 0.0 + assert comparison.einstein_radius_parent == pytest.approx( + record.procedure["einstein_radius_parent"], rel=1e-12 + ) + assert comparison.einstein_radius_twin == pytest.approx( + record.procedure["einstein_radius_twin"], rel=1e-12 + ) + assert len(comparison.tangential_critical_curves_parent) > 0 + assert len(comparison.tangential_critical_curves_twin) > 0 + assert np.any(comparison.convergence_difference.array != 0.0) + + # The aperture mass inside theta_E,parent is the transformation's one preservation. + index = int(np.argmin(np.abs(comparison.radii - comparison.einstein_radius_parent))) + assert comparison.radii[index] == pytest.approx(comparison.einstein_radius_parent) + assert comparison.mass_enclosed_twin[index] == pytest.approx( + comparison.mass_enclosed_parent[index], rel=1.0e-3 + ) + + +def test__compare_to_data(witness_baseline, witness_twin, masked_imaging_7x7): + delta = al.twins.compare_to_data(witness_twin, masked_imaging_7x7) + assert np.isfinite(delta) + assert delta != 0.0 + + assert al.twins.compare_to_data(witness_baseline, masked_imaging_7x7) == 0.0 + + +def test__compare_to_data__pixelized_source_raises( + witness_baseline, masked_imaging_7x7, tracer_x2_plane_inversion_7x7 +): + twin = al.twins.Twin( + reference=witness_baseline.reference, tracer=tracer_x2_plane_inversion_7x7 + ) + with pytest.raises(NotImplementedError, match="pixelized"): + al.twins.compare_to_data(twin, masked_imaging_7x7) + + +def test__subplot_twin_comparison__writes_png( + witness_baseline, witness_twin, tmp_path, monkeypatch +): + import autoarray.plot.utils as plot_utils + + monkeypatch.setattr(plot_utils, "_FAST_PLOTS", False) + monkeypatch.delenv("PYAUTO_OUTPUT_MODE", raising=False) + + comparison = al.twins.compare( + witness_baseline, witness_twin, aperture_settings=COARSE_APERTURE + ) + aplt.subplot_twin_comparison( + comparison, output_path=str(tmp_path), output_format="png" + ) + + assert (tmp_path / "twin_comparison.png").exists() diff --git a/test_autolens/twins/test_stellar_dark.py b/test_autolens/twins/test_stellar_dark.py new file mode 100644 index 000000000..0d3fd581c --- /dev/null +++ b/test_autolens/twins/test_stellar_dark.py @@ -0,0 +1,392 @@ +import json + +import numpy as np +import pytest + +import autogalaxy as ag +import autolens as al +from autolens.twins._galaxies import galaxy_dict_from +from autolens.twins.transform import Transformation + +from test_autolens.twins.conftest import COARSE_APERTURE, GNFW_APERTURE + + +def _canonical(obj) -> str: + return json.dumps(al.to_dict(obj), sort_keys=True) + + +def _aperture_from(record, settings=COARSE_APERTURE): + return ag.CircularAperture( + radius=record.procedure["aperture"]["radius"], + centre=tuple(record.procedure["aperture"]["centre"]), + **settings, + ) + + +def _setup_grid(twin): + mask = twin.setup.mask + return al.Grid2D.uniform( + shape_native=mask.shape_native, + pixel_scales=mask.pixel_scales, + origin=mask.origin, + over_sample_size=1, + ) + + +def test__galaxies_are_addressed_by_tracer_json_keys(witness_reference): + # The tracer reloaded from files/tracer.json keeps the model's galaxy names. + assert list(galaxy_dict_from(witness_reference.tracer)) == ["lens", "source"] + + +def test__galaxies_of_a_list_tracer_are_addressed_by_position(tracer_x2_plane_7x7): + assert list(galaxy_dict_from(tracer_x2_plane_7x7)) == [ + "galaxy_0", + "galaxy_1", + "galaxy_2", + ] + + +def test__witness__50_50_budget_inside_parent_einstein_radius( + witness_baseline, witness_twin +): + assert isinstance( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", light="bulge", scale_radius_kpc=10.0 + ), + Transformation, + ) + + record = witness_twin.records[-1] + parent_lens = witness_baseline.tracer.galaxies.lens + lens = witness_twin.tracer.galaxies.lens + + # The SIE is removed; the light basis is now a stellar lmp.Gaussian basis; a dark halo is added. + assert not hasattr(lens, "mass") + assert all( + isinstance(profile, al.lmp.Gaussian) for profile in lens.bulge.profile_list + ) + assert isinstance(lens.dark, al.mp.NFWSph) + + grid = _setup_grid(witness_twin) + aperture = _aperture_from(record) + mass_reference = aperture.mass_from(mass_obj=parent_lens.mass, grid=grid) + mass_stellar = aperture.mass_from(mass_obj=lens.bulge, grid=grid) + mass_dark = aperture.mass_from(mass_obj=lens.dark, grid=grid) + + assert mass_reference == pytest.approx( + record.procedure["derived"]["mass_reference"], rel=1e-12 + ) + assert mass_stellar / mass_reference == pytest.approx(0.5, abs=1.0e-3) + assert mass_dark / mass_reference == pytest.approx(0.5, abs=1.0e-3) + + # theta_E,parent is the tracer-level Einstein radius of the whole parent. + assert record.procedure["einstein_radius_parent"] == pytest.approx( + al.LensCalc.from_tracer(witness_baseline.tracer).einstein_radius_from( + grid=witness_baseline.setup.grid + ), + rel=1e-12, + ) + assert record.procedure["aperture"]["centre"] == [0.0, 0.0] + + # Shear, light, and the source are unchanged. + assert _canonical(lens.shear) == _canonical(parent_lens.shear) + assert _canonical(witness_twin.tracer.galaxies.source) == _canonical( + witness_baseline.tracer.galaxies.source + ) + assert np.array_equal( + lens.image_2d_from(grid=grid).array, parent_lens.image_2d_from(grid=grid).array + ) + + +def test__witness__record(witness_twin): + record = witness_twin.records[-1] + + assert record.name == "ReplaceTotalMassWithStellarPlusDark" + assert record.preserved == ["projected mass of lens inside theta_E,parent"] + assert record.changed == [ + "galaxies.lens.mass", + "galaxies.lens.bulge", + "galaxies.lens.dark", + ] + assert any("external shear kept" in effect for effect in record.side_effects) + assert any( + "Einstein radius of the twin" in effect for effect in record.side_effects + ) + + assert record.requested["galaxy"] == "lens" + assert record.requested["scale_radius_kpc"] == 10.0 + assert record.requested["fractions"] == {"stellar": 0.5, "dark": 0.5} + assert record.requested["dark_cls"] == "autogalaxy.profiles.mass.dark.nfw.NFWSph" + + procedure = record.procedure + assert "mass_budget.rst" in procedure["text"] + for key in ( + "mass_reference", + "mass_light", + "mass_to_light_ratio", + "kappa_s", + "f_dark", + "scale_radius", + ): + assert key in procedure["derived"] + assert procedure["redshift"] == 0.5 + assert procedure["einstein_radius_twin"] > 0.0 + + truth = record.truth + lens = witness_twin.tracer.galaxies.lens + assert ( + truth["mass_to_light_ratio"] == lens.bulge.profile_list[0].mass_to_light_ratio + ) + assert truth["kappa_s"] == lens.dark.kappa_s + assert truth["scale_radius_arcsec"] == lens.dark.scale_radius + assert truth["inner_slope"] == 1.0 + assert truth["f_dark"] == 0.5 + + def has_none(value): + if value is None: + return True + if isinstance(value, dict): + return any(has_none(item) for item in value.values()) + if isinstance(value, list): + return any(has_none(item) for item in value) + return False + + assert not has_none(record.requested) + assert not has_none(record.procedure) + + +def test__galaxy_addressed_by_role(witness_baseline, witness_twin): + by_role = witness_baseline.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="main_deflector", + light="bulge", + scale_radius_kpc=10.0, + aperture_settings=COARSE_APERTURE, + ) + ) + assert _canonical(by_role.tracer) == _canonical(witness_twin.tracer) + + with pytest.raises(ValueError, match="No galaxy named"): + witness_baseline.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="dwarf", light="bulge", scale_radius_kpc=10.0 + ) + ) + + +def test__does_not_mutate(witness_reference, witness_baseline, stellar_dark): + tracer_before = _canonical(witness_baseline.tracer) + reference_before = _canonical(witness_reference.tracer) + setup_before = _canonical(witness_baseline.setup) + + first = witness_baseline.apply(stellar_dark) + second = witness_baseline.apply(stellar_dark) + + assert _canonical(witness_baseline.tracer) == tracer_before + assert _canonical(witness_reference.tracer) == reference_before + assert _canonical(witness_baseline.setup) == setup_before + assert witness_baseline.records == () + assert first.tracer is not witness_baseline.tracer + assert first.setup is witness_baseline.setup + assert _canonical(first.tracer) == _canonical(second.tracer) + assert first.records == second.records + + +def test__record__round_trips_dict(witness_twin): + record = witness_twin.records[-1] + reloaded = al.from_dict(json.loads(json.dumps(al.to_dict(record)))) + assert reloaded == record + assert reloaded.truth == record.truth + + reloaded_tracer = al.from_dict(json.loads(_canonical(witness_twin.tracer))) + assert _canonical(reloaded_tracer) == _canonical(witness_twin.tracer) + + +def test__record__round_trips_simulated_twin_directory(witness_twin, tmp_path): + simulated = witness_twin.simulate(seed=3) + output = tmp_path / "simulated" + simulated.output_to(output) + + assert (output / "reference_tracer.json").exists() + + with open(output / "twin.json") as f: + twin_json = json.load(f) + truth = twin_json["truth"] + assert len(truth) == 1 + assert truth[0]["transformation"] == "ReplaceTotalMassWithStellarPlusDark" + assert truth[0]["record_index"] == 0 + for key in ( + "mass_to_light_ratio", + "kappa_s", + "scale_radius_arcsec", + "scale_radius_kpc", + "inner_slope", + "f_dark", + "einstein_radius_parent", + ): + assert isinstance(truth[0][key], float) + + reloaded = al.twins.SimulatedTwin.from_directory(output) + assert reloaded.twin.records == witness_twin.records + assert _canonical(reloaded.truth) == _canonical(witness_twin.tracer) + assert np.array_equal( + reloaded.resimulate().dataset.data.native.array, + simulated.dataset.data.native.array, + ) + + +def test__scale_radius_kpc_uses_setup_cosmology(witness_baseline): + cosmology = al.cosmo.FlatLambdaCDM(H0=50.0, Om0=0.4) + twin = witness_baseline.with_setup( + witness_baseline.setup.with_(cosmology=cosmology) + ) + + transformed = twin.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", + light="bulge", + scale_radius_kpc=8.0, + aperture_settings=COARSE_APERTURE, + ) + ) + procedure = transformed.records[-1].procedure + expected = 8.0 * cosmology.arcsec_per_kpc_from(redshift=0.5) + + assert procedure["cosmology_source"] == "setup" + assert procedure["cosmology"].endswith("FlatLambdaCDM") + assert procedure["scale_radius_arcsec"] == pytest.approx(expected, rel=1e-12) + assert transformed.tracer.galaxies.lens.dark.scale_radius == pytest.approx( + expected, rel=1e-12 + ) + + default = witness_baseline.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", + light="bulge", + scale_radius_kpc=8.0, + aperture_settings=COARSE_APERTURE, + ) + ) + default_procedure = default.records[-1].procedure + assert default_procedure["cosmology_source"] == "tracer" + assert default_procedure["scale_radius_arcsec"] == pytest.approx( + 8.0 * al.cosmo.Planck15().arcsec_per_kpc_from(redshift=0.5), rel=1e-12 + ) + assert default_procedure["scale_radius_arcsec"] != pytest.approx(expected) + + +def test__gnfw_inner_slope(witness_baseline): + transformed = witness_baseline.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", + light="bulge", + scale_radius_kpc=10.0, + f_dark=0.3, + dark_cls=al.mp.gNFWSph, + aperture_settings=GNFW_APERTURE, + inner_slope=1.5, + ) + ) + record = transformed.records[-1] + dark = transformed.tracer.galaxies.lens.dark + + assert isinstance(dark, al.mp.gNFWSph) + assert dark.inner_slope == 1.5 + assert record.requested["dark_fixed"] == {"inner_slope": 1.5} + assert record.procedure["derived"]["fixed"] == {"inner_slope": 1.5} + assert record.truth["inner_slope"] == 1.5 + assert record.truth["f_dark"] == 0.3 + + grid = _setup_grid(transformed) + aperture = _aperture_from(record, settings=GNFW_APERTURE) + mass_reference = record.procedure["derived"]["mass_reference"] + assert aperture.mass_from( + mass_obj=dark, grid=grid + ) / mass_reference == pytest.approx(0.3, abs=1.0e-3) + + reloaded = al.from_dict(json.loads(json.dumps(al.to_dict(record)))) + assert reloaded == record + reloaded_tracer = al.from_dict(json.loads(_canonical(transformed.tracer))) + assert reloaded_tracer.galaxies.lens.dark.inner_slope == 1.5 + + +def test__budget(witness_baseline): + with pytest.raises(NotImplementedError): + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", + light="bulge", + scale_radius_kpc=10.0, + budget={ + "lens": {"stellar": 0.5, "dark": 0.5}, + "dwarf": {"stellar": 0.5, "dark": 0.5}, + }, + ) + + with pytest.raises(ValueError): + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", light="bulge", scale_radius_kpc=10.0, f_dark=1.5 + ) + + transformed = witness_baseline.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", + light="bulge", + scale_radius_kpc=10.0, + budget=ag.MassBudget( + aperture=None, fractions={"lens": {"stellar": 0.8, "dark": 0.2}} + ), + aperture_settings=COARSE_APERTURE, + ) + ) + assert transformed.records[-1].truth["f_dark"] == 0.2 + + with pytest.raises(ValueError, match="budget names"): + witness_baseline.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", + light="bulge", + scale_radius_kpc=10.0, + budget={"source": {"stellar": 0.5, "dark": 0.5}}, + ) + ) + + +def test__keep_shear_false_removes_shear(witness_baseline): + transformed = witness_baseline.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", + light="bulge", + scale_radius_kpc=10.0, + keep_shear=False, + aperture_settings=COARSE_APERTURE, + ) + ) + record = transformed.records[-1] + assert not hasattr(transformed.tracer.galaxies.lens, "shear") + assert "galaxies.lens.shear" in record.changed + assert hasattr(witness_baseline.tracer.galaxies.lens, "shear") + + +def test__invalid_light_raises(witness_baseline): + tracer = witness_baseline.tracer + galaxies = galaxy_dict_from(tracer) + sersic_basis = al.lp_basis.Basis(profile_list=[al.lp.SersicSph()]) + lens = al.Galaxy(redshift=0.5, bulge=sersic_basis, mass=galaxies["lens"].mass) + twin = al.twins.Twin( + reference=witness_baseline.reference, + tracer=al.Tracer(galaxies=[lens, galaxies["source"]]), + ) + + with pytest.raises(ag.exc.ProfileException): + twin.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="galaxy_0", light="bulge", scale_radius_kpc=10.0 + ) + ) + + with pytest.raises(ValueError, match="must be an MGE light Basis"): + twin.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="galaxy_0", light="mass", scale_radius_kpc=10.0 + ) + ) diff --git a/test_autolens/twins/test_twins.py b/test_autolens/twins/test_twins.py index fc1889d80..ca3163467 100644 --- a/test_autolens/twins/test_twins.py +++ b/test_autolens/twins/test_twins.py @@ -248,6 +248,14 @@ def test__simulate__seed_determinism_and_sensitivity(twin): def test__simulate__observed_route_leaves_global_random_state(twin): + # Warm the route first. The first `import jax` in a process draws 13 numbers from + # NumPy's global random state: jax 0.11 `jax/_src/clusters/k8s_cluster.py` evaluates + # `np.random.rand(5)` (a default argument of `retry`) and `np.random.rand(8)` (a + # `@retry(wait=...)` decorator) at import. The observed route triggers that import + # lazily, via autoarray `AbstractNDArray.__getitem__` (`import jax.numpy`). This is a + # one-time upstream import side effect, not a draw by `Twin.simulate`. + twin.simulate(seed=1) + np.random.seed(5) expected = np.random.rand()