From e32a5abe830b0529cd7585294ee925abcb1f2ddf Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 17 Jul 2026 11:50:31 +0100 Subject: [PATCH] Use paths.preserve_in_zip for the galaxy-image cache Switches the resume fast-path cache write to the public AbstractPaths.preserve_in_zip (PyAutoFit#1390) and deletes the private _append_to_search_zip helper it replaced (phase 2 of PyAutoFit#1389). Behaviour identical; the test stub gains the no-op method. Co-Authored-By: Claude Fable 5 --- .../analysis/adapt_images/adapt_images.py | 881 +++++++++--------- test_autogalaxy/analysis/test_adapt_images.py | 219 ++--- 2 files changed, 540 insertions(+), 560 deletions(-) diff --git a/autogalaxy/analysis/adapt_images/adapt_images.py b/autogalaxy/analysis/adapt_images/adapt_images.py index 09ff8b12..56901f00 100644 --- a/autogalaxy/analysis/adapt_images/adapt_images.py +++ b/autogalaxy/analysis/adapt_images/adapt_images.py @@ -1,452 +1,429 @@ -from __future__ import annotations -import numpy as np -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple - -from autoconf import conf -from autoconf import cached_property - -import autoarray as aa - -if TYPE_CHECKING: - from autogalaxy.galaxy.galaxy import Galaxy - - -def _galaxy_images_cache_path(result, use_model_images: bool): - """ - The on-disk cache file for the raw per-galaxy images of a result, inside the - result's own ``files/`` folder, or ``None`` when the result has no on-disk - output (e.g. ``NullPaths``) and therefore cannot cache. - """ - from pathlib import Path - - paths = getattr(result, "paths", None) - files_path = getattr(paths, "_files_path", None) - if files_path is None or not Path(files_path).is_dir(): - return None - name = "galaxy_images_model" if use_model_images else "galaxy_images_snr" - return Path(files_path) / f"{name}.fits" - - -def _galaxy_image_dict_from_cache(cache_path) -> Optional[Dict]: - """ - Load the raw (pre minimum-percent clip) per-galaxy image dictionary from a - result's cache file, or ``None`` when the file does not exist (the first - arrival at this result computes and writes it). - - The FITS layout mirrors the ``adapt_images.fits`` artifact the aggregator - reads (``agg_util.adapt_images_from``): HDU 0 is the mask (header carries - pixel scales and origin), HDU 1+ are one image per galaxy with the galaxy - path as ``EXTNAME``. - """ - from astropy.io import fits as astropy_fits - - from autoarray.mask.mask_2d import Mask2DKeys - from autoconf.fitsable import ndarray_via_hdu_from - - if cache_path is None or not cache_path.exists(): - return None - - with astropy_fits.open(cache_path) as hdu_list: - header = hdu_list[0].header - pixel_scales = ( - header[Mask2DKeys.PIXSCAY.value], - header[Mask2DKeys.PIXSCAX.value], - ) - origin = ( - header[Mask2DKeys.ORIGINY.value], - header[Mask2DKeys.ORIGINX.value], - ) - mask = aa.Mask2D( - mask=ndarray_via_hdu_from(hdu_list[0]), - pixel_scales=pixel_scales, - origin=origin, - ) - - galaxy_name_image_dict = {} - for hdu in hdu_list[1:]: - image = aa.Array2D.no_mask( - values=ndarray_via_hdu_from(hdu), - pixel_scales=mask.pixel_scales, - origin=mask.origin, - ) - galaxy_name_image_dict[hdu.header["EXTNAME"].lower()] = image.apply_mask( - mask=mask - ) - - return galaxy_name_image_dict - - -def _append_to_search_zip(paths, file_path): - """ - Also add a cache file into the search's ``.zip`` archive. - - A resumed search's ``paths.restore()`` deletes the output directory and - re-extracts the zip, so a cache written only to ``files/`` after the search - completed would be destroyed by the next resume. Appending it to the zip - makes it a permanent part of the completed output (each later resume - re-extracts and re-zips it with everything else). - """ - import zipfile - from pathlib import Path - - zip_path = getattr(paths, "_zip_path", None) - output_path = getattr(paths, "output_path", None) - if zip_path is None or output_path is None or not Path(zip_path).exists(): - return - arcname = str(Path(file_path).relative_to(output_path)) - with zipfile.ZipFile(zip_path, "a") as f: - if arcname not in f.namelist(): - f.write(file_path, arcname) - - -def _galaxy_image_dict_to_cache(cache_path, galaxy_name_image_dict: Dict, paths): - """ - Persist the raw per-galaxy image dictionary to the result's cache file, in - the same FITS layout ``_galaxy_image_dict_from_cache`` reads, and preserve - it in the search's zip archive so later resumes keep it. - """ - from autoconf.fitsable import hdu_list_for_output_from - - image_list = [ - galaxy_name_image_dict[name].native_for_fits - for name in galaxy_name_image_dict - ] - hdu_list = hdu_list_for_output_from( - values_list=[image_list[0].mask.astype("float")] + image_list, - ext_name_list=["mask"] + list(galaxy_name_image_dict.keys()), - header_dict=next(iter(galaxy_name_image_dict.values())).mask.header_dict, - ) - hdu_list.writeto(cache_path, overwrite=True) - _append_to_search_zip(paths, cache_path) - - -def galaxy_name_image_dict_via_result_from( - result, use_model_images: bool = False -) -> "AdaptImages": - """ - Returns the adapt-images from a non-linear search result. - - For model-fitting, the adapt-images are typically setup using the maximum log likelihood model of the - previous model-fit. This means the model-fitting is used to cleanly deblend the light of the different - galaxies in the image (e.g. separate the lens light from the source light). - - This method uses attributes of a result (e.g. dictionary mapping galaxy instances to their model-images) - to create the adapt-images. - - This can use either: - - - The model image of each galaxy in the best-fit model. - - The subtracted image of each galaxy in the best-fit model, where the subtracted image is the dataset - minus the model images of all other galaxies. - - Certain models produce galaxy-images with negative flux values (e.g. a pixelization), which can cause - numerical issues with the adaptive schemes. To prevent this, we set a minimum flux value for each - galaxy-image, which is a fraction of the maximum flux value of that image defined via a config file. - - The raw per-galaxy images are cached to the result's own ``files/`` folder on first computation - (``galaxy_images_model.fits`` / ``galaxy_images_snr.fits``) and loaded from there on every later call — - computing them rebuilds the result's maximum log likelihood fit, which on a resumed pipeline pays a fresh - JIT compile plus (for pixelized fits) an inversion, and dominates SLaM resume overhead - (autolens_profiling#70). Staleness is structurally guarded: changing the upstream model or search - produces a new search identifier and therefore a fresh output directory with no cache file. Results with - no on-disk output (e.g. ``NullPaths``) always compute. - - Parameters - ---------- - result - The result of a previous model-fit, which contains the model-image of each galaxy. - use_model_images - If True, the model images of the galaxies are used to create the adapt images. If False, the subtracted - images of the galaxies are used. - - Returns - ------- - The adapt-images, which are the model-image of each galaxy inferred via the previous model-fit. - """ - adapt_minimum_percent = conf.instance["general"]["adapt"]["adapt_minimum_percent"] - - cache_path = _galaxy_images_cache_path(result, use_model_images=use_model_images) - raw_image_dict = _galaxy_image_dict_from_cache(cache_path) - - if raw_image_dict is None: - raw_image_dict = {} - - for path, galaxy in result.path_galaxy_tuples: - if use_model_images: - raw_image_dict[path] = result.model_image_galaxy_dict[path] - else: - raw_image_dict[path] = result.subtracted_signal_to_noise_map_galaxy_dict[ - path - ] - - if cache_path is not None: - _galaxy_image_dict_to_cache(cache_path, raw_image_dict, paths=result.paths) - - galaxy_name_image_dict = {} - - for path, galaxy_image in raw_image_dict.items(): - minimum_galaxy_value = adapt_minimum_percent * np.max(galaxy_image.array) - galaxy_image[galaxy_image < minimum_galaxy_value] = minimum_galaxy_value - - galaxy_name_image_dict[path] = galaxy_image - - return galaxy_name_image_dict - - -class AdaptImages: - def __init__( - self, - galaxy_image_dict: Optional[Dict[Galaxy, aa.Array2D]] = None, - galaxy_name_image_dict: Optional[Dict[Tuple[str, ...], aa.Array2D]] = None, - galaxy_image_plane_mesh_grid_dict: Optional[Dict[Galaxy, aa.Array2D]] = None, - galaxy_name_image_plane_mesh_grid_dict: Optional[ - Dict[Tuple[str, ...], aa.Grid2DIrregular] - ] = None, - galaxy_path_list: Optional[List[str]] = None, - ): - """ - Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the - reconstructed galaxy's morphology. - - Pixelization image-mesh objects (e.g. `KMeans`, `Hilbert`) adapt the distribution of pixels to the observed - image's brightness and therefore to the reconstructed source's morphology. - - Certain regularization schemes (e.g. `Adapt`) adapt their regularization coefficients to the - reconstructed source's morphology. - - These adaptive schemes use "adapt-images", which are images of each galaxy (e.g. the lens and source of a - strong lens) estimated via an earlier model-fit. - - The adapt-images are stored as the model-image of each galaxy in a model (e.g. the lens and source for a - strong lens). They are stored as a dictionary mapping each instance of the galaxy to its model-image. - - For model-fitting, the galaxy instances are updated for every iteration of the non-linear search. This means - an `AdaptImages` instance cannot be passed directly to an `Analysis` class, as the galaxy instances need to be - updated for every iteration of the non-linear search. - - A dictionary mapping the path name of each galaxy (e.g. "galaxies.lens") to its model-image is therefore used - which is called inside the `log_likelihood_function` o map the model-image of each galaxy to the galaxy - instance of that iteration's specific model. - - Parameters - ---------- - galaxy_image_dict - A dictionary associating each galaxy instance to an image of only that galaxy (e.g. for a strong lens - one entry will map an instance of the source galaxy entry to an image of the lensed source. - galaxy_name_image_dict - A dictionary associating each galaxy path name (e.g. "galaxies.source") to an image of only that - galaxy (e.g. for a strong lens the `source` entry is an image of the lensed source, without the lens light). - """ - - self.galaxy_image_dict = galaxy_image_dict - self.galaxy_name_image_dict = galaxy_name_image_dict - - self.galaxy_image_plane_mesh_grid_dict = galaxy_image_plane_mesh_grid_dict - self.galaxy_name_image_plane_mesh_grid_dict = ( - galaxy_name_image_plane_mesh_grid_dict - ) - - # Parallel to the analysis-time galaxies list (as built by - # ``Analysis.galaxies_via_instance_from``). Populated by - # ``updated_via_instance_from`` and used by ``image_for_galaxy`` to - # recover the galaxy's path-tuple key after a JAX unflatten has produced - # fresh ``Galaxy`` objects whose hashes no longer match - # ``galaxy_image_dict`` keys. - self.galaxy_path_list = galaxy_path_list - - @property - def mask(self) -> aa.Mask2D: - """ - The mask of the adapt images. - """ - try: - return list(self.galaxy_image_dict.values())[0].mask - except AttributeError: - return list(self.galaxy_name_image_dict.values())[0].mask - - @cached_property - def model_image(self) -> aa.Array2D: - """ - The model-image is the sum of all individual galaxy images in the image dictionary. - - This is computed by summing the model-image of each individual adapt galaxy contained in the dictionary. - """ - adapt_model_image = aa.Array2D( - values=np.zeros(self.mask.pixels_in_mask), - mask=self.mask, - ) - - try: - for path in self.galaxy_image_dict.keys(): - adapt_model_image += self.galaxy_image_dict[path] - except AttributeError: - for path in self.galaxy_name_image_dict.keys(): - adapt_model_image += self.galaxy_name_image_dict[path] - - return adapt_model_image - - def updated_via_instance_from( - self, - instance, - dataset_model: Optional["aa.DatasetModel"] = None, - mask=None, - galaxies: Optional[List["Galaxy"]] = None, - xp=np, - ) -> "AdaptImages": - """ - Returns adapt-images which have been updated to map galaxy instances instead of galaxy names. - - For model-fitting, the galaxy instances are updated for every iteration of the non-linear search. This means - an `AdaptImages` instance cannot be passed directly to an `Analysis` class, as the galaxy instances need to be - updated for every iteration of the non-linear search. - - A dictionary mapping the path name of each galaxy (e.g. "galaxies.lens") to its model-image is therefore used - which is called inside the `log_likelihood_function` o map the model-image of each galaxy to the galaxy - instance of that iteration's specific model. - - This function is also called when loading an `AdaptImages` instance from a PyAutoFit database, as the - galaxy instances are also created on-fly from the database. Database images do not have a mask, so it is - also applied to the adapt images on-the-fly during database loading. - - When a ``dataset_model`` is supplied with a non-trivial ``grid_offset`` or ``grid_rotation_angle``, the cached - ``galaxy_name_image_plane_mesh_grid_dict`` entries are transformed into the same frame as the dataset's - image-plane grid (which ``FitDataset.grids`` rotates by the same amount). Without this transform the cached - mesh and the data grid would sit in different frames, producing a misaligned source reconstruction. - - Parameters - ---------- - instance - The instance of the model-fit (e.g. in a non-linear search) which is used to update the adapt images. - dataset_model - The dataset model whose ``grid_offset`` and ``grid_rotation_angle`` are applied to cached mesh grids so - they remain consistent with the rotated/shifted data grid produced by ``FitDataset.grids``. If ``None``, - the cached mesh grids are passed through unchanged. - mask - A mask which can be applied to the adapt images, which is used when setting up the adaptive images - via the aggregator and autofit database tools. - galaxies - Optional list of galaxies in the order used by the calling ``Analysis`` (i.e. the list passed to - ``FitImaging`` / ``Tracer``). When provided, a parallel ``galaxy_path_list`` is populated so that - ``image_for_galaxy`` can recover the path-tuple key for each galaxy after JAX has unflattened the - galaxy instances into fresh objects. When ``None`` the path list is populated in ``path_instance_tuples_for_class`` - order, which matches ``Analysis.galaxies_via_instance_from`` for the common case (no - ``extra_galaxies`` / ``scaling_galaxies``). - xp - Array backend (``numpy`` or ``jax.numpy``) used when transforming cached mesh grids. - - Returns - ------- - - """ - from autogalaxy.galaxy.galaxy import Galaxy - - path_by_id = { - id(galaxy): str(galaxy_name) - for galaxy_name, galaxy in instance.path_instance_tuples_for_class(Galaxy) - } - - galaxy_image_dict = None - - if self.galaxy_name_image_dict is not None: - - galaxy_image_dict = {} - - for galaxy_name, galaxy in instance.path_instance_tuples_for_class(Galaxy): - galaxy_name = str(galaxy_name) - - if galaxy_name in self.galaxy_name_image_dict: - galaxy_image_dict[galaxy] = self.galaxy_name_image_dict[galaxy_name] - - if mask is not None: - for key, image in galaxy_image_dict.items(): - galaxy_image_dict[key] = aa.Array2D(values=image, mask=mask) - - galaxy_image_plane_mesh_grid_dict = None - - if self.galaxy_name_image_plane_mesh_grid_dict is not None: - - galaxy_image_plane_mesh_grid_dict = {} - - for galaxy_name, galaxy in instance.path_instance_tuples_for_class(Galaxy): - galaxy_name = str(galaxy_name) - - if galaxy_name in self.galaxy_name_image_plane_mesh_grid_dict: - cached_mesh = self.galaxy_name_image_plane_mesh_grid_dict[galaxy_name] - if dataset_model is not None: - cached_mesh = cached_mesh.subtracted_and_rotated_from( - offset=dataset_model.grid_offset, - angle=dataset_model.grid_rotation_angle, - xp=xp, - ) - galaxy_image_plane_mesh_grid_dict[galaxy] = cached_mesh - - if galaxies is not None: - galaxy_path_list = [path_by_id.get(id(g)) for g in galaxies] - else: - galaxy_path_list = [ - str(galaxy_name) - for galaxy_name, _ in instance.path_instance_tuples_for_class(Galaxy) - ] - - return AdaptImages( - galaxy_image_dict=galaxy_image_dict, - galaxy_image_plane_mesh_grid_dict=galaxy_image_plane_mesh_grid_dict, - galaxy_name_image_dict=self.galaxy_name_image_dict, - galaxy_name_image_plane_mesh_grid_dict=self.galaxy_name_image_plane_mesh_grid_dict, - galaxy_path_list=galaxy_path_list, - ) - - def image_for_galaxy( - self, galaxy: "Galaxy", galaxies: Optional[List["Galaxy"]] = None - ) -> Optional[aa.Array2D]: - """ - Return the adapt image for ``galaxy``, robust to JAX ``jit`` boundaries. - - ``galaxy_image_dict`` is keyed by the trace-time ``Galaxy`` instances. After ``jax.jit`` has flattened - and unflattened a ``FitImaging``, the galaxies inside it are fresh Python objects whose ``__hash__`` - differs from the trace-time keys, so a direct lookup misses. This helper falls back to the path-tuple - keyed ``galaxy_name_image_dict`` using ``galaxy_path_list`` to map the post-unflatten galaxy back to - its trace-time path. - - Returns ``None`` when no adapt image is associated with the galaxy. - """ - try: - return self.galaxy_image_dict[galaxy] - except (AttributeError, KeyError, TypeError): - pass - - path = self._path_for_galaxy(galaxy, galaxies) - if path is None or self.galaxy_name_image_dict is None: - return None - return self.galaxy_name_image_dict.get(path) - - def image_plane_mesh_grid_for_galaxy( - self, galaxy: "Galaxy", galaxies: Optional[List["Galaxy"]] = None - ) -> Optional[aa.Grid2DIrregular]: - """ - Return the image-plane mesh grid for ``galaxy``, robust to JAX ``jit`` boundaries. - - Companion to :meth:`image_for_galaxy` for ``galaxy_image_plane_mesh_grid_dict`` / - ``galaxy_name_image_plane_mesh_grid_dict``. - """ - try: - return self.galaxy_image_plane_mesh_grid_dict[galaxy] - except (AttributeError, KeyError, TypeError): - pass - - path = self._path_for_galaxy(galaxy, galaxies) - if path is None or self.galaxy_name_image_plane_mesh_grid_dict is None: - return None - return self.galaxy_name_image_plane_mesh_grid_dict.get(path) - - def _path_for_galaxy( - self, galaxy: "Galaxy", galaxies: Optional[List["Galaxy"]] - ) -> Optional[str]: - if not self.galaxy_path_list or galaxies is None: - return None - for index, candidate in enumerate(galaxies): - if candidate is galaxy: - if index < len(self.galaxy_path_list): - return self.galaxy_path_list[index] - return None - return None +from __future__ import annotations +import numpy as np +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple + +from autoconf import conf +from autoconf import cached_property + +import autoarray as aa + +if TYPE_CHECKING: + from autogalaxy.galaxy.galaxy import Galaxy + + +def _galaxy_images_cache_path(result, use_model_images: bool): + """ + The on-disk cache file for the raw per-galaxy images of a result, inside the + result's own ``files/`` folder, or ``None`` when the result has no on-disk + output (e.g. ``NullPaths``) and therefore cannot cache. + """ + from pathlib import Path + + paths = getattr(result, "paths", None) + files_path = getattr(paths, "_files_path", None) + if files_path is None or not Path(files_path).is_dir(): + return None + name = "galaxy_images_model" if use_model_images else "galaxy_images_snr" + return Path(files_path) / f"{name}.fits" + + +def _galaxy_image_dict_from_cache(cache_path) -> Optional[Dict]: + """ + Load the raw (pre minimum-percent clip) per-galaxy image dictionary from a + result's cache file, or ``None`` when the file does not exist (the first + arrival at this result computes and writes it). + + The FITS layout mirrors the ``adapt_images.fits`` artifact the aggregator + reads (``agg_util.adapt_images_from``): HDU 0 is the mask (header carries + pixel scales and origin), HDU 1+ are one image per galaxy with the galaxy + path as ``EXTNAME``. + """ + from astropy.io import fits as astropy_fits + + from autoarray.mask.mask_2d import Mask2DKeys + from autoconf.fitsable import ndarray_via_hdu_from + + if cache_path is None or not cache_path.exists(): + return None + + with astropy_fits.open(cache_path) as hdu_list: + header = hdu_list[0].header + pixel_scales = ( + header[Mask2DKeys.PIXSCAY.value], + header[Mask2DKeys.PIXSCAX.value], + ) + origin = ( + header[Mask2DKeys.ORIGINY.value], + header[Mask2DKeys.ORIGINX.value], + ) + mask = aa.Mask2D( + mask=ndarray_via_hdu_from(hdu_list[0]), + pixel_scales=pixel_scales, + origin=origin, + ) + + galaxy_name_image_dict = {} + for hdu in hdu_list[1:]: + image = aa.Array2D.no_mask( + values=ndarray_via_hdu_from(hdu), + pixel_scales=mask.pixel_scales, + origin=mask.origin, + ) + galaxy_name_image_dict[hdu.header["EXTNAME"].lower()] = image.apply_mask( + mask=mask + ) + + return galaxy_name_image_dict + + +def _galaxy_image_dict_to_cache(cache_path, galaxy_name_image_dict: Dict, paths): + """ + Persist the raw per-galaxy image dictionary to the result's cache file, in + the same FITS layout ``_galaxy_image_dict_from_cache`` reads, and preserve + it in the search's zip archive so later resumes keep it. + """ + from autoconf.fitsable import hdu_list_for_output_from + + image_list = [ + galaxy_name_image_dict[name].native_for_fits + for name in galaxy_name_image_dict + ] + hdu_list = hdu_list_for_output_from( + values_list=[image_list[0].mask.astype("float")] + image_list, + ext_name_list=["mask"] + list(galaxy_name_image_dict.keys()), + header_dict=next(iter(galaxy_name_image_dict.values())).mask.header_dict, + ) + hdu_list.writeto(cache_path, overwrite=True) + paths.preserve_in_zip(cache_path) + + +def galaxy_name_image_dict_via_result_from( + result, use_model_images: bool = False +) -> "AdaptImages": + """ + Returns the adapt-images from a non-linear search result. + + For model-fitting, the adapt-images are typically setup using the maximum log likelihood model of the + previous model-fit. This means the model-fitting is used to cleanly deblend the light of the different + galaxies in the image (e.g. separate the lens light from the source light). + + This method uses attributes of a result (e.g. dictionary mapping galaxy instances to their model-images) + to create the adapt-images. + + This can use either: + + - The model image of each galaxy in the best-fit model. + - The subtracted image of each galaxy in the best-fit model, where the subtracted image is the dataset + minus the model images of all other galaxies. + + Certain models produce galaxy-images with negative flux values (e.g. a pixelization), which can cause + numerical issues with the adaptive schemes. To prevent this, we set a minimum flux value for each + galaxy-image, which is a fraction of the maximum flux value of that image defined via a config file. + + The raw per-galaxy images are cached to the result's own ``files/`` folder on first computation + (``galaxy_images_model.fits`` / ``galaxy_images_snr.fits``) and loaded from there on every later call — + computing them rebuilds the result's maximum log likelihood fit, which on a resumed pipeline pays a fresh + JIT compile plus (for pixelized fits) an inversion, and dominates SLaM resume overhead + (autolens_profiling#70). Staleness is structurally guarded: changing the upstream model or search + produces a new search identifier and therefore a fresh output directory with no cache file. Results with + no on-disk output (e.g. ``NullPaths``) always compute. + + Parameters + ---------- + result + The result of a previous model-fit, which contains the model-image of each galaxy. + use_model_images + If True, the model images of the galaxies are used to create the adapt images. If False, the subtracted + images of the galaxies are used. + + Returns + ------- + The adapt-images, which are the model-image of each galaxy inferred via the previous model-fit. + """ + adapt_minimum_percent = conf.instance["general"]["adapt"]["adapt_minimum_percent"] + + cache_path = _galaxy_images_cache_path(result, use_model_images=use_model_images) + raw_image_dict = _galaxy_image_dict_from_cache(cache_path) + + if raw_image_dict is None: + raw_image_dict = {} + + for path, galaxy in result.path_galaxy_tuples: + if use_model_images: + raw_image_dict[path] = result.model_image_galaxy_dict[path] + else: + raw_image_dict[path] = result.subtracted_signal_to_noise_map_galaxy_dict[ + path + ] + + if cache_path is not None: + _galaxy_image_dict_to_cache(cache_path, raw_image_dict, paths=result.paths) + + galaxy_name_image_dict = {} + + for path, galaxy_image in raw_image_dict.items(): + minimum_galaxy_value = adapt_minimum_percent * np.max(galaxy_image.array) + galaxy_image[galaxy_image < minimum_galaxy_value] = minimum_galaxy_value + + galaxy_name_image_dict[path] = galaxy_image + + return galaxy_name_image_dict + + +class AdaptImages: + def __init__( + self, + galaxy_image_dict: Optional[Dict[Galaxy, aa.Array2D]] = None, + galaxy_name_image_dict: Optional[Dict[Tuple[str, ...], aa.Array2D]] = None, + galaxy_image_plane_mesh_grid_dict: Optional[Dict[Galaxy, aa.Array2D]] = None, + galaxy_name_image_plane_mesh_grid_dict: Optional[ + Dict[Tuple[str, ...], aa.Grid2DIrregular] + ] = None, + galaxy_path_list: Optional[List[str]] = None, + ): + """ + Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the + reconstructed galaxy's morphology. + + Pixelization image-mesh objects (e.g. `KMeans`, `Hilbert`) adapt the distribution of pixels to the observed + image's brightness and therefore to the reconstructed source's morphology. + + Certain regularization schemes (e.g. `Adapt`) adapt their regularization coefficients to the + reconstructed source's morphology. + + These adaptive schemes use "adapt-images", which are images of each galaxy (e.g. the lens and source of a + strong lens) estimated via an earlier model-fit. + + The adapt-images are stored as the model-image of each galaxy in a model (e.g. the lens and source for a + strong lens). They are stored as a dictionary mapping each instance of the galaxy to its model-image. + + For model-fitting, the galaxy instances are updated for every iteration of the non-linear search. This means + an `AdaptImages` instance cannot be passed directly to an `Analysis` class, as the galaxy instances need to be + updated for every iteration of the non-linear search. + + A dictionary mapping the path name of each galaxy (e.g. "galaxies.lens") to its model-image is therefore used + which is called inside the `log_likelihood_function` o map the model-image of each galaxy to the galaxy + instance of that iteration's specific model. + + Parameters + ---------- + galaxy_image_dict + A dictionary associating each galaxy instance to an image of only that galaxy (e.g. for a strong lens + one entry will map an instance of the source galaxy entry to an image of the lensed source. + galaxy_name_image_dict + A dictionary associating each galaxy path name (e.g. "galaxies.source") to an image of only that + galaxy (e.g. for a strong lens the `source` entry is an image of the lensed source, without the lens light). + """ + + self.galaxy_image_dict = galaxy_image_dict + self.galaxy_name_image_dict = galaxy_name_image_dict + + self.galaxy_image_plane_mesh_grid_dict = galaxy_image_plane_mesh_grid_dict + self.galaxy_name_image_plane_mesh_grid_dict = ( + galaxy_name_image_plane_mesh_grid_dict + ) + + # Parallel to the analysis-time galaxies list (as built by + # ``Analysis.galaxies_via_instance_from``). Populated by + # ``updated_via_instance_from`` and used by ``image_for_galaxy`` to + # recover the galaxy's path-tuple key after a JAX unflatten has produced + # fresh ``Galaxy`` objects whose hashes no longer match + # ``galaxy_image_dict`` keys. + self.galaxy_path_list = galaxy_path_list + + @property + def mask(self) -> aa.Mask2D: + """ + The mask of the adapt images. + """ + try: + return list(self.galaxy_image_dict.values())[0].mask + except AttributeError: + return list(self.galaxy_name_image_dict.values())[0].mask + + @cached_property + def model_image(self) -> aa.Array2D: + """ + The model-image is the sum of all individual galaxy images in the image dictionary. + + This is computed by summing the model-image of each individual adapt galaxy contained in the dictionary. + """ + adapt_model_image = aa.Array2D( + values=np.zeros(self.mask.pixels_in_mask), + mask=self.mask, + ) + + try: + for path in self.galaxy_image_dict.keys(): + adapt_model_image += self.galaxy_image_dict[path] + except AttributeError: + for path in self.galaxy_name_image_dict.keys(): + adapt_model_image += self.galaxy_name_image_dict[path] + + return adapt_model_image + + def updated_via_instance_from( + self, + instance, + dataset_model: Optional["aa.DatasetModel"] = None, + mask=None, + galaxies: Optional[List["Galaxy"]] = None, + xp=np, + ) -> "AdaptImages": + """ + Returns adapt-images which have been updated to map galaxy instances instead of galaxy names. + + For model-fitting, the galaxy instances are updated for every iteration of the non-linear search. This means + an `AdaptImages` instance cannot be passed directly to an `Analysis` class, as the galaxy instances need to be + updated for every iteration of the non-linear search. + + A dictionary mapping the path name of each galaxy (e.g. "galaxies.lens") to its model-image is therefore used + which is called inside the `log_likelihood_function` o map the model-image of each galaxy to the galaxy + instance of that iteration's specific model. + + This function is also called when loading an `AdaptImages` instance from a PyAutoFit database, as the + galaxy instances are also created on-fly from the database. Database images do not have a mask, so it is + also applied to the adapt images on-the-fly during database loading. + + When a ``dataset_model`` is supplied with a non-trivial ``grid_offset`` or ``grid_rotation_angle``, the cached + ``galaxy_name_image_plane_mesh_grid_dict`` entries are transformed into the same frame as the dataset's + image-plane grid (which ``FitDataset.grids`` rotates by the same amount). Without this transform the cached + mesh and the data grid would sit in different frames, producing a misaligned source reconstruction. + + Parameters + ---------- + instance + The instance of the model-fit (e.g. in a non-linear search) which is used to update the adapt images. + dataset_model + The dataset model whose ``grid_offset`` and ``grid_rotation_angle`` are applied to cached mesh grids so + they remain consistent with the rotated/shifted data grid produced by ``FitDataset.grids``. If ``None``, + the cached mesh grids are passed through unchanged. + mask + A mask which can be applied to the adapt images, which is used when setting up the adaptive images + via the aggregator and autofit database tools. + galaxies + Optional list of galaxies in the order used by the calling ``Analysis`` (i.e. the list passed to + ``FitImaging`` / ``Tracer``). When provided, a parallel ``galaxy_path_list`` is populated so that + ``image_for_galaxy`` can recover the path-tuple key for each galaxy after JAX has unflattened the + galaxy instances into fresh objects. When ``None`` the path list is populated in ``path_instance_tuples_for_class`` + order, which matches ``Analysis.galaxies_via_instance_from`` for the common case (no + ``extra_galaxies`` / ``scaling_galaxies``). + xp + Array backend (``numpy`` or ``jax.numpy``) used when transforming cached mesh grids. + + Returns + ------- + + """ + from autogalaxy.galaxy.galaxy import Galaxy + + path_by_id = { + id(galaxy): str(galaxy_name) + for galaxy_name, galaxy in instance.path_instance_tuples_for_class(Galaxy) + } + + galaxy_image_dict = None + + if self.galaxy_name_image_dict is not None: + + galaxy_image_dict = {} + + for galaxy_name, galaxy in instance.path_instance_tuples_for_class(Galaxy): + galaxy_name = str(galaxy_name) + + if galaxy_name in self.galaxy_name_image_dict: + galaxy_image_dict[galaxy] = self.galaxy_name_image_dict[galaxy_name] + + if mask is not None: + for key, image in galaxy_image_dict.items(): + galaxy_image_dict[key] = aa.Array2D(values=image, mask=mask) + + galaxy_image_plane_mesh_grid_dict = None + + if self.galaxy_name_image_plane_mesh_grid_dict is not None: + + galaxy_image_plane_mesh_grid_dict = {} + + for galaxy_name, galaxy in instance.path_instance_tuples_for_class(Galaxy): + galaxy_name = str(galaxy_name) + + if galaxy_name in self.galaxy_name_image_plane_mesh_grid_dict: + cached_mesh = self.galaxy_name_image_plane_mesh_grid_dict[galaxy_name] + if dataset_model is not None: + cached_mesh = cached_mesh.subtracted_and_rotated_from( + offset=dataset_model.grid_offset, + angle=dataset_model.grid_rotation_angle, + xp=xp, + ) + galaxy_image_plane_mesh_grid_dict[galaxy] = cached_mesh + + if galaxies is not None: + galaxy_path_list = [path_by_id.get(id(g)) for g in galaxies] + else: + galaxy_path_list = [ + str(galaxy_name) + for galaxy_name, _ in instance.path_instance_tuples_for_class(Galaxy) + ] + + return AdaptImages( + galaxy_image_dict=galaxy_image_dict, + galaxy_image_plane_mesh_grid_dict=galaxy_image_plane_mesh_grid_dict, + galaxy_name_image_dict=self.galaxy_name_image_dict, + galaxy_name_image_plane_mesh_grid_dict=self.galaxy_name_image_plane_mesh_grid_dict, + galaxy_path_list=galaxy_path_list, + ) + + def image_for_galaxy( + self, galaxy: "Galaxy", galaxies: Optional[List["Galaxy"]] = None + ) -> Optional[aa.Array2D]: + """ + Return the adapt image for ``galaxy``, robust to JAX ``jit`` boundaries. + + ``galaxy_image_dict`` is keyed by the trace-time ``Galaxy`` instances. After ``jax.jit`` has flattened + and unflattened a ``FitImaging``, the galaxies inside it are fresh Python objects whose ``__hash__`` + differs from the trace-time keys, so a direct lookup misses. This helper falls back to the path-tuple + keyed ``galaxy_name_image_dict`` using ``galaxy_path_list`` to map the post-unflatten galaxy back to + its trace-time path. + + Returns ``None`` when no adapt image is associated with the galaxy. + """ + try: + return self.galaxy_image_dict[galaxy] + except (AttributeError, KeyError, TypeError): + pass + + path = self._path_for_galaxy(galaxy, galaxies) + if path is None or self.galaxy_name_image_dict is None: + return None + return self.galaxy_name_image_dict.get(path) + + def image_plane_mesh_grid_for_galaxy( + self, galaxy: "Galaxy", galaxies: Optional[List["Galaxy"]] = None + ) -> Optional[aa.Grid2DIrregular]: + """ + Return the image-plane mesh grid for ``galaxy``, robust to JAX ``jit`` boundaries. + + Companion to :meth:`image_for_galaxy` for ``galaxy_image_plane_mesh_grid_dict`` / + ``galaxy_name_image_plane_mesh_grid_dict``. + """ + try: + return self.galaxy_image_plane_mesh_grid_dict[galaxy] + except (AttributeError, KeyError, TypeError): + pass + + path = self._path_for_galaxy(galaxy, galaxies) + if path is None or self.galaxy_name_image_plane_mesh_grid_dict is None: + return None + return self.galaxy_name_image_plane_mesh_grid_dict.get(path) + + def _path_for_galaxy( + self, galaxy: "Galaxy", galaxies: Optional[List["Galaxy"]] + ) -> Optional[str]: + if not self.galaxy_path_list or galaxies is None: + return None + for index, candidate in enumerate(galaxies): + if candidate is galaxy: + if index < len(self.galaxy_path_list): + return self.galaxy_path_list[index] + return None + return None diff --git a/test_autogalaxy/analysis/test_adapt_images.py b/test_autogalaxy/analysis/test_adapt_images.py index 337d4236..69116816 100644 --- a/test_autogalaxy/analysis/test_adapt_images.py +++ b/test_autogalaxy/analysis/test_adapt_images.py @@ -1,111 +1,111 @@ -import pytest -import numpy as np - -import autofit as af -import autogalaxy as ag - - -def test__instance_with_associated_adapt_images_from(masked_imaging_7x7): - g0 = ag.Galaxy(redshift=0.5) - g1 = ag.Galaxy(redshift=1.0) - - galaxy_image_dict = { - g0: ag.Array2D.ones(shape_native=(3, 3), pixel_scales=1.0), - g1: ag.Array2D.full(fill_value=2.0, shape_native=(3, 3), pixel_scales=1.0), - } - - adapt_images = ag.AdaptImages( - galaxy_image_dict=galaxy_image_dict, - ) - - assert adapt_images.model_image.native == pytest.approx( - 3.0 * np.ones((3, 3)), 1.0e-4 - ) - - -def test__image_for_galaxy__resolves_after_galaxy_identity_changes(): - """ - Simulates the post-``jax.jit`` unflatten boundary: ``adapt_images.galaxy_image_dict`` is keyed by the - trace-time ``Galaxy`` instances, but the lookup at ``GalaxiesToInversion.mapper_galaxy_dict`` is performed - against fresh ``Galaxy`` objects whose ``__hash__`` differs. The path-tuple lookup via - ``galaxy_name_image_dict`` must still resolve to the right adapt image. - """ - galaxies = af.ModelInstance() - galaxies.lens = ag.Galaxy(redshift=0.5) - galaxies.source = ag.Galaxy(redshift=1.0) - - instance = af.ModelInstance() - instance.galaxies = galaxies - - galaxy_name_image_dict = { - str(("galaxies", "lens")): ag.Array2D.ones(shape_native=(3, 3), pixel_scales=1.0), - str(("galaxies", "source")): ag.Array2D.full( - fill_value=2.0, shape_native=(3, 3), pixel_scales=1.0 - ), - } - - trace_galaxies = [galaxies.lens, galaxies.source] - - adapt_images = ag.AdaptImages( - galaxy_name_image_dict=galaxy_name_image_dict, - ).updated_via_instance_from(instance=instance, galaxies=trace_galaxies) - - assert adapt_images.galaxy_path_list == [ - str(("galaxies", "lens")), - str(("galaxies", "source")), - ] - - # Fast path: by-instance lookup still works for the trace-time galaxies. - assert adapt_images.image_for_galaxy( - trace_galaxies[0], trace_galaxies - ).native == pytest.approx(np.ones((3, 3)), 1.0e-4) - - # Simulate post-unflatten: fresh ``Galaxy`` objects with new ``.id`` values - # placed at the same positions as the trace-time list. ``galaxy_image_dict`` - # cannot resolve them (hash mismatch) so the helper must fall back to - # ``galaxy_name_image_dict`` via ``galaxy_path_list``. - fresh_galaxies = [ag.Galaxy(redshift=0.5), ag.Galaxy(redshift=1.0)] - - assert adapt_images.galaxy_image_dict.get(fresh_galaxies[0]) is None - assert adapt_images.image_for_galaxy( - fresh_galaxies[0], fresh_galaxies - ).native == pytest.approx(np.ones((3, 3)), 1.0e-4) - assert adapt_images.image_for_galaxy( - fresh_galaxies[1], fresh_galaxies - ).native == pytest.approx(2.0 * np.ones((3, 3)), 1.0e-4) - - -def test__image_plane_mesh_grid_for_galaxy__resolves_after_galaxy_identity_changes(): - """ - Companion to :func:`test__image_for_galaxy__resolves_after_galaxy_identity_changes` for the mesh-grid - lookup path used by ``GalaxiesToInversion.image_plane_mesh_grid_list``. - """ - galaxies = af.ModelInstance() - galaxies.lens = ag.Galaxy(redshift=0.5) - galaxies.source = ag.Galaxy(redshift=1.0) - - instance = af.ModelInstance() - instance.galaxies = galaxies - - galaxy_name_image_plane_mesh_grid_dict = { - str(("galaxies", "lens")): ag.Grid2DIrregular(values=[(3.0, 3.0), (3.0, 3.0)]), - str(("galaxies", "source")): ag.Grid2DIrregular(values=[(4.0, 4.0), (4.0, 4.0)]), - } - - trace_galaxies = [galaxies.lens, galaxies.source] - - adapt_images = ag.AdaptImages( - galaxy_name_image_plane_mesh_grid_dict=galaxy_name_image_plane_mesh_grid_dict, - ).updated_via_instance_from(instance=instance, galaxies=trace_galaxies) - - fresh_galaxies = [ag.Galaxy(redshift=0.5), ag.Galaxy(redshift=1.0)] - - assert adapt_images.image_plane_mesh_grid_for_galaxy( - fresh_galaxies[0], fresh_galaxies - ) == pytest.approx(3.0 * np.ones((2, 2)), 1.0e-4) - assert adapt_images.image_plane_mesh_grid_for_galaxy( - fresh_galaxies[1], fresh_galaxies - ) == pytest.approx(4.0 * np.ones((2, 2)), 1.0e-4) +import pytest +import numpy as np + +import autofit as af +import autogalaxy as ag + + +def test__instance_with_associated_adapt_images_from(masked_imaging_7x7): + g0 = ag.Galaxy(redshift=0.5) + g1 = ag.Galaxy(redshift=1.0) + + galaxy_image_dict = { + g0: ag.Array2D.ones(shape_native=(3, 3), pixel_scales=1.0), + g1: ag.Array2D.full(fill_value=2.0, shape_native=(3, 3), pixel_scales=1.0), + } + + adapt_images = ag.AdaptImages( + galaxy_image_dict=galaxy_image_dict, + ) + + assert adapt_images.model_image.native == pytest.approx( + 3.0 * np.ones((3, 3)), 1.0e-4 + ) + + +def test__image_for_galaxy__resolves_after_galaxy_identity_changes(): + """ + Simulates the post-``jax.jit`` unflatten boundary: ``adapt_images.galaxy_image_dict`` is keyed by the + trace-time ``Galaxy`` instances, but the lookup at ``GalaxiesToInversion.mapper_galaxy_dict`` is performed + against fresh ``Galaxy`` objects whose ``__hash__`` differs. The path-tuple lookup via + ``galaxy_name_image_dict`` must still resolve to the right adapt image. + """ + galaxies = af.ModelInstance() + galaxies.lens = ag.Galaxy(redshift=0.5) + galaxies.source = ag.Galaxy(redshift=1.0) + + instance = af.ModelInstance() + instance.galaxies = galaxies + + galaxy_name_image_dict = { + str(("galaxies", "lens")): ag.Array2D.ones(shape_native=(3, 3), pixel_scales=1.0), + str(("galaxies", "source")): ag.Array2D.full( + fill_value=2.0, shape_native=(3, 3), pixel_scales=1.0 + ), + } + + trace_galaxies = [galaxies.lens, galaxies.source] + + adapt_images = ag.AdaptImages( + galaxy_name_image_dict=galaxy_name_image_dict, + ).updated_via_instance_from(instance=instance, galaxies=trace_galaxies) + + assert adapt_images.galaxy_path_list == [ + str(("galaxies", "lens")), + str(("galaxies", "source")), + ] + + # Fast path: by-instance lookup still works for the trace-time galaxies. + assert adapt_images.image_for_galaxy( + trace_galaxies[0], trace_galaxies + ).native == pytest.approx(np.ones((3, 3)), 1.0e-4) + + # Simulate post-unflatten: fresh ``Galaxy`` objects with new ``.id`` values + # placed at the same positions as the trace-time list. ``galaxy_image_dict`` + # cannot resolve them (hash mismatch) so the helper must fall back to + # ``galaxy_name_image_dict`` via ``galaxy_path_list``. + fresh_galaxies = [ag.Galaxy(redshift=0.5), ag.Galaxy(redshift=1.0)] + + assert adapt_images.galaxy_image_dict.get(fresh_galaxies[0]) is None + assert adapt_images.image_for_galaxy( + fresh_galaxies[0], fresh_galaxies + ).native == pytest.approx(np.ones((3, 3)), 1.0e-4) + assert adapt_images.image_for_galaxy( + fresh_galaxies[1], fresh_galaxies + ).native == pytest.approx(2.0 * np.ones((3, 3)), 1.0e-4) + + +def test__image_plane_mesh_grid_for_galaxy__resolves_after_galaxy_identity_changes(): + """ + Companion to :func:`test__image_for_galaxy__resolves_after_galaxy_identity_changes` for the mesh-grid + lookup path used by ``GalaxiesToInversion.image_plane_mesh_grid_list``. + """ + galaxies = af.ModelInstance() + galaxies.lens = ag.Galaxy(redshift=0.5) + galaxies.source = ag.Galaxy(redshift=1.0) + + instance = af.ModelInstance() + instance.galaxies = galaxies + + galaxy_name_image_plane_mesh_grid_dict = { + str(("galaxies", "lens")): ag.Grid2DIrregular(values=[(3.0, 3.0), (3.0, 3.0)]), + str(("galaxies", "source")): ag.Grid2DIrregular(values=[(4.0, 4.0), (4.0, 4.0)]), + } + + trace_galaxies = [galaxies.lens, galaxies.source] + + adapt_images = ag.AdaptImages( + galaxy_name_image_plane_mesh_grid_dict=galaxy_name_image_plane_mesh_grid_dict, + ).updated_via_instance_from(instance=instance, galaxies=trace_galaxies) + + fresh_galaxies = [ag.Galaxy(redshift=0.5), ag.Galaxy(redshift=1.0)] + + assert adapt_images.image_plane_mesh_grid_for_galaxy( + fresh_galaxies[0], fresh_galaxies + ) == pytest.approx(3.0 * np.ones((2, 2)), 1.0e-4) + assert adapt_images.image_plane_mesh_grid_for_galaxy( + fresh_galaxies[1], fresh_galaxies + ) == pytest.approx(4.0 * np.ones((2, 2)), 1.0e-4) class _StubCachePaths: @@ -114,6 +114,9 @@ class _StubCachePaths: def __init__(self, files_path): self._files_path = files_path + def preserve_in_zip(self, file_path): + """No zip in the stub — mirrors AbstractPaths' no-zip no-op.""" + class _StubCacheResult: """Duck-typed result for `galaxy_name_image_dict_via_result_from`."""