From e1e5fcfe918c8054b11a267dee2655bc4ac0feeb Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 18 Sep 2026 10:01:27 +0100 Subject: [PATCH] feat: accept bare MassField in the fields slot --- autolens/aggregator/tracer.py | 6 +- autolens/analysis/analysis/lens.py | 6 +- autolens/analysis/model_util.py | 9 +- autolens/analysis/result.py | 8 +- autolens/lens/tracer.py | 20 ++- autolens/lens/tracer_util.py | 13 ++ docs/api/galaxy.rst | 3 +- docs/api/modeling.rst | 4 + test_autolens/aggregator/test_tracer.py | 41 ++++++ .../analysis/analysis/test_analysis_lens.py | 40 ++++++ test_autolens/analysis/test_result.py | 20 +++ test_autolens/lens/test_tracer_fields.py | 119 ++++++++++++++++++ 12 files changed, 265 insertions(+), 24 deletions(-) create mode 100644 test_autolens/aggregator/test_tracer.py diff --git a/autolens/aggregator/tracer.py b/autolens/aggregator/tracer.py index 404e19264..2d61180d9 100644 --- a/autolens/aggregator/tracer.py +++ b/autolens/aggregator/tracer.py @@ -100,11 +100,7 @@ def _tracer_from( # 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 - ) + fields = tracer_util.fields_list_from(getattr(instance, "fields", None)) tracer = Tracer(galaxies=galaxy_list, cosmology=cosmology, fields=fields) diff --git a/autolens/analysis/analysis/lens.py b/autolens/analysis/analysis/lens.py index 03fc9e6a1..0ca354f72 100644 --- a/autolens/analysis/analysis/lens.py +++ b/autolens/analysis/analysis/lens.py @@ -135,11 +135,7 @@ def tracer_via_instance_from( # 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 - ) + fields = tracer_util.fields_list_from(getattr(instance, "fields", None)) return Tracer( galaxies=galaxy_list, diff --git a/autolens/analysis/model_util.py b/autolens/analysis/model_util.py index ef20d53d7..debb42438 100644 --- a/autolens/analysis/model_util.py +++ b/autolens/analysis/model_util.py @@ -147,6 +147,12 @@ def mass_field_from( fields=af.Collection(field=al.model_util.mass_field_from(lens=lens)), ) + For a single field, ``fields=al.model_util.mass_field_from(lens=lens)`` is also + supported. This shortens prior paths (``fields.shear.gamma_1`` instead of + ``fields.field.shear.gamma_1``), so the two forms have different result identifiers + by design. The collection form above remains the primary example and retains its + existing identifier. + 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 @@ -177,7 +183,8 @@ def mass_field_from( Returns ------- - The ``af.Model`` of the ``MassField``, to be placed in the model's ``fields`` collection. + The ``af.Model`` of the ``MassField``, to be placed in the model's ``fields`` collection + or supplied directly as ``fields`` for a single field. """ if not hasattr(lens, "mass"): raise ValueError( diff --git a/autolens/analysis/result.py b/autolens/analysis/result.py index 13d345e09..770dded66 100644 --- a/autolens/analysis/result.py +++ b/autolens/analysis/result.py @@ -31,6 +31,7 @@ SourceMaxSeparation, ) from autolens.lens.tracer import Tracer +from autolens.lens import tracer_util from autolens.point.solver import PointSolver from autonerves.test_mode import is_test_mode, skip_checks @@ -54,12 +55,7 @@ def max_log_likelihood_fields(self) -> List[ag.MassField]: 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) + return tracer_util.fields_list_from(getattr(self.instance, "fields", None)) def source_plane_light_profile_centre_from( self, plane_redshift: Optional[float] = None diff --git a/autolens/lens/tracer.py b/autolens/lens/tracer.py index 355e99ae3..6246e17b4 100644 --- a/autolens/lens/tracer.py +++ b/autolens/lens/tracer.py @@ -242,7 +242,8 @@ 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 + A bare ``MassField`` is accepted as a one-element collection. Otherwise this + mirrors ``_validate_galaxies`` — 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. @@ -259,6 +260,9 @@ def _validate_fields(fields): if isinstance(fields, af.ModelInstance): return + if isinstance(fields, MassField): + fields = [fields] + if isinstance(fields, (str, bytes)): raise TypeError( f"fields must be an iterable of MassField objects, but a " @@ -297,7 +301,7 @@ def __init__( self, galaxies: Union[List[ag.Galaxy], af.ModelInstance], cosmology: ag.cosmo.LensingCosmology = None, - fields: Optional[Union[List[MassField], af.ModelInstance]] = None, + fields: Optional[Union[MassField, List[MassField], af.ModelInstance]] = None, ): """ Performs gravitational lensing ray-tracing calculations based on an input list of galaxies and a cosmology. @@ -330,7 +334,8 @@ def __init__( 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 + external shear, a mass sheet, an external potential). Pass a list or a single + `MassField`; the tracer stores a fresh list in either case. 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. """ @@ -343,13 +348,15 @@ def __init__( if fields is not None: _validate_fields(fields=fields) + fields = tracer_util.fields_list_from(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.fields = fields self.cosmology = cosmology or ag.cosmo.Planck15() @@ -492,7 +499,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, + fields: Optional[Union[MassField, List[MassField]]] = None, ): """ Returns a tracer where the lens system is split into planes with specified redshift distances between them. @@ -541,6 +548,7 @@ def sliced_tracer_from( 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. + Accepts a list or a single `MassField`. """ cosmology = cosmology or ag.cosmo.Planck15() @@ -556,7 +564,7 @@ def sliced_tracer_from( galaxies = lens_galaxies + line_of_sight_galaxies + source_galaxies - fields = list(fields) if fields is not None else [] + fields = tracer_util.fields_list_from(fields) # 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. diff --git a/autolens/lens/tracer_util.py b/autolens/lens/tracer_util.py index 9b9186fed..a9b162357 100644 --- a/autolens/lens/tracer_util.py +++ b/autolens/lens/tracer_util.py @@ -4,6 +4,7 @@ This module contains lower-level helpers that are called by ``Tracer`` but kept separate to avoid cluttering the main class. Key functions: +- ``fields_list_from`` — normalizes an optional single field or collection to a fresh list. - ``plane_redshifts_from`` — derives the list of unique plane redshifts from a list of galaxies, collapsing multiple galaxies at the same redshift into a single plane. - ``ordered_plane_redshifts_with_slicing_from`` — extends the above with optional @@ -21,6 +22,18 @@ from autolens import exc +def fields_list_from(fields) -> List[ag.MassField]: + """Return a fresh list from no fields, a single MassField, or a field collection. + + Container and entry validation remain the responsibility of the caller. + """ + if fields is None: + return [] + if isinstance(fields, ag.MassField): + return [fields] + return list(fields) + + def _redshift_is_traced(redshift) -> bool: """ Return True if ``redshift`` is a JAX traced scalar that cannot be coerced to a diff --git a/docs/api/galaxy.rst b/docs/api/galaxy.rst index da6550690..1145ad727 100644 --- a/docs/api/galaxy.rst +++ b/docs/api/galaxy.rst @@ -15,7 +15,8 @@ shear, a mass sheet, an external potential — at its own redshift and with no l 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. +instead remains fully supported. For one field, ``Tracer(galaxies=..., fields=field)`` +is also accepted; ``tracer.fields`` is always stored as a list. .. currentmodule:: autolens diff --git a/docs/api/modeling.rst b/docs/api/modeling.rst index c261804f7..0941e8365 100644 --- a/docs/api/modeling.rst +++ b/docs/api/modeling.rst @@ -73,6 +73,10 @@ 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. +The primary form is ``fields=af.Collection(field=field)``. For a single field, +``fields=field`` is also supported. It shortens prior paths from +``fields.field.shear.gamma_1`` to ``fields.shear.gamma_1`` and therefore gives a +different result identifier by design. Existing collection-form identifiers are unchanged. .. currentmodule:: autolens.analysis.model_util diff --git a/test_autolens/aggregator/test_tracer.py b/test_autolens/aggregator/test_tracer.py new file mode 100644 index 000000000..25a3abad4 --- /dev/null +++ b/test_autolens/aggregator/test_tracer.py @@ -0,0 +1,41 @@ +from types import SimpleNamespace + +import pytest +import autofit as af +import autolens as al + +from autolens.aggregator.tracer import _tracer_from + + +@pytest.mark.parametrize("form", ["flat", "collection", "absent", "none", "empty"]) +@pytest.mark.parametrize("override", [False, True]) +def test__tracer_from__normalizes_saved_and_override_fields(form, override): + field = af.Model(al.MassField, redshift=0.5, shear=al.mp.ExternalShear) + model = af.Collection( + galaxies=af.Collection( + lens=al.Galaxy(redshift=0.5, mass=al.mp.IsothermalSph()), + source=al.Galaxy(redshift=1.0, bulge=al.lp.SersicSph()), + ), + cosmology=al.cosmo.Planck15(), + ) + if form == "flat": + model.fields = field + elif form == "collection": + model.fields = af.Collection(field=field) + elif form == "none": + model.fields = None + elif form == "empty": + model.fields = af.Collection() + instance = model.instance_from_prior_medians() + fit = SimpleNamespace(children=[], instance=instance) + + tracers = _tracer_from(fit=fit, instance=instance if override else None) + + assert len(tracers) == 1 + expected = ( + [instance.fields] + if form == "flat" + else [instance.fields.field] if form == "collection" else [] + ) + assert tracers[0].fields == expected + assert len(tracers[0].galaxies) == 2 diff --git a/test_autolens/analysis/analysis/test_analysis_lens.py b/test_autolens/analysis/analysis/test_analysis_lens.py index 1926a8b5b..9abfa52d9 100644 --- a/test_autolens/analysis/analysis/test_analysis_lens.py +++ b/test_autolens/analysis/analysis/test_analysis_lens.py @@ -149,3 +149,43 @@ def test__log_likelihood_function__a_field_shear_matches_the_galaxy_attached_she ) assert field_likelihood == pytest.approx(galaxy_likelihood, 1.0e-8) + + +def test__tracer_via_instance_from__bare_field_is_folded_into_tracer_fields( + analysis_imaging_7x7, +): + model = af.Collection( + galaxies=af.Collection( + lens=af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal), + source=af.Model(al.Galaxy, redshift=1.0, bulge=al.lp.Sersic), + ), + fields=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 tracer.fields == [instance.fields] + assert all(field is not instance.fields for field in tracer.galaxies) + + +def test__log_likelihood__bare_and_collection_fields_agree(masked_imaging_7x7): + analysis = al.AnalysisImaging(dataset=masked_imaging_7x7, use_jax=False) + galaxies = af.Collection( + lens=al.Galaxy(redshift=0.5, mass=al.mp.IsothermalSph(einstein_radius=1.0)), + source=al.Galaxy(redshift=1.0, bulge=al.lp.SersicSph(intensity=1.0)), + ) + field = af.Model( + al.MassField, + redshift=0.5, + shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.06), + ) + flat = af.Collection(galaxies=galaxies, fields=field).instance_from_prior_medians() + collection = af.Collection( + galaxies=galaxies, fields=af.Collection(field=field) + ).instance_from_prior_medians() + assert analysis.log_likelihood_function(flat) == pytest.approx( + analysis.log_likelihood_function(collection), rel=0, abs=1.0e-8 + ) diff --git a/test_autolens/analysis/test_result.py b/test_autolens/analysis/test_result.py index c502038bd..e57e4caf0 100644 --- a/test_autolens/analysis/test_result.py +++ b/test_autolens/analysis/test_result.py @@ -440,3 +440,23 @@ def _poison(*args, **kwargs): assert isinstance(second, al.PositionsLH) assert second.positions.array == pytest.approx(first.positions.array, 1.0e-8) assert second.threshold == pytest.approx(first.threshold, 1.0e-8) + + +@pytest.mark.parametrize("form", ["flat", "collection", "absent", "none", "empty"]) +def test__max_log_likelihood_fields__normalizes_model_slot(form, analysis_imaging_7x7): + field = al.MassField(redshift=0.5, shear=al.mp.ExternalShear(gamma_1=0.05)) + instance = af.ModelInstance() + if form == "flat": + instance.fields = field + elif form == "collection": + instance.fields = af.ModelInstance({"field": field}) + elif form == "none": + instance.fields = None + elif form == "empty": + instance.fields = af.ModelInstance() + result = res.Result( + samples_summary=al.m.MockSamplesSummary(max_log_likelihood_instance=instance), + analysis=analysis_imaging_7x7, + ) + expected = [field] if form in ("flat", "collection") else [] + assert result.max_log_likelihood_fields == expected diff --git a/test_autolens/lens/test_tracer_fields.py b/test_autolens/lens/test_tracer_fields.py index 4c7de9435..449774875 100644 --- a/test_autolens/lens/test_tracer_fields.py +++ b/test_autolens/lens/test_tracer_fields.py @@ -530,3 +530,122 @@ def test__regression__the_model_identifier_is_unchanged_by_the_fields_slot(): assert result.returncode == 0, result.stderr assert result.stdout.strip().splitlines()[-1] == "fef2697b5c32ba56bb18a7baecb7b0f6" + + +@pytest.fixture(name="bare_field") +def make_bare_field(): + return al.MassField( + redshift=0.75, shear=_shear(), mass_sheet=al.mp.MassSheet(kappa=0.1) + ) + + +def test__bare_field__stored_as_single_member_and_equivalent_to_list(bare_field, grid): + galaxies = [_lens(), _source()] + bare = al.Tracer(galaxies=galaxies, fields=bare_field) + listed = al.Tracer(galaxies=galaxies, fields=[bare_field]) + + assert bare.fields == [bare_field] + assert bare.fields[0] is bare_field + assert bare.members == listed.members + assert bare.plane_redshifts == listed.plane_redshifts == [0.5, 0.75, 1.0] + for method in ("deflections_yx_2d_from", "convergence_2d_from"): + assert np.asarray(getattr(bare, method)(grid)) == pytest.approx( + np.asarray(getattr(listed, method)(grid)), abs=1.0e-8 + ) + assert np.any(np.asarray(bare.convergence_2d_from(grid)) != 0) + + +def test__fields_list__storage_does_not_alias_input(bare_field): + fields = [bare_field] + tracer = al.Tracer(galaxies=[_lens(), _source()], fields=fields) + fields.clear() + assert tracer.fields == [bare_field] + + +def test__bare_field__geometry_warning_is_not_skipped(): + 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__bare_field__json_matches_list_and_round_trips(bare_field, tmp_path): + galaxies = [_lens(), _source()] + bare = al.Tracer(galaxies=galaxies, fields=bare_field) + listed = al.Tracer(galaxies=galaxies, fields=[bare_field]) + bare_path, list_path = tmp_path / "bare.json", tmp_path / "list.json" + output_to_json(bare, file_path=bare_path) + output_to_json(listed, file_path=list_path) + assert bare_path.read_bytes() == list_path.read_bytes() + restored = from_json(file_path=bare_path) + assert isinstance(restored.fields, list) + assert isinstance(restored.fields[0], al.MassField) + assert al.to_dict(restored) == al.to_dict(listed) + + +def test__sliced_tracer_from__bare_field_redshift_is_snapped(): + field = al.MassField(redshift=0.42, shear=_shear()) + tracer = al.Tracer.sliced_tracer_from( + lens_galaxies=[al.Galaxy(redshift=0.5)], + line_of_sight_galaxies=[], + source_galaxies=[al.Galaxy(redshift=2.0)], + planes_between_lenses=[1, 1], + fields=field, + ) + assert tracer.fields == [field] + assert field.redshift == 0.5 + + +def test__bare_galaxy_in_fields__still_rejected(): + with pytest.raises(TypeError, match="Pass a list of MassField objects"): + al.Tracer(galaxies=[_lens()], fields=_source()) + + +@pytest.mark.parametrize("fields", [None, [], (), af.ModelInstance()]) +def test__empty_fields__stored_as_empty_list(fields): + assert al.Tracer(galaxies=[_lens(), _source()], fields=fields).fields == [] + + +def test__generator_fields__rejected_without_consuming_entries(bare_field): + fields = (field for field in [bare_field]) + with pytest.raises(TypeError, match="Pass a list of MassField objects"): + al.Tracer(galaxies=[_lens(), _source()], fields=fields) + assert next(fields) is bare_field + + +_FIELDS_IDENTIFIER_SNIPPET = """ +import autofit as af +import autolens as al + +paths = [] +for flat in (False, True): + field = af.Model(al.MassField, redshift=0.5, shear=af.Model(al.mp.ExternalShear)) + model = af.Collection( + galaxies=af.Collection( + lens=af.Model(al.Galaxy, redshift=0.5, mass=af.Model(al.mp.Isothermal)), + source=af.Model(al.Galaxy, redshift=1.0, bulge=af.Model(al.lp.Sersic)), + ), + fields=field if flat else af.Collection(field=field), + ) + assert model.prior_count == 14 + field_paths = {path for path in model.unique_prior_paths if path[0] == "fields"} + prefix = ("fields", "shear") if flat else ("fields", "field", "shear") + assert field_paths == {prefix + ("gamma_1",), prefix + ("gamma_2",)} + paths.append({path for path in model.unique_prior_paths if path[0] != "fields"}) + # Captured before the bare-field implementation, under shipped configuration. + expected = "36a0be37c9667958bafca2f01e487f80" if flat else "c5cf98ea7733689dc8c0ede6939c1d83" + assert model.identifier == expected, model.identifier +assert paths[0] == paths[1] +""" + + +def test__fields_model_forms__prior_paths_counts_and_identifiers_are_preserved(): + result = subprocess.run( + [sys.executable, "-c", _FIELDS_IDENTIFIER_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.stdout + result.stderr