diff --git a/autoarray/util/dataset_util.py b/autoarray/util/dataset_util.py index 0e8888888..ede4f03eb 100644 --- a/autoarray/util/dataset_util.py +++ b/autoarray/util/dataset_util.py @@ -69,25 +69,121 @@ def cap_array_2d_for_small_datasets(array_2d, pixel_scales): ) +def _on_disk_shape_native(data_path): + """ + Returns the ``(rows, columns)`` shape of the first 2D image in the FITS file + at ``data_path``, or ``None`` if that cannot be determined. + + Only the headers are read, never the pixel data, so this costs a single + small read regardless of dataset size. + + ``None`` means "unknown", and every caller must treat it as "leave the + dataset alone" — this function feeds a destructive predicate, so an + unreadable or unconventional file must never be grounds for deleting it. + """ + from astropy.io import fits + + try: + with fits.open(data_path) as hdu_list: + for hdu in hdu_list: + header = hdu.header + if header.get("NAXIS") == 2: + # NAXIS1 is the fastest-varying axis (columns), NAXIS2 the + # rows, so the numpy-order shape is (NAXIS2, NAXIS1). + return (header["NAXIS2"], header["NAXIS1"]) + except Exception: + return None + + return None + + +def _is_small_datasets_on_disk(dataset_path): + """ + Returns True if the dataset on disk at ``dataset_path`` was written by a + simulator running under ``PYAUTO_SMALL_DATASETS=1``. + + The regime is not recorded anywhere on disk, so it is inferred from the + shape of ``data.fits``: the cap in ``Mask2D.circular`` / ``Grid2D.uniform`` + rewrites anything larger than ``SMALL_DATASETS_SHAPE_NATIVE`` to *exactly* + that shape, so an on-disk ``data.fits`` at exactly (16, 16) can only have + come from a capped run. + + Three deliberate narrownesses, all of them because this predicate ends in + ``shutil.rmtree`` and a false positive silently deletes a user's data: + + - **Exactly** the cap shape, never "at or below" it. The cap cannot emit + 12x12, so widening the test buys no detection and only adds risk. + - **``data.fits`` by name**, never "the first FITS in the directory". PSF + kernels are legitimately tiny at full resolution (11x11 is common), and a + glob would regenerate every dataset carrying one on every single run. + - **Unknown means no.** A missing, unreadable or non-2D ``data.fits`` + returns False, preserving the existence-only behaviour for the dataset + families this cannot speak about (see the caveat in ``should_simulate``). + """ + data_path = Path(dataset_path) / "data.fits" + + if not data_path.exists(): + return False + + return _on_disk_shape_native(data_path) == SMALL_DATASETS_SHAPE_NATIVE + + def should_simulate(dataset_path): """ Returns True if the dataset at ``dataset_path`` needs to be simulated. - When ``PYAUTO_SMALL_DATASETS=1`` is active, any existing dataset - is deleted so the simulator re-creates it at the reduced resolution. This - avoids shape mismatches between full-resolution FITS files on disk and the - 15x15 mask/grid cap applied by the env var. + A dataset is invalid when it was simulated under a different resolution + regime than the one in force now, because ``PYAUTO_SMALL_DATASETS=1`` caps + masks and grids to ``SMALL_DATASETS_SHAPE_NATIVE``. Both directions are + checked: + + - Entering the **small** regime, any existing dataset is deleted so the + simulator re-creates it at the reduced resolution, avoiding shape + mismatches between full-resolution FITS on disk and the capped + mask/grid. + - Entering the **full** regime, a dataset left behind by an earlier capped + run is likewise deleted. Existence alone cannot distinguish the two, so + the regime is inferred from the data on disk + (``_is_small_datasets_on_disk``). + + That second check is what makes a local FAIL mean something. ``dataset/`` + is gitignored in the workspaces, so CI clones fresh and always simulates, + while a local checkout keeps its dataset indefinitely — and since + ``PYAUTO_SMALL_DATASETS=1`` is the default for most harness runs, a single + earlier run would leave capped FITS that every later full-resolution run + then loaded silently, producing deterministic, environment-only failures + that could not be reproduced in CI (autolens_workspace_test#260). Use this as a drop-in replacement for ``not path.exists(dataset_path)`` in the workspace auto-simulation pattern:: if aa.util.dataset.should_simulate(dataset_path): subprocess.run([sys.executable, "scripts/.../simulator.py"], check=True) + + Known gap + --------- + The full-regime check reads ``data.fits``, so it covers imaging-style + datasets only. It cannot see a stale capped dataset whose corruption is not + visible in that file's shape: + + - point-source and weak-lensing datasets, which are JSON with no FITS; + - interferometer datasets, whose visibility count is fixed by the uv file + while the real-space grid behind it is capped, so the capped and full + files share a shape and differ only in values. + + Those regress to the previous existence-only behaviour rather than being + fixed here. Closing them needs the regime recorded at write time rather + than inferred at read time. """ if os.environ.get("PYAUTO_SMALL_DATASETS") == "1": if Path(dataset_path).exists(): shutil.rmtree(dataset_path) + return not Path(dataset_path).exists() + + if Path(dataset_path).exists() and _is_small_datasets_on_disk(dataset_path): + shutil.rmtree(dataset_path) + return not Path(dataset_path).exists() SMALL_DATASETS_N_CATALOGUE = 25 diff --git a/test_autoarray/util/test_dataset_util.py b/test_autoarray/util/test_dataset_util.py index 1d9c678e7..71fa9024f 100644 --- a/test_autoarray/util/test_dataset_util.py +++ b/test_autoarray/util/test_dataset_util.py @@ -90,3 +90,186 @@ def test__env_set__non_square_above_cap__center_crops_to_16x16(monkeypatch): assert result.shape_native == SMALL_DATASETS_SHAPE_NATIVE assert pixel_scales == SMALL_DATASETS_PIXEL_SCALES + + +""" +__should_simulate — regime transitions__ + +`should_simulate` must regenerate a dataset whenever the resolution regime on +disk differs from the one in force. Before autolens_workspace_test#260 only the +full->small transition was implemented; small->full silently reused capped FITS +at full resolution, producing deterministic failures that no CI run could +reproduce (CI clones fresh, so it never has a stale dataset). + +All four transitions are covered below because the bug was precisely that one +of the four was never exercised. +""" + +import json + +from autoarray.util.dataset_util import ( + should_simulate, + _is_small_datasets_on_disk, + _on_disk_shape_native, +) + +def _write_dataset(dataset_path, shape, extra_files=()): + """Write a minimal dataset directory containing a `data.fits` of `shape`.""" + dataset_path.mkdir(parents=True, exist_ok=True) + + aa.output_to_fits( + values=np.ones(shape), + file_path=str(dataset_path / "data.fits"), + overwrite=True, + ) + + for name, file_shape in extra_files: + aa.output_to_fits( + values=np.ones(file_shape), + file_path=str(dataset_path / name), + overwrite=True, + ) + + return dataset_path + + +def test__small_regime__existing_full_dataset__is_deleted_and_resimulated( + monkeypatch, tmp_path +): + monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1") + dataset_path = _write_dataset(tmp_path / "dataset", (180, 180)) + + assert should_simulate(str(dataset_path)) is True + assert not dataset_path.exists() + + +def test__small_regime__existing_small_dataset__is_still_deleted_and_resimulated( + monkeypatch, tmp_path +): + # The small path is unconditional by design: it cannot know the capped + # dataset on disk was produced by the SAME cap, so it always regenerates. + monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1") + dataset_path = _write_dataset(tmp_path / "dataset", SMALL_DATASETS_SHAPE_NATIVE) + + assert should_simulate(str(dataset_path)) is True + assert not dataset_path.exists() + + +def test__full_regime__stale_small_dataset__is_deleted_and_resimulated( + monkeypatch, tmp_path +): + # THE REGRESSION TEST. Before the fix this returned False and the capped + # FITS were loaded at full resolution. + monkeypatch.delenv("PYAUTO_SMALL_DATASETS", raising=False) + dataset_path = _write_dataset(tmp_path / "dataset", SMALL_DATASETS_SHAPE_NATIVE) + + assert should_simulate(str(dataset_path)) is True + assert not dataset_path.exists() + + +def test__full_regime__full_dataset__is_kept(monkeypatch, tmp_path): + monkeypatch.delenv("PYAUTO_SMALL_DATASETS", raising=False) + dataset_path = _write_dataset(tmp_path / "dataset", (180, 180)) + + assert should_simulate(str(dataset_path)) is False + assert (dataset_path / "data.fits").exists() + + +def test__full_regime__absent_dataset__simulates(monkeypatch, tmp_path): + monkeypatch.delenv("PYAUTO_SMALL_DATASETS", raising=False) + + assert should_simulate(str(tmp_path / "does_not_exist")) is True + + +""" +__should_simulate — false-positive guards__ + +The full-regime branch ends in `shutil.rmtree`, so every one of these asserts +that a dataset is PRESERVED. A regression here silently deletes real data. +""" + + +def test__full_regime__tiny_psf_alongside_full_data__is_kept(monkeypatch, tmp_path): + # PSF kernels are legitimately tiny at full resolution (11x11 is the common + # workspace value, and a 16x16 PSF is a plausible one). The check must key + # on `data.fits` by name — a "first FITS in the directory" implementation + # would delete this dataset on every run forever. + monkeypatch.delenv("PYAUTO_SMALL_DATASETS", raising=False) + dataset_path = _write_dataset( + tmp_path / "dataset", + (180, 180), + extra_files=(("psf.fits", (11, 11)), ("noise_map.fits", (180, 180))), + ) + + assert should_simulate(str(dataset_path)) is False + assert (dataset_path / "data.fits").exists() + assert (dataset_path / "psf.fits").exists() + + +def test__full_regime__psf_at_exactly_the_cap_shape__does_not_trigger_deletion( + monkeypatch, tmp_path +): + monkeypatch.delenv("PYAUTO_SMALL_DATASETS", raising=False) + dataset_path = _write_dataset( + tmp_path / "dataset", + (180, 180), + extra_files=(("psf.fits", SMALL_DATASETS_SHAPE_NATIVE),), + ) + + assert should_simulate(str(dataset_path)) is False + assert (dataset_path / "data.fits").exists() + + +def test__full_regime__below_cap_data__is_kept_because_the_cap_emits_exactly_16x16( + monkeypatch, tmp_path +): + # The cap rewrites anything larger to EXACTLY (16, 16) and never produces + # 12x12, so a 12x12 dataset was not capped and must be left alone. This is + # why the predicate is `==` and not `<=`. + monkeypatch.delenv("PYAUTO_SMALL_DATASETS", raising=False) + dataset_path = _write_dataset(tmp_path / "dataset", (12, 12)) + + assert should_simulate(str(dataset_path)) is False + assert (dataset_path / "data.fits").exists() + + +def test__full_regime__json_only_dataset__is_kept(monkeypatch, tmp_path): + # Point-source and weak-lensing datasets carry no FITS at all. The check + # cannot speak about them, so it must fall back to existence-only rather + # than delete. (These remain exposed to the underlying bug — see the + # "Known gap" section of should_simulate's docstring.) + monkeypatch.delenv("PYAUTO_SMALL_DATASETS", raising=False) + dataset_path = tmp_path / "dataset" + dataset_path.mkdir() + (dataset_path / "point_dataset.json").write_text(json.dumps({"positions": []})) + + assert should_simulate(str(dataset_path)) is False + assert (dataset_path / "point_dataset.json").exists() + + +def test__full_regime__unreadable_data_fits__is_kept(monkeypatch, tmp_path): + # "Unknown regime" must never mean "delete". + monkeypatch.delenv("PYAUTO_SMALL_DATASETS", raising=False) + dataset_path = tmp_path / "dataset" + dataset_path.mkdir() + (dataset_path / "data.fits").write_bytes(b"not a fits file") + + assert _on_disk_shape_native(dataset_path / "data.fits") is None + assert should_simulate(str(dataset_path)) is False + assert (dataset_path / "data.fits").exists() + + +def test__is_small_datasets_on_disk__reads_shape_from_header(tmp_path): + small = _write_dataset(tmp_path / "small", SMALL_DATASETS_SHAPE_NATIVE) + full = _write_dataset(tmp_path / "full", (180, 180)) + + assert _is_small_datasets_on_disk(str(small)) is True + assert _is_small_datasets_on_disk(str(full)) is False + + +def test__on_disk_shape_native__is_row_column_ordered(tmp_path): + # NAXIS1 is columns and NAXIS2 is rows, so a non-square array must come + # back in numpy (rows, columns) order rather than transposed. + dataset_path = _write_dataset(tmp_path / "dataset", (30, 50)) + + assert _on_disk_shape_native(dataset_path / "data.fits") == (30, 50)