diff --git a/autolens/twins/__init__.py b/autolens/twins/__init__.py index 91074aff3..673caf171 100644 --- a/autolens/twins/__init__.py +++ b/autolens/twins/__init__.py @@ -24,6 +24,14 @@ Einstein radius (``einstein_radius_aperture_from``); ``compare`` and ``compare_to_data`` report what a transformation changed, and ``aplt.subplot_twin_comparison`` plots it. +``model.stellar_dark_model_from`` builds the model that fits a stellar-plus-dark twin +with the parameterization it was made with (fixed MGE light, one shared mass-to-light +ratio, a dark halo centred on the light, free external shear; broad fixed priors that +never know the truth), and ``recovery.recovery_table_from`` compares a fit's posterior +medians and ``sigma`` bounds with the twin's ``truth`` (a ``RecoveryTable`` with +``all_inside`` and ``to_csv``). Both map ``truth`` keys onto model paths through +``model.truth_paths_from``. + NumPy only: nothing here supports JAX. """ @@ -34,7 +42,9 @@ 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 +from autolens.twins.model import stellar_dark_model_from, truth_paths_from +from autolens.twins.recovery import RecoveryRow, RecoveryTable, recovery_table_from +from autolens.twins import model, recovery, transform __all__ = [ "ObservationSetup", @@ -46,5 +56,12 @@ "compare", "compare_to_data", "einstein_radius_aperture_from", + "stellar_dark_model_from", + "truth_paths_from", + "RecoveryRow", + "RecoveryTable", + "recovery_table_from", + "model", + "recovery", "transform", ] diff --git a/autolens/twins/model.py b/autolens/twins/model.py new file mode 100644 index 000000000..e5d43bded --- /dev/null +++ b/autolens/twins/model.py @@ -0,0 +1,371 @@ +""" +``stellar_dark_model_from``: the inference model of a stellar-plus-dark twin, with no +misspecification. + +A twin made by ``transform.ReplaceTotalMassWithStellarPlusDark`` has a lens galaxy whose +mass is a ``Basis`` of ``lmp.Gaussian`` (its MGE light times one mass-to-light ratio) +plus a dark halo. ``stellar_dark_model_from`` returns the ``af.Collection`` that fits +that twin with the same parameterization: the light of each Gaussian is fixed to the +twin's, and only the mass-to-light ratio, the dark halo and the external shear are +free. It is the twin-based counterpart of +``ag.analysis.chaining_util.mass_light_dark_basis_from``, which builds the same basis +from a light-pipeline ``Result``. + +``truth_paths_from`` maps the transformation's ``truth`` keys onto the model paths of +the free parameters; ``al.twins.recovery.recovery_table_from`` uses it, so the model +and the recovery table cannot disagree about where a parameter lives. + +NumPy only: nothing here supports JAX. +""" + +import inspect +from typing import Dict, Optional, Tuple + +import autofit as af +import autogalaxy as ag +from autogalaxy.profiles.mass.abstract.abstract import MassProfile + +from autolens.twins._galaxies import galaxy_dict_from, galaxy_name_from +from autolens.twins.transform.stellar_dark import NAME as STELLAR_DARK_NAME + +# The default priors of the free parameters, by ``truth`` key. They are broad, fixed +# and never derived from the truth (a prior centred on the answer defeats the +# experiment): +# +# - ``mass_to_light_ratio``: log-uniform over 1e-6 to 1e6 (the PyAutoGalaxy default +# for stellar profiles). Its units are convergence per unit image intensity, so its +# scale depends on the data units; twelve decades cover any of them. +# - ``kappa_s``: log-uniform over 1e-4 to 10. +# - ``scale_radius_arcsec``: uniform over 0.1 to 30 arc-seconds (the NFW default's +# upper limit; roughly 1 to 100 kpc at the redshifts of galaxy-scale lenses). +# - ``inner_slope``: uniform over 0 to 2 (cored to steeper than isothermal). +DEFAULT_PRIORS = { + "mass_to_light_ratio": lambda: af.LogUniformPrior( + lower_limit=1.0e-6, upper_limit=1.0e6 + ), + "kappa_s": lambda: af.LogUniformPrior(lower_limit=1.0e-4, upper_limit=10.0), + "scale_radius_arcsec": lambda: af.UniformPrior(lower_limit=0.1, upper_limit=30.0), + "inner_slope": lambda: af.UniformPrior(lower_limit=0.0, upper_limit=2.0), +} + +# The ``truth`` key of each recoverable parameter and the name of the model parameter +# that holds it; ``truth_paths_from`` finds each name under ``galaxies.``. +TRUTH_PARAMETER_NAMES = { + "mass_to_light_ratio": "mass_to_light_ratio", + "kappa_s": "kappa_s", + "scale_radius_arcsec": "scale_radius", + "inner_slope": "inner_slope", +} + +_DARK_FREE_NAMES = ("kappa_s", "scale_radius", "inner_slope") + + +def _has_free_inner_slope(dark_cls) -> bool: + """Whether ``inner_slope`` is a constructor argument (``NFWSph`` fixes it to 1).""" + return "inner_slope" in inspect.signature(dark_cls.__init__).parameters + + +def _is_external_shear(profile) -> bool: + return isinstance(profile, ag.mp.ExternalShear) + + +def stellar_dark_record_from(twin, galaxy: Optional[str] = None): + """ + The last ``ReplaceTotalMassWithStellarPlusDark`` record of ``twin`` (for ``galaxy``, + a galaxy name or role, when given). + + Raises + ------ + ValueError + If the twin has no such record. + """ + name = None if galaxy is None else galaxy_name_from(twin, galaxy) + for record in reversed(twin.records): + if record.name != STELLAR_DARK_NAME: + continue + if name is None or record.truth.get("galaxy") == name: + return record + raise ValueError( + f"The twin has no {STELLAR_DARK_NAME} record" + + ("" if name is None else f" for galaxy {name!r}") + + f" (records: {[record.name for record in twin.records]}); apply " + f"al.twins.transform.ReplaceTotalMassWithStellarPlusDark first." + ) + + +def _prior_from(key: str, priors: Optional[Dict]): + if priors is not None and key in priors: + return priors[key] + return DEFAULT_PRIORS[key]() + + +def _stellar_model_from(light_basis, mass_to_light_ratio_prior) -> af.Model: + """ + A ``Basis`` of ``lmp.Gaussian`` models with every member's centre, ``ell_comps``, + ``sigma`` and intensity fixed to the twin's and one mass-to-light ratio shared by + all members (as ``mass_light_dark_basis_from`` ties it). + """ + model_list = [] + for profile in light_basis.profile_list: + model = af.Model(ag.lmp.Gaussian) + model.centre = tuple(float(value) for value in profile.centre) + model.ell_comps = tuple(float(value) for value in profile.ell_comps) + model.intensity = float(profile.intensity) + model.sigma = float(profile.sigma) + model.mass_to_light_ratio = mass_to_light_ratio_prior + model_list.append(model) + return af.Model(ag.lp_basis.Basis, profile_list=model_list) + + +def _dark_model_from(dark, truth, free_scale_radius, free_inner_slope, priors): + """ + The dark profile's model: centre fixed to the light centre, free ``kappa_s``, free + (or truth-fixed) ``scale_radius`` and, if the class has one, a free (or + truth-fixed) ``inner_slope``. Any other constructor argument is fixed to the twin's. + """ + dark_cls = type(dark) + model = af.Model(dark_cls) + + for name in inspect.signature(dark_cls.__init__).parameters: + if name in ("self", "centre") or name in _DARK_FREE_NAMES: + continue + if hasattr(dark, name): + setattr(model, name, getattr(dark, name)) + + model.centre = tuple(float(value) for value in truth["dark_centre"]) + model.kappa_s = _prior_from("kappa_s", priors) + model.scale_radius = ( + _prior_from("scale_radius_arcsec", priors) + if free_scale_radius + else float(truth["scale_radius_arcsec"]) + ) + if _has_free_inner_slope(dark_cls): + model.inner_slope = ( + _prior_from("inner_slope", priors) + if free_inner_slope + else float(truth["inner_slope"]) + ) + return model + + +def _field_model_from(field, free_shear): + """A ``MassField`` model whose ``ExternalShear`` components are free if asked.""" + components = { + key: ( + af.Model(ag.mp.ExternalShear) + if free_shear and _is_external_shear(profile) + else profile + ) + for key, profile in field.profile_dict.items() + } + return af.Model(ag.MassField, redshift=field.redshift, **components) + + +def _source_name_from(galaxies: Dict, lens_name: str) -> str: + """``"source"`` when the tracer has one, else the single highest-redshift galaxy.""" + if "source" in galaxies: + return "source" + redshifts = { + name: galaxy.redshift for name, galaxy in galaxies.items() if name != lens_name + } + if not redshifts: + raise ValueError("The twin's tracer has no galaxy besides the lens.") + highest = max(redshifts.values()) + names = [name for name, redshift in redshifts.items() if redshift == highest] + if len(names) > 1: + raise ValueError( + f"Several galaxies ({names}) share the highest redshift; pass source_model " + f"as a dictionary keyed by galaxy name." + ) + return names[0] + + +def stellar_dark_model_from( + twin, + galaxy: Optional[str] = None, + *, + source_model, + free_scale_radius: bool = True, + free_inner_slope: Optional[bool] = None, + free_shear: bool = True, + priors: Optional[Dict] = None, +) -> af.Collection: + """ + The model that fits a stellar-plus-dark twin with the parameterization it was made + with, as ``af.Collection(galaxies=af.Collection(=..., ...), fields=...)``. + + Everything is read from the last ``ReplaceTotalMassWithStellarPlusDark`` record of + the twin (its ``requested`` light and dark attribute names, and its ``truth``) and + from the twin's tracer, not by guessing attribute names: + + - the lens galaxy's light attribute (``requested["light"]``, e.g. ``bulge``) is a + ``Basis`` of ``af.Model(lmp.Gaussian)``, each member's centre, ``ell_comps``, + ``sigma`` and intensity fixed to the twin's, all sharing ONE free + ``mass_to_light_ratio`` (the same prior object, so the model has one parameter); + - the dark attribute (``requested["dark_name"]``) is ``af.Model`` of the twin's dark + class, centred on the light (``truth["dark_centre"]``, fixed), with free + ``kappa_s``, free ``scale_radius`` (fixed to the truth when + ``free_scale_radius=False``) and, for a class with an ``inner_slope`` (e.g. + ``gNFWSph``), a free ``inner_slope`` unless ``free_inner_slope=False``; + - an ``ExternalShear`` attribute of the lens galaxy, or an ``ExternalShear`` + component of a tracer ``MassField`` (the workspace convention, + ``fields=af.Model(al.MassField, redshift=..., shear=af.Model(al.mp.ExternalShear))``), + becomes a free ``af.Model(ExternalShear)`` when ``free_shear`` (else it is fixed to + the twin's). One field is passed as ``fields=``; several as + ``fields=af.Collection(field_0=..., ...)``. The field's component names are kept; + the tracer does not keep field names, so the single-field slot is ``fields``; + - any other attribute of the lens galaxy is fixed to the twin's; + - ``source_model`` replaces the source galaxy (the galaxy named ``"source"``, else + the single highest-redshift galaxy), or, as a dictionary, the galaxies it names; + - every other galaxy is kept as its fixed instance under its name. + + Default priors (``DEFAULT_PRIORS``, by ``truth`` key) are broad, fixed and never + derived from the truth: ``mass_to_light_ratio`` log-uniform over [1e-6, 1e6], + ``kappa_s`` log-uniform over [1e-4, 10], ``scale_radius_arcsec`` uniform over + [0.1, 30] arc-seconds and ``inner_slope`` uniform over [0, 2]. + + Parameters + ---------- + twin + A ``Twin`` with a ``ReplaceTotalMassWithStellarPlusDark`` record. + galaxy + The transformed galaxy (name or role); defaults to the galaxy of the last + stellar-dark record. + source_model + The source galaxy model (e.g. ``af.Model(al.Galaxy, redshift=1.0, bulge=...)``) + or a dictionary ``{galaxy name: model}``. + free_scale_radius + Fit the dark scale radius (default) or fix it to the truth. + free_inner_slope + Fit the inner slope of a dark class that has one. ``None`` (default) means free + when the class has an ``inner_slope``. ``True`` for a class without one raises. + free_shear + Fit the external shear (default) or fix it to the twin's. + priors + Overrides of ``DEFAULT_PRIORS``, by ``truth`` key, e.g. + ``{"kappa_s": af.LogUniformPrior(lower_limit=1e-3, upper_limit=1.0)}``. + + Raises + ------ + ValueError + If the twin has no ``ReplaceTotalMassWithStellarPlusDark`` record, or a + ``priors`` key is not a ``DEFAULT_PRIORS`` key. + """ + if priors is not None: + unknown = sorted(set(priors) - set(DEFAULT_PRIORS)) + if unknown: + raise ValueError( + f"Unknown priors keys {unknown}; the keys are {list(DEFAULT_PRIORS)}." + ) + + record = stellar_dark_record_from(twin, galaxy) + truth = record.truth + lens_name = truth["galaxy"] + light_name = record.requested["light"] + dark_name = record.requested["dark_name"] + + tracer = twin.tracer + galaxies = galaxy_dict_from(tracer) + lens = galaxies[lens_name] + + dark = getattr(lens, dark_name, None) + if free_inner_slope is None: + free_inner_slope = dark is not None and _has_free_inner_slope(type(dark)) + elif free_inner_slope and not _has_free_inner_slope(type(dark)): + raise ValueError( + f"free_inner_slope=True but the dark profile " + f"{type(dark).__name__} has no inner_slope." + ) + + lens_components = {} + for key, profile in lens.profile_dict.items(): + if key == light_name and all( + isinstance(member, ag.lmp.Gaussian) + for member in getattr(profile, "profile_list", []) + ): + lens_components[key] = _stellar_model_from( + light_basis=profile, + mass_to_light_ratio_prior=_prior_from("mass_to_light_ratio", priors), + ) + elif key == dark_name and isinstance(profile, MassProfile): + lens_components[key] = _dark_model_from( + dark=profile, + truth=truth, + free_scale_radius=free_scale_radius, + free_inner_slope=free_inner_slope, + priors=priors, + ) + elif free_shear and _is_external_shear(profile): + lens_components[key] = af.Model(ag.mp.ExternalShear) + else: + lens_components[key] = profile + + if isinstance(source_model, dict): + source_models = dict(source_model) + else: + source_models = {_source_name_from(galaxies, lens_name): source_model} + unknown = sorted(set(source_models) - set(galaxies)) + if unknown: + raise ValueError( + f"source_model names galaxies {unknown} that are not in the twin's tracer " + f"({list(galaxies)})." + ) + + galaxy_models = {} + for name, instance in galaxies.items(): + if name == lens_name: + galaxy_models[name] = af.Model( + ag.Galaxy, redshift=lens.redshift, **lens_components + ) + elif name in source_models: + galaxy_models[name] = source_models[name] + else: + galaxy_models[name] = instance + + model = af.Collection(galaxies=af.Collection(**galaxy_models)) + + fields = list(getattr(tracer, "fields", None) or []) + if len(fields) == 1: + model.fields = _field_model_from(fields[0], free_shear=free_shear) + elif len(fields) > 1: + model.fields = af.Collection( + **{ + f"field_{index}": _field_model_from(field, free_shear=free_shear) + for index, field in enumerate(fields) + } + ) + + return model + + +def truth_paths_from(model, galaxy: str) -> Dict[str, Tuple[str, ...]]: + """ + The model path of each recoverable ``truth`` key (``TRUTH_PARAMETER_NAMES``) that is + a free parameter of ``model`` under ``galaxies.``. + + A tied parameter (the shared ``mass_to_light_ratio``) has one path per member; the + first is returned, and any of them indexes the same sample column. A truth key whose + parameter is fixed (e.g. ``scale_radius_arcsec`` with ``free_scale_radius=False``, + or ``inner_slope`` of an ``NFWSph``) is left out. + + Raises + ------ + ValueError + If a parameter name is held by more than one free parameter of the galaxy. + """ + paths = {} + for key, parameter in TRUTH_PARAMETER_NAMES.items(): + matches = [ + prior_paths[0] + for prior_paths in model.all_paths + if prior_paths[0][:2] == ("galaxies", galaxy) + and prior_paths[0][-1] == parameter + ] + if len(matches) > 1: + raise ValueError( + f"Several free parameters of galaxies.{galaxy} are named {parameter!r} " + f"({matches}); pass the paths explicitly." + ) + if matches: + paths[key] = tuple(matches[0]) + return paths diff --git a/autolens/twins/recovery.py b/autolens/twins/recovery.py new file mode 100644 index 000000000..470dbce04 --- /dev/null +++ b/autolens/twins/recovery.py @@ -0,0 +1,215 @@ +""" +``recovery_table_from``: compare a fit's posterior with the ground truth a twin was +simulated with. + +For each recoverable ``truth`` key (``model.TRUTH_PARAMETER_NAMES``) that is a free +parameter of the fit, a ``RecoveryTable`` row lists the truth, the median of the 1D +marginalized PDF, its lower and upper values at ``sigma`` and whether the truth lies +between them. The truth keys are mapped onto model paths by +``model.truth_paths_from``, the same map the model builder publishes. + +NumPy only: nothing here supports JAX. +""" + +import csv +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +from autolens.twins.model import truth_paths_from +from autolens.twins.transform.stellar_dark import NAME as STELLAR_DARK_NAME + +CSV_COLUMNS = ("name", "path", "truth", "median", "lower", "upper", "inside") + + +@dataclass(frozen=True) +class RecoveryRow: + """One parameter: its truth, PDF median and the bounds at the table's sigma.""" + + name: str + path: Tuple[str, ...] + truth: float + median: float + lower: float + upper: float + + @property + def inside(self) -> bool: + """Whether the truth lies within ``[lower, upper]``.""" + return self.lower <= self.truth <= self.upper + + +class RecoveryTable: + def __init__(self, rows: List[RecoveryRow], sigma: float): + """ + The recovery of a twin's truth by a fit, one ``RecoveryRow`` per parameter. + + Parameters + ---------- + rows + The rows, in ``truth`` key order. + sigma + The sigma at which ``lower`` and ``upper`` are taken. + """ + self.rows = list(rows) + self.sigma = float(sigma) + + def __iter__(self): + return iter(self.rows) + + def __len__(self): + return len(self.rows) + + def __getitem__(self, name: str) -> RecoveryRow: + """The row of truth key ``name``.""" + for row in self.rows: + if row.name == name: + return row + raise KeyError(name) + + @property + def all_inside(self) -> bool: + """Whether every truth lies within its bounds.""" + return all(row.inside for row in self.rows) + + def to_csv(self, file_path: Union[str, Path]): + """Write the rows to ``file_path`` (columns ``CSV_COLUMNS``, path dot-joined).""" + with open(file_path, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow(CSV_COLUMNS) + for row in self.rows: + writer.writerow( + [ + row.name, + ".".join(row.path), + repr(row.truth), + repr(row.median), + repr(row.lower), + repr(row.upper), + row.inside, + ] + ) + + def __str__(self): + header = ( + f"{'name':<22}{'truth':>13}{'median':>13}{'lower':>13}{'upper':>13} " + f"inside ({self.sigma:g} sigma)" + ) + lines = [header] + for row in self.rows: + lines.append( + f"{row.name:<22}{row.truth:>13.5g}{row.median:>13.5g}" + f"{row.lower:>13.5g}{row.upper:>13.5g} {'yes' if row.inside else 'NO'}" + ) + return "\n".join(lines) + + def __repr__(self): + return ( + f"RecoveryTable(sigma={self.sigma:g}, names={[row.name for row in self.rows]}, " + f"all_inside={self.all_inside})" + ) + + +def truth_from(truth: Union[Dict, str, Path]) -> Dict: + """ + The flat truth dictionary: ``truth`` itself, or the ``"truth"`` block of a + ``twin.json`` file (or of the ``twin.json`` in a directory written by + ``SimulatedTwin.output_to``). That block lists one entry per transformation; the + last ``ReplaceTotalMassWithStellarPlusDark`` entry is returned (else the last). + """ + if isinstance(truth, dict): + return dict(truth) + + file_path = Path(truth) + if file_path.is_dir(): + file_path = file_path / "twin.json" + with open(file_path) as f: + entries = json.load(f)["truth"] + if not entries: + raise ValueError(f"{file_path} has an empty truth block.") + for entry in reversed(entries): + if entry.get("transformation") == STELLAR_DARK_NAME: + return dict(entry) + return dict(entries[-1]) + + +def _column_from(samples, path: Tuple[str, ...]) -> int: + for index, prior_paths in enumerate(samples.paths): + if tuple(path) in (tuple(prior_path) for prior_path in prior_paths): + return index + raise ValueError(f"{'.'.join(path)} is not a free parameter of the samples' model.") + + +def _values_at_sigma_from(samples, sigma: float) -> List[Tuple[float, float]]: + if hasattr(samples, "values_at_sigma"): + return samples.values_at_sigma(sigma=sigma, as_instance=False) + + # A SamplesSummary stores the values at 1 and 3 sigma only, keyed by path. + stored = {1.0: "values_at_sigma_1", 3.0: "values_at_sigma_3"}.get(float(sigma)) + values = None if stored is None else getattr(samples, stored, None) + if values is None: + raise ValueError( + f"A {type(samples).__name__} stores values at sigma 1 and 3 only; got " + f"sigma={sigma}. Pass the full samples instead." + ) + return [tuple(values[".".join(prior_paths[0])]) for prior_paths in samples.paths] + + +def recovery_table_from( + samples, + truth: Union[Dict, str, Path], + paths: Optional[Dict[str, Tuple[str, ...]]] = None, + sigma: float = 3.0, +) -> RecoveryTable: + """ + Compare the posterior in ``samples`` with ``truth``. + + The median is ``samples.median_pdf(as_instance=False)`` and the bounds are + ``samples.values_at_sigma(sigma, as_instance=False)`` (the 1D marginalized + quantiles; parameter covariance is ignored). Sample columns are found through + ``samples.paths``, which groups every path of a tied parameter, so any member's + path of the shared ``mass_to_light_ratio`` finds its column. A ``SamplesSummary`` + is accepted for ``sigma`` 1 or 3. + + Parameters + ---------- + samples + The ``Samples`` (e.g. ``result.samples``) or ``SamplesSummary`` of a fit of a + model made by ``stellar_dark_model_from``. + truth + A flat truth dictionary (``TransformationRecord.truth``) or the path of a + ``twin.json`` (or of its directory), whose ``"truth"`` block is read. + paths + ``{truth key: model path}``; defaults to + ``truth_paths_from(samples.model, truth["galaxy"])``. Only truth keys with a + path become rows. + sigma + The sigma of the bounds (default 3). + """ + truth = truth_from(truth) + if paths is None: + if samples.model is None: + raise ValueError("The samples have no model; pass paths explicitly.") + paths = truth_paths_from(samples.model, truth["galaxy"]) + + median = samples.median_pdf(as_instance=False) + bounds = _values_at_sigma_from(samples, sigma) + + rows = [] + for name, path in paths.items(): + if name not in truth: + continue + column = _column_from(samples, tuple(path)) + lower, upper = bounds[column] + rows.append( + RecoveryRow( + name=name, + path=tuple(path), + truth=float(truth[name]), + median=float(median[column]), + lower=float(lower), + upper=float(upper), + ) + ) + return RecoveryTable(rows=rows, sigma=sigma) diff --git a/autolens/twins/transform/stellar_dark.py b/autolens/twins/transform/stellar_dark.py index b50a16b28..f513b0318 100644 --- a/autolens/twins/transform/stellar_dark.py +++ b/autolens/twins/transform/stellar_dark.py @@ -297,11 +297,21 @@ def apply(self, twin): 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) - ] + # A tracer field is either a bare ExternalShear or a MassField holding named + # components (the DR1 pipeline layout, `fields.shear`). + shear_field_indexes = [] + shear_field_keys = {} + for index, field in enumerate(getattr(tracer, "fields", [])): + if _is_external_shear(field): + shear_field_indexes.append(index) + elif isinstance(field, ag.MassField): + keys = [ + key + for key, value in vars(field).items() + if _is_external_shear(value) + ] + if keys: + shear_field_keys[index] = keys mass_centre = tuple( float(value) for value in next(iter(replaced.values())).centre @@ -366,17 +376,31 @@ def apply(self, twin): 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 - ] + shear_paths = ( + [f"{path}.{key}" for key in shear_keys] + + [f"fields.{index}" for index in shear_field_indexes] + + [ + f"fields.{index}.{key}" + for index, keys in shear_field_keys.items() + for key in keys + ] + ) 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 - ] + fields = [] + for index, field in enumerate(new_tracer.fields): + if index in shear_field_indexes: + continue + for key in shear_field_keys.get(index, []): + delattr(field, key) + # A MassField left with no mass component is dropped with its shear. + if isinstance(field, ag.MassField) and not field.cls_list_from( + cls=MassProfile + ): + continue + fields.append(field) + new_tracer.fields = fields changed.extend(shear_paths) for shear_path in shear_paths: removed[shear_path] = get_class_path(ag.mp.ExternalShear) diff --git a/test_autolens/twins/test_model.py b/test_autolens/twins/test_model.py new file mode 100644 index 000000000..62bd9a487 --- /dev/null +++ b/test_autolens/twins/test_model.py @@ -0,0 +1,312 @@ +import numpy as np +import pytest + +import autofit as af +import autolens as al +from autolens.twins.model import DEFAULT_PRIORS, stellar_dark_record_from + +from test_autolens.twins.conftest import COARSE_APERTURE, GNFW_APERTURE + + +def _truth_vector_from(model, twin, record): + """ + The parameter vector of ``model`` at the twin's values: the record's truth for the + stellar-dark parameters and the twin's shear for the free shear components. + """ + truth = record.truth + tracer = twin.tracer + values = [] + for paths in model.all_paths: + path = paths[0] + name = path[-1] + if name == "mass_to_light_ratio": + values.append(truth["mass_to_light_ratio"]) + elif name == "kappa_s": + values.append(truth["kappa_s"]) + elif name == "scale_radius": + values.append(truth["scale_radius_arcsec"]) + elif name == "inner_slope": + values.append(truth["inner_slope"]) + elif name in ("gamma_1", "gamma_2"): + if path[0] == "fields": + shear = getattr(tracer.fields[0], path[1]) + else: + shear = getattr(getattr(tracer.galaxies, path[1]), path[2]) + values.append(getattr(shear, name)) + else: + raise AssertionError(f"unexpected free parameter {path}") + return values + + +def _grid(): + return al.Grid2D.uniform(shape_native=(8, 8), pixel_scales=0.3, over_sample_size=1) + # An even shape keeps (0, 0), where a gNFW convergence diverges, off the grid. + + +def _assert_lens_reproduced(model, twin, record): + instance = model.instance_from_vector( + vector=_truth_vector_from(model, twin, record) + ) + name = record.truth["galaxy"] + grid = _grid() + + twin_lens = getattr(twin.tracer.galaxies, name) + model_lens = getattr(instance.galaxies, name) + + convergence_twin = twin_lens.convergence_2d_from(grid=grid).array + convergence_model = model_lens.convergence_2d_from(grid=grid).array + assert np.max(np.abs(convergence_model - convergence_twin)) < 1.0e-10 + + deflections_twin = twin_lens.deflections_yx_2d_from(grid=grid).array + deflections_model = model_lens.deflections_yx_2d_from(grid=grid).array + assert np.max(np.abs(deflections_model - deflections_twin)) < 1.0e-10 + + return instance + + +def _reference_from(tracer, masked_imaging_7x7): + return al.twins.ReferenceModel.from_result( + result=al.m.MockResult(max_log_likelihood_tracer=tracer), + dataset=masked_imaging_7x7, + lens_id="model_test", + ) + + +def _field_shear_tracer(): + """The witness lens with its external shear held in a tracer MassField.""" + from test_autolens.twins.conftest import make_witness_tracer + + witness = make_witness_tracer() + lens = witness.galaxies.lens + lens = al.Galaxy(redshift=lens.redshift, bulge=lens.bulge, mass=lens.mass) + field = al.MassField( + redshift=0.5, shear=al.mp.ExternalShear(gamma_1=0.03, gamma_2=-0.02) + ) + return al.Tracer( + galaxies=af.ModelInstance({"lens": lens, "source": witness.galaxies.source}), + fields=[field], + ) + + +def test__stellar_dark_model__layout_and_prior_count(witness_twin): + source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp.SersicSph) + model = al.twins.stellar_dark_model_from(witness_twin, source_model=source) + + # M/L (one, shared) + kappa_s + scale_radius + shear (2) + source Sersic (5). + assert model.prior_count == 1 + 2 + 2 + 5 + + lens = model.galaxies.lens + members = lens.bulge.profile_list + assert len(members) == 3 + assert all(member.cls is al.lmp.Gaussian for member in members) + assert len({id(member.mass_to_light_ratio) for member in members}) == 1 + assert isinstance(members[0].mass_to_light_ratio, af.LogUniformPrior) + + twin_bulge = witness_twin.tracer.galaxies.lens.bulge + for member, profile in zip(members, twin_bulge.profile_list): + assert member.sigma == profile.sigma + assert member.intensity == profile.intensity + assert tuple(member.centre) == tuple(profile.centre) + + assert lens.dark.cls is al.mp.NFWSph + assert tuple(lens.dark.centre) == (0.0, 0.0) + assert lens.shear.cls is al.mp.ExternalShear + assert not hasattr(lens, "mass") + assert model.galaxies.source is source + + # Prior medians build an instance and a tracer. + instance = model.instance_from_prior_medians() + al.Tracer(galaxies=list(instance.galaxies)) + + +def test__stellar_dark_model__instance_at_truth_reproduces_the_twin(witness_twin): + source = witness_twin.tracer.galaxies.source + model = al.twins.stellar_dark_model_from(witness_twin, "lens", source_model=source) + record = stellar_dark_record_from(witness_twin) + + assert model.prior_count == 1 + 2 + 2 + + instance = _assert_lens_reproduced(model, witness_twin, record) + + grid = _grid() + tracer = al.Tracer(galaxies=list(instance.galaxies)) + assert ( + np.max( + np.abs( + tracer.deflections_yx_2d_from(grid=grid).array + - witness_twin.tracer.deflections_yx_2d_from(grid=grid).array + ) + ) + < 1.0e-10 + ) + + +def test__stellar_dark_model__fixed_scale_radius_shear_and_prior_overrides( + witness_twin, +): + source = witness_twin.tracer.galaxies.source + kappa_s_prior = af.LogUniformPrior(lower_limit=1.0e-3, upper_limit=1.0) + model = al.twins.stellar_dark_model_from( + witness_twin, + source_model=source, + free_scale_radius=False, + free_shear=False, + priors={"kappa_s": kappa_s_prior}, + ) + record = stellar_dark_record_from(witness_twin) + + assert model.prior_count == 2 + assert model.galaxies.lens.dark.kappa_s is kappa_s_prior + assert model.galaxies.lens.dark.scale_radius == record.truth["scale_radius_arcsec"] + assert isinstance(model.galaxies.lens.shear, al.mp.ExternalShear) + + with pytest.raises(ValueError, match="Unknown priors keys"): + al.twins.stellar_dark_model_from( + witness_twin, source_model=source, priors={"kapa_s": kappa_s_prior} + ) + + +def test__stellar_dark_model__default_priors_do_not_know_the_truth(witness_twin): + source = witness_twin.tracer.galaxies.source + model = al.twins.stellar_dark_model_from(witness_twin, source_model=source) + dark = model.galaxies.lens.dark + + for prior, key in ( + ( + model.galaxies.lens.bulge.profile_list[0].mass_to_light_ratio, + "mass_to_light_ratio", + ), + (dark.kappa_s, "kappa_s"), + (dark.scale_radius, "scale_radius_arcsec"), + ): + default = DEFAULT_PRIORS[key]() + assert type(prior) is type(default) + assert prior.lower_limit == default.lower_limit + assert prior.upper_limit == default.upper_limit + + +def test__stellar_dark_model__gnfw_free_inner_slope(witness_baseline): + twin = witness_baseline.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", + light="bulge", + scale_radius_kpc=10.0, + f_dark=0.5, + dark_cls=al.mp.gNFWSph, + inner_slope=1.3, + aperture_settings=GNFW_APERTURE, + ) + ) + source = twin.tracer.galaxies.source + record = stellar_dark_record_from(twin) + + model = al.twins.stellar_dark_model_from(twin, source_model=source) + assert model.galaxies.lens.dark.cls is al.mp.gNFWSph + assert isinstance(model.galaxies.lens.dark.inner_slope, af.UniformPrior) + assert model.prior_count == 1 + 3 + 2 + assert al.twins.truth_paths_from(model, "lens")["inner_slope"] == ( + "galaxies", + "lens", + "dark", + "inner_slope", + ) + + _assert_lens_reproduced(model, twin, record) + + fixed = al.twins.stellar_dark_model_from( + twin, source_model=source, free_inner_slope=False + ) + assert fixed.galaxies.lens.dark.inner_slope == 1.3 + assert fixed.prior_count == 1 + 2 + 2 + + +def test__stellar_dark_model__nfw_rejects_free_inner_slope(witness_twin): + with pytest.raises(ValueError, match="no inner_slope"): + al.twins.stellar_dark_model_from( + witness_twin, + source_model=witness_twin.tracer.galaxies.source, + free_inner_slope=True, + ) + + +def test__stellar_dark_model__field_shear(masked_imaging_7x7): + twin = al.twins.Twin.baseline( + _reference_from(_field_shear_tracer(), masked_imaging_7x7) + ).apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", + light="bulge", + scale_radius_kpc=10.0, + aperture_settings=COARSE_APERTURE, + ) + ) + source = twin.tracer.galaxies.source + model = al.twins.stellar_dark_model_from(twin, source_model=source) + + # The shear is a free ExternalShear of a MassField under `fields`, as in the + # workspace's `fields=af.Model(al.MassField, redshift=..., shear=...)`. + assert model.fields.cls is al.MassField + assert model.fields.shear.cls is al.mp.ExternalShear + assert ("fields", "shear", "gamma_1") in [paths[0] for paths in model.all_paths] + assert model.prior_count == 1 + 2 + 2 + + record = stellar_dark_record_from(twin) + vector = _truth_vector_from(model, twin, record) + instance = model.instance_from_vector(vector=vector) + + # The tracer AnalysisImaging builds from the instance matches the twin's. + analysis = al.AnalysisImaging(dataset=masked_imaging_7x7, use_jax=False) + tracer = analysis.tracer_via_instance_from(instance=instance) + grid = _grid() + assert ( + np.max( + np.abs( + tracer.deflections_yx_2d_from(grid=grid).array + - twin.tracer.deflections_yx_2d_from(grid=grid).array + ) + ) + < 1.0e-10 + ) + + fixed = al.twins.stellar_dark_model_from( + twin, source_model=source, free_shear=False + ) + assert isinstance(fixed.fields.shear, al.mp.ExternalShear) + assert fixed.prior_count == 1 + 2 + + +def test__stellar_dark_model__source_model_dict_and_errors(witness_twin): + source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp.SersicSph) + model = al.twins.stellar_dark_model_from( + witness_twin, source_model={"source": source} + ) + assert model.galaxies.source is source + + with pytest.raises(ValueError, match="not in the twin's tracer"): + al.twins.stellar_dark_model_from(witness_twin, source_model={"src": source}) + + +def test__stellar_dark_model__no_record_raises(witness_baseline): + with pytest.raises(ValueError, match="no ReplaceTotalMassWithStellarPlusDark"): + al.twins.stellar_dark_model_from( + witness_baseline, + source_model=witness_baseline.tracer.galaxies.source, + ) + + +def test__truth_paths_from(witness_twin): + model = al.twins.stellar_dark_model_from( + witness_twin, source_model=witness_twin.tracer.galaxies.source + ) + assert al.twins.truth_paths_from(model, "lens") == { + "mass_to_light_ratio": ( + "galaxies", + "lens", + "bulge", + "profile_list", + "0", + "mass_to_light_ratio", + ), + "kappa_s": ("galaxies", "lens", "dark", "kappa_s"), + "scale_radius_arcsec": ("galaxies", "lens", "dark", "scale_radius"), + } diff --git a/test_autolens/twins/test_recovery.py b/test_autolens/twins/test_recovery.py new file mode 100644 index 000000000..21880720d --- /dev/null +++ b/test_autolens/twins/test_recovery.py @@ -0,0 +1,210 @@ +import csv +import json + +import numpy as np +import pytest + +import autofit as af +import autolens as al +from autolens.twins.model import stellar_dark_record_from + + +def _samples_from(model, centres, widths, total=4000, seed=1): + """ + Synthetic equal-weight samples: each parameter Gaussian about ``centres[name]`` + with standard deviation ``widths[name]``, by the last element of its path. + """ + rng = np.random.default_rng(seed) + columns = [ + rng.normal(centres[paths[0][-1]], widths[paths[0][-1]], size=total) + for paths in model.all_paths + ] + parameter_lists = np.stack(columns, axis=1).tolist() + sample_list = af.Sample.from_lists( + model=model, + parameter_lists=parameter_lists, + log_likelihood_list=[0.0] * total, + log_prior_list=[0.0] * total, + weight_list=[1.0 / total] * total, + ) + return af.SamplesPDF(model=model, sample_list=sample_list) + + +@pytest.fixture(name="model_and_truth") +def make_model_and_truth(witness_twin): + model = al.twins.stellar_dark_model_from( + witness_twin, source_model=witness_twin.tracer.galaxies.source + ) + return model, stellar_dark_record_from(witness_twin).truth + + +def _centres(truth, offset=0.0): + return { + "mass_to_light_ratio": truth["mass_to_light_ratio"] * (1.0 + offset), + "kappa_s": truth["kappa_s"] * (1.0 + offset), + "scale_radius": truth["scale_radius_arcsec"] * (1.0 + offset), + "gamma_1": 0.03, + "gamma_2": -0.02, + } + + +def _widths(truth): + return { + "mass_to_light_ratio": 0.01 * truth["mass_to_light_ratio"], + "kappa_s": 0.01 * truth["kappa_s"], + "scale_radius": 0.01 * truth["scale_radius_arcsec"], + "gamma_1": 0.001, + "gamma_2": 0.001, + } + + +def test__recovery_table__truth_inside(model_and_truth, tmp_path): + model, truth = model_and_truth + samples = _samples_from(model, _centres(truth), _widths(truth)) + + table = al.twins.recovery_table_from(samples, truth) + + assert [row.name for row in table] == [ + "mass_to_light_ratio", + "kappa_s", + "scale_radius_arcsec", + ] + assert table.all_inside + row = table["kappa_s"] + assert row.path == ("galaxies", "lens", "dark", "kappa_s") + assert row.truth == truth["kappa_s"] + assert row.median == pytest.approx(truth["kappa_s"], rel=2.0e-3) + assert row.lower < row.median < row.upper + # 3 sigma of a 1% Gaussian. + assert row.upper - row.lower == pytest.approx(0.06 * truth["kappa_s"], rel=0.2) + + assert "kappa_s" in str(table) + assert "yes" in str(table) + + table.to_csv(tmp_path / "recovery.csv") + with open(tmp_path / "recovery.csv") as f: + rows = list(csv.DictReader(f)) + assert [row["name"] for row in rows] == [row.name for row in table] + assert rows[1]["path"] == "galaxies.lens.dark.kappa_s" + assert float(rows[1]["truth"]) == truth["kappa_s"] + assert rows[1]["inside"] == "True" + + +def test__recovery_table__truth_outside(model_and_truth): + model, truth = model_and_truth + samples = _samples_from(model, _centres(truth, offset=0.2), _widths(truth)) + + table = al.twins.recovery_table_from(samples, truth, sigma=3.0) + + assert not table.all_inside + assert not any(row.inside for row in table) + assert "NO" in str(table) + + # The same offset is inside a posterior ten times wider. + widths = {name: 10.0 * width for name, width in _widths(truth).items()} + wide = _samples_from(model, _centres(truth, offset=0.2), widths) + assert al.twins.recovery_table_from(wide, truth, sigma=3.0).all_inside + + +def test__recovery_table__shared_mass_to_light_ratio_any_member_path( + model_and_truth, +): + model, truth = model_and_truth + samples = _samples_from(model, _centres(truth), _widths(truth)) + + member_path = ( + "galaxies", + "lens", + "bulge", + "profile_list", + "2", + "mass_to_light_ratio", + ) + table = al.twins.recovery_table_from( + samples, truth, paths={"mass_to_light_ratio": member_path} + ) + assert len(table) == 1 + assert table["mass_to_light_ratio"].median == pytest.approx( + truth["mass_to_light_ratio"], rel=2.0e-3 + ) + + with pytest.raises(ValueError, match="not a free parameter"): + al.twins.recovery_table_from( + samples, truth, paths={"kappa_s": ("galaxies", "lens", "kappa_s")} + ) + + +def test__recovery_table__truth_from_twin_json_and_summary(model_and_truth, tmp_path): + model, truth = model_and_truth + samples = _samples_from(model, _centres(truth), _widths(truth)) + + twin_json = { + "truth": [ + {"transformation": "Identity", "record_index": 0, "galaxy": "other"}, + { + "transformation": "ReplaceTotalMassWithStellarPlusDark", + "record_index": 1, + **truth, + }, + ] + } + with open(tmp_path / "twin.json", "w") as f: + json.dump(twin_json, f) + + from_dict = al.twins.recovery_table_from(samples, truth) + from_file = al.twins.recovery_table_from(samples, tmp_path / "twin.json") + from_directory = al.twins.recovery_table_from(samples, str(tmp_path)) + for table in (from_file, from_directory): + assert [(row.name, row.truth, row.median) for row in table] == [ + (row.name, row.truth, row.median) for row in from_dict + ] + + summary = samples.summary() + from_summary = al.twins.recovery_table_from(summary, truth, sigma=3.0) + for row, expected in zip(from_summary, from_dict): + assert row.median == pytest.approx(expected.median, rel=1.0e-12) + assert row.lower == pytest.approx(expected.lower, rel=1.0e-12) + assert row.upper == pytest.approx(expected.upper, rel=1.0e-12) + + with pytest.raises(ValueError, match="sigma 1 and 3"): + al.twins.recovery_table_from(summary, truth, sigma=2.0) + + +def test__end_to_end__simulate_model_nautilus_recovery( + witness_twin, monkeypatch, tmp_path, plot_patch +): + """ + simulate -> twin -> model -> Nautilus -> recovery table, with PYAUTO_TEST_MODE=1 + (the real Nautilus backend on a reduced budget), on the 7x7 witness twin. The + budget is far too small to converge, so only the plumbing is asserted. Figures + are not rendered (``plot_patch``), which is most of the run time otherwise. + """ + monkeypatch.setenv("PYAUTO_TEST_MODE", "1") + + simulated = witness_twin.simulate(seed=1) + simulated.output_to(tmp_path / "twin") + + model = al.twins.stellar_dark_model_from( + simulated.twin, source_model=simulated.twin.tracer.galaxies.source + ) + analysis = al.AnalysisImaging(dataset=simulated.dataset, use_jax=False) + search = af.Nautilus( + name="twin_end_to_end", path_prefix=str(tmp_path / "output"), n_live=50 + ) + result = search.fit(model=model, analysis=analysis) + + table = al.twins.recovery_table_from(result.samples, tmp_path / "twin") + + assert [row.name for row in table] == [ + "mass_to_light_ratio", + "kappa_s", + "scale_radius_arcsec", + ] + assert ( + table["kappa_s"].truth + == stellar_dark_record_from(witness_twin).truth["kappa_s"] + ) + for row in table: + assert np.isfinite([row.median, row.lower, row.upper]).all() + assert row.lower <= row.upper + table.to_csv(tmp_path / "recovery.csv") diff --git a/test_autolens/twins/test_stellar_dark.py b/test_autolens/twins/test_stellar_dark.py index 0d3fd581c..a36b381e4 100644 --- a/test_autolens/twins/test_stellar_dark.py +++ b/test_autolens/twins/test_stellar_dark.py @@ -1,8 +1,10 @@ +import copy import json import numpy as np import pytest +import autofit as af import autogalaxy as ag import autolens as al from autolens.twins._galaxies import galaxy_dict_from @@ -390,3 +392,67 @@ def test__invalid_light_raises(witness_baseline): galaxy="galaxy_0", light="mass", scale_radius_kpc=10.0 ) ) + + +def _with_field_shear(twin, extra=None): + galaxies = galaxy_dict_from(twin.tracer) + lens = copy.deepcopy(galaxies["lens"]) + shear = lens.shear + delattr(lens, "shear") + field = al.MassField(redshift=lens.redshift, shear=shear, **(extra or {})) + tracer = al.Tracer( + galaxies=af.ModelInstance( + {"lens": lens, "source": copy.deepcopy(galaxies["source"])} + ), + fields=[field], + ) + return al.twins.Twin(reference=twin.reference, tracer=tracer) + + +def test__keep_shear_massfield_layout(witness_baseline): + twin = _with_field_shear(witness_baseline) + + kept = twin.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", + light="bulge", + scale_radius_kpc=10.0, + aperture_settings=COARSE_APERTURE, + ) + ) + assert isinstance(kept.tracer.fields[0].shear, al.mp.ExternalShear) + assert any("fields.0.shear" in effect for effect in kept.records[-1].side_effects) + + removed = twin.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", + light="bulge", + scale_radius_kpc=10.0, + keep_shear=False, + aperture_settings=COARSE_APERTURE, + ) + ) + record = removed.records[-1] + assert removed.tracer.fields == [] + assert "fields.0.shear" in record.changed + assert record.procedure["removed"]["fields.0.shear"].endswith("ExternalShear") + assert isinstance(twin.tracer.fields[0].shear, al.mp.ExternalShear) + + +def test__keep_shear_false_keeps_other_field_components(witness_baseline): + twin = _with_field_shear( + witness_baseline, extra={"sheet": al.mp.MassSheet(kappa=0.01)} + ) + + removed = twin.apply( + al.twins.transform.ReplaceTotalMassWithStellarPlusDark( + galaxy="lens", + light="bulge", + scale_radius_kpc=10.0, + keep_shear=False, + aperture_settings=COARSE_APERTURE, + ) + ) + field = removed.tracer.fields[0] + assert not hasattr(field, "shear") + assert isinstance(field.sheet, al.mp.MassSheet)