fix: regenerate datasets left on disk by a small-datasets run

should_simulate() was existence-only and asymmetric. It force-regenerated when
PYAUTO_SMALL_DATASETS=1, but had no corresponding check on the full-resolution
path, so a dataset written by an earlier capped run was reused indefinitely.

Because dataset/ is gitignored in the workspaces, CI clones fresh and always
simulates, and so can never hit this. Locally the directory persists, and since
PYAUTO_SMALL_DATASETS=1 is the default for most harness runs, a single earlier
run leaves 16x16 FITS that every later full-resolution run then loads silently.
The result is deterministic, bit-identical assertion failures that pass in CI
on the same commit (autolens_workspace_test#260) — an active trap, since the
failures look exactly like fresh correctness regressions.

The regime is recorded nowhere on disk, so it is inferred from data.fits: the
cap in Mask2D.circular / Grid2D.uniform rewrites anything larger to exactly
SMALL_DATASETS_SHAPE_NATIVE, so a data.fits at exactly (16, 16) can only have
come from a capped run.

The predicate ends in shutil.rmtree, so it is deliberately narrow:

- exactly the cap shape, never "at or below" — the cap cannot emit 12x12, so
  widening the test buys no detection and only risks real data;
- data.fits by name, never "the first FITS in the directory" — PSF kernels are
  legitimately tiny at full resolution (11x11 is common, and one on disk is
  already byte-identical in size to a capped data.fits), so a glob would
  regenerate every dataset carrying one on every run;
- unknown means no — a missing, unreadable or non-2D data.fits preserves the
  existing existence-only behaviour rather than deleting.

Scope: this covers the imaging manifestation only. Point-source and weak-lensing
datasets are JSON with no FITS, and interferometer datasets keep their shape
under the cap while their values change, so both regress to existence-only and
remain exposed. Closing those needs the regime recorded at write time rather
than inferred at read time; filed separately.

Cost is one FITS header read (~0.6ms warm), and only when a data.fits exists.
No re-simulation in the steady state.

Tests cover all four regime transitions — the bug was precisely that one of the
four was never exercised — plus the false-positive guards, every one of which
asserts a dataset is PRESERVED. Control-tested: the small->full regression test
returns False against the unfixed body. Full suite 1168 passed.

Verified end to end against the original failure: poison the dataset with one
PYAUTO_SMALL_DATASETS=1 run, then run imaging/jax_grad/lp.py under the
full_datasets profile. Before: AssertionError, source gradients ~1e-12. After:
regenerated at full resolution, checks pass.

Co-Authored-By: Claude <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VEHLT33XpVcRt5YCJGLRMJ

Author: Claude

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

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)