diff --git a/autolens/analysis/result.py b/autolens/analysis/result.py index b3b8fb21f..3e65aac16 100644 --- a/autolens/analysis/result.py +++ b/autolens/analysis/result.py @@ -15,6 +15,7 @@ These results feed directly into downstream pipeline stages and post-processing scripts. """ +import json import logging import os import numpy as np @@ -267,6 +268,59 @@ def positions_threshold_from( return threshold + def _cached_multiple_image_positions_from( + self, plane_redshift: Optional[float] = None + ) -> aa.Grid2DIrregular: + """ + The multiple image positions of the maximum log likelihood lens model, loaded + from this result's own ``files/`` folder when previously solved, else solved + via the point solver and persisted there. + + Solving the multiple image positions runs a point-solver grid search over the + maximum log likelihood tracer, which on a resumed pipeline (e.g. SLaM) pays a + fresh JIT compile and dominates resume overhead (autolens_profiling#70) — the + solved positions themselves are a pure product of the completed fit, so they + are cached as ``files/multiple_image_positions[_plane_].json``. Staleness + is structurally guarded: a changed model or search produces a new search + identifier and a fresh output directory with no cache file. Results with no + on-disk output (e.g. ``NullPaths``) always solve. + """ + from pathlib import Path + + name = "multiple_image_positions" + if plane_redshift is not None: + name += f"_plane_{str(plane_redshift).replace('.', '_')}" + + files_path = getattr(getattr(self, "paths", None), "_files_path", None) + cache_path = ( + Path(files_path) / f"{name}.json" + if files_path is not None and Path(files_path).is_dir() + else None + ) + + if cache_path is not None and cache_path.exists(): + with open(cache_path) as f: + return aa.Grid2DIrregular(values=[tuple(p) for p in json.load(f)]) + + positions = self.image_plane_multiple_image_positions( + plane_redshift=plane_redshift + ) + + if cache_path is not None: + with open(cache_path, "w") as f: + json.dump(np.asarray(positions.array).tolist(), f) + + # Preserve the cache in the search's zip — a resumed search's + # paths.restore() wipes the output dir and re-extracts the zip, + # destroying any file written only to files/ after completion. + from autogalaxy.analysis.adapt_images.adapt_images import ( + _append_to_search_zip, + ) + + _append_to_search_zip(self.paths, cache_path) + + return positions + def positions_likelihood_from( self, factor=1.0, @@ -355,11 +409,10 @@ def positions_likelihood_from( ) return - positions = ( - self.image_plane_multiple_image_positions(plane_redshift=plane_redshift) - if positions is None - else positions - ) + if positions is None: + positions = self._cached_multiple_image_positions_from( + plane_redshift=plane_redshift + ) if mass_centre_radial_distance_min is not None: mass_centre = self.max_log_likelihood_tracer.extract_attribute( diff --git a/test_autolens/analysis/test_result.py b/test_autolens/analysis/test_result.py index b595a930b..528e2b5cf 100644 --- a/test_autolens/analysis/test_result.py +++ b/test_autolens/analysis/test_result.py @@ -380,3 +380,50 @@ def test___image_dict(analysis_imaging_7x7): assert (image_dict[str(("galaxies", "lens"))].native == np.zeros((7, 7))).all() assert isinstance(image_dict[str(("galaxies", "source"))], Array2D) + + +def test__positions_likelihood_from__loads_cached_positions_on_second_call( + tmp_path, monkeypatch, analysis_imaging_7x7 +): + class _StubPaths: + def __init__(self, files_path): + self._files_path = files_path + + tracer = al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal( + centre=(0.1, 0.0), einstein_radius=1.0, ell_comps=(0.0, 0.0) + ), + ), + al.Galaxy(redshift=1.0, bulge=al.lp.SersicSph(centre=(0.0, 0.0))), + ] + ) + + samples_summary = al.m.MockSamplesSummary(max_log_likelihood_instance=tracer) + + result = res.Result(samples_summary=samples_summary, analysis=analysis_imaging_7x7) + result.paths = _StubPaths(files_path=tmp_path) + + first = result.positions_likelihood_from(factor=0.1, minimum_threshold=0.2) + + assert (tmp_path / "multiple_image_positions.json").exists() + + # The second call must load the cached positions — solving again raises. + def _poison(*args, **kwargs): + raise AssertionError("point solver re-ran — cached positions not used") + + result_cached = res.Result( + samples_summary=samples_summary, analysis=analysis_imaging_7x7 + ) + result_cached.paths = _StubPaths(files_path=tmp_path) + monkeypatch.setattr( + result_cached, "image_plane_multiple_image_positions", _poison + ) + + second = result_cached.positions_likelihood_from(factor=0.1, minimum_threshold=0.2) + + assert isinstance(second, al.PositionsLH) + assert second.positions.array == pytest.approx(first.positions.array, 1.0e-8) + assert second.threshold == pytest.approx(first.threshold, 1.0e-8)