From 609338fc26bdd3be19e1c082314a45feb70c3953 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sat, 5 Sep 2026 13:21:44 +0000 Subject: [PATCH] fix: magnification latent finite for pixelized sources (#726) `total_source_flux` / `total_source_flux_mujy` read the source galaxy's `image_2d_from`, which is zeros for a galaxy whose only light model is a `Pixelization`, so every pixelized fit had `total_source_flux = 0.0` and `magnification = inf` (a dropped column, or an archived `0.0` sentinel). Add `_pixelized_source_flux`: the inversion's reconstruction integrated over the source mesh with its `areas_for_magnification`, summed for every mapper that `inversion.linear_obj_galaxy_dict` assigns to the source galaxy. Light-profile flux plus pixelized flux is now the source-plane flux; `magnification` keeps its form and becomes finite. A mesh geometry with no `areas_for_magnification` gives NaN with a one-time warning, never a silent zero. All branching is on Python structure, so the per-sample `jax.jit` latent path is unaffected. Tests: Delaunay and rectangular pixelized fits give a finite magnification equal to sum(image) / sum(reconstruction * areas); linear light profile + pixelization sums both terms; light-profile-only fits are unchanged; the no-inversion, no-source-mapper and missing-areas (NaN + one warning) branches are pinned. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011kyfKgDB1rMkcQsn19ow4T --- autolens/analysis/latent.py | 137 ++++++++++++++-- test_autolens/analysis/test_latent.py | 225 ++++++++++++++++++++++++++ 2 files changed, 351 insertions(+), 11 deletions(-) diff --git a/autolens/analysis/latent.py b/autolens/analysis/latent.py index e48352bc7..e4598bfcf 100644 --- a/autolens/analysis/latent.py +++ b/autolens/analysis/latent.py @@ -23,6 +23,7 @@ import numpy as np +import autoarray as aa import autofit as af from autonerves import conf from autogalaxy.imaging.model.latent import ( @@ -43,6 +44,11 @@ # single search performs. _JAX_ZERO_CONTOUR_FALLBACK_WARNED: bool = False +# Set to True the first time ``_pixelized_source_flux`` meets a mapper whose +# mesh geometry exposes no ``areas_for_magnification``. Deduplicates the +# warning across the many fit evaluations a single search performs. +_MESH_AREAS_WARNED: bool = False + def _jax_zero_contour_available() -> bool: """ @@ -87,6 +93,89 @@ def _maybe_magzero_warn(magzero, name) -> bool: return False +def _pixelized_source_flux(fit, xp=np): + """ + Source-plane integrated flux of the source galaxy's *pixelized* + (inversion) component, in raw image units. + + ``Galaxy.image_2d_from`` returns zeros for a galaxy whose only light + model is a ``Pixelization`` — the reconstruction lives on the inversion's + mesh, not on a light profile. This helper integrates that reconstruction + over the mesh: ``Σ_i s_i A_i``, where ``s_i`` are the reconstructed + source-pixel values (``inversion.reconstruction_dict[mapper]``) and + ``A_i`` the exact per-pixel areas the mesh geometry publishes as + ``areas_for_magnification`` (barycentric dual areas for a Delaunay mesh — + PyAutoArray#524 — transformed cell areas for a rectangular mesh). + + The mapper → galaxy association is read from + ``inversion.linear_obj_galaxy_dict`` — the dict the inversion is given at + construction (``autolens/lens/to_inversion.py``), whose galaxy values are + the tracer's own galaxy objects. ``fit.tracer_to_inversion`` must *not* be + used here: on ``FitImaging`` it is an uncached property that rebuilds a + ``TracerToInversion`` per call, whose mappers are new objects and + therefore not keys of ``reconstruction_dict``. + + Returns ``0.0`` when the fit has no inversion or no mapper belongs to the + source galaxy (the light-profile path already covers those fits), and + ``NaN`` — with a one-time-per-process warning — when a mapper's mesh + geometry publishes no exact quadrature areas. + + All branching is on Python structure (``None`` checks, ``isinstance``, + dict membership), never on array values, so the function is safe inside + the per-sample ``jax.jit`` the latent batch path uses. + """ + global _MESH_AREAS_WARNED + + inversion = getattr(fit, "inversion", None) + if inversion is None: + return 0.0 + + linear_obj_galaxy_dict = getattr(inversion, "linear_obj_galaxy_dict", None) + if not linear_obj_galaxy_dict: + return 0.0 + + try: + source_galaxy = fit.tracer.galaxies[-1] + except (AttributeError, IndexError): + return 0.0 + + total = 0.0 + found_mapper = False + + for linear_obj, galaxy in linear_obj_galaxy_dict.items(): + if galaxy is not source_galaxy: + continue + if not isinstance(linear_obj, aa.Mapper): + continue + + areas = getattr( + getattr(linear_obj, "mesh_geometry", None), + "areas_for_magnification", + None, + ) + if areas is None: + if not _MESH_AREAS_WARNED: + logger.warning( + "The mesh geometry '%s' publishes no " + "'areas_for_magnification', so the source-plane flux of " + "the pixelized source cannot be integrated; " + "'total_source_flux' and 'magnification' will be NaN.", + type(getattr(linear_obj, "mesh_geometry", None)).__name__, + ) + _MESH_AREAS_WARNED = True + return xp.nan + + found_mapper = True + total = total + xp.sum( + xp.asarray(inversion.reconstruction_dict[linear_obj]) * xp.asarray(areas) + ) + + if not found_mapper: + return 0.0 + + return total + + def total_lens_flux(fit, magzero=None, xp=np): """ Total integrated flux of the lens galaxy (``fit.tracer.galaxies[0]``), @@ -122,10 +211,20 @@ 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`. + Two contributions are summed, so that a source modelled with light + profiles, a pixelization, or both is measured correctly: + + - **Light profiles** — read 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. + - **Pixelization** — the inversion's reconstruction integrated over the + source mesh with its exact per-pixel ``areas_for_magnification`` + (barycentric dual areas for a Delaunay mesh, PyAutoArray#524; + transformed cell areas for a rectangular mesh); see :func:`_pixelized_source_flux`. Without + this term a purely pixelized source has ``image_2d_from`` zeros and a + zero source flux, making :func:`magnification` infinite + (PyAutoLens#726). ``magzero`` is accepted but ignored. """ @@ -136,7 +235,7 @@ def total_source_flux(fit, magzero=None, xp=np): ) except (AttributeError, IndexError): return xp.nan - return xp.sum(source_image.array) + return xp.sum(source_image.array) + _pixelized_source_flux(fit=fit, xp=xp) def total_lens_flux_mujy(fit, magzero, xp=np): @@ -190,11 +289,19 @@ def total_source_flux_mujy(fit, magzero, xp=np): """ Source-plane intrinsic flux of the source galaxy, in microjanskies. - Reads from ``fit.tracer_linear_light_profiles_to_light_profiles`` rather - than ``fit.tracer`` so that linear light profiles (whose ``intensity`` - 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 + Same definition as :func:`total_source_flux` — light-profile flux plus + the pixelized (inversion) flux — magzero-converted. + + The light-profile term reads from + ``fit.tracer_linear_light_profiles_to_light_profiles`` rather than + ``fit.tracer`` so that linear light profiles (whose ``intensity`` 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. + The pixelized term integrates the inversion's reconstruction over the + source mesh with its ``areas_for_magnification`` (barycentric dual + areas for a Delaunay mesh — PyAutoArray#524 — transformed cell areas for + a rectangular mesh); see :func:`_pixelized_source_flux`. Returns NaN + one warning when ``magzero`` is missing; see :func:`total_lens_flux_mujy` for the rationale. @@ -208,7 +315,7 @@ def total_source_flux_mujy(fit, magzero, xp=np): ) except (AttributeError, IndexError): return xp.nan - total_flux = xp.sum(source_image.array) + total_flux = xp.sum(source_image.array) + _pixelized_source_flux(fit=fit, xp=xp) return flux_mujy_via_ab_mag_from( ab_mag=ab_mag_via_flux_from(flux=total_flux, magzero=magzero, xp=xp), xp=xp, @@ -218,7 +325,15 @@ def total_source_flux_mujy(fit, magzero, xp=np): def magnification(fit, magzero, xp=np): """ Ratio of image-plane to source-plane source flux — the integrated - magnification implied by the lens model and source light profile. + magnification implied by the lens model and the source's light. + + The denominator is :func:`total_source_flux_mujy`: light-profile flux + plus, for a pixelized source, the inversion's reconstruction integrated + over the source mesh with its ``areas_for_magnification`` (barycentric + dual areas for a Delaunay mesh — PyAutoArray#524 — transformed cell + areas for a rectangular mesh). A source whose only light model is a ``Pixelization`` + therefore has a finite magnification rather than an infinite one + (PyAutoLens#726). ``magzero`` is accepted but unused (the µJy conversions cancel in the ratio). It's still required in the signature so the dispatcher can diff --git a/test_autolens/analysis/test_latent.py b/test_autolens/analysis/test_latent.py index eff5f5414..f954d6671 100644 --- a/test_autolens/analysis/test_latent.py +++ b/test_autolens/analysis/test_latent.py @@ -6,12 +6,18 @@ import numpy as np import pytest +import autoarray as aa import autofit as af import autolens as al +from autogalaxy.imaging.model.latent import ( + ab_mag_via_flux_from, + flux_mujy_via_ab_mag_from, +) from autolens.analysis import latent as _latent_module from autolens.analysis.latent import ( LATENT_FUNCTIONS, LatentLens, + _pixelized_source_flux, effective_einstein_radius, latent_keys_enabled, magnification, @@ -356,6 +362,225 @@ def _raise(cls, tracer): assert np.isnan(effective_einstein_radius(fit=fit, magzero=None)) +# --------------------------------------------------------------------------- +# Pixelized (inversion) source flux — PyAutoLens#726 +# --------------------------------------------------------------------------- + +def _lens_galaxy(): + return al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal(centre=(0.1, 0.1), einstein_radius=1.0), + ) + + +def _delaunay_fit(dataset): + """A real ``FitImaging`` whose source galaxy's only light model is a + Delaunay pixelization. The Delaunay mesh does not build its own + image-plane mesh grid, so one is supplied via ``adapt_images``.""" + pixelization = al.Pixelization( + mesh=al.mesh.Delaunay(pixels=9, zeroed_pixels=0), + regularization=al.reg.Constant(coefficient=0.01), + ) + source = al.Galaxy(redshift=1.0, pixelization=pixelization) + + image_plane_mesh_grid = al.image_mesh.Overlay( + shape=(3, 3) + ).image_plane_mesh_grid_from(mask=dataset.mask) + + adapt_images = al.AdaptImages( + galaxy_image_dict={source: dataset.data}, + galaxy_image_plane_mesh_grid_dict={source: image_plane_mesh_grid}, + ) + + fit = al.FitImaging( + dataset=dataset, + tracer=al.Tracer(galaxies=[_lens_galaxy(), source]), + adapt_images=adapt_images, + ) + return fit, source + + +def _rectangular_fit(dataset): + """A real ``FitImaging`` whose source galaxy's only light model is a + rectangular pixelization (which builds its own mesh grid).""" + pixelization = al.Pixelization( + mesh=al.mesh.RectangularUniform(shape=(3, 3)), + regularization=al.reg.Constant(coefficient=0.01), + ) + source = al.Galaxy(redshift=1.0, pixelization=pixelization) + + fit = al.FitImaging( + dataset=dataset, + tracer=al.Tracer(galaxies=[_lens_galaxy(), source]), + ) + return fit, source + + +def _mesh_integrated_flux(fit, source): + """Σ_mappers Σ_i s_i A_i for the mappers belonging to ``source`` — + computed here independently of the library implementation.""" + inversion = fit.inversion + total = 0.0 + for linear_obj, galaxy in inversion.linear_obj_galaxy_dict.items(): + if galaxy is not source or not isinstance(linear_obj, aa.Mapper): + continue + total += np.sum( + np.asarray(inversion.reconstruction_dict[linear_obj]) + * np.asarray(linear_obj.mesh_geometry.areas_for_magnification) + ) + return total + + +def test_pixelized_source_flux_zero_when_fit_has_no_inversion(): + """The light-profile path already covers a fit with no inversion, so the + pixelized term must contribute exactly nothing (not NaN, not a sentinel).""" + fit = SimpleNamespace( + tracer=SimpleNamespace(galaxies=[object(), object()]), + inversion=None, + ) + assert _pixelized_source_flux(fit=fit) == 0.0 + + # A fit object with no ``inversion`` attribute at all behaves identically. + fit_no_attr = SimpleNamespace( + tracer=SimpleNamespace(galaxies=[object(), object()]) + ) + assert _pixelized_source_flux(fit=fit_no_attr) == 0.0 + + +def test_delaunay_pixelized_source_magnification_is_finite(masked_imaging_7x7): + """Regression for PyAutoLens#726: a source whose only light model is a + pixelization had zero source-plane flux and therefore infinite + magnification.""" + fit, source = _delaunay_fit(masked_imaging_7x7) + + denominator = _mesh_integrated_flux(fit=fit, source=source) + lensed = np.sum(fit.galaxy_image_dict[source].array) + + assert denominator != 0.0 + assert total_source_flux(fit=fit) == pytest.approx(denominator) + assert _pixelized_source_flux(fit=fit) == pytest.approx(denominator) + + value = magnification(fit=fit, magzero=25.0) + assert np.isfinite(value) + assert value == pytest.approx(lensed / denominator) + + +def test_rectangular_pixelized_source_magnification_is_finite(masked_imaging_7x7): + fit, source = _rectangular_fit(masked_imaging_7x7) + + denominator = _mesh_integrated_flux(fit=fit, source=source) + lensed = np.sum(fit.galaxy_image_dict[source].array) + + assert denominator != 0.0 + assert total_source_flux(fit=fit) == pytest.approx(denominator) + + value = magnification(fit=fit, magzero=25.0) + assert np.isfinite(value) + assert value == pytest.approx(lensed / denominator) + + +def test_total_source_flux_sums_light_profile_and_pixelization(masked_imaging_7x7): + """A source carrying BOTH a linear light profile and a pixelization must + contribute both terms — neither may swallow the other.""" + pixelization = al.Pixelization( + mesh=al.mesh.RectangularUniform(shape=(3, 3)), + regularization=al.reg.Constant(coefficient=0.01), + ) + source = al.Galaxy( + redshift=1.0, + light=al.lp_linear.Sersic(), + pixelization=pixelization, + ) + fit = al.FitImaging( + dataset=masked_imaging_7x7, + tracer=al.Tracer(galaxies=[_lens_galaxy(), source]), + ) + + tracer = fit.tracer_linear_light_profiles_to_light_profiles + light_profile_flux = np.sum( + tracer.galaxies[-1].image_2d_from(grid=fit.dataset.grids.lp).array + ) + pixelized_flux = _mesh_integrated_flux(fit=fit, source=source) + + assert light_profile_flux != 0.0 + assert pixelized_flux != 0.0 + assert _pixelized_source_flux(fit=fit) == pytest.approx(pixelized_flux) + assert total_source_flux(fit=fit) == pytest.approx( + light_profile_flux + pixelized_flux + ) + + +def test_light_profile_only_fit_source_flux_unchanged(fit_imaging_x2_plane_7x7): + """The light-profile-only path is untouched: the pixelized term is exactly + zero, so ``total_source_flux`` is still the plain light-profile sum.""" + fit = fit_imaging_x2_plane_7x7 + + tracer = fit.tracer_linear_light_profiles_to_light_profiles + light_profile_flux = np.sum( + tracer.galaxies[-1].image_2d_from(grid=fit.dataset.grids.lp).array + ) + + assert fit.inversion is None + assert _pixelized_source_flux(fit=fit) == 0.0 + assert total_source_flux(fit=fit) == pytest.approx(light_profile_flux) + assert total_source_flux_mujy(fit=fit, magzero=25.0) == pytest.approx( + flux_mujy_via_ab_mag_from( + ab_mag=ab_mag_via_flux_from(flux=light_profile_flux, magzero=25.0) + ) + ) + assert np.isfinite(magnification(fit=fit, magzero=25.0)) + + +class _NoAreasMapper(aa.Mapper): + """A mapper whose mesh geometry publishes no ``areas_for_magnification`` + (no exact quadrature areas), bypassing ``Mapper.__init__``.""" + + def __init__(self): + pass + + @property + def mesh_geometry(self): + return SimpleNamespace() + + +def test_pixelized_source_flux_nan_and_one_warning_without_mesh_areas(caplog): + _latent_module._MESH_AREAS_WARNED = False + caplog.set_level(logging.WARNING) + + source = object() + mapper = _NoAreasMapper() + fit = SimpleNamespace( + tracer=SimpleNamespace(galaxies=[object(), source]), + inversion=SimpleNamespace( + linear_obj_galaxy_dict={mapper: source}, + reconstruction_dict={mapper: np.ones(4)}, + ), + ) + + values = [_pixelized_source_flux(fit=fit) for _ in range(3)] + + assert all(np.isnan(v) for v in values) + matching = [r for r in caplog.records if "areas_for_magnification" in r.message] + assert len(matching) == 1 + + +def test_pixelized_source_flux_zero_when_no_mapper_belongs_to_source(): + """A mapper belonging to another galaxy (e.g. a pixelized lens) must not + leak into the source-plane flux.""" + source = object() + other_galaxy = object() + mapper = _NoAreasMapper() + fit = SimpleNamespace( + tracer=SimpleNamespace(galaxies=[other_galaxy, source]), + inversion=SimpleNamespace( + linear_obj_galaxy_dict={mapper: other_galaxy}, + reconstruction_dict={mapper: np.ones(4)}, + ), + ) + + assert _pixelized_source_flux(fit=fit) == 0.0 + + # --------------------------------------------------------------------------- # Registry / config-reader tests # ---------------------------------------------------------------------------