perf: reuse an already-correctly-capped dataset instead of re-simulating it

autoarray/util/dataset_util.py

@@ -214,6 +214,44 @@ def _stamp_contradicted_by_shape(dataset_path):
     )
 
 
+def _is_capped_at_the_current_cap(dataset_path):
+    """
+    Returns True only when the dataset on disk was written by a capped run
+    **and** was capped to the size in force right now.
+
+    Both halves are required, and the second is the one that is easy to drop.
+    ``SMALLDAT = T`` does not mean "capped at today's cap" -- it means "capped at
+    whatever ``SMALL_DATASETS_SHAPE_NATIVE`` was when this file was written". If
+    that constant is ever changed, every dataset already on disk goes on claiming
+    ``T`` at the old size, and reusing them on the stamp alone would silently
+    feed stale, wrong-sized data to a run that asked for the new cap -- the exact
+    class of silent-stale-dataset bug the stamp was introduced to prevent,
+    reintroduced through the opposite branch.
+
+    This is the mirror image of :func:`_stamp_contradicted_by_shape` on the
+    full-resolution branch, and it exists for the same reason: the stamp records
+    the writer's *environment*, not a measured property of the data. Neither
+    branch may treat it as unfalsifiable.
+
+    Interferometer datasets deliberately never satisfy this. Their ``data.fits``
+    is ``(n_visibilities, 2)`` -- its shape is fixed by the committed uv file and
+    does not change under the cap -- so shape cannot corroborate their stamp, and
+    a capped interferometer dataset is regenerated on every run exactly as it was
+    before. That is the conservative choice and it is deliberate: the alternative
+    is trusting the stamp alone for precisely the family whose corruption is
+    invisible.
+
+    Anything without a readable ``data.fits`` at the top level -- JSON-only
+    datasets, datacubes nesting theirs in ``channel_XXX/``, multi_dataset's
+    prefixed names -- also returns False and is regenerated, preserving today's
+    behaviour for the families this cannot speak about.
+    """
+    return (
+        _small_datasets_stamp_on_disk(dataset_path) is True
+        and _is_small_datasets_on_disk(dataset_path)
+    )
+
+
 def should_simulate(dataset_path):
     """
     Returns True if the dataset at ``dataset_path`` needs to be simulated.
@@ -223,10 +261,16 @@ def should_simulate(dataset_path):
     masks and grids to ``SMALL_DATASETS_SHAPE_NATIVE``. Both directions are
     checked:
 
-    - Entering the **small** regime, any existing dataset is deleted so the
+    - Entering the **small** regime, an 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.
+      mismatches between full-resolution FITS on disk and the capped mask/grid.
+      It is **kept** only when it is already capped to the size in force now --
+      stamped ``SMALLDAT = T`` *and* measuring exactly
+      ``SMALL_DATASETS_SHAPE_NATIVE`` (:func:`_is_capped_at_the_current_cap`).
+      Before the stamp this branch was unconditional because it had no way to
+      tell a dataset produced by the same cap from one produced by a different
+      run; it does now, and re-simulating an already-correct dataset is pure
+      cost on every smoke run across the workspaces.
     - 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 taken from the ``SMALLDAT`` header card that
@@ -316,7 +360,9 @@ def should_simulate(dataset_path):
     that is true of weak lensing only.
     """
     if os.environ.get("PYAUTO_SMALL_DATASETS") == "1":
-        if Path(dataset_path).exists():
+        if Path(dataset_path).exists() and not _is_capped_at_the_current_cap(
+            dataset_path
+        ):
             shutil.rmtree(dataset_path)
 
         return not Path(dataset_path).exists()

test_autoarray/util/test_dataset_util.py

@@ -113,6 +113,7 @@ def test__env_set__non_square_above_cap__center_crops_to_16x16(monkeypatch):
     _on_disk_shape_native,
     _small_datasets_stamp_on_disk,
     _stamp_contradicted_by_shape,
+    _is_capped_at_the_current_cap,
     SMALL_DATASETS_HEADER_KEY,
 )
 
@@ -188,13 +189,88 @@ def test__small_regime__existing_full_dataset__is_deleted_and_resimulated(
     assert not dataset_path.exists()
 
 
-def test__small_regime__existing_small_dataset__is_still_deleted_and_resimulated(
+def test__small_regime__dataset_already_at_the_current_cap__is_kept(
     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.
+    # This branch used to be unconditional because it had no way to tell a
+    # dataset produced by the SAME cap from any other. The stamp gives it one,
+    # so an already-correct dataset is reused instead of re-simulated -- pure
+    # cost otherwise, on every smoke run across ~253 call sites.
     monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1")
-    dataset_path = _write_dataset(tmp_path / "dataset", SMALL_DATASETS_SHAPE_NATIVE)
+    dataset_path = _write_dataset(
+        tmp_path / "dataset", SMALL_DATASETS_SHAPE_NATIVE, stamp=True
+    )
+
+    assert should_simulate(str(dataset_path)) is False
+    assert (dataset_path / "data.fits").exists()
+
+
+def test__small_regime__stamped_capped_but_at_a_DIFFERENT_cap__is_regenerated(
+    monkeypatch, tmp_path
+):
+    # THE TRAP. `SMALLDAT = T` means "capped at whatever the cap was when this
+    # was written", NOT "capped at today's cap". Reusing on the stamp alone
+    # would silently feed stale wrong-sized data to a run that asked for a
+    # different cap -- the same silent-stale-dataset bug the stamp exists to
+    # prevent, reintroduced through the opposite branch.
+    monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1")
+    other_cap = (SMALL_DATASETS_SHAPE_NATIVE[0] * 2, SMALL_DATASETS_SHAPE_NATIVE[1] * 2)
+    dataset_path = _write_dataset(tmp_path / "dataset", other_cap, stamp=True)
+
+    assert _small_datasets_stamp_on_disk(str(dataset_path)) is True  # claims capped
+    assert should_simulate(str(dataset_path)) is True  # but not at THIS cap
+    assert not dataset_path.exists()
+
+
+def test__small_regime__unstamped_legacy_dataset__is_regenerated(
+    monkeypatch, tmp_path
+):
+    # Every dataset written before the stamp landed is unstamped. Reuse requires
+    # positive evidence, so these keep the old always-regenerate behaviour.
+    monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1")
+    dataset_path = _write_dataset(
+        tmp_path / "dataset", SMALL_DATASETS_SHAPE_NATIVE, stamp=None
+    )
+
+    assert should_simulate(str(dataset_path)) is True
+    assert not dataset_path.exists()
+
+
+def test__small_regime__full_resolution_dataset__is_regenerated(monkeypatch, tmp_path):
+    # A dataset stamped F is full resolution; the capped run needs a capped one.
+    monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1")
+    dataset_path = _write_dataset(tmp_path / "dataset", (180, 180), stamp=False)
+
+    assert should_simulate(str(dataset_path)) is True
+    assert not dataset_path.exists()
+
+
+def test__small_regime__interferometer_dataset__is_always_regenerated(
+    monkeypatch, tmp_path
+):
+    # Deliberate and conservative. An interferometer data.fits is
+    # (n_visibilities, 2) -- its shape is fixed by the committed uv file and does
+    # not change under the cap -- so shape cannot corroborate its stamp. Rather
+    # than trust the stamp alone for precisely the family whose corruption is
+    # invisible, this family keeps regenerating every run, as before.
+    monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1")
+    dataset_path = _write_dataset(tmp_path / "dataset", (360, 2), stamp=True)
+
+    assert _small_datasets_stamp_on_disk(str(dataset_path)) is True
+    assert should_simulate(str(dataset_path)) is True
+    assert not dataset_path.exists()
+
+
+def test__small_regime__dataset_with_no_top_level_data_fits__is_regenerated(
+    monkeypatch, tmp_path
+):
+    # JSON-only datasets, datacubes nesting FITS in channel_XXX/, multi_dataset's
+    # prefixed names: no readable data.fits means no positive evidence, so they
+    # regenerate exactly as they did before.
+    monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1")
+    dataset_path = tmp_path / "dataset"
+    dataset_path.mkdir()
+    (dataset_path / "dataset.json").write_text("{}")
 
     assert should_simulate(str(dataset_path)) is True
     assert not dataset_path.exists()