feat: cache per-galaxy result images for SLaM resume fast-path (#502)

autogalaxy/analysis/adapt_images/adapt_images.py

@@ -11,6 +11,115 @@
     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":
@@ -34,6 +143,14 @@ def galaxy_name_image_dict_via_result_from(
     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
@@ -48,14 +165,26 @@ def galaxy_name_image_dict_via_result_from(
     """
     adapt_minimum_percent = conf.instance["general"]["adapt"]["adapt_minimum_percent"]
 
-    galaxy_name_image_dict = {}
+    cache_path = _galaxy_images_cache_path(result, use_model_images=use_model_images)
+    raw_image_dict = _galaxy_image_dict_from_cache(cache_path)
 
-    for path, galaxy in result.path_galaxy_tuples:
-        if use_model_images:
-            galaxy_image = result.model_image_galaxy_dict[path]
-        else:
-            galaxy_image = result.subtracted_signal_to_noise_map_galaxy_dict[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
 

test_autogalaxy/analysis/test_adapt_images.py

@@ -106,3 +106,100 @@ def test__image_plane_mesh_grid_for_galaxy__resolves_after_galaxy_identity_chang
     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:
+    """Duck-typed paths: only what the galaxy-image cache helpers touch."""
+
+    def __init__(self, files_path):
+        self._files_path = files_path
+
+
+class _StubCacheResult:
+    """Duck-typed result for `galaxy_name_image_dict_via_result_from`."""
+
+    def __init__(self, files_path, image_dict, poisoned=False):
+        self.paths = _StubCachePaths(files_path)
+        self._image_dict = image_dict
+        self._poisoned = poisoned
+
+    @property
+    def path_galaxy_tuples(self):
+        return [(name, None) for name in self._image_dict]
+
+    @property
+    def subtracted_signal_to_noise_map_galaxy_dict(self):
+        if self._poisoned:
+            raise AssertionError(
+                "recompute path taken — the cache should have been loaded"
+            )
+        return self._image_dict
+
+
+def _cache_test_image_dict():
+    mask = ag.Mask2D.circular(shape_native=(7, 7), pixel_scales=0.1, radius=0.3)
+    image_0 = ag.Array2D.ones(shape_native=(7, 7), pixel_scales=0.1).apply_mask(
+        mask=mask
+    )
+    image_1 = ag.Array2D.full(
+        fill_value=2.0, shape_native=(7, 7), pixel_scales=0.1
+    ).apply_mask(mask=mask)
+    return {
+        "('galaxies', 'lens')": image_0,
+        "('galaxies', 'source')": image_1,
+    }
+
+
+def test__galaxy_image_dict_cache__round_trip(tmp_path):
+    from autogalaxy.analysis.adapt_images.adapt_images import (
+        _galaxy_image_dict_from_cache,
+        _galaxy_image_dict_to_cache,
+    )
+
+    image_dict = _cache_test_image_dict()
+    cache_path = tmp_path / "galaxy_images_snr.fits"
+
+    _galaxy_image_dict_to_cache(cache_path, image_dict, paths=_StubCachePaths(tmp_path))
+    loaded = _galaxy_image_dict_from_cache(cache_path)
+
+    assert set(loaded.keys()) == set(image_dict.keys())
+    for name in image_dict:
+        assert loaded[name].array == pytest.approx(image_dict[name].array, 1.0e-8)
+        assert (loaded[name].mask == image_dict[name].mask).all()
+
+
+def test__galaxy_name_image_dict_via_result_from__loads_cache_on_second_call(tmp_path):
+    from autogalaxy.analysis.adapt_images.adapt_images import (
+        galaxy_name_image_dict_via_result_from,
+    )
+
+    image_dict = _cache_test_image_dict()
+
+    result = _StubCacheResult(tmp_path, image_dict)
+    first = galaxy_name_image_dict_via_result_from(result=result)
+
+    assert (tmp_path / "galaxy_images_snr.fits").exists()
+
+    # A poisoned result raises if the compute path is taken — the second call
+    # must come entirely from the cache written by the first.
+    poisoned = _StubCacheResult(tmp_path, image_dict, poisoned=True)
+    second = galaxy_name_image_dict_via_result_from(result=poisoned)
+
+    assert set(second.keys()) == set(first.keys())
+    for name in first:
+        assert second[name].array == pytest.approx(first[name].array, 1.0e-8)
+
+
+def test__galaxy_name_image_dict_via_result_from__no_paths_always_computes():
+    from autogalaxy.analysis.adapt_images.adapt_images import (
+        galaxy_name_image_dict_via_result_from,
+    )
+
+    image_dict = _cache_test_image_dict()
+
+    result = _StubCacheResult(files_path=None, image_dict=image_dict)
+    result.paths = None
+
+    galaxy_name_image_dict = galaxy_name_image_dict_via_result_from(result=result)
+
+    assert set(galaxy_name_image_dict.keys()) == set(image_dict.keys())