diff --git a/autolens/analysis/latent.py b/autolens/analysis/latent.py index e4598bfcf..3b3afcd1f 100644 --- a/autolens/analysis/latent.py +++ b/autolens/analysis/latent.py @@ -49,6 +49,12 @@ # warning across the many fit evaluations a single search performs. _MESH_AREAS_WARNED: bool = False +# Set to True the first time ``_pixelized_source_flux`` cannot resolve the +# data's pixel area, and so cannot express the mesh-integrated flux in this +# module's per-data-pixel convention. Deduplicates that warning across the +# many fit evaluations a single search performs. +_PIXEL_AREA_WARNED: bool = False + def _jax_zero_contour_available() -> bool: """ @@ -93,19 +99,61 @@ def _maybe_magzero_warn(magzero, name) -> bool: return False +def _data_pixel_area(fit) -> Optional[float]: + """ + The data's pixel area ``A_pix`` in arcsec², or ``None`` if it cannot be + resolved from the fit. + + Read from ``fit.dataset.grids.lp`` — the *same* grid the light-profile + flux term of :func:`total_source_flux` is evaluated on — so both terms of + that sum share one convention. ``Grid2D.pixel_area`` is the product of the + grid's ``pixel_scales`` and is a plain Python float, so dividing by it + introduces no array op and is trace-safe inside ``jax.jit``. A grid type + that publishes no ``pixel_area`` falls back to multiplying out + ``pixel_scales``, on the grid and then on the dataset itself. + """ + dataset = getattr(fit, "dataset", None) + grid = getattr(getattr(dataset, "grids", None), "lp", None) + + pixel_area = getattr(grid, "pixel_area", None) + if pixel_area is not None: + return float(pixel_area) + + for obj in (grid, dataset): + pixel_scales = getattr(obj, "pixel_scales", None) + if pixel_scales is not None: + return float(pixel_scales[0]) * float(pixel_scales[1]) + + return None + + def _pixelized_source_flux(fit, xp=np): """ Source-plane integrated flux of the source galaxy's *pixelized* - (inversion) component, in raw image units. + (inversion) component, in raw image units **per data pixel**. ``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 + over the mesh and re-expresses it in this module's per-data-pixel flux + convention: ``Σ_i s_i A_i / A_pix``, where ``s_i`` are the reconstructed + source-pixel values (``inversion.reconstruction_dict[mapper]``), ``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). + PyAutoArray#524 — transformed cell areas for a rectangular mesh) and + ``A_pix`` the data pixel area (:func:`_data_pixel_area`). + + **Why the division by ``A_pix``.** Every other flux latent here is a sum + over *data pixels* — ``Σ galaxy_image_dict[galaxy]``, + ``Σ galaxy.image_2d_from(grid)`` — i.e. ``∫ I d²θ / A_pix``. The + reconstruction ``s_i`` is in those same per-data-pixel units (the mapping + matrix maps ``s`` onto image-pixel values), so ``Σ_i s_i A_i`` is the + plain surface integral ``∫ S d²β`` and must be divided by ``A_pix`` to + join the convention. Only then is :func:`magnification` — + ``Σ image / (Σ s A / A_pix)`` — the physical ``∫ I d²θ / ∫ S d²β``. + Undivided it scales with the data's pixel scale: on a 0.1"/pixel dataset + the pixelized magnification came out exactly 100× too large. The 1" + fixtures have ``A_pix = 1``, which is why this was invisible in tests. The mapper → galaxy association is read from ``inversion.linear_obj_galaxy_dict`` — the dict the inversion is given at @@ -118,13 +166,16 @@ def _pixelized_source_flux(fit, xp=np): 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. + geometry publishes no exact quadrature areas, or when the data pixel area + cannot be resolved (an unscaled number would silently be wrong by the + pixel area). 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 + global _PIXEL_AREA_WARNED inversion = getattr(fit, "inversion", None) if inversion is None: @@ -173,7 +224,19 @@ def _pixelized_source_flux(fit, xp=np): if not found_mapper: return 0.0 - return total + pixel_area = _data_pixel_area(fit) + if pixel_area is None: + if not _PIXEL_AREA_WARNED: + logger.warning( + "The data pixel area could not be resolved from the fit " + "(no 'fit.dataset.grids.lp.pixel_area' and no 'pixel_scales'), " + "so the pixelized source flux cannot be expressed per data " + "pixel; 'total_source_flux' and 'magnification' will be NaN." + ) + _PIXEL_AREA_WARNED = True + return xp.nan + + return total / pixel_area def total_lens_flux(fit, magzero=None, xp=np): @@ -218,13 +281,14 @@ def total_source_flux(fit, magzero=None, xp=np): ``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). + - **Pixelization** — the inversion's reconstruction integrated with its + exact per-pixel ``areas_for_magnification`` (barycentric dual areas for + a Delaunay mesh, PyAutoArray#524; transformed cell areas for a + rectangular mesh) and expressed per data pixel (divided by the data + pixel area), which is the convention the light-profile term above is + already in; 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. """ @@ -298,10 +362,11 @@ def total_source_flux_mujy(fit, magzero, xp=np): 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`. + The pixelized term integrates the inversion's reconstruction with its + ``areas_for_magnification`` (barycentric dual areas for a Delaunay mesh — + PyAutoArray#524 — transformed cell areas for a rectangular mesh) and + expresses it per data pixel (divided by the data pixel area), matching the + light-profile term's convention; see :func:`_pixelized_source_flux`. Returns NaN + one warning when ``magzero`` is missing; see :func:`total_lens_flux_mujy` for the rationale. @@ -329,11 +394,13 @@ def magnification(fit, magzero, xp=np): 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). + with its ``areas_for_magnification`` (barycentric dual areas for a + Delaunay mesh — PyAutoArray#524 — transformed cell areas for a rectangular + mesh) and expressed per data pixel (divided by the data pixel area). Both + numerator and denominator are then per-data-pixel sums, so the ratio is + the physical ``∫ I d²θ / ∫ S d²β`` and does not depend on the data's pixel + scale. 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 f954d6671..67f2b7e34 100644 --- a/test_autolens/analysis/test_latent.py +++ b/test_autolens/analysis/test_latent.py @@ -373,6 +373,22 @@ def _lens_galaxy(): ) +# The scene the pixel-scale-stability test re-observes at two pixel scales. +_SCALE_TEST_LENS = al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.0), +) +_SCALE_TEST_SOURCE = al.Galaxy( + redshift=1.0, + light=al.lp.Sersic( + centre=(0.05, 0.05), + intensity=1.0, + effective_radius=0.2, + sersic_index=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 @@ -416,9 +432,43 @@ def _rectangular_fit(dataset): return fit, source +_MASK_7X7 = np.array( + [ + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + [True, True, False, False, False, True, True], + [True, True, False, False, False, True, True], + [True, True, False, False, False, True, True], + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + ] +) + + +def _masked_imaging_7x7_at(pixel_scales): + """The ``masked_imaging_7x7`` fixture's geometry at arbitrary pixel + scales. The fixture itself is 1.0"/pixel, where a division by the data + pixel area is numerically invisible.""" + dataset = aa.Imaging( + data=aa.Array2D.no_mask(values=np.ones((7, 7)), pixel_scales=pixel_scales), + psf=aa.Convolver( + kernel=aa.Array2D.no_mask( + values=[[0.0, 0.5, 0.0], [0.5, 1.0, 0.5], [0.0, 0.5, 0.0]], + pixel_scales=pixel_scales, + ) + ), + noise_map=aa.Array2D.no_mask( + values=2.0 * np.ones((7, 7)), pixel_scales=pixel_scales + ), + over_sample_size_lp=1, + ) + return dataset.apply_mask(mask=aa.Mask2D(mask=_MASK_7X7, pixel_scales=pixel_scales)) + + 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.""" + """Σ_mappers Σ_i s_i A_i for the mappers belonging to ``source`` — the + plain source-plane surface integral ∫ S d²β, computed here independently + of the library implementation.""" inversion = fit.inversion total = 0.0 for linear_obj, galaxy in inversion.linear_obj_galaxy_dict.items(): @@ -431,6 +481,14 @@ def _mesh_integrated_flux(fit, source): return total +def _expected_pixelized_flux(fit, source): + """Σ_i s_i A_i / A_pix — the pixelized source flux in the module's + per-data-pixel convention, the one every other flux latent uses. ``A_pix`` + is read from the same grid the light-profile term is evaluated on.""" + mesh_integral = _mesh_integrated_flux(fit=fit, source=source) + return mesh_integral / fit.dataset.grids.lp.pixel_area + + 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).""" @@ -453,7 +511,7 @@ def test_delaunay_pixelized_source_magnification_is_finite(masked_imaging_7x7): magnification.""" fit, source = _delaunay_fit(masked_imaging_7x7) - denominator = _mesh_integrated_flux(fit=fit, source=source) + denominator = _expected_pixelized_flux(fit=fit, source=source) lensed = np.sum(fit.galaxy_image_dict[source].array) assert denominator != 0.0 @@ -468,7 +526,7 @@ def test_delaunay_pixelized_source_magnification_is_finite(masked_imaging_7x7): 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) + denominator = _expected_pixelized_flux(fit=fit, source=source) lensed = np.sum(fit.galaxy_image_dict[source].array) assert denominator != 0.0 @@ -500,7 +558,7 @@ def test_total_source_flux_sums_light_profile_and_pixelization(masked_imaging_7x 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) + pixelized_flux = _expected_pixelized_flux(fit=fit, source=source) assert light_profile_flux != 0.0 assert pixelized_flux != 0.0 @@ -510,6 +568,132 @@ def test_total_source_flux_sums_light_profile_and_pixelization(masked_imaging_7x ) +def test_pixelized_source_flux_is_per_data_pixel(): + """The pixelized term is Σ_i s_i A_i / A_pix, not Σ_i s_i A_i — it must + join the per-data-pixel convention every other flux latent uses. + + At the 1" fixtures A_pix = 1 and the two are indistinguishable, so this + test builds the same fit on 0.5"/pixel data (A_pix = 0.25), where the + division is a factor of 4.""" + dataset = _masked_imaging_7x7_at(pixel_scales=(0.5, 0.5)) + assert dataset.grids.lp.pixel_area == pytest.approx(0.25) + + fit, source = _rectangular_fit(dataset) + + mesh_integral = _mesh_integrated_flux(fit=fit, source=source) + assert mesh_integral != 0.0 + + expected = mesh_integral / 0.25 + assert _pixelized_source_flux(fit=fit) == pytest.approx(expected) + assert total_source_flux(fit=fit) == pytest.approx(expected) + + # The division is real, not a no-op inherited from a 1" fixture. + assert total_source_flux(fit=fit) != pytest.approx(mesh_integral) + + lensed = np.sum(fit.galaxy_image_dict[source].array) + assert magnification(fit=fit, magzero=25.0) == pytest.approx(lensed / expected) + + +def _simulated_lens_dataset(pixel_scales, shape_native): + """A noise-free image of ``_SCALE_TEST_LENS`` lensing ``_SCALE_TEST_SOURCE``, + masked to a 2" circle — the same physical scene at whatever pixel scale is + asked for, so long as ``pixel_scales * shape_native`` is held fixed.""" + grid = al.Grid2D.uniform(shape_native=shape_native, pixel_scales=pixel_scales) + tracer = al.Tracer(galaxies=[_SCALE_TEST_LENS, _SCALE_TEST_SOURCE]) + + dataset = aa.Imaging( + data=aa.Array2D( + values=tracer.image_2d_from(grid=grid).native, + mask=aa.Mask2D.all_false( + shape_native=shape_native, pixel_scales=pixel_scales + ), + ), + psf=aa.Convolver( + kernel=aa.Array2D.no_mask( + values=[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]], + pixel_scales=pixel_scales, + ) + ), + noise_map=aa.Array2D.no_mask( + values=0.01 * np.ones(shape_native), pixel_scales=pixel_scales + ), + over_sample_size_lp=1, + ) + return dataset.apply_mask( + mask=aa.Mask2D.circular( + shape_native=shape_native, pixel_scales=pixel_scales, radius=2.0 + ) + ) + + +def test_pixelized_magnification_is_stable_across_pixel_scales(): + """The physical statement behind the per-data-pixel convention: fit the + *same* scene at two pixel scales and the pixelized magnification stays put. + + What it pins is deliberately loose — a rectangular reconstruction of a + compact Sersic over a 6.4" field is coarse, so only the order of magnitude + is meaningful and the window here is a factor of two. It is exactly what + the unscaled ``Sum_i s_i A_i`` violated: undivided, the two answers differ + by the pixel-area ratio, 4x, which the last assertion checks is outside + the window (so the window is not vacuous). + + It pins nothing about agreement with a light-profile source's + magnification, which at this coarseness is not itself scale-stable. + """ + pixelization = al.Pixelization( + mesh=al.mesh.RectangularUniform(shape=(20, 20)), + regularization=al.reg.Constant(coefficient=1.0), + ) + + values = {} + for pixel_scales, shape_native in (((0.4, 0.4), (16, 16)), ((0.2, 0.2), (32, 32))): + dataset = _simulated_lens_dataset( + pixel_scales=pixel_scales, shape_native=shape_native + ) + fit = al.FitImaging( + dataset=dataset, + tracer=al.Tracer( + galaxies=[ + _SCALE_TEST_LENS, + al.Galaxy(redshift=1.0, pixelization=pixelization), + ] + ), + ) + values[pixel_scales[0]] = magnification(fit=fit, magzero=25.0) + + assert all(np.isfinite(v) for v in values.values()) + + ratio = values[0.4] / values[0.2] + assert 0.5 < ratio < 2.0 + + # Without the division by the data pixel area the same two fits would + # differ by the ratio of pixel areas (0.16 / 0.04), i.e. off the window. + unscaled_ratio = ratio * (0.16 / 0.04) + assert not 0.5 < unscaled_ratio < 2.0 + + +def test_pixelized_source_flux_nan_when_pixel_area_unresolvable(caplog): + """A fit whose dataset publishes neither a grid ``pixel_area`` nor + ``pixel_scales`` yields NaN plus one warning per process — never a + silently unscaled number.""" + _latent_module._PIXEL_AREA_WARNED = False + caplog.set_level(logging.WARNING) + + dataset = _masked_imaging_7x7_at(pixel_scales=(0.5, 0.5)) + fit, _ = _rectangular_fit(dataset) + fit_no_pixel_area = SimpleNamespace( + tracer=fit.tracer, + inversion=fit.inversion, + dataset=SimpleNamespace(grids=SimpleNamespace(lp=SimpleNamespace())), + ) + + values = [_pixelized_source_flux(fit=fit_no_pixel_area) for _ in range(3)] + + assert all(np.isnan(v) for v in values) + matching = [r for r in caplog.records if "data pixel area" in r.message] + assert len(matching) == 1 + + 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."""