From 5e434136ad4012c7b0ef0b763f6599bb57aa505b Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 17 Sep 2026 20:15:58 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20MassField=20in=20PyAutoLens=20=E2=80=94?= =?UTF-8?q?=20Tracer(fields=3D),=20analysis=20slot,=20COOLEST,=20LOS,=20mo?= =?UTF-8?q?del=5Futil?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 2 of the mass-field epic. `Tracer(galaxies, cosmology=None, fields=None)` stores external fields in `tracer.fields`; only the planes merge galaxies and fields (`members` / `members_ascending_redshift`), so `tracer.galaxies` never holds a field and per-galaxy surfaces (the galaxy image dictionary, fit plane images, subplots) skip fields explicitly while every lensing quantity runs over members. The analysis folds a model's `fields` collection (AnalysisLens, the aggregator tracer builder, Result.max_log_likelihood_fields); COOLEST exports fields as their own MassField entities and imports MassField entities into `tracer.fields` (the one import behaviour change; export of legacy galaxy-attached sheets is pinned byte-identical); the LOS sampler gains fields_from()/galaxies_and_fields_from() with its legacy return kept as the default; model_util.mass_field_from ties the MassSheet / ExternalPotential centre to the lens mass centre as a shared prior. al.MassField is re-exported; API docs and the overview updated. Legacy galaxy-attached sheets: no warning, no behaviour change, identifier pinned (fef2697b5c32ba56bb18a7baecb7b0f6). No prior config changes. Closes #741. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01PqfNuorRbom8ZjD9i9hvWo --- autolens/__init__.py | 1 + autolens/aggregator/tracer.py | 10 +- autolens/analysis/analysis/lens.py | 10 + autolens/analysis/model_util.py | 92 +++ autolens/analysis/result.py | 21 +- autolens/imaging/fit_imaging.py | 5 + autolens/interferometer/fit_interferometer.py | 5 + autolens/interop/coolest.py | 88 ++- autolens/jax/registration.py | 19 +- autolens/lens/los.py | 104 +++- autolens/lens/tracer.py | 200 ++++++- docs/api/galaxy.rst | 8 + docs/api/modeling.rst | 21 +- docs/overview/overview_1_start_here.md | 20 + .../analysis/analysis/test_analysis_lens.py | 79 +++ test_autolens/analysis/test_model_util.py | 117 ++++ .../fixtures/legacy_galaxy_attached.json | 459 +++++++++++++++ test_autolens/interop/test_coolest.py | 233 +++++++- test_autolens/lens/test_los.py | 114 ++++ test_autolens/lens/test_tracer_fields.py | 532 ++++++++++++++++++ 20 files changed, 2089 insertions(+), 49 deletions(-) create mode 100644 test_autolens/analysis/test_model_util.py create mode 100644 test_autolens/interop/fixtures/legacy_galaxy_attached.json create mode 100644 test_autolens/lens/test_tracer_fields.py diff --git a/autolens/__init__.py b/autolens/__init__.py index 66b6c8b6a..5ffbd1898 100644 --- a/autolens/__init__.py +++ b/autolens/__init__.py @@ -60,6 +60,7 @@ from autogalaxy.gui.scribbler import Scribbler from autogalaxy.galaxy.galaxy import Galaxy from autogalaxy.galaxy.galaxies import Galaxies +from autogalaxy.galaxy.mass_field import MassField from autogalaxy.galaxy.galaxy_table import GalaxyTable from autogalaxy.galaxy.galaxy_table import galaxy_table_from_csv from autogalaxy.galaxy.galaxy_table import galaxy_table_to_csv diff --git a/autolens/aggregator/tracer.py b/autolens/aggregator/tracer.py index 7a246f31d..404e19264 100644 --- a/autolens/aggregator/tracer.py +++ b/autolens/aggregator/tracer.py @@ -98,7 +98,15 @@ def _tracer_from( if instance.scaling_galaxies is not None: galaxy_list += list(instance.scaling_galaxies) - tracer = Tracer(galaxies=galaxy_list, cosmology=cosmology) + # Mirrors the `fields` fold in `AnalysisLens.tracer_via_instance_from`: a model may declare an + # external field alongside its galaxies, and a `MassField` is not a galaxy. + fields = ( + list(instance.fields) + if getattr(instance, "fields", None) is not None + else None + ) + + tracer = Tracer(galaxies=galaxy_list, cosmology=cosmology, fields=fields) tracer_list.append(tracer) diff --git a/autolens/analysis/analysis/lens.py b/autolens/analysis/analysis/lens.py index 70e54571a..03fc9e6a1 100644 --- a/autolens/analysis/analysis/lens.py +++ b/autolens/analysis/analysis/lens.py @@ -132,9 +132,19 @@ def tracer_via_instance_from( if getattr(instance, "scaling_galaxies", None) is not None: galaxy_list += list(instance.scaling_galaxies) + # A model may declare an external field alongside its galaxies, as + # `af.Collection(galaxies=..., fields=af.Collection(field=af.Model(al.MassField, ...)))`. + # It is folded here rather than into `galaxy_list`, because a `MassField` is not a galaxy. + fields = ( + list(instance.fields) + if getattr(instance, "fields", None) is not None + else None + ) + return Tracer( galaxies=galaxy_list, cosmology=cosmology, + fields=fields, ) def log_likelihood_penalty_from( diff --git a/autolens/analysis/model_util.py b/autolens/analysis/model_util.py index aeca624f2..ef20d53d7 100644 --- a/autolens/analysis/model_util.py +++ b/autolens/analysis/model_util.py @@ -15,11 +15,15 @@ PyAutoLens-specific: - ``random_galaxies_for_simulation_from`` — sample concrete (lens, source) ``Galaxy`` instances for synthetic-data generation in ``start_here`` scripts. +- ``mass_field_from`` — build the ``af.Model(al.MassField, ...)`` external-field model that + goes in a model's ``fields`` collection. """ from typing import Optional, Tuple import numpy as np +import autofit as af + import autolens as al from autogalaxy.analysis.model_util import mge_model_from @@ -119,3 +123,91 @@ def random_galaxies_for_simulation_from( source = al.Galaxy(redshift=1.0, bulge=source_bulge) return lens, source + + +def mass_field_from( + lens: af.Model, + shear: bool = True, + mass_sheet: bool = False, + potential: bool = False, + redshift: Optional[float] = None, +) -> af.Model: + """ + Build the ``af.Model(al.MassField, ...)`` describing the tidal field of everything + outside the modelled system, with its geometry tied to a lens galaxy model. + + The external field is composed of the components no galaxy owns — an ``ExternalShear``, + a ``MassSheet``, an ``ExternalPotential`` — and is placed in a model's ``fields`` + collection rather than bolted onto a galaxy:: + + lens = af.Model(al.Galaxy, redshift=0.5, mass=af.Model(al.mp.Isothermal)) + + model = af.Collection( + galaxies=af.Collection(lens=lens, source=source), + fields=af.Collection(field=al.model_util.mass_field_from(lens=lens)), + ) + + The ``MassSheet`` and ``ExternalPotential`` have a meaningful centre (unlike an + ``ExternalShear``, whose deflection field is constant), and that centre is physically + the centre of the system the field is expanded about. It is therefore **shared** with + the lens galaxy's mass centre — the same prior object, not a copy — so the two move + together in the non-linear search and cost no extra dimensions. + + In a multi-deflector system, ``lens`` is the *primary* deflector: the galaxy whose + centre the external expansion is about. Pass that galaxy's model, not a secondary one. + + The galaxy-attached form (``af.Model(al.Galaxy, ..., shear=af.Model(al.mp.ExternalShear))``) + remains fully supported and is not deprecated; this is a second way to write the same + physics, one which keeps the external field out of the galaxy. + + Parameters + ---------- + lens + The ``af.Model`` of the (primary) lens galaxy. Its ``mass.centre`` is shared with the + field's centred components and, unless ``redshift`` is given, its redshift is the + field's redshift. + shear + If True (default), the field includes an ``ExternalShear``. + mass_sheet + If True, the field includes a ``MassSheet``, centred on the lens mass. + potential + If True, the field includes an ``ExternalPotential``, centred on the lens mass. + redshift + The redshift of the field. Defaults to the lens galaxy model's redshift. + + Returns + ------- + The ``af.Model`` of the ``MassField``, to be placed in the model's ``fields`` collection. + """ + if not hasattr(lens, "mass"): + raise ValueError( + f"The `lens` model passed to `mass_field_from` has no `mass` attribute: " + f"{lens!r}. The external field's centred components are tied to the lens " + f"galaxy's mass centre, so the lens model must have a `mass`, e.g. " + f"af.Model(al.Galaxy, redshift=0.5, mass=af.Model(al.mp.Isothermal))." + ) + + redshift = lens.redshift if redshift is None else redshift + + components = {} + + if shear: + components["shear"] = af.Model(al.mp.ExternalShear) + + if mass_sheet: + components["mass_sheet"] = af.Model(al.mp.MassSheet) + + if potential: + components["potential"] = af.Model(al.mp.ExternalPotential) + + field = af.Model(al.MassField, redshift=redshift, **components) + + # The *same* prior object, not a copy: `field.potential.centre is lens.mass.centre` must + # hold, so the external expansion and the lens mass share one centre in the search. + if mass_sheet: + field.mass_sheet.centre = lens.mass.centre + + if potential: + field.potential.centre = lens.mass.centre + + return field diff --git a/autolens/analysis/result.py b/autolens/analysis/result.py index 6f3c81810..13d345e09 100644 --- a/autolens/analysis/result.py +++ b/autolens/analysis/result.py @@ -45,6 +45,22 @@ def max_log_likelihood_tracer(self) -> Tracer: """ return self.analysis.tracer_via_instance_from(instance=self.instance) + @property + def max_log_likelihood_fields(self) -> List[ag.MassField]: + """ + An instance of the list of `MassField` objects corresponding to the maximum log likelihood model inferred by + the non-linear search. + + Mirrors `max_log_likelihood_galaxies` (which reads the galaxies off the same instance), and returns an empty + list for the overwhelming majority of models, which declare no `fields` collection at all. + """ + fields = getattr(self.instance, "fields", None) + + if fields is None: + return [] + + return list(fields) + def source_plane_light_profile_centre_from( self, plane_redshift: Optional[float] = None ) -> aa.Grid2DIrregular: @@ -251,7 +267,10 @@ def positions_threshold_from( by `factor` and rounded up to the `threshold`. """ - tracer = Tracer(galaxies=self.max_log_likelihood_galaxies) + tracer = Tracer( + galaxies=self.max_log_likelihood_galaxies, + fields=self.max_log_likelihood_fields, + ) positions = ( self.image_plane_multiple_image_positions(plane_redshift=plane_redshift) diff --git a/autolens/imaging/fit_imaging.py b/autolens/imaging/fit_imaging.py index 8da5ebfd1..88f79847f 100644 --- a/autolens/imaging/fit_imaging.py +++ b/autolens/imaging/fit_imaging.py @@ -321,6 +321,11 @@ def model_images_of_planes_list(self) -> List[aa.Array2D]: for plane_index, galaxies in enumerate(self.tracer.planes): for galaxy in galaxies: + # A plane may hold a `MassField`, which is not a key of the galaxy dictionary (it has no + # light, so `Tracer.galaxy_image_2d_dict_from` skips it); indexing it here would KeyError. + if isinstance(galaxy, ag.MassField): + continue + model_images_of_planes_list[plane_index] += galaxy_model_image_dict[ galaxy ] diff --git a/autolens/interferometer/fit_interferometer.py b/autolens/interferometer/fit_interferometer.py index dbe659354..99060015e 100644 --- a/autolens/interferometer/fit_interferometer.py +++ b/autolens/interferometer/fit_interferometer.py @@ -292,6 +292,11 @@ def model_visibilities_of_planes_list(self) -> List[aa.Visibilities]: for plane_index, galaxies in enumerate(self.tracer.planes): for galaxy in galaxies: + # A plane may hold a `MassField`, which is not a key of the galaxy dictionary (it has no + # light, so `Tracer.galaxy_image_2d_dict_from` skips it); indexing it here would KeyError. + if isinstance(galaxy, ag.MassField): + continue + model_visibilities_of_planes_list[ plane_index ] += galaxy_model_visibilities_dict[galaxy] diff --git a/autolens/interop/coolest.py b/autolens/interop/coolest.py index 2c634e85d..f7e1acd31 100644 --- a/autolens/interop/coolest.py +++ b/autolens/interop/coolest.py @@ -186,14 +186,19 @@ def to_coolest( Every galaxy becomes a COOLEST ``Galaxy`` lensing entity with its light and mass profiles converted to COOLEST conventions. ``ExternalShear`` and - ``MassSheet`` profiles are exported as COOLEST ``MassField`` entities (the - standard treats external fields separately from galaxies). + ``MassSheet`` profiles attached to a galaxy are peeled off and exported as + COOLEST ``MassField`` entities named ``mass_field_{i}`` (the standard treats + external fields separately from galaxies). + + A ``Tracer``'s own ``fields`` -- its ``ag.MassField`` objects -- are exported + as COOLEST ``MassField`` entities too, named ``field_{j}``, carrying every + mass profile they hold (not only shears and sheets). Parameters ---------- galaxies The galaxies of the lens model (each with a ``redshift``), or a - ``Tracer`` whose galaxies (and cosmology) are used. + ``Tracer`` whose galaxies, fields and cosmology are used. file_path The output path of the template; a ``.json`` extension is appended by the COOLEST serializer if not present. @@ -241,16 +246,21 @@ def to_coolest( if isinstance(galaxies, Tracer): cosmology = cosmology or galaxies.cosmology + # Read `fields` before rebinding the name: a list of galaxies has no fields. + fields = list(galaxies.fields) galaxies = list(galaxies.galaxies) else: galaxies = list(galaxies) + fields = [] cosmology = cosmology or ag.cosmo.Planck15() if dataset is not None: pixel_size = float(dataset.pixel_scales[0]) shape_native = tuple(dataset.shape_native) - redshift_max = max(galaxy.redshift for galaxy in galaxies) + redshift_max = max( + [galaxy.redshift for galaxy in galaxies] + [field.redshift for field in fields] + ) metadata = dict(metadata or {}) @@ -360,6 +370,56 @@ def to_coolest( ) ) + # The tracer's own `MassField` objects. Unlike the legacy peel above, which pulls only + # `ExternalShear` / `MassSheet` off a galaxy, a field carries *every* mass profile it holds -- + # it is already the container for external mass, so there is nothing to select. + for j, field in enumerate(fields): + field_name = f"field_{j}" + + field_mass_profiles = field.cls_list_from(cls=MassProfile) + + sigma_crit = None + if any(isinstance(profile, NFW) for profile in field_mass_profiles): + if field.redshift < redshift_max: + sigma_crit = _sigma_crit_from( + cosmology=cosmology, + redshift_0=field.redshift, + redshift_1=redshift_max, + ) + + field_dicts = [] + + for profile in field_mass_profiles: + try: + field_dicts.append( + coolest_dict_from_mass(profile=profile, sigma_crit=sigma_crit) + ) + except (ag.exc.ProfileException, ValueError): + if on_unsupported == "raise": + raise + _record_skipped( + metadata=metadata, + galaxy_name=field_name, + profile_name=_profile_name_from(profile), + ) + + field_mass_model = lazy.MassModel(*[d["type"] for d in field_dicts]) + + for coolest_profile, profile_dict in zip(list(field_mass_model), field_dicts): + _set_parameters( + coolest_profile=coolest_profile, + parameters=profile_dict["parameters"], + point_estimate_cls=PointEstimate, + ) + + entities.append( + lazy.MassField( + field_name, + float(field.redshift), + mass_model=field_mass_model, + ) + ) + h0 = getattr(cosmology.H0, "value", cosmology.H0) om0 = getattr(cosmology.Om0, "value", cosmology.Om0) @@ -413,9 +473,11 @@ def from_coolest( Every COOLEST ``Galaxy`` entity becomes a PyAutoLens galaxy with its light profiles (``light_0``, ``light_1``, ...) and mass profiles (``mass_0``, - ...); every ``MassField`` entity becomes a galaxy holding its external - mass profiles. All parameters are converted from COOLEST conventions to - PyAutoLens conventions (see ``autogalaxy.interop.coolest``). + ...); every ``MassField`` entity becomes an ``ag.MassField`` holding its + external mass profiles (``mass_0``, ...), returned in the tracer's + ``fields`` rather than its ``galaxies``. All parameters are converted from + COOLEST conventions to PyAutoLens conventions (see + ``autogalaxy.interop.coolest``). Parameters ---------- @@ -450,8 +512,13 @@ def from_coolest( redshift_max = max(float(entity.redshift) for entity in entities) galaxies = [] + fields = [] for entity in entities: + # A COOLEST `MassField` is the standard's own external-field entity, so it maps onto + # PyAutoLens's `MassField` rather than onto a galaxy holding external mass profiles. + is_field = type(entity).__name__ == "MassField" + profiles = {} light_model = getattr(entity, "light_model", None) or [] @@ -483,6 +550,9 @@ def from_coolest( intermediate=intermediate, ) - galaxies.append(Galaxy(redshift=float(entity.redshift), **profiles)) + if is_field: + fields.append(ag.MassField(redshift=float(entity.redshift), **profiles)) + else: + galaxies.append(Galaxy(redshift=float(entity.redshift), **profiles)) - return Tracer(galaxies=galaxies, cosmology=cosmology) + return Tracer(galaxies=galaxies, fields=fields, cosmology=cosmology) diff --git a/autolens/jax/registration.py b/autolens/jax/registration.py index 77b63ed67..8485d49ee 100644 --- a/autolens/jax/registration.py +++ b/autolens/jax/registration.py @@ -27,10 +27,11 @@ def register_tracer_classes(tracer) -> bool: """Register every concrete class reachable from ``tracer`` as a JAX pytree. - Walks ``tracer.galaxies`` and registers ``Galaxy`` plus each light / - mass / point profile class encountered. Also registers ``Tracer`` - itself with ``no_flatten=("cosmology",)`` so the cosmology rides as - aux data across the JIT boundary (it is a per-fit constant). + Walks ``tracer.galaxies`` and then ``tracer.fields``, registering + ``Galaxy`` and ``MassField`` plus each light / mass / point profile class + encountered. Also registers ``Tracer`` itself with + ``no_flatten=("cosmology",)`` so the cosmology rides as aux data across the + JIT boundary (it is a per-fit constant). Returns ``True`` if registration ran (or was already complete), ``False`` if JAX is not installed (in which case the call is a silent @@ -43,6 +44,7 @@ def register_tracer_classes(tracer) -> bool: from autoarray.abstract_ndarray import register_instance_pytree from autogalaxy.galaxy.galaxy import Galaxy + from autogalaxy.galaxy.mass_field import MassField from autolens.lens.tracer import Tracer register_instance_pytree(Tracer, no_flatten=("cosmology",)) @@ -57,9 +59,18 @@ def register_tracer_classes(tracer) -> bool: # redshifts constant. register_instance_pytree(Galaxy, no_flatten=("redshift",)) + # A ``MassField`` is a redshift-bearing member of the tracer exactly as a + # ``Galaxy`` is, and its redshift is read by the same plane bookkeeping, so + # it gets the same ``no_flatten=("redshift",)`` treatment for the same + # reason. + register_instance_pytree(MassField, no_flatten=("redshift",)) + for galaxy in tracer.galaxies: _register_object_classes(galaxy) + for field in tracer.fields: + _register_object_classes(field) + return True diff --git a/autolens/lens/los.py b/autolens/lens/los.py index 0f3f35a95..871a09636 100644 --- a/autolens/lens/los.py +++ b/autolens/lens/los.py @@ -614,10 +614,39 @@ def __init__( self.mass_concentration_coefficients = mass_concentration_coefficients self.seed = seed - def galaxies_from(self) -> List[ag.Galaxy]: + # The sampled population, cached so `galaxies_from()` and `fields_from()` describe the *same* + # line of sight. Without it the two would resample -- and with the default `seed=None` the + # second call would draw an entirely different population, which is not a line of sight at all. + self._population = None + + def _sampled_population(self): + """ + Returns the sampled line-of-sight population, sampling it on first call and caching it. + + Returns + ------- + A tuple ``(halo_galaxies_per_plane, sheet_kappas_per_plane)``, where the first is one list of + halo `Galaxy` objects per plane and the second one ``(z_centre, kappa_negative)`` tuple per + plane. They are held per plane, rather than as one flat list of halos, purely so + `galaxies_from` can rebuild its historical plane-by-plane ordering exactly. + """ + if self._population is None: + self._population = self._sample_population() + + return self._population + + def galaxies_from(self, sheets_as_fields: bool = False) -> List[ag.Galaxy]: """ Sample LOS halos and negative kappa sheets. + Parameters + ---------- + sheets_as_fields + If True, only the halo galaxies are returned and the negative-kappa sheets are left to + `fields_from`, which returns them as `MassField` objects for `Tracer(fields=...)`. The + default is False, which keeps the historical return -- halos *and* sheet galaxies -- so + existing scripts (e.g. the workspace `los_halos` simulators) are unaffected. + Returns ------- galaxies : list of Galaxy @@ -625,7 +654,61 @@ def galaxies_from(self) -> List[ag.Galaxy]: - One ``Galaxy`` per sampled halo, each with an ``NFWTruncatedSph`` mass profile. - One ``Galaxy`` per plane with a ``MassSheet`` for the - negative kappa correction. + negative kappa correction (omitted if ``sheets_as_fields=True``). + """ + halo_galaxies_per_plane, sheet_kappas_per_plane = self._sampled_population() + + if sheets_as_fields: + return [galaxy for halos in halo_galaxies_per_plane for galaxy in halos] + + galaxies = [] + + for halos, (z_cen, kappa_neg) in zip( + halo_galaxies_per_plane, sheet_kappas_per_plane + ): + galaxies += halos + galaxies.append( + ag.Galaxy( + redshift=z_cen, + mass_sheet=ag.mp.MassSheet(kappa=kappa_neg), + ) + ) + + return galaxies + + def fields_from(self) -> List[ag.MassField]: + """ + Returns the negative-kappa sheet of every plane as a `MassField`. + + A negative-kappa sheet is not a galaxy: it is the mean-density correction of the line of + sight, which is exactly the external mass a `MassField` exists to hold. Pass the result to + `Tracer(fields=...)` beside `galaxies_from(sheets_as_fields=True)`. + + The population is shared with `galaxies_from`, so the sheets returned here are the ones + belonging to the same sampled line of sight. + """ + _, sheet_kappas_per_plane = self._sampled_population() + + return [ + ag.MassField( + redshift=z_cen, + mass_sheet=ag.mp.MassSheet(kappa=kappa_neg), + ) + for z_cen, kappa_neg in sheet_kappas_per_plane + ] + + def galaxies_and_fields_from(self) -> Tuple[List[ag.Galaxy], List[ag.MassField]]: + """ + Returns ``(halo_galaxies, sheet_fields)`` -- the pair `Tracer(galaxies=..., fields=...)` takes. + + This is the convenience door for the `MassField` form; `galaxies_from()` on its own remains + the galaxy-attached form and is not deprecated. + """ + return self.galaxies_from(sheets_as_fields=True), self.fields_from() + + def _sample_population(self): + """ + Samples the line-of-sight population once. See `_sampled_population` for the return shape. """ cosmology = self.cosmology rng = np.random.RandomState(self.seed) @@ -680,9 +763,12 @@ def galaxies_from(self) -> List[ag.Galaxy]: ) mc_coeffs[i] = [A, B] - galaxies = [] + halo_galaxies_per_plane = [] + sheet_kappas_per_plane = [] for i in range(n_planes): + plane_halo_galaxies = [] + z_lo = boundaries[i] z_hi = boundaries[i + 1] z_cen = centres[i] @@ -757,7 +843,7 @@ def galaxies_from(self) -> List[ag.Galaxy]: cosmology=cosmology, truncation_factor=self.truncation_factor, ) - galaxies.append( + plane_halo_galaxies.append( ag.Galaxy(redshift=z_cen, mass=halo) ) @@ -785,11 +871,7 @@ def galaxies_from(self) -> List[ag.Galaxy]: quad_limit=quad_limit, quad_epsrel=quad_epsrel, ) - galaxies.append( - ag.Galaxy( - redshift=z_cen, - mass_sheet=ag.mp.MassSheet(kappa=kappa_neg), - ) - ) + halo_galaxies_per_plane.append(plane_halo_galaxies) + sheet_kappas_per_plane.append((z_cen, kappa_neg)) - return galaxies + return halo_galaxies_per_plane, sheet_kappas_per_plane diff --git a/autolens/lens/tracer.py b/autolens/lens/tracer.py index da243131d..355e99ae3 100644 --- a/autolens/lens/tracer.py +++ b/autolens/lens/tracer.py @@ -33,6 +33,7 @@ from autogalaxy.profiles.geometry_profiles import GeometryProfile from autogalaxy.profiles.light.abstract import LightProfile from autogalaxy.profiles.light.snr import LightProfileSNR +from autogalaxy.galaxy.mass_field import MassField from autogalaxy.profiles.mass.abstract.abstract import MassProfile from autogalaxy.profiles.point_sources import Point, PointSolved @@ -72,7 +73,7 @@ class MultiPlaneRedshiftWarning(UserWarning): """ -def _warn_if_no_light_is_behind_any_mass(galaxies): +def _warn_if_no_light_is_behind_any_mass(galaxies, fields=None): """ Warn if every light-bearing galaxy lies in front of (or level with) every mass-bearing galaxy, across more than one redshift plane. @@ -105,12 +106,23 @@ def _warn_if_no_light_is_behind_any_mass(galaxies): ---------- galaxies The galaxies the tracer was constructed with. + fields + The `MassField` objects the tracer was constructed with, if any. A field is + mass-bearing exactly like a galaxy's mass profiles, so it participates in the + geometry check; passing no fields leaves the check byte-identical to the + galaxies-only behaviour it had before fields existed. """ try: galaxy_list = list(galaxies) except TypeError: return + if fields is not None: + try: + galaxy_list = galaxy_list + list(fields) + except TypeError: + return + mass_redshifts = [] light_redshifts = [] all_redshifts = [] @@ -225,11 +237,67 @@ def _validate_galaxies(galaxies): ) +def _validate_fields(fields): + """ + Raise if ``fields`` is not something the tracer can treat as a collection of + ``MassField`` objects. + + This mirrors ``_validate_galaxies`` exactly — same container checks, same + string trap, same ``af.ModelInstance`` escape hatch — but names ``MassField`` + in every message, because the two arguments are not interchangeable and the + most likely mistake is passing one where the other belongs. + + A ``Galaxy`` in ``fields`` is called out by name: it is a valid lensing object + in the wrong argument, and saying "not a MassField" would not tell the caller + where it should have gone. + + Parameters + ---------- + fields + The input to validate. + """ + if isinstance(fields, af.ModelInstance): + return + + if isinstance(fields, (str, bytes)): + raise TypeError( + f"fields must be an iterable of MassField objects, but a " + f"{type(fields).__name__} was input: {fields!r}. A string is itself " + f"iterable, so this is not caught by an iterability check — pass a list " + f"of MassField objects, e.g. Tracer(galaxies=[lens, source], " + f"fields=[field])" + ) + + if not isinstance(fields, (list, tuple)): + raise TypeError( + f"fields must be an iterable of MassField objects, but a " + f"{type(fields).__name__} was input: {fields!r}. Pass a list of " + f"MassField objects, e.g. Tracer(galaxies=[lens, source], fields=[field])" + ) + + for index, field in enumerate(fields): + if isinstance(field, ag.Galaxy): + raise TypeError( + f"fields must be an iterable of MassField objects, but the entry at " + f"index {index} is a Galaxy: {field!r}. Galaxies go in the galaxies " + f"argument, e.g. Tracer(galaxies=[lens, source], fields=[field])" + ) + + if not hasattr(field, "redshift"): + raise TypeError( + f"fields must be an iterable of MassField objects, but the entry at " + f"index {index} is a {type(field).__name__}: {field!r}. Every entry " + f"needs a redshift, which is what the ray-tracing calculation groups " + f"the tracer's members into planes by" + ) + + class Tracer(ABC, ag.OperateImageGalaxies): def __init__( self, galaxies: Union[List[ag.Galaxy], af.ModelInstance], cosmology: ag.cosmo.LensingCosmology = None, + fields: Optional[Union[List[MassField], af.ModelInstance]] = None, ): """ Performs gravitational lensing ray-tracing calculations based on an input list of galaxies and a cosmology. @@ -260,18 +328,47 @@ def __init__( The list of galaxies which make up the gravitational lensing ray-tracing system. cosmology The cosmology used to perform ray-tracing calculations. + fields + The `MassField` objects describing the tidal field of everything outside the modelled system (e.g. an + external shear, a mass sheet, an external potential). Each is placed at its own redshift in the + multi-plane calculation exactly like a galaxy's mass, but contributes no light. Attaching these + components to a `Galaxy` instead remains fully supported and is not deprecated. """ # if isinstance(galaxies, af.ModelInstance): # galaxies = list(galaxies.values()) _validate_galaxies(galaxies=galaxies) - _warn_if_no_light_is_behind_any_mass(galaxies=galaxies) + + if fields is not None: + _validate_fields(fields=fields) + + _warn_if_no_light_is_behind_any_mass(galaxies=galaxies, fields=fields) self.galaxies = galaxies + # Stored as a plain list (never None) so every member-walking property below can concatenate it + # unconditionally, and so a tracer with no fields serialises as `fields: []` rather than a null. + self.fields = list(fields) if fields is not None else [] + self.cosmology = cosmology or ag.cosmo.Planck15() + @property + def members(self) -> List: + """ + Returns every object the tracer places in a plane: its galaxies followed by its fields. + + A `MassField` is not a galaxy and never appears in `galaxies`, but it *is* a redshift-bearing, + mass-bearing member of the lensing system. Every lensing-quantity calculation (deflections, + convergence, potential, ray-tracing) therefore runs over `members`, whereas every surface which + enumerates *galaxies* (e.g. the per-galaxy image dictionary) runs over `galaxies` alone. + + Returns + ------- + The galaxies and fields of the tracer, in input order, galaxies first. + """ + return list(self.galaxies) + list(self.fields) + @property def galaxies_ascending_redshift(self) -> List[ag.Galaxy]: """ @@ -299,6 +396,32 @@ def galaxies_ascending_redshift(self) -> List[ag.Galaxy]: return list(self.galaxies) + @property + def members_ascending_redshift(self) -> List: + """ + Returns every member of the tracer -- its galaxies and its fields -- in ascending redshift order. + + This is the list the plane calculation is built from, so that a `MassField` sits in the plane at its + own redshift and contributes its mass to the multi-plane ray-tracing exactly as a galaxy's mass does. + + `galaxies_ascending_redshift` keeps its galaxies-only meaning and is unchanged; the two properties + agree whenever a tracer has no fields. + + The traced-redshift bypass matches `galaxies_ascending_redshift`: if any member has a JAX-traced + `redshift`, the input order is trusted (galaxies first, then fields), because a traced value cannot + be used as a sort key. + + Returns + ------- + The galaxies and fields in the tracer in ascending redshift order. + """ + members = self.members + + if not tracer_util._any_traced(members): + return sorted(members, key=lambda member: member.redshift) + + return members + @property def plane_redshifts(self) -> List[float]: """ @@ -321,8 +444,9 @@ def plane_redshifts(self) -> List[float]: ------- The list of unique redshifts of the planes. """ + # Members, not galaxies: a `MassField` at a redshift no galaxy occupies is a plane of its own. return tracer_util.plane_redshifts_from( - galaxies=self.galaxies_ascending_redshift + galaxies=self.members_ascending_redshift ) @property @@ -353,8 +477,10 @@ def planes(self) -> List[ag.Galaxies]: ------- The list of list of galaxies grouped into their planes. """ + # Members, not galaxies: a plane may hold a `MassField` beside its galaxies, so that its mass enters + # every plane-level lensing calculation while every light calculation sees the field's zeros. return tracer_util.planes_from( - galaxies=self.galaxies_ascending_redshift, + galaxies=self.members_ascending_redshift, plane_redshifts=self.plane_redshifts, ) @@ -366,6 +492,7 @@ def sliced_tracer_from( source_galaxies: List[ag.Galaxy], planes_between_lenses: List[int], cosmology: ag.cosmo.LensingCosmology = None, + fields: Optional[List[MassField]] = None, ): """ Returns a tracer where the lens system is split into planes with specified redshift distances between them. @@ -411,6 +538,9 @@ def sliced_tracer_from( between Earth (redshift 0.0) and the first lens galaxy, the next between the lens and source, etc. cosmology The cosmology used to perform ray-tracing calculations. + fields + The `MassField` objects describing the external tidal field. Each has its redshift snapped to the + nearest sliced plane redshift, exactly as the `line_of_sight_galaxies` do. """ cosmology = cosmology or ag.cosmo.Planck15() @@ -426,7 +556,11 @@ def sliced_tracer_from( galaxies = lens_galaxies + line_of_sight_galaxies + source_galaxies - for galaxy in galaxies: + fields = list(fields) if fields is not None else [] + + # Fields are snapped exactly as `line_of_sight_galaxies` are: a field is a redshift-bearing member of + # the system, so leaving it off a sliced plane would defeat the whole point of the slicing. + for galaxy in galaxies + fields: redshift_differences = list( map(lambda z: abs(z - galaxy.redshift), plane_redshifts) ) @@ -434,7 +568,7 @@ def sliced_tracer_from( redshift_differences.index(min(redshift_differences)) ] - return Tracer(galaxies=galaxies, cosmology=cosmology) + return Tracer(galaxies=galaxies, cosmology=cosmology, fields=fields) @property def total_planes(self) -> int: @@ -577,9 +711,11 @@ def grid_2d_at_redshift_from( cosmology The cosmology used for ray-tracing from which angular diameter distances between planes are computed. """ + # A lensing-quantity path: the deflections used to trace to the input redshift must include the + # tracer's fields, so this walks members rather than galaxies. return tracer_util.grid_2d_at_redshift_from( redshift=redshift, - galaxies=self.galaxies_ascending_redshift, + galaxies=self.members_ascending_redshift, grid=grid, cosmology=self.cosmology, ) @@ -848,16 +984,23 @@ def galaxy_image_2d_dict_from( traced_grid_list = self.traced_grid_2d_list_from(grid=grid, xp=xp) for plane_index, galaxies in enumerate(self.planes): + # A plane may hold a `MassField` beside its galaxies. This dictionary is keyed by *galaxy*, and + # consumed by adaptive features and per-galaxy visualisation, so a field -- whose image is zeros + # by construction -- is skipped rather than added as a key with a zeros image. + plane_galaxies = [ + galaxy for galaxy in galaxies if not isinstance(galaxy, MassField) + ] + image_2d_list = [ galaxy.image_2d_from( grid=traced_grid_list[plane_index], operated_only=operated_only, xp=xp, ) - for galaxy in galaxies + for galaxy in plane_galaxies ] - for galaxy_index, galaxy in enumerate(galaxies): + for galaxy_index, galaxy in enumerate(plane_galaxies): galaxy_image_2d_dict[galaxy] = image_2d_list[galaxy_index] return galaxy_image_2d_dict @@ -923,11 +1066,10 @@ def deflections_of_planes_summed_from( grid The 2D (y, x) coordinates where values of the deflections are evaluated. """ + # Members, not galaxies: the deflections of the tracer are the deflections of every mass it holds, + # and a `MassField`'s mass is mass. return sum( - [ - galaxy.deflections_yx_2d_from(grid=grid, xp=xp) - for galaxy in self.galaxies - ] + [galaxy.deflections_yx_2d_from(grid=grid, xp=xp) for galaxy in self.members] ) @aa.decorators.to_vector_yx @@ -982,8 +1124,9 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np) -> aa.Array2D: grid The 2D (y, x) coordinates where values of the convergence are evaluated. """ + # Members, not galaxies: a field's convergence (e.g. a `MassSheet`) is part of the tracer's convergence. return sum( - [galaxy.convergence_2d_from(grid=grid, xp=xp) for galaxy in self.galaxies] + [galaxy.convergence_2d_from(grid=grid, xp=xp) for galaxy in self.members] ) @aa.decorators.to_array @@ -1009,8 +1152,9 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np) -> aa.Array2D: grid The 2D (y, x) coordinates where values of the potential are evaluated. """ + # Members, not galaxies: a field's potential (e.g. an `ExternalPotential`) is part of the tracer's. return sum( - [galaxy.potential_2d_from(grid=grid, xp=xp) for galaxy in self.galaxies] + [galaxy.potential_2d_from(grid=grid, xp=xp) for galaxy in self.members] ) @aa.decorators.to_array @@ -1031,7 +1175,8 @@ def time_delays_from(self, grid: aa.type.Grid2DLike, xp=np) -> aa.Array2D: function, which performs the calculation and has full latex documentation of the equations used. """ return tracer_util.time_delays_from( - galaxies=ag.Galaxies(self.galaxies_ascending_redshift), + # Members, not galaxies: a time delay is a lensing quantity computed from the tracer's full mass. + galaxies=ag.Galaxies(self.members_ascending_redshift), grid=grid, xp=xp, cosmology=self.cosmology, @@ -1055,7 +1200,9 @@ def has(self, cls: Type) -> bool: ------- True if any galaxy in the tracer has the input class type, else False. """ - return any(map(lambda galaxy: galaxy.has(cls=cls), self.galaxies)) + # Members, not galaxies: `tracer.has(cls=al.mp.ExternalShear)` must be True for a shear held by a + # field, exactly as it is for a shear attached to a galaxy. + return any(map(lambda galaxy: galaxy.has(cls=cls), self.members)) def cls_list_from(self, cls: Type) -> List: """ @@ -1071,7 +1218,8 @@ def cls_list_from(self, cls: Type) -> List: """ cls_list = [] - for galaxy in self.galaxies: + # Members, not galaxies: this lists *profiles*, and a field's profiles are profiles of the tracer. + for galaxy in self.members: if galaxy.has(cls=cls): for cls_galaxy in galaxy.cls_list_from(cls=cls): cls_list.append(cls_galaxy) @@ -1206,9 +1354,11 @@ def extract(value, name): except (AttributeError, IndexError): return None + # Members, not galaxies: this lists the *attributes of profiles* (e.g. every mass profile centre), + # so a field's profiles are included alongside the galaxies'. attributes = [ extract(value, attr_name) - for galaxy in self.galaxies + for galaxy in self.members for value in galaxy.__dict__.values() if isinstance(value, cls) ] @@ -1266,6 +1416,10 @@ def extract_attributes_of_planes( attribute_list = [] + # This walks *plane members*, so a `MassField` in a plane contributes its own entry (its mass profile + # attributes for `cls=MassProfile`, a `None` for a class it does not hold, filtered out by + # `filter_nones=True`). That is the member-aware half of the pair: `extract_attributes_of_galaxies` + # below enumerates galaxies and so stays galaxy-only. for plane in self.planes: for galaxy in plane: attribute_list += [ @@ -1330,6 +1484,9 @@ def extract_attributes_of_galaxies( The name of the attribute which is extracted from the class type (e.g. axis_ratio, centre). filter_nones """ + # Galaxies, not members: the returned list is indexed *by galaxy*, so a `MassField` would shift every + # index and break the caller's pairing. Use `extract_attribute` (whole tracer) or + # `extract_attributes_of_planes` (per plane) when a field's profiles must be included. if filter_nones: return [ galaxy.extract_attribute(cls=cls, attr_name=attr_name) @@ -1363,7 +1520,10 @@ def extract_profile(self, profile_name: str) -> GeometryProfile: profile_name The name of the profile component in the tracer. """ - for galaxy in self.galaxies: + # Members, not galaxies: a profile is looked up by its component name, and a field's components are + # named exactly as a galaxy's are. `extract_plane_index_of_profile` already walks planes (and so + # already sees fields), so the two must agree. + for galaxy in self.members: try: return galaxy.__dict__[profile_name] except KeyError: diff --git a/docs/api/galaxy.rst b/docs/api/galaxy.rst index 48cd03e86..da6550690 100644 --- a/docs/api/galaxy.rst +++ b/docs/api/galaxy.rst @@ -10,6 +10,13 @@ given redshift. ``Tracer`` groups galaxies by redshift into planes and performs multi-plane gravitational lensing ray-tracing, computing lensed images, convergence, deflection angles, magnification, critical curves, and caustics. +``MassField`` holds the mass of everything *outside* the modelled system — an external +shear, a mass sheet, an external potential — at its own redshift and with no light of its +own; pass a list of them as ``Tracer(galaxies=..., fields=[field])`` and each is placed in +the plane at its redshift, contributing its mass to every lensing calculation while every +per-galaxy surface still sees galaxies alone. Attaching these components to a ``Galaxy`` +instead remains fully supported. + .. currentmodule:: autolens .. autosummary:: @@ -19,6 +26,7 @@ deflection angles, magnification, critical curves, and caustics. Galaxy Galaxies + MassField Tracer To treat the redshift of a galaxy as a free parameter in a model, the ``Redshift`` object must diff --git a/docs/api/modeling.rst b/docs/api/modeling.rst index b5b8bcc1d..c261804f7 100644 --- a/docs/api/modeling.rst +++ b/docs/api/modeling.rst @@ -63,4 +63,23 @@ Adapt .. autosummary:: :toctree: generated/ - AdaptImages \ No newline at end of file + AdaptImages + +Model Utilities +--------------- + +Convenience functions which compose commonly used ``af.Model`` objects. + +``mass_field_from`` builds the ``af.Model(al.MassField, ...)`` external-field model that +goes in a model's ``fields`` collection, sharing the centre of its ``MassSheet`` / +``ExternalPotential`` components with the lens galaxy model's mass centre. + +.. currentmodule:: autolens.analysis.model_util + +.. autosummary:: + :toctree: generated/ + + mass_field_from + mge_model_from + mge_point_model_from + hilbert_pixels_from_pixel_scale diff --git a/docs/overview/overview_1_start_here.md b/docs/overview/overview_1_start_here.md index 4f67feba8..bce3ce6de 100644 --- a/docs/overview/overview_1_start_here.md +++ b/docs/overview/overview_1_start_here.md @@ -253,6 +253,26 @@ aplt.plot_array(array=tracer.image_2d_from(grid=grid), title="Image") :width: 600 ``` +### External fields + +Not all of the mass that lenses a source belongs to a galaxy: line-of-sight structure and +the group or cluster the lens sits in produce an external shear, a mass sheet and +higher-order terms which describe the tidal field of everything outside the system being +modelled. A `MassField` is the home for those components — mass at a redshift, with no +light of its own — and a tracer takes a list of them: + +```python +field = al.MassField(redshift=0.5, shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)) + +tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy], fields=[field]) +``` + +Each field is placed in the plane at its redshift and contributes its mass to every lensing +calculation, exactly as a galaxy's mass profiles do; it simply contributes no light. The +galaxy-attached form — writing `al.Galaxy(redshift=0.5, shear=al.mp.ExternalShear(...))` — +remains fully supported and is not deprecated, so nothing you have already written needs to +change. + ## Units The units used throughout the strong lensing literature vary, therefore lets quickly describe the units used in diff --git a/test_autolens/analysis/analysis/test_analysis_lens.py b/test_autolens/analysis/analysis/test_analysis_lens.py index 754b49868..1926a8b5b 100644 --- a/test_autolens/analysis/analysis/test_analysis_lens.py +++ b/test_autolens/analysis/analysis/test_analysis_lens.py @@ -70,3 +70,82 @@ def test__tracer_for_instance__subhalo_redshift_rescale_used(analysis_imaging_7x tracer = analysis_imaging_7x7.tracer_via_instance_from(instance=instance) assert tracer.galaxies[1].mass.centre == pytest.approx((-0.19959, -0.39919), 1.0e-4) + + +def test__tracer_for_instance__fields_collection_is_folded_into_tracer_fields( + analysis_imaging_7x7, +): + """ + A model may declare an external field beside its galaxies, as + `af.Collection(galaxies=..., fields=af.Collection(field=af.Model(al.MassField, ...)))`. + + The field must land in `tracer.fields`, and never in `tracer.galaxies`: a `MassField` is + not a galaxy, and every per-galaxy surface downstream is keyed on that distinction. + """ + model = af.Collection( + galaxies=af.Collection( + lens=al.Galaxy( + redshift=0.5, + mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.0), + ), + source=al.Galaxy(redshift=1.0, bulge=al.lp.SersicSph(intensity=1.0)), + ), + fields=af.Collection( + field=af.Model( + al.MassField, redshift=0.5, shear=af.Model(al.mp.ExternalShear) + ) + ), + ) + + instance = model.instance_from_prior_medians() + tracer = analysis_imaging_7x7.tracer_via_instance_from(instance=instance) + + assert len(tracer.fields) == 1 + assert isinstance(tracer.fields[0], al.MassField) + assert tracer.fields[0].redshift == 0.5 + assert all(not isinstance(galaxy, al.MassField) for galaxy in tracer.galaxies) + assert len(list(tracer.galaxies)) == 2 + + +def test__log_likelihood_function__a_field_shear_matches_the_galaxy_attached_shear( + analysis_imaging_7x7, +): + """ + The point of the whole feature: an external shear held by a `MassField` must give the + same likelihood as the identical shear attached to the lens galaxy. Both instances are + built at the same prior-median parameter values, so the two models are the same physics + expressed two ways. + """ + shear = af.Model(al.mp.ExternalShear) + + field_model = af.Collection( + galaxies=af.Collection( + lens=al.Galaxy( + redshift=0.5, + mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.0), + ), + source=al.Galaxy(redshift=1.0, bulge=al.lp.SersicSph(intensity=1.0)), + ), + fields=af.Collection(field=af.Model(al.MassField, redshift=0.5, shear=shear)), + ) + + galaxy_model = af.Collection( + galaxies=af.Collection( + lens=af.Model( + al.Galaxy, + redshift=0.5, + mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.0), + shear=shear, + ), + source=al.Galaxy(redshift=1.0, bulge=al.lp.SersicSph(intensity=1.0)), + ), + ) + + field_likelihood = analysis_imaging_7x7.log_likelihood_function( + instance=field_model.instance_from_prior_medians() + ) + galaxy_likelihood = analysis_imaging_7x7.log_likelihood_function( + instance=galaxy_model.instance_from_prior_medians() + ) + + assert field_likelihood == pytest.approx(galaxy_likelihood, 1.0e-8) diff --git a/test_autolens/analysis/test_model_util.py b/test_autolens/analysis/test_model_util.py new file mode 100644 index 000000000..8e60adb04 --- /dev/null +++ b/test_autolens/analysis/test_model_util.py @@ -0,0 +1,117 @@ +""" +Tests for `al.model_util.mass_field_from` — the composer for the external-field model that +goes in a lens model's `fields` collection. + +The behaviour worth pinning is the *shared* centre: the `MassSheet` and `ExternalPotential` +in the field are expanded about the primary deflector's centre, and they must share the lens +mass's centre prior object rather than get a copy of it. A copy would silently add two free +parameters and let the two centres drift apart in the search. +""" + +import pytest + +import autofit as af +import autolens as al + + +@pytest.fixture(name="lens") +def make_lens(): + return af.Model(al.Galaxy, redshift=0.5, mass=af.Model(al.mp.Isothermal)) + + +def test__shear_only_is_the_default(lens): + field = al.model_util.mass_field_from(lens=lens) + + assert field.cls is al.MassField + assert hasattr(field, "shear") + assert not hasattr(field, "mass_sheet") + assert not hasattr(field, "potential") + + +def test__mass_sheet_alone(lens): + field = al.model_util.mass_field_from(lens=lens, shear=False, mass_sheet=True) + + assert not hasattr(field, "shear") + assert hasattr(field, "mass_sheet") + assert not hasattr(field, "potential") + + +def test__potential_alone(lens): + field = al.model_util.mass_field_from(lens=lens, shear=False, potential=True) + + assert not hasattr(field, "shear") + assert not hasattr(field, "mass_sheet") + assert hasattr(field, "potential") + + +def test__all_three_together(lens): + field = al.model_util.mass_field_from( + lens=lens, shear=True, mass_sheet=True, potential=True + ) + + assert hasattr(field, "shear") + assert hasattr(field, "mass_sheet") + assert hasattr(field, "potential") + + +def test__the_centres_are_the_same_prior_object_as_the_lens_mass_centre(lens): + field = al.model_util.mass_field_from( + lens=lens, shear=True, mass_sheet=True, potential=True + ) + + assert field.mass_sheet.centre is lens.mass.centre + assert field.potential.centre is lens.mass.centre + + +def test__the_shared_centre_costs_no_extra_dimensions(lens): + tied = af.Collection( + galaxies=af.Collection(lens=lens), + fields=af.Collection( + field=al.model_util.mass_field_from( + lens=lens, shear=False, mass_sheet=True, potential=True + ) + ), + ) + + untied = af.Collection( + galaxies=af.Collection(lens=lens), + fields=af.Collection( + field=af.Model( + al.MassField, + redshift=0.5, + mass_sheet=af.Model(al.mp.MassSheet), + potential=af.Model(al.mp.ExternalPotential), + ) + ), + ) + + # Two centres, two components each, shared away. + assert tied.prior_count == untied.prior_count - 4 + + +def test__a_composed_collection_builds_an_instance_with_matching_centres(lens): + model = af.Collection( + galaxies=af.Collection(lens=lens), + fields=af.Collection( + field=al.model_util.mass_field_from(lens=lens, potential=True) + ), + ) + + instance = model.instance_from_prior_medians() + + assert isinstance(instance.fields.field, al.MassField) + assert instance.fields.field.potential.centre == pytest.approx( + instance.galaxies.lens.mass.centre, 1.0e-12 + ) + + +def test__redshift_defaults_to_the_lens_and_can_be_overridden(lens): + assert al.model_util.mass_field_from(lens=lens).redshift == 0.5 + assert al.model_util.mass_field_from(lens=lens, redshift=0.8).redshift == 0.8 + + +def test__a_lens_model_with_no_mass_raises_and_names_the_model(): + lens = af.Model(al.Galaxy, redshift=0.5, bulge=af.Model(al.lp.Sersic)) + + with pytest.raises(ValueError, match="no `mass` attribute"): + al.model_util.mass_field_from(lens=lens) diff --git a/test_autolens/interop/fixtures/legacy_galaxy_attached.json b/test_autolens/interop/fixtures/legacy_galaxy_attached.json new file mode 100644 index 000000000..c0d08c4f7 --- /dev/null +++ b/test_autolens/interop/fixtures/legacy_galaxy_attached.json @@ -0,0 +1,459 @@ +{ + "coordinates_origin": { + "dec": "+00d00m00s", + "ra": "00h00m00s" + }, + "cosmology": { + "H0": 67.74, + "Om0": 0.3075, + "astropy_name": "FlatLambdaCDM" + }, + "instrument": { + "band": "", + "name": "", + "pixel_size": 0.12, + "psf": { + "description": null, + "type": null + }, + "readout_noise": 0.0 + }, + "lensing_entities": [ + { + "light_model": [ + { + "id": "0-galaxy-light-0-Sersic", + "parameters": { + "I_eff": { + "fixed": false, + "id": "0-galaxy-light-0-Sersic-I_eff", + "point_estimate": { + "value": 1.2 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "center_x": { + "fixed": false, + "id": "0-galaxy-light-0-Sersic-center_x", + "point_estimate": { + "value": -0.03 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "center_y": { + "fixed": false, + "id": "0-galaxy-light-0-Sersic-center_y", + "point_estimate": { + "value": 0.05 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "n": { + "fixed": false, + "id": "0-galaxy-light-0-Sersic-n", + "point_estimate": { + "value": 3.5 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "phi": { + "fixed": false, + "id": "0-galaxy-light-0-Sersic-phi", + "point_estimate": { + "value": -20.0 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "q": { + "fixed": false, + "id": "0-galaxy-light-0-Sersic-q", + "point_estimate": { + "value": 0.8 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "theta_eff": { + "fixed": false, + "id": "0-galaxy-light-0-Sersic-theta_eff", + "point_estimate": { + "value": 0.9 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + } + }, + "type": "Sersic" + } + ], + "mass_model": [ + { + "id": "0-galaxy-mass-0-SIE", + "parameters": { + "center_x": { + "fixed": false, + "id": "0-galaxy-mass-0-SIE-center_x", + "point_estimate": { + "value": -0.03 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "center_y": { + "fixed": false, + "id": "0-galaxy-mass-0-SIE-center_y", + "point_estimate": { + "value": 0.05 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "phi": { + "fixed": false, + "id": "0-galaxy-mass-0-SIE-phi", + "point_estimate": { + "value": -45.0 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "q": { + "fixed": false, + "id": "0-galaxy-mass-0-SIE-q", + "point_estimate": { + "value": 0.7 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "theta_E": { + "fixed": false, + "id": "0-galaxy-mass-0-SIE-theta_E", + "point_estimate": { + "value": 1.279597687640351 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + } + }, + "type": "SIE" + } + ], + "name": "galaxy_0", + "redshift": 0.5, + "type": "Galaxy" + }, + { + "mass_model": [ + { + "id": "1-massfield-mass-0-ExternalShear", + "parameters": { + "gamma_ext": { + "fixed": false, + "id": "1-massfield-mass-0-ExternalShear-gamma_ext", + "point_estimate": { + "value": 0.03605551275463989 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "phi_ext": { + "fixed": false, + "id": "1-massfield-mass-0-ExternalShear-phi_ext", + "point_estimate": { + "value": 61.845033762989885 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + } + }, + "type": "ExternalShear" + }, + { + "id": "1-massfield-mass-1-ConvergenceSheet", + "parameters": { + "kappa_s": { + "fixed": false, + "id": "1-massfield-mass-1-ConvergenceSheet-kappa_s", + "point_estimate": { + "value": 0.05 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + } + }, + "type": "ConvergenceSheet" + } + ], + "name": "mass_field_0", + "redshift": 0.5, + "type": "MassField" + }, + { + "light_model": [ + { + "id": "2-galaxy-light-0-Sersic", + "parameters": { + "I_eff": { + "fixed": false, + "id": "2-galaxy-light-0-Sersic-I_eff", + "point_estimate": { + "value": 0.7 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "center_x": { + "fixed": false, + "id": "2-galaxy-light-0-Sersic-center_x", + "point_estimate": { + "value": 0.2 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "center_y": { + "fixed": false, + "id": "2-galaxy-light-0-Sersic-center_y", + "point_estimate": { + "value": 0.1 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "n": { + "fixed": false, + "id": "2-galaxy-light-0-Sersic-n", + "point_estimate": { + "value": 1.2 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "phi": { + "fixed": false, + "id": "2-galaxy-light-0-Sersic-phi", + "point_estimate": { + "value": 60.0 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "q": { + "fixed": false, + "id": "2-galaxy-light-0-Sersic-q", + "point_estimate": { + "value": 0.6 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + }, + "theta_eff": { + "fixed": false, + "id": "2-galaxy-light-0-Sersic-theta_eff", + "point_estimate": { + "value": 0.3 + }, + "posterior_stats": { + "mean": null, + "median": null, + "percentile_16th": null, + "percentile_84th": null + }, + "prior": { + "type": null + } + } + }, + "type": "Sersic" + } + ], + "mass_model": [], + "name": "galaxy_1", + "redshift": 1.5, + "type": "Galaxy" + } + ], + "likelihoods": null, + "meta": { + "coolest.__version__": "0.1.11", + "generated_by": "PyAutoLens (autolens.interop.coolest); theta_E is the profile's intermediate-axis Einstein radius, NFW rho_c is in solar masses per arcsec**3 with Sigma_crit in solar masses per arcsec**2." + }, + "mode": "MAP", + "observation": { + "exposure_time": null, + "mag_sky_brightness": null, + "mag_zero_point": null, + "noise": { + "type": null + }, + "pixels": { + "field_of_view_x": [ + -0.6, + 0.6 + ], + "field_of_view_y": [ + -0.6, + 0.6 + ], + "fits_file": { + "path": null + }, + "num_pix_x": 10, + "num_pix_y": 10 + } + }, + "standard": "COOLEST" +} \ No newline at end of file diff --git a/test_autolens/interop/test_coolest.py b/test_autolens/interop/test_coolest.py index 28766fd00..a700779ac 100644 --- a/test_autolens/interop/test_coolest.py +++ b/test_autolens/interop/test_coolest.py @@ -1,4 +1,5 @@ import json +from pathlib import Path import numpy as np import pytest @@ -191,7 +192,11 @@ def test__round_trip__relative_file_path(tmp_path, monkeypatch): tracer_back = interop_coolest.from_coolest(file_path="template.json") - assert len(tracer_back.galaxies) == 3 + # Two galaxies and one field: a COOLEST `MassField` entity (here the shear peeled off + # the lens galaxy on export) is imported as an `al.MassField` in `tracer.fields`, not as + # a third galaxy. The three lensing entities are unchanged; only where they land is. + assert len(tracer_back.galaxies) == 2 + assert len(tracer_back.fields) == 1 def test__from_coolest__missing_point_estimate_raises(tmp_path): @@ -406,7 +411,10 @@ def test__from_coolest__skipped_template_returns_mass_profiles_only(tmp_path): light_profiles = [] mass_profiles = [] - for galaxy in tracer_back.galaxies: + # `members`, not `galaxies`: the exported `MassField` entity now imports as an + # `al.MassField` in `tracer.fields`, so the shear is a member of the tracer rather than + # a profile of a galaxy. The set of imported profiles is unchanged. + for galaxy in tracer_back.members: light_profiles += galaxy.cls_list_from(cls=LightProfile) mass_profiles += galaxy.cls_list_from(cls=MassProfile) @@ -415,3 +423,224 @@ def test__from_coolest__skipped_template_returns_mass_profiles_only(tmp_path): types = sorted(type(profile).__name__ for profile in mass_profiles) assert types == ["ExternalShear", "Isothermal"] + + +# ====================================================================================== +# MassField +# ====================================================================================== + + +def legacy_galaxy_attached_tracer(): + """ + The galaxy-attached form: an `ExternalShear` and a `MassSheet` bolted onto the lens + galaxy. This form remains fully supported and its export must never change, which is + what `test__to_coolest__legacy_galaxy_attached_export_is_byte_identical` pins. + """ + lens = al.Galaxy( + redshift=0.5, + bulge=al.lp.Sersic( + centre=(0.05, -0.03), + ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=70.0), + intensity=1.2, + effective_radius=0.9, + sersic_index=3.5, + ), + mass=al.mp.Isothermal( + centre=(0.05, -0.03), + ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=45.0), + einstein_radius=1.3, + ), + shear=al.mp.ExternalShear(gamma_1=0.02, gamma_2=-0.03), + mass_sheet=al.mp.MassSheet(centre=(0.0, 0.0), kappa=0.05), + ) + source = al.Galaxy( + redshift=1.5, + bulge=al.lp.Sersic( + centre=(0.1, 0.2), + ell_comps=al.convert.ell_comps_from(axis_ratio=0.6, angle=-30.0), + intensity=0.7, + effective_radius=0.3, + sersic_index=1.2, + ), + ) + return al.Tracer(galaxies=[lens, source], cosmology=al.cosmo.Planck15()) + + +def test__to_coolest__legacy_galaxy_attached_export_is_byte_identical(tmp_path): + """ + Regression: the template written for a galaxy-attached external field must be exactly + what it was before the `fields` slot existed — same entity names (`mass_field_{i}`), + same ordering, same numbers. + + The fixture was written from the pre-change tree. The comparison is of the JSON text, + normalised only through `json.load` (the template's `meta` carries no timestamp, only + the coolest version and a fixed `generated_by` string, so nothing volatile is elided). + """ + file_path = interop_coolest.to_coolest( + galaxies=legacy_galaxy_attached_tracer(), + file_path=str(tmp_path / "template"), + shape_native=(10, 10), + pixel_size=0.12, + ) + + # `fixtures/`, not `files/`: the repository's root `.gitignore` excludes every directory + # named `files`, so a fixture written there would never reach CI. + fixture_path = ( + Path(__file__).resolve().parent / "fixtures" / "legacy_galaxy_attached.json" + ) + + with open(file_path) as f: + written = f.read() + + with open(fixture_path) as f: + expected = f.read() + + assert written == expected + + names = [entity["name"] for entity in json.loads(written)["lensing_entities"]] + + assert names == ["galaxy_0", "mass_field_0", "galaxy_1"] + + +def test__from_coolest__legacy_export_imports_its_peeled_sheets_as_a_mass_field( + tmp_path, +): + """ + The one import behaviour change: a COOLEST `MassField` entity — including the one the + legacy peel writes for a galaxy-attached shear/sheet — comes back as an `al.MassField` + in `tracer.fields`, not as an extra galaxy in `tracer.galaxies`. + """ + file_path = interop_coolest.to_coolest( + galaxies=legacy_galaxy_attached_tracer(), + file_path=str(tmp_path / "template"), + shape_native=(10, 10), + ) + + tracer_back = interop_coolest.from_coolest(file_path=file_path) + + assert len(tracer_back.galaxies) == 2 + assert len(tracer_back.fields) == 1 + assert isinstance(tracer_back.fields[0], al.MassField) + assert tracer_back.fields[0].redshift == 0.5 + + types = sorted( + type(profile).__name__ + for profile in tracer_back.fields[0].cls_list_from(cls=al.mp.MassProfile) + ) + + assert types == ["ExternalShear", "MassSheet"] + + +def test__to_coolest__tracer_fields_are_written_as_their_own_entities(tmp_path): + tracer = al.Tracer( + galaxies=lens_model_galaxies(), + fields=[ + al.MassField( + redshift=0.5, + shear=al.mp.ExternalShear(gamma_1=0.01, gamma_2=0.02), + mass_sheet=al.mp.MassSheet(centre=(0.0, 0.0), kappa=0.04), + ) + ], + ) + + file_path = interop_coolest.to_coolest( + galaxies=tracer, file_path=str(tmp_path / "template"), shape_native=(10, 10) + ) + + with open(file_path) as f: + template = json.load(f) + + entities = template["lensing_entities"] + names = [entity["name"] for entity in entities] + + # `mass_field_0` is the legacy peel of the lens galaxy's own shear; `field_0` is the + # tracer's `MassField`, carrying both of its profiles. + assert "field_0" in names + + field_entity = [e for e in entities if e["name"] == "field_0"][0] + + assert field_entity["type"] == "MassField" + assert field_entity["redshift"] == 0.5 + assert sorted(p["type"] for p in field_entity["mass_model"]) == [ + "ConvergenceSheet", + "ExternalShear", + ] + + +def test__round_trip__a_tracer_with_fields_returns_its_field(tmp_path): + shear = al.mp.ExternalShear(gamma_1=0.01, gamma_2=0.02) + mass_sheet = al.mp.MassSheet(centre=(0.0, 0.0), kappa=0.04) + + tracer = al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal( + centre=(0.05, -0.03), + ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=45.0), + einstein_radius=1.3, + ), + ), + al.Galaxy(redshift=1.5, bulge=al.lp.Sersic(intensity=0.7)), + ], + fields=[al.MassField(redshift=0.5, shear=shear, mass_sheet=mass_sheet)], + ) + + file_path = interop_coolest.to_coolest( + galaxies=tracer, file_path=str(tmp_path / "template"), shape_native=(10, 10) + ) + tracer_back = interop_coolest.from_coolest(file_path=file_path) + + assert len(tracer_back.fields) == 1 + assert all(not isinstance(galaxy, al.MassField) for galaxy in tracer_back.galaxies) + + profiles = { + type(profile).__name__: profile + for profile in tracer_back.fields[0].cls_list_from(cls=al.mp.MassProfile) + } + + assert profiles["ExternalShear"].gamma_1 == pytest.approx(shear.gamma_1, abs=1e-12) + assert profiles["ExternalShear"].gamma_2 == pytest.approx(shear.gamma_2, abs=1e-12) + assert profiles["MassSheet"].kappa == pytest.approx(mass_sheet.kappa, abs=1e-12) + + deflections = tracer.deflections_yx_2d_from(grid=grid()) + deflections_back = tracer_back.deflections_yx_2d_from(grid=grid()) + + assert np.asarray(deflections_back) == pytest.approx( + np.asarray(deflections), rel=1e-6, abs=1e-10 + ) + + +def test__to_coolest__a_field_with_an_unsupported_profile_raises_or_is_recorded( + tmp_path, +): + tracer = al.Tracer( + galaxies=lens_model_galaxies(), + fields=[ + al.MassField( + redshift=0.5, + potential=al.mp.ExternalPotential(gamma_1=0.01, tau_1=0.02), + ) + ], + ) + + with pytest.raises((al.exc.ProfileException, ValueError)): + interop_coolest.to_coolest( + galaxies=tracer, + file_path=str(tmp_path / "template_raise"), + shape_native=(10, 10), + ) + + file_path = interop_coolest.to_coolest( + galaxies=tracer, + file_path=str(tmp_path / "template_skip"), + shape_native=(10, 10), + on_unsupported="skip", + ) + + with open(file_path) as f: + template = json.load(f) + + skipped = template["meta"]["skipped_profiles"] + + assert {"galaxy": "field_0", "profile": "ExternalPotential"} in skipped diff --git a/test_autolens/lens/test_los.py b/test_autolens/lens/test_los.py index 85176e59c..90befd52e 100644 --- a/test_autolens/lens/test_los.py +++ b/test_autolens/lens/test_los.py @@ -108,3 +108,117 @@ def test__galaxies_from__test_mode_caps_halos_per_plane(monkeypatch): assert len(counts) > 0 assert all(count <= 3 for count in counts.values()) + + +def _test_mode_sampler(monkeypatch): + monkeypatch.setenv("PYAUTO_TEST_MODE", "2") + + cosmology = Planck15() + _, centres = los.los_planes_from( + z_lens=0.5, z_source=1.0, planes_before_lens=4, planes_after_lens=4 + ) + n_planes = len(centres) + mf, mc = _approx_coefficients(n_planes) + + sampler = los.LOSSampler( + z_lens=0.5, + z_source=1.0, + planes_before_lens=4, + planes_after_lens=4, + m_min=1e7, + m_max=1e10, + cone_radius_arcsec=5.0, + c_scatter=0.15, + truncation_factor=100.0, + cosmology=cosmology, + mass_function_coefficients=mf, + mass_concentration_coefficients=mc, + seed=42, + ) + + return sampler, n_planes + + +def test__fields_from__sheets_are_mass_fields_with_the_legacy_kappa_per_plane( + monkeypatch, +): + """ + A negative-kappa sheet is external mass, not a galaxy. `fields_from` returns exactly the + sheets `galaxies_from()` has always returned — same redshift, same kappa, same order — + wrapped in a `MassField` instead of a `Galaxy`. + """ + sampler, n_planes = _test_mode_sampler(monkeypatch) + + galaxies = sampler.galaxies_from() + fields = sampler.fields_from() + + sheet_galaxies = [ + g + for g in galaxies + if hasattr(g, "mass_sheet") and isinstance(g.mass_sheet, al.mp.MassSheet) + ] + + assert len(fields) == n_planes + assert all(isinstance(field, al.MassField) for field in fields) + + for field, sheet_galaxy in zip(fields, sheet_galaxies): + assert field.redshift == sheet_galaxy.redshift + assert field.mass_sheet.kappa == pytest.approx( + sheet_galaxy.mass_sheet.kappa, rel=1e-12 + ) + + +def test__galaxies_from__sheets_as_fields_returns_halos_only(monkeypatch): + sampler, n_planes = _test_mode_sampler(monkeypatch) + + everything = sampler.galaxies_from() + halos_only = sampler.galaxies_from(sheets_as_fields=True) + + assert len(halos_only) == len(everything) - n_planes + assert all( + hasattr(g, "mass") and isinstance(g.mass, al.mp.NFWTruncatedSph) + for g in halos_only + ) + + # The test-mode halo cap still applies on this path. + counts = {} + for g in halos_only: + counts[g.redshift] = counts.get(g.redshift, 0) + 1 + + assert len(counts) > 0 + assert all(count <= 3 for count in counts.values()) + + +def test__galaxies_and_fields_from__is_the_pair_the_tracer_takes(monkeypatch): + sampler, n_planes = _test_mode_sampler(monkeypatch) + + galaxies, fields = sampler.galaxies_and_fields_from() + + assert galaxies == sampler.galaxies_from(sheets_as_fields=True) + assert len(fields) == n_planes + + tracer = al.Tracer( + galaxies=galaxies + + [ + al.Galaxy(redshift=0.5, mass=al.mp.IsothermalSph(einstein_radius=1.0)), + al.Galaxy(redshift=1.0, bulge=al.lp.SersicSph(intensity=1.0)), + ], + fields=fields, + ) + + assert len(tracer.fields) == n_planes + assert all(not isinstance(galaxy, al.MassField) for galaxy in tracer.galaxies) + + +def test__the_population_is_sampled_once_and_shared(monkeypatch): + """ + `galaxies_from()` and `fields_from()` must describe the same line of sight, so the + sample is cached on the instance rather than redrawn per call. + """ + sampler, _ = _test_mode_sampler(monkeypatch) + + first = sampler.galaxies_from() + second = sampler.galaxies_from() + + assert [g.redshift for g in first] == [g.redshift for g in second] + assert first[0] is second[0] diff --git a/test_autolens/lens/test_tracer_fields.py b/test_autolens/lens/test_tracer_fields.py new file mode 100644 index 000000000..4c7de9435 --- /dev/null +++ b/test_autolens/lens/test_tracer_fields.py @@ -0,0 +1,532 @@ +""" +Tests for `Tracer(fields=...)` — the `MassField` slot on the tracer. + +A `MassField` is the redshift-bearing container for the mass of everything *outside* the +modelled system (an external shear, a mass sheet, an external potential). It is not a +`Galaxy`: it has no light, and it never appears in `tracer.galaxies`. + +The contract these tests pin down is a single sentence: **a field is a member of the +lensing system for every mass calculation and invisible to every galaxy-listing surface.** +So `tracer.fields=[MassField(shear=...)]` must give bit-for-bit the same deflections, +convergence, potential and traced grids as attaching that same shear to the lens galaxy +(the galaxy-attached form, which remains fully supported), while `galaxy_image_2d_dict_from` +keys and `extract_attributes_of_galaxies` stay galaxies-only. +""" + +import os +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest + +import autofit as af +import autolens as al + +from autonerves.dictable import from_json, output_to_json + +from autolens.lens.tracer import MultiPlaneRedshiftWarning + + +@pytest.fixture(name="grid") +def make_grid(): + return al.Grid2D.uniform(shape_native=(10, 10), pixel_scales=0.1) + + +def _shear(): + return al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.06) + + +def _lens(with_shear: bool = False): + kwargs = dict(redshift=0.5, mass=al.mp.Isothermal(einstein_radius=1.0)) + if with_shear: + kwargs["shear"] = _shear() + return al.Galaxy(**kwargs) + + +def _source(): + return al.Galaxy(redshift=1.0, bulge=al.lp.Sersic(intensity=1.0)) + + +# ====================================================================================== +# storage + the galaxies list is untouched +# ====================================================================================== + + +def test__fields_default_to_an_empty_list(): + tracer = al.Tracer(galaxies=[_lens(), _source()]) + + assert tracer.fields == [] + + +def test__fields_are_stored_and_galaxies_are_unchanged(): + field = al.MassField(redshift=0.5, shear=_shear()) + lens, source = _lens(), _source() + + tracer = al.Tracer(galaxies=[lens, source], fields=[field]) + + assert tracer.fields == [field] + assert list(tracer.galaxies) == [lens, source] + assert all(not isinstance(galaxy, al.MassField) for galaxy in tracer.galaxies) + + +def test__fields_is_the_last_positional_argument__so_tracer_galaxies_cosmology_still_works(): + """`Tracer(galaxies, cosmology)` is a positional call in the wild; it must keep working.""" + cosmology = al.cosmo.Planck15() + + tracer = al.Tracer([_lens(), _source()], cosmology) + + assert tracer.cosmology is cosmology + assert tracer.fields == [] + + +def test__members_is_galaxies_then_fields__and_members_ascending_redshift_sorts(): + field = al.MassField(redshift=0.75, shear=_shear()) + lens, source = _lens(), _source() + + tracer = al.Tracer(galaxies=[source, lens], fields=[field]) + + assert tracer.members == [source, lens, field] + assert tracer.members_ascending_redshift == [lens, field, source] + # `galaxies_ascending_redshift` keeps its galaxies-only meaning. + assert tracer.galaxies_ascending_redshift == [lens, source] + + +# ====================================================================================== +# lensing quantities: a field equals the same profile attached to the lens galaxy +# ====================================================================================== + + +def test__traced_grid_2d_list_from__field_shear_equals_galaxy_attached_shear(grid): + via_field = al.Tracer( + galaxies=[_lens(), _source()], + fields=[al.MassField(redshift=0.5, shear=_shear())], + ) + via_galaxy = al.Tracer(galaxies=[_lens(with_shear=True), _source()]) + + for grid_field, grid_galaxy in zip( + via_field.traced_grid_2d_list_from(grid=grid), + via_galaxy.traced_grid_2d_list_from(grid=grid), + ): + assert grid_field == pytest.approx(np.array(grid_galaxy), 1.0e-8) + + +@pytest.mark.parametrize( + "method", ["deflections_yx_2d_from", "convergence_2d_from", "potential_2d_from"] +) +def test__lensing_quantities__field_shear_equals_galaxy_attached_shear(grid, method): + via_field = al.Tracer( + galaxies=[_lens(), _source()], + fields=[al.MassField(redshift=0.5, shear=_shear())], + ) + via_galaxy = al.Tracer(galaxies=[_lens(with_shear=True), _source()]) + + assert np.array(getattr(via_field, method)(grid=grid)) == pytest.approx( + np.array(getattr(via_galaxy, method)(grid=grid)), 1.0e-8 + ) + + +def test__lensing_quantities__a_field_actually_changes_the_answer(grid): + """Guards against the tests above passing because the field is being ignored.""" + without = al.Tracer(galaxies=[_lens(), _source()]) + with_field = al.Tracer( + galaxies=[_lens(), _source()], + fields=[al.MassField(redshift=0.5, shear=_shear())], + ) + + assert not np.allclose( + np.array(without.deflections_yx_2d_from(grid=grid)), + np.array(with_field.deflections_yx_2d_from(grid=grid)), + ) + + # A `MassSheet` has convergence where a shear does not, so this pins the convergence path too. + with_sheet = al.Tracer( + galaxies=[_lens(), _source()], + fields=[al.MassField(redshift=0.5, mass_sheet=al.mp.MassSheet(kappa=0.1))], + ) + + assert not np.allclose( + np.array(without.convergence_2d_from(grid=grid)), + np.array(with_sheet.convergence_2d_from(grid=grid)), + ) + + +def test__potential_2d_from__external_potential_field_contributes(grid): + field = al.MassField( + redshift=0.5, potential=al.mp.ExternalPotential(gamma_1=0.05, tau_1=0.02) + ) + + without = al.Tracer(galaxies=[_lens(), _source()]) + with_field = al.Tracer(galaxies=[_lens(), _source()], fields=[field]) + + assert not np.allclose( + np.array(without.potential_2d_from(grid=grid)), + np.array(with_field.potential_2d_from(grid=grid)), + ) + + +# ====================================================================================== +# planes +# ====================================================================================== + + +def test__a_field_at_a_galaxy_redshift_joins_that_plane(): + field = al.MassField(redshift=0.5, shear=_shear()) + lens, source = _lens(), _source() + + tracer = al.Tracer(galaxies=[lens, source], fields=[field]) + + assert tracer.plane_redshifts == [0.5, 1.0] + assert tracer.total_planes == 2 + assert list(tracer.planes[0]) == [lens, field] + assert list(tracer.planes[1]) == [source] + + +def test__a_field_at_its_own_redshift_is_a_plane_of_its_own(): + field = al.MassField(redshift=0.75, mass_sheet=al.mp.MassSheet(kappa=0.1)) + lens, source = _lens(), _source() + + tracer = al.Tracer(galaxies=[lens, source], fields=[field]) + + assert tracer.plane_redshifts == [0.5, 0.75, 1.0] + assert tracer.total_planes == 3 + assert list(tracer.planes[1]) == [field] + + +def test__a_field_only_plane_is_ray_traced__grid_list_has_an_entry_for_it(grid): + tracer = al.Tracer( + galaxies=[_lens(), _source()], + fields=[ + al.MassField(redshift=0.75, mass=al.mp.IsothermalSph(einstein_radius=0.5)) + ], + ) + + traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid) + + assert len(traced_grid_list) == 3 + + +def test__plane_index_via_redshift_from__finds_a_field_only_plane(): + tracer = al.Tracer( + galaxies=[_lens(), _source()], + fields=[al.MassField(redshift=0.75, shear=_shear())], + ) + + assert tracer.plane_index_via_redshift_from(redshift=0.75) == 1 + + +# ====================================================================================== +# has / cls_list_from / image lists +# ====================================================================================== + + +def test__has__is_true_for_a_profile_held_only_by_a_field(): + tracer = al.Tracer( + galaxies=[_lens(), _source()], + fields=[al.MassField(redshift=0.5, shear=_shear())], + ) + + assert tracer.has(cls=al.mp.ExternalShear) is True + assert ( + al.Tracer(galaxies=[_lens(), _source()]).has(cls=al.mp.ExternalShear) is False + ) + + +def test__cls_list_from__includes_a_fields_profiles(): + shear = _shear() + tracer = al.Tracer( + galaxies=[_lens(), _source()], fields=[al.MassField(redshift=0.5, shear=shear)] + ) + + assert shear in tracer.cls_list_from(cls=al.mp.ExternalShear) + assert len(tracer.cls_list_from(cls=al.mp.MassProfile)) == 2 + + +def test__image_2d_list_from__is_unaffected_by_a_field(grid): + without = al.Tracer(galaxies=[_lens(), _source()]) + with_field = al.Tracer( + galaxies=[_lens(), _source()], + fields=[al.MassField(redshift=0.5, shear=_shear())], + ) + + # The image of a plane holding a field is the image of its galaxies alone: the field's + # contribution is zeros, which is exactly what makes it safe to sum over plane members. + assert len(with_field.image_2d_list_from(grid=grid)) == 2 + assert np.array(with_field.image_2d_list_from(grid=grid)[0]) == pytest.approx( + np.array(without.image_2d_list_from(grid=grid)[0]), 1.0e-8 + ) + + +def test__upper_plane_index_with_light_profile__a_field_only_top_plane_does_not_raise_it(): + tracer = al.Tracer( + galaxies=[_lens(), _source()], + fields=[al.MassField(redshift=2.0, mass_sheet=al.mp.MassSheet(kappa=0.1))], + ) + + assert tracer.upper_plane_index_with_light_profile == 1 + + +# ====================================================================================== +# per-galaxy surfaces stay galaxies-only +# ====================================================================================== + + +def test__galaxy_image_2d_dict_from__keys_are_galaxies_only(grid): + lens, source = _lens(), _source() + field = al.MassField(redshift=0.5, shear=_shear()) + + tracer = al.Tracer(galaxies=[lens, source], fields=[field]) + + assert list(tracer.galaxy_image_2d_dict_from(grid=grid).keys()) == [lens, source] + + +def test__extract_attributes_of_galaxies__stays_galaxies_only(): + tracer = al.Tracer( + galaxies=[_lens(), _source()], + fields=[al.MassField(redshift=0.5, shear=_shear())], + ) + + centres = tracer.extract_attributes_of_galaxies( + cls=al.mp.MassProfile, attr_name="centre" + ) + + assert len(centres) == 2 + + +def test__extract_attribute__is_member_aware(): + """The whole-tracer attribute list is of *profiles*, so a field's profiles count.""" + tracer = al.Tracer( + galaxies=[_lens(), _source()], + fields=[al.MassField(redshift=0.5, shear=_shear())], + ) + + assert len(tracer.extract_attribute(cls=al.mp.MassProfile, attr_name="centre")) == 2 + + +def test__extract_profile_and_plane_index__find_a_fields_component(): + field = al.MassField(redshift=0.5, shear=_shear()) + tracer = al.Tracer(galaxies=[_lens(), _source()], fields=[field]) + + assert tracer.extract_profile(profile_name="shear") is field.shear + assert tracer.extract_plane_index_of_profile(profile_name="shear") == 0 + + +# ====================================================================================== +# FitImaging +# ====================================================================================== + + +def test__fit_imaging__galaxy_dict_and_plane_images_match_the_galaxy_attached_form( + masked_imaging_7x7, +): + lens, source = _lens(), _source() + + fit_field = al.FitImaging( + dataset=masked_imaging_7x7, + tracer=al.Tracer( + galaxies=[lens, source], fields=[al.MassField(redshift=0.5, shear=_shear())] + ), + ) + fit_galaxy = al.FitImaging( + dataset=masked_imaging_7x7, + tracer=al.Tracer(galaxies=[_lens(with_shear=True), _source()]), + ) + + assert list(fit_field.galaxy_model_image_dict.keys()) == [lens, source] + + for image_field, image_galaxy in zip( + fit_field.model_images_of_planes_list, fit_galaxy.model_images_of_planes_list + ): + assert np.array(image_field) == pytest.approx(np.array(image_galaxy), 1.0e-8) + + assert fit_field.log_likelihood == pytest.approx(fit_galaxy.log_likelihood, 1.0e-8) + + +def test__fit_imaging__a_field_only_plane_does_not_break_the_plane_image_list( + masked_imaging_7x7, +): + fit = al.FitImaging( + dataset=masked_imaging_7x7, + tracer=al.Tracer( + galaxies=[_lens(), _source()], + fields=[al.MassField(redshift=0.75, shear=_shear())], + ), + ) + + assert len(fit.model_images_of_planes_list) == 3 + assert np.array(fit.model_images_of_planes_list[1]) == pytest.approx(0.0, 1.0e-8) + + +# ====================================================================================== +# sliced_tracer_from +# ====================================================================================== + + +def test__sliced_tracer_from__snaps_a_fields_redshift_to_the_nearest_plane(): + lens_g0 = al.Galaxy(redshift=0.5) + source_g0 = al.Galaxy(redshift=2.0) + field = al.MassField(redshift=0.42, shear=_shear()) + + tracer = al.Tracer.sliced_tracer_from( + lens_galaxies=[lens_g0], + line_of_sight_galaxies=[], + source_galaxies=[source_g0], + planes_between_lenses=[1, 1], + cosmology=al.cosmo.Planck15(), + fields=[field], + ) + + assert field.redshift == 0.5 + assert tracer.fields == [field] + assert field in list(tracer.planes[tracer.plane_index_via_redshift_from(0.5)]) + + +# ====================================================================================== +# json round trip +# ====================================================================================== + + +def test__output_to_and_load_from_json__rebuilds_the_fields(): + json_file = Path(__file__).resolve().parent / "files" / "tracer_fields.json" + + tracer = al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, mass_profile=al.mp.IsothermalSph(einstein_radius=1.0) + ), + al.Galaxy(redshift=1.0), + ], + cosmology=al.cosmo.Planck15(), + fields=[al.MassField(redshift=0.5, shear=_shear())], + ) + + output_to_json(tracer, file_path=json_file) + + tracer_from_json = from_json(file_path=json_file) + + assert len(tracer_from_json.fields) == 1 + assert isinstance(tracer_from_json.fields[0], al.MassField) + assert tracer_from_json.fields[0].redshift == 0.5 + assert tracer_from_json.fields[0].shear.gamma_1 == pytest.approx(0.05, 1.0e-12) + assert tracer_from_json.fields[0].shear.gamma_2 == pytest.approx(0.06, 1.0e-12) + + +# ====================================================================================== +# validation +# ====================================================================================== + + +def test__fields_a_string_is_rejected_even_though_a_string_is_iterable(): + with pytest.raises(TypeError, match="MassField"): + al.Tracer(galaxies=[_lens()], fields="not a list") + + +@pytest.mark.parametrize("fields", [42, {"a": 1}, 1.5, b"bytes"]) +def test__fields_non_iterable_and_wrong_container_inputs_are_rejected(fields): + with pytest.raises(TypeError, match="fields"): + al.Tracer(galaxies=[_lens()], fields=fields) + + +def test__fields_a_list_entry_without_a_redshift_is_rejected_and_names_the_index(): + with pytest.raises(TypeError, match="index 1"): + al.Tracer(galaxies=[_lens()], fields=[al.MassField(redshift=0.5), "nope"]) + + +def test__fields_a_galaxy_is_rejected_with_a_message_saying_where_galaxies_go(): + with pytest.raises(TypeError, match="Galaxies go in the galaxies argument"): + al.Tracer(galaxies=[_lens()], fields=[_source()]) + + +def test__fields_a_model_instance_is_accepted(): + model = af.Collection( + fields=af.Collection( + field=af.Model( + al.MassField, redshift=0.5, shear=af.Model(al.mp.ExternalShear) + ) + ) + ) + instance = model.instance_from_prior_medians() + + tracer = al.Tracer(galaxies=[_lens(), _source()], fields=instance.fields) + + assert len(tracer.fields) == 1 + assert isinstance(tracer.fields[0], al.MassField) + + +# ====================================================================================== +# MultiPlaneRedshiftWarning +# ====================================================================================== + + +def test__warning__a_tracer_whose_only_mass_is_a_field_behind_all_light_warns(): + with pytest.warns(MultiPlaneRedshiftWarning): + al.Tracer( + galaxies=[al.Galaxy(redshift=0.5, bulge=al.lp.Sersic(intensity=1.0))], + fields=[al.MassField(redshift=1.0, shear=_shear())], + ) + + +def test__warning__a_field_at_the_lens_redshift_with_light_behind_stays_quiet(recwarn): + al.Tracer( + galaxies=[al.Galaxy(redshift=0.5), _source()], + fields=[al.MassField(redshift=0.5, shear=_shear())], + ) + + assert [w for w in recwarn if w.category is MultiPlaneRedshiftWarning] == [] + + +# ====================================================================================== +# identifier pin +# ====================================================================================== + + +_IDENTIFIER_PIN_SNIPPET = """ +import autofit as af +import autolens as al + +model = af.Collection( + galaxies=af.Collection( + lens=af.Model( + al.Galaxy, + redshift=0.5, + mass=af.Model(al.mp.Isothermal), + shear=af.Model(al.mp.ExternalShear), + ), + source=af.Model(al.Galaxy, redshift=1.0, bulge=af.Model(al.lp.Sersic)), + ) +) + +print(model.identifier) +""" + + +def test__regression__the_model_identifier_is_unchanged_by_the_fields_slot(): + """ + The identifier of a lens/source model with no fields must not move when the `fields` + slot is added to `Tracer`: a changed identifier silently orphans every existing + results folder on disk. + + The model is the one in `_IDENTIFIER_PIN_SNIPPET` above — a lens `Galaxy` with an + `Isothermal` and an `ExternalShear` at z=0.5 and a `Sersic` source at z=1.0 — and the + expected value was captured on PyAutoLens `main` at 223132c, BEFORE any of this work + was written. Never update it to match a new value. + + It runs in a subprocess, deliberately. A model identifier is a function of the model + *and the active PyAutoFit configuration*, and `test_autolens/conftest.py` pushes the + suite's own `test_autolens/config` onto `conf.instance` for every test. The pinned + value was captured under the library's shipped configuration, so the only way to + assert that exact constant is to evaluate it under that configuration — which is what + a fresh interpreter with no conftest gives. + """ + result = subprocess.run( + [sys.executable, "-c", _IDENTIFIER_PIN_SNIPPET], + capture_output=True, + text=True, + cwd=Path(__file__).resolve().parents[2], + env={**os.environ, "PYAUTO_SKIP_WORKSPACE_VERSION_CHECK": "1"}, + ) + + assert result.returncode == 0, result.stderr + + assert result.stdout.strip().splitlines()[-1] == "fef2697b5c32ba56bb18a7baecb7b0f6"