diff --git a/autolens/analysis/latent.py b/autolens/analysis/latent.py index 0ef44bbd7..e1c314d29 100644 --- a/autolens/analysis/latent.py +++ b/autolens/analysis/latent.py @@ -9,10 +9,13 @@ wiring can reuse it without code duplication. User-level enable/disable: each key in ``autolens/config/latent.yaml`` maps -to a bool. All five default ``false`` because ``compute_latent_samples`` -runs on every fit (``latent_after_fit: true`` in autofit's default -``output.yaml``) and the latents that require ``magzero`` would otherwise -crash existing fits where ``magzero`` is not passed. +to a bool. The raw-flux latents (``total_lens_flux``, +``total_lensed_source_flux``, ``total_source_flux``) require no instrument +inputs and default ``true``. The microjansky variants require ``magzero`` +on the Analysis; they default ``false`` and return NaN with a single +warning per process if enabled without ``magzero`` (rather than raising, +which would kill the post-fit metric write of an otherwise-converged +search). """ import logging from typing import Callable, Dict, List, Optional @@ -27,14 +30,84 @@ logger = logging.getLogger(__name__) +# Latent names that have already emitted a "magzero missing" warning in this +# process. Used by ``_maybe_magzero_warn`` to deduplicate the message across +# the many fit evaluations a single search performs. +_MAGZERO_WARNED: set = set() -def _require_magzero(magzero, name): + +def _maybe_magzero_warn(magzero, name) -> bool: + """ + Return True when ``magzero`` is missing (and emit a one-time-per-process + warning for ``name``); False otherwise. + + Callers that get True must early-return ``xp.nan`` — the µJy conversion + is meaningless without a zero-point, but a search-killing raise here + would discard otherwise-converged fits. + """ if magzero is None: - raise ValueError( - f"magzero must be passed to the Analysis via kwargs to compute " - f"the '{name}' latent. Disable it in config/latent.yaml or " - f"pass magzero=." + if name not in _MAGZERO_WARNED: + logger.warning( + "magzero not set on Analysis; '%s' latent will be NaN. " + "Pass magzero= to AnalysisImaging to enable it, " + "or disable in config/latent.yaml to silence this warning.", + name, + ) + _MAGZERO_WARNED.add(name) + return True + return False + + +def total_lens_flux(fit, magzero=None, xp=np): + """ + Total integrated flux of the lens galaxy (``fit.tracer.galaxies[0]``), + in the raw image units the fit was performed in. + + Requires no instrument inputs — ``magzero`` is accepted for uniform + dispatcher context but ignored. See the workspace flux guide + (``scripts/guides/units/flux.py``) for how to convert to microjanskies. + + Returns NaN when galaxy 0 has no light profile. + """ + try: + image = fit.galaxy_image_dict[fit.tracer.galaxies[0]] + except (AttributeError, KeyError, IndexError): + return xp.nan + return xp.sum(image.array) + + +def total_lensed_source_flux(fit, magzero=None, xp=np): + """ + Image-plane integrated flux of the source galaxy after lensing + (``fit.galaxy_image_dict[fit.tracer.galaxies[-1]]``), in raw image + units. ``magzero`` is accepted but ignored. + """ + try: + image = fit.galaxy_image_dict[fit.tracer.galaxies[-1]] + except (AttributeError, KeyError, IndexError): + return xp.nan + return xp.sum(image.array) + + +def total_source_flux(fit, magzero=None, xp=np): + """ + Source-plane intrinsic flux of the source galaxy, in raw image units. + + Reads from ``fit.tracer_linear_light_profiles_to_light_profiles`` so + that linear light profiles (whose ``intensity`` is solved by the + inversion) contribute the correct image — same tracer-conversion + handling as :func:`total_source_flux_mujy`. + + ``magzero`` is accepted but ignored. + """ + try: + tracer = fit.tracer_linear_light_profiles_to_light_profiles + source_image = tracer.galaxies[-1].image_2d_from( + grid=fit.dataset.grids.lp, xp=xp ) + except (AttributeError, IndexError): + return xp.nan + return xp.sum(source_image.array) def total_lens_flux_mujy(fit, magzero, xp=np): @@ -42,10 +115,16 @@ def total_lens_flux_mujy(fit, magzero, xp=np): Total integrated flux of the lens galaxy (``fit.tracer.galaxies[0]``), magzero-converted to microjanskies. - Returns NaN when galaxy 0 has no light profile (raises ``KeyError`` / - ``AttributeError`` inside ``fit.galaxy_image_dict``). + Returns NaN — with a one-time-per-process warning — when ``magzero`` + is missing, rather than raising. The µJy conversion is meaningless + without a zero-point, but a hard raise during post-fit latent + computation would discard the result of an otherwise-converged + multi-hour search. + + Also returns NaN when galaxy 0 has no light profile. """ - _require_magzero(magzero, "total_lens_flux_mujy") + if _maybe_magzero_warn(magzero, "total_lens_flux_mujy"): + return xp.nan try: image = fit.galaxy_image_dict[fit.tracer.galaxies[0]] except (AttributeError, KeyError, IndexError): @@ -60,9 +139,13 @@ def total_lens_flux_mujy(fit, magzero, xp=np): def total_lensed_source_flux_mujy(fit, magzero, xp=np): """ Image-plane integrated flux of the source galaxy after lensing - (``fit.galaxy_image_dict[fit.tracer.galaxies[-1]]``). + (``fit.galaxy_image_dict[fit.tracer.galaxies[-1]]``), in microjanskies. + + Returns NaN + one warning when ``magzero`` is missing; see + :func:`total_lens_flux_mujy` for the rationale. """ - _require_magzero(magzero, "total_lensed_source_flux_mujy") + if _maybe_magzero_warn(magzero, "total_lensed_source_flux_mujy"): + return xp.nan try: image = fit.galaxy_image_dict[fit.tracer.galaxies[-1]] except (AttributeError, KeyError, IndexError): @@ -83,8 +166,12 @@ def total_source_flux_mujy(fit, magzero, xp=np): is solved by the inversion at fit time) contribute the correct image. For non-linear fits this property is a no-op pass-through (returns ``fit.tracer``), so the numpy-only and JAX paths both work uniformly. + + Returns NaN + one warning when ``magzero`` is missing; see + :func:`total_lens_flux_mujy` for the rationale. """ - _require_magzero(magzero, "total_source_flux_mujy") + if _maybe_magzero_warn(magzero, "total_source_flux_mujy"): + return xp.nan try: tracer = fit.tracer_linear_light_profiles_to_light_profiles source_image = tracer.galaxies[-1].image_2d_from( @@ -142,6 +229,9 @@ def effective_einstein_radius(fit, magzero, xp=np): LATENT_FUNCTIONS: Dict[str, Callable] = { + "total_lens_flux": total_lens_flux, + "total_lensed_source_flux": total_lensed_source_flux, + "total_source_flux": total_source_flux, "total_lens_flux_mujy": total_lens_flux_mujy, "total_lensed_source_flux_mujy": total_lensed_source_flux_mujy, "total_source_flux_mujy": total_source_flux_mujy, diff --git a/autolens/config/latent.yaml b/autolens/config/latent.yaml index a7a3b8680..48ef66aa4 100644 --- a/autolens/config/latent.yaml +++ b/autolens/config/latent.yaml @@ -9,17 +9,35 @@ # Workspaces should mirror this file in their own `config/latent.yaml` to # override defaults locally (workspace values shadow library values). # -# All keys ship `false` because: -# - `compute_latent_samples` runs on every fit (`latent_after_fit: true` -# in autofit's default output.yaml), so on-by-default would crash any -# existing fit that doesn't pass `magzero` to the Analysis. -# - autoconf lowercases yaml keys at read time, so the registry/yaml -# names must be snake_case-lowercase (this leaks into the `latent.csv` -# column header — e.g. `total_lens_flux_mujy`, not `_muJy`). +# Raw-flux keys (`total_lens_flux`, `total_lensed_source_flux`, +# `total_source_flux`) require no instrument inputs and default `true`. +# The `_mujy` variants require `magzero` on the Analysis; they default +# `false` and return NaN + one warning per process if enabled without +# `magzero` (rather than raising, which would discard a converged search). +# +# autoconf lowercases yaml keys at read time, so the registry/yaml names +# must be snake_case-lowercase (this leaks into the `latent.csv` column +# header — e.g. `total_lens_flux_mujy`, not `_muJy`). + +# `total_lens_flux` — total integrated flux of the lens galaxy +# (`fit.tracer.galaxies[0]`), in the raw image units the fit was performed +# in. No instrument inputs required. +total_lens_flux: true + +# `total_lensed_source_flux` — image-plane integrated flux of the source +# galaxy after lensing (`fit.galaxy_image_dict[tracer.galaxies[-1]]`), in +# raw image units. +total_lensed_source_flux: true + +# `total_source_flux` — source-plane intrinsic flux of the source galaxy, +# in raw image units. Reads from +# `tracer_linear_light_profiles_to_light_profiles` so linear-profile fits +# get the correct (inversion-solved) flux. +total_source_flux: true # `total_lens_flux_mujy` — total integrated flux of the lens galaxy # (`fit.tracer.galaxies[0]`) in microjanskies. Requires `magzero` via -# Analysis kwargs. Returns NaN when the lens has no light profile. +# Analysis kwargs. Returns NaN + one warning if `magzero` is missing. total_lens_flux_mujy: false # `total_lensed_source_flux_mujy` — image-plane integrated flux of the @@ -34,6 +52,11 @@ total_source_flux_mujy: false # `magnification` — ratio of image-plane lensed source flux to source-plane # intrinsic source flux. Dimensionless; `magzero` is accepted but unused. +# Default `false` because it routes through the `_mujy` latents internally +# (the µJy conversions cancel in the ratio) — to flip on, also flip on +# `total_lensed_source_flux_mujy` and `total_source_flux_mujy` and supply +# a `magzero`. A follow-up could rewire `magnification` to the raw-flux +# latents so it's universally enable-able. magnification: false # `effective_einstein_radius` — effective Einstein radius in arcseconds, diff --git a/test_autolens/analysis/test_latent.py b/test_autolens/analysis/test_latent.py index 848f79353..73f6d2951 100644 --- a/test_autolens/analysis/test_latent.py +++ b/test_autolens/analysis/test_latent.py @@ -7,13 +7,17 @@ import autofit as af import autolens as al +from autolens.analysis import latent as _latent_module from autolens.analysis.latent import ( LATENT_FUNCTIONS, effective_einstein_radius, latent_keys_enabled, magnification, + total_lens_flux, total_lens_flux_mujy, + total_lensed_source_flux, total_lensed_source_flux_mujy, + total_source_flux, total_source_flux_mujy, ) @@ -56,9 +60,34 @@ def test_total_lens_flux_mujy_returns_nan_when_no_light_profile(): assert np.isnan(total_lens_flux_mujy(fit=fit, magzero=25.0)) -def test_total_lens_flux_mujy_missing_magzero_raises(): - with pytest.raises(ValueError, match="magzero"): - total_lens_flux_mujy(fit=MagicMock(), magzero=None) +def test_total_lens_flux_mujy_missing_magzero_returns_nan_and_warns(caplog): + _latent_module._MAGZERO_WARNED.discard("total_lens_flux_mujy") + caplog.set_level(logging.WARNING) + fit = _fit_with_galaxy_images({0: [1.0, 2.0, 3.0, 4.0]}) + + value = total_lens_flux_mujy(fit=fit, magzero=None) + + assert np.isnan(value) + assert any( + "magzero" in rec.message and "total_lens_flux_mujy" in rec.message + for rec in caplog.records + ) + + +def test_total_lens_flux_against_known_image(): + fit = _fit_with_galaxy_images({0: [1.0, 2.0, 3.0, 4.0]}) + + # Raw sum — no AB-mag conversion. `magzero` accepted but ignored. + assert total_lens_flux(fit=fit) == pytest.approx(10.0) + assert total_lens_flux(fit=fit, magzero=25.0) == pytest.approx(10.0) + + +def test_total_lens_flux_returns_nan_when_no_light_profile(): + fit = MagicMock() + fit.galaxy_image_dict.__getitem__.side_effect = KeyError("no light") + fit.tracer.galaxies = [object()] + + assert np.isnan(total_lens_flux(fit=fit)) def test_total_lensed_source_flux_mujy_against_known_image(): @@ -71,9 +100,24 @@ def test_total_lensed_source_flux_mujy_against_known_image(): assert value == pytest.approx(expected_muJy) -def test_total_lensed_source_flux_mujy_missing_magzero_raises(): - with pytest.raises(ValueError, match="magzero"): - total_lensed_source_flux_mujy(fit=MagicMock(), magzero=None) +def test_total_lensed_source_flux_mujy_missing_magzero_returns_nan_and_warns(caplog): + _latent_module._MAGZERO_WARNED.discard("total_lensed_source_flux_mujy") + caplog.set_level(logging.WARNING) + fit = _fit_with_galaxy_images({0: [0.0, 0.0], 1: [5.0, 5.0]}) + + value = total_lensed_source_flux_mujy(fit=fit, magzero=None) + + assert np.isnan(value) + assert any( + "magzero" in rec.message and "total_lensed_source_flux_mujy" in rec.message + for rec in caplog.records + ) + + +def test_total_lensed_source_flux_against_known_image(): + fit = _fit_with_galaxy_images({0: [0.0, 0.0], 1: [5.0, 5.0]}) + + assert total_lensed_source_flux(fit=fit) == pytest.approx(10.0) def test_total_source_flux_mujy_against_known_image(): @@ -130,9 +174,62 @@ def test_total_source_flux_mujy_uses_converted_tracer_for_linear_profiles(): assert value != 0.0 -def test_total_source_flux_mujy_missing_magzero_raises(): - with pytest.raises(ValueError, match="magzero"): - total_source_flux_mujy(fit=MagicMock(), magzero=None) +def test_total_source_flux_mujy_missing_magzero_returns_nan_and_warns(caplog): + _latent_module._MAGZERO_WARNED.discard("total_source_flux_mujy") + caplog.set_level(logging.WARNING) + + # Source-plane fit fixture (matches the positive test above). + source = SimpleNamespace( + image_2d_from=lambda grid, xp=np: SimpleNamespace( + array=np.array([2.0, 3.0, 5.0]) + ) + ) + galaxies_namespace = SimpleNamespace(galaxies=[object(), source]) + fit = SimpleNamespace( + tracer=galaxies_namespace, + tracer_linear_light_profiles_to_light_profiles=galaxies_namespace, + dataset=SimpleNamespace(grids=SimpleNamespace(lp=object())), + ) + + value = total_source_flux_mujy(fit=fit, magzero=None) + + assert np.isnan(value) + assert any( + "magzero" in rec.message and "total_source_flux_mujy" in rec.message + for rec in caplog.records + ) + + +def test_total_source_flux_against_known_image(): + source = SimpleNamespace( + image_2d_from=lambda grid, xp=np: SimpleNamespace( + array=np.array([2.0, 3.0, 5.0]) + ) + ) + galaxies_namespace = SimpleNamespace(galaxies=[object(), source]) + fit = SimpleNamespace( + tracer=galaxies_namespace, + tracer_linear_light_profiles_to_light_profiles=galaxies_namespace, + dataset=SimpleNamespace(grids=SimpleNamespace(lp=object())), + ) + + assert total_source_flux(fit=fit) == pytest.approx(10.0) + + +def test_maybe_magzero_warn_logs_only_once_per_name(caplog): + """Sibling of the per-latent NaN/warn tests — asserts the module-level + dedup set really suppresses repeat warnings for the same name.""" + _latent_module._MAGZERO_WARNED.discard("total_lens_flux_mujy") + caplog.set_level(logging.WARNING) + + fit = _fit_with_galaxy_images({0: [1.0, 2.0, 3.0, 4.0]}) + for _ in range(3): + total_lens_flux_mujy(fit=fit, magzero=None) + + matching = [ + r for r in caplog.records if "total_lens_flux_mujy" in r.message + ] + assert len(matching) == 1 def test_magnification_is_lensed_over_intrinsic(): @@ -242,6 +339,9 @@ def test_latent_keys_enabled_drops_unknown_with_warning(caplog): def test_latent_functions_registry_keys(): assert set(LATENT_FUNCTIONS) == { + "total_lens_flux", + "total_lensed_source_flux", + "total_source_flux", "total_lens_flux_mujy", "total_lensed_source_flux_mujy", "total_source_flux_mujy", @@ -272,9 +372,15 @@ def test_analysis_imaging_compute_latent_variables_aligns_with_keys( assert isinstance(values, tuple) assert len(values) == len(analysis.LATENT_KEYS) - # Only total_lens_flux_mujy is enabled in test_autolens/config/latent.yaml. - assert analysis.LATENT_KEYS == ["total_lens_flux_mujy"] - assert np.isfinite(values[0]) + # test_autolens/config/latent.yaml enables the three raw-flux keys plus + # total_lens_flux_mujy. magzero=25.0 above keeps the µJy column finite. + assert analysis.LATENT_KEYS == [ + "total_lens_flux", + "total_lensed_source_flux", + "total_source_flux", + "total_lens_flux_mujy", + ] + assert all(np.isfinite(v) for v in values) def test_analysis_imaging_compute_latent_variables_raises_when_empty(monkeypatch): @@ -291,7 +397,12 @@ def test_analysis_imaging_compute_latent_variables_raises_when_empty(monkeypatch def test_analysis_imaging_latent_keys_property_reads_config(): # The autouse `set_config_path` fixture in test_autolens/conftest.py - # pushes test_autolens/config/latent.yaml — only total_lens_flux_mujy - # is enabled there. + # pushes test_autolens/config/latent.yaml — the three raw-flux keys + # and total_lens_flux_mujy are enabled there. analysis = al.AnalysisImaging(dataset=MagicMock(), use_jax=False) - assert analysis.LATENT_KEYS == ["total_lens_flux_mujy"] + assert analysis.LATENT_KEYS == [ + "total_lens_flux", + "total_lensed_source_flux", + "total_source_flux", + "total_lens_flux_mujy", + ] diff --git a/test_autolens/config/latent.yaml b/test_autolens/config/latent.yaml index f0024d11c..1c3bf124e 100644 --- a/test_autolens/config/latent.yaml +++ b/test_autolens/config/latent.yaml @@ -2,11 +2,12 @@ # autouse `set_config_path` fixture in `test_autolens/conftest.py`, # shadowing the library defaults during the test suite. # -# Only `total_lens_flux_mujy` is enabled so the end-to-end test through -# `AnalysisImaging.compute_latent_variables` has the smallest possible -# fixture surface (lens galaxy with a single light profile). The other -# latent functions are exercised directly via their own unit tests with -# `SimpleNamespace`/`MagicMock`-built fits. +# All flux latents (raw and µJy) are `true` here (vs `_mujy: false` in the +# library default) so the test suite exercises both code paths through +# `AnalysisImaging.compute_latent_variables`. +total_lens_flux: true +total_lensed_source_flux: true +total_source_flux: true total_lens_flux_mujy: true total_lensed_source_flux_mujy: false total_source_flux_mujy: false