[cophus]

The quantem test suite grew from approximately 2k lines in June to 22k lines in October, mostly from tests added during Claude Code sessions. This PR consolidates the suite following the approach @TomaSusi used for abTEM in abTEM/abTEM#589: we removed parametrizations that do not change the code path, merged duplicate tests, rewrote tests that could not fail so that they check computed values, and shared expensive inputs between tests. We also added a CLAUDE.md with rules for writing tests, so that future sessions extend existing tests instead of adding new ones. Each commit message lists every removed or merged test and the test that still covers that behavior.

Results

	dev	this PR
Test IDs	1,460	1,055
Test lines	22,419	20,561
Serial run time	366 s	275 s
Run time with 4 workers	approximately 122 s	approximately 97 s

Tests	Test IDs	Serial run time
Direct ptychography	539 → 342	276 → 200 s
Other diffractive imaging	136 → 98	26 → 16 s
Datastructures	134 → 81	3.2 → 1.5 s
Visualization	88 → 31	< 1 s
Tomography	127 → 104	< 1 s
Utils, core, spectroscopy and config	106 → 69	< 1 s

We measured all run times on an Apple silicon laptop, where repeated runs with 4 workers varied by approximately 10 s. The diffraction, imaging and ML tests are unchanged (see below).

Changes

Over-parametrization. Many tests were parametrized over values that take the same code path, for example kernel name aliases, both direct ptychography classes for code in their shared base class, both store="zip" and store="dir" for every serialization test, and a 4 × 6 grid of upsampling factors and shifts. We kept one case per code path, and where several inexpensive cases share one path, we loop over them inside one test and name the failing case in the assert message.

Duplicate tests. The Fourier and montage direct ptychography files contained copies of the tests for their shared helper functions and 17 separate save and load round trips, which are now 5. The INR and tensor decomposition tests contained identical copies of approximately 140 lines of helper functions, which now live in diffractive_imaging/conftest.py.

Tests that could not fail. The normalization stretch tests asserted 0 <= f(x) <= 1 for inputs that each stretch first clips to [0, 1], and the histogram tests asserted ~isnan(h[h > 0]). The tomography logger tests checked for an events file that SummaryWriter.__init__ always writes, the pure-phase double-mask test used a constant phase that gives zero whether the mask is applied once or twice, and the GPU origin-fit test compared two parallax images that were identically zero. These tests now check computed values.

Run time. The direct ptychography tests now share module-scoped fixtures for the Fourier and montage reconstruction pair, the drift frames and the tilt series acquisitions. The sparse event tests use a mean count of 4 instead of 40, which reduces each from_vector build from 2.2 s to 0.2 s and leaves the largest dense-to-sparse error at 1.3e-5, below the 1e-4 tolerance. The target-residency tests in test_ptychography.py previously rebuilt a 64⁴ simulated dataset for each of 9 test IDs, and now use a single 8 × 8 × 16 × 16 dataset.

CI scheduling. The CI workflow now runs pytest with --dist loadscope, which assigns each test class (or each module's free functions) to a single xdist worker, so that each module- and class-scoped fixture is built once instead of once per worker.

Device leak. On Apple silicon, test_config_set_device left the global device set to mps, so later tests in the same worker ran on MPS. The test now restores the original device.

CLAUDE.md. The new CLAUDE.md asks Claude Code sessions to extend an existing test before adding a new one, to test each code path once, to write asserts that can fail, and to use the smallest problem that exercises the code path. .gitignore previously excluded CLAUDE.md, so we replaced that entry with CLAUDE.local.md, which keeps personal notes out of the repository.

Remaining work

To avoid merge conflicts, we did not modify test files that open PRs also modify. These include tests/diffraction/, which #305 modifies; we audited it and expect a reduction from 256 to approximately 206 test IDs and from 56 s to approximately 32 s once #305 merges. The others are test_optimizermixin.py and four tomography test files (#246), test_vector.py (#273) and test_lattice.py (#236). Approximately 50 of the 64 test IDs in test_optimizermixin.py check dataclass default values and can be removed after #246 merges.

Other issues found

CI does not run the 31 tests marked @pytest.mark.slow, because the workflow does not pass --runslow. These include test_tomography_cpu_recon.py, whose docstring states that it exists to cover the learned reconstruction loop in CI. We should either run the slow tests in a separate scheduled CI job, or remove the slow marker from the tests that CI must cover.

tomography_opt.py:29 reads the pose learning rate from DEFAULT_LRS["pose"], but dataset_models.py:162 defines the key as "pose_lr", so the code always uses the fallback value, which is currently equal to the intended value of 5e-2.

Testing

The full suite passes both serially and with 4 workers (1,023 passed, 32 skipped). We ran the direct ptychography tests three times with 4 workers to check for flaky tests. test_multi_gpu.py requires two CUDA devices, so we linted it but did not run it.

🤖 Generated with Claude Code

[gvarnavi]

lgtmc (looks good to my Claude) 😅

[arthurmccray]

Okay confirmed that this works for multi-gpus. The only issues I found were in diffraction and other areas you deliberately didn't modify, so I don't think there's anything else to add for now.