From 0ee1146749e001788c90cf042733db19d79afb4e Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 16 Sep 2026 23:24:42 +0100 Subject: [PATCH] feat(interop): COOLEST observation pixel grid + on_unsupported="skip" to_coolest wrote every template with an empty observation.pixels block (field_of_view 0, num_pix 0) because it built lazy.Observation() with no pixel grid; COOLEST's plotting API cannot render that. It now writes a PixelatedRegularGrid from shape_native + pixel_size or from dataset= (anything with shape_native and pixel_scales), and warns when neither is given. on_unsupported="skip" exports the full mass model and any supported light while omitting profiles COOLEST cannot represent (an MGE Basis, a Pixelization), recording each under meta.skipped_profiles. The default "raise" is unchanged. Refs #739 Co-Authored-By: Claude Fable 5.1 --- autolens/interop/coolest.py | 180 +++++++++++++++++++--- test_autolens/interop/test_coolest.py | 214 ++++++++++++++++++++++++-- 2 files changed, 358 insertions(+), 36 deletions(-) diff --git a/autolens/interop/coolest.py b/autolens/interop/coolest.py index 726b98af0..2c634e85d 100644 --- a/autolens/interop/coolest.py +++ b/autolens/interop/coolest.py @@ -37,8 +37,10 @@ """ import os -from typing import Dict, List, Optional, Union +import warnings +from typing import Dict, List, Optional, Tuple, Union +import autoarray as aa import autogalaxy as ag from autogalaxy.galaxy.galaxy import Galaxy @@ -92,6 +94,74 @@ def _sigma_crit_from(cosmology, redshift_0: float, redshift_1: float) -> float: return float(sigma_crit_kpc2 * kpc_per_arcsec**2) +def _profile_name_from(profile) -> str: + """ + The name a profile is recorded under in the template's + ``meta.skipped_profiles`` list. + + A ``Basis`` (e.g. an MGE) is a single profile holding many inner profiles, + so it is named ``Basis( x N)`` to make clear how many + components were omitted. + """ + name = type(profile).__name__ + + profile_list = getattr(profile, "profile_list", None) + + if profile_list is None: + return name + + inner = [type(inner_profile).__name__ for inner_profile in profile_list] + unique = sorted(set(inner)) + + if len(unique) == 1: + return f"{name}({unique[0]} x {len(inner)})" + return f"{name}({len(inner)} profiles)" + + +def _record_skipped(metadata: Dict, galaxy_name: str, profile_name: str): + metadata.setdefault("skipped_profiles", []).append( + {"galaxy": galaxy_name, "profile": profile_name} + ) + + +def _observation_from(lazy, shape_native, pixel_size: float): + """ + The COOLEST ``Observation`` block, whose pixel grid gives the template the + field of view and pixel count that COOLEST's plotting API requires. + + A template written without it has a ``PixelatedRegularGrid`` of zeros, + which COOLEST's plotting cannot use, so a warning is raised. + """ + if shape_native is None: + warnings.warn( + "The COOLEST template is being written with no observation pixel " + "grid, so its `observation.pixels` field of view and pixel counts " + "are all 0 and COOLEST's plotting API cannot render it. Pass " + "`shape_native=(y_pixels, x_pixels)` (with `pixel_size`) or " + "`dataset=` the imaging the model was fitted to.", + UserWarning, + ) + return lazy.Observation() + + pixels_y, pixels_x = int(shape_native[0]), int(shape_native[1]) + pixel_size = float(pixel_size) + + grid = lazy.PixelatedRegularGrid( + field_of_view_x=( + -pixels_x * pixel_size / 2.0, + pixels_x * pixel_size / 2.0, + ), + field_of_view_y=( + -pixels_y * pixel_size / 2.0, + pixels_y * pixel_size / 2.0, + ), + num_pix_x=pixels_x, + num_pix_y=pixels_y, + ) + + return lazy.Observation(pixels=grid) + + def _set_parameters(coolest_profile, parameters: Dict, point_estimate_cls): for name, value in parameters.items(): coolest_profile.parameters[name].set_point_estimate( @@ -105,6 +175,9 @@ def to_coolest( cosmology: Optional[ag.cosmo.LensingCosmology] = None, mode: str = "MAP", pixel_size: float = 0.1, + shape_native: Optional[Tuple[int, int]] = None, + dataset=None, + on_unsupported: str = "raise", metadata: Optional[Dict] = None, ) -> str: """ @@ -133,14 +206,37 @@ def to_coolest( for a simulated one. pixel_size The instrument pixel size (arcsec) written to the template's minimal - instrument block. + instrument block and used to compute the observation grid's field of + view. Overridden by ``dataset`` if one is passed. + shape_native + The ``(y_pixels, x_pixels)`` shape of the observation, which with + ``pixel_size`` writes the template's ``observation.pixels`` grid (its + field of view and pixel counts). COOLEST's plotting API requires this + grid; if neither ``shape_native`` nor ``dataset`` is passed the grid is + written as zeros and a ``UserWarning`` is raised. + dataset + The dataset the model was fitted to (e.g. an ``al.Imaging``), or any + object with ``shape_native`` and ``pixel_scales`` attributes, from + which ``shape_native`` and ``pixel_size`` are taken. + on_unsupported + What to do with profiles COOLEST cannot represent (e.g. a ``Basis`` of + Gaussians, an MGE, or a ``Pixelization``). "raise" (default) raises the + converter's error naming the profile; "skip" omits the profile from the + template and records it in the template's + ``meta["skipped_profiles"]`` as ``{"galaxy": ..., "profile": ...}``, + so the supported parts of the model still export. metadata - Extra metadata stored in the template. + Extra metadata stored in the template's ``meta`` block. Returns ------- The path of the written ``.json`` template file. """ + if on_unsupported not in ("raise", "skip"): + raise ValueError( + f"The `on_unsupported` input of `to_coolest` must be 'raise' or " + f"'skip', but '{on_unsupported}' was passed." + ) PointEstimate, lazy, COOLEST, JSONSerializer = _coolest_modules() if isinstance(galaxies, Tracer): @@ -150,15 +246,45 @@ def to_coolest( galaxies = list(galaxies) 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) + metadata = dict(metadata or {}) + entities = [] for i, galaxy in enumerate(galaxies): - light_dicts = [ - coolest_dict_from_light(profile=profile) - for profile in galaxy.cls_list_from(cls=LightProfile) - ] + galaxy_name = f"galaxy_{i}" + + # A `Basis` is both a `LightProfile` and a `MassProfile`, so it appears + # in both lists below and is recorded only once. + skipped_ids = set() + + light_dicts = [] + + for profile in galaxy.cls_list_from(cls=LightProfile): + try: + light_dicts.append(coolest_dict_from_light(profile=profile)) + except (ag.exc.ProfileException, ValueError): + if on_unsupported == "raise": + raise + skipped_ids.add(id(profile)) + _record_skipped( + metadata=metadata, + galaxy_name=galaxy_name, + profile_name=_profile_name_from(profile), + ) + + if on_unsupported == "skip": + for pixelization in galaxy.cls_list_from(cls=aa.Pixelization): + _record_skipped( + metadata=metadata, + galaxy_name=galaxy_name, + profile_name=(f"Pixelization({type(pixelization.mesh).__name__})"), + ) mass_profiles = galaxy.cls_list_from(cls=MassProfile) @@ -175,12 +301,24 @@ def to_coolest( mass_dicts = [] for profile in mass_profiles: + if id(profile) in skipped_ids: + continue if isinstance(profile, (ExternalShear, MassSheet)): field_profiles.append(profile) else: - mass_dicts.append( - coolest_dict_from_mass(profile=profile, sigma_crit=sigma_crit) - ) + try: + mass_dicts.append( + coolest_dict_from_mass(profile=profile, sigma_crit=sigma_crit) + ) + except (ag.exc.ProfileException, ValueError): + if on_unsupported == "raise": + raise + skipped_ids.add(id(profile)) + _record_skipped( + metadata=metadata, + galaxy_name=galaxy_name, + profile_name=_profile_name_from(profile), + ) light_model = lazy.LightModel(*[d["type"] for d in light_dicts]) mass_model = lazy.MassModel(*[d["type"] for d in mass_dicts]) @@ -196,7 +334,7 @@ def to_coolest( entities.append( lazy.Galaxy( - f"galaxy_{i}", + galaxy_name, float(galaxy.redshift), light_model=light_model, mass_model=mass_model, @@ -205,13 +343,10 @@ def to_coolest( if field_profiles: field_dicts = [ - coolest_dict_from_mass(profile=profile) - for profile in field_profiles + coolest_dict_from_mass(profile=profile) for profile in field_profiles ] field_model = lazy.MassModel(*[d["type"] for d in field_dicts]) - for coolest_profile, profile_dict in zip( - list(field_model), field_dicts - ): + for coolest_profile, profile_dict in zip(list(field_model), field_dicts): _set_parameters( coolest_profile=coolest_profile, parameters=profile_dict["parameters"], @@ -228,7 +363,6 @@ def to_coolest( h0 = getattr(cosmology.H0, "value", cosmology.H0) om0 = getattr(cosmology.Om0, "value", cosmology.Om0) - metadata = dict(metadata or {}) metadata.setdefault( "generated_by", "PyAutoLens (autolens.interop.coolest); theta_E is the profile's " @@ -240,7 +374,7 @@ def to_coolest( mode, lazy.CoordinatesOrigin(), lazy.LensingEntityList(*entities), - lazy.Observation(), + _observation_from(lazy=lazy, shape_native=shape_native, pixel_size=pixel_size), lazy.Instrument(pixel_size), cosmology=lazy.Cosmology(H0=float(h0), Om0=float(om0)), metadata=metadata, @@ -299,9 +433,9 @@ def from_coolest( """ _, _, _, JSONSerializer = _coolest_modules() - root = JSONSerializer( - _path_no_ext(file_path), check_external_files=False - ).load(verbose=False) + root = JSONSerializer(_path_no_ext(file_path), check_external_files=False).load( + verbose=False + ) if cosmology is None: if root.cosmology is not None: @@ -349,8 +483,6 @@ def from_coolest( intermediate=intermediate, ) - galaxies.append( - Galaxy(redshift=float(entity.redshift), **profiles) - ) + galaxies.append(Galaxy(redshift=float(entity.redshift), **profiles)) return Tracer(galaxies=galaxies, cosmology=cosmology) diff --git a/test_autolens/interop/test_coolest.py b/test_autolens/interop/test_coolest.py index 920ccab65..28766fd00 100644 --- a/test_autolens/interop/test_coolest.py +++ b/test_autolens/interop/test_coolest.py @@ -127,13 +127,11 @@ def test__round_trip__nfw_uses_sigma_crit_from_cosmology(tmp_path): file_path = interop_coolest.to_coolest( galaxies=galaxies, file_path=str(tmp_path / "template"), cosmology=cosmology ) - tracer_back = interop_coolest.from_coolest( - file_path=file_path, cosmology=cosmology - ) + tracer_back = interop_coolest.from_coolest(file_path=file_path, cosmology=cosmology) - nfw_back = [ - galaxy for galaxy in tracer_back.galaxies if galaxy.redshift == 0.3 - ][0].mass_0 + nfw_back = [galaxy for galaxy in tracer_back.galaxies if galaxy.redshift == 0.3][ + 0 + ].mass_0 assert nfw_back.kappa_s == pytest.approx(0.15, rel=1e-8) assert nfw_back.scale_radius == pytest.approx(6.0, rel=1e-8) @@ -166,9 +164,7 @@ def test__from_coolest__intermediate_returns_power_law_intermediate(tmp_path): galaxies=galaxies, file_path=str(tmp_path / "template") ) - tracer_back = interop_coolest.from_coolest( - file_path=file_path, intermediate=True - ) + tracer_back = interop_coolest.from_coolest(file_path=file_path, intermediate=True) mass_back = [g for g in tracer_back.galaxies if g.redshift == 0.5][0].mass_0 @@ -177,9 +173,9 @@ def test__from_coolest__intermediate_returns_power_law_intermediate(tmp_path): grid = al.Grid2DIrregular([[1.0, 0.5], [-0.3, 0.7]]) tracer = al.Tracer(galaxies=galaxies) - assert np.asarray( - tracer_back.deflections_yx_2d_from(grid=grid) - ) == pytest.approx(np.asarray(tracer.deflections_yx_2d_from(grid=grid)), rel=1e-8) + assert np.asarray(tracer_back.deflections_yx_2d_from(grid=grid)) == pytest.approx( + np.asarray(tracer.deflections_yx_2d_from(grid=grid)), rel=1e-8 + ) def test__round_trip__relative_file_path(tmp_path, monkeypatch): @@ -225,3 +221,197 @@ def test__from_coolest__missing_point_estimate_raises(tmp_path): with pytest.raises(ValueError): interop_coolest.from_coolest(file_path=f"{path}.json") + + +def unsupported_model_galaxies(): + """ + A model COOLEST cannot fully represent: an MGE (`Basis` of Gaussians) lens + light, an isothermal mass with shear and a pixelized (Delaunay) source. + """ + lens = al.Galaxy( + redshift=0.5, + bulge=al.lp_basis.Basis( + profile_list=[ + al.lp.GaussianSph(intensity=1.0, sigma=0.1), + al.lp.GaussianSph(intensity=0.5, sigma=0.3), + ] + ), + mass=al.mp.Isothermal( + centre=(0.0, 0.0), + ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=45.0), + einstein_radius=1.6, + ), + shear=al.mp.ExternalShear(gamma_1=0.02, gamma_2=-0.03), + ) + source = al.Galaxy( + redshift=1.0, + pixelization=al.Pixelization( + mesh=al.mesh.Delaunay(pixels=100), + regularization=al.reg.Constant(coefficient=1.0), + ), + ) + return [lens, source] + + +def test__to_coolest__shape_native_writes_observation_pixel_grid(tmp_path): + file_path = interop_coolest.to_coolest( + galaxies=lens_model_galaxies(), + file_path=str(tmp_path / "template"), + shape_native=(10, 12), + pixel_size=0.05, + ) + + with open(file_path) as f: + template = json.load(f) + + pixels = template["observation"]["pixels"] + + assert pixels["field_of_view_x"] == pytest.approx([-0.3, 0.3]) + assert pixels["field_of_view_y"] == pytest.approx([-0.25, 0.25]) + assert pixels["num_pix_x"] == 12 + assert pixels["num_pix_y"] == 10 + + # The grid's pixel scale is the instrument pixel size, which is what + # COOLEST's plotting API uses to lay the image out. + pixel_size = (pixels["field_of_view_x"][1] - pixels["field_of_view_x"][0]) / pixels[ + "num_pix_x" + ] + + assert pixel_size == pytest.approx(template["instrument"]["pixel_size"]) + + +def test__to_coolest__dataset_writes_observation_pixel_grid(tmp_path): + dataset = al.Imaging( + data=al.Array2D.no_mask( + values=np.ones((8, 6)), + pixel_scales=0.2, + ), + noise_map=al.Array2D.no_mask( + values=np.ones((8, 6)), + pixel_scales=0.2, + ), + ) + + file_path = interop_coolest.to_coolest( + galaxies=lens_model_galaxies(), + file_path=str(tmp_path / "template"), + dataset=dataset, + ) + + with open(file_path) as f: + template = json.load(f) + + pixels = template["observation"]["pixels"] + + assert pixels["num_pix_y"] == 8 + assert pixels["num_pix_x"] == 6 + assert pixels["field_of_view_y"] == pytest.approx([-0.8, 0.8]) + assert pixels["field_of_view_x"] == pytest.approx([-0.6, 0.6]) + assert template["instrument"]["pixel_size"] == pytest.approx(0.2) + + +def test__to_coolest__no_grid_warns_and_writes_zeros(tmp_path): + with pytest.warns(UserWarning): + file_path = interop_coolest.to_coolest( + galaxies=lens_model_galaxies(), file_path=str(tmp_path / "template") + ) + + with open(file_path) as f: + template = json.load(f) + + pixels = template["observation"]["pixels"] + + assert pixels["num_pix_x"] == 0 + assert pixels["num_pix_y"] == 0 + assert pixels["field_of_view_x"] == pytest.approx([0.0, 0.0]) + + +def test__to_coolest__unsupported_profiles_raise_by_default(tmp_path): + import autogalaxy as ag + + with pytest.raises(ag.exc.ProfileException): + interop_coolest.to_coolest( + galaxies=unsupported_model_galaxies(), + file_path=str(tmp_path / "template"), + shape_native=(10, 10), + ) + + +def test__to_coolest__on_unsupported_skip_records_skipped_profiles(tmp_path): + file_path = interop_coolest.to_coolest( + galaxies=unsupported_model_galaxies(), + file_path=str(tmp_path / "template"), + shape_native=(10, 10), + on_unsupported="skip", + ) + + with open(file_path) as f: + template = json.load(f) + + skipped = template["meta"]["skipped_profiles"] + + profiles = [entry["profile"] for entry in skipped] + + assert any( + profile.startswith("Basis(") and "Gaussian" in profile for profile in profiles + ) + assert "Pixelization(Delaunay)" in profiles + + # The `Basis` is both a light and a mass profile but is recorded once. + assert len([p for p in profiles if p.startswith("Basis(")]) == 1 + + assert [entry["galaxy"] for entry in skipped if "Basis" in entry["profile"]] == [ + "galaxy_0" + ] + + entities = template["lensing_entities"] + + lens_entity = [ + entity + for entity in entities + if entity["type"] == "Galaxy" and entity["redshift"] == 0.5 + ][0] + + # The supported mass profile is still exported, with no light profiles. + assert [profile["type"] for profile in lens_entity["mass_model"]] == ["SIE"] + assert lens_entity["light_model"] == [] + + # The shear is still written as its own MassField entity. + assert [entity["type"] for entity in entities].count("MassField") == 1 + + +def test__to_coolest__on_unsupported_invalid_raises(tmp_path): + with pytest.raises(ValueError): + interop_coolest.to_coolest( + galaxies=lens_model_galaxies(), + file_path=str(tmp_path / "template"), + shape_native=(10, 10), + on_unsupported="ignore", + ) + + +def test__from_coolest__skipped_template_returns_mass_profiles_only(tmp_path): + file_path = interop_coolest.to_coolest( + galaxies=unsupported_model_galaxies(), + file_path=str(tmp_path / "template"), + shape_native=(10, 10), + on_unsupported="skip", + ) + + tracer_back = interop_coolest.from_coolest(file_path=file_path) + + from autogalaxy.profiles.light.abstract import LightProfile + from autogalaxy.profiles.mass.abstract.abstract import MassProfile + + light_profiles = [] + mass_profiles = [] + + for galaxy in tracer_back.galaxies: + light_profiles += galaxy.cls_list_from(cls=LightProfile) + mass_profiles += galaxy.cls_list_from(cls=MassProfile) + + assert light_profiles == [] + + types = sorted(type(profile).__name__ for profile in mass_profiles) + + assert types == ["ExternalShear", "Isothermal"]