feat: MassField in PyAutoLens — Tracer(fields=), analysis slot, COOLEST, LOS, model_util

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

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)
 

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(

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

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)

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
                 ]

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]

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)

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