diff --git a/autogalaxy/analysis/adapt_images/adapt_images.py b/autogalaxy/analysis/adapt_images/adapt_images.py index 5ae82792..d5f55888 100644 --- a/autogalaxy/analysis/adapt_images/adapt_images.py +++ b/autogalaxy/analysis/adapt_images/adapt_images.py @@ -27,7 +27,24 @@ def _galaxy_images_cache_path(result, use_model_images: bool): return Path(files_path) / f"{name}.fits" -def _galaxy_image_dict_from_cache(cache_path) -> Optional[Dict]: +def _masks_are_equal(mask_0, mask_1) -> bool: + """ + Returns whether two masks describe the same set of unmasked pixels on the same + geometry, comparing shape, pixel scales, origin and the mask array itself. + + Used to decide whether an on-disk adapt-image cache was written under the mask + the caller is now fitting with. + """ + if mask_0.shape_native != mask_1.shape_native: + return False + if not np.allclose(mask_0.pixel_scales, mask_1.pixel_scales): + return False + if not np.allclose(mask_0.origin, mask_1.origin): + return False + return bool(np.array_equal(np.asarray(mask_0), np.asarray(mask_1))) + + +def _galaxy_image_dict_from_cache(cache_path, mask=None) -> 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 @@ -37,6 +54,24 @@ def _galaxy_image_dict_from_cache(cache_path) -> Optional[Dict]: 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``. + + ``None`` is also returned when ``mask`` is supplied and does not match the + mask the cache was written under, which makes the mismatch a cache miss so + the caller recomputes and overwrites the stale file. The cache lives in a + directory keyed by the search identifier, which encodes the model and the + search but **not** the dataset — so a run whose dataset shaping changed while + the model stayed identical lands on the previous run's cache. Without this + check the stale images load silently and surface much later as an + ``IndexError`` deep in the inversion, when adapt data sized for the old mask + is indexed by the new mask's slim indexes (PyAutoGalaxy#516). + + Parameters + ---------- + cache_path + The cache file to load, or ``None`` when the result cannot cache. + mask + The mask the caller expects the cached images to be defined on. When + ``None`` no validation is performed. """ from astropy.io import fits as astropy_fits @@ -56,21 +91,24 @@ def _galaxy_image_dict_from_cache(cache_path) -> Optional[Dict]: header[Mask2DKeys.ORIGINY.value], header[Mask2DKeys.ORIGINX.value], ) - mask = aa.Mask2D( + cached_mask = aa.Mask2D( mask=ndarray_via_hdu_from(hdu_list[0]), pixel_scales=pixel_scales, origin=origin, ) + if mask is not None and not _masks_are_equal(cached_mask, mask): + return None + 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, + pixel_scales=cached_mask.pixel_scales, + origin=cached_mask.origin, ) galaxy_name_image_dict[hdu.header["EXTNAME"].lower()] = image.apply_mask( - mask=mask + mask=cached_mask ) return galaxy_name_image_dict @@ -124,9 +162,22 @@ def galaxy_name_image_dict_via_result_from( (``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. + (autolens_profiling#70). Results with no on-disk output (e.g. ``NullPaths``) always compute. + + Two things can make a cache file stale, and they are guarded differently. Changing the upstream model or + search is guarded structurally: it produces a new search identifier and therefore a fresh output directory + with no cache file. Changing the *dataset* is not, because the identifier does not encode it — a rerun + whose mask changed while the model stayed identical lands on the previous run's directory and its cache. + That case is caught by validating the cached mask against ``result.mask`` on load and treating a mismatch + as a cache miss, so the images are recomputed on the current mask. Reading ``result.mask`` resolves to + ``analysis.dataset.mask`` and does not rebuild the maximum log likelihood fit, so the check costs nothing + the cache was saving (PyAutoGalaxy#516). + + Note that a stale cache is not *repaired* on disk: the recomputed images are written to the loose + ``files/`` folder, but ``Paths.preserve_in_zip`` only adds a member that is absent from the search's zip + and never replaces one, so the next ``restore()`` re-extracts the stale copy. Such a search therefore + misses the cache on every run rather than once — correct, but without the caching win until its output is + cleared (PyAutoFit#1414). Parameters ---------- @@ -143,7 +194,9 @@ def galaxy_name_image_dict_via_result_from( 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) + raw_image_dict = _galaxy_image_dict_from_cache( + cache_path, mask=getattr(result, "mask", None) + ) if raw_image_dict is None: raw_image_dict = {} diff --git a/test_autogalaxy/analysis/test_adapt_images.py b/test_autogalaxy/analysis/test_adapt_images.py index 69116816..808b6fd2 100644 --- a/test_autogalaxy/analysis/test_adapt_images.py +++ b/test_autogalaxy/analysis/test_adapt_images.py @@ -206,3 +206,81 @@ def test__galaxy_name_image_dict_via_result_from__no_paths_always_computes(): galaxy_name_image_dict = galaxy_name_image_dict_via_result_from(result=result) assert set(galaxy_name_image_dict.keys()) == set(image_dict.keys()) + + +def test__galaxy_image_dict_cache__mask_mismatch_is_a_cache_miss(tmp_path): + """ + The cache lives in a directory keyed by the search identifier, which encodes the model and the search + but not the dataset. A rerun whose mask changed while the model stayed identical therefore lands on the + previous run's cache, whose images are sized for the old mask (PyAutoGalaxy#516). + """ + 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)) + + cached_mask = next(iter(image_dict.values())).mask + + # The mask the cache was written under is a hit. + + assert _galaxy_image_dict_from_cache(cache_path, mask=cached_mask) is not None + + # A mask enclosing a different number of pixels is a miss, so the caller recomputes. + + other_mask = ag.Mask2D.circular( + shape_native=(7, 7), pixel_scales=0.1, radius=0.2 + ) + + assert other_mask.pixels_in_mask != cached_mask.pixels_in_mask + assert _galaxy_image_dict_from_cache(cache_path, mask=other_mask) is None + + # So is a mask of the same shape and pixel count on a different geometry. + + rescaled_mask = ag.Mask2D( + mask=np.asarray(cached_mask), pixel_scales=0.2, origin=cached_mask.origin + ) + + assert _galaxy_image_dict_from_cache(cache_path, mask=rescaled_mask) is None + + +def test__galaxy_name_image_dict_via_result_from__stale_cache_is_recomputed(tmp_path): + """ + End-to-end: a result whose mask differs from the one the cache was written under must recompute rather + than silently load adapt images sized for the old mask. + """ + from autogalaxy.analysis.adapt_images.adapt_images import ( + galaxy_name_image_dict_via_result_from, + ) + + image_dict = _cache_test_image_dict() + mask = next(iter(image_dict.values())).mask + + result = _StubCacheResult(tmp_path, image_dict) + result.mask = mask + + galaxy_name_image_dict_via_result_from(result=result) + + assert (tmp_path / "galaxy_images_snr.fits").exists() + + # Same output directory, same model — but a different mask. The poisoned stub raises if the cache is + # loaded instead of recomputed, so reaching the compute path is the assertion. + + stale = _StubCacheResult(tmp_path, image_dict, poisoned=True) + stale.mask = ag.Mask2D.circular( + shape_native=(7, 7), pixel_scales=0.1, radius=0.2 + ) + + with pytest.raises(AssertionError): + galaxy_name_image_dict_via_result_from(result=stale) + + # The unchanged mask still hits the cache, so the fix does not cost the caching win. + + unchanged = _StubCacheResult(tmp_path, image_dict, poisoned=True) + unchanged.mask = mask + + assert galaxy_name_image_dict_via_result_from(result=unchanged) is not None