small datasets: honour the cap card, the JAX switch and the cap in mesh shapes

Three findings of the phase-8 diagnosis of the user-workspace smoke timings,
all of them the same shape: a harness-level fact the library did not act on.

1. Capped datasets whose shape cannot corroborate their stamp were deleted and
   re-simulated on every run. `_is_capped_at_the_current_cap` required
   `SMALLDAT = T` *and* a measured `data.fits` at exactly the cap, because
   `SMALLDAT` records "the env var was set at write time", not "capped at
   today's cap". Interferometer data is `(n_visibilities, 2)`, multi_dataset
   prefixes its FITS and datacubes nest theirs per channel, so none of them
   could ever corroborate: 5.2-6.5 s per script, ~26 s per autolens_workspace
   CI run and ~28 s per autogalaxy_workspace one, plus every local run.

   The `SMALLSHP` card (PyAutoNerves#159) states the proposition directly, so
   corroboration is no longer needed and those families are reachable. Two
   things keep the safety property intact. The card is read *with* the
   both-axes contradiction guard the full-resolution branch already applies --
   both cards record the writing *process*, so a 180x180 image written in a
   shell exporting `PYAUTO_SMALL_DATASETS=1` carries them while being full
   resolution, and reusing it in a capped run is the shape-mismatch class the
   capped branch exists to prevent. And an absent or malformed card falls back
   to the shape heuristic rather than being coerced, so every pre-card dataset
   behaves exactly as it did.

   The path resolution (`_capped_data_paths`) is widened for the capped branch
   only, by exact suffix at two known levels -- `data.fits`, then
   `{waveband}_data.fits`, then `channel_*/data.fits`, never a `*.fits` glob.
   The full-resolution branch still reads `data.fits` alone: widening the
   capped branch can only move a dataset from delete to keep, whereas widening
   the destructive one is its own change with its own review.

   `test__small_regime__interferometer_dataset__is_always_regenerated` asserted
   the conservative behaviour this issue exists to change, and is rewritten to
   pin the half that stands: without a cap card, that family still regenerates.

2. `apply_sparse_operator(use_jax=True)` ignored `PYAUTO_DISABLE_JAX=1`, so the
   smoke profile paid 2.3-3.2 s of JIT for a backend it had asked to disable.
   The interferometer method now lets the harness win over an explicit `True`,
   through `autonerves.test_mode.disable_jax()`.

   The imaging method is deliberately left alone, and says why in its docstring:
   it has no `use_jax` argument to override -- it *is* the JAX implementation,
   and its NumPy sibling `apply_sparse_operator_cpu` returns a different
   operator class and requires numba. Routing to that under an env var would
   change the returned type based on the environment, which is a larger and
   (on these measurements) unmotivated change.

3. Pixelization mesh `shape=` did not honour the cap the data does, so a capped
   run reconstructed 1600-2500 source pixels from ~80 image pixels. The
   rectangular meshes and `image_mesh.Overlay` now cap per axis under
   `PYAUTO_SMALL_DATASETS=1`, with the same `respect_small_datasets=True`
   escape hatch `Grid2D.uniform` carries. Measured on the four HowTo chapter-3
   scripts whose guards this subsumes: 43.6 s -> 17.3 s.

`disable_jax` is imported defensively: an autonerves too old to carry the
predicate is exactly the case the sibling literal in `dataset_util` exists for,
and a hard import would turn a safe degradation into an `ImportError` at module
load.

Refs #528

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

Author: claude

autoarray/dataset/imaging/dataset.py

@@ -602,6 +602,20 @@ def apply_sparse_operator(
         Imaging
             A new `Imaging` dataset with the precomputed `ImagingSparseOperator` attached, enabling
             efficient pixelized source reconstruction via the sparse linear algebra formalism.
+
+        Notes
+        -----
+        `PYAUTO_DISABLE_JAX=1` is *not* honoured here, unlike
+        `Interferometer.apply_sparse_operator`, and the asymmetry is deliberate. There the
+        variable overrides a `use_jax` argument that already selects between two backends
+        computing the same operator. Here there is no such argument: this method is the JAX
+        implementation, and the NumPy/CPU alternative is the separately named
+        `apply_sparse_operator_cpu`, which returns a different operator class
+        (`SparseLinAlgImagingNumba`) and requires numba. Silently returning that under an
+        environment variable would change the type of the returned object based on the
+        environment, which is a larger change than honouring a switch -- and an unmeasured
+        one: every JIT cost the phase-8 workspace timings attribute to this variable
+        (2.3-3.2 s per script) was on the interferometer path.
         """
 
         if self.psf is not None and self.psf.convolve_over_sample_size > 1:

autoarray/dataset/interferometer/dataset.py

@@ -5,6 +5,24 @@
 from autonerves.fitsable import ndarray_via_fits_from
 from autonerves import cached_property
 
+try:
+    from autonerves.test_mode import disable_jax
+except ImportError:
+    # `disable_jax()` arrives in the autonerves release that closes
+    # PyAutoNerves#159. Importing it unconditionally would make an autonerves
+    # older than that release an `ImportError` at module load -- and a
+    # `--no-deps` install, an editable checkout or a hand-built virtualenv can
+    # all put one on the path regardless of the floor in `pyproject.toml`,
+    # which constrains resolution only. This is the same trade `dataset_util`
+    # records against `SMALL_DATASETS_HEADER_KEY`: degrade to the predicate's
+    # own one-line body rather than fail hard to avoid restating it. Delete the
+    # fallback when the floor names a release carrying the predicate.
+    import os
+
+    def disable_jax():
+        return os.environ.get("PYAUTO_DISABLE_JAX") == "1"
+
+
 from autoarray.dataset.abstract.dataset import AbstractDataset
 from autoarray.dataset.grids import GridsDataset
 from autoarray.inversion.inversion.interferometer.inversion_interferometer_util import (
@@ -266,6 +284,14 @@ def apply_sparse_operator(
         use_jax
             If `True`, JAX is used to accelerate the NUFFT precision matrix computation.
 
+            `PYAUTO_DISABLE_JAX=1` overrides this to `False`. That variable is a
+            harness-level switch, not a preference: it is the documented way to force the
+            NumPy path (the workspace `start_here` guides name it beside `use_jax=False`),
+            and the smoke profiles set it so a fast run does not pay a JIT compile. An
+            explicit `use_jax=True` in a script -- which is the right thing for a script
+            demonstrating the production path to say -- must therefore not defeat it, or
+            the harness pays 2.3-3.2 s of compile for a backend it asked to disable.
+
         Precondition
         ------------
         Every visibility must have equal real and imaginary noise sigma
@@ -289,6 +315,9 @@ def apply_sparse_operator(
             If any visibility has unequal real and imaginary noise sigma.
         """
 
+        if disable_jax():
+            use_jax = False
+
         noise_map_real = np.asarray(self.noise_map.real)
         noise_map_imag = np.asarray(self.noise_map.imag)
 

autoarray/inversion/mesh/image_mesh/overlay.py

@@ -8,6 +8,7 @@
 
 from autoarray.geometry import geometry_util
 from autoarray.structures.grids import grid_2d_util
+from autoarray.util.dataset_util import cap_mesh_shape_for_small_datasets
 from autoarray import numba_util
 
 
@@ -158,7 +159,7 @@ def overlay_via_unmasked_overlaid_from(
 
 
 class Overlay(AbstractImageMesh):
-    def __init__(self, shape=(3, 3)):
+    def __init__(self, shape=(3, 3), respect_small_datasets: bool = True):
         """
         Computes an image-mesh by overlaying a uniform grid of (y,x) coordinates over the masked image that the
         pixelization is fitting.
@@ -176,10 +177,18 @@ def __init__(self, shape=(3, 3)):
         ----------
         shape
             The 2D shape of the grid which is overlaid over the grid to determine the image mesh.
+        respect_small_datasets
+            When `PYAUTO_SMALL_DATASETS=1` is set, `shape` is capped per axis to the small-datasets
+            cap, matching the cap `Grid2D.uniform` and `Mask2D.circular` apply to the data. Pass
+            `False` to opt out for an image mesh whose resolution is load-bearing for the script.
         """
 
         super().__init__()
 
+        shape = cap_mesh_shape_for_small_datasets(
+            shape, respect_small_datasets=respect_small_datasets
+        )
+
         self.shape = (int(shape[0]), int(shape[1]))
 
     def image_plane_mesh_grid_from(

autoarray/inversion/mesh/mesh/rectangular_bilinear_adapt_density.py

@@ -10,6 +10,7 @@ class RectangularBilinearAdaptDensity(RectangularRTUAdaptDensity):
     def __init__(
         self,
         shape: Tuple[int, int] = (3, 3),
+        respect_small_datasets: bool = True,
     ):
         """
         A rectangular mesh of pixels used to reconstruct a source on a regular
@@ -71,7 +72,9 @@ def __init__(
             If either dimension is less than 3, as a minimum of 3×3 pixels
             is required to define interior and boundary structure.
         """
-        super().__init__(shape=shape)
+        super().__init__(
+            shape=shape, respect_small_datasets=respect_small_datasets
+        )
 
     @property
     def interpolator_kwargs(self) -> dict:

autoarray/inversion/mesh/mesh/rectangular_bilinear_adapt_image.py

@@ -12,6 +12,7 @@ def __init__(
         shape: Tuple[int, int] = (3, 3),
         weight_power: float = 1.0,
         weight_floor: float = 0.0,
+        respect_small_datasets: bool = True,
     ):
         """
         A rectangular mesh of pixels used to reconstruct a source on a regular
@@ -69,6 +70,7 @@ def __init__(
             shape=shape,
             weight_power=weight_power,
             weight_floor=weight_floor,
+            respect_small_datasets=respect_small_datasets,
         )
 
     @property

autoarray/inversion/mesh/mesh/rectangular_rtu_adapt_density.py

@@ -9,6 +9,7 @@
 from autoarray.inversion.mesh.border_relocator import BorderRelocator
 
 from autoarray.structures.grids import grid_2d_util
+from autoarray.util.dataset_util import cap_mesh_shape_for_small_datasets
 
 from autoarray import exc
 
@@ -71,6 +72,7 @@ def __init__(
         shape: Tuple[int, int] = (3, 3),
         bandwidth: Optional[float] = None,
         n_knots: Optional[int] = None,
+        respect_small_datasets: bool = True,
     ):
         """
         A rectangular mesh of pixels used to reconstruct a source on a regular
@@ -132,6 +134,11 @@ def __init__(
         n_knots
             Size of the fixed knot table used to invert the CDF. Defaults to
             the kernel default.
+        respect_small_datasets
+            When ``PYAUTO_SMALL_DATASETS=1`` is set, `shape` is capped per axis
+            to the small-datasets cap, matching the cap `Grid2D.uniform` and
+            `Mask2D.circular` apply to the data. Pass ``False`` to opt out for a
+            mesh whose resolution is load-bearing for the script.
 
         Raises
         ------
@@ -144,6 +151,10 @@ def __init__(
             KERNEL_CDF_DEFAULT_KNOTS,
         )
 
+        shape = cap_mesh_shape_for_small_datasets(
+            shape, respect_small_datasets=respect_small_datasets
+        )
+
         if shape[0] <= 2 or shape[1] <= 2:
             raise exc.MeshException(
                 "The rectangular pixelization must be at least dimensions 3x3"

autoarray/inversion/mesh/mesh/rectangular_rtu_adapt_image.py

@@ -15,6 +15,7 @@ def __init__(
         weight_floor: float = 0.0,
         bandwidth: Optional[float] = None,
         n_knots: Optional[int] = None,
+        respect_small_datasets: bool = True,
     ):
         """
         A uniform rectangular mesh of pixels used to reconstruct a source on a
@@ -82,7 +83,12 @@ def __init__(
             the kernel default.
         """
 
-        super().__init__(shape=shape, bandwidth=bandwidth, n_knots=n_knots)
+        super().__init__(
+            shape=shape,
+            bandwidth=bandwidth,
+            n_knots=n_knots,
+            respect_small_datasets=respect_small_datasets,
+        )
 
         self.weight_power = weight_power
         self.weight_floor = weight_floor