From 6b29d3bfd5f9fd7b212d3dc676d0be347d069318 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Thu, 9 Jul 2026 11:37:21 +0100
Subject: [PATCH] Cap SimulatorShearYX random-position count under
 PYAUTO_SMALL_DATASETS

via_tracer_random_positions_from routes n_galaxies through autoarray's new
cap_catalogue_size_for_small_datasets (200 -> 25 in smoke mode); the explicit
user grid of via_tracer_from is deliberately never mutated. With the
should_simulate adoption already shipped (#581), scripts/weak/modeling.py under
PYAUTO_SMALL_DATASETS=1 + PYAUTO_TEST_MODE=1 now regenerates a 25-galaxy
catalogue and completes in ~30s (vs ~7 min).

Step 9 of the weak-lensing series (PyAutoLabs/PyAutoLens#583).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 autolens/weak/simulator.py                    |  6 +++
 .../weak/test_simulator_small_datasets.py     | 43 +++++++++++++++++++
 2 files changed, 49 insertions(+)
 create mode 100644 test_autolens/weak/test_simulator_small_datasets.py

diff --git a/autolens/weak/simulator.py b/autolens/weak/simulator.py
index 490a38c68..292fd62ae 100644
--- a/autolens/weak/simulator.py
+++ b/autolens/weak/simulator.py
@@ -113,6 +113,12 @@ def via_tracer_random_positions_from(
         name
             Optional label passed through to ``WeakDataset.name``.
         """
+        from autoarray.util import dataset_util
+
+        # Smoke-mode cap (PYAUTO_SMALL_DATASETS=1): generated catalogue sizes only —
+        # via_tracer_from's user-provided grid is never mutated.
+        n_galaxies = dataset_util.cap_catalogue_size_for_small_datasets(n_galaxies)
+
         positions = self._rng.uniform(
             low=-grid_extent, high=grid_extent, size=(n_galaxies, 2)
         )
diff --git a/test_autolens/weak/test_simulator_small_datasets.py b/test_autolens/weak/test_simulator_small_datasets.py
new file mode 100644
index 000000000..f013908f1
--- /dev/null
+++ b/test_autolens/weak/test_simulator_small_datasets.py
@@ -0,0 +1,43 @@
+import autolens as al
+
+
+def _tracer():
+    lens = al.Galaxy(
+        redshift=0.5,
+        mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.6),
+    )
+    return al.Tracer(galaxies=[lens, al.Galaxy(redshift=1.0)])
+
+
+def test__random_positions__capped_under_small_datasets(monkeypatch):
+    monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1")
+
+    dataset = al.SimulatorShearYX(noise_sigma=0.3, seed=1).via_tracer_random_positions_from(
+        tracer=_tracer(), n_galaxies=200, grid_extent=3.0
+    )
+
+    assert dataset.n_galaxies == 25
+
+
+def test__random_positions__uncapped_without_env_var(monkeypatch):
+    monkeypatch.delenv("PYAUTO_SMALL_DATASETS", raising=False)
+
+    dataset = al.SimulatorShearYX(noise_sigma=0.3, seed=1).via_tracer_random_positions_from(
+        tracer=_tracer(), n_galaxies=40, grid_extent=3.0
+    )
+
+    assert dataset.n_galaxies == 40
+
+
+def test__explicit_grid__never_capped(monkeypatch):
+    import autoarray as aa
+    import numpy as np
+
+    monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1")
+
+    grid = aa.Grid2DIrregular(values=np.random.default_rng(1).uniform(-3, 3, (60, 2)))
+    dataset = al.SimulatorShearYX(noise_sigma=0.3, seed=1).via_tracer_from(
+        tracer=_tracer(), grid=grid
+    )
+
+    assert dataset.n_galaxies == 60