Work Classification

Library (PyAutoLens only). Feature Agent sized it large and suggested splitting; the epic's human ruling (2026-09-17) sized phase 2 as one prompt, so it stays whole. Workspaces follow as phases 3–5.

Affected Repositories

• PyAutoLens (primary). Depends on PyAutoGalaxy main ≥ 33714b80 (ag.MassField); CI runs the same-named library branch, so no release is needed.

Branch Survey

Remote branches: main, dev_Q, two archive/condemned/* — none touch lens/, analysis/ or interop/. No active.md task claims PyAutoLens.

Suggested branch: feature/mass-field-integration
Worktree root (local-dev): ~/Code/PyAutoLabs-wt/mass-field-integration/ (a web session works on its clone instead)

Facts established on main

• Identifier pin, captured before any edit: 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)))) → fef2697b5c32ba56bb18a7baecb7b0f6. Must hold on the branch.
• Tracer.__init__(galaxies, cosmology=None) runs _validate_galaxies (duck-typed on redshift) and _warn_if_no_light_is_behind_any_mass, then stores self.galaxies. galaxies_ascending_redshift sorts unless any redshift is traced; plane_redshifts / planes call tracer_util.plane_redshifts_from / planes_from on that list — both are duck-typed on .redshift and .has, so a MassField member works there unchanged.
• Tracer has no custom to_dict; autonerves.dictable serialises __init__ arguments, so a fields argument stored as self.fields round-trips once it exists (test_tracer.py already has from_json/output_to_json tests to mirror).
• Tracer.galaxy_image_2d_dict_from iterates planes (for galaxies in self.planes: for galaxy in galaxies), so without a filter a field becomes a dict key with a zeros image and reaches FitImaging.galaxy_model_image_dict, the galaxies subplots and model_images_of_planes_list. The prompt's assumption that per-galaxy surfaces only walk tracer.galaxies is false here — handle it (step 1e).
• autolens/jax/registration.py walks tracer.galaxies and registers Galaxy with no_flatten=("redshift",) (unlike the autogalaxy walker); it is a deliberate copy, so patch it, do not share.
• AnalysisLens.tracer_via_instance_from (autolens/analysis/analysis/lens.py:83) folds extra_galaxies / scaling_galaxies; AnalysisImaging, AnalysisInterferometer (via AnalysisDataset) and AnalysisPoint all inherit it, so one edit covers the three. autolens/aggregator/tracer.py:~90 and autolens/analysis/result.py:254 build tracers independently.
• COOLEST (autolens/interop/coolest.py): to_coolest accepts a Tracer or a list, peels ExternalShear / MassSheet off each galaxy into lazy.MassField(f"mass_field_{i}", z, mass_model=...); from_coolest turns every entity, MassField included, into a Galaxy. ExternalPotential is not in the peel today and has no COOLEST profile.
• LOS: LOSSampler.galaxies_from returns one list: ag.Galaxy(redshift=z_cen, mass=halo) per halo (line ~761) plus ag.Galaxy(redshift=..., mass_sheet=ag.mp.MassSheet(kappa=kappa_neg)) per plane (line ~789). Only test_autolens/lens/test_los.py calls it in-repo; the workspace los_halos scripts call it too (outside this repo — keep the shim).
• al.model_util is autolens/analysis/model_util.py, importing mge_model_from etc. from autogalaxy; al.MassField is not yet re-exported.

Implementation Steps

1. Tracer(galaxies, fields=None, cosmology=None) — autolens/lens/tracer.py
   a. Signature gains fields: Optional[List[ag.MassField]] = None; self.fields = list(fields) if fields is not None else []. Validate with a _validate_fields mirroring _validate_galaxies (list/tuple or af.ModelInstance; each entry has redshift; messages name MassField, e.g. fields must be an iterable of MassField objects, e.g. Tracer(galaxies=[lens, source], fields=[field])). Extend _warn_if_no_light_is_behind_any_mass to take galaxies + fields (fields are mass-bearing, never light-bearing).
   b. Add @property members_ascending_redshift (name it clearly; galaxies_ascending_redshift keeps its meaning and return value — galaxies only): sorted(self.galaxies + self.fields, key=redshift) with the same traced-redshift bypass (tracer_util._any_traced). plane_redshifts and planes use it, so a plane (ag.Galaxies) may hold a MassField. Audit every other use of galaxies_ascending_redshift in the file (grid_2d_at_redshift_from ~582, time_delays_from ~1034, plane_index_via_redshift_from) and switch those that compute lensing quantities to the members list; leave galaxy-only surfaces on galaxies.
   c. total_planes follows from plane_redshifts, so a field-only plane counts. upper_plane_index_with_light_profile, image_2d_list_from, traced_grid_2d_list_from, deflections_*, convergence_2d_from, potential_2d_from, time_delays_from, has, cls_list_from, perform_inversion, plane_indexes_with_pixelizations, the positions solver, magnification and critical curves all flow through planes — do not change them, but write one test per surface proving a field contributes (mass) or is invisible (light).
   d. sliced_tracer_from(lens_galaxies, line_of_sight_galaxies, source_galaxies, planes_between_lenses, cosmology=None, fields=None): snap each field's redshift to the nearest sliced plane exactly like the line-of-sight galaxies, and pass fields= through to the constructor.
   e. Galaxy-only per-galaxy surfaces. In galaxy_image_2d_dict_from, skip members that are not ag.Galaxy (isinstance(galaxy, ag.MassField) → continue) so fields never become dict keys. Then audit every consumer that walks plane members and indexes that dict: autolens/imaging/fit_imaging.py ~306–330 (model_images_of_planes_list loops for galaxy in galaxies over planes), galaxy_model_image_dict, subplot_galaxies_images / the fit plotters, lens/to_inversion.py (uses tracer.galaxies, fine), the aggregator galaxy tables, extract_attributes_of_galaxies. Each one either skips MassField members or already only reads tracer.galaxies. Write the assertion test: fit a Tracer(galaxies=[lens, source], fields=[field]) through FitImaging and check galaxy_model_image_dict.keys() == {lens, source} and the plane images equal those of the galaxy-attached equivalent.
   f. Dict round trip: output_to_json / from_json of a tracer with fields rebuilds tracer.fields as al.MassField instances (mirror the existing test).
   g. JAX — autolens/jax/registration.py: register MassField with no_flatten=("redshift",) beside Galaxy (same reasoning: plane bookkeeping compares redshifts), walk tracer.fields after tracer.galaxies, and extend the docstring. Probe: the three-way check (NumPy / xp=jnp / jax.jit with the tracer traced and the grid closed over, as in PyAutoGalaxy#621's probe) on traced_grid_2d_list_from with a field present; unit tests stay NumPy-only.
   h. MultiPlaneRedshiftWarning: a tracer whose only mass is a field behind all light must warn exactly as a mass galaxy would; a field at the lens redshift with light behind must not.

2. Analysis fields slot — autolens/analysis/analysis/lens.py tracer_via_instance_from: after the scaling_galaxies fold, fields = list(instance.fields) if getattr(instance, "fields", None) is not None else None; Tracer(galaxies=galaxy_list, fields=fields, cosmology=cosmology). Same fold in autolens/aggregator/tracer.py (_tracer_from / the list builder ~line 90) and in autolens/analysis/result.py:254 only if max_log_likelihood_galaxies sits beside a max_log_likelihood_fields (add that property on Result if the instance carries fields, else leave the positions helper galaxy-only and say so). The subhalo perturb path is unchanged. Test: an AnalysisImaging fit of af.Collection(galaxies=..., fields=af.Collection(field=af.Model(al.MassField, redshift=0.5, shear=af.Model(al.mp.ExternalShear)))) under PYAUTO_TEST_MODE=2 (or the conftest's test-mode fixture) produces a Tracer whose fields has one MassField and whose log likelihood equals the galaxy-attached model's at the same parameter values.

3. COOLEST — autolens/interop/coolest.py:

   • to_coolest: when given a Tracer, also take fields = list(galaxies.fields); after the galaxy loop, each field becomes one lazy.MassField(f"field_{j}", float(field.redshift), mass_model=...) carrying all its profiles (ExternalShear → ExternalShear, MassSheet → ConvergenceSheet via coolest_dict_from_mass); an ExternalPotential or any other unsupported profile goes through the existing on_unsupported branch naming the profile (_record_skipped with the field's name). The legacy peel of galaxy-attached sheets is untouched — write a byte-identity test: export a galaxy-attached tracer on this branch and compare to a fixture JSON captured on main (or the existing round-trip tests plus a stored expected file), and keep autolens_workspace_test's coolest_round_trip.py as the downstream regression.
   • from_coolest: a COOLEST MassField entity becomes al.MassField(redshift=..., **{f"mass_{i}": ...}) appended to a fields list, and the returned tracer is Tracer(galaxies=galaxies, fields=fields, cosmology=cosmology). Update the docstring (it currently says every MassField entity becomes a galaxy). Test: round trip of a tracer with fields=[MassField(shear, mass_sheet)] returns a MassField with parameters equal to 1e-12; a legacy galaxy-attached export now imports as a field too (document this as the one behaviour change on import, and say so in the PR body).

4. LOS sampler — autolens/lens/los.py: keep galaxies_from() returning the halo galaxies and, for compatibility, add fields_from() returning one ag.MassField(redshift=z_cen, mass_sheet=ag.mp.MassSheet(kappa=kappa_neg)) per plane; refactor the shared sampling into a private helper so the two do not sample twice (cache the sampled population on the instance, or return both from one galaxies_and_fields_from() -> Tuple[List[Galaxy], List[MassField]]). Shim: galaxies_from() keeps returning halos + sheet galaxies exactly as today unless called with sheets_as_fields=True (default False) — name the workspace los_halos scripts as the callers the shim protects, and note in the PR that phase 3 flips them. Test in test_los.py: the new return's sheets are MassFields with the same kappa per plane as the legacy sheet galaxies; the test-mode halo cap still applies.

5. model_util.mass_field_from — autolens/analysis/model_util.py: mass_field_from(lens: af.Model, shear: bool = True, mass_sheet: bool = False, potential: bool = False, redshift: Optional[float] = None) -> af.Model. Raise ValueError naming the galaxy if lens has no mass; redshift defaults to lens.redshift; build af.Model(al.MassField, redshift=..., shear=af.Model(al.mp.ExternalShear)) etc.; tie field.potential.centre = lens.mass.centre and field.mass_sheet.centre = lens.mass.centre as shared priors (assignment of the same prior object, so is holds). Docstring: multi-deflector callers name the primary galaxy; the result goes in fields=af.Collection(field=...). Tests in test_autolens/analysis/test_model_util.py (create it if absent): each flag, the is identity of the tied centres, the no-mass error, the redshift override, and that the composed collection's prior_count matches expectation.

6. Exports + docs — autolens/__init__.py: from autogalaxy.galaxy.mass_field import MassField beside Galaxy; docs/api/galaxy.rst lists MassField and documents Tracer.fields; docs/api/modeling.rst (or wherever model_util lives) gains mass_field_from. A short "External fields" paragraph in the overview page that introduces the tracer, stating the galaxy-attached form remains supported. No PyAutoMemory citations.

7. Tests + gate — pytest test_autolens -q -n auto green; the identifier pin as a regression test; git diff --name-only shows nothing under autolens/config/; no warnings.warn added for legacy forms.

Key Files

• autolens/lens/tracer.py — constructor, validation, ascending members, planes, sliced_tracer_from, galaxy_image_2d_dict_from.
• autolens/lens/tracer_util.py — read-only; duck-typed plane builders.
• autolens/imaging/fit_imaging.py — plane-image loop must skip fields.
• autolens/analysis/analysis/lens.py, autolens/aggregator/tracer.py, autolens/analysis/result.py — tracer builders.
• autolens/interop/coolest.py — export/import.
• autolens/lens/los.py — sampler return shape.
• autolens/analysis/model_util.py — helper.
• autolens/jax/registration.py — walker.
• autolens/__init__.py, docs/api/galaxy.rst, docs/api/modeling.rst.
• Tests: test_autolens/lens/test_tracer.py, test_tracer_validation.py, test_los.py, test_autolens/interop/test_coolest.py, test_autolens/analysis/analysis/test_analysis_lens.py, new test_autolens/analysis/test_model_util.py.

Hard constraints

• Legacy galaxy-attached sheets and shear-only galaxies in galaxies: no warning, no behaviour change, COOLEST export byte-identical.
• tracer.galaxies never contains a field; per-galaxy surfaces never see one.
• No prior config changes; identifier pin holds.

Out of scope

Workspace adoption (guides/profiles/mass.py, multi_galaxy/, group/, siblings) is phases 3–5 behind a release.

MassField in PyAutoLens: Tracer(fields=), the analysis fields slot, COOLEST 1:1, LOS sheets, model_util helper

Type: feature
Target: PyAutoLens
Repos:

• PyAutoLens
  Themes:
• cluster
  Difficulty: large
  Autonomy: supervised
  Priority: normal
  Consequence: judge
  Witness: al.Tracer(galaxies=[lens, source], fields=[al.MassField(redshift=0.5, shear=...)]).traced_grid_2d_list_from(grid) is np.allclose to the same shear on lens; tracer.galaxies is [lens, source] (no field in it); a model with fields=af.Collection(field=af.Model(al.MassField, ...)) fits through AnalysisImaging under PYAUTO_TEST_MODE=2; COOLEST round trip of that tracer returns an al.MassField with the same parameters to 1e-12 and a legacy al.Galaxy(..., shear=...) still exports byte-identically; al.model_util.mass_field_from(lens=lens_model, potential=True).potential.centre is lens_model.mass.centre; identifier pin (galaxy-attached model, equal on main and branch); pytest test_autolens -q -n auto green.
  Review-minutes: 10
  Unattended: ready
  Epic: mass-field
  Phase: 2
  Filed: 2026-09-17

Gate cleared 2026-09-17: phase 1 (complete/2026/09/mass-field-class.md,
PyAutoGalaxy#621) is merged on PyAutoGalaxy main; CI runs the same-named
library branch, so main is enough and a release is not required for this
phase.

Teach PyAutoLens about ag.MassField (phase 1 of the epic
draft/feature/autogalaxy/mass_field_epic.md; read it first). A field is its
own thing: it gets its own tracer argument and its own model slot, and only the
tracer's planes merge it with the galaxies. tracer.galaxies never holds a
field, so every script that indexes galaxies positionally keeps working.

What to build

1. Tracer(galaxies, fields=None, cosmology=...) — @PyAutoLens/autolens/lens/tracer.py:
   • store self.fields (a list; [] when None); validate members on
     redshift the way _validate_galaxies does, with a message naming
     MassField.
   • galaxies_ascending_redshift, plane_redshifts and planes are built
     over galaxies and fields; a plane (an ag.Galaxies) may therefore
     hold a MassField, which phase 1 made safe. Everything downstream
     (traced_grid_2d_list_from, deflections_yx_2d_from,
     convergence_2d_from, potential_2d_from, has, cls_list_from, the
     positions solver, magnification, critical curves) flows through planes
     and needs no change — verify each with a test rather than assume.
   • tracer.galaxies is unchanged; tracer.fields is the new accessor;
     total_planes counts field-only planes (a LOS sheet plane is a plane).
   • sliced_tracer_from places fields in the planes nearest their redshift
     like line_of_sight_galaxies; to_dict/from_dict carries fields;
     the JAX pytree flatten of Tracer includes fields
     (@PyAutoLens/autolens/jax/registration.py walks tracer.galaxies — walk
     tracer.fields too and register MassField; the walker is a copy of the
     autogalaxy one by a recorded human decision, so patch both, do not share).
   • MultiPlaneRedshiftWarning: fields are mass-bearing for the light/mass
     ordering check.
   • The per-galaxy surfaces (FitImaging.galaxy_model_image_dict,
     subplot_galaxies_images, the aggregator's galaxy tables) iterate
     tracer.galaxies and therefore never see a field — assert that in a test.
2. Analysis fields slot — @PyAutoLens/autolens/analysis/analysis/lens.py
tracer_via_instance_from: beside the extra_galaxies / scaling_galaxies
   folds, fields = list(instance.fields) if getattr(instance, "fields", None) is not None else None
   and pass Tracer(galaxies=galaxy_list, fields=fields, cosmology=cosmology).
   Same in the interferometer / point-source analyses if they build their own
   tracer. The subhalo perturb path is unchanged.
3. COOLEST — @PyAutoLens/autolens/interop/coolest.py:
   • to_coolest: each tracer.fields entry becomes one COOLEST MassField
     entity carrying all its profiles (ExternalShear → ExternalShear,
     MassSheet → ConvergenceSheet; ExternalPotential has no COOLEST
     profile — route it through the existing on_unsupported path with the
     profile named). Keep the legacy peel: a Galaxy carrying sheets still
     exports them as a separate MassField entity, byte-identical to today's
     output (the workspace_test round-trip script is the regression).
   • from_coolest: a COOLEST MassField entity becomes an al.MassField
     in the returned tracer's fields, not a Galaxy. Update the docstring.
4. Line-of-sight sampler — @PyAutoLens/autolens/lens/los.py builds
   ag.Galaxy(redshift=z_cen, mass_sheet=...) per plane; build ag.MassField
   and return the sheets separately (galaxies_from → halos as galaxies,
   fields_from → sheets, or one call returning both) so the caller writes
   Tracer(galaxies=[lens, source] + halos, fields=sheets). Keep a shim for
   the old single-list return if any caller in the tests or the workspace
   los_halos scripts depends on it, and say which in the PR.
5. model_util helper — beside mge_model_from (the module the workspace
   reaches as al.model_util): mass_field_from(lens, shear=True, mass_sheet=False, potential=False, redshift=None) -> af.Model. Builds
   af.Model(al.MassField, redshift=<lens.redshift unless given>, ...) with
   af.Model(al.mp.ExternalShear) / af.Model(al.mp.MassSheet) /
   af.Model(al.mp.ExternalPotential) as requested, and ties
   potential.centre (and mass_sheet.centre) to lens.mass.centre as a shared
   prior — the human's convention for where an external potential is centred.
   lens is an af.Model of a Galaxy with a mass; if it has none, raise
   naming the galaxy. Document that for a multi-deflector system the caller
   picks the primary galaxy explicitly, and that the result goes in
   fields=af.Collection(field=...).
6. Tests + docs — tracer tests for each bullet of step 1; an analysis
   test fitting a model with fields under test mode; round-trip and
   legacy-peel COOLEST tests; a test_model_util case for each helper flag and
   the shared-prior identity; los tests for the new return; the identifier
   pin; API docs for MassField (re-export in @PyAutoLens/autolens/__init__.py
   beside Galaxy), Tracer.fields and model_util.mass_field_from.

Hard constraints

• Legacy galaxy-attached sheets and shear-only galaxies in galaxies: no
  warning, no behaviour change, COOLEST export unchanged (regression-test it).
• tracer.galaxies never contains a field.
• No prior config changes.