From 713be1939c0e68bd49cf07347615186be4f7ecda Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 14:45:25 -0700 Subject: [PATCH 01/19] Consolidate datastructures tests test_dataset.py (30 -> 16 IDs): - Drop the array/name/signal_units getter-setter and repr/str tests; the origin/sampling/units setters and their length validation stay, as one loop. - test_mean/max/min only checked types and shapes; replaced by one test that compares each reduction (all axes, both axes, axis 0) against numpy. - The four 4D/single-axis crop tests become one table that checks shape and origin, both out of place and in place; test_crop is kept. - test_bin and TestBinMetadata merged (shape, sampling, origin, in place). - Down/upsample resample tests merged into one table with their own center tolerances; odd/even and API-error tests kept. - Removed the unused 3D fixtures. test_dataset4dstem.py (15 -> 8 IDs): - The six dp_mean/max/median property and get_* tests checked only shapes; one loop now checks each against numpy and checks the attach/cache contract. - getitem 2D/3D merged; virtual image now checks the summed values. - Deleted test_show: its monkeypatch targeted visualization.show_2d, which dataset2d had already imported, so it rendered a real figure and only checked for a 2-tuple. Dataset4d.show is a one-line forward. test_dataset3d_show.py (25 -> 3 IDs): - Parametrized cases become loops that name the failing kwargs. Duplicate branches dropped (ncols/max -1, step 2). The long dataset is 25 frames instead of 100, so max=None draws 25 subplots, not 100; the default max=20 truncation is still exercised. test_autoserialize.py (16 -> 6 IDs): - store=zip/dir only changes the container, so only the Dummy test loops over both stores; the other tests use zip. - Nested roundtrip and nested skip merged (same class, two saves). - test_torch_tensor_roundtrip removed; test_torch_tensor_and_optimizer checks the same tensor. Its optimizer check compared state_dict keys, which are always {state, param_groups}; it now checks the restored lr. Co-Authored-By: Claude Opus 5.5 --- tests/datastructures/test_autoserialize.py | 241 ++++++---------- tests/datastructures/test_dataset.py | 301 +++++--------------- tests/datastructures/test_dataset3d_show.py | 101 ++++--- tests/datastructures/test_dataset4dstem.py | 140 ++------- 4 files changed, 228 insertions(+), 555 deletions(-) diff --git a/tests/datastructures/test_autoserialize.py b/tests/datastructures/test_autoserialize.py index d4abf0e1f..35e007e86 100644 --- a/tests/datastructures/test_autoserialize.py +++ b/tests/datastructures/test_autoserialize.py @@ -1,5 +1,4 @@ import numpy as np -import pytest import torch import torch.nn as nn @@ -24,54 +23,57 @@ def __init__(self): self._private = "secret" -@pytest.mark.parametrize("store", ["zip", "dir"]) -def test_dummy_save_load_and_skip(tmp_path, store): +def test_dummy_save_load_and_skip(tmp_path): """ + For both the zip and the directory store: 1. Save/load Dummy with no skips: all fields rehydrate. 2. Then save but skip 'count' and skip 'child', verifying each is dropped. 3. Then save but skip '_private', verifying only _private is dropped. """ - # ——— (a) no skips ——— - d = Dummy() - tgt = tmp_path / ("dummy_full.zip" if store == "zip" else "dummy_full_dir") - d.save(str(tgt), mode="w", store=store) - loaded_full = load(str(tgt)) + for store in ("zip", "dir"): + suffix = ".zip" if store == "zip" else "_dir" + d = Dummy() - # types - assert isinstance(loaded_full, Dummy) - assert isinstance(loaded_full.child, Child) + # ——— (a) no skips ——— + tgt = tmp_path / f"dummy_full{suffix}" + d.save(str(tgt), mode="w", store=store) + loaded_full = load(str(tgt)) - # values round-trip - np.testing.assert_allclose(loaded_full.arr, d.arr) - assert loaded_full.count == d.count - np.testing.assert_allclose(loaded_full.child.values, d.child.values) - assert loaded_full.child.flag is True - assert loaded_full._private == "secret" + # types + assert isinstance(loaded_full, Dummy), store + assert isinstance(loaded_full.child, Child), store - # ——— (b) skip 'count' and skip 'child' ——— - tgt2 = tmp_path / ("dummy_skip_cd.zip" if store == "zip" else "dummy_skip_cd_dir") - d.save(str(tgt2), mode="w", store=store, skip=["count", "child"]) - loaded_skip_cd = load(str(tgt2)) + # values round-trip + np.testing.assert_allclose(loaded_full.arr, d.arr, err_msg=store) + assert loaded_full.count == d.count, store + np.testing.assert_allclose(loaded_full.child.values, d.child.values, err_msg=store) + assert loaded_full.child.flag is True, store + assert loaded_full._private == "secret", store - # arr still there - np.testing.assert_allclose(loaded_skip_cd.arr, d.arr) + # ——— (b) skip 'count' and skip 'child' ——— + tgt2 = tmp_path / f"dummy_skip_cd{suffix}" + d.save(str(tgt2), mode="w", store=store, skip=["count", "child"]) + loaded_skip_cd = load(str(tgt2)) - # count & child should be gone - assert not hasattr(loaded_skip_cd, "count") - assert not hasattr(loaded_skip_cd, "child") + # arr still there + np.testing.assert_allclose(loaded_skip_cd.arr, d.arr, err_msg=store) - # _private still exists - assert loaded_skip_cd._private == "secret" + # count & child should be gone + assert not hasattr(loaded_skip_cd, "count"), store + assert not hasattr(loaded_skip_cd, "child"), store - # ——— (c) skip "_private" only ——— - tgt3 = tmp_path / ("dummy_skip_priv.zip" if store == "zip" else "dummy_skip_priv_dir") - d.save(str(tgt3), mode="w", store=store, skip=["_private"]) - loaded_skip_priv = load(str(tgt3)) + # _private still exists + assert loaded_skip_cd._private == "secret", store - np.testing.assert_allclose(loaded_skip_priv.arr, d.arr) - assert loaded_skip_priv.count == d.count - assert isinstance(loaded_skip_priv.child, Child) - assert not hasattr(loaded_skip_priv, "_private") + # ——— (c) skip "_private" only ——— + tgt3 = tmp_path / f"dummy_skip_priv{suffix}" + d.save(str(tgt3), mode="w", store=store, skip=["_private"]) + loaded_skip_priv = load(str(tgt3)) + + np.testing.assert_allclose(loaded_skip_priv.arr, d.arr, err_msg=store) + assert loaded_skip_priv.count == d.count, store + assert isinstance(loaded_skip_priv.child, Child), store + assert not hasattr(loaded_skip_priv, "_private"), store # ──────────────────────────── # @@ -110,8 +112,7 @@ def method(self): return self.data.sum() -@pytest.mark.parametrize("store", ["zip", "dir"]) -def test_nested_save_load_roundtrip(tmp_path, store): +def test_nested_save_load_and_skip(tmp_path): """ 1. Save/load TestingClass with no skips: - `data` (ndarray) @@ -119,87 +120,46 @@ def test_nested_save_load_roundtrip(tmp_path, store): - nested list + nested tuple (with ndarrays) - top‐level child.info - nested subClass inside list_with_child and tuple_with_child - 2. Verify each nested structure (including the nested subClass objects) round‐trips. + 2. Save with skip=["info","child"]: `info` and top‐level `child` are removed entirely, + but the *nested* subClass instances inside list_with_child/tuple_with_child remain. """ obj = TestingClass() - tgt = tmp_path / ("nested.zip" if store == "zip" else "nested_dir") - obj.save(str(tgt), mode="w", store=store) - loaded = load(str(tgt)) - - # — Top‐level class — - assert isinstance(loaded, TestingClass) - - # — data array — - np.testing.assert_allclose(loaded.data, obj.data) - - # — top‐level dicts and child — - assert loaded.info == obj.info - assert isinstance(loaded.child, subClass) - assert loaded.child.info == obj.child.info - - # — list_a: first two entries + ndarray — - assert loaded.list_a[:2] == obj.list_a[:2] - np.testing.assert_allclose(loaded.list_a[2], obj.list_a[2]) - - # — tuple_a: first two entries + ndarray — - assert loaded.tuple_a[:2] == obj.tuple_a[:2] - np.testing.assert_allclose(loaded.tuple_a[2], obj.tuple_a[2]) - - # — nested list & nested tuple of primitives — - assert loaded.list_nested == obj.list_nested - assert loaded.tuple_nested == obj.tuple_nested - - # — list_with_child: check nested subClass — - assert loaded.list_with_child[0] == "start" - assert isinstance(loaded.list_with_child[1], subClass) - assert loaded.list_with_child[1].info == obj.list_with_child[1].info - - # — tuple_with_child: check nested subClass — - assert loaded.tuple_with_child[0] == "start" - assert isinstance(loaded.tuple_with_child[1], subClass) - assert loaded.tuple_with_child[1].info == obj.tuple_with_child[1].info - - -@pytest.mark.parametrize("store", ["zip", "dir"]) -def test_nested_skip_fields(tmp_path, store): - """ - 1. Save with skip=["info","child"], load back. - 2. Only `data`, `list_a`, `list_nested`, `tuple_a`, `tuple_nested`, - `list_with_child` and `tuple_with_child` survive. - 3. `info` and top‐level `child` are removed entirely, but the *nested* - subClass instances inside list_with_child/tuple_with_child should remain. - """ - obj = TestingClass() - tgt = tmp_path / ("skip_nested.zip" if store == "zip" else "skip_nested_dir") - obj.save(str(tgt), mode="w", store=store, skip=["info", "child"]) - loaded = load(str(tgt)) - - # — data must still be present — - np.testing.assert_allclose(loaded.data, obj.data) - - # — top‐level info & child must be gone — - assert not hasattr(loaded, "info") - assert not hasattr(loaded, "child") - - # — list_a and its ndarray still rehydrate — - assert loaded.list_a[:2] == obj.list_a[:2] - np.testing.assert_allclose(loaded.list_a[2], obj.list_a[2]) - - # — tuple_a and its ndarray still rehydrate — - assert loaded.tuple_a[:2] == obj.tuple_a[:2] - np.testing.assert_allclose(loaded.tuple_a[2], obj.tuple_a[2]) - - # — nested list & nested tuple still rehydrate — - assert loaded.list_nested == obj.list_nested - assert loaded.tuple_nested == obj.tuple_nested - - # — list_with_child: nested subClass should survive — - assert loaded.list_with_child[0] == "start" - assert isinstance(loaded.list_with_child[1], subClass) - - # — tuple_with_child: nested subClass should survive — - assert loaded.tuple_with_child[0] == "start" - assert isinstance(loaded.tuple_with_child[1], subClass) + for skip in ([], ["info", "child"]): + tgt = tmp_path / f"nested_{len(skip)}.zip" + obj.save(str(tgt), mode="w", skip=skip) + loaded = load(str(tgt)) + + # — Top‐level class — + assert isinstance(loaded, TestingClass), skip + + # — data array — + np.testing.assert_allclose(loaded.data, obj.data, err_msg=str(skip)) + + # — top‐level dicts and child: kept, or gone when skipped — + if skip: + assert not hasattr(loaded, "info") + assert not hasattr(loaded, "child") + else: + assert loaded.info == obj.info + assert isinstance(loaded.child, subClass) + assert loaded.child.info == obj.child.info + + # — list_a / tuple_a: first two entries + ndarray — + assert loaded.list_a[:2] == obj.list_a[:2], skip + np.testing.assert_allclose(loaded.list_a[2], obj.list_a[2], err_msg=str(skip)) + assert loaded.tuple_a[:2] == obj.tuple_a[:2], skip + np.testing.assert_allclose(loaded.tuple_a[2], obj.tuple_a[2], err_msg=str(skip)) + + # — nested list & nested tuple of primitives — + assert loaded.list_nested == obj.list_nested, skip + assert loaded.tuple_nested == obj.tuple_nested, skip + + # — list_with_child / tuple_with_child: nested subClass survives — + for name in ("list_with_child", "tuple_with_child"): + assert getattr(loaded, name)[0] == "start", (name, skip) + assert isinstance(getattr(loaded, name)[1], subClass), (name, skip) + if not skip: + assert getattr(loaded, name)[1].info == getattr(obj, name)[1].info, name # ──────────────────────────── # @@ -207,23 +167,6 @@ def test_nested_skip_fields(tmp_path, store): # ──────────────────────────── # -class TorchOnly(AutoSerialize): - def __init__(self): - self.data = torch.rand(100, 100) - - def method(self): - return self.data.sum() - - -@pytest.mark.parametrize("store", ["zip", "dir"]) -def test_torch_tensor_roundtrip(tmp_path, store): - obj = TorchOnly() - tgt = tmp_path / ("tensor.zip" if store == "zip" else "tensor_dir") - obj.save(str(tgt), mode="w", store=store) - loaded = load(str(tgt)) - assert torch.allclose(obj.data, loaded.data) - - class TorchWithOpt(AutoSerialize): def __init__(self): self.data = torch.rand(100, 100) @@ -234,16 +177,15 @@ def method(self): return self.data.sum() -@pytest.mark.parametrize("store", ["zip", "dir"]) -def test_torch_tensor_and_optimizer(tmp_path, store): +def test_torch_tensor_and_optimizer(tmp_path): obj = TorchWithOpt() - tgt = tmp_path / ("opt.zip" if store == "zip" else "opt_dir") - obj.save(str(tgt), mode="w", store=store) + tgt = tmp_path / "opt.zip" + obj.save(str(tgt), mode="w") loaded = load(str(tgt)) assert torch.allclose(obj.data, loaded.data) assert isinstance(loaded.opt, torch.optim.Adam) - assert set(obj.opt.state_dict().keys()) == set(loaded.opt.state_dict().keys()) + assert loaded.opt.param_groups[0]["lr"] == 0.01 class TestModule(torch.nn.Module, AutoSerialize): @@ -258,11 +200,10 @@ def forward(self, x): return self.linear(x) -@pytest.mark.parametrize("store", ["zip", "dir"]) -def test_torch_nn_module_roundtrip(tmp_path, store): +def test_torch_nn_module_roundtrip(tmp_path): mod = TestModule() - tgt = tmp_path / ("mod.zip" if store == "zip" else "mod_dir") - mod.save(str(tgt), mode="w", store=store) + tgt = tmp_path / "mod.zip" + mod.save(str(tgt), mode="w") mod_reload = load(str(tgt)) assert isinstance(mod_reload, TestModule) @@ -291,11 +232,10 @@ def p1(self, value): self._p1.copy_(value) -@pytest.mark.parametrize("store", ["zip", "dir"]) -def test_hybrid_module_roundtrip(tmp_path, store): +def test_hybrid_module_roundtrip(tmp_path): mod = TestHybrid() - tgt = tmp_path / ("hybrid.zip" if store == "zip" else "hybrid_dir") - mod.save(str(tgt), mode="w", store=store) + tgt = tmp_path / "hybrid.zip" + mod.save(str(tgt), mode="w") mod_reload = load(str(tgt)) for key in mod.state_dict(): @@ -309,11 +249,10 @@ def __init__(self): self.empty = set() -@pytest.mark.parametrize("store", ["zip", "dir"]) -def test_set_roundtrip(tmp_path, store): +def test_set_roundtrip(tmp_path): """Sets load back as sets, whether their items are strings or numbers.""" - target = tmp_path / ("sets.zip" if store == "zip" else "sets_dir") - WithSets().save(str(target), mode="w", store=store) + target = tmp_path / "sets.zip" + WithSets().save(str(target), mode="w") loaded = load(str(target)) assert loaded.names == {"C10", "C12", "rotation_angle"} diff --git a/tests/datastructures/test_dataset.py b/tests/datastructures/test_dataset.py index 201fe92be..5598ba80a 100644 --- a/tests/datastructures/test_dataset.py +++ b/tests/datastructures/test_dataset.py @@ -14,12 +14,6 @@ def sample_2d_array(): return np.random.rand(10, 10) -@pytest.fixture -def sample_3d_array(): - """Create a sample 3D array for testing.""" - return np.random.rand(5, 10, 15) - - @pytest.fixture def sample_dataset_2d(sample_2d_array): """Create a sample 2D dataset for testing.""" @@ -33,17 +27,8 @@ def sample_dataset_2d(sample_2d_array): ) -@pytest.fixture -def sample_dataset_3d(sample_3d_array): - """Create a sample 3D dataset for testing.""" - return Dataset.from_array( - array=sample_3d_array, - name="test_3d_dataset", - origin=(0, 0, 0), - sampling=(1, 1, 1), - units=["nm", "nm", "nm"], - signal_units="counts", - ) +def center_world(d): + return d.origin + (np.array(d.shape) - 1) / 2.0 * d.sampling class TestDatasetInitialization: @@ -93,74 +78,17 @@ def test_direct_initialization_error(self, sample_2d_array): class TestDatasetProperties: - """Test Dataset properties and setters.""" - - def test_array_property(self, sample_dataset_2d, sample_2d_array): - """Test array property getter and setter.""" - # Test getter - assert np.array_equal(sample_dataset_2d.array, sample_2d_array) - - # Test setter - new_array = np.ones((10, 10)) - sample_dataset_2d.array = new_array - assert np.array_equal(sample_dataset_2d.array, new_array) - - def test_name_property(self, sample_dataset_2d): - """Test name property getter and setter.""" - # Test getter - assert sample_dataset_2d.name == "test_2d_dataset" - - # Test setter - sample_dataset_2d.name = "new_name" - assert sample_dataset_2d.name == "new_name" - - def test_origin_property(self, sample_dataset_2d): - """Test origin property getter and setter.""" - # Test getter - assert np.array_equal(sample_dataset_2d.origin, np.array([0, 0])) - - # Test setter - sample_dataset_2d.origin = (1, 1) - assert np.array_equal(sample_dataset_2d.origin, np.array([1, 1])) - - # Test validation - with pytest.raises(ValueError): - sample_dataset_2d.origin = (1, 1, 1) # Wrong number of dimensions - - def test_sampling_property(self, sample_dataset_2d): - """Test sampling property getter and setter.""" - # Test getter - assert np.array_equal(sample_dataset_2d.sampling, np.array([1, 1])) - - # Test setter - sample_dataset_2d.sampling = (2, 2) - assert np.array_equal(sample_dataset_2d.sampling, np.array([2, 2])) - - # Test validation - with pytest.raises(ValueError): - sample_dataset_2d.sampling = (2, 2, 2) # Wrong number of dimensions - - def test_units_property(self, sample_dataset_2d): - """Test units property getter and setter.""" - # Test getter - assert sample_dataset_2d.units == ["nm", "nm"] - - # Test setter - sample_dataset_2d.units = ["Å", "Å"] - assert sample_dataset_2d.units == ["Å", "Å"] - - # Test validation - with pytest.raises(ValueError): - sample_dataset_2d.units = ["Å", "Å", "Å"] # Wrong number of dimensions - - def test_signal_units_property(self, sample_dataset_2d): - """Test signal_units property getter and setter.""" - # Test getter - assert sample_dataset_2d.signal_units == "counts" - - # Test setter - sample_dataset_2d.signal_units = "electrons" - assert sample_dataset_2d.signal_units == "electrons" + def test_axis_metadata_setters_validate_length(self, sample_dataset_2d): + """origin, sampling and units accept one value per axis and reject any other length.""" + for name, valid, expected, wrong in [ + ("origin", (1, 1), np.array([1, 1]), (1, 1, 1)), + ("sampling", (2, 2), np.array([2, 2]), (2, 2, 2)), + ("units", ["Å", "Å"], ["Å", "Å"], ["Å", "Å", "Å"]), + ]: + setattr(sample_dataset_2d, name, valid) + assert np.array_equal(getattr(sample_dataset_2d, name), expected), name + with pytest.raises(ValueError): + setattr(sample_dataset_2d, name, wrong) class TestDatasetMethods: @@ -185,40 +113,12 @@ def test_copy(self, sample_dataset_2d): copied_dataset.name = "modified_copy" assert sample_dataset_2d.name == "test_2d_dataset" - def test_mean(self, sample_dataset_2d): - """Test mean method.""" - # Mean of all elements - mean_value = sample_dataset_2d.mean() - assert isinstance(mean_value, (float, np.float64)) - - # Mean along axis 0 - mean_along_axis = sample_dataset_2d.mean(axes=(0,)) - assert mean_along_axis.shape == (10,) - - # Mean along both axes (should be the same as mean of all elements) - mean_along_both = sample_dataset_2d.mean(axes=(0, 1)) - assert isinstance(mean_along_both, (float, np.float64)) - assert np.isclose(mean_value, mean_along_both) - - def test_max(self, sample_dataset_2d): - """Test max method.""" - # Max of all elements - max_value = sample_dataset_2d.max() - assert isinstance(max_value, (float, np.float64)) - - # Max along axis 0 - max_along_axis = sample_dataset_2d.max(axes=(0,)) - assert max_along_axis.shape == (10,) - - def test_min(self, sample_dataset_2d): - """Test min method.""" - # Min of all elements - min_value = sample_dataset_2d.min() - assert isinstance(min_value, (float, np.float64)) - - # Min along axis 0 - min_along_axis = sample_dataset_2d.min(axes=(0,)) - assert min_along_axis.shape == (10,) + def test_reductions_match_numpy(self, sample_dataset_2d, sample_2d_array): + for name, reference in [("mean", np.mean), ("max", np.max), ("min", np.min)]: + reduce = getattr(sample_dataset_2d, name) + assert np.isclose(reduce(), reference(sample_2d_array)), name + assert np.isclose(reduce(axes=(0, 1)), reference(sample_2d_array)), name + assert np.allclose(reduce(axes=(0,)), reference(sample_2d_array, axis=0)), name def test_pad(self, sample_dataset_2d): """Test pad method.""" @@ -252,40 +152,29 @@ def test_crop(self, sample_dataset_2d): assert sample_dataset_2d.shape == (8, 8) assert np.array_equal(sample_dataset_2d.origin, np.array([1, 1])) - def test_crop_4dstem_kspace(self): - """Test cropping k-space axes of a 4D-STEM dataset.""" - dset = Dataset.from_array(np.random.rand(8, 8, 96, 96)) - cropped = dset.crop(crop_widths=((4, 92), (4, 92)), axes=(2, 3)) - assert cropped.shape == (8, 8, 88, 88) - assert dset.shape == (8, 8, 96, 96) - - def test_crop_4dstem_realspace_in_place(self): - """Test in-place real-space crop of a 4D-STEM dataset.""" - dset = Dataset.from_array( - np.random.rand(16, 16, 32, 32), - origin=(10, 20, 30, 40), - sampling=(0.5, 0.25, 2, 3), - ) - dset.crop(crop_widths=((4, 12), (4, 12)), axes=(0, 1), modify_in_place=True) - assert dset.shape == (8, 8, 32, 32) - assert np.array_equal(dset.origin, np.array([12, 21, 30, 40])) - - def test_crop_4dstem_stop_zero(self): - """Test that stop=0 keeps all remaining elements.""" - dset = Dataset.from_array(np.random.rand(8, 8, 96, 96)) - cropped = dset.crop(crop_widths=((10, 0), (10, 0)), axes=(2, 3)) - assert cropped.shape == (8, 8, 86, 86) - - def test_crop_single_axis_updates_origin_in_place(self): - """Test in-place single-axis crop updates origin.""" - dset = Dataset.from_array( - np.random.rand(2048, 64), - origin=(5, 7), - sampling=(0.5, 2), - ) - dset.crop(crop_widths=((80, 2000),), axes=0, modify_in_place=True) - assert dset.shape == (1920, 64) - assert np.array_equal(dset.origin, np.array([45, 7])) + def test_crop_selected_axes(self): + """Cropping a subset of axes, including stop=0 (keep the rest) and a single int axis.""" + cases = [ + # shape, origin, sampling, crop_widths, axes, expected shape, expected origin + ((8, 8, 96, 96), (0, 0, 0, 0), (1, 1, 1, 1), ((4, 92), (4, 92)), (2, 3), + (8, 8, 88, 88), (0, 0, 4, 4)), + ((8, 8, 96, 96), (0, 0, 0, 0), (1, 1, 1, 1), ((10, 0), (10, 0)), (2, 3), + (8, 8, 86, 86), (0, 0, 10, 10)), + ((16, 16, 32, 32), (10, 20, 30, 40), (0.5, 0.25, 2, 3), ((4, 12), (4, 12)), (0, 1), + (8, 8, 32, 32), (12, 21, 30, 40)), + ((2048, 64), (5, 7), (0.5, 2), ((80, 2000),), 0, + (1920, 64), (45, 7)), + ] # fmt: skip + for shape, origin, sampling, crop_widths, axes, out_shape, out_origin in cases: + dset = Dataset.from_array(np.random.rand(*shape), origin=origin, sampling=sampling) + cropped = dset.crop(crop_widths=crop_widths, axes=axes) + assert cropped.shape == out_shape, (shape, crop_widths, axes) + assert np.allclose(cropped.origin, out_origin), (shape, crop_widths, axes) + assert dset.shape == shape, (shape, crop_widths, axes) + + dset.crop(crop_widths=crop_widths, axes=axes, modify_in_place=True) + assert dset.shape == out_shape, (shape, crop_widths, axes) + assert np.allclose(dset.origin, out_origin), (shape, crop_widths, axes) def test_getitem_slice_updates_origin_like_crop(self): """Test slicing keeps origin aligned with crop semantics.""" @@ -322,94 +211,39 @@ def test_getitem_slice_reduced_rank_updates_origin_and_sampling(self): def test_bin(self, sample_dataset_2d): """Test bin method.""" - # Bin by factor of 2 - binned_dataset = sample_dataset_2d.bin(bin_factors=2) - # Check shape - assert binned_dataset.shape == (5, 5) # Original (10, 10) / 2 - # Check that the original dataset is unchanged - assert sample_dataset_2d.shape == (10, 10) - # Test modify_in_place - sample_dataset_2d.bin(bin_factors=2, modify_in_place=True) - assert sample_dataset_2d.shape == (5, 5) - - -class TestDatasetRepresentation: - """Test Dataset string representation.""" - - def test_repr(self, sample_dataset_2d): - """Test __repr__ method.""" - repr_str = repr(sample_dataset_2d) - assert "Dataset" in repr_str - assert "shape=(10, 10)" in repr_str - assert "name='test_2d_dataset'" in repr_str - - def test_str(self, sample_dataset_2d): - """Test __str__ method.""" - str_str = str(sample_dataset_2d) - assert "quantem Dataset" in str_str - assert "shape: (10, 10)" in str_str - assert "name: 'test_2d_dataset'" in str_str or "named 'test_2d_dataset'" in str_str - - -class TestBinMetadata: - def test_bin_updates_sampling_and_origin(self, sample_dataset_2d): # Original metadata assert np.array_equal(sample_dataset_2d.origin, np.array([0.0, 0.0])) assert np.array_equal(sample_dataset_2d.sampling, np.array([1.0, 1.0])) - # Bin by 2 on both axes + # Bin by factor of 2 binned = sample_dataset_2d.bin(bin_factors=2) - - # Shape reduced - assert binned.shape == (5, 5) - + assert binned.shape == (5, 5) # Original (10, 10) / 2 # Sampling doubled assert np.allclose(binned.sampling, np.array([2.0, 2.0])) - # Origin shifted to center of first 2x2 block: +0.5 * (2-1) * old_sampling assert np.allclose(binned.origin, np.array([0.5, 0.5])) + # Check that the original dataset is unchanged + assert sample_dataset_2d.shape == (10, 10) + # Test modify_in_place + sample_dataset_2d.bin(bin_factors=2, modify_in_place=True) + assert sample_dataset_2d.shape == (5, 5) -class TestFourierResample: - def test_downsample_mean_and_sampling(self): - # Constant array → easy mean check; 10x10 → factors 0.5,0.5 → 5x5 - arr = np.ones((10, 10), dtype=float) - ds = Dataset.from_array(arr, origin=(0, 0), sampling=(1.0, 1.0), units=["px", "px"]) - - ds2 = ds.fourier_resample(factors=0.5) # apply to all axes - assert ds2.shape == (5, 5) - - # Mean preserved ~ 1 - assert np.isclose(ds2.array.mean(), 1.0, rtol=1e-6, atol=1e-6) - - # Sampling doubled - assert np.allclose(ds2.sampling, np.array([2.0, 2.0])) - - # Center preserved in world coords - def center_world(d): - return d.origin + (np.array(d.shape) - 1) / 2.0 * d.sampling - - assert np.allclose(center_world(ds), center_world(ds2), rtol=1e-12, atol=1e-12) - - def test_upsample_mean_and_sampling(self): - # 5x6 → factors 2 → 10x12 - arr = np.ones((5, 6), dtype=float) - ds = Dataset.from_array(arr, origin=(1.5, -2.0), sampling=(0.2, 0.5), units=["nm", "nm"]) - - ds2 = ds.fourier_resample(factors=2.0) - assert ds2.shape == (10, 12) - - # Mean preserved ~ 1 - assert np.isclose(ds2.array.mean(), 1.0, rtol=1e-6, atol=1e-6) - - # Sampling halved - assert np.allclose(ds2.sampling, np.array([0.1, 0.25])) - - # Center preserved in world coords (with nonzero origin/sampling) - def center_world(d): - return d.origin + (np.array(d.shape) - 1) / 2.0 * d.sampling - assert np.allclose(center_world(ds), center_world(ds2), rtol=1e-10, atol=1e-10) +class TestFourierResample: + def test_resample_preserves_mean_and_center(self): + cases = [ + # shape, origin, sampling, kwargs, expected shape, expected sampling, center tol + ((10, 10), (0, 0), (1.0, 1.0), {"factors": 0.5}, (5, 5), (2.0, 2.0), 1e-12), + ((5, 6), (1.5, -2.0), (0.2, 0.5), {"factors": 2.0}, (10, 12), (0.1, 0.25), 1e-10), + ] + for shape, origin, sampling, kwargs, out_shape, out_sampling, tol in cases: + ds = Dataset.from_array(np.ones(shape), origin=origin, sampling=sampling) + ds2 = ds.fourier_resample(**kwargs) + assert ds2.shape == out_shape, kwargs + assert np.isclose(ds2.array.mean(), 1.0, rtol=1e-6, atol=1e-6), kwargs + assert np.allclose(ds2.sampling, np.array(out_sampling)), kwargs + assert np.allclose(center_world(ds), center_world(ds2), rtol=tol, atol=tol), kwargs def test_odd_even_center_preservation(self): # Odd→even sizes; check center preservation rather than sampling integers @@ -419,11 +253,6 @@ def test_odd_even_center_preservation(self): # Target out_shape (5, 4) explicitly ds2 = ds.fourier_resample(out_shape=(5, 4)) assert ds2.shape == (5, 4) - - # Center preserved - def center_world(d): - return d.origin + (np.array(d.shape) - 1) / 2.0 * d.sampling - assert np.allclose(center_world(ds), center_world(ds2), rtol=1e-12, atol=1e-12) def test_api_errors(self, sample_dataset_2d): @@ -444,11 +273,11 @@ def test_numpy_copy_is_readonly(self): so writes raise instead of silently updating only the detached copy. """ import torch + from quantem.core.datastructures.dataset4dstem import Dataset4dstem + ds = Dataset4dstem.from_tensor(torch.zeros(2, 2, 2, 2)) arr = ds.numpy() assert arr.flags.writeable is False with pytest.raises(ValueError, match="read-only"): arr[0, 0, 0, 0] = 99.0 - - diff --git a/tests/datastructures/test_dataset3d_show.py b/tests/datastructures/test_dataset3d_show.py index 7938c66f7..cf9b627a8 100644 --- a/tests/datastructures/test_dataset3d_show.py +++ b/tests/datastructures/test_dataset3d_show.py @@ -14,9 +14,9 @@ def dataset_with_10_frames(): @pytest.fixture -def dataset_with_100_frames(): - """Create a Dataset3d with 100 frames for testing default max behavior.""" - return Dataset3d.from_array(np.random.rand(100, 8, 8)) +def dataset_with_25_frames(): + """Create a Dataset3d with more frames than the default max of 20.""" + return Dataset3d.from_array(np.random.rand(25, 8, 8)) def extract_frame_indices_from_figure(fig): @@ -28,68 +28,63 @@ def extract_frame_indices_from_figure(fig): ] +def assert_shows_frames(dataset, cases): + for kwargs, expected_indices in cases: + fig, _ = dataset.show(returnfig=True, **kwargs) + assert extract_frame_indices_from_figure(fig) == expected_indices, kwargs + plt.close(fig) + + class TestShowInputValidation: """Test that invalid inputs raise clear errors.""" - @pytest.mark.parametrize( - "kwargs,match", - [ + def test_raises_value_error(self, dataset_with_10_frames): + for kwargs, match in [ ({"step": 0}, "cannot be zero"), ({"start": 100}, "out of bounds"), ({"start": -100}, "out of bounds"), ({"start": 5, "end": 5}, "No frames to display"), ({"ncols": 0}, "ncols must be >= 1"), - ({"ncols": -1}, "ncols must be >= 1"), ({"max": 0}, "max must be >= 1"), - ({"max": -1}, "max must be >= 1"), - ], - ) - def test_raises_value_error(self, dataset_with_10_frames, kwargs, match): - with pytest.raises(ValueError, match=match): - dataset_with_10_frames.show(**kwargs) + ]: + with pytest.raises(ValueError, match=match): + dataset_with_10_frames.show(**kwargs) class TestShowFrameSelection: """Test frame selection with start, end, step, max combinations.""" - @pytest.mark.parametrize( - "kwargs,expected_indices", - [ - ({}, list(range(20))), - ({"max": 5}, [0, 1, 2, 3, 4]), - ({"max": None}, list(range(100))), - ({"start": 90}, list(range(90, 100))), - ({"start": 95, "max": 3}, [95, 96, 97]), - ], - ) - def test_large_dataset(self, dataset_with_100_frames, kwargs, expected_indices): - fig, _ = dataset_with_100_frames.show(returnfig=True, **kwargs) - assert extract_frame_indices_from_figure(fig) == expected_indices - plt.close(fig) + def test_large_dataset(self, dataset_with_25_frames): + assert_shows_frames( + dataset_with_25_frames, + [ + ({}, list(range(20))), + ({"max": 5}, [0, 1, 2, 3, 4]), + ({"max": None}, list(range(25))), + ({"start": 20}, list(range(20, 25))), + ({"start": 20, "max": 3}, [20, 21, 22]), + ], + ) - @pytest.mark.parametrize( - "kwargs,expected_indices", - [ - # Default shows all frames (< max) - ({}, list(range(10))), - # Start and end - ({"start": 5}, [5, 6, 7, 8, 9]), - ({"end": 5}, [0, 1, 2, 3, 4]), - # Step - ({"step": 2}, [0, 2, 4, 6, 8]), - ({"step": 3}, [0, 3, 6, 9]), - ({"start": 2, "end": 8, "step": 2}, [2, 4, 6]), - # Negative start index - ({"start": -1, "max": 1}, [9]), - ({"start": -3, "max": 2}, [7, 8]), - # Negative step (reverse order) - ({"start": 9, "step": -1}, [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]), - ({"start": 9, "end": 4, "step": -1}, [9, 8, 7, 6, 5]), - ({"start": 9, "step": -2}, [9, 7, 5, 3, 1]), - ({"start": 9, "step": -1, "max": 3}, [9, 8, 7]), - ], - ) - def test_small_dataset(self, dataset_with_10_frames, kwargs, expected_indices): - fig, _ = dataset_with_10_frames.show(returnfig=True, **kwargs) - assert extract_frame_indices_from_figure(fig) == expected_indices - plt.close(fig) + def test_small_dataset(self, dataset_with_10_frames): + assert_shows_frames( + dataset_with_10_frames, + [ + # Default shows all frames (< max) + ({}, list(range(10))), + # Start and end + ({"start": 5}, [5, 6, 7, 8, 9]), + ({"end": 5}, [0, 1, 2, 3, 4]), + # Step + ({"step": 3}, [0, 3, 6, 9]), + ({"start": 2, "end": 8, "step": 2}, [2, 4, 6]), + # Negative start index + ({"start": -1, "max": 1}, [9]), + ({"start": -3, "max": 2}, [7, 8]), + # Negative step (reverse order) + ({"start": 9, "step": -1}, [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]), + ({"start": 9, "end": 4, "step": -1}, [9, 8, 7, 6, 5]), + ({"start": 9, "step": -2}, [9, 7, 5, 3, 1]), + ({"start": 9, "step": -1, "max": 3}, [9, 8, 7]), + ], + ) diff --git a/tests/datastructures/test_dataset4dstem.py b/tests/datastructures/test_dataset4dstem.py index 7e53d5769..a55d41b33 100644 --- a/tests/datastructures/test_dataset4dstem.py +++ b/tests/datastructures/test_dataset4dstem.py @@ -68,113 +68,42 @@ def test_from_array_defaults(self, sample_4d_array): class TestDataset4dstemGetItem: """Test Dataset4dstem __getitem__ method.""" - def test_getitem_2d(self, sample_dataset4dstem): - """Test __getitem__ method to get a 2D slice (single diffraction pattern).""" - # Get a 2D slice (single diffraction pattern) + def test_getitem(self, sample_dataset4dstem): + """A single diffraction pattern and a line of diffraction patterns.""" dp = sample_dataset4dstem[0, 0] assert dp.ndim == 2 assert dp.shape == (10, 10) - - # Check that the data is correct assert np.array_equal(dp.array, sample_dataset4dstem.array[0, 0]) - def test_getitem_3d(self, sample_dataset4dstem): - """Test __getitem__ method to get a 3D slice (line of diffraction patterns).""" - # Get a 3D slice (line of diffraction patterns) line = sample_dataset4dstem[0, :] assert line.ndim == 3 assert line.shape == (5, 10, 10) - - # Check that the data is correct assert np.array_equal(line.array, sample_dataset4dstem.array[0, :]) class TestDataset4dstemDpMethods: """Test Dataset4dstem diffraction pattern methods.""" - def test_dp_mean_property(self, sample_dataset4dstem): - """Test dp_mean property.""" - # Initially, dp_mean should not be attached - assert not hasattr(sample_dataset4dstem, "_dp_mean") - - # Accessing dp_mean should trigger calculation but not attach - dp_mean = sample_dataset4dstem.dp_mean - assert not hasattr(sample_dataset4dstem, "_dp_mean") - assert dp_mean.ndim == 2 - assert dp_mean.shape == (10, 10) - - # Accessing again should recalculate - dp_mean2 = sample_dataset4dstem.dp_mean - assert dp_mean2 is not dp_mean # Should be a new calculation each time - - def test_get_dp_mean(self, sample_dataset4dstem): - """Test get_dp_mean method.""" - # Get dp_mean without attaching - dp_mean = sample_dataset4dstem.get_dp_mean(attach=False) - assert not hasattr(sample_dataset4dstem, "_dp_mean") - assert dp_mean.ndim == 2 - assert dp_mean.shape == (10, 10) - - # Get dp_mean with attaching - dp_mean2 = sample_dataset4dstem.get_dp_mean(attach=True) - assert hasattr(sample_dataset4dstem, "_dp_mean") - assert dp_mean2 is sample_dataset4dstem._dp_mean - - def test_dp_max_property(self, sample_dataset4dstem): - """Test dp_max property.""" - # Initially, dp_max should not be attached - assert not hasattr(sample_dataset4dstem, "_dp_max") - - # Accessing dp_max should trigger calculation but not attach - dp_max = sample_dataset4dstem.dp_max - assert not hasattr(sample_dataset4dstem, "_dp_max") - assert dp_max.ndim == 2 - assert dp_max.shape == (10, 10) - - # Accessing again should recalculate - dp_max2 = sample_dataset4dstem.dp_max - assert dp_max2 is not dp_max # Should be a new calculation each time - - def test_get_dp_max(self, sample_dataset4dstem): - """Test get_dp_max method.""" - # Get dp_max without attaching - dp_max = sample_dataset4dstem.get_dp_max(attach=False) - assert not hasattr(sample_dataset4dstem, "_dp_max") - assert dp_max.ndim == 2 - assert dp_max.shape == (10, 10) - - # Get dp_max with attaching - dp_max2 = sample_dataset4dstem.get_dp_max(attach=True) - assert hasattr(sample_dataset4dstem, "_dp_max") - assert dp_max2 is sample_dataset4dstem._dp_max - - def test_dp_median_property(self, sample_dataset4dstem): - """Test dp_median property.""" - # Initially, dp_median should not be attached - assert not hasattr(sample_dataset4dstem, "_dp_median") - - # Accessing dp_median should trigger calculation but not attach - dp_median = sample_dataset4dstem.dp_median - assert not hasattr(sample_dataset4dstem, "_dp_median") - assert dp_median.ndim == 2 - assert dp_median.shape == (10, 10) - - # Accessing again should recalculate - dp_median2 = sample_dataset4dstem.dp_median - assert dp_median2 is not dp_median # Should be a new calculation each time - - def test_get_dp_median(self, sample_dataset4dstem): - """Test get_dp_median method.""" - # Get dp_median without attaching - dp_median = sample_dataset4dstem.get_dp_median(attach=False) - assert not hasattr(sample_dataset4dstem, "_dp_median") - assert dp_median.ndim == 2 - assert dp_median.shape == (10, 10) - - # Get dp_median with attaching - dp_median2 = sample_dataset4dstem.get_dp_median(attach=True) - assert hasattr(sample_dataset4dstem, "_dp_median") - assert dp_median2 is sample_dataset4dstem._dp_median + def test_dp_reductions_and_attach(self, sample_dataset4dstem): + """dp_* computes without attaching; get_dp_*(attach=True) caches what dp_* then returns.""" + array = sample_dataset4dstem.array + for name, reference in [("mean", np.mean), ("max", np.max), ("median", np.median)]: + attr = f"_dp_{name}" + expected = reference(array, axis=(0, 1)) + + dp = getattr(sample_dataset4dstem, f"dp_{name}") + assert not hasattr(sample_dataset4dstem, attr), name + assert np.allclose(dp.array, expected), name + assert np.allclose(dp.sampling, [0.1, 0.1]), name + assert getattr(sample_dataset4dstem, f"dp_{name}") is not dp, name + + get = getattr(sample_dataset4dstem, f"get_dp_{name}") + assert np.allclose(get(attach=False).array, expected), name + assert not hasattr(sample_dataset4dstem, attr), name + + attached = get(attach=True) + assert getattr(sample_dataset4dstem, attr) is attached, name + assert getattr(sample_dataset4dstem, f"dp_{name}") is attached, name def test_get_virtual_image(self, sample_dataset4dstem): """Test get_virtual_image method.""" @@ -186,8 +115,10 @@ def test_get_virtual_image(self, sample_dataset4dstem): virtual_image = sample_dataset4dstem.get_virtual_image( mask=mask, name="test_virtual_image", attach=False ) - assert virtual_image.ndim == 2 assert virtual_image.shape == (5, 5) + assert np.allclose( + virtual_image.array, sample_dataset4dstem.array[:, :, 4:6, 4:6].sum(axis=(-2, -1)) + ) assert "test_virtual_image" not in sample_dataset4dstem.virtual_images # Get virtual image with attaching @@ -198,27 +129,6 @@ def test_get_virtual_image(self, sample_dataset4dstem): assert virtual_image2 is sample_dataset4dstem.virtual_images["test_virtual_image"] -class TestDataset4dstemShow: - """Test Dataset4dstem show method.""" - - def test_show(self, sample_dataset4dstem, monkeypatch): - """Test show method.""" - - # Mock the show_2d function to avoid actual plotting - def mock_show_2d(*args, **kwargs): - return None - - # Apply the mock - monkeypatch.setattr("quantem.core.visualization.visualization.show_2d", mock_show_2d) - - # Call show method - result = sample_dataset4dstem.show() - - # Check that the method returns a tuple of (fig, axs) - assert isinstance(result, tuple) - assert len(result) == 2 - - class TestProbePositions: """`probe_positions` hands the implicit raster grid to consumers that need it explicitly.""" From c17501127af5e9fb04b99c860d10ddd89508c775 Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 14:49:15 -0700 Subject: [PATCH 02/19] Consolidate visualization tests and assert real values test_custom_normalizations.py (37 -> 9 IDs): - Intervals/stretches clip to [0, 1], so the "0 <= result <= 1" checks on the log/inverse-log/asinh/sinh stretches and on every CustomNormalization variant could not fail (the 1..10 input clipped to all ones). Replaced by: one loop asserting f(0)=0, f(1)=1, monotonicity, inverse round trip, the documented log/inverse-log formulas and asinh/sinh symmetry; and one CustomNormalization table checking exact values for the default, manual, centered, power, logarithmic and asinh variants. - Interval limit tests become one table; ManualInterval.__call__ now checks every value. Linear/power exact checks kept in one test. - Invalid-parameter tests become one loop. - NormalizationConfig default/custom tests deleted (dataclass fields). - _resolve_normalization: the preset test only checked isinstance; now one test checks explicit configs for None, the vmin/vmax and quantile kwargs, dict, two presets and pass-through; invalid preset/type merged. - Removed the unused outlier fixture. test_visualization_utils.py (31 -> 12 IDs): - The three "custom origin/sampling/statistic" histogram tests asserted ~isnan(hist[hist > 0]), which cannot fail, and the others only the output shape. One table now places known points and checks bin values for sum, mean, origin and sampling; the column-vector test checks equality with the 1D call and the mismatch error. - array_to_rgba: cmap path compared against matplotlib; angle path checks black at zero amplitude, lightness order, hue change and the chroma boost. combine_arrays_to_rgba checks black/no-signal, distinct hues and boost. - ScalebarConfig default/custom tests deleted; _resolve_scalebar's six tests become one, adding the sampling-kwarg path. - estimate_scalebar_length checks the actual lengths and the non-uniform sampling error instead of isinstance/positivity. - Colorbar tests check the norm and the arg ticks; the chroma-boost variant only asserted "is not None" and now checks the colormap changes. - turbo_black: fade_len was never reflected in any assertion; one loop now checks N, the black start and the unfaded tail. test_visualization.py (20 -> 10 IDs): - isinstance(fig, Figure) could not fail with figax passed in. Smoke cases for real/complex/cmap/chroma (and the combined variants) become loops asserting one image each; scalebar/colorbar/title checks kept in one call. - New value check: manual norm + viridis is what imshow receives. - Grid normalization tests merged; pyplot figures closed. - Deleted three commented-out tests. Co-Authored-By: Claude Opus 5.5 --- .../test_custom_normalizations.py | 314 +++++------------ tests/visualization/test_visualization.py | 129 +++---- .../visualization/test_visualization_utils.py | 331 ++++++------------ 3 files changed, 233 insertions(+), 541 deletions(-) diff --git a/tests/visualization/test_custom_normalizations.py b/tests/visualization/test_custom_normalizations.py index 0201f9978..9b9ee9be9 100644 --- a/tests/visualization/test_custom_normalizations.py +++ b/tests/visualization/test_custom_normalizations.py @@ -2,7 +2,6 @@ import pytest from quantem.core.visualization.custom_normalizations import ( - NORMALIZATION_PRESETS, CenteredInterval, CustomNormalization, HyperbolicSineStretch, @@ -23,247 +22,120 @@ def sample_data(): return np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], dtype=np.float64) -@pytest.fixture -def sample_data_with_outliers(): - return np.array( - [-10.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 100.0], - dtype=np.float64, - ) - - -class TestManualInterval: - def test_manual_interval_default(self, sample_data): - interval = ManualInterval() - vmin, vmax = interval.get_limits(sample_data) - assert vmin == 1.0 - assert vmax == 10.0 - - def test_manual_interval_with_limits(self, sample_data): - interval = ManualInterval(vmin=2.0, vmax=8.0) - vmin, vmax = interval.get_limits(sample_data) - assert vmin == 2.0 - assert vmax == 8.0 +class TestIntervals: + def test_interval_limits(self, sample_data): + for interval, expected in [ + (ManualInterval(), (1.0, 10.0)), + (ManualInterval(vmin=2.0, vmax=8.0), (2.0, 8.0)), + # With default center=0, the range should be [-10, 10] for sample_data [1, 10] + (CenteredInterval(), (-10.0, 10.0)), + (CenteredInterval(vcenter=5.0), (0.0, 10.0)), + (CenteredInterval(vcenter=5.0, half_range=3.0), (2.0, 8.0)), + (QuantileInterval(), (1.18, 9.82)), # 2nd and 98th percentiles + (QuantileInterval(lower_quantile=0.1, upper_quantile=0.9), (1.9, 9.1)), + ]: + assert np.allclose(interval.get_limits(sample_data), expected), interval def test_manual_interval_call(self, sample_data): interval = ManualInterval(vmin=2.0, vmax=8.0) result = interval(sample_data) - assert np.all(result >= 0.0) - assert np.all(result <= 1.0) + assert np.allclose(result, np.clip((sample_data - 2.0) / 6.0, 0.0, 1.0)) assert result[0] == 0.0 # value below vmin assert result[-1] == 1.0 # value above vmax -class TestCenteredInterval: - def test_centered_interval_default(self, sample_data): - interval = CenteredInterval() - vmin, vmax = interval.get_limits(sample_data) - # With default center=0, the range should be [-10, 10] for sample_data [1, 10] - assert vmin == -10.0 - assert vmax == 10.0 - - def test_centered_interval_with_center(self, sample_data): - interval = CenteredInterval(vcenter=5.0) - vmin, vmax = interval.get_limits(sample_data) - # With center=5, the range should be [0, 10] for sample_data [1, 10] - assert vmin == 0.0 - assert vmax == 10.0 - - def test_centered_interval_with_half_range(self, sample_data): - interval = CenteredInterval(vcenter=5.0, half_range=3.0) - vmin, vmax = interval.get_limits(sample_data) - assert vmin == 2.0 - assert vmax == 8.0 - - -class TestQuantileInterval: - def test_quantile_interval_default(self, sample_data): - interval = QuantileInterval() - vmin, vmax = interval.get_limits(sample_data) - assert np.isclose(vmin, 1.18) # 2nd percentile - assert np.isclose(vmax, 9.82) # 98th percentile - - def test_quantile_interval_custom(self, sample_data): - interval = QuantileInterval(lower_quantile=0.1, upper_quantile=0.9) - vmin, vmax = interval.get_limits(sample_data) - assert np.isclose(vmin, 1.9) # 10th percentile - assert np.isclose(vmax, 9.1) # 90th percentile - - -class TestLinearStretch: - def test_linear_stretch_default(self, sample_data): - stretch = LinearStretch() - result = stretch(sample_data) - assert np.allclose(result, sample_data) - - def test_linear_stretch_with_slope(self, sample_data): - stretch = LinearStretch(slope=2) +class TestStretches: + def test_linear_and_power_stretch(self, sample_data): normalized_data = sample_data / 10.0 # Normalize to [0, 1] - result = stretch(normalized_data) - assert np.allclose(result, 2 * normalized_data) - - def test_linear_stretch_with_intercept(self, sample_data): - stretch = LinearStretch(intercept=1) - normalized_data = sample_data / 10.0 # Normalize to [0, 1] - result = stretch(normalized_data) - assert np.allclose(result, normalized_data + 1) - - -class TestPowerLawStretch: - def test_power_law_stretch_default(self, sample_data): - stretch = PowerLawStretch() - result = stretch(sample_data) - assert np.allclose(result, sample_data) - - def test_power_law_stretch_square(self, sample_data): - stretch = PowerLawStretch(power=2) - normalized_data = sample_data / 10.0 # Normalize to [0, 1] - result = stretch(normalized_data) - assert np.allclose(result, normalized_data**2) - - def test_power_law_stretch_invalid_power(self): - with pytest.raises(ValueError): - PowerLawStretch(power=0) - - -class TestLogarithmicStretch: - def test_logarithmic_stretch_default(self, sample_data): - stretch = LogarithmicStretch() - result = stretch(sample_data) - assert np.all(result >= 0) - assert np.all(result <= 1) - - def test_logarithmic_stretch_invalid_a(self): - with pytest.raises(ValueError): - LogarithmicStretch(a=0) - - -class TestInverseLogarithmicStretch: - def test_inverse_logarithmic_stretch_default(self, sample_data): - stretch = InverseLogarithmicStretch() - result = stretch(sample_data) - assert np.all(result >= 0) - assert np.all(result <= 1) - - def test_inverse_logarithmic_stretch_invalid_a(self): - with pytest.raises(ValueError): - InverseLogarithmicStretch(a=0) - - -class TestInverseHyperbolicSineStretch: - def test_inverse_hyperbolic_sine_stretch_default(self, sample_data): - stretch = InverseHyperbolicSineStretch() - result = stretch(sample_data) - assert np.all(result >= 0) - assert np.all(result <= 1) - - def test_inverse_hyperbolic_sine_stretch_invalid_a(self): - with pytest.raises(ValueError): - InverseHyperbolicSineStretch(a=0) - - -class TestHyperbolicSineStretch: - def test_hyperbolic_sine_stretch_default(self, sample_data): - stretch = HyperbolicSineStretch() - result = stretch(sample_data) - assert np.all(result >= 0) - assert np.all(result <= 1) - - def test_hyperbolic_sine_stretch_invalid_a(self): - with pytest.raises(ValueError): - HyperbolicSineStretch(a=0) + # the identity stretches return their input untouched, even outside [0, 1] + assert np.allclose(LinearStretch()(sample_data), sample_data) + assert np.allclose(PowerLawStretch()(sample_data), sample_data) + assert np.allclose(LinearStretch(slope=2)(normalized_data), 2 * normalized_data) + assert np.allclose(LinearStretch(intercept=1)(normalized_data), normalized_data + 1) + assert np.allclose(PowerLawStretch(power=2)(normalized_data), normalized_data**2) + + def test_nonlinear_stretches_map_unit_interval_onto_itself(self): + x = np.linspace(0.0, 1.0, 11) + a = 1000.0 + expected = { + "log": np.log(a * x + 1) / np.log(a + 1), + "inverse_log": (np.exp(x * np.log(a + 1)) - 1) / a, + } + for name, stretch in [ + ("log", LogarithmicStretch(a)), + ("inverse_log", InverseLogarithmicStretch(a)), + ("asinh", InverseHyperbolicSineStretch()), + ("sinh", HyperbolicSineStretch()), + ("power", PowerLawStretch(power=2)), + ]: + y = stretch(x) + assert y[0] == pytest.approx(0.0, abs=1e-12), name + assert y[-1] == pytest.approx(1.0), name + assert np.all(np.diff(y) > 0), name + assert np.allclose(stretch.inverse(y), x), name + if name in expected: + assert np.allclose(y, expected[name]), name + if name in ("asinh", "sinh"): + assert y[5] == pytest.approx(0.5), name # symmetric about the midpoint + + def test_invalid_parameters_raise(self): + for build in ( + lambda: PowerLawStretch(power=0), + lambda: LogarithmicStretch(a=0), + lambda: InverseLogarithmicStretch(a=0), + lambda: InverseHyperbolicSineStretch(a=0), + lambda: HyperbolicSineStretch(a=0), + ): + with pytest.raises(ValueError): + build() class TestCustomNormalization: - def test_custom_normalization_default(self, sample_data): - norm = CustomNormalization() - result = norm(sample_data) - assert np.all(result >= 0) - assert np.all(result <= 1) - - def test_custom_normalization_manual(self, sample_data): - norm = CustomNormalization(interval_type="manual", vmin=2, vmax=8) - result = norm(sample_data) - assert np.all(result >= 0) - assert np.all(result <= 1) - - def test_custom_normalization_centered(self, sample_data): - norm = CustomNormalization(interval_type="centered", vcenter=5) - result = norm(sample_data) - assert np.all(result >= 0) - assert np.all(result <= 1) - - def test_custom_normalization_power(self, sample_data): - norm = CustomNormalization(stretch_type="power", power=2) - result = norm(sample_data) - assert np.all(result >= 0) - assert np.all(result <= 1) - - def test_custom_normalization_logarithmic(self, sample_data): - norm = CustomNormalization(stretch_type="logarithmic") - result = norm(sample_data) - assert np.all(result >= 0) - assert np.all(result <= 1) - - def test_custom_normalization_asinh(self, sample_data): - norm = CustomNormalization(stretch_type="asinh") - result = norm(sample_data) - assert np.all(result >= 0) - assert np.all(result <= 1) - - def test_custom_normalization_invalid_interval_type(self, sample_data): + def test_interval_then_stretch(self, sample_data): + lo, hi = np.quantile(sample_data, (0.02, 0.98)) + quantile = np.clip((sample_data - lo) / (hi - lo), 0.0, 1.0) + for kwargs, expected in [ + ({}, quantile), + ( + {"interval_type": "manual", "vmin": 2, "vmax": 8}, + np.clip((sample_data - 2) / 6, 0, 1), + ), + ({"interval_type": "centered", "vcenter": 5}, sample_data / 10.0), + ({"stretch_type": "power", "power": 2}, quantile**2), + ({"stretch_type": "logarithmic"}, np.log(1000 * quantile + 1) / np.log(1001)), + ( + {"stretch_type": "asinh"}, + np.arcsinh((2 * quantile - 1) / 0.1) / (2 * np.arcsinh(10.0)) + 0.5, + ), + ]: + assert np.allclose(CustomNormalization(**kwargs)(sample_data), expected), kwargs + + def test_invalid_types_raise(self): with pytest.raises(ValueError): CustomNormalization(interval_type="invalid") - - def test_custom_normalization_invalid_stretch_type(self, sample_data): with pytest.raises(ValueError): CustomNormalization(stretch_type="invalid") -class TestNormalizationConfig: - def test_normalization_config_default(self): - config = NormalizationConfig() - assert config.interval_type == "quantile" - assert config.stretch_type == "linear" - assert config.lower_quantile == 0.02 - assert config.upper_quantile == 0.98 - - def test_normalization_config_custom(self): - config = NormalizationConfig( - interval_type="manual", - stretch_type="power", - vmin=1, - vmax=10, - power=2, - ) - assert config.interval_type == "manual" - assert config.stretch_type == "power" - assert config.vmin == 1 - assert config.vmax == 10 - assert config.power == 2 - - class TestResolveNormalization: - def test_resolve_normalization_none(self): - config = _resolve_normalization(None) - assert isinstance(config, NormalizationConfig) - assert config.interval_type == "quantile" - - def test_resolve_normalization_dict(self): + def test_resolve_normalization(self): + assert _resolve_normalization(None) == NormalizationConfig() + assert _resolve_normalization(None, vmin=1, vmax=10) == NormalizationConfig( + interval_type="manual", vmin=1, vmax=10 + ) + assert _resolve_normalization(None, upper_quantile=0.9) == NormalizationConfig( + upper_quantile=0.9 + ) config = _resolve_normalization({"interval_type": "manual", "vmin": 1, "vmax": 10}) - assert isinstance(config, NormalizationConfig) - assert config.interval_type == "manual" - assert config.vmin == 1 - assert config.vmax == 10 - - def test_resolve_normalization_preset(self): - for preset_name in NORMALIZATION_PRESETS: - config = _resolve_normalization(preset_name) - assert isinstance(config, NormalizationConfig) - - def test_resolve_normalization_invalid_preset(self): + assert (config.interval_type, config.vmin, config.vmax) == ("manual", 1, 10) + config = _resolve_normalization("log_minmax") + assert (config.interval_type, config.stretch_type) == ("manual", "logarithmic") + assert _resolve_normalization("power_sqrt").power == 0.5 + original = NormalizationConfig(power=3.0) + assert _resolve_normalization(original) is original + + def test_resolve_normalization_invalid(self): with pytest.raises(ValueError): _resolve_normalization("invalid_preset") - - def test_resolve_normalization_invalid_type(self): with pytest.raises(TypeError): _resolve_normalization(123) # type: ignore diff --git a/tests/visualization/test_visualization.py b/tests/visualization/test_visualization.py index adaa75de7..3dcb551b5 100644 --- a/tests/visualization/test_visualization.py +++ b/tests/visualization/test_visualization.py @@ -1,5 +1,7 @@ +import matplotlib.pyplot as plt import numpy as np import pytest +from matplotlib import colormaps from matplotlib.axes import Axes from matplotlib.figure import Figure @@ -40,131 +42,86 @@ def mock_fig_ax(): class TestShow2DFunction: - def test_show_2d_basic(self, sample_array, mock_fig_ax): - fig, ax = _show_2d_array(sample_array, figax=mock_fig_ax) - assert isinstance(fig, Figure) - assert len(ax.images) == 1 - - def test_show_2d_with_complex(self, sample_complex_array, mock_fig_ax): - fig, ax = _show_2d_array(sample_complex_array, figax=mock_fig_ax) - assert isinstance(fig, Figure) - assert len(ax.images) == 1 - - def test_show_2d_with_norm(self, sample_array, mock_fig_ax): + def test_show_2d_renders_one_image(self, sample_array, sample_complex_array): + for array, kwargs in [ + (sample_array, {}), + (sample_array, {"cmap": "viridis"}), + (sample_complex_array, {}), + (sample_complex_array, {"chroma_boost": 2.0}), + ]: + fig = Figure() + ax = fig.add_subplot(111) + returned_fig, returned_ax = _show_2d_array(array, figax=(fig, ax), **kwargs) + assert returned_fig is fig and returned_ax is ax, kwargs + assert len(ax.images) == 1, (array.dtype, kwargs) + + def test_show_2d_maps_values_through_norm_and_cmap(self, sample_array, mock_fig_ax): norm = NormalizationConfig(interval_type="manual", vmin=1, vmax=9) - fig, ax = _show_2d_array(sample_array, norm=norm, figax=mock_fig_ax) - assert isinstance(fig, Figure) - assert len(ax.images) == 1 + _, ax = _show_2d_array(sample_array, norm=norm, cmap="viridis", figax=mock_fig_ax) + expected = colormaps["viridis"]((sample_array - 1) / 8) + assert np.allclose(ax.images[0].get_array(), expected) - def test_show_2d_with_scalebar(self, sample_array, mock_fig_ax): + def test_show_2d_decorations(self, sample_array, mock_fig_ax): scalebar = ScalebarConfig(sampling=1.0, units="nm") - fig, ax = _show_2d_array(sample_array, scalebar=scalebar, figax=mock_fig_ax) - assert isinstance(fig, Figure) - assert len(ax.images) == 1 + fig, ax = _show_2d_array( + sample_array, scalebar=scalebar, cbar=True, title="Test Title", figax=mock_fig_ax + ) assert len(ax.artists) > 0 # scalebar should be added - - def test_show_2d_with_cbar(self, sample_array, mock_fig_ax): - fig, ax = _show_2d_array(sample_array, cbar=True, figax=mock_fig_ax) - assert isinstance(fig, Figure) - assert len(ax.images) == 1 assert len(fig.axes) > 1 # colorbar should be added - - def test_show_2d_with_custom_cmap(self, sample_array, mock_fig_ax): - fig, ax = _show_2d_array(sample_array, cmap="viridis", figax=mock_fig_ax) - assert isinstance(fig, Figure) - assert len(ax.images) == 1 - - def test_show_2d_with_chroma_boost(self, sample_complex_array, mock_fig_ax): - fig, ax = _show_2d_array(sample_complex_array, chroma_boost=2.0, figax=mock_fig_ax) - assert isinstance(fig, Figure) - assert len(ax.images) == 1 - - def test_show_2d_with_title(self, sample_array, mock_fig_ax): - fig, ax = _show_2d_array(sample_array, title="Test Title", figax=mock_fig_ax) - assert isinstance(fig, Figure) - assert len(ax.images) == 1 assert ax.get_title() == "Test Title" class TestShow2DCombined: - def test_show_2d_combined_basic(self, sample_arrays, mock_fig_ax): - fig, ax = _show_2d_combined(sample_arrays, figax=mock_fig_ax) - assert isinstance(fig, Figure) - assert len(ax.images) == 1 - - def test_show_2d_combined_with_norm(self, sample_arrays, mock_fig_ax): - # Convert arrays to float type before normalization - sample_arrays = [arr.astype(float) for arr in sample_arrays] - norm = NormalizationConfig(interval_type="manual", vmin=1, vmax=12) - fig, ax = _show_2d_combined(sample_arrays, norm=norm, figax=mock_fig_ax) - assert isinstance(fig, Figure) - assert isinstance(ax, Axes) + def test_show_2d_combined_renders_one_image(self, sample_arrays): + float_arrays = [arr.astype(float) for arr in sample_arrays] + for arrays, kwargs in [ + (sample_arrays, {}), + (float_arrays, {"norm": NormalizationConfig(interval_type="manual", vmin=1, vmax=12)}), + (sample_arrays, {"cmap": "viridis"}), + (sample_arrays, {"chroma_boost": 2.0}), + ]: + fig = Figure() + ax = fig.add_subplot(111) + _show_2d_combined(arrays, figax=(fig, ax), **kwargs) + assert len(ax.images) == 1, kwargs def test_show_2d_combined_with_scalebar(self, sample_arrays, mock_fig_ax): scalebar = ScalebarConfig(sampling=1.0, units="nm") fig, ax = _show_2d_combined(sample_arrays, scalebar=scalebar, figax=mock_fig_ax) - assert isinstance(fig, Figure) assert len(ax.images) == 1 assert len(ax.artists) > 0 # scalebar should be added - def test_show_2d_combined_with_custom_cmap(self, sample_arrays, mock_fig_ax): - fig, ax = _show_2d_combined(sample_arrays, cmap="viridis", figax=mock_fig_ax) - assert isinstance(fig, Figure) - assert len(ax.images) == 1 - - def test_show_2d_combined_with_chroma_boost(self, sample_arrays, mock_fig_ax): - fig, ax = _show_2d_combined(sample_arrays, chroma_boost=2.0, figax=mock_fig_ax) - assert isinstance(fig, Figure) - assert len(ax.images) == 1 - class TestNormalizeShowInputToGrid: - def test_normalize_show_input_to_grid_single_array(self, sample_array): + def test_normalize_show_input_to_grid(self, sample_array, sample_arrays): grid = _normalize_show_input_to_grid(sample_array) assert len(grid) == 1 assert len(grid[0]) == 1 assert np.array_equal(grid[0][0], sample_array) - def test_normalize_show_input_to_grid_list(self, sample_arrays): grid = _normalize_show_input_to_grid(sample_arrays) assert len(grid) == 1 assert len(grid[0]) == len(sample_arrays) for arr1, arr2 in zip(grid[0], sample_arrays): assert np.array_equal(arr1, arr2) - def test_normalize_show_input_to_grid_nested_list(self, sample_arrays): - nested = [sample_arrays[:2], sample_arrays[1:]] - grid = _normalize_show_input_to_grid(nested) - assert len(grid) == 2 - assert len(grid[0]) == 2 - assert len(grid[1]) == 2 + grid = _normalize_show_input_to_grid([sample_arrays[:2], sample_arrays[1:]]) + assert [len(row) for row in grid] == [2, 2] + assert np.array_equal(grid[1][1], sample_arrays[2]) class TestShow2DMain: - # def test_show_2d_single_array(self, sample_array): - # fig, axs = show_2d(sample_array) - # assert isinstance(fig, Figure) - # assert len(axs.shape) == 2 - # assert axs.shape == (1, 1) - - # def test_show_2d_list(self, sample_arrays): - # fig, axs = show_2d(sample_arrays) - # assert isinstance(fig, Figure) - # assert len(axs.shape) == 2 - # assert axs.shape == (1, len(sample_arrays)) - def test_show_2d_nested_list(self, sample_arrays): nested = [sample_arrays[:2], sample_arrays[1:]] fig, axs = show_2d(nested) assert isinstance(fig, Figure) - assert len(axs.shape) == 2 assert axs.shape == (2, 2) + plt.close(fig) def test_show_2d_with_combine_images(self, sample_arrays, mock_fig_ax): # Convert arrays to float type and ensure proper shape handling sample_arrays = [arr.astype(float) for arr in sample_arrays] fig, ax = show_2d(sample_arrays, combine_images=True, figax=mock_fig_ax) - assert isinstance(fig, Figure) assert isinstance(ax, Axes) assert len(ax.images) == 1 # Combined images should result in a single image @@ -173,12 +130,6 @@ def test_show_2d_with_invalid_combine_images(self, sample_arrays): with pytest.raises(ValueError): show_2d(nested, combine_images=True) - # def test_show_2d_with_figax(self, sample_array, mock_fig_ax): - # fig, axs = show_2d(sample_array, figax=mock_fig_ax) - # assert isinstance(fig, Figure) - # assert len(axs.shape) == 2 - # assert axs.shape == (1, 1) - def test_show_2d_with_invalid_figax_shape(self, sample_arrays, mock_fig_ax): with pytest.raises(ValueError): show_2d(sample_arrays, figax=mock_fig_ax) diff --git a/tests/visualization/test_visualization_utils.py b/tests/visualization/test_visualization_utils.py index 1617f0a4b..b7c3356ee 100644 --- a/tests/visualization/test_visualization_utils.py +++ b/tests/visualization/test_visualization_utils.py @@ -1,11 +1,12 @@ +import matplotlib as mpl import matplotlib.pyplot as plt import numpy as np import pytest -from matplotlib import colors -from matplotlib.axes import Axes from matplotlib.colorbar import Colorbar from matplotlib.colors import LinearSegmentedColormap, Normalize +from quantem.core import config +from quantem.core.visualization.custom_normalizations import CustomNormalization from quantem.core.visualization.visualization_utils import ( ScalebarConfig, _resolve_scalebar, @@ -20,39 +21,6 @@ ) -@pytest.fixture -def sample_array(): - return np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], dtype=np.float64) - - -@pytest.fixture -def sample_angle_array(): - return np.array( - [ - [0.0, np.pi / 4, np.pi / 2], - [-np.pi / 2, 0.0, np.pi / 4], - [-np.pi, -3 * np.pi / 4, -np.pi / 2], - ], - dtype=np.float64, - ) - - -@pytest.fixture -def sample_complex_array(): - return np.array( - [[1 + 1j, 2 + 2j, 3 + 3j], [4 + 4j, 5 + 5j, 6 + 6j], [7 + 7j, 8 + 8j, 9 + 9j]], - dtype=np.complex128, - ) - - -@pytest.fixture -def sample_arrays(): - return [ - np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float64), - np.array([[5.0, 6.0], [7.0, 8.0]], dtype=np.float64), - ] - - @pytest.fixture def mock_fig_ax(): fig, ax = plt.subplots() @@ -61,113 +29,73 @@ def mock_fig_ax(): class TestArrayToRGBA: - def test_array_to_rgba_without_angle(self, sample_array): - rgba = array_to_rgba(sample_array) - assert rgba.shape == (*sample_array.shape, 4) + def test_array_to_rgba_without_angle_applies_cmap(self): + amplitude = np.linspace(0.0, 1.0, 9).reshape(3, 3) + for cmap in ("gray", "viridis"): + rgba = array_to_rgba(amplitude, cmap=cmap) + assert np.allclose(rgba, mpl.colormaps[cmap](amplitude)), cmap + + def test_array_to_rgba_with_angle(self): + amplitude = np.array([[0.0, 0.5, 1.0], [0.5, 0.5, 0.5]]) + angle = np.array([[0.0, 0.0, 0.0], [0.0, np.pi / 2, -np.pi / 2]]) + rgba = array_to_rgba(amplitude, angle) + boosted = array_to_rgba(amplitude, angle, chroma_boost=2.0) + + assert rgba.shape == (2, 3, 4) assert np.all(rgba[..., 3] == 1) # alpha channel should be 1 - - def test_array_to_rgba_with_angle(self, sample_array, sample_angle_array): - rgba = array_to_rgba(sample_array, sample_angle_array) - assert rgba.shape == (*sample_array.shape, 4) - assert np.all(rgba[..., 3] == 1) # alpha channel should be 1 - - def test_array_to_rgba_with_custom_cmap(self, sample_array): - rgba = array_to_rgba(sample_array, cmap="viridis") - assert rgba.shape == (*sample_array.shape, 4) - - def test_array_to_rgba_with_chroma_boost(self, sample_array, sample_angle_array): - rgba = array_to_rgba(sample_array, sample_angle_array, chroma_boost=2.0) - assert rgba.shape == (*sample_array.shape, 4) - - def test_array_to_rgba_shape_mismatch(self, sample_array): - wrong_shape = np.array([[1, 2], [3, 4]]) + assert np.allclose(rgba[0, 0, :3], 0.0) # zero amplitude is black + # lightness grows with amplitude at a fixed angle + assert rgba[0, 0, :3].sum() < rgba[0, 1, :3].sum() < rgba[0, 2, :3].sum() + # equal amplitude, different angles -> different hues + assert not np.allclose(rgba[1, 0, :3], rgba[1, 1, :3]) + assert not np.allclose(rgba[1, 1, :3], rgba[1, 2, :3]) + # a chroma boost spreads the channels further apart + assert np.all(np.ptp(boosted[1, :, :3], axis=-1) > np.ptp(rgba[1, :, :3], axis=-1)) + + def test_array_to_rgba_shape_mismatch(self): with pytest.raises(ValueError): - array_to_rgba(sample_array, wrong_shape) + array_to_rgba(np.ones((3, 3)), np.ones((2, 2))) class TestCombineArraysToRGBA: - def test_combine_arrays_to_rgba(self, sample_arrays): - rgba = combine_arrays_to_rgba(sample_arrays) - assert rgba.shape == (*sample_arrays[0].shape, 4) + def test_combine_arrays_to_rgba(self): + norm = CustomNormalization(interval_type="manual", vmin=0, vmax=1) + first = np.array([[1.0, 0.0], [0.0, 0.0]]) + second = np.array([[0.0, 1.0], [0.0, 0.0]]) + rgba = combine_arrays_to_rgba([first, second], norm=norm) + boosted = combine_arrays_to_rgba([first, second], norm=norm, chroma_boost=2.0) + + assert rgba.shape == (2, 2, 4) assert np.all(rgba[..., 3] == 1) # alpha channel should be 1 - - def test_combine_arrays_to_rgba_with_chroma_boost(self, sample_arrays): - rgba = combine_arrays_to_rgba(sample_arrays, chroma_boost=2.0) - assert rgba.shape == (*sample_arrays[0].shape, 4) - - -class TestScalebarConfig: - def test_scalebar_config_default(self): - config = ScalebarConfig() - assert config.sampling == 1.0 - assert config.units == "pixels" - assert config.length is None - assert config.width_px == 1 - assert config.pad_px == 0.5 - assert config.color == "white" - assert config.loc == "lower right" - - def test_scalebar_config_custom(self): - config = ScalebarConfig( - sampling=2.0, - units="nm", - length=10.0, - width_px=2, - pad_px=1.0, - color="black", - loc="upper left", + assert np.allclose(rgba[1, :, :3], 0.0) # no signal is black + assert not np.allclose(rgba[0, 0, :3], rgba[0, 1, :3]) # each array gets its own hue + assert not np.allclose(boosted[0, :, :3], rgba[0, :, :3]) + + +class TestScalebar: + def test_resolve_scalebar(self): + assert _resolve_scalebar(None) is None + assert _resolve_scalebar(False) is None + assert _resolve_scalebar(True) == ScalebarConfig() + assert _resolve_scalebar({"sampling": 2.0, "units": "nm"}) == ScalebarConfig( + sampling=2.0, units="nm" + ) + assert _resolve_scalebar(None, sampling=2.0) == ScalebarConfig( + sampling=2.0, units=config.get("viz.real_space_units") ) - assert config.sampling == 2.0 - assert config.units == "nm" - assert config.length == 10.0 - assert config.width_px == 2 - assert config.pad_px == 1.0 - assert config.color == "black" - assert config.loc == "upper left" - - -class TestResolveScalebar: - def test_resolve_scalebar_none(self): - config = _resolve_scalebar(None) - assert config is None - - def test_resolve_scalebar_false(self): - config = _resolve_scalebar(False) - assert config is None - - def test_resolve_scalebar_true(self): - config = _resolve_scalebar(True) - assert isinstance(config, ScalebarConfig) - - def test_resolve_scalebar_dict(self): - config = _resolve_scalebar({"sampling": 2.0, "units": "nm"}) - assert isinstance(config, ScalebarConfig) - assert config.sampling == 2.0 - assert config.units == "nm" - - def test_resolve_scalebar_config(self): original = ScalebarConfig(sampling=2.0, units="nm") - config = _resolve_scalebar(original) - assert config is original - - def test_resolve_scalebar_invalid_type(self): + assert _resolve_scalebar(original) is original with pytest.raises(TypeError): _resolve_scalebar(123) # type: ignore - -class TestEstimateScalebarLength: def test_estimate_scalebar_length(self): - length, length_px = estimate_scalebar_length(100, 1.0) - assert isinstance(length, float) - assert isinstance(length_px, float) - - def test_estimate_scalebar_length_with_sampling(self): - length, length_px = estimate_scalebar_length(100, 2.0) - assert length > 0 - assert length_px > 0 - + # a "nice" length near a quarter of the field of view, and its size in pixels + assert estimate_scalebar_length(100, 1.0) == (20.0, 20.0) + assert estimate_scalebar_length(100, 2.0) == (50.0, 25.0) + assert estimate_scalebar_length(100, (2.0, 2.0)) == (50.0, 25.0) + with pytest.raises(ValueError): + estimate_scalebar_length(100, (1.0, 2.0)) -class TestAddScalebarToAx: def test_add_scalebar_to_ax(self, mock_fig_ax): fig, ax = mock_fig_ax add_scalebar_to_ax( @@ -185,123 +113,64 @@ def test_add_scalebar_to_ax(self, mock_fig_ax): assert len(ax.artists) > 0 -class TestAddCbarToAx: +class TestColorbars: def test_add_cbar_to_ax(self, mock_fig_ax): fig, ax = mock_fig_ax - data = np.random.rand(10, 10) - norm = Normalize(vmin=data.min(), vmax=data.max()) + norm = Normalize(vmin=0.0, vmax=1.0) cmap = LinearSegmentedColormap.from_list("test", ["blue", "red"]) - cax = fig.add_axes([0.85, 0.15, 0.05, 0.7]) # [left, bottom, width, height] + cax = fig.add_axes((0.85, 0.15, 0.05, 0.7)) # [left, bottom, width, height] cbar = add_cbar_to_ax(fig, cax=cax, norm=norm, cmap=cmap) assert isinstance(cbar, Colorbar) + assert cbar.norm is norm - -class TestAddArgCbarToAx: def test_add_arg_cbar_to_ax(self, mock_fig_ax): fig, ax = mock_fig_ax - # Create sample complex data - data = np.random.rand(10, 10) + 1j * np.random.rand(10, 10) - cmap = LinearSegmentedColormap.from_list("test", ["blue", "red"]) - norm = Normalize(vmin=0, vmax=2 * np.pi) - ax.imshow(np.angle(data), cmap=cmap, norm=norm) - cax = fig.add_axes([0.85, 0.15, 0.05, 0.7]) # [left, bottom, width, height] - cbar = add_arg_cbar_to_ax(fig, cax=cax) + cbar = add_arg_cbar_to_ax(fig, cax=fig.add_axes((0.85, 0.15, 0.05, 0.7))) + boosted = add_arg_cbar_to_ax(fig, ax, chroma_boost=2.0) assert isinstance(cbar, Colorbar) - - def test_add_arg_cbar_to_ax_with_chroma_boost(self, mock_fig_ax): - fig, ax = mock_fig_ax - cbar = add_arg_cbar_to_ax(fig, ax, chroma_boost=2.0) - assert cbar is not None - assert hasattr(cbar, "ax") - assert isinstance(cbar.ax, Axes) + assert np.allclose(cbar.get_ticks(), [-np.pi, -np.pi / 2, 0, np.pi / 2, np.pi]) + assert not np.allclose(boosted.cmap.colors, cbar.cmap.colors) class TestTurboBlack: - def test_turbo_black_default(self): - cmap = turbo_black() - assert isinstance(cmap, colors.ListedColormap) - assert cmap.N == 256 - - def test_turbo_black_custom_num_colors(self): - cmap = turbo_black(num_colors=128) - assert isinstance(cmap, colors.ListedColormap) - assert cmap.N == 128 - - def test_turbo_black_custom_fade_len(self): - cmap = turbo_black(fade_len=32) - assert isinstance(cmap, colors.ListedColormap) - assert cmap.N == 256 + def test_turbo_black(self): + for num_colors, fade_len in [(256, None), (128, None), (256, 32)]: + cmap = turbo_black(num_colors=num_colors, fade_len=fade_len) + fade = num_colors // 8 if fade_len is None else fade_len + turbo = mpl.colormaps["turbo"].resampled(num_colors)(np.linspace(0, 1, num_colors)) + colors = np.asarray(cmap.colors) + assert cmap.N == num_colors, (num_colors, fade_len) + assert np.allclose(colors[0, :3], 0.0), (num_colors, fade_len) # starts at black + assert np.allclose(colors[fade:], turbo[fade:]), (num_colors, fade_len) + assert np.allclose(colors[fade - 1], turbo[fade - 1]), (num_colors, fade_len) class TestBilinearHistogram2D: - def test_basic_histogram(self): - x = np.random.rand(100) - y = np.random.rand(100) - weight = np.ones_like(x) - shape = (10, 10) - hist = bilinear_histogram_2d(x=x, y=y, weight=weight, shape=shape) - assert hist.shape == shape - - def test_with_custom_weights(self): - x = np.random.rand(100) - y = np.random.rand(100) - weight = np.random.rand(100) - shape = (10, 10) - hist = bilinear_histogram_2d(x=x, y=y, weight=weight, shape=shape) - assert hist.shape == shape - - def test_bilinear_histogram_2d(self, sample_array): - # Create a grid of points that matches the sample array size - x = np.linspace(0, 1, 3, dtype=np.float64) # 3 points to match sample_array shape - y = np.linspace(0, 1, 3, dtype=np.float64) - X, Y = np.meshgrid(x, y) - Z = sample_array.flatten() - weight = np.ones_like(Z) - - hist = bilinear_histogram_2d( - (10, 10), - X.flatten(), - Y.flatten(), - weight, - origin=(0.0, 0.0), - sampling=(1.0, 1.0), - statistic="sum", - ) - assert hist.shape == (10, 10) - assert not np.any(np.isnan(hist)) - assert np.all(hist >= 0) - - def test_bilinear_histogram_2d_with_custom_origin(self): - shape = (10, 10) - x = np.array([1, 2, 3]) - y = np.array([1, 2, 3]) - weight = np.array([1, 1, 1]) - hist = bilinear_histogram_2d(shape, x, y, weight, origin=(0.5, 0.5)) - assert hist.shape == shape - assert np.all(~np.isnan(hist[hist > 0])) # Only check non-zero values for NaN - - def test_bilinear_histogram_2d_with_custom_sampling(self): - shape = (10, 10) - x = np.array([1, 2, 3]) - y = np.array([1, 2, 3]) - weight = np.array([1, 1, 1]) - hist = bilinear_histogram_2d(shape, x, y, weight, sampling=(2.0, 2.0)) - assert hist.shape == shape - assert np.all(~np.isnan(hist[hist > 0])) # Only check non-zero values for NaN - - def test_bilinear_histogram_2d_with_custom_statistic(self): - shape = (10, 10) - x = np.array([1, 2, 3]) - y = np.array([1, 2, 3]) - weight = np.array([1, 1, 1]) - hist = bilinear_histogram_2d(shape, x, y, weight, statistic="mean") - assert hist.shape == shape - assert np.all(~np.isnan(hist[hist > 0])) # Only check non-zero values for NaN - - def test_bilinear_histogram_2d_accepts_column_vectors(self): - shape = (8, 8) + def test_bins_weights(self): + x = np.array([0.5, 0.5, 2.5]) + y = np.array([0.5, 0.5, 3.5]) + weight = np.array([1.0, 2.0, 4.0]) + for kwargs, expected in [ + ({}, {(0, 0): 3.0, (2, 3): 4.0}), + ({"statistic": "mean"}, {(0, 0): 1.5, (2, 3): 4.0}), + # the point at y = 3.5 falls outside [-1, 3) and is dropped + ({"origin": (-1.0, -1.0)}, {(1, 1): 3.0}), + ({"sampling": (2.0, 2.0)}, {(0, 0): 3.0, (1, 1): 4.0}), + ]: + hist = bilinear_histogram_2d((4, 4), x, y, weight, **kwargs) + assert hist.shape == (4, 4), kwargs + empty = 0.0 if kwargs.get("statistic", "sum") == "sum" else np.nan + want = np.full((4, 4), empty) + for index, value in expected.items(): + want[index] = value + np.testing.assert_allclose(hist, want, err_msg=str(kwargs)) + + def test_accepts_column_vectors(self): x = np.array([[1.0], [2.0], [3.0]]) y = np.array([[1.5], [2.5], [3.5]]) weight = np.array([[2.0], [3.0], [4.0]]) - hist = bilinear_histogram_2d(shape, x, y, weight) - assert hist.shape == shape + hist = bilinear_histogram_2d((8, 8), x, y, weight) + assert np.array_equal(hist, bilinear_histogram_2d((8, 8), x[:, 0], y[:, 0], weight[:, 0])) + assert hist.sum() == 9.0 + with pytest.raises(ValueError): + bilinear_histogram_2d((8, 8), x, y[:2], weight) From 428615ec62ef65a830b2882fd24b76110699700e Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 14:49:48 -0700 Subject: [PATCH 03/19] Fold imaging_utils parametrization into loops test_imaging_utils.py (37 -> 7 IDs): the 4 x 6 upsample/shift product, the 6 torch-parity shifts and the 3 center-tap factors all exercise one code path each. They are now loops inside the same tests, with every case kept and named in the assertion message. Co-Authored-By: Claude Opus 5.5 --- tests/utils/test_imaging_utils.py | 39 ++++++++++++++++--------------- 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/tests/utils/test_imaging_utils.py b/tests/utils/test_imaging_utils.py index b8688de24..cfc826473 100644 --- a/tests/utils/test_imaging_utils.py +++ b/tests/utils/test_imaging_utils.py @@ -37,27 +37,28 @@ def test_agrees_with_ifft2_at_unit_upsampling(self, image): # `dft_upsample` omits the 1/(M*N) normalization `ifft2` applies assert np.allclose(local, expected * spectrum.size, atol=1e-8) - @pytest.mark.parametrize("up", [2, 4, 8]) - def test_peak_lands_on_the_center_tap(self, image, up): + def test_peak_lands_on_the_center_tap(self, image): """Centering the fine grid on the true peak must put the maximum at its middle. The middle is `ceil(1.5 * up)`, since the grid spans `arange(-du, du+1) / up`. """ cc = np.fft.fft2(image) * np.conj(np.fft.fft2(image)) - local = dft_upsample(cc, up, (0.0, 0.0)) + for up in (2, 4, 8): + local = dft_upsample(cc, up, (0.0, 0.0)) - du = int(np.ceil(1.5 * up)) - assert local.shape == (2 * du + 1, 2 * du + 1) - assert np.unravel_index(np.argmax(local), local.shape) == (du, du) + du = int(np.ceil(1.5 * up)) + assert local.shape == (2 * du + 1, 2 * du + 1), up + assert np.unravel_index(np.argmax(local), local.shape) == (du, du), up class TestCrossCorrelationShift: - @pytest.mark.parametrize("shift", SHIFTS) - @pytest.mark.parametrize("up", [2, 4, 8, 16]) - def test_recovers_a_known_shift(self, image, shift, up): + def test_recovers_a_known_shift(self, image): """The returned shift realigns `im` onto `im_ref`, so it is the negated input.""" - measured = np.asarray(cross_correlation_shift(image, _shifted(image, shift), up)) - assert measured == pytest.approx(-np.asarray(shift), abs=0.02) + for shift in SHIFTS: + shifted = _shifted(image, shift) + for up in (2, 4, 8, 16): + measured = np.asarray(cross_correlation_shift(image, shifted, up)) + assert measured == pytest.approx(-np.asarray(shift), abs=0.02), (shift, up) def test_upsampling_beats_no_upsampling(self, image): """Sub-pixel refinement has to be an improvement on the integer peak.""" @@ -76,18 +77,18 @@ def test_identical_images_do_not_shift(self, image): measured = np.asarray(cross_correlation_shift(image, image, upsample_factor=up)) assert measured == pytest.approx((0.0, 0.0), abs=1e-6) - @pytest.mark.parametrize("shift", SHIFTS) - def test_matches_the_torch_implementation(self, image, shift): + def test_matches_the_torch_implementation(self, image): """The two paths are independent ports; they must not disagree.""" import torch - shifted = _shifted(image, shift) - numpy_shift = np.asarray(cross_correlation_shift(image, shifted, upsample_factor=8)) - torch_shift = cross_correlation_shift_torch( - torch.as_tensor(image), torch.as_tensor(shifted), upsample_factor=8 - ).numpy() + for shift in SHIFTS: + shifted = _shifted(image, shift) + numpy_shift = np.asarray(cross_correlation_shift(image, shifted, upsample_factor=8)) + torch_shift = cross_correlation_shift_torch( + torch.as_tensor(image), torch.as_tensor(shifted), upsample_factor=8 + ).numpy() - assert numpy_shift == pytest.approx(torch_shift, abs=0.02) + assert numpy_shift == pytest.approx(torch_shift, abs=0.02), shift def test_shifted_image_is_realigned(self, image): """`return_shifted_image` must undo the displacement it just measured.""" From 9278111d63fc5867f8da0bb69042deaff67b872b Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 14:50:46 -0700 Subject: [PATCH 04/19] Restore the device after test_config_set_device On Apple silicon the MPS branch left the global device set to "mps", so every later test in the same worker that reads config.get_device() ran on MPS. The test now restores the device it started with in a finally block. Also: - test_check_key_val_known_key_silent and ..._nested_key_silent merged into one loop over both keys. - test_version_exists removed; test_version_is_string implies it. Co-Authored-By: Claude Opus 5.5 --- tests/test_config.py | 76 +++++++++++++++++++++---------------------- tests/test_version.py | 4 --- 2 files changed, 37 insertions(+), 43 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 7dc1037c0..0caa4dd8d 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -7,35 +7,39 @@ def test_config_set_device(): - config.set_device("cpu") - assert config.get_device() == "cpu" - if torch.cuda.is_available(): - config.set_device("gpu") - assert config.get_device() == "cuda:0" - config.set_device("GPU") - assert config.get_device() == "cuda:0" - config.set_device(0) - assert config.get_device() == "cuda:0" - NUM_DEVICES = torch.cuda.device_count() - if NUM_DEVICES > 0: - config.set_device(NUM_DEVICES - 1) - assert config.get_device() == f"cuda:{NUM_DEVICES - 1}" - config.refresh() + start = config.get_device() + try: + config.set_device("cpu") assert config.get_device() == "cpu" - else: - with pytest.raises(RuntimeError): - config.set_device("cuda:0") - - if torch.mps.is_available(): - config.set_device("mps") - assert config.get_device() == "mps" - config.set_device("gpu") - assert config.get_device() == "mps" - config.set_device("GPU") - assert config.get_device() == "mps" - else: - with pytest.raises(RuntimeError): + if torch.cuda.is_available(): + config.set_device("gpu") + assert config.get_device() == "cuda:0" + config.set_device("GPU") + assert config.get_device() == "cuda:0" + config.set_device(0) + assert config.get_device() == "cuda:0" + NUM_DEVICES = torch.cuda.device_count() + if NUM_DEVICES > 0: + config.set_device(NUM_DEVICES - 1) + assert config.get_device() == f"cuda:{NUM_DEVICES - 1}" + config.refresh() + assert config.get_device() == "cpu" + else: + with pytest.raises(RuntimeError): + config.set_device("cuda:0") + + if torch.mps.is_available(): config.set_device("mps") + assert config.get_device() == "mps" + config.set_device("gpu") + assert config.get_device() == "mps" + config.set_device("GPU") + assert config.get_device() == "mps" + else: + with pytest.raises(RuntimeError): + config.set_device("mps") + finally: + config.set_device(start) def test_config_update_defaults(): @@ -90,19 +94,13 @@ def test_collect_env_overrides_yaml(tmp_path): assert merged == {"verbose": 9, "viz": {"cmap": "viridis"}} -def test_check_key_val_known_key_silent(): - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - config.check_key_val("verbose", 1) - assert caught == [] - - -def test_check_key_val_nested_key_silent(): +def test_check_key_val_known_and_nested_keys_silent(): """Only the top-level key is checked, a key under it is not.""" - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - config.check_key_val("viz.not_a_real_key", 1) - assert caught == [] + for key in ("verbose", "viz.not_a_real_key"): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + config.check_key_val(key, 1) + assert caught == [], key def test_check_key_val_unknown_key_warns_and_still_sets(): diff --git a/tests/test_version.py b/tests/test_version.py index b39acf0b4..9b9d3cbc5 100644 --- a/tests/test_version.py +++ b/tests/test_version.py @@ -1,9 +1,5 @@ import quantem -def test_version_exists(): - assert hasattr(quantem, "__version__") - - def test_version_is_string(): assert isinstance(quantem.__version__, str) From ff0da2b0b06cf4affc12a9e6dfedccdb6d2bea73 Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 14:51:25 -0700 Subject: [PATCH 05/19] Loop invalid-input cases in spectroscopy and bundle tests - test_eds_background.py: the four invalid-input parametrize cases become a loop in one test (6 -> 3 IDs), each case kept. - test_xeds_fitting.py: test_loss_and_tv_helpers_are_finite only checked finiteness; it now checks the exact MSE, Poisson and in-plane TV values for the same inputs, and the unknown-loss error. - test_serialize_bundle.py: the three reserved-name cases become a loop (6 -> 4 IDs). Co-Authored-By: Claude Opus 5.5 --- tests/core/test_serialize_bundle.py | 8 ++++---- tests/spectroscopy/test_eds_background.py | 21 ++++++++------------- tests/spectroscopy/test_xeds_fitting.py | 13 +++++++++---- 3 files changed, 21 insertions(+), 21 deletions(-) diff --git a/tests/core/test_serialize_bundle.py b/tests/core/test_serialize_bundle.py index 8f3af7e0f..6962e5c31 100644 --- a/tests/core/test_serialize_bundle.py +++ b/tests/core/test_serialize_bundle.py @@ -68,7 +68,7 @@ def test_bundle_round_trip(tmp_path): assert np.allclose(back.atoms.arrays["occupancy"], [0.4, 1.0, 0.6]) -@pytest.mark.parametrize("name", ["save", "print_tree", "_recursive_save"]) -def test_bundle_rejects_reserved_names(name): - with pytest.raises(ValueError, match="shadow"): - Bundle(**{name: 1}) +def test_bundle_rejects_reserved_names(): + for name in ("save", "print_tree", "_recursive_save"): + with pytest.raises(ValueError, match="shadow"): + Bundle(**{name: 1}) diff --git a/tests/spectroscopy/test_eds_background.py b/tests/spectroscopy/test_eds_background.py index bb6bf9bc4..195bc72b0 100644 --- a/tests/spectroscopy/test_eds_background.py +++ b/tests/spectroscopy/test_eds_background.py @@ -31,23 +31,18 @@ def test_calculate_background_polynomial_returns_finite_nonnegative_background( background[[0, -1]], true_background[[0, -1]], rtol=0.3, atol=0.8 ) - @pytest.mark.parametrize( - "kwargs,exc_type,match", - [ + def test_calculate_background_polynomial_validates_inputs(self, xeds_factory): + dataset = xeds_factory(np.zeros((1, 1, 64), dtype=np.float32)) + spectrum = np.linspace(1.0, 2.0, 64, dtype=np.float32) + + for kwargs, exc_type, match in [ ({"degree": True}, TypeError, "degree"), ({"degree": -1}, ValueError, "degree"), ({"percentile": 120}, ValueError, "percentile"), ({"window_size": 0}, ValueError, "window_size"), - ], - ) - def test_calculate_background_polynomial_validates_inputs( - self, xeds_factory, kwargs, exc_type, match - ): - dataset = xeds_factory(np.zeros((1, 1, 64), dtype=np.float32)) - spectrum = np.linspace(1.0, 2.0, 64, dtype=np.float32) - - with pytest.raises(exc_type, match=match): - dataset.calculate_background_polynomial(spectrum, **kwargs) + ]: + with pytest.raises(exc_type, match=match): + dataset.calculate_background_polynomial(spectrum, **kwargs) def test_calculate_background_powerlaw_is_alias(self, xeds_factory): dataset = xeds_factory(np.zeros((1, 1, 64), dtype=np.float32)) diff --git a/tests/spectroscopy/test_xeds_fitting.py b/tests/spectroscopy/test_xeds_fitting.py index 72ef1eef2..abbc6f336 100644 --- a/tests/spectroscopy/test_xeds_fitting.py +++ b/tests/spectroscopy/test_xeds_fitting.py @@ -1,4 +1,5 @@ import numpy as np +import pytest import torch from matplotlib import pyplot as plt @@ -134,14 +135,18 @@ def test_full_cube_fit_returns_maps_and_fit_images(self, xeds_factory): class TestXEDSFitHelpers: - def test_loss_and_tv_helpers_are_finite(self): + def test_loss_and_tv_helper_values(self): predicted = torch.tensor([1.0, 2.0, 3.0]) target = torch.tensor([0.0, 2.0, 4.0]) maps = torch.tensor([[[1.0, 2.0], [3.0, 5.0]]]) - assert torch.isfinite(xeds_data_loss(predicted, target, loss="mse")) - assert torch.isfinite(xeds_data_loss(predicted, target, loss="poisson")) - assert torch.isfinite(total_variation_in_plane(maps)) + assert float(xeds_data_loss(predicted, target, loss="mse")) == pytest.approx(2 / 3) + poisson = (6 - 2 * np.log(2) - 4 * np.log(3)) / 3 # mean(pred - target * log(pred)) + assert float(xeds_data_loss(predicted, target, loss="poisson")) == pytest.approx(poisson) + # mean |dy| = (2 + 3) / 2 plus mean |dx| = (1 + 2) / 2 + assert float(total_variation_in_plane(maps)) == pytest.approx(4.0) + with pytest.raises(ValueError, match="loss must be"): + xeds_data_loss(predicted, target, loss="l1") def test_plot_loss_history_shifts_negative_values_for_log_plot(self): fig, ax = plt.subplots(1, 1) From 2d9f316d40c604d4dacdb24f28afc9d6e59ac6f6 Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 14:54:49 -0700 Subject: [PATCH 06/19] Consolidate tomography orchestrator, radon and logger tests test_logger_tomography.py (5 -> 3 IDs): SummaryWriter writes the events file in __init__, so "an events file exists" could not fail, and test_log_iter_without_val_loss asserted nothing. The scalar tests are one test that reads the events back and checks every tag, value and step (including that loss/val is skipped when val_loss is None); the image test checks the logged image tags. test_tomography_opt.py (16 -> 11 IDs): - test_setter_getter_roundtrip and test_current_lrs folded into test_set_optimizers_builds_object_and_pose (same state, all assertions). - test_scheduler_setter_getter folded into test_set_schedulers_builds. - The step/zero_grad once-per-call regressions become one loop. - _get_default_lr checked only isinstance(float); it now checks the object and pose values, which also pins the routing, and the unknown-key error. test_tomography.py (15 -> 9 IDs): - The conventional factory and empty-loss-accessor tests merged. - The three constraints-setter tests and the two wrong-type setter tests merged. - test_to_updates_device deleted (moves cpu to cpu). - TestInrFactory::test_from_models_builds deleted (isinstance only; every INR test here and in test_tomography_opt.py builds through from_models). - Both plot_losses tests asserted nothing; they now check the plotted losses (and learning rates) and close the figure. test_radon.py (19 -> 9 IDs): - Shape tests merged into one radon and one iradon test. - The named-filter parametrize only checked shape/finiteness; one loop now checks each window lies between 0 and the ramp and attenuates Nyquist, and that the ramp peaks there. Odd-size and unknown-name errors merged. - The explicit- and default-theta round trips are one loop. Co-Authored-By: Claude Opus 5.5 --- tests/tomography/test_logger_tomography.py | 57 ++++++----- tests/tomography/test_radon.py | 77 ++++++--------- tests/tomography/test_tomography.py | 52 ++++------ tests/tomography/test_tomography_opt.py | 110 ++++++--------------- 4 files changed, 104 insertions(+), 192 deletions(-) diff --git a/tests/tomography/test_logger_tomography.py b/tests/tomography/test_logger_tomography.py index 88d6bb265..9ece5f053 100644 --- a/tests/tomography/test_logger_tomography.py +++ b/tests/tomography/test_logger_tomography.py @@ -3,15 +3,16 @@ ``LoggerTomography`` is a thin tensorboard wrapper that the reconstruction loop only drives when a ``log_dir`` is passed, so the end-to-end recon tests never exercise it. These CPU, always-on tests construct a logger against a ``tmp_path`` and drive each method with small -stubs that expose only the attributes the logger reads, asserting the calls run and write -event files. Matplotlib backend is ``Agg`` (set in the root conftest), so figure logging is +stubs that expose only the attributes the logger reads, then read the event files back. Matplotlib backend is ``Agg`` (set in the root conftest), so figure logging is headless. """ from types import SimpleNamespace import numpy as np +import pytest import torch +from tensorboard.backend.event_processing.event_accumulator import EventAccumulator from quantem.tomography.logger_tomography import LoggerTomography @@ -34,23 +35,19 @@ def test_init_creates_log_dir(tmp_path): logger.close() -def test_log_epoch_writes_events(tmp_path): - logger = _make_logger(tmp_path) - try: - logger.log_epoch(epoch=0, loss=1.0, tilt_series_loss=0.8, soft_loss=0.2) - logger.flush() - events = list(logger.log_dir.glob("events.out.tfevents.*")) - assert events, "log_epoch should have written a tensorboard event file" - finally: - logger.close() +def _read_events(logger) -> EventAccumulator: + logger.flush() + events = EventAccumulator(str(logger.log_dir)) + events.Reload() + return events -def test_log_iter_unpacks_learning_rates(tmp_path): +def test_log_epoch_and_iter_write_scalars(tmp_path): logger = _make_logger(tmp_path) - obj_model = SimpleNamespace(_soft_constraint_losses=[0.3]) try: + logger.log_epoch(epoch=0, loss=1.0, tilt_series_loss=0.8, soft_loss=0.2) logger.log_iter( - object_model=obj_model, + object_model=SimpleNamespace(_soft_constraint_losses=[0.3]), iter=2, consistency_loss=0.5, total_loss=0.7, @@ -58,26 +55,27 @@ def test_log_iter_unpacks_learning_rates(tmp_path): num_samples_per_ray=16, val_loss=0.4, ) - logger.flush() - assert list(logger.log_dir.glob("events.out.tfevents.*")) - finally: - logger.close() - - -def test_log_iter_without_val_loss(tmp_path): - logger = _make_logger(tmp_path) - obj_model = SimpleNamespace(_soft_constraint_losses=[0.1]) - try: # val_loss defaults to None -> the val branch must be skipped without error. logger.log_iter( - object_model=obj_model, - iter=0, + object_model=SimpleNamespace(_soft_constraint_losses=[0.1]), + iter=3, consistency_loss=0.5, total_loss=0.6, learning_rates={}, num_samples_per_ray=8, ) - logger.flush() + events = _read_events(logger) + + def scalars(tag): + return {e.step: e.value for e in events.Scalars(tag)} + + assert scalars("loss/total") == pytest.approx({0: 1.0, 2: 0.7, 3: 0.6}) + assert scalars("loss/tilt_series") == pytest.approx({0: 0.8}) + assert scalars("loss/soft") == pytest.approx({0: 0.2, 2: 0.3, 3: 0.1}) + assert scalars("learning_rate/object") == pytest.approx({2: 1e-3}) + assert scalars("learning_rate/pose") == pytest.approx({2: 1e-2}) + assert scalars("num_samples_per_ray") == pytest.approx({2: 16, 3: 8}) + assert scalars("loss/val") == pytest.approx({2: 0.4}) finally: logger.close() @@ -97,7 +95,8 @@ def test_log_iter_images(tmp_path): dataset_model=dataset_model, iter=1, ) - logger.flush() - assert list(logger.log_dir.glob("events.out.tfevents.*")) + images = set(_read_events(logger).Tags()["images"]) + expected = {f"volume/sum_{axis}_{c}" for axis in "zyx" for c in range(2)} + assert images == expected | {"z1_z3_angles", "shifts"} finally: logger.close() diff --git a/tests/tomography/test_radon.py b/tests/tomography/test_radon.py index 41516f039..a863471ab 100644 --- a/tests/tomography/test_radon.py +++ b/tests/tomography/test_radon.py @@ -20,43 +20,24 @@ def _disk(n: int, cy: int, cx: int, r: int) -> torch.Tensor: class TestRadonShapes: - def test_2d_input_returns_angles_by_pixels(self): - img = _disk(64, 32, 32, 10) + def test_radon_shapes(self): theta = torch.linspace(0, 180, 30) - sino = radon_torch(img, theta=theta) - assert sino.shape == (30, 64) - - def test_batched_input_returns_batch_angles_pixels(self): + assert radon_torch(_disk(64, 32, 32, 10), theta=theta).shape == (30, 64) imgs = torch.stack([_disk(48, 24, 20, 8), _disk(48, 24, 28, 8)]) - theta = torch.linspace(0, 180, 20) - sino = radon_torch(imgs, theta=theta) - assert sino.shape == (2, 20, 48) - - def test_default_theta_is_180_angles(self): - sino = radon_torch(_disk(32, 16, 16, 6)) - assert sino.shape == (180, 32) + assert radon_torch(imgs, theta=theta).shape == (2, 30, 48) + assert radon_torch(_disk(32, 16, 16, 6)).shape == (180, 32) # default theta def test_iradon_shapes(self): - sino = radon_torch(_disk(40, 20, 20, 8), theta=torch.linspace(0, 180, 25)) - rec = iradon_torch(sino, theta=torch.linspace(0, 180, 25)) - assert rec.shape == (40, 40) - - def test_iradon_output_size_override(self): - sino = radon_torch(_disk(40, 20, 20, 8), theta=torch.linspace(0, 180, 25)) - rec = iradon_torch(sino, theta=torch.linspace(0, 180, 25), output_size=32) - assert rec.shape == (32, 32) + theta = torch.linspace(0, 180, 25) + sino = radon_torch(_disk(40, 20, 20, 8), theta=theta) + assert iradon_torch(sino, theta=theta).shape == (40, 40) + assert iradon_torch(sino, theta=theta, output_size=32).shape == (32, 32) class TestFourierFilter: - def test_even_size_ok(self): - f = get_fourier_filter_torch(64, "ramp") - assert f.shape == (1, 64) - - def test_odd_size_raises(self): + def test_bad_arguments_raise(self): with pytest.raises(ValueError): get_fourier_filter_torch(63, "ramp") - - def test_unknown_filter_raises(self): with pytest.raises(ValueError): get_fourier_filter_torch(64, "not-a-filter") @@ -64,11 +45,15 @@ def test_none_filter_is_all_ones(self): f = get_fourier_filter_torch(64, None) assert torch.allclose(f, torch.ones_like(f)) - @pytest.mark.parametrize("name", ["ramp", "shepp-logan", "cosine", "hamming", "hann"]) - def test_named_filters_run(self, name): - f = get_fourier_filter_torch(64, name) - assert f.shape == (1, 64) - assert torch.isfinite(f).all() + def test_named_filters_window_the_ramp(self): + ramp = get_fourier_filter_torch(64, "ramp") + assert ramp.shape == (1, 64) + assert int(torch.argmax(ramp)) == 32 # the ramp peaks at Nyquist + for name in ["shepp-logan", "cosine", "hamming", "hann"]: + f = get_fourier_filter_torch(64, name) + assert f.shape == (1, 64), name + assert torch.all(f >= -1e-6) and torch.all(f <= ramp + 1e-6), name + assert f[0, 32] < ramp[0, 32], name # attenuated at Nyquist class TestRadonBehaviour: @@ -89,22 +74,18 @@ def test_iradon_circle_zeros_outside(self): assert torch.allclose(rec[outside], torch.zeros(int(outside.sum()))) def test_roundtrip_recovers_structure(self): - disk = _disk(64, 32, 24, 9) - theta = torch.linspace(0, 180, 60) - rec = iradon_torch(radon_torch(disk, theta=theta), theta=theta, filter_name="ramp") - corr = np.corrcoef(disk.numpy().ravel(), rec.numpy().ravel())[0, 1] - assert corr > 0.9 - - def test_default_theta_roundtrip_is_consistent(self): - """radon and iradon must share an angle convention when ``theta`` is defaulted. - - iradon's default previously included the 180-degree endpoint while radon's did not, - so a default-theta round-trip sampled mismatched angles. + """Also with ``theta`` defaulted on both sides: radon and iradon must share an angle + convention. iradon's default previously included the 180-degree endpoint while radon's + did not, so a default-theta round-trip sampled mismatched angles. """ - disk = _disk(64, 32, 28, 10) - rec = iradon_torch(radon_torch(disk), filter_name="ramp") # both default theta - corr = np.corrcoef(disk.numpy().ravel(), rec.numpy().ravel())[0, 1] - assert corr > 0.9 + theta = torch.linspace(0, 180, 60) + for disk, kwargs in [ + (_disk(64, 32, 24, 9), {"theta": theta}), + (_disk(64, 32, 28, 10), {}), + ]: + rec = iradon_torch(radon_torch(disk, **kwargs), filter_name="ramp", **kwargs) + corr = np.corrcoef(disk.numpy().ravel(), rec.numpy().ravel())[0, 1] + assert corr > 0.9, kwargs class TestRadonVsSkimage: diff --git a/tests/tomography/test_tomography.py b/tests/tomography/test_tomography.py index da71f9c9e..cd3064089 100644 --- a/tests/tomography/test_tomography.py +++ b/tests/tomography/test_tomography.py @@ -8,9 +8,9 @@ ``requires_torch`` (build on CUDA when present; see conftest). """ +import matplotlib.pyplot as plt import numpy as np import pytest -import torch from quantem.tomography.dataset_models import TomographyINRDataset, TomographyPixDataset from quantem.tomography.object_models import ObjConstraintParams, ObjectINR, ObjectPixelated @@ -35,12 +35,15 @@ def _conventional(n=12): class TestConventionalFactory: - def test_from_models_builds(self): + def test_from_models_builds_with_empty_history(self): tomo = _conventional() assert isinstance(tomo, TomographyConventional) assert isinstance(tomo.obj_model, ObjectPixelated) assert isinstance(tomo.dset, TomographyPixDataset) assert tomo.num_epochs == 0 + assert tomo.epoch_losses.shape == (0,) + assert tomo.consistency_losses.shape == (0,) + assert tomo.learning_rates == {} def test_direct_init_requires_token(self): with pytest.raises(RuntimeError): @@ -53,7 +56,7 @@ def test_direct_init_requires_token(self): class TestBaseProperties: - def test_constraints_setter_dict_and_object(self): + def test_constraints_setter(self): tomo = _conventional() tomo.constraints = {"name": "obj_pixelated", "tv_vol": 0.02} assert isinstance(tomo.constraints, ObjConstraintParams.ObjPixelatedConstraints) @@ -61,34 +64,18 @@ def test_constraints_setter_dict_and_object(self): obj_c = ObjConstraintParams.ObjPixelatedConstraints(positivity=True) tomo.constraints = obj_c assert tomo.constraints is obj_c - - def test_constraints_setter_none_is_noop(self): - tomo = _conventional() - before = tomo.constraints - tomo.constraints = None - assert tomo.constraints is before - - def test_constraints_setter_invalid_raises(self): - tomo = _conventional() + tomo.constraints = None # no-op + assert tomo.constraints is obj_c with pytest.raises(ValueError): tomo.constraints = 1.0 - def test_logger_setter_rejects_wrong_type(self): + def test_setters_reject_wrong_types(self): tomo = _conventional() with pytest.raises(TypeError): tomo.logger = "not a logger" - - def test_dset_setter_rejects_wrong_type(self): - tomo = _conventional() with pytest.raises(TypeError): tomo.dset = object() - def test_loss_accessors_start_empty(self): - tomo = _conventional() - assert tomo.epoch_losses.shape == (0,) - assert tomo.consistency_losses.shape == (0,) - assert tomo.learning_rates == {} - def test_append_learning_rates_accumulates(self): tomo = _conventional() tomo.append_learning_rates({"object": 1e-3, "pose": 1e-2}) @@ -96,15 +83,13 @@ def test_append_learning_rates_accumulates(self): assert tomo.learning_rates["object"] == [1e-3, 5e-4] assert tomo.learning_rates["pose"] == [1e-2, 5e-3] - def test_to_updates_device(self): - tomo = _conventional() - tomo.to("cpu") - assert torch.device(tomo.device) == torch.device("cpu") - - def test_plot_losses_runs(self): + def test_plot_losses(self): tomo = _conventional() tomo._epoch_losses.extend([1.0, 0.5, 0.25]) tomo.plot_losses() # Agg backend; plt.show() is a no-op + fig = plt.gcf() + assert np.allclose(fig.axes[0].lines[0].get_ydata(), [1.0, 0.5, 0.25]) + plt.close(fig) @requires_torch @@ -119,16 +104,15 @@ def _inr_tomo(self, device, n=16): ) return Tomography.from_models(dset=dset, obj_model=obj, device=device, verbose=False) - def test_from_models_builds(self, torch_device): - tomo = self._inr_tomo(torch_device) - assert isinstance(tomo, Tomography) - assert isinstance(tomo.obj_model, ObjectINR) - - def test_plot_losses_runs(self, torch_device): + def test_plot_losses(self, torch_device): tomo = self._inr_tomo(torch_device) tomo._epoch_losses.extend([1.0, 0.5]) tomo._lrs["object"] = [1e-3, 5e-4] tomo.plot_losses() + loss_ax, lr_ax = plt.gcf().axes[:2] + assert np.allclose(loss_ax.lines[0].get_ydata(), [1.0, 0.5]) + assert np.allclose(lr_ax.lines[0].get_ydata(), [1e-3, 5e-4]) + plt.close(plt.gcf()) def test_save_volume_overwrite_guard(self, torch_device, tmp_path): tomo = self._inr_tomo(torch_device) diff --git a/tests/tomography/test_tomography_opt.py b/tests/tomography/test_tomography_opt.py index bae00ad9a..c9898156d 100644 --- a/tests/tomography/test_tomography_opt.py +++ b/tests/tomography/test_tomography_opt.py @@ -49,16 +49,15 @@ def _td_tomography(device): return Tomography.from_models(dset=dset, obj_model=obj, device=device, verbose=False) +def _object_and_pose_params(): + return { + "object": OptimizerParams.Adam(lr=1e-3), + "pose": OptimizerParams.Adam(lr=1e-2), + } + + @requires_torch class TestOptimizerParams: - def test_setter_getter_roundtrip(self, torch_device): - tomo = _inr_tomography(torch_device) - tomo.optimizer_params = { - "object": OptimizerParams.Adam(lr=1e-3), - "pose": OptimizerParams.Adam(lr=1e-2), - } - assert set(tomo.optimizer_params.keys()) == {"object", "pose"} - def test_unknown_key_raises(self, torch_device): tomo = _inr_tomography(torch_device) with pytest.raises(ValueError): @@ -67,20 +66,10 @@ def test_unknown_key_raises(self, torch_device): def test_set_optimizers_builds_object_and_pose(self, torch_device): """Regression guard: the pose path must not raise (see module docstring).""" tomo = _inr_tomography(torch_device) - tomo.optimizer_params = { - "object": OptimizerParams.Adam(lr=1e-3), - "pose": OptimizerParams.Adam(lr=1e-2), - } + tomo.optimizer_params = _object_and_pose_params() + assert set(tomo.optimizer_params.keys()) == {"object", "pose"} tomo.set_optimizers() assert set(tomo.optimizers.keys()) == {"object", "pose"} - - def test_current_lrs(self, torch_device): - tomo = _inr_tomography(torch_device) - tomo.optimizer_params = { - "object": OptimizerParams.Adam(lr=1e-3), - "pose": OptimizerParams.Adam(lr=1e-2), - } - tomo.set_optimizers() lrs = tomo.get_current_lrs() assert set(lrs.keys()) == {"object", "pose"} assert lrs["object"] == pytest.approx(1e-3) @@ -88,10 +77,7 @@ def test_current_lrs(self, torch_device): def test_remove_optimizer(self, torch_device): tomo = _inr_tomography(torch_device) - tomo.optimizer_params = { - "object": OptimizerParams.Adam(lr=1e-3), - "pose": OptimizerParams.Adam(lr=1e-2), - } + tomo.optimizer_params = _object_and_pose_params() tomo.set_optimizers() tomo.remove_optimizer("object") assert "object" not in tomo.optimizers @@ -112,64 +98,36 @@ def test_pplr_object_groups(self, torch_device): assert len(tomo.optimizers["object"].param_groups) == 3 assert "pose" in tomo.optimizers - def test_step_optimizers_steps_each_once(self, torch_device, monkeypatch): - """Regression: with object+pose, each optimizer must step exactly once per call. + def test_step_and_zero_grad_touch_each_optimizer_once(self, torch_device, monkeypatch): + """Regression: with object+pose, each optimizer must step (and zero) exactly once per call. ``step_optimizers`` previously looped over both keys and stepped *both* optimizers on every pass, so each took two Adam steps per batch. """ tomo = _inr_tomography(torch_device) - tomo.optimizer_params = { - "object": OptimizerParams.Adam(lr=1e-3), - "pose": OptimizerParams.Adam(lr=1e-2), - } + tomo.optimizer_params = _object_and_pose_params() tomo.set_optimizers() assert set(tomo.optimizers.keys()) == {"object", "pose"} # both optimizers live - counts = {"object": 0, "pose": 0} - monkeypatch.setattr( - tomo.obj_model, - "step_optimizer", - lambda: counts.__setitem__("object", counts["object"] + 1), - ) - monkeypatch.setattr( - tomo.dset, "step_optimizer", lambda: counts.__setitem__("pose", counts["pose"] + 1) - ) - tomo.step_optimizers() - assert counts == {"object": 1, "pose": 1} - - def test_zero_grad_all_zeros_each_once(self, torch_device, monkeypatch): - """Companion to the step regression: zero_grad_all must touch each optimizer once.""" - tomo = _inr_tomography(torch_device) - tomo.optimizer_params = { - "object": OptimizerParams.Adam(lr=1e-3), - "pose": OptimizerParams.Adam(lr=1e-2), - } - tomo.set_optimizers() - counts = {"object": 0, "pose": 0} - monkeypatch.setattr( - tomo.obj_model, - "zero_optimizer_grad", - lambda: counts.__setitem__("object", counts["object"] + 1), - ) - monkeypatch.setattr( - tomo.dset, - "zero_optimizer_grad", - lambda: counts.__setitem__("pose", counts["pose"] + 1), - ) - tomo.zero_grad_all() - assert counts == {"object": 1, "pose": 1} + for tomo_method, model_method in [ + ("step_optimizers", "step_optimizer"), + ("zero_grad_all", "zero_optimizer_grad"), + ]: + counts = {"object": 0, "pose": 0} + for key, model in (("object", tomo.obj_model), ("pose", tomo.dset)): + monkeypatch.setattr( + model, model_method, lambda key=key: counts.__setitem__(key, counts[key] + 1) + ) + getattr(tomo, tomo_method)() + assert counts == {"object": 1, "pose": 1}, tomo_method @requires_torch class TestOptHelpers: - def test_get_default_lr_object_and_pose(self, torch_device): - tomo = _inr_tomography(torch_device) - assert isinstance(tomo._get_default_lr("object"), float) - assert isinstance(tomo._get_default_lr("pose"), float) - - def test_get_default_lr_unknown_raises(self, torch_device): + def test_get_default_lr(self, torch_device): tomo = _inr_tomography(torch_device) + assert tomo._get_default_lr("object") == pytest.approx(8e-6) + assert tomo._get_default_lr("pose") == pytest.approx(5e-2) with pytest.raises(ValueError): tomo._get_default_lr("banana") @@ -185,25 +143,15 @@ def test_current_lrs_zero_without_optimizers(self, torch_device): @requires_torch class TestSchedulerParams: - def test_scheduler_setter_getter(self, torch_device): - tomo = _inr_tomography(torch_device) - tomo.scheduler_params = { - "object": SchedulerParams.CosineAnnealing(T_max=10), - "pose": SchedulerParams.CosineAnnealing(T_max=10), - } - assert set(tomo.scheduler_params.keys()) == {"object", "pose"} - def test_set_schedulers_builds(self, torch_device): tomo = _inr_tomography(torch_device) - tomo.optimizer_params = { - "object": OptimizerParams.Adam(lr=1e-3), - "pose": OptimizerParams.Adam(lr=1e-2), - } + tomo.optimizer_params = _object_and_pose_params() tomo.set_optimizers() tomo.scheduler_params = { "object": SchedulerParams.CosineAnnealing(T_max=10), "pose": SchedulerParams.CosineAnnealing(T_max=10), } + assert set(tomo.scheduler_params.keys()) == {"object", "pose"} tomo.set_schedulers(tomo.scheduler_params, num_iter=10) assert set(tomo.schedulers.keys()) == {"object", "pose"} From 89c93f2196eb1ce55abdd57f925688003aeb5bb9 Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 14:59:38 -0700 Subject: [PATCH 07/19] Share the implicit-object test data and drop inherited duplicates The constants, probe, smooth phase, dataset simulation, ptychography builder and shift-tolerant correlation were copied verbatim into test_object_inr.py and test_object_tensor_decomp.py. They now live once, as an INR_-prefixed block appended to diffractive_imaging/conftest.py (nothing above it changed; the prefix avoids the direct-ptychography N/Q_PROBE/... names). The importlib concern in the old docstring does not apply: other modules here already import from .conftest. test_object_inr.py (23 -> 21 IDs): - test_multislice_z_coordinates and test_single_slice_z_is_zero merged. - test_autograd_only_backward_raises and test_complex_obj_type_not_supported merged into one test. test_object_tensor_decomp.py (27 -> 20 IDs). ObjectTensorDecomp overrides only its factories, name, model, optimizer grouping and pretrain, so these re-tested ObjectINR code and are removed (covered in test_object_inr.py): - test_off_object_is_vacuum (ObjectINR.forward masking) - test_multislice_z_coordinates (ObjectINR._z_coords); the 3-slice K-Planes forward shape is now part of test_forward_shape_dtype_and_vacuum - test_autograd_only_backward_raises (ObjectINR.backward) - TestImplicitSync (same assertion as the INR test) - slow test_multislice_runs (INR slow multislice test plus the unit forward) Also: test_is_implicit_and_name folded into the forward test, and the two KPlanesHead tests are one loop over use_hybrid_mlp. Both files pass with --runslow (thresholds unchanged). Co-Authored-By: Claude Opus 5.5 --- tests/diffractive_imaging/conftest.py | 133 +++++++++ tests/diffractive_imaging/test_object_inr.py | 174 ++--------- .../test_object_tensor_decomp.py | 270 +++--------------- 3 files changed, 196 insertions(+), 381 deletions(-) diff --git a/tests/diffractive_imaging/conftest.py b/tests/diffractive_imaging/conftest.py index f9a12034e..22595eb4e 100644 --- a/tests/diffractive_imaging/conftest.py +++ b/tests/diffractive_imaging/conftest.py @@ -820,3 +820,136 @@ def layer_purity(slices, truths) -> tuple[float, float]: purities.append(c[layer] ** 2 / (c**2).sum()) owns.append(c[layer]) return float(np.mean(purities)), float(np.mean(owns)) + + +# --- Implicit-object (ObjectINR / ObjectTensorDecomp) reconstruction data --------------------- +# A smooth phase object larger than the scanned region, scanned in its interior with padding +# >= roi // 2, so patches never reach the object boundary (the non-toroidal regime the +# implicit models assume). Used by test_object_inr.py and test_object_tensor_decomp.py. + +INR_N = 40 # detector / roi size (px) +INR_OGT = 64 # ground-truth object size (px); larger than the scanned region +INR_PAD = 20 # obj padding (>= roi // 2 so interior patches never hit the boundary) +INR_Q_MAX = 0.5 # inverse Angstroms +INR_Q_PROBE = INR_Q_MAX / 2 +INR_ENERGY = 300e3 # eV +INR_C10 = 50.0 # defocus (Angstrom) +INR_STEP = 2 # scan step (px) +INR_SCAN_START = 20 # first scan position (px); INR_SCAN_START - roi//2 >= 0 +INR_SCAN_STOP = 44 # exclusive; INR_SCAN_STOP - 1 + roi//2 - 1 < INR_OGT -> no wrap + + +def inr_smooth_phase() -> np.ndarray: + """A smooth, band-limited phase object (friendly to a SIREN, unlike white noise).""" + yy, xx = np.meshgrid(np.arange(INR_OGT), np.arange(INR_OGT), indexing="ij") + phase = 0.7 * np.sin(2 * np.pi * xx / INR_OGT * 4) * np.cos(2 * np.pi * yy / INR_OGT * 3) + return phase.astype(np.float32) + + +def inr_probe_array() -> np.ndarray: + sampling = 1 / INR_Q_MAX / 2 + reciprocal_sampling = 2 * INR_Q_MAX / INR_N + qx = qy = np.fft.fftfreq(INR_N, sampling) + q = np.sqrt(qx[:, None] ** 2 + qy[None, :] ** 2) + aperture = np.sqrt(np.clip((INR_Q_PROBE - q) / reciprocal_sampling + 0.5, 0, 1)) + chi = q**2 * electron_wavelength_angstrom(INR_ENERGY) * np.pi * INR_C10 + probe_fourier = aperture * np.exp(-1j * chi) + probe_fourier /= np.sqrt(np.sum(np.abs(probe_fourier) ** 2)) + return (np.fft.ifft2(probe_fourier) * INR_N).astype(np.complex64) + + +def inr_semiangle_mrad() -> float: + return electron_wavelength_angstrom(INR_ENERGY) * INR_Q_PROBE * 1e3 + + +def inr_probe_params() -> dict: + return {"energy": INR_ENERGY, "C10": INR_C10, "semiangle_cutoff": inr_semiangle_mrad()} + + +def inr_synthetic_dataset(): + """Simulate the non-toroidal 4D-STEM dataset; return ``(dataset_model, gt_phase, probe)``.""" + from quantem.core.datastructures.dataset4dstem import Dataset4dstem + from quantem.diffractive_imaging.dataset_models import PtychographyDatasetRaster + + phase = inr_smooth_phase() + complex_obj = np.exp(1j * phase) + probe = inr_probe_array() + reciprocal_sampling = 2 * INR_Q_MAX / INR_N + + gpos = np.arange(INR_SCAN_START, INR_SCAN_STOP, INR_STEP) + xx, yy = np.meshgrid(gpos, gpos, indexing="ij") + positions = np.stack((xx.ravel(), yy.ravel()), axis=-1) + x0 = positions[:, 0].astype(int) + y0 = positions[:, 1].astype(int) + x_ind = np.fft.fftfreq(INR_N, d=1 / INR_N).astype(int) + y_ind = np.fft.fftfreq(INR_N, d=1 / INR_N).astype(int) + row = x0[:, None, None] + x_ind[None, :, None] # no modulo: interior scan, no wrap + col = y0[:, None, None] + y_ind[None, None, :] + assert row.min() >= 0 and row.max() < INR_OGT and col.min() >= 0 and col.max() < INR_OGT + exit_waves = complex_obj[row, col] * probe + intensities = np.abs(np.fft.fft2(exit_waves)) ** 2 + + sxy = len(gpos) + dset = Dataset4dstem.from_array( + array=np.fft.fftshift(intensities * 100, axes=(-2, -1)).reshape((sxy, sxy, INR_N, INR_N)), + sampling=(INR_STEP, INR_STEP, reciprocal_sampling, reciprocal_sampling), + units=("A", "A", "A^-1", "A^-1"), + ) + pdset = PtychographyDatasetRaster.from_dataset4dstem(dset) + pdset.learn_scan_positions = False + pdset.learn_descan = False + pdset.preprocess( + com_fit_function="constant", + plot_rotation=False, + plot_com=False, + probe_energy=INR_ENERGY, + ) + return pdset, phase, probe + + +def inr_ptychography(obj_model, pdset=None, probe_model=None): + """``Ptychography`` on the synthetic dataset; the exact simulation probe, frozen, by default.""" + from quantem.diffractive_imaging.detector_models import DetectorPixelated + from quantem.diffractive_imaging.probe_models import ProbePixelated + from quantem.diffractive_imaging.ptychography import Ptychography + + gt_phase = None + if pdset is None: + pdset, gt_phase, probe = inr_synthetic_dataset() + if probe_model is None: + probe_model = ProbePixelated.from_array( + num_probes=1, probe_params=inr_probe_params(), probe_array=probe.copy() + ) + ptycho = Ptychography.from_models( + dset=pdset, + obj_model=obj_model, + probe_model=probe_model, + detector_model=DetectorPixelated(), + rng=0, + verbose=False, + ) + ptycho.preprocess(obj_padding_px=(INR_PAD, INR_PAD)) + return ptycho, gt_phase + + +def _inr_corr(a: np.ndarray, b: np.ndarray) -> float: + a = a - a.mean() + b = b - b.mean() + return float((a * b).sum() / np.sqrt((a**2).sum() * (b**2).sum() + 1e-12)) + + +def _inr_center_crop(a: np.ndarray, s: int) -> np.ndarray: + r0 = (a.shape[0] - s) // 2 + c0 = (a.shape[1] - s) // 2 + return a[r0 : r0 + s, c0 : c0 + s] + + +def inr_best_corr(recon: np.ndarray, gt: np.ndarray, s: int = 32) -> float: + """Best |correlation| over small integer shifts (handles ~1px registration offset).""" + g = _inr_center_crop(gt, s) + best = -1.0 + for dr in range(-3, 4): + for dc in range(-3, 4): + r = _inr_center_crop(np.roll(recon, (dr, dc), (0, 1)), s) + best = max(best, abs(_inr_corr(r, g))) + return best diff --git a/tests/diffractive_imaging/test_object_inr.py b/tests/diffractive_imaging/test_object_inr.py index 2e3c50b07..e69a4f531 100644 --- a/tests/diffractive_imaging/test_object_inr.py +++ b/tests/diffractive_imaging/test_object_inr.py @@ -5,11 +5,11 @@ masking, multislice z-coordinates, gradient flow), the dataset's implicit-object coordinate production, and an end-to-end reconstruction on a smooth synthetic object. -The reconstruction fixtures deliberately differ from ``test_ptychography.py``: that -fixture scans a torus-wrapped object (edge positions, zero padding), which the -pixelated path reproduces via ``% obj_shape`` indexing but the INR (vacuum outside the -object) cannot. Here the ground-truth object is larger than the scanned region and the -scan is confined to the interior with padding >= roi // 2, so patches never reach the +The reconstruction data (``inr_synthetic_dataset`` in ``conftest.py``) deliberately differs +from ``test_ptychography.py``: that fixture scans a torus-wrapped object (edge positions, zero +padding), which the pixelated path reproduces via ``% obj_shape`` indexing but the INR (vacuum +outside the object) cannot. Here the ground-truth object is larger than the scanned region and +the scan is confined to the interior with padding >= roi // 2, so patches never reach the object boundary -- the physically realistic, non-toroidal regime both models agree on. """ @@ -18,32 +18,20 @@ import torch from quantem.core import config -from quantem.core.datastructures.dataset4dstem import Dataset4dstem from quantem.core.io.serialize import load as autoserialize_load from quantem.core.ml import OptimizerParams, SchedulerParams from quantem.core.ml.cnn import CNN2d -from quantem.core.utils.utils import electron_wavelength_angstrom -from quantem.diffractive_imaging.dataset_models import PtychographyDatasetRaster -from quantem.diffractive_imaging.detector_models import DetectorPixelated from quantem.diffractive_imaging.object_models import ObjectINR, ObjectPixelated from quantem.diffractive_imaging.probe_models import ProbeDIP, ProbeParametric, ProbePixelated -from quantem.diffractive_imaging.ptychography import Ptychography + +from .conftest import INR_N as N +from .conftest import INR_PAD as PAD +from .conftest import inr_best_corr, inr_probe_params, inr_ptychography, inr_synthetic_dataset if config.NUM_DEVICES > 0: config.set_device("gpu") -N = 40 # detector / roi size (px) -OGT = 64 # ground-truth object size (px); larger than the scanned region -PAD = 20 # obj padding (>= roi // 2 so interior patches never hit the boundary) -Q_MAX = 0.5 # inverse Angstroms -Q_PROBE = Q_MAX / 2 -PROBE_ENERGY = 300e3 # eV -C10 = 50.0 # defocus (Angstrom) -STEP = 2 # scan step (px) -SCAN_START = 20 # first scan position (px); SCAN_START - roi//2 >= 0 -SCAN_STOP = 44 # exclusive; SCAN_STOP - 1 + roi//2 - 1 < OGT -> no wrap - -# Reconstruction config (validated against the fixture below: pixelated reaches corr~0.89, +# Reconstruction config (validated against the conftest dataset: pixelated reaches corr~0.89, # the INR reaches corr~0.91 at omega_0=5 / lr=1e-2). omega_0=5 suits this smooth object; # the default of 10 (and the image-fitting default of 30) stall here. _INR_OMEGA = 5.0 @@ -52,73 +40,10 @@ _CORR_THRESHOLD = 0.7 -def _smooth_phase() -> np.ndarray: - """A smooth, band-limited phase object (friendly to a SIREN, unlike white noise).""" - yy, xx = np.meshgrid(np.arange(OGT), np.arange(OGT), indexing="ij") - return (0.7 * np.sin(2 * np.pi * xx / OGT * 4) * np.cos(2 * np.pi * yy / OGT * 3)).astype( - np.float32 - ) - - -def _probe_array() -> np.ndarray: - sampling = 1 / Q_MAX / 2 - reciprocal_sampling = 2 * Q_MAX / N - qx = qy = np.fft.fftfreq(N, sampling) - q = np.sqrt(qx[:, None] ** 2 + qy[None, :] ** 2) - aperture = np.sqrt(np.clip((Q_PROBE - q) / reciprocal_sampling + 0.5, 0, 1)) - chi = q**2 * electron_wavelength_angstrom(PROBE_ENERGY) * np.pi * C10 - probe_fourier = aperture * np.exp(-1j * chi) - probe_fourier /= np.sqrt(np.sum(np.abs(probe_fourier) ** 2)) - return (np.fft.ifft2(probe_fourier) * N).astype(np.complex64) - - -def _semiangle_mrad() -> float: - return electron_wavelength_angstrom(PROBE_ENERGY) * Q_PROBE * 1e3 - - -def _build_synthetic_dataset() -> tuple[PtychographyDatasetRaster, np.ndarray, np.ndarray]: - """Simulate a non-toroidal 4D-STEM dataset; return (dataset_model, gt_phase, probe).""" - phase = _smooth_phase() - complex_obj = np.exp(1j * phase) - probe = _probe_array() - reciprocal_sampling = 2 * Q_MAX / N - - gpos = np.arange(SCAN_START, SCAN_STOP, STEP) - xx, yy = np.meshgrid(gpos, gpos, indexing="ij") - positions = np.stack((xx.ravel(), yy.ravel()), axis=-1) - x0 = positions[:, 0].astype(int) - y0 = positions[:, 1].astype(int) - x_ind = np.fft.fftfreq(N, d=1 / N).astype(int) - y_ind = np.fft.fftfreq(N, d=1 / N).astype(int) - row = x0[:, None, None] + x_ind[None, :, None] # no modulo: interior scan, no wrap - col = y0[:, None, None] + y_ind[None, None, :] - assert row.min() >= 0 and row.max() < OGT and col.min() >= 0 and col.max() < OGT - exit_waves = complex_obj[row, col] * probe - intensities = np.abs(np.fft.fft2(exit_waves)) ** 2 - - sxy = len(gpos) - dset = Dataset4dstem.from_array( - array=np.fft.fftshift(intensities * 100, axes=(-2, -1)).reshape((sxy, sxy, N, N)), - sampling=(STEP, STEP, reciprocal_sampling, reciprocal_sampling), - units=("A", "A", "A^-1", "A^-1"), - ) - pdset = PtychographyDatasetRaster.from_dataset4dstem(dset) - pdset.learn_scan_positions = False - pdset.learn_descan = False - pdset.preprocess( - com_fit_function="constant", - plot_rotation=False, - plot_com=False, - probe_energy=PROBE_ENERGY, - ) - return pdset, phase, probe - - def _build_inr_ptycho( num_slices: int = 1, slice_thicknesses=None, first_omega_0: float = _INR_OMEGA -) -> tuple[Ptychography, np.ndarray]: +): """Build an ObjectINR ptychography with the exact (frozen) simulation probe.""" - pdset, gt_phase, probe = _build_synthetic_dataset() obj = ObjectINR.from_uniform( num_slices=num_slices, slice_thicknesses=slice_thicknesses, @@ -128,44 +53,7 @@ def _build_inr_ptycho( hidden_omega_0=first_omega_0, rng=0, ) - probe_model = ProbePixelated.from_array( - num_probes=1, - probe_params={"energy": PROBE_ENERGY, "C10": C10, "semiangle_cutoff": _semiangle_mrad()}, - probe_array=probe.copy(), - ) - ptycho = Ptychography.from_models( - dset=pdset, - obj_model=obj, - probe_model=probe_model, - detector_model=DetectorPixelated(), - rng=0, - verbose=False, - ) - ptycho.preprocess(obj_padding_px=(PAD, PAD)) - return ptycho, gt_phase - - -def _corr(a: np.ndarray, b: np.ndarray) -> float: - a = a - a.mean() - b = b - b.mean() - return float((a * b).sum() / np.sqrt((a**2).sum() * (b**2).sum() + 1e-12)) - - -def _center_crop(a: np.ndarray, s: int) -> np.ndarray: - r0 = (a.shape[0] - s) // 2 - c0 = (a.shape[1] - s) // 2 - return a[r0 : r0 + s, c0 : c0 + s] - - -def _best_corr(recon: np.ndarray, gt: np.ndarray, s: int = 32) -> float: - """Best |correlation| over small integer shifts (handles ~1px registration offset).""" - g = _center_crop(gt, s) - best = -1.0 - for dr in range(-3, 4): - for dc in range(-3, 4): - r = _center_crop(np.roll(recon, (dr, dc), (0, 1)), s) - best = max(best, abs(_corr(r, g))) - return best + return inr_ptychography(obj) # --------------------------------------------------------------------------- # @@ -197,7 +85,7 @@ def test_off_object_is_vacuum(self): out = obj.forward(off) assert torch.allclose(out, torch.ones_like(out), atol=1e-6) - def test_multislice_z_coordinates(self): + def test_z_coordinates(self): obj = ObjectINR.from_uniform( num_slices=3, slice_thicknesses=2.0, hidden_features=32, rng=0 ) @@ -209,10 +97,9 @@ def test_multislice_z_coordinates(self): patches = obj.forward(torch.rand(4, 6, 6, 2) * 2 - 1) assert patches.shape == (3, 4, 6, 6) - def test_single_slice_z_is_zero(self): - obj = ObjectINR.from_uniform(num_slices=1, hidden_features=32, rng=0) - obj._initialize_obj((1, 8, 8)) - assert torch.allclose(obj._z_coords, torch.zeros(1)) + single = ObjectINR.from_uniform(num_slices=1, hidden_features=32, rng=0) + single._initialize_obj((1, 8, 8)) + assert torch.allclose(single._z_coords, torch.zeros(1)) def test_obj_materialization(self): obj = ObjectINR.from_uniform( @@ -253,12 +140,10 @@ def test_gradients_flow_and_fit_smooth_phase(self): losses.append(loss.item()) assert losses[-1] < 0.05 * losses[0] - def test_autograd_only_backward_raises(self): + def test_unsupported_paths_raise(self): obj = ObjectINR.from_uniform(num_slices=1, hidden_features=32, rng=0) - with pytest.raises(NotImplementedError): + with pytest.raises(NotImplementedError): # autograd only obj.backward() - - def test_complex_obj_type_not_supported(self): with pytest.raises(NotImplementedError): ObjectINR.from_uniform(num_slices=1, obj_type="complex") @@ -384,7 +269,7 @@ def _corr_to_pix(arr: np.ndarray) -> float: class TestImplicitDatasetCoords: def test_scan_coords_normalization(self): """An integer scan position maps to linspace(-1,1,N) grid nodes; spacing = 1 px.""" - pdset, _, _ = _build_synthetic_dataset() + pdset, _, _ = inr_synthetic_dataset() pdset.implicit_object = True padding = (PAD, PAD) h_full, w_full = pdset._obj_shape_full_2d(padding) @@ -405,7 +290,7 @@ def test_scan_coords_normalization(self): assert d_col == pytest.approx(2 / (int(w_full) - 1), abs=1e-5) def test_forward_returns_coords_and_zero_fractional_when_implicit(self): - pdset, _, _ = _build_synthetic_dataset() + pdset, _, _ = inr_synthetic_dataset() pdset.implicit_object = True batch = torch.arange(4) patch_data, positions_px, fractional, _descan = pdset.forward(batch, (PAD, PAD)) @@ -435,7 +320,7 @@ def test_loss_decreases_and_recovers_object(self): ) losses = np.array(ptycho._iter_losses) assert losses[-1] < _LOSS_RATIO * losses[0] - assert _best_corr(ptycho.obj[0], gt_phase) > _CORR_THRESHOLD + assert inr_best_corr(ptycho.obj[0], gt_phase) > _CORR_THRESHOLD def test_multislice_runs(self): ptycho, _ = _build_inr_ptycho(num_slices=2, slice_thicknesses=20.0) @@ -496,12 +381,9 @@ class TestObjectINRProbeTypes: reconstruction runs end-to-end and the loss stays finite. """ - def _probe_params(self) -> dict: - return {"energy": PROBE_ENERGY, "C10": C10, "semiangle_cutoff": _semiangle_mrad()} - def test_inr_runs_with_each_probe_type(self): - pdset, _gt, _probe = _build_synthetic_dataset() - pp = self._probe_params() + pdset, _gt, _probe = inr_synthetic_dataset() + pp = inr_probe_params() probe_builders = { "pixelated": lambda: ProbePixelated.from_params(probe_params=pp), "parametric": lambda: ProbeParametric.from_params(probe_params=pp), @@ -521,15 +403,7 @@ def test_inr_runs_with_each_probe_type(self): hidden_omega_0=_INR_OMEGA, rng=0, ) - ptycho = Ptychography.from_models( - dset=pdset, - obj_model=obj, - probe_model=build(), - detector_model=DetectorPixelated(), - rng=0, - verbose=False, - ) - ptycho.preprocess(obj_padding_px=(PAD, PAD)) + ptycho, _ = inr_ptychography(obj, pdset=pdset, probe_model=build()) assert ptycho.dset.implicit_object is True, name ptycho.reconstruct( num_iters=10, diff --git a/tests/diffractive_imaging/test_object_tensor_decomp.py b/tests/diffractive_imaging/test_object_tensor_decomp.py index 17d87055c..15136a304 100644 --- a/tests/diffractive_imaging/test_object_tensor_decomp.py +++ b/tests/diffractive_imaging/test_object_tensor_decomp.py @@ -7,9 +7,9 @@ shared-core ``KPlanes`` head fix, the object model in isolation (forward/vacuum/PPLR parameter groups, bad-key rejection), the ptychography PPLR optimizer pass-through, and an end-to-end reconstruction on the same non-toroidal synthetic object used by -``test_object_inr.py`` (fixtures duplicated here, matching that file's own duplication -from ``test_ptychography.py`` — ``--import-mode=importlib`` makes cross-test imports -unreliable). +``test_object_inr.py`` (``inr_synthetic_dataset`` in ``conftest.py``). The coordinate-query +paths it inherits (off-object masking, z coordinates, autograd-only backward, the implicit +flag) are tested once, in ``test_object_inr.py``. """ import numpy as np @@ -18,36 +18,21 @@ import torch.nn as nn from quantem.core import config -from quantem.core.datastructures.dataset4dstem import Dataset4dstem from quantem.core.io.serialize import load as autoserialize_load from quantem.core.ml import OptimizerParams, SchedulerParams from quantem.core.ml.models.kplanes import CPTilted, KPlanes -from quantem.core.utils.utils import electron_wavelength_angstrom -from quantem.diffractive_imaging.dataset_models import PtychographyDatasetRaster -from quantem.diffractive_imaging.detector_models import DetectorPixelated from quantem.diffractive_imaging.object_models import ( ObjectModelType, ObjectPixelated, ObjectTensorDecomp, ) -from quantem.diffractive_imaging.probe_models import ProbePixelated -from quantem.diffractive_imaging.ptychography import Ptychography + +from .conftest import inr_best_corr, inr_ptychography if config.NUM_DEVICES > 0: config.set_device("gpu") -N = 40 # detector / roi size (px) -OGT = 64 # ground-truth object size (px); larger than the scanned region -PAD = 20 # obj padding (>= roi // 2 so interior patches never hit the boundary) -Q_MAX = 0.5 # inverse Angstroms -Q_PROBE = Q_MAX / 2 -PROBE_ENERGY = 300e3 # eV -C10 = 50.0 # defocus (Angstrom) -STEP = 2 # scan step (px) -SCAN_START = 20 # first scan position (px); SCAN_START - roi//2 >= 0 -SCAN_STOP = 44 # exclusive; SCAN_STOP - 1 + roi//2 - 1 < OGT -> no wrap - -# K-Planes recon config (validated against the fixture below). grids learn faster than the +# K-Planes recon config (validated against the conftest dataset). grids learn faster than the # decoder head; both run under a single cosine-annealing schedule. _LR_GRIDS = 1e-2 _LR_SIGMA = 1e-3 @@ -55,76 +40,13 @@ _CORR_THRESHOLD = 0.6 -def _smooth_phase() -> np.ndarray: - """A smooth, band-limited phase object.""" - yy, xx = np.meshgrid(np.arange(OGT), np.arange(OGT), indexing="ij") - return (0.7 * np.sin(2 * np.pi * xx / OGT * 4) * np.cos(2 * np.pi * yy / OGT * 3)).astype( - np.float32 - ) - - -def _probe_array() -> np.ndarray: - sampling = 1 / Q_MAX / 2 - reciprocal_sampling = 2 * Q_MAX / N - qx = qy = np.fft.fftfreq(N, sampling) - q = np.sqrt(qx[:, None] ** 2 + qy[None, :] ** 2) - aperture = np.sqrt(np.clip((Q_PROBE - q) / reciprocal_sampling + 0.5, 0, 1)) - chi = q**2 * electron_wavelength_angstrom(PROBE_ENERGY) * np.pi * C10 - probe_fourier = aperture * np.exp(-1j * chi) - probe_fourier /= np.sqrt(np.sum(np.abs(probe_fourier) ** 2)) - return (np.fft.ifft2(probe_fourier) * N).astype(np.complex64) - - -def _semiangle_mrad() -> float: - return electron_wavelength_angstrom(PROBE_ENERGY) * Q_PROBE * 1e3 - - -def _build_synthetic_dataset() -> tuple[PtychographyDatasetRaster, np.ndarray, np.ndarray]: - """Simulate a non-toroidal 4D-STEM dataset; return (dataset_model, gt_phase, probe).""" - phase = _smooth_phase() - complex_obj = np.exp(1j * phase) - probe = _probe_array() - reciprocal_sampling = 2 * Q_MAX / N - - gpos = np.arange(SCAN_START, SCAN_STOP, STEP) - xx, yy = np.meshgrid(gpos, gpos, indexing="ij") - positions = np.stack((xx.ravel(), yy.ravel()), axis=-1) - x0 = positions[:, 0].astype(int) - y0 = positions[:, 1].astype(int) - x_ind = np.fft.fftfreq(N, d=1 / N).astype(int) - y_ind = np.fft.fftfreq(N, d=1 / N).astype(int) - row = x0[:, None, None] + x_ind[None, :, None] # no modulo: interior scan, no wrap - col = y0[:, None, None] + y_ind[None, None, :] - assert row.min() >= 0 and row.max() < OGT and col.min() >= 0 and col.max() < OGT - exit_waves = complex_obj[row, col] * probe - intensities = np.abs(np.fft.fft2(exit_waves)) ** 2 - - sxy = len(gpos) - dset = Dataset4dstem.from_array( - array=np.fft.fftshift(intensities * 100, axes=(-2, -1)).reshape((sxy, sxy, N, N)), - sampling=(STEP, STEP, reciprocal_sampling, reciprocal_sampling), - units=("A", "A", "A^-1", "A^-1"), - ) - pdset = PtychographyDatasetRaster.from_dataset4dstem(dset) - pdset.learn_scan_positions = False - pdset.learn_descan = False - pdset.preprocess( - com_fit_function="constant", - plot_rotation=False, - plot_com=False, - probe_energy=PROBE_ENERGY, - ) - return pdset, phase, probe - - def _build_kplanes_ptycho( num_slices: int = 1, slice_thicknesses=None, M_features: int = 24, resolution: tuple[int, int, int] = (16, 64, 64), -) -> tuple[Ptychography, np.ndarray]: +): """Build an ObjectTensorDecomp (K-Planes) ptychography with the exact (frozen) probe.""" - pdset, gt_phase, probe = _build_synthetic_dataset() obj = ObjectTensorDecomp.from_uniform( num_slices=num_slices, slice_thicknesses=slice_thicknesses, @@ -135,44 +57,7 @@ def _build_kplanes_ptycho( obj_type="pure_phase", rng=0, ) - probe_model = ProbePixelated.from_array( - num_probes=1, - probe_params={"energy": PROBE_ENERGY, "C10": C10, "semiangle_cutoff": _semiangle_mrad()}, - probe_array=probe.copy(), - ) - ptycho = Ptychography.from_models( - dset=pdset, - obj_model=obj, - probe_model=probe_model, - detector_model=DetectorPixelated(), - rng=0, - verbose=False, - ) - ptycho.preprocess(obj_padding_px=(PAD, PAD)) - return ptycho, gt_phase - - -def _corr(a: np.ndarray, b: np.ndarray) -> float: - a = a - a.mean() - b = b - b.mean() - return float((a * b).sum() / np.sqrt((a**2).sum() * (b**2).sum() + 1e-12)) - - -def _center_crop(a: np.ndarray, s: int) -> np.ndarray: - r0 = (a.shape[0] - s) // 2 - c0 = (a.shape[1] - s) // 2 - return a[r0 : r0 + s, c0 : c0 + s] - - -def _best_corr(recon: np.ndarray, gt: np.ndarray, s: int = 32) -> float: - """Best |correlation| over small integer shifts (handles ~1px registration offset).""" - g = _center_crop(gt, s) - best = -1.0 - for dr in range(-3, 4): - for dc in range(-3, 4): - r = _center_crop(np.roll(recon, (dr, dc), (0, 1)), s) - best = max(best, abs(_corr(r, g))) - return best + return inr_ptychography(obj) def _pplr_params() -> dict: @@ -188,31 +73,23 @@ def _pplr_params() -> dict: # Shared-core KPlanes head fix (use_hybrid_mlp=False must build a decoder head) # --------------------------------------------------------------------------- # class TestKPlanesHead: - def test_non_hybrid_builds_head_and_runs(self): - m = KPlanes( - M_features=8, - resolution=(16, 16, 16), - multiscale_res_multipliers=[0.5, 1.0], - density_activation=nn.Identity(), - use_hybrid_mlp=False, - ) - assert isinstance(m.sigma_net, nn.Linear) - out = m(torch.rand(20, 3) * 2 - 1) - assert out.shape == (20, 1) - assert set(m.get_params().keys()) == {"grids", "sigma_net"} - - def test_hybrid_still_builds_mlp_head(self): - m = KPlanes( - M_features=8, - resolution=(16, 16, 16), - multiscale_res_multipliers=[1.0], - density_activation=nn.Identity(), - use_hybrid_mlp=True, - hybrid_hidden_dim=32, - hybrid_num_layers=2, - ) - assert isinstance(m.sigma_net, nn.Sequential) - assert m(torch.rand(7, 3) * 2 - 1).shape == (7, 1) + def test_head_type_follows_use_hybrid_mlp(self): + """use_hybrid_mlp=False builds a linear decoder head, True an MLP; both run.""" + for use_hybrid_mlp, head, extra in [ + (False, nn.Linear, {}), + (True, nn.Sequential, {"hybrid_hidden_dim": 32, "hybrid_num_layers": 2}), + ]: + m = KPlanes( + M_features=8, + resolution=(16, 16, 16), + multiscale_res_multipliers=[0.5, 1.0], + density_activation=nn.Identity(), + use_hybrid_mlp=use_hybrid_mlp, + **extra, + ) + assert isinstance(m.sigma_net, head), use_hybrid_mlp + assert m(torch.rand(20, 3) * 2 - 1).shape == (20, 1), use_hybrid_mlp + assert set(m.get_params().keys()) == {"grids", "sigma_net"}, use_hybrid_mlp # --------------------------------------------------------------------------- # @@ -230,44 +107,19 @@ def _obj(self, **kw): **kw, ) - def test_is_implicit_and_name(self): - obj = self._obj() - assert obj.is_implicit is True - assert obj.name == "ObjectTensorDecomp" - assert isinstance(obj, ObjectModelType) - def test_forward_shape_dtype_and_vacuum(self): - obj = self._obj() - obj._initialize_obj((1, 24, 28)) - coords = torch.rand(5, 8, 8, 2) * 2 - 1 - patches = obj.forward(coords) - assert patches.shape == (1, 5, 8, 8) - assert patches.is_complex() - # vacuum init (zeroed decoder head) -> unit transmission everywhere - assert torch.allclose(patches, torch.ones_like(patches), atol=1e-6) - - def test_off_object_is_vacuum(self): - obj = self._obj(resolution=(16, 16, 16)) - obj._initialize_obj((1, 16, 16)) - # train the head a little so the model is not identically zero - opt = torch.optim.Adam(obj.model.parameters(), lr=1e-2) - inside = torch.rand(3, 4, 4, 2) * 2 - 1 - for _ in range(5): - opt.zero_grad() - obj.forward(inside).imag.sum().backward() - opt.step() - off = torch.full((1, 2, 2, 2), 5.0) - out = obj.forward(off) - assert torch.allclose(out, torch.ones_like(out), atol=1e-6) - - def test_multislice_z_coordinates(self): - obj = self._obj(num_slices=3, slice_thicknesses=2.0, resolution=(16, 16, 16)) - obj._initialize_obj((3, 16, 16)) - z = obj._z_coords - assert z.shape == (3,) - assert torch.allclose(z, torch.tensor([-1.0, 0.0, 1.0]).to(z), atol=1e-6) - patches = obj.forward(torch.rand(4, 6, 6, 2) * 2 - 1) - assert patches.shape == (3, 4, 6, 6) + for num_slices, thickness in [(1, None), (3, 2.0)]: + obj = self._obj(num_slices=num_slices, slice_thicknesses=thickness) + assert obj.is_implicit is True + assert obj.name == "ObjectTensorDecomp" + assert isinstance(obj, ObjectModelType) + obj._initialize_obj((num_slices, 24, 28)) + coords = torch.rand(5, 8, 8, 2) * 2 - 1 + patches = obj.forward(coords) + assert patches.shape == (num_slices, 5, 8, 8), num_slices + assert patches.is_complex(), num_slices + # vacuum init (zeroed decoder head) -> unit transmission everywhere + assert torch.allclose(patches, torch.ones_like(patches), atol=1e-6), num_slices def test_get_optimization_parameters_keys(self): obj = self._obj() @@ -322,11 +174,6 @@ def test_disable_sentinel_accepted(self): obj.set_optimizer({"default": OptimizerParams.NoneOptimizer()}) assert obj.optimizer is None - def test_autograd_only_backward_raises(self): - obj = self._obj() - with pytest.raises(NotImplementedError): - obj.backward() - def test_from_model_warns_on_lambda_activation(self): # a lambda activation cannot be pickled by AutoSerialize -> from_model must warn model = KPlanes( @@ -429,16 +276,6 @@ def test_potential_identity_default_and_positivity_penalty(self): assert float(obj._materialize_obj().min()) > -1e-2 # driven non-negative -# --------------------------------------------------------------------------- # -# Implicit-flag sync (ObjectTensorDecomp is implicit, like ObjectINR) -# --------------------------------------------------------------------------- # -class TestImplicitSync: - def test_implicit_flag_synced_from_obj_model(self): - ptycho, _ = _build_kplanes_ptycho() - assert ptycho.obj_model.is_implicit is True - assert ptycho.dset.implicit_object is True - - # --------------------------------------------------------------------------- # # PPLR optimizer pass-through through the ptychography optimizer layer # --------------------------------------------------------------------------- # @@ -477,7 +314,7 @@ def test_loss_decreases_and_recovers_object(self): ) losses = np.array(ptycho._iter_losses) assert losses[-1] < _LOSS_RATIO * losses[0] - assert _best_corr(ptycho.obj[0], gt_phase) > _CORR_THRESHOLD + assert inr_best_corr(ptycho.obj[0], gt_phase) > _CORR_THRESHOLD def test_reconstruct_with_reset_runs(self): """reset=True replays reset_optimizer with the disabled sentinel before applying PPLR.""" @@ -497,7 +334,6 @@ def test_reconstruct_with_reset_runs(self): def test_tilted_reconstruct_runs(self): """A tilted K-Planes object reconstructs end-to-end with a 3-group PPLR optimizer.""" - pdset, gt_phase, probe = _build_synthetic_dataset() obj = ObjectTensorDecomp.from_uniform( num_slices=1, M_features=12, @@ -508,24 +344,7 @@ def test_tilted_reconstruct_runs(self): obj_type="pure_phase", rng=0, ) - probe_model = ProbePixelated.from_array( - num_probes=1, - probe_params={ - "energy": PROBE_ENERGY, - "C10": C10, - "semiangle_cutoff": _semiangle_mrad(), - }, - probe_array=probe.copy(), - ) - ptycho = Ptychography.from_models( - dset=pdset, - obj_model=obj, - probe_model=probe_model, - detector_model=DetectorPixelated(), - rng=0, - verbose=False, - ) - ptycho.preprocess(obj_padding_px=(PAD, PAD)) + ptycho, _ = inr_ptychography(obj) ptycho.reconstruct( num_iters=60, optimizer_params={ @@ -554,17 +373,6 @@ def test_tv_constraint_reconstruct_runs(self): assert np.isfinite(losses).all() assert ptycho.obj_model.constraints.tv_weight_xy == 1e-3 - def test_multislice_runs(self): - ptycho, _ = _build_kplanes_ptycho( - num_slices=2, slice_thicknesses=20.0, resolution=(8, 64, 64) - ) - ptycho.reconstruct( - num_iters=5, - optimizer_params=_pplr_params(), - batch_size=200, - ) - assert ptycho.obj.shape[0] == 2 - def test_scheduler_scales_pplr_groups(self): ptycho, _ = _build_kplanes_ptycho() ptycho.reconstruct( From aabe7c94c89f428b54f698586e4f6229f4699a78 Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 15:01:10 -0700 Subject: [PATCH 08/19] Build target-residency tests on a small dataset; merge ptycho state tests test_ptychography.py (18 -> 4 IDs): - TestTargetResidency ran on the 64^4 white-noise fixture (4096 simulated 64x64 patterns, preprocessed with rotation/CoM plots), rebuilt for each of 9 IDs, about 0.8 s each, to test a string property. It now uses an 8x8x16x16 random dataset; default/valid/invalid-setter checks are one test with the six invalid values looped and named; the save/load check is kept. - TestAspectRatioRotatedFOV: the positions-inside-object and obj_cropped checks built the same rotated 32x64 scan separately; one test now builds it once per transpose and runs both checks (4 builds -> 2). - Slow TestPtychographySaveLoadRoundtrip: the function-scoped trained fixture reran the 50-iteration 64^4 reconstruction for each of 5 tests. One test now reconstructs once and checks the reloaded losses, positions and object (exact), the strict 0.3 loss threshold, and the continue-training behaviour (about 24 s once instead of five times). Co-Authored-By: Claude Opus 5.5 --- .../diffractive_imaging/test_ptychography.py | 132 ++++++++---------- 1 file changed, 61 insertions(+), 71 deletions(-) diff --git a/tests/diffractive_imaging/test_ptychography.py b/tests/diffractive_imaging/test_ptychography.py index c722e1d55..5c248cf8c 100644 --- a/tests/diffractive_imaging/test_ptychography.py +++ b/tests/diffractive_imaging/test_ptychography.py @@ -183,29 +183,38 @@ def single_probe_ptycho_model(ptycho_dataset, probe_array): return ptycho -class TestTargetResidency: - """Property + serialization behavior for the streaming-target knob.""" - - def test_default_is_device(self, ptycho_dataset): - assert ptycho_dataset.target_residency == "device" +@pytest.fixture +def small_ptycho_dataset(): + """A small random dataset model, enough for property and serialization checks.""" + array = np.random.default_rng(42).random((8, 8, 16, 16)).astype(np.float32) + dset = Dataset4dstem.from_array( + array, sampling=(1.0, 1.0, 0.05, 0.05), units=("A", "A", "A^-1", "A^-1") + ) + pdset = PtychographyDatasetRaster.from_dataset4dstem(dset) + pdset.preprocess(com_fit_function="constant", plot_rotation=False, plot_com=False) + return pdset - def test_setter_accepts_valid(self, ptycho_dataset): - ptycho_dataset.target_residency = "cpu" - assert ptycho_dataset.target_residency == "cpu" - ptycho_dataset.target_residency = "device" - assert ptycho_dataset.target_residency == "device" - @pytest.mark.parametrize("bad", ["gpu", "GPU", "CPU", "", "cuda", "Device"]) - def test_setter_rejects_invalid(self, ptycho_dataset, bad): - with pytest.raises(ValueError, match="target_residency"): - ptycho_dataset.target_residency = bad - # value should be unchanged after a rejected set - assert ptycho_dataset.target_residency == "device" +class TestTargetResidency: + """Property + serialization behavior for the streaming-target knob.""" - def test_save_load_roundtrip(self, ptycho_dataset, tmp_path): - ptycho_dataset.target_residency = "cpu" + def test_default_and_setter(self, small_ptycho_dataset): + assert small_ptycho_dataset.target_residency == "device" + small_ptycho_dataset.target_residency = "cpu" + assert small_ptycho_dataset.target_residency == "cpu" + small_ptycho_dataset.target_residency = "device" + assert small_ptycho_dataset.target_residency == "device" + + for bad in ["gpu", "GPU", "CPU", "", "cuda", "Device"]: + with pytest.raises(ValueError, match="target_residency"): + small_ptycho_dataset.target_residency = bad + # value should be unchanged after a rejected set + assert small_ptycho_dataset.target_residency == "device", bad + + def test_save_load_roundtrip(self, small_ptycho_dataset, tmp_path): + small_ptycho_dataset.target_residency = "cpu" path = tmp_path / "pdset.zip" - ptycho_dataset.save(str(path)) + small_ptycho_dataset.save(str(path)) reloaded = autoserialize_load(str(path)) assert reloaded.target_residency == "cpu" @@ -268,26 +277,20 @@ class TestAspectRatioRotatedFOV: ``_crop_rotate_obj_fov`` because the un-rotated FOV could not fit the object frame. """ - @pytest.mark.parametrize("transpose", [False, True]) - def test_positions_inside_object(self, complex_obj, probe_array, transpose): - ptycho = _build_aspect_ratio_ptycho( - complex_obj, probe_array, gpts=(32, 64), com_rotation=89, transpose=transpose - ) - obj_full = ptycho.dset._obj_shape_full_2d(ptycho.obj_padding_px) - pos = ptycho.dset.initial_scan_positions_px.cpu().detach().numpy() - # positions must sit strictly inside the object (no wrap / no edge pile-up) - assert pos[:, 0].min() >= 0 and pos[:, 1].min() >= 0 - assert pos[:, 0].max() < obj_full[-2] - assert pos[:, 1].max() < obj_full[-1] - - @pytest.mark.parametrize("transpose", [False, True]) - def test_obj_cropped_shape(self, complex_obj, probe_array, transpose): - ptycho = _build_aspect_ratio_ptycho( - complex_obj, probe_array, gpts=(32, 64), com_rotation=89, transpose=transpose - ) - # obj_cropped must not raise and returns the non-transposed display FOV shape - cropped = ptycho.obj_cropped - assert tuple(cropped.shape) == tuple(ptycho.obj_shape_crop) + def test_positions_inside_object_and_obj_cropped_shape(self, complex_obj, probe_array): + for transpose in (False, True): + ptycho = _build_aspect_ratio_ptycho( + complex_obj, probe_array, gpts=(32, 64), com_rotation=89, transpose=transpose + ) + obj_full = ptycho.dset._obj_shape_full_2d(ptycho.obj_padding_px) + pos = ptycho.dset.initial_scan_positions_px.cpu().detach().numpy() + # positions must sit strictly inside the object (no wrap / no edge pile-up) + assert pos[:, 0].min() >= 0 and pos[:, 1].min() >= 0, transpose + assert pos[:, 0].max() < obj_full[-2], transpose + assert pos[:, 1].max() < obj_full[-1], transpose + # obj_cropped must not raise and returns the non-transposed display FOV shape + cropped = ptycho.obj_cropped + assert tuple(cropped.shape) == tuple(ptycho.obj_shape_crop), transpose @pytest.mark.slow @@ -303,8 +306,10 @@ class TestPtychographySaveLoadRoundtrip: NUM_ITERS = 50 # enough headroom for the strict 0.3 threshold at lr=5e-3 - @pytest.fixture - def trained_ptycho(self, single_probe_ptycho_model): + def test_reconstruct_save_load_and_continue(self, single_probe_ptycho_model, tmp_path): + """Reconstruct, then reload: losses, positions and object match exactly, and continued + training extends the history without regressing -- confirms optimizer state and + parameter bindings survive the save/load roundtrip end-to-end.""" ptycho = single_probe_ptycho_model ptycho.reconstruct( num_iters=self.NUM_ITERS, @@ -316,47 +321,32 @@ def trained_ptycho(self, single_probe_ptycho_model): batch_size=N**2, device=config.get_device(), ) - return ptycho - - def test_iter_losses_preserved(self, trained_ptycho, tmp_path): - path = tmp_path / "ptycho.zip" - trained_ptycho.save(str(path), save_raw_data=True) - reloaded = autoserialize_load(str(path)) - np.testing.assert_array_equal(reloaded._iter_losses, trained_ptycho._iter_losses) - - def test_scan_positions_preserved(self, trained_ptycho, tmp_path): path = tmp_path / "ptycho.zip" - trained_ptycho.save(str(path), save_raw_data=True) + ptycho.save(str(path), save_raw_data=True) reloaded = autoserialize_load(str(path)) - original = trained_ptycho.dset.scan_positions_px.detach().cpu().numpy() - new = reloaded.dset.scan_positions_px.detach().cpu().numpy() - np.testing.assert_allclose(new, original, rtol=0, atol=0) - def test_object_preserved(self, trained_ptycho, tmp_path): - path = tmp_path / "ptycho.zip" - trained_ptycho.save(str(path), save_raw_data=True) - reloaded = autoserialize_load(str(path)) - original = trained_ptycho.obj_model._obj.detach().cpu().numpy() - new = reloaded.obj_model._obj.detach().cpu().numpy() - np.testing.assert_allclose(new, original, rtol=0, atol=0) + np.testing.assert_array_equal(reloaded._iter_losses, ptycho._iter_losses) + np.testing.assert_allclose( + reloaded.dset.scan_positions_px.detach().cpu().numpy(), + ptycho.dset.scan_positions_px.detach().cpu().numpy(), + rtol=0, + atol=0, + ) + np.testing.assert_allclose( + reloaded.obj_model._obj.detach().cpu().numpy(), + ptycho.obj_model._obj.detach().cpu().numpy(), + rtol=0, + atol=0, + ) - def test_loss_decreases_below_strict_threshold(self, trained_ptycho): - losses = trained_ptycho._iter_losses + losses = ptycho._iter_losses assert losses[-1] < 0.3 * losses[0], ( f"loss should drop below 30% of initial in {self.NUM_ITERS} iters: " f"initial={losses[0]:.3e}, final={losses[-1]:.3e}, " f"ratio={losses[-1] / losses[0]:.2f}" ) - def test_continue_training_after_reload(self, trained_ptycho, tmp_path): - """Reload a trained ptycho, continue training, and verify the loss keeps - decreasing — confirms optimizer state and parameter bindings survive the - save/load roundtrip end-to-end.""" - path = tmp_path / "ptycho.zip" - trained_ptycho.save(str(path), save_raw_data=True) - reloaded = autoserialize_load(str(path)) loss_after_reload = reloaded._iter_losses[-1] - n_continue = 10 reloaded.reconstruct( num_iters=n_continue, From ff17a33cfe8fb0cc04fee7c97556f13c1f5a8637 Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 15:03:19 -0700 Subject: [PATCH 09/19] Consolidate ptychography constraint tests and fix the pure-phase mask case test_constraints.py (24 -> 12 IDs): - test_mask_applied_once[pure_phase] used a constant phase, which the recentering zeroes, so the expected 0 held whether the FOV mask was applied once or twice. It now uses a zero-mean +-0.7 checkerboard: once gives |phase| 0.35, twice would give 0.175. - TestParseDict: the four successful parses are one test, the unknown and missing-name errors another. - TestTypoCatching: the dict and add_constraint paths are one test. - TestRoundtrip: the object/probe/dataset dataclass identity checks are a loop; the dict form is kept. - TestReconstructKwargs: test_dataclass_leaf_applied folded into test_mixed_dataclass_and_dict_leaves (the object dataclass leaf now also sets tv_weight_z=1.5 and checks it); the partial dict update is kept. - TestPurePhaseRealValued: the three storage-dtype tests and the hard-constraints-stay-real test are one loop; the phase-wrap warning presence/absence pair is another. Co-Authored-By: Claude Opus 5.5 --- tests/diffractive_imaging/test_constraints.py | 152 +++++++----------- 1 file changed, 62 insertions(+), 90 deletions(-) diff --git a/tests/diffractive_imaging/test_constraints.py b/tests/diffractive_imaging/test_constraints.py index eeb7e1a0b..ab275c334 100644 --- a/tests/diffractive_imaging/test_constraints.py +++ b/tests/diffractive_imaging/test_constraints.py @@ -61,25 +61,15 @@ def ptycho(): class TestParseDict: - def test_object_raster_by_name(self): + def test_parse_dict(self): c = PtychoObjConstraintParams.parse_dict({"name": "raster", "tv_weight_z": 5.0}) assert isinstance(c, PtychoObjConstraintParams.Raster) assert c.tv_weight_z == 5.0 assert c.positivity is True # default preserved - def test_object_inr_by_type(self): c = PtychoObjConstraintParams.parse_dict({"type": "inr"}) assert isinstance(c, PtychoObjConstraintParams.INR) - def test_object_unknown_raises(self): - with pytest.raises(ValueError, match="Unknown object constraint type"): - PtychoObjConstraintParams.parse_dict({"name": "nope"}) - - def test_object_missing_name_raises(self): - with pytest.raises(ValueError, match="Must provide either 'name' or 'type'"): - PtychoObjConstraintParams.parse_dict({"tv_weight_z": 5.0}) - - def test_probe_raster_with_fields(self): c = PtychoProbeConstraintParams.parse_dict( {"name": "raster", "center_probe": True, "tv_weight": 0.1} ) @@ -87,21 +77,24 @@ def test_probe_raster_with_fields(self): assert c.center_probe is True assert c.tv_weight == 0.1 - def test_dataset_raster_default(self): c = PtychoDatasetConstraintParams.parse_dict({"name": "raster"}) assert isinstance(c, PtychoDatasetConstraintParams.Raster) assert c.clip_scan_positions is True # default preserved + def test_parse_dict_errors(self): + with pytest.raises(ValueError, match="Unknown object constraint type"): + PtychoObjConstraintParams.parse_dict({"name": "nope"}) + with pytest.raises(ValueError, match="Must provide either 'name' or 'type'"): + PtychoObjConstraintParams.parse_dict({"tv_weight_z": 5.0}) + # --- Constraint typo catching ------------------------------------------------- class TestTypoCatching: - def test_setting_unknown_field_via_dict_raises(self, ptycho): + def test_unknown_constraint_key_raises(self, ptycho): with pytest.raises(KeyError, match="Invalid constraint key"): ptycho.obj_model.constraints = {"not_a_real_field": True} - - def test_add_constraint_unknown_key_raises(self, ptycho): with pytest.raises(KeyError, match="Invalid constraint key"): ptycho.obj_model.add_constraint("not_a_real_field", True) @@ -110,23 +103,29 @@ def test_add_constraint_unknown_key_raises(self, ptycho): class TestRoundtrip: - def test_obj_constraints_dataclass(self, ptycho): - obj_c = PtychoObjConstraintParams.Raster(tv_weight_z=2.5, identical_slices=True) - ptycho.constraints = {"object": obj_c} - assert ptycho.obj_model.constraints is obj_c + def test_constraint_dataclasses_are_stored(self, ptycho): + for key, constraints, model in [ + ( + "object", + PtychoObjConstraintParams.Raster(tv_weight_z=2.5, identical_slices=True), + lambda: ptycho.obj_model, + ), + ( + "probe", + PtychoProbeConstraintParams.Raster(center_probe=True, tv_weight=0.05), + lambda: ptycho.probe_model, + ), + ( + "dataset", + PtychoDatasetConstraintParams.Raster(descan_tv_weight=0.01), + lambda: ptycho.dset, + ), + ]: + ptycho.constraints = {key: constraints} + assert model().constraints is constraints, key assert ptycho.obj_model.constraints.tv_weight_z == 2.5 assert ptycho.obj_model.constraints.identical_slices is True - def test_probe_constraints_dataclass(self, ptycho): - probe_c = PtychoProbeConstraintParams.Raster(center_probe=True, tv_weight=0.05) - ptycho.constraints = {"probe": probe_c} - assert ptycho.probe_model.constraints is probe_c - - def test_dataset_constraints_dataclass(self, ptycho): - dset_c = PtychoDatasetConstraintParams.Raster(descan_tv_weight=0.01) - ptycho.constraints = {"dataset": dset_c} - assert ptycho.dset.constraints is dset_c - def test_dict_form_still_works(self, ptycho): """Backward compatibility: nested-dict form sets individual fields.""" ptycho.constraints = { @@ -142,20 +141,6 @@ def test_dict_form_still_works(self, ptycho): class TestReconstructKwargs: - def test_dataclass_leaf_applied(self, ptycho): - from quantem.core.ml import OptimizerParams - - obj_c = PtychoObjConstraintParams.Raster(tv_weight_z=1.5) - ptycho.reconstruct( - num_iters=1, - reset=True, - optimizer_params={"object": OptimizerParams.Adam(lr=1e-2)}, - constraints={"object": obj_c}, - batch_size=4, - device="cpu", - ) - assert ptycho.obj_model.constraints.tv_weight_z == 1.5 - def test_dict_leaf_partial_update(self, ptycho): from quantem.core.ml import OptimizerParams @@ -179,63 +164,47 @@ def test_mixed_dataclass_and_dict_leaves(self, ptycho): reset=True, optimizer_params={"object": OptimizerParams.Adam(lr=1e-2)}, constraints={ - "object": PtychoObjConstraintParams.Raster(tv_weight_xy=0.4), + "object": PtychoObjConstraintParams.Raster(tv_weight_xy=0.4, tv_weight_z=1.5), "probe": {"center_probe": True}, }, batch_size=4, device="cpu", ) assert ptycho.obj_model.constraints.tv_weight_xy == 0.4 + assert ptycho.obj_model.constraints.tv_weight_z == 1.5 assert ptycho.probe_model.constraints.center_probe is True # --- Real-valued pure_phase representation ----------------------------------- -class TestPurePhaseRealValued: - def test_pure_phase_pixelated_obj_is_real(self): - obj = ObjectPixelated.from_uniform(obj_type="pure_phase", num_slices=1) - obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1)) - assert not obj._obj.is_complex(), f"pure_phase _obj should be real, got {obj._obj.dtype}" - - def test_complex_pixelated_obj_is_complex(self): - obj = ObjectPixelated.from_uniform(obj_type="complex", num_slices=1) - obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1)) - assert obj._obj.is_complex() - - def test_potential_pixelated_obj_is_real(self): - obj = ObjectPixelated.from_uniform(obj_type="potential", num_slices=1) - obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1)) - assert not obj._obj.is_complex() - - def test_pure_phase_tv_emits_no_phase_warning(self): - obj = ObjectPixelated.from_uniform(obj_type="pure_phase", num_slices=1) - obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1)) - obj.constraints.tv_weight_xy = 0.1 - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - obj.get_tv_loss(obj._obj) - phase_warnings = [w for w in caught if "phase wrapping" in str(w.message)] - assert not phase_warnings, ( - f"pure_phase should not emit phase-wrap warning, got {phase_warnings}" - ) +def _initialized_obj(obj_type) -> ObjectPixelated: + obj = ObjectPixelated.from_uniform(obj_type=obj_type, num_slices=1) + obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1)) + return obj - def test_complex_tv_still_emits_phase_warning(self): - obj = ObjectPixelated.from_uniform(obj_type="complex", num_slices=1) - obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1)) - obj.constraints.tv_weight_xy = 0.1 - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - obj.get_tv_loss(obj._obj) - assert any("phase wrapping" in str(w.message) for w in caught), ( - "complex obj_type should still emit phase-wrap warning" - ) - def test_pure_phase_apply_hard_constraints_stays_real(self): - obj = ObjectPixelated.from_uniform(obj_type="pure_phase", num_slices=1) - obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1)) - out = obj.apply_hard_constraints(obj._obj) - assert not out.is_complex() +class TestPurePhaseRealValued: + def test_storage_dtype_per_obj_type(self): + for obj_type, is_complex in [ + ("pure_phase", False), + ("complex", True), + ("potential", False), + ]: + obj = _initialized_obj(obj_type) + assert obj._obj.is_complex() is is_complex, f"{obj_type}: got {obj._obj.dtype}" + if obj_type == "pure_phase": + assert not obj.apply_hard_constraints(obj._obj).is_complex() + + def test_tv_phase_warning_only_for_complex(self): + for obj_type, warns in [("pure_phase", False), ("complex", True)]: + obj = _initialized_obj(obj_type) + obj.constraints.tv_weight_xy = 0.1 + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + obj.get_tv_loss(obj._obj) + phase_warnings = [w for w in caught if "phase wrapping" in str(w.message)] + assert bool(phase_warnings) is warns, f"{obj_type}: {phase_warnings}" # --- FOV-mask single application --------------------------------------------- @@ -243,14 +212,17 @@ def test_pure_phase_apply_hard_constraints_stays_real(self): class TestFovMaskSingleApplication: def _make_obj(self, obj_type) -> ObjectPixelated: - obj = ObjectPixelated.from_uniform(obj_type=obj_type, num_slices=1) - obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1)) + obj = _initialized_obj(obj_type) obj.constraints.apply_fov_mask = True # Force a non-trivial _obj so masking is observable if obj_type == "complex": obj._obj = torch.nn.Parameter( torch.ones(1, 16, 16, dtype=torch.complex64) * (0.5 + 0.3j) ) + elif obj_type == "pure_phase": + # zero-mean +-0.7 checkerboard: the phase recentering leaves it unchanged + yy, xx = torch.meshgrid(torch.arange(16), torch.arange(16), indexing="ij") + obj._obj = torch.nn.Parameter(0.7 * (1 - 2 * ((yy + xx) % 2)).float()[None]) else: obj._obj = torch.nn.Parameter(torch.full((1, 16, 16), 0.7)) return obj @@ -270,8 +242,8 @@ def test_mask_applied_once(self, obj_type): # of the unmasked value; once it scales by 0.5. We compare to the # per-obj-type expected post-constraint value. if obj_type == "pure_phase": - # phase recentered to zero mean, then *= 0.5 mask - expected_mag = 0.0 # phase=constant -> recenter to 0 -> *0.5 = 0 + # zero-mean phase survives recentering, then *= 0.5 mask -> |phase| 0.35 + expected_mag = 0.35 elif obj_type == "potential": # positivity clamp keeps 0.7, * 0.5 -> 0.35 (one application) expected_mag = 0.35 From 830fb1437ba41755d1d57286576053445e2370f2 Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 15:04:11 -0700 Subject: [PATCH 10/19] Merge the multi-GPU spawn-lifecycle tests TestMultiGPUDeviceRestoration, TestMultiGPULossTracking, TestMultiGPULRTracking and TestMultiGPUOptimizerState each built a fresh model and ran the same reset / continue / reset spawn sequence (9 spawns). One test now runs a single sequence (4 spawns): the first call checks the device list, CUDA placement, loss and lr history and Adam state; then a continuation without device= (the device must persist), a continuation with device=, and a reset, checking that the loss and lr histories extend and reset. The single-GPU and DIP tests are unchanged. 6 -> 3 IDs. This test requires two CUDA devices, so it was linted but not run. Co-Authored-By: Claude Opus 5.5 --- tests/diffractive_imaging/test_multi_gpu.py | 65 ++++++--------------- 1 file changed, 18 insertions(+), 47 deletions(-) diff --git a/tests/diffractive_imaging/test_multi_gpu.py b/tests/diffractive_imaging/test_multi_gpu.py index 1fb9b4dc5..10c3aabf7 100644 --- a/tests/diffractive_imaging/test_multi_gpu.py +++ b/tests/diffractive_imaging/test_multi_gpu.py @@ -136,64 +136,21 @@ def test_int_device_persists(self): assert p.device == DEVICE_0 -class TestMultiGPUDeviceRestoration: - """device=[…] is stored and restored across spawn boundaries.""" +class TestMultiGPULifecycle: + """device=[…], _iter_losses, iter_lrs and Adam state across spawn boundaries.""" - def test_gpu_list_persists(self): + def test_device_losses_lrs_and_optimizer_state(self): p = _make_ptycho() p.reconstruct( num_iters=N_ITERS, reset=True, optimizer_params=_opt(), batch_size=16, device=GPU_IDS ) assert p.device == GPU_IDS assert p.obj_model._obj.device.type == "cuda" - - p.reconstruct(num_iters=N_ITERS, batch_size=16) - assert p.device == GPU_IDS - - -class TestMultiGPULossTracking: - """_iter_losses extends correctly across spawn(reset=False)/reset=True.""" - - def test_losses_length_lifecycle(self): - p = _make_ptycho() - p.reconstruct( - num_iters=N_ITERS, reset=True, optimizer_params=_opt(), batch_size=16, device=GPU_IDS - ) assert len(p._iter_losses) == N_ITERS - - p.reconstruct(num_iters=N_ITERS, reset=False, batch_size=16, device=GPU_IDS) - assert len(p._iter_losses) == 2 * N_ITERS, "continuation must not double-count" - - p.reconstruct(num_iters=N_ITERS, reset=True, batch_size=16, device=GPU_IDS) - assert len(p._iter_losses) == N_ITERS - - -class TestMultiGPULRTracking: - """iter_lrs extends correctly across spawn boundaries.""" - - def test_iter_lrs_lifecycle(self): - p = _make_ptycho() - p.reconstruct( - num_iters=N_ITERS, reset=True, optimizer_params=_opt(), batch_size=16, device=GPU_IDS - ) assert "object" in p.iter_lrs assert len(p.iter_lrs["object"]) == N_ITERS - p.reconstruct(num_iters=N_ITERS, reset=False, batch_size=16, device=GPU_IDS) - assert len(p.iter_lrs["object"]) == 2 * N_ITERS - - p.reconstruct(num_iters=N_ITERS, reset=True, batch_size=16, device=GPU_IDS) - assert len(p.iter_lrs["object"]) == N_ITERS - - -class TestMultiGPUOptimizerState: - """Adam state survives the save/restore around the spawn worker.""" - - def test_adam_state_restored_on_device(self): - p = _make_ptycho() - p.reconstruct( - num_iters=N_ITERS, reset=True, optimizer_params=_opt(), batch_size=16, device=GPU_IDS - ) + # Adam state survives the save/restore around the spawn worker obj_opt = p.optimizers.get("object") assert obj_opt is not None assert len(obj_opt.state) > 0 @@ -201,6 +158,20 @@ def test_adam_state_restored_on_device(self): assert "exp_avg" in first, "Adam moments missing" assert first["exp_avg"].device.type == "cuda" + # follow-up call without device= keeps the GPU list + p.reconstruct(num_iters=N_ITERS, batch_size=16) + assert p.device == GPU_IDS + assert len(p._iter_losses) == 2 * N_ITERS, "continuation must not double-count" + assert len(p.iter_lrs["object"]) == 2 * N_ITERS + + p.reconstruct(num_iters=N_ITERS, reset=False, batch_size=16, device=GPU_IDS) + assert len(p._iter_losses) == 3 * N_ITERS, "continuation must not double-count" + assert len(p.iter_lrs["object"]) == 3 * N_ITERS + + p.reconstruct(num_iters=N_ITERS, reset=True, batch_size=16, device=GPU_IDS) + assert len(p._iter_losses) == N_ITERS + assert len(p.iter_lrs["object"]) == N_ITERS + class TestDIPMultiGPU: """DIP path mirrors the pixelated multi-GPU contract.""" From 3105aa07806d812b7dbbae17d98a941188178ca9 Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 14:53:32 -0700 Subject: [PATCH 11/19] Consolidate the DirectPtychography tests Removed (cannot fail, or implied by another test): - test_one_slice_is_the_single_plane: num_slices=1 is the default, so both calls were identical. - test_no_per_pixel_stack_is_kept, test_variance_loss_has_been_removed: tombstones for removed attributes. - test_astigmatism_breaks_the_radial_symmetry: covered by test_the_shift_fit_returns_the_twin_nearest_its_reference, which recovers C12 from the same shifts. - test_the_two_classes_agree: rms_gradient_loss lives on the base class, and prlx equality is pinned by the montage TestFourierEquivalence. - test_a_spec_built_anywhere_is_accepted_everywhere: implied by the `is` identity in test_every_import_path_is_the_same_class. - test_the_noise_variance (Wiener): asserted the stored value equals the one passed. - test_a_finer_scan_sampling_refines_the_grid: implied by test_a_finer_scan_sampling_recovers_it. - test_upsampling_preserves_the_field_of_view[1,2,3]: same base _obj_sampling as the montage TestVisualization sampling test. - TestDatasetVariants: the prlx residual minimum is covered by test_is_a_fraction_minimized_at_the_true_defocus. - test_upsampling_an_irregular_regrid_warns: the warning is required by test_upsampling_factor_degrades_on_an_irregular_scan and test_a_finer_scan_sampling_recovers_it; its lattice_rms check moved into test_no_unfolding_warning_without_an_irregular_regrid. - test_twisted_layers_separate[honeycomb_scene], test_odd_aberrations_..._sign[C23]: same code paths as the square and C21 cases; twisted_honeycomb_cell is dropped from conftest. Merged (same branch or same setup): - test_aliases_resolve[4], test_the_matched_filter_has_been_removed[2], test_unknown_kernel_raises -> test_kernel_names. - test_is_a_fraction[4], test_is_minimized_at_the_true_defocus[3] -> test_is_a_fraction_minimized_at_the_true_defocus; the two "is None before reconstructing" tests -> test_is_none_before_reconstructing. - test_grid_search, test_grid_search_with_the_rms_gradient, test_grid_search_accepts_a_callable, test_recorded_losses_are_the_requested_objective -> test_grid_search[residual, rms_gradient, callable]. - test_unknown_losses_are_rejected[3] -> one loop on the shared reconstruction. - Six round-trip tests (TestSerialization x2, TestFromDataset3d, TestComplexObject, TestSectioning, TestPartialCoherence) -> TestSerialization: one gridded reconstruction carrying a complex sectioned object, partial coherence and a stored fit, plus one ungridded reconstruction. - TestReconstructAllPermutations (4) and test_permutations_leave_out_only_icom -> one table. - test_lattice_positions_reproduce_from_dataset4d[4 kernels]: regridding does not depend on the kernel, so the pair is built once and the kernels looped. - The four no-warning tests of TestFromDataset3d -> two. - test_lattice_positions_unfold_through_the_ungridded_path -> an absolute check in test_upsampling_factor_degrades_on_an_irregular_scan, which computes the same image. - test_sections_are_single_plane_reconstructions[5 kernels] and test_a_complex_object_has_amplitude_sections -> one complex obf case: _reconstruct_sections is kernel-generic apart from the iCoM refusal, which keeps its own test. - TestSectioning::test_aliases[4], test_rejects_an_unknown_mode -> test_mode_names. - test_kernels_without_a_joint_solve_raise[3] and the identical ssb call in TestSectioning -> one test matching both halves of the message. - test_only_the_normal_equation_kernels_take_it[3], TestWienerKernel's and WDD's ARD refusals -> one case in TestArdPrior::test_rejects_bad_arguments (one branch of _validate_prior). - The empirical-probe reproductions (complex ssb/obf, joint obf, and the pure-phase cases of the montage file) -> TestEmpiricalProbe, sharing one analytic/empirical pair. - The two visualize tests, test_obj_is_the_complex_transmission and test_pure_phase_has_no_amplitude, test_no_aberrations_means_no_shift into test_pure_defocus_matches_the_analytic_shift. - Cheap parametrizations sharing one object become loops naming the case: test_one_mode_is_the_coherent_probe, test_a_double_precision_probe_..., test_one_plane, test_slice_thicknesses_are_validated, test_a_finite_source_leaves_obf_..., test_rejects_a_negative_spread. Co-Authored-By: Claude Opus 5.5 --- tests/diffractive_imaging/conftest.py | 22 - .../test_direct_ptychography.py | 919 ++++++------------ 2 files changed, 321 insertions(+), 620 deletions(-) diff --git a/tests/diffractive_imaging/conftest.py b/tests/diffractive_imaging/conftest.py index 22595eb4e..bda2cc9aa 100644 --- a/tests/diffractive_imaging/conftest.py +++ b/tests/diffractive_imaging/conftest.py @@ -784,28 +784,6 @@ def twisted_square_cell(): return shape, layers, float(np.degrees(twist)) -def twisted_honeycomb_cell(a: float = 2.46): - """Two honeycombs twisted by 21.79 degrees, the (1, 2) commensurate cell. - - Both layers are periodic in the orthogonal cell of their ``sqrt(7) a`` coincidence lattice, - ``sqrt(21) a x sqrt(7) a``, snapped to whole pixels (a 0.1-0.2% strain). - """ - a1, a2 = np.array([a, 0.0]), np.array([a / 2, a * np.sqrt(3) / 2]) - t, t_image = a1 + 2 * a2, 2 * a1 + a2 - first = -np.arctan2(t[1], t[0]) - second = first + np.arctan2(t[1], t[0]) - np.arctan2(t_image[1], t_image[0]) - side = np.sqrt(7) * a - box_true = (side, np.sqrt(3) * side) # built as (col, row), transposed below - shape = (int(round(box_true[1] / ML_SAMPLING)), int(round(box_true[0] / ML_SAMPLING))) - strain = np.asarray(shape[::-1]) * ML_SAMPLING / np.asarray(box_true) - motif = [(np.zeros(2), 1.0), ((a1 + a2) / 3, 1.0)] - layers = [] - for angle in (first, second): - sites, weights = _lattice_in_box((a1, a2), motif, angle, box_true) - layers.append(periodic_gaussians((sites * strain)[:, ::-1], shape, weights=weights)) - return shape, layers, float(np.degrees(np.arccos(13 / 14))) - - def layer_purity(slices, truths) -> tuple[float, float]: """``(purity, correlation)`` averaged over layers, each true layer scored at its own slice. diff --git a/tests/diffractive_imaging/test_direct_ptychography.py b/tests/diffractive_imaging/test_direct_ptychography.py index 41d3e23cf..0fbf0d21e 100644 --- a/tests/diffractive_imaging/test_direct_ptychography.py +++ b/tests/diffractive_imaging/test_direct_ptychography.py @@ -49,12 +49,10 @@ integer_shift_defocus, layer_purity, library_probe_fourier, - make_dataset4d, measured_ctf, ml_kwargs, scan_positions_px, thin_layer_4dstem, - twisted_honeycomb_cell, twisted_square_cell, white_noise_object_2D, ) @@ -167,42 +165,19 @@ def test_recovers_the_band_limited_object(self, recon, kernel): assert np.isfinite(recon.obj).all() assert abs(correlation(recon.obj, band_limited_phase())) > 0.7 - @pytest.mark.parametrize("kernel", DECONVOLUTION_KERNELS) - def test_aliases_resolve(self, recon, kernel): + def test_kernel_names(self, recon): aliases = { - "ssb": "single-sideband", - "obf": "optimum-bright-field", - "prlx": "tilt-corrected-bright-field", - "icom": "center-of-mass", + "single-sideband": "ssb", + "optimum-bright-field": "obf", + "tilt-corrected-bright-field": "prlx", + "center-of-mass": "icom", + "Wigner-Distribution-Deconvolution": "wdd", } - by_short = recon.reconstruct(deconvolution_kernel=kernel, verbose=False).obj.copy() - by_alias = recon.reconstruct(deconvolution_kernel=aliases[kernel], verbose=False).obj - - assert np.array_equal(by_short, by_alias) + for alias, kernel in aliases.items(): + assert recon._normalize_kernel_name(alias) == kernel, f"{alias=}" - @pytest.mark.parametrize("kernel", ["mf", "matched-filter"]) - def test_the_matched_filter_has_been_removed(self, recon, kernel): with pytest.raises(ValueError, match="matched-filter kernel has been removed"): - recon.reconstruct(deconvolution_kernel=kernel, verbose=False) - - @pytest.mark.parametrize("upsampling_factor", [1, 2, 3]) - def test_upsampling_preserves_the_field_of_view(self, recon, upsampling_factor): - recon.reconstruct( - deconvolution_kernel="ssb", upsampling_factor=upsampling_factor, verbose=False - ) - - assert recon.obj.shape == (N * upsampling_factor, N * upsampling_factor) - assert recon.obj.shape[0] * recon._obj_sampling[0] == pytest.approx(N * SCAN_SAMPLING) - - def test_no_per_pixel_stack_is_kept(self, recon): - """The kernels are summed over bright-field pixels in Fourier space, never stored.""" - recon.reconstruct(deconvolution_kernel="ssb", verbose=False) - - assert recon.corrected_bf.shape == (N, N) - assert not hasattr(recon, "corrected_stack") - assert not hasattr(recon, "variance_loss") - - def test_unknown_kernel_raises(self, recon): + recon.reconstruct(deconvolution_kernel="matched-filter", verbose=False) with pytest.raises(ValueError, match="Unknown deconvolution kernel"): recon.reconstruct(deconvolution_kernel="maximum-entropy", verbose=False) @@ -222,11 +197,7 @@ def test_pure_defocus_matches_the_analytic_shift(self, recon): ) assert torch.allclose(shifts, expected, rtol=1e-5, atol=1e-6) - - def test_no_aberrations_means_no_shift(self, recon): - shifts = recon._return_lateral_shifts(0.0, {}, recon.bf_mask) - - assert torch.count_nonzero(shifts) == 0 + assert torch.count_nonzero(recon._return_lateral_shifts(0.0, {}, recon.bf_mask)) == 0 def test_rotation_rotates_the_shifts(self, recon): """`_passively_rotate_grid` sends (kx, ky) -> (-ky, kx) at 90 degrees.""" @@ -237,14 +208,6 @@ def test_rotation_rotates_the_shifts(self, recon): expected = torch.stack((-unrotated[:, 1], unrotated[:, 0]), dim=-1) assert torch.allclose(rotated, expected, atol=1e-5) - def test_astigmatism_breaks_the_radial_symmetry(self, recon): - radial = recon._return_lateral_shifts(0.0, {"C10": TRUE_C10}, recon.bf_mask) - astigmatic = recon._return_lateral_shifts( - 0.0, {"C10": TRUE_C10, "C12": 0.3 * TRUE_C10}, recon.bf_mask - ) - - assert not torch.allclose(radial, astigmatic) - class TestBrightFieldSubsets: def test_checkerboard_halves_are_additive(self, recon): @@ -335,26 +298,25 @@ class TestResidualLoss: """The least-squares residual of each kernel's forward model, as a hyperparameter objective.""" def test_is_none_before_reconstructing(self, dataset4d): - assert _build(dataset4d).residual_loss() is None - - @pytest.mark.parametrize("kernel", DECONVOLUTION_KERNELS) - def test_is_a_fraction(self, recon, kernel): - recon.reconstruct(deconvolution_kernel=kernel, verbose=False) - - assert 0.0 < recon.residual_loss() < 1.0 - - @pytest.mark.parametrize("kernel", ("ssb", "obf", "prlx")) - def test_is_minimized_at_the_true_defocus(self, recon, kernel): - losses = {} - for scale in (0.5, 0.8, 1.0, 1.2, 1.5): - recon.reconstruct( - deconvolution_kernel=kernel, - override_aberration_coefs={"C10": scale * TRUE_C10}, - verbose=False, - ) - losses[scale] = recon.residual_loss() + fresh = _build(dataset4d) + assert fresh.residual_loss() is None + assert fresh.rms_gradient_loss() is None + + def test_is_a_fraction_minimized_at_the_true_defocus(self, recon): + for kernel in DECONVOLUTION_KERNELS: + # iCoM reads no aberrations, so it has no defocus to prefer + scales = (1.0,) if kernel == "icom" else (0.5, 0.8, 1.0, 1.2, 1.5) + losses = {} + for scale in scales: + recon.reconstruct( + deconvolution_kernel=kernel, + override_aberration_coefs={"C10": scale * TRUE_C10}, + verbose=False, + ) + losses[scale] = recon.residual_loss() - assert min(losses, key=losses.get) == pytest.approx(1.0) + assert 0.0 < losses[1.0] < 1.0, f"{kernel=}" + assert min(losses, key=losses.get) == pytest.approx(1.0), f"{kernel=} {losses=}" def test_the_weak_object_kernels_share_one_model(self, recon): """ssb, obf and mf differ in how they normalize, not in what they model.""" @@ -382,21 +344,27 @@ def test_parallax_residual_is_the_spread_of_the_corrected_images(self, recon): expected = float((deviation * mask).sum() / (energy * mask).sum()) assert recon.residual_loss() == pytest.approx(expected, rel=1e-4) - def test_variance_loss_has_been_removed(self, recon): - recon.reconstruct(deconvolution_kernel="prlx", verbose=False) - with pytest.raises(ValueError, match="loss='residual'"): - recon._return_loss_value("variance") - class TestHyperparameterFitting: """All four routines must recover a seeded defocus from a blind start.""" - def test_grid_search(self, dataset4d): + @pytest.mark.parametrize( + "loss, sign", + [ + ("residual", 1), + ("rms_gradient", -1), + # anything that scores a reconstruction can drive a search + (lambda r: -float(np.std(r.obj)), -1), + ], + ids=["residual", "rms_gradient", "callable"], + ) + def test_grid_search(self, dataset4d, loss, sign): recon = _build(dataset4d, aberration_coefs={}) recon.grid_search_hyperparameters( aberration_coefs={ "C10": OptimizationParameter(low=0.4 * TRUE_C10, high=1.6 * TRUE_C10, n_points=7) }, + loss=loss, deconvolution_kernel="prlx", verbose=False, ) @@ -404,6 +372,8 @@ def test_grid_search(self, dataset4d): step = (1.6 - 0.4) * TRUE_C10 / 6 assert abs(fitted - TRUE_C10) <= step + # the recorded losses are the objective actually optimized, not the default + assert all(np.sign(value) == sign for _, value in recon._grid_search_results) def test_cross_correlation(self, dataset4d): recon = _build(dataset4d, aberration_coefs={}) @@ -507,13 +477,8 @@ def test_fitting_leaves_the_initial_state_intact(self, dataset4d): class TestLossFunctions: """The searches accept an objective by name or as a callable, and minimize it.""" - def test_rms_gradient_is_none_before_reconstruct(self, dataset4d): - assert _build(dataset4d).rms_gradient_loss() is None - - def test_rms_gradient_peaks_at_the_true_defocus(self, dataset4d): + def test_rms_gradient_peaks_at_the_true_defocus(self, recon): """It is a sharpness metric, so the correct deconvolution must maximize it.""" - recon = _build(dataset4d) - losses = {} for scale in (0.5, 1.0, 1.5): recon.reconstruct( @@ -527,7 +492,7 @@ def test_rms_gradient_peaks_at_the_true_defocus(self, dataset4d): assert losses[1.0] < losses[0.5] assert losses[1.0] < losses[1.5] - def test_rms_gradient_follows_the_object_sampling(self, dataset4d): + def test_rms_gradient_follows_the_object_sampling(self, recon): """Per Angstrom, not per pixel, so upsampling does not rescale it. ``band_limit=False`` keeps the tiled replicas, so the two canvases hold the same @@ -535,7 +500,6 @@ def test_rms_gradient_follows_the_object_sampling(self, dataset4d): band limit on, upsampling a parallax reconstruction removes those replicas and the image amplitude genuinely changes, so this would be measuring that instead. """ - recon = _build(dataset4d) kwargs = dict(deconvolution_kernel="prlx", band_limit=False, verbose=False) recon.reconstruct(upsampling_factor=1, **kwargs) coarse = recon.rms_gradient_loss() @@ -546,70 +510,11 @@ def test_rms_gradient_follows_the_object_sampling(self, dataset4d): # is the same quantity; a per-pixel metric would differ by the sampling ratio assert fine == pytest.approx(coarse, rel=0.35) - def test_the_two_classes_agree(self, dataset4d): - """`prlx` is the same operator in both, so the objective must match.""" - from quantem.diffractive_imaging import DirectPtychographyMontage - - kwargs = dict(direct_ptycho_kwargs(TRUE_C10), edge_blend_pixels=0) - fourier = DirectPtychography.from_dataset4d(dataset4d, **kwargs) - montage = DirectPtychographyMontage.from_dataset4d(dataset4d, boundary="wrap", **kwargs) - for recon in (fourier, montage): - recon.reconstruct(deconvolution_kernel="prlx", verbose=False) - - assert montage.rms_gradient_loss() == pytest.approx(fourier.rms_gradient_loss(), rel=1e-5) - - def test_grid_search_with_the_rms_gradient(self, dataset4d): - recon = _build(dataset4d, aberration_coefs={}) - recon.grid_search_hyperparameters( - aberration_coefs={ - "C10": OptimizationParameter(low=0.4 * TRUE_C10, high=1.6 * TRUE_C10, n_points=7) - }, - loss="rms_gradient", - deconvolution_kernel="prlx", - verbose=False, - ) - fitted = recon.hyperparameter_state.optimized_aberrations["C10"] - - step = (1.6 - 0.4) * TRUE_C10 / 6 - assert abs(fitted - TRUE_C10) <= step - - def test_grid_search_accepts_a_callable(self, dataset4d): - """Anything that scores a reconstruction can drive a search.""" - recon = _build(dataset4d, aberration_coefs={}) - recon.grid_search_hyperparameters( - aberration_coefs={ - "C10": OptimizationParameter(low=0.4 * TRUE_C10, high=1.6 * TRUE_C10, n_points=7) - }, - loss=lambda r: -float(np.std(r.obj)), - deconvolution_kernel="prlx", - verbose=False, - ) - fitted = recon.hyperparameter_state.optimized_aberrations["C10"] - - step = (1.6 - 0.4) * TRUE_C10 / 6 - assert abs(fitted - TRUE_C10) <= step - - def test_recorded_losses_are_the_requested_objective(self, dataset4d): - """`_grid_search_results` must hold the loss actually optimized, not the default.""" - recon = _build(dataset4d, aberration_coefs={}) - recon.grid_search_hyperparameters( - aberration_coefs={ - "C10": OptimizationParameter(low=0.4 * TRUE_C10, high=1.6 * TRUE_C10, n_points=3) - }, - loss="rms_gradient", - deconvolution_kernel="prlx", - verbose=False, - ) - - # the RMS gradient loss is negative; the residual loss is positive - assert all(loss < 0 for _, loss in recon._grid_search_results) - - @pytest.mark.parametrize("bad", ["sharpness", 3, None]) - def test_unknown_losses_are_rejected(self, dataset4d, bad): - recon = _build(dataset4d) + def test_unknown_losses_are_rejected(self, recon): recon.reconstruct(deconvolution_kernel="prlx", verbose=False) - with pytest.raises(ValueError, match="must be a callable or one of"): - recon._return_loss_value(bad) + for bad in ("sharpness", 3, None): + with pytest.raises(ValueError, match=f"must be a callable or one of .*got {bad!r}"): + recon._return_loss_value(bad) class TestHyperparameterState: @@ -640,43 +545,54 @@ def test_rejects_unknown_aberrations(self): class TestSerialization: - def test_round_trip_preserves_the_reconstruction(self, dataset4d, tmp_path): - recon = _build(dataset4d) - recon.reconstruct(deconvolution_kernel="ssb", verbose=False) + def test_round_trip_preserves_the_reconstruction(self, dataset4d, tmp_path, capsys): + from quantem.diffractive_imaging.direct_ptychography_base import HyperparameterState + + recon = _build(dataset4d, focal_spread=30.0, source_size=0.2) + state = recon.hyperparameter_state + state.optimized_aberrations = {"C10": TRUE_C10, "C12": 3.0} + state.optimized_rotation_angle = 1.5 + state.optimized_keys = {"C10", "C12", "rotation_angle"} + state.optimized_slice_thicknesses = [6.0] + kwargs = dict(deconvolution_kernel="obf", obj_type="complex", **TestSectioning.SECTIONS) + recon.reconstruct(**kwargs) before = recon.obj.copy() path = str(tmp_path / "direct.zip") recon.save(path, mode="o") + # nothing fell back to pickling, which the serializer announces on stdout + assert "falling back" not in capsys.readouterr().out restored = load(path) assert isinstance(restored, DirectPtychography) assert np.array_equal(restored.obj, before) + assert restored.obj_type == "complex" + assert restored.slice_mode == "sectioned" + assert np.allclose(restored.section_residuals, recon.section_residuals) assert restored.num_bf == recon.num_bf assert restored.gpts == recon.gpts + assert restored.focal_spread == 30.0 + assert restored.source_size == 0.2 - restored.reconstruct(deconvolution_kernel="ssb", verbose=False) - assert np.allclose(restored.obj, before) - - def test_the_hyperparameter_state_saves_natively(self, dataset4d, tmp_path, capsys): - from quantem.diffractive_imaging.direct_ptychography_base import HyperparameterState + restored_state = restored.hyperparameter_state + assert isinstance(restored_state, HyperparameterState) + assert restored_state.current_aberrations() == state.current_aberrations() + assert restored_state.current_rotation_angle() == 1.5 + assert restored_state.optimized_keys == {"C10", "C12", "rotation_angle"} + assert restored_state.current_slice_thicknesses() == [6.0] - recon = _build(dataset4d) - recon.hyperparameter_state.optimized_aberrations = {"C10": 42.0, "C12": 3.0} - recon.hyperparameter_state.optimized_rotation_angle = 1.5 - recon.hyperparameter_state.optimized_keys = {"C10", "C12", "rotation_angle"} - recon.hyperparameter_state.optimized_slice_thicknesses = [6.0] + restored.reconstruct(**kwargs) + assert np.allclose(restored.obj, before) - path = str(tmp_path / "direct.zip") - recon.save(path, mode="o") - # nothing fell back to pickling, which the serializer announces on stdout - assert "falling back" not in capsys.readouterr().out + def test_an_ungridded_reconstruction_round_trips(self, dataset4d, tmp_path): + reconstruction = TestFromDataset3d._build( + TestFromDataset3d._dataset3d(dataset4d), scan_positions_px() * SCAN_SAMPLING + ) + reconstruction.reconstruct(deconvolution_kernel="prlx", verbose=False) + path = tmp_path / "ungridded.zip" + reconstruction.save(path, mode="o") - state = load(path).hyperparameter_state - assert isinstance(state, HyperparameterState) - assert state.current_aberrations() == recon.hyperparameter_state.current_aberrations() - assert state.current_rotation_angle() == 1.5 - assert state.optimized_keys == {"C10", "C12", "rotation_angle"} - assert state.current_slice_thicknesses() == [6.0] + assert np.allclose(load(path).obj, reconstruction.obj) class TestRotationSensitivity: @@ -709,23 +625,20 @@ def test_detector_rotation_is_estimated_when_not_supplied(self, dataset4d): class TestReconstructAllPermutations: - def test_returns_one_image_per_kernel(self, recon): - images = recon._reconstruct_all_permutations(verbose=False) - - assert len(images) == len(DECONVOLUTION_KERNELS) + 2 # with "wiener" and "wdd" - assert all(image.shape == (N, N) for image in images) - assert all(np.isfinite(image).all() for image in images) - - @pytest.mark.parametrize( - "kwargs, count", - [ - (dict(obj_type="complex"), 5), - (dict(num_slices=2, slice_mode="joint", slice_thicknesses=5.0), 2), - (dict(prior="ard"), 1), - ], - ) - def test_leaves_out_the_kernels_that_cannot_take_the_settings(self, recon, kwargs, count): - assert len(recon._reconstruct_all_permutations(verbose=False, **kwargs)) == count + def test_leaves_out_the_kernels_that_cannot_take_the_settings(self, recon): + cases = [ + (dict(), 6, (N, N)), + (dict(obj_type="complex"), 5, (N, N)), + (dict(num_slices=2, slice_mode="joint", slice_thicknesses=5.0), 2, (2, N, N)), + # sectioned, the default for several slices: all but iCoM + (dict(num_slices=2, slice_thicknesses=50.0), 5, (2, N, N)), + (dict(prior="ard"), 1, (N, N)), + ] + for kwargs, count, shape in cases: + images = recon._reconstruct_all_permutations(verbose=False, **kwargs) + assert len(images) == count, f"{kwargs=}" + assert all(image.shape == shape for image in images), f"{kwargs=}" + assert all(np.isfinite(image).all() for image in images), f"{kwargs=}" class TestNormalizationOrder: @@ -751,28 +664,6 @@ def test_edge_blending_tapers_the_stack(self, dataset4d): assert (edge_blended - 1).abs().mean() < (edge_sharp - 1).abs().mean() -class TestDatasetVariants: - def test_reconstruction_tracks_the_simulated_defocus(self): - """Data simulated at a different defocus must prefer that defocus.""" - other_C10 = 2 * TRUE_C10 - dataset = make_dataset4d(defocus=other_C10) - recon = DirectPtychography.from_dataset4d( - dataset, edge_blend_pixels=0, **direct_ptycho_kwargs(other_C10) - ) - - losses = {} - for scale in (0.5, 1.0, 1.5): - recon.reconstruct( - deconvolution_kernel="prlx", - override_aberration_coefs={"C10": scale * other_C10}, - parallax_flip_phase=False, - verbose=False, - ) - losses[scale] = recon.residual_loss() - - assert min(losses, key=losses.get) == pytest.approx(1.0) - - class TestFromDataset3d: """Ungridded scans, by resampling the bright-field stack onto a grid first. @@ -804,18 +695,18 @@ def _build(dataset3d, positions, **overrides): kwargs.update(overrides) return DirectPtychography.from_dataset3d(dataset3d, positions, **kwargs) - @pytest.mark.parametrize("kernel", DECONVOLUTION_KERNELS) - def test_lattice_positions_reproduce_from_dataset4d(self, dataset4d, kernel): + def test_lattice_positions_reproduce_from_dataset4d(self, dataset4d): """On a lattice the splat is an identity map, so this must be exact.""" gridded = _build(dataset4d) ungridded = self._build(self._dataset3d(dataset4d), scan_positions_px() * SCAN_SAMPLING) assert ungridded.scan_gpts == gridded.scan_gpts - assert np.allclose( - ungridded.reconstruct(deconvolution_kernel=kernel, verbose=False).obj, - gridded.reconstruct(deconvolution_kernel=kernel, verbose=False).obj, - atol=1e-6, - ) + for kernel in DECONVOLUTION_KERNELS: + assert np.allclose( + ungridded.reconstruct(deconvolution_kernel=kernel, verbose=False).obj, + gridded.reconstruct(deconvolution_kernel=kernel, verbose=False).obj, + atol=1e-6, + ), f"{kernel=}" def test_position_axis_order_is_row_col(self, dataset4d): """Swapping the columns must transpose the regridded stack, not scramble it. @@ -872,24 +763,18 @@ def test_warns_when_positions_are_clustered(self, dataset4d): verbose=False, ) - def test_a_finer_grid_does_not_warn_about_its_gaps(self, dataset4d): + def test_a_covered_or_finer_grid_does_not_warn(self, dataset4d): """Fewer positions than pixels is a comb, not missing data.""" positions = scan_positions_px() * SCAN_SAMPLING - with warnings.catch_warnings(): - warnings.simplefilter("error", UserWarning) - self._build( - self._dataset3d(dataset4d), - positions, - scan_sampling=(SCAN_SAMPLING / 2, SCAN_SAMPLING / 2), - ) - - def test_no_warning_on_a_fully_covered_grid(self, dataset4d): - positions = scan_positions_px() * SCAN_SAMPLING - - with warnings.catch_warnings(): - warnings.simplefilter("error", UserWarning) - self._build(self._dataset3d(dataset4d), positions) + for scan_sampling in (SCAN_SAMPLING, SCAN_SAMPLING / 2): + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + self._build( + self._dataset3d(dataset4d), + positions, + scan_sampling=(scan_sampling, scan_sampling), + ) def test_scan_gpts_pads_without_rescaling_the_positions(self, dataset4d): """`scan_sampling` stays authoritative: a bigger grid pads, it does not rescale.""" @@ -948,16 +833,6 @@ def test_rejects_linear_normalization(self, dataset4d): normalization_order=1, ) - def test_survives_a_serialization_round_trip(self, dataset4d, tmp_path): - reconstruction = self._build( - self._dataset3d(dataset4d), scan_positions_px() * SCAN_SAMPLING - ) - reconstruction.reconstruct(deconvolution_kernel="prlx", verbose=False) - path = tmp_path / "ungridded.zip" - reconstruction.save(path, mode="o") - - assert np.allclose(load(path).obj, reconstruction.obj) - @staticmethod def _disk_masked(dataset4d, radius=14.0): """A non-rectangular scan subset -- the case that exposes hole handling.""" @@ -1035,18 +910,8 @@ def test_rejects_an_unknown_hole_fill(self, dataset4d): with pytest.raises(ValueError, match="`hole_fill` must be"): self._build(dataset3d, positions, hole_fill="interpolate") - def test_upsampling_an_irregular_regrid_warns(self, dataset4d): - """Regridding discards the sub-pixel positions upsampling needs to unfold.""" - rng = np.random.default_rng(3) - positions = scan_positions_px() * SCAN_SAMPLING - positions = positions + rng.uniform(-0.4, 0.4, positions.shape) * SCAN_SAMPLING - recon = self._build(self._dataset3d(dataset4d), positions) - - assert recon._regrid_info["lattice_rms_px"] > 0.1 - with pytest.warns(UserWarning, match="cannot unfold"): - recon.reconstruct(deconvolution_kernel="prlx", upsampling_factor=2, verbose=False) - - def test_no_unfolding_warning_on_a_lattice_or_without_upsampling(self, dataset4d): + def test_no_unfolding_warning_without_an_irregular_regrid(self, dataset4d): + """Only upsampling an irregular regrid warns; `from_dataset4d` has no regrid info.""" dataset3d = self._dataset3d(dataset4d) on_lattice = self._build(dataset3d, scan_positions_px() * SCAN_SAMPLING) @@ -1059,12 +924,11 @@ def test_no_unfolding_warning_on_a_lattice_or_without_upsampling(self, dataset4d positions = scan_positions_px() * SCAN_SAMPLING positions = positions + rng.uniform(-0.4, 0.4, positions.shape) * SCAN_SAMPLING irregular = self._build(dataset3d, positions) + assert irregular._regrid_info["lattice_rms_px"] > 0.1 with warnings.catch_warnings(): warnings.simplefilter("error", UserWarning) irregular.reconstruct(deconvolution_kernel="prlx", verbose=False) - def test_gridded_reconstructions_never_warn_about_unfolding(self, dataset4d): - """`from_dataset4d` has no regrid info, so the check must be a no-op there.""" gridded = _build(dataset4d) with warnings.catch_warnings(): warnings.simplefilter("error", UserWarning) @@ -1118,22 +982,11 @@ def test_gridded_upsampling_unfolds_the_ctf(self, ctf_scene): assert up_in > 0.95 assert up_ext > 0.95 # genuinely extended, not replicated - def test_lattice_positions_unfold_through_the_ungridded_path(self, ctf_scene): - _, _, _, ctf_full, _ = ctf_scene - recon = self._ungridded(ctf_scene, ctf_raster_positions()) - - obj = recon.reconstruct( - deconvolution_kernel="prlx", upsampling_factor=2, verbose=False - ).obj - - assert self._scores(obj, ctf_full, 2)[1] > 0.95 - def test_upsampling_factor_degrades_on_an_irregular_scan(self, ctf_scene): """Binning first discards the sub-pixel positions upsampling needs. - Compared against the same call on a lattice rather than an absolute threshold: what - matters is that scattering the positions costs the extension band, not where the - number happens to land. + Lattice positions unfold through the ungridded path; scattering them costs the + extension band. """ _, _, _, ctf_full, _ = ctf_scene @@ -1148,6 +1001,7 @@ def test_upsampling_factor_degrades_on_an_irregular_scan(self, ctf_scene): deconvolution_kernel="prlx", upsampling_factor=2, verbose=False ).obj + assert self._scores(on_lattice, ctf_full, 2)[1] > 0.95 assert ( self._scores(scattered, ctf_full, 2)[1] < self._scores(on_lattice, ctf_full, 2)[1] - 0.1 @@ -1174,15 +1028,6 @@ def test_a_finer_scan_sampling_recovers_it(self, ctf_scene): assert fine_corr > 0.99 assert fine_corr > coarse_corr - def test_a_finer_scan_sampling_refines_the_grid(self, ctf_scene): - plain = self._ungridded(ctf_scene, ctf_raster_positions()) - finer = self._ungridded( - ctf_scene, ctf_raster_positions(), scan_sampling=(CTF_SCAN_SAMPLING / 2,) * 2 - ) - - assert finer.scan_sampling[0] == pytest.approx(plain.scan_sampling[0] / 2) - assert finer.scan_gpts[0] >= 2 * plain.scan_gpts[0] - 2 - def test_nearest_beats_bilinear_on_a_finer_grid(self, ctf_scene): """Bilinear smears each measurement over four pixels and blurs out the detail.""" _, _, _, ctf_full, _ = ctf_scene @@ -1264,15 +1109,6 @@ def test_every_import_path_is_the_same_class(self): assert package is direct is base is iterative is canonical assert canonical.__module__ == "quantem.diffractive_imaging.ptycho_utils" - def test_a_spec_built_anywhere_is_accepted_everywhere(self): - """The failure the duplicate caused: isinstance across the two searches.""" - from quantem.diffractive_imaging import OptimizationParameter - from quantem.diffractive_imaging.optimize_hyperparameters import ( - OptimizationParameter as iterative, - ) - - assert isinstance(OptimizationParameter(0.0, 1.0), iterative) - # --------------------------------------------------------------------------- # complex-valued objects @@ -1404,6 +1240,11 @@ def test_icom_amplitude_is_the_bright_field_sum(self, amplitude_scene): def test_obj_is_the_complex_transmission(self, amplitude_scene): recon, _ = amplitude_scene + recon.reconstruct(deconvolution_kernel="obf", verbose=False) + assert recon.corrected_amplitude is None + assert not np.iscomplexobj(recon.obj) + assert recon.obj_type == "pure_phase" + recon.reconstruct(deconvolution_kernel="obf", obj_type="complex", verbose=False) obj = recon.obj @@ -1445,49 +1286,24 @@ def test_unknown_obj_type_raises(self, recon): with pytest.raises(ValueError, match="obj_type"): recon.reconstruct(deconvolution_kernel="obf", obj_type="potential", verbose=False) - def test_pure_phase_has_no_amplitude(self, recon): - recon.reconstruct(deconvolution_kernel="obf", verbose=False) - - assert recon.corrected_amplitude is None - assert not np.iscomplexobj(recon.obj) - assert recon.obj_type == "pure_phase" - - def test_visualize_shows_both_channels(self, amplitude_scene): + def test_visualize_shows_both_channels_of_every_slice(self, amplitude_scene): import matplotlib matplotlib.use("Agg") recon, _ = amplitude_scene recon.reconstruct(deconvolution_kernel="obf", obj_type="complex", verbose=False) - fig, axs = recon.visualize(return_fig=True) - + _, axs = recon.visualize(return_fig=True) assert np.size(axs) == 4 - def test_round_trip_preserves_both_channels(self, amplitude_scene, tmp_path): - recon, _ = amplitude_scene - recon.reconstruct(deconvolution_kernel="obf", obj_type="complex", verbose=False) - before = recon.obj.copy() - - path = str(tmp_path / "complex.zip") - recon.save(path, mode="o") - restored = load(path) - - assert np.array_equal(restored.obj, before) - assert restored.obj_type == "complex" - - @pytest.mark.parametrize("kernel", ["ssb", "obf"]) - def test_an_empirical_probe_reproduces_the_analytic_one(self, dataset4d, kernel): - from quantem.diffractive_imaging.complex_probe import FourierProbe - - analytic = _build(dataset4d) - psi = analytic_probe_array(analytic, {"C10": TRUE_C10}) - probe = FourierProbe.from_array( - psi, analytic.reciprocal_sampling, analytic.wavelength, normalize=False + recon.reconstruct( + deconvolution_kernel="obf", + num_slices=2, + slice_mode="joint", + slice_thicknesses=6.0, + verbose=False, ) - empirical = _build(dataset4d, fourier_probe=probe) - - kwargs = dict(deconvolution_kernel=kernel, obj_type="complex", verbose=False) - expected = analytic.reconstruct(**kwargs).obj.copy() - assert np.allclose(empirical.reconstruct(**kwargs).obj, expected, rtol=1e-5, atol=1e-6) + _, axs = recon.visualize(return_fig=True) + assert np.shape(axs) == (2, 2) # --------------------------------------------------------------------------- @@ -1495,49 +1311,37 @@ def test_an_empirical_probe_reproduces_the_analytic_one(self, dataset4d, kernel) # --------------------------------------------------------------------------- -def _twisted_scene(kind): +@pytest.fixture(scope="module") +def square_scene(): """``(reconstruction, true layers, spacing)`` for two twisted layers about the C10 plane.""" - if kind == "square": - _, layers, _ = twisted_square_cell() - layers, spacing = [0.3 * layer for layer in layers], 6.0 - else: - _, layers, _ = twisted_honeycomb_cell() - layers, spacing = [0.1 * layer for layer in layers], 3.35 + _, layers, _ = twisted_square_cell() + layers, spacing = [0.3 * layer for layer in layers], 6.0 depths = [-spacing / 2, spacing / 2] dataset = thin_layer_4dstem([1j * layer for layer in layers], depths, ML_ABERRATIONS) recon = DirectPtychography.from_dataset4d(dataset, **ml_kwargs(ML_ABERRATIONS)) return recon, layers, spacing -@pytest.fixture(scope="module") -def square_scene(): - return _twisted_scene("square") - - -@pytest.fixture(scope="module") -def honeycomb_scene(): - return _twisted_scene("honeycomb") - - class TestSectioning: """``slice_mode="sectioned"``: one single-plane reconstruction per slice depth.""" SECTIONS = dict(num_slices=3, slice_thicknesses=50.0, slice_mode="sectioned", verbose=False) - @pytest.mark.parametrize("kernel", ["ssb", "obf", "wiener", "wdd", "prlx"]) - def test_sections_are_single_plane_reconstructions(self, recon, kernel): - recon.reconstruct(deconvolution_kernel=kernel, **self.SECTIONS) - stack = to_numpy(recon.corrected_bf).copy() + def test_sections_are_single_plane_reconstructions(self, recon): + kwargs = dict(deconvolution_kernel="obf", obj_type="complex") + recon.reconstruct(**kwargs, **self.SECTIONS) + phases = to_numpy(recon.corrected_bf).copy() + amplitudes = to_numpy(recon.corrected_amplitude).copy() assert recon.slice_mode == "sectioned" + assert recon.obj.shape == phases.shape == amplitudes.shape == (3, N, N) assert np.allclose(recon.slice_depths, [-50.0, 0.0, 50.0]) - for section, depth in zip(stack, recon.slice_depths): + for phase, amplitude, depth in zip(phases, amplitudes, recon.slice_depths): recon.reconstruct( - deconvolution_kernel=kernel, - override_aberration_coefs={"C10": TRUE_C10 + depth}, - verbose=False, + override_aberration_coefs={"C10": TRUE_C10 + depth}, verbose=False, **kwargs ) - assert np.array_equal(section, to_numpy(recon.corrected_bf)) + assert np.array_equal(phase, to_numpy(recon.corrected_bf)), f"{depth=}" + assert np.array_equal(amplitude, to_numpy(recon.corrected_amplitude)), f"{depth=}" def test_sections_start_from_the_initial_state_when_asked(self, dataset4d): recon = _build(dataset4d) @@ -1549,11 +1353,6 @@ def test_sections_start_from_the_initial_state_when_asked(self, dataset4d): assert np.array_equal(from_initial, to_numpy(recon.corrected_bf)) - def test_a_complex_object_has_amplitude_sections(self, recon): - recon.reconstruct(deconvolution_kernel="obf", obj_type="complex", **self.SECTIONS) - assert recon.corrected_amplitude.shape == recon.corrected_bf.shape == (3, N, N) - assert recon.obj.shape == (3, N, N) - def test_each_section_has_its_own_residual(self): """A layer at the reference plane is explained best by the section through it.""" _, layers, _ = twisted_square_cell() @@ -1575,28 +1374,16 @@ def test_icom_cannot_be_sectioned(self, recon): with pytest.raises(ValueError, match="iCoM does not depend on the aberrations"): recon.reconstruct(deconvolution_kernel="icom", **self.SECTIONS) - def test_kernels_without_a_joint_solve_point_to_sectioning(self, recon): - with pytest.raises(ValueError, match="slice_mode='sectioned'"): - recon.reconstruct( - deconvolution_kernel="ssb", - num_slices=2, - slice_thicknesses=5.0, - slice_mode="joint", - verbose=False, - ) - def test_sectioning_is_the_default(self, recon): recon.reconstruct( deconvolution_kernel="ssb", num_slices=2, slice_thicknesses=5.0, verbose=False ) assert recon.slice_mode == "sectioned" - @pytest.mark.parametrize("alias", ["sectioned", "depth-section", "sweep", "Section"]) - def test_aliases(self, recon, alias): - recon.reconstruct(deconvolution_kernel="ssb", **dict(self.SECTIONS, slice_mode=alias)) - assert recon.slice_mode == "sectioned" - - def test_rejects_an_unknown_mode(self, recon): + def test_mode_names(self, recon): + for alias in ("sectioned", "depth-section", "sweep", "Section"): + assert recon._normalize_slice_mode(alias) == "sectioned", f"{alias=}" + assert recon._normalize_slice_mode("joint-solve") == "joint" with pytest.raises(ValueError, match="`slice_mode`"): recon.reconstruct(deconvolution_kernel="obf", **dict(self.SECTIONS, slice_mode="x")) @@ -1607,11 +1394,6 @@ def test_a_joint_reconstruction_resets_the_mode(self, recon): assert recon.section_residuals is None assert recon.residual_loss() is not None - def test_permutations_leave_out_only_icom(self, recon): - images = recon._reconstruct_all_permutations(**self.SECTIONS) - assert len(images) == 5 - assert all(image.shape == (3, N, N) for image in images) - def test_fits_refuse_the_residual_of_a_sectioned_stack(self, recon): sections = dict(num_slices=3, slice_thicknesses=50.0, slice_mode="sectioned") with pytest.raises(ValueError, match="no residual to minimize"): @@ -1634,26 +1416,14 @@ def test_sections_carry_gradients(self, recon): recon.rms_gradient_loss().backward() assert c10.grad is not None and torch.isfinite(c10.grad) - def test_survives_a_round_trip(self, dataset4d, tmp_path): - recon = _build(dataset4d) - recon.reconstruct(deconvolution_kernel="obf", **self.SECTIONS) - path = str(tmp_path / "sections.zip") - recon.save(path, mode="o") - restored = load(path) - - assert restored.slice_mode == "sectioned" - assert np.allclose(restored.section_residuals, recon.section_residuals) - assert np.array_equal(restored.obj, recon.obj) - class TestMultislice: """Several thin slices solved jointly at each spatial frequency.""" - @pytest.mark.parametrize("scene", ["square_scene", "honeycomb_scene"]) - def test_twisted_layers_separate(self, scene, request): + def test_twisted_layers_separate(self, square_scene): """Layers far closer than ``lambda / alpha**2`` (51 A here) come apart, which a C10 sweep cannot do at any defocus.""" - recon, layers, spacing = request.getfixturevalue(scene) + recon, layers, spacing = square_scene recon.reconstruct( deconvolution_kernel="obf", num_slices=2, @@ -1740,106 +1510,36 @@ def test_shapes_and_depths(self, square_scene): assert np.allclose(recon.slice_depths, [-8.0, -2.0, 10.0]) assert recon.num_slices == 3 - def test_one_slice_is_the_single_plane(self, recon): - recon.reconstruct(deconvolution_kernel="obf", verbose=False) - single = recon.obj.copy() - recon.reconstruct(deconvolution_kernel="obf", num_slices=1, verbose=False) - - assert np.array_equal(recon.obj, single) - - @pytest.mark.parametrize("kernel", ["ssb", "prlx", "icom"]) - def test_kernels_without_a_joint_solve_raise(self, recon, kernel): - with pytest.raises(ValueError, match="only 'obf'"): + def test_kernels_without_a_joint_solve_point_to_sectioning(self, recon): + with pytest.raises(ValueError, match="only 'obf'.*slice_mode='sectioned'"): recon.reconstruct( - deconvolution_kernel=kernel, + deconvolution_kernel="ssb", num_slices=2, slice_mode="joint", slice_thicknesses=5.0, verbose=False, ) - @pytest.mark.parametrize( - ("thicknesses", "match"), - [(None, "cannot be None"), ([1.0, 2.0, 3.0], "num_slices - 1"), (-1.0, "positive")], - ) - def test_slice_thicknesses_are_validated(self, recon, thicknesses, match): - with pytest.raises(ValueError, match=match): - recon.reconstruct( - deconvolution_kernel="obf", - num_slices=2, - slice_mode="joint", - slice_thicknesses=thicknesses, - verbose=False, - ) - - def test_an_empirical_probe_reproduces_the_analytic_one(self, dataset4d): - from quantem.diffractive_imaging.complex_probe import FourierProbe - - analytic = _build(dataset4d) - psi = analytic_probe_array(analytic, {"C10": TRUE_C10}) - probe = FourierProbe.from_array( - psi, analytic.reciprocal_sampling, analytic.wavelength, normalize=False - ) - empirical = _build(dataset4d, fourier_probe=probe) - - kwargs = dict( - deconvolution_kernel="obf", - obj_type="complex", - num_slices=2, - slice_mode="joint", - slice_thicknesses=20.0, - verbose=False, - ) - expected = analytic.reconstruct(**kwargs).obj.copy() - # two descriptions of one probe differ by float32 rounding, which the near-parallel - # slice columns at low q amplify up to the eigenvalue floor - error = np.abs(empirical.reconstruct(**kwargs).obj - expected).max() - assert error < 1e-4 * np.abs(expected).max() - - def test_the_residual_ignores_the_probe_normalization(self, dataset4d): - """A measured probe normalized to unit power gives the analytic probe's residual.""" - from quantem.diffractive_imaging.complex_probe import FourierProbe - - analytic = _build(dataset4d) - psi = analytic_probe_array(analytic, {"C10": TRUE_C10}) - psi = psi / psi.abs().square().sum().sqrt() - probe = FourierProbe.from_array( - psi, analytic.reciprocal_sampling, analytic.wavelength, normalize=False - ) - empirical = _build(dataset4d, fourier_probe=probe) - - kwargs = dict( - deconvolution_kernel="obf", - num_slices=2, - slice_mode="joint", - slice_thicknesses=20.0, - verbose=False, - ) - expected = analytic.reconstruct(**kwargs).residual_loss() - assert empirical.reconstruct(**kwargs).residual_loss() == pytest.approx(expected, abs=1e-4) - - def test_visualize_shows_one_row_per_slice(self, square_scene): - import matplotlib - - matplotlib.use("Agg") - recon, _, spacing = square_scene - recon.reconstruct( - deconvolution_kernel="obf", - num_slices=2, - slice_mode="joint", - slice_thicknesses=spacing, - verbose=False, - ) - fig, axs = recon.visualize(return_fig=True) - - assert np.shape(axs) == (2, 2) + def test_slice_thicknesses_are_validated(self, recon): + for thicknesses, match in ( + (None, "cannot be None"), + ([1.0, 2.0, 3.0], "num_slices - 1"), + (-1.0, "positive"), + ): + with pytest.raises(ValueError, match=match): + recon.reconstruct( + deconvolution_kernel="obf", + num_slices=2, + slice_mode="joint", + slice_thicknesses=thicknesses, + verbose=False, + ) - @pytest.mark.parametrize("kernel", ["obf"]) - def test_a_large_ridge_is_the_sweep(self, square_scene, kernel): + def test_a_large_ridge_is_the_sweep(self, square_scene): """With the ridge far above every eigenvalue each slice is its own plane alone.""" recon, _, spacing = square_scene recon.reconstruct( - deconvolution_kernel=kernel, + deconvolution_kernel="obf", num_slices=2, slice_mode="joint", slice_thicknesses=spacing, @@ -1850,7 +1550,7 @@ def test_a_large_ridge_is_the_sweep(self, square_scene, kernel): for image, z in zip(slices, (-spacing / 2, spacing / 2)): override = {"C10": ML_ABERRATIONS["C10"] + z} recon.reconstruct( - deconvolution_kernel=kernel, override_aberration_coefs=override, verbose=False + deconvolution_kernel="obf", override_aberration_coefs=override, verbose=False ) plane = to_numpy(recon.corrected_bf) # the two differ only in how the soft aperture edge is weighted @@ -2029,37 +1729,93 @@ def test_the_residual_ignores_the_prior(self, noisy_square_scenes): assert recon.noise_variance is None assert recon.reconstruct(prior="ard", **kwargs).residual_loss() == pytest.approx(expected) - @pytest.mark.parametrize("kernel", ["obf"]) - @pytest.mark.parametrize("obj_type", ["pure_phase", "complex"]) - def test_one_plane(self, noisy_square_scenes, kernel, obj_type): + def test_one_plane(self, noisy_square_scenes): recons, _, _ = noisy_square_scenes recon = recons[300.0] - kwargs = dict(deconvolution_kernel=kernel, obj_type=obj_type, verbose=False) - expected = to_numpy(recon.reconstruct(**kwargs).corrected_bf) - image = to_numpy(recon.reconstruct(prior="ard", **kwargs).corrected_bf) + for obj_type in ("pure_phase", "complex"): + kwargs = dict(deconvolution_kernel="obf", obj_type=obj_type, verbose=False) + expected = to_numpy(recon.reconstruct(**kwargs).corrected_bf) + image = to_numpy(recon.reconstruct(prior="ard", **kwargs).corrected_bf) + + assert image.shape == expected.shape, f"{obj_type=}" + assert np.isfinite(image).all(), f"{obj_type=}" + assert correlation(image, expected) > 0.9, f"{obj_type=}" + if obj_type == "complex": + assert np.isfinite(to_numpy(recon.corrected_amplitude)).all() + + def test_rejects_bad_arguments(self, recon): + for kwargs, match in ( + # only the normal equations of 'obf' take it; 'wiener' has its own prior + (dict(deconvolution_kernel="ssb", prior="ard"), "normal equations"), + (dict(deconvolution_kernel="obf", prior="lasso"), "prior"), + (dict(deconvolution_kernel="obf", noise_variance=1.0), "only used"), + (dict(deconvolution_kernel="obf", prior="ard", noise_variance=-1.0), "positive"), + ): + with pytest.raises(ValueError, match=match): + recon.reconstruct(verbose=False, **kwargs) - assert image.shape == expected.shape - assert np.isfinite(image).all() - assert correlation(image, expected) > 0.9 - if obj_type == "complex": - assert np.isfinite(to_numpy(recon.corrected_amplitude)).all() - @pytest.mark.parametrize("kernel", ["ssb", "prlx", "icom"]) - def test_only_the_normal_equation_kernels_take_it(self, recon, kernel): - with pytest.raises(ValueError, match="normal equations"): - recon.reconstruct(deconvolution_kernel=kernel, prior="ard", verbose=False) +class TestEmpiricalProbe: + """A probe given as an array must reconstruct as its analytic description does.""" - @pytest.mark.parametrize( - "kwargs, match", - [ - (dict(prior="lasso"), "prior"), - (dict(noise_variance=1.0), "only used"), - (dict(prior="ard", noise_variance=-1.0), "positive"), - ], - ) - def test_rejects_bad_arguments(self, recon, kwargs, match): - with pytest.raises(ValueError, match=match): - recon.reconstruct(deconvolution_kernel="obf", verbose=False, **kwargs) + def test_an_empirical_probe_reproduces_the_analytic_one(self, dataset4d): + from quantem.diffractive_imaging.complex_probe import FourierProbe + + analytic = _build(dataset4d) + psi = analytic_probe_array(analytic, {"C10": TRUE_C10}) + probe = FourierProbe.from_array( + psi, analytic.reciprocal_sampling, analytic.wavelength, normalize=False + ) + empirical = _build(dataset4d, fourier_probe=probe) + + for kernel in ("ssb", "obf"): + for obj_type in ("pure_phase", "complex"): + kwargs = dict(deconvolution_kernel=kernel, obj_type=obj_type, verbose=False) + expected = analytic.reconstruct(**kwargs).obj.copy() + image = empirical.reconstruct(**kwargs).obj + if obj_type == "pure_phase": + error = np.abs(image - expected).max() / np.abs(expected).max() + assert error < 1e-5, f"{kernel=} {obj_type=}" + else: + assert np.allclose(image, expected, rtol=1e-5, atol=1e-6), ( + f"{kernel=} {obj_type=}" + ) + + kwargs = dict( + deconvolution_kernel="obf", + obj_type="complex", + num_slices=2, + slice_mode="joint", + slice_thicknesses=20.0, + verbose=False, + ) + expected = analytic.reconstruct(**kwargs).obj.copy() + # two descriptions of one probe differ by float32 rounding, which the near-parallel + # slice columns at low q amplify up to the eigenvalue floor + error = np.abs(empirical.reconstruct(**kwargs).obj - expected).max() + assert error < 1e-4 * np.abs(expected).max() + + def test_the_residual_ignores_the_probe_normalization(self, dataset4d): + """A measured probe normalized to unit power gives the analytic probe's residual.""" + from quantem.diffractive_imaging.complex_probe import FourierProbe + + analytic = _build(dataset4d) + psi = analytic_probe_array(analytic, {"C10": TRUE_C10}) + psi = psi / psi.abs().square().sum().sqrt() + probe = FourierProbe.from_array( + psi, analytic.reciprocal_sampling, analytic.wavelength, normalize=False + ) + empirical = _build(dataset4d, fourier_probe=probe) + + kwargs = dict( + deconvolution_kernel="obf", + num_slices=2, + slice_mode="joint", + slice_thicknesses=20.0, + verbose=False, + ) + expected = analytic.reconstruct(**kwargs).residual_loss() + assert empirical.reconstruct(**kwargs).residual_loss() == pytest.approx(expected, abs=1e-4) class TestOverlapConvention: @@ -2087,10 +1843,8 @@ def test_a_shifted_probe_only_shifts_the_object(self, weak_scene): shifted = np.roll(expected, (-2, 3), axis=(0, 1)) assert correlation(to_numpy(recon.corrected_bf), shifted) > 0.9999 - @pytest.mark.parametrize("aberration", ["C21", "C23"]) - def test_odd_aberrations_are_fitted_with_their_sign(self, aberration): - angle = aberration.replace("C", "phi") - truth = {"C10": 50.0, aberration: 300.0, angle: 0.4} + def test_odd_aberrations_are_fitted_with_their_sign(self): + truth = {"C10": 50.0, "C21": 300.0, "phi21": 0.4} dataset = thin_layer_4dstem( [1j * 0.02 * white_noise_object_2D(n=64, phi0=1.0)], [0.0], @@ -2100,8 +1854,8 @@ def test_odd_aberrations_are_fitted_with_their_sign(self, aberration): losses = {} for name, override in ( ("true", truth), - ("inverted", dict(truth, **{angle: 0.4 + np.pi})), - ("absent", {"C10": 50.0, aberration: 0.0}), + ("inverted", dict(truth, phi21=0.4 + np.pi)), + ("absent", {"C10": 50.0, "C21": 0.0}), ): recon.reconstruct( deconvolution_kernel="obf", override_aberration_coefs=override, verbose=False @@ -2129,30 +1883,31 @@ def test_a_refined_probe_matches_the_native_one(self, weak_scene): assert np.allclose(recon.obj, expected, atol=1e-5 * np.abs(expected).max()) - @pytest.mark.parametrize("num_modes", [1, 2]) - def test_a_double_precision_probe_is_stored_in_single_precision(self, weak_scene, num_modes): + def test_a_double_precision_probe_is_stored_in_single_precision(self, weak_scene): """The reconstructions accumulate in complex64, and MPS has no complex128.""" from quantem.diffractive_imaging.complex_probe import FourierProbe - array = library_probe_fourier(ML_ABERRATIONS) - if num_modes > 1: - array = np.stack([array] * num_modes) / np.sqrt(num_modes) recon = _full_detector(weak_scene) - recon.fourier_probe = _empirical(array) - recon.reconstruct(deconvolution_kernel="obf", verbose=False) - expected = recon.obj.copy() + single = library_probe_fourier(ML_ABERRATIONS) + for array in (single, np.stack([single] * 2) / np.sqrt(2)): + num_modes = 1 if array.ndim == 2 else len(array) + recon.fourier_probe = _empirical(array) + recon.reconstruct(deconvolution_kernel="obf", verbose=False) + expected = recon.obj.copy() - probe = FourierProbe.from_array( - np.asarray(array, dtype=np.complex128), - (ML_RECIPROCAL_SAMPLING, ML_RECIPROCAL_SAMPLING), - electron_wavelength_angstrom(PROBE_ENERGY), - ) - assert probe.array.dtype == torch.complex64 - assert probe.modes is None or probe.modes.dtype == torch.complex64 - recon.fourier_probe = probe - recon.reconstruct(deconvolution_kernel="obf", verbose=False) + probe = FourierProbe.from_array( + np.asarray(array, dtype=np.complex128), + (ML_RECIPROCAL_SAMPLING, ML_RECIPROCAL_SAMPLING), + electron_wavelength_angstrom(PROBE_ENERGY), + ) + assert probe.array.dtype == torch.complex64, f"{num_modes=}" + assert probe.modes is None or probe.modes.dtype == torch.complex64, f"{num_modes=}" + recon.fourier_probe = probe + recon.reconstruct(deconvolution_kernel="obf", verbose=False) - assert np.allclose(recon.obj, expected, atol=1e-5 * np.abs(expected).max()) + assert np.allclose(recon.obj, expected, atol=1e-5 * np.abs(expected).max()), ( + f"{num_modes=}" + ) def test_refinement_makes_rotated_sampling_accurate(self, weak_scene): """A rotation puts ``k - q`` between grid points, where bilinear sampling is only as @@ -2242,25 +1997,22 @@ def test_resampling_refines_every_mode(self): expected = mode.resampled_to((step, step)).array assert torch.allclose(fine.array, expected) - @pytest.mark.parametrize( - "kwargs", - [ + def test_one_mode_is_the_coherent_probe(self, weak_scene): + recon = _full_detector(weak_scene) + psi = library_probe_fourier(ML_ABERRATIONS) + for kwargs in ( {}, {"obj_type": "complex"}, {"num_slices": 2, "slice_mode": "joint", "slice_thicknesses": 10.0}, - ], - ) - def test_one_mode_is_the_coherent_probe(self, weak_scene, kwargs): - recon = _full_detector(weak_scene) - psi = library_probe_fourier(ML_ABERRATIONS) - results = [] - for array in (psi, psi[None]): - recon.fourier_probe = _empirical(array) - recon.reconstruct(deconvolution_kernel="obf", verbose=False, **kwargs) - results.append((recon.obj.copy(), recon.residual_loss())) + ): + results = [] + for array in (psi, psi[None]): + recon.fourier_probe = _empirical(array) + recon.reconstruct(deconvolution_kernel="obf", verbose=False, **kwargs) + results.append((recon.obj.copy(), recon.residual_loss())) - assert np.allclose(results[0][0], results[1][0]) - assert results[0][1] == pytest.approx(results[1][1]) + assert np.allclose(results[0][0], results[1][0]), f"{kwargs=}" + assert results[0][1] == pytest.approx(results[1][1]), f"{kwargs=}" def test_the_mode_sum_explains_a_focal_spread(self): """No single focus describes two incoherent ones; their mode sum does exactly.""" @@ -2566,24 +2318,25 @@ def probe_at(c10): scale = float(expected.abs().max()) assert float((actual - expected).abs().max()) < 1e-6 * scale - @pytest.mark.parametrize("obj_type", ["pure_phase", "complex"]) - def test_a_finite_source_leaves_obf_and_the_residual_unchanged(self, weak_scene, obj_type): + def test_a_finite_source_leaves_obf_and_the_residual_unchanged(self, weak_scene): recon = DirectPtychography.from_dataset4d(weak_scene, **ml_kwargs()) - kwargs = dict(deconvolution_kernel="obf", obj_type=obj_type, verbose=False) - expected = to_numpy(recon.reconstruct(**kwargs).corrected_bf) - residual = recon.residual_loss() - # small enough that no frequency inside the band falls below the round-off floor - recon.source_size = 0.2 - image = to_numpy(recon.reconstruct(**kwargs).corrected_bf) - - # the source multiplies every column at q by the same factor, which OBF's gain and the - # least-squares projection both divide out - assert np.abs(image - expected).max() < 1e-4 * np.abs(expected).max() - # the data are noiseless and the residual is 1e-7 to 1e-5 of the energy, so only its - # absolute round-off is meaningful. That round-off differs between platforms (up to 2e-7 - # on Linux CI), while mishandling the source would move the residual by far more, since - # its envelope falls well below one inside the band - assert recon.residual_loss() == pytest.approx(residual, abs=1e-6) + for obj_type in ("pure_phase", "complex"): + kwargs = dict(deconvolution_kernel="obf", obj_type=obj_type, verbose=False) + recon.source_size = 0.0 + expected = to_numpy(recon.reconstruct(**kwargs).corrected_bf) + residual = recon.residual_loss() + # small enough that no frequency inside the band falls below the round-off floor + recon.source_size = 0.2 + image = to_numpy(recon.reconstruct(**kwargs).corrected_bf) + + # the source multiplies every column at q by the same factor, which OBF's gain and + # the least-squares projection both divide out + assert np.abs(image - expected).max() < 1e-4 * np.abs(expected).max(), f"{obj_type=}" + # the data are noiseless and the residual is 1e-7 to 1e-5 of the energy, so only its + # absolute round-off is meaningful. That round-off differs between platforms (up to + # 2e-7 on Linux CI), while mishandling the source would move the residual by far + # more, since its envelope falls well below one inside the band + assert recon.residual_loss() == pytest.approx(residual, abs=1e-6), f"{obj_type=}" def test_ssb_keeps_its_coherent_weights(self, weak_scene): recon = DirectPtychography.from_dataset4d(weak_scene, **ml_kwargs()) @@ -2612,21 +2365,10 @@ def test_the_residual_finds_the_focal_spread(self): # 0.13, 5e-7 and 0.16: the envelope is the nine modes' average to quadrature accuracy assert residuals[40.0] < 1e-3 * min(residuals[0.0], residuals[80.0]) - def test_survives_a_round_trip(self, weak_scene, tmp_path): - recon = DirectPtychography.from_dataset4d( - weak_scene, focal_spread=30.0, source_size=0.2, **ml_kwargs() - ) - path = str(tmp_path / "coherence.zip") - recon.save(path, mode="o") - restored = load(path) - - assert restored.focal_spread == 30.0 - assert restored.source_size == 0.2 - - @pytest.mark.parametrize("name", ["focal_spread", "source_size"]) - def test_rejects_a_negative_spread(self, recon, name): - with pytest.raises(ValueError, match="non-negative"): - setattr(recon, name, -1.0) + def test_rejects_a_negative_spread(self, recon): + for name in ("focal_spread", "source_size"): + with pytest.raises(ValueError, match=f"`{name}` must be a non-negative"): + setattr(recon, name, -1.0) class TestWienerKernel: @@ -2682,27 +2424,12 @@ def test_it_solves_several_slices(self, noisy_square_scenes): assert corr > corr_ls + 0.03 assert purity > 0.95 - def test_the_noise_variance(self, calibration_scene): - recon, _, _ = calibration_scene - recon.reconstruct(deconvolution_kernel="wiener", verbose=False) - estimated = recon.noise_variance - recon.reconstruct( - deconvolution_kernel="wiener", noise_variance=2 * estimated, verbose=False - ) - - assert estimated > 0 - assert recon.noise_variance == 2 * estimated - def test_shares_the_residual_of_obf(self, calibration_scene): recon, _, _ = calibration_scene expected = recon.reconstruct(deconvolution_kernel="obf", verbose=False).residual_loss() recon.reconstruct(deconvolution_kernel="wiener", verbose=False) assert recon.residual_loss() == pytest.approx(expected) - def test_refuses_the_ard_prior(self, recon): - with pytest.raises(ValueError, match="its own prior"): - recon.reconstruct(deconvolution_kernel="wiener", prior="ard", verbose=False) - @pytest.fixture(scope="module") def strong_scene(): @@ -2793,19 +2520,15 @@ def test_upsampling(self, strong_scene): assert fine.shape == (128, 128) assert correlation(fine[::2, ::2], coarse) > 0.99 - @pytest.mark.parametrize( - "kwargs, match", - [ - ( - dict(num_slices=2, slice_mode="joint", slice_thicknesses=5.0), - "single multiplicative object", - ), - (dict(prior="ard"), "normal equations"), - ], - ) - def test_rejects(self, recon, kwargs, match): - with pytest.raises(ValueError, match=match): - recon.reconstruct(deconvolution_kernel="wdd", verbose=False, **kwargs) + def test_rejects_several_joint_slices(self, recon): + with pytest.raises(ValueError, match="single multiplicative object"): + recon.reconstruct( + deconvolution_kernel="wdd", + num_slices=2, + slice_mode="joint", + slice_thicknesses=5.0, + verbose=False, + ) @pytest.fixture(scope="module") From 044d92f24b152c9af70effe727aee5227f2cd651 Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 15:07:18 -0700 Subject: [PATCH 12/19] Consolidate the DirectPtychographyMontage and GPU tests Removed (cannot fail, or exercise no library code): - TestBandLimit::test_deconvolution_kernels_are_untouched[ssb,obf]: band_limit only reaches _resolve_q_lowpass, which returns the cutoff unchanged for ssb/obf either way, so both runs were the same computation. test_an_explicit_cutoff_below_the_band_limit_ still_wins asserts that branch directly. - test_the_icom_kernel_is_the_ricom_kernel (numpy-only identity), test_photon_wavelength_is_not_the_electron_one (core util and literal arithmetic), test_nearest_is_a_roll_of_the_bright_field_images (rounding identity), test_interpolation_defaults_to_nearest (identical call). - isfinite-only or "not allclose" checks: test_accepts_the_deconvolution_kernels (covered by TestRealSpaceKernels::test_converges_to_the_fourier_kernel), test_icom_runs_in_either_mode (test_icom_matches_the_fourier_class, test_truncated_icom_is_ricom, TestSparseEvents::test_icom_matches_dense), test_scattered_positions_reconstruct (test_scattered_positions_match_dense), test_position_axis_order_is_row_col (test_matches_from_dataset4d, test_scan_origin_recovers_the_input_positions), test_rotation_changes_the_reconstruction (replaced by an exact quarter-turn case in test_matches_parallax_kernel). - test_fft_mode_beats_a_truncated_stencil (implied by test_fft_mode_is_exact and the convergence test), test_torch_size_attributes_round_trip_as_tuples (never round-tripped), test_is_analytic_reports_which_path_is_live, test_fft_mode_leaves_no_stencil_info, test_dense_and_sparse_are_mutually_exclusive. - TestIntegerShiftConstruction: its integrality check is now a precondition inside test_matches_parallax_kernel. - Copies of the Fourier file's from_dataset3d tests, which go through the same build_vbf_stack_from_dataset3d: auto scan sampling, Dataset2d positions, wrong units, mismatched count. - test_grid_search_recovers_the_seeded_defocus: same base driver as the Fourier test_grid_search; the montage objective is covered by test_residual_loss_is_minimized_at_the_true_defocus. - TestComplexAndMultislice::test_parallax_rejects_complex: same _validate_reconstruction_mode as the Fourier test. - The Fourier cases of TestFourierProbe::test_empirical_probe_reproduces_the_analytic_one (now TestEmpiricalProbe in the Fourier file) and the montage case of test_prlx_runs_on_a_modelled_empirical_probe (test_fitted_coefficients_drive_parallax). - test_semiangle_cutoff_becomes_optional: its reconstruction moved into test_an_empirical_probe_reads_its_extent_off_the_bf_mask. Merged: - test_matches_parallax_kernel[pixel_shift x U], test_matches_with_phase_flip, test_matches_with_butterworth_filters -> one loop on a shared pair (pixel_shift reaches no other path), plus the quarter turn. - Sampling and field-of-view tests over both classes (scalebar[6], resets, before reconstructing, fov, wrapped canvas, wrap default, zero scan origin) -> test_scalebar_follows_the_upsampling; padded canvas at U=2 only. - TestWavelength and TestSemiangleCutoff: the checks live in the base class, so one class; six parametrized refusals -> two loops. - Eight round-trip tests -> TestSerialization::test_round_trip (an ungridded montage with scan offset, defocus gradient and a stored fit, and an empirical-probe montage), TestSparseEvents::test_survives_a_round_trip, and the wavelength round trip. - TestDefocusGradient: the two rate tests, the delta tests, the update/leave defocus fits and the three defocus_map argument tests each become one; the gradient sweeps keep one oblique gradient (plus the zero control for the fit). - TestRealSpaceKernels: the residual-undefined tests and the auto-radius tests become one each. TestConvolutionModes: conv2d[2x2] and the mode-resolution tests become loops. - TestSparseEvents: the wrap-nearest case was the same call as test_deconvolution_stencils_match_dense[ssb]; the pad-bilinear case reuses the class pair through reconstruct(boundary="pad"), which is bit-identical to building a pad pair; the complex sparse tests move here onto the class pair. The counts are scaled to a mean of 4 rather than 40, which leaves every dense/sparse error unchanged (largest 1.3e-5 against 1e-4) and cuts each from_vector build from 2.2 s to 0.2 s. - TestFrameDrift: the drift is estimated once per class; the four rejection tests share one set of frames. - The three stencil refusals of TestComplexAndMultislice -> one test. - A module fixture shares the Fourier/montage pair among the tests that only reconstruct. GPU: - TestAccelerators moves to test_direct_ptychography_gpu.py and compares every route with the CPU, residual included; the ungridded origin fit now has aberrations, since with none the parallax image was identically zero. test_padded_canvas_shape_is_device_ independent is the pad case of that loop. - test_sections_match_the_cpu[3 kernels] -> obf: sectioning is a CPU-side loop over the single-plane kernels, each compared in test_every_kernel_matches_the_cpu. - conftest: ACCELERATORS is unused; tilted_layers_4dstem replaces the layered-dataset helpers duplicated in the GPU and tilt-series files. Co-Authored-By: Claude Opus 5.5 --- tests/diffractive_imaging/conftest.py | 34 +- .../test_direct_ptychography_gpu.py | 103 +- .../test_direct_ptychography_montage.py | 1346 +++++------------ 3 files changed, 462 insertions(+), 1021 deletions(-) diff --git a/tests/diffractive_imaging/conftest.py b/tests/diffractive_imaging/conftest.py index bda2cc9aa..92c6687b0 100644 --- a/tests/diffractive_imaging/conftest.py +++ b/tests/diffractive_imaging/conftest.py @@ -583,17 +583,6 @@ def ctf_interleaved_frames( return datasets, positions_A -#: every accelerator present, so the same assertions run on CPU, CUDA and MPS -ACCELERATORS = [ - device - for device, present in ( - ("cuda", torch.cuda.is_available()), - ("mps", torch.backends.mps.is_available()), - ) - if present -] - - # --------------------------------------------------------------------------- # complex-valued and multislice objects # @@ -784,6 +773,29 @@ def twisted_square_cell(): return shape, layers, float(np.degrees(twist)) +def periodic_shift(image: np.ndarray, shift, sampling: float = ML_SAMPLING) -> np.ndarray: + """``image`` translated by ``shift`` Angstrom, periodically, by the Fourier shift theorem.""" + q = np.fft.fftfreq(image.shape[0], sampling) + ramp = np.exp(-2j * np.pi * (q[:, None] * shift[0] + q[None, :] * shift[1])) + return np.fft.ifft2(np.fft.fft2(image) * ramp).real + + +def tilted_layers_4dstem( + tilt_mrad=(0.0, 0.0), spacing: float = 6.0, dose=None, seed=0, drift=(0.0, 0.0) +) -> Dataset4d: + """The twisted square layers, ``0.3`` rad strong and ``spacing`` apart about the C10 plane, + seen at a specimen tilt: each displaced by its depth times the tilt, and both by ``drift``. + """ + _, layers, _ = twisted_square_cell() + theta = np.asarray(tilt_mrad, dtype=float) * 1e-3 + depths = (-spacing / 2, spacing / 2) + logs = [ + 1j * periodic_shift(0.3 * layer, z * theta + np.asarray(drift)) + for layer, z in zip(layers, depths) + ] + return thin_layer_4dstem(logs, list(depths), ML_ABERRATIONS, dose=dose, seed=seed) + + def layer_purity(slices, truths) -> tuple[float, float]: """``(purity, correlation)`` averaged over layers, each true layer scored at its own slice. diff --git a/tests/diffractive_imaging/test_direct_ptychography_gpu.py b/tests/diffractive_imaging/test_direct_ptychography_gpu.py index 028654ea2..43af7916f 100644 --- a/tests/diffractive_imaging/test_direct_ptychography_gpu.py +++ b/tests/diffractive_imaging/test_direct_ptychography_gpu.py @@ -4,10 +4,13 @@ these tests run only where a GPU is present. """ +import warnings + import numpy as np import pytest import torch +from quantem.core.datastructures import Dataset3d from quantem.core.utils.utils import to_numpy from quantem.diffractive_imaging import ( DirectPtychography, @@ -18,13 +21,17 @@ ) from .conftest import ( - ML_ABERRATIONS, ML_N_DET, - ML_SAMPLING, + PROBE_ENERGY, + SCAN_SAMPLING, + SEMIANGLE_CUTOFF, + N, direct_ptycho_kwargs, integer_shift_defocus, ml_kwargs, + scan_positions_px, thin_layer_4dstem, + tilted_layers_4dstem, twisted_square_cell, ) @@ -49,10 +56,11 @@ def pair(dataset4d): return _build(dataset4d, "cpu"), _build(dataset4d, "gpu") -def _assert_close(image, expected, rtol=1e-4): +def _assert_close(image, expected, rtol=1e-4, case=""): image, expected = to_numpy(image), to_numpy(expected) - assert np.isfinite(image).all() - assert np.abs(image - expected).max() <= rtol * np.abs(expected).max() + assert image.shape == expected.shape, case + assert np.isfinite(image).all(), case + assert np.abs(image - expected).max() <= rtol * np.abs(expected).max(), case def test_gpu_resolves_to_a_device(pair): @@ -81,24 +89,9 @@ def test_a_complex_object_matches_the_cpu(pair, kernel): SPACING = 6.0 -def _shifted(image, shift): - q = np.fft.fftfreq(image.shape[0], ML_SAMPLING) - ramp = np.exp(-2j * np.pi * (q[:, None] * shift[0] + q[None, :] * shift[1])) - return np.fft.ifft2(np.fft.fft2(image) * ramp).real - - -def _layered(tilt_mrad=(0, 0)): - """Two twisted layers 6 A apart, each displaced by its depth times the tilt.""" - _, layers, _ = twisted_square_cell() - theta = np.asarray(tilt_mrad, dtype=float) * 1e-3 - depths = (-SPACING / 2, SPACING / 2) - logs = [1j * _shifted(0.3 * layer, z * theta) for layer, z in zip(layers, depths)] - return thin_layer_4dstem(logs, list(depths), ML_ABERRATIONS) - - @pytest.fixture(scope="module") def layered_pair(): - dataset = _layered() + dataset = tilted_layers_4dstem(spacing=SPACING) return tuple( DirectPtychography.from_dataset4d(dataset, device=device, **ml_kwargs()) for device in ("cpu", "gpu") @@ -146,8 +139,69 @@ def test_the_montage_fft_route_matches_the_cpu(dataset4d): _assert_close(images["gpu"], images["cpu"]) +def test_the_montage_splat_and_stencil_routes_match_the_cpu(dataset4d): + """MPS has no float64 at all, so every positional tensor -- coordinates, shifts, canvas + origins -- has to take its dtype from the device rather than hardcode one.""" + for constructor, reconstruct in ( + (dict(boundary="wrap"), {}), + # float noise must not push a canvas bound across an integer on one device only + (dict(boundary="pad"), {}), + (dict(boundary="wrap"), dict(upsampling_factor=2, interpolation="bilinear")), + (dict(boundary="pad", defocus_gradient=(30.0, -10.0)), {}), + (dict(boundary="wrap"), dict(deconvolution_kernel="ssb", stencil_radius=4)), + ): + case = f"{constructor=} {reconstruct=}" + results = {} + for device in ("cpu", "gpu"): + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, + device=device, + edge_blend_pixels=0, + **direct_ptycho_kwargs(C10), + **constructor, + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + montage.reconstruct(verbose=False, **reconstruct) + residual = None if "deconvolution_kernel" in reconstruct else montage.residual_loss() + results[device] = montage.obj, residual + _assert_close(results["gpu"][0], results["cpu"][0], case=case) + if results["cpu"][1] is not None: + assert results["gpu"][1] == pytest.approx(results["cpu"][1], rel=1e-4), case + + +def test_an_ungridded_scan_fits_its_origin_on_the_device(dataset4d): + """Regression: probe positions arrive as numpy and met the measured origin on MPS. + + Forcing the origin skips the centre-of-mass fit entirely, which is where the device + mismatch lived, so this deliberately lets it run. + """ + dataset3d = Dataset3d.from_array( + np.asarray(dataset4d.array).reshape(-1, N, N), + name="ungridded", + sampling=(1.0, dataset4d.sampling[-2], dataset4d.sampling[-1]), + units=("index", "A^-1", "A^-1"), + ) + images = {} + for device in ("cpu", "gpu"): + reconstruction = DirectPtychographyMontage.from_dataset3d( + dataset3d, + scan_positions_px() * SCAN_SAMPLING, + energy=PROBE_ENERGY, + semiangle_cutoff=SEMIANGLE_CUTOFF, + rotation_angle=0.0, + aberration_coefs={"C10": C10}, + scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), + device=device, + verbose=False, + ) + images[device] = reconstruction.reconstruct(deconvolution_kernel="prlx", verbose=False).obj + assert np.abs(images["cpu"]).max() > 0 + _assert_close(images["gpu"], images["cpu"]) + + def test_a_tilt_series_matches_the_cpu(): - datasets = [_layered((0, 0)), _layered((10, 0))] + datasets = [tilted_layers_4dstem(tilt, SPACING) for tilt in ((0, 0), (10, 0))] images = {} for device in ("cpu", "gpu"): series = DirectPtychography.from_tilt_series( @@ -183,11 +237,10 @@ def test_tilt_corrected_dark_field_matches_the_cpu(): _assert_close(images["gpu"][1], images["cpu"][1]) -@pytest.mark.parametrize("kernel", ["ssb", "obf", "prlx"]) -def test_sections_match_the_cpu(pair, kernel): +def test_sections_match_the_cpu(pair): cpu, gpu = pair kwargs = dict( - deconvolution_kernel=kernel, num_slices=3, slice_thicknesses=50.0, slice_mode="sectioned" + deconvolution_kernel="obf", num_slices=3, slice_thicknesses=50.0, slice_mode="sectioned" ) cpu.reconstruct(**kwargs) gpu.reconstruct(**kwargs) diff --git a/tests/diffractive_imaging/test_direct_ptychography_montage.py b/tests/diffractive_imaging/test_direct_ptychography_montage.py index 20145001e..41182f669 100644 --- a/tests/diffractive_imaging/test_direct_ptychography_montage.py +++ b/tests/diffractive_imaging/test_direct_ptychography_montage.py @@ -31,7 +31,6 @@ ) from .conftest import ( - ACCELERATORS, CTF_SAMPLING, CTF_SCAN_SAMPLING, ORIGIN, @@ -80,6 +79,12 @@ def _build_pair(dataset4d, defocus): return fourier, montage +@pytest.fixture(scope="module") +def fourier_montage(dataset4d): + """`_build_pair` at the integer-shift defocus, for the tests that only reconstruct.""" + return _build_pair(dataset4d, integer_shift_defocus(1)) + + def _relative_error(a, b): return float(np.abs(a - b).max() / np.abs(b).max()) @@ -87,25 +92,6 @@ def _relative_error(a, b): class TestSplatKernel: """Unit tests for `scatter_add_splat`, independent of any reconstruction.""" - def test_integer_coordinates_deposit_one_pixel(self): - values = torch.tensor([[2.0]]) - coords = torch.tensor([[[3.0, 4.0]]]) - sum_w, sum_wv, sum_wv2 = scatter_add_splat(values, coords, (8, 8)) - - assert sum_w.sum().item() == pytest.approx(1.0) - assert sum_w.reshape(8, 8)[3, 4].item() == pytest.approx(1.0) - assert sum_wv.reshape(8, 8)[3, 4].item() == pytest.approx(2.0) - assert sum_wv2.reshape(8, 8)[3, 4].item() == pytest.approx(4.0) - - def test_half_pixel_splits_four_ways(self): - values = torch.tensor([[1.0]]) - coords = torch.tensor([[[3.5, 4.5]]]) - sum_w, _, _ = scatter_add_splat(values, coords, (8, 8)) - - nonzero = sum_w[sum_w > 0] - assert nonzero.numel() == 4 - assert torch.allclose(nonzero, torch.full((4,), 0.25, dtype=nonzero.dtype)) - def test_wrap_is_periodic(self): values = torch.tensor([[1.0]]) coords = torch.tensor([[[-0.5, 8.25]]]) @@ -125,14 +111,6 @@ def test_pad_drops_out_of_bounds_without_corrupting_the_edge(self): assert sum_wv.sum().item() == pytest.approx(2.0) assert sum_w.reshape(8, 8)[0].max().item() == pytest.approx(0.0) - def test_weights_are_a_partition_of_unity(self): - generator = torch.Generator().manual_seed(0) - values = torch.randn(4, 100, generator=generator) - coords = torch.rand(4, 100, 2, generator=generator) * 8 - sum_w, _, _ = scatter_add_splat(values, coords, (8, 8), boundary="wrap") - - assert sum_w.sum().item() == pytest.approx(400.0) - def test_matches_brute_force(self): generator = torch.Generator().manual_seed(1) values = torch.randn(3, 50, generator=generator) @@ -189,21 +167,6 @@ def test_rejects_unknown_modes(self): with pytest.raises(ValueError, match="interpolation"): scatter_add_splat(values, coords, (8, 8), interpolation="cubic") - def test_convolve_with_a_unit_tap_is_a_plain_splat(self): - values = torch.tensor([[1.0, 2.0]]) - coords = torch.tensor([[[3.0, 4.0], [5.0, 6.0]]]) - - out = scatter_add_convolve( - values, - coords, - (8, 8), - torch.tensor([[0, 0]]), - torch.ones(1, 1, dtype=torch.complex64), - ).reshape(8, 8) - - assert out[3, 4].item() == pytest.approx(1.0) - assert out[5, 6].item() == pytest.approx(2.0) - def test_convolve_places_each_tap_at_its_offset(self): offsets = torch.tensor([[0, 0], [1, 0], [0, -2]]) weights = torch.tensor([[1.0, 2.0j, -3.0]], dtype=torch.complex64) @@ -230,83 +193,47 @@ def test_convolve_wraps_taps_at_the_boundary(self): assert out[7, 7].item() == pytest.approx(1.0) -class TestIntegerShiftConstruction: - """The defocus used below must put every BF pixel on an exact canvas pixel.""" - - @pytest.mark.parametrize("upsampling_factor", [1, 2]) - @pytest.mark.parametrize("pixel_shift", [1, 2]) - def test_shifts_land_on_pixel_centers(self, dataset4d, upsampling_factor, pixel_shift): - defocus = integer_shift_defocus(pixel_shift, upsampling_factor) - montage = DirectPtychographyMontage.from_dataset4d( - dataset4d, boundary="wrap", **_common_kwargs(defocus) - ) - shifts, _ = montage._return_shifts_px( - 0.0, {"C10": defocus}, montage.bf_mask, upsampling_factor - ) - residual = (shifts - shifts.round()).abs().max().item() - - assert residual < 1e-4, f"shifts are not integral: max residual {residual:.2e} px" - assert shifts.abs().max().item() > 1.0, "test would be vacuous with zero shifts" - - class TestFourierEquivalence: """`DirectPtychographyMontage` must reproduce `DirectPtychography`'s parallax kernel.""" - @pytest.mark.parametrize("upsampling_factor", [1, 2]) - @pytest.mark.parametrize("pixel_shift", [1, 2]) - def test_matches_parallax_kernel(self, dataset4d, upsampling_factor, pixel_shift): - defocus = integer_shift_defocus(pixel_shift, upsampling_factor) - fourier, montage = _build_pair(dataset4d, defocus) - - recon_kwargs = dict( - deconvolution_kernel="prlx", - parallax_flip_phase=False, - upsampling_factor=upsampling_factor, - verbose=False, - ) - fourier.reconstruct(**recon_kwargs) - montage.reconstruct( - boundary="wrap", interpolation="bilinear", weight_normalize=False, **recon_kwargs - ) - - assert montage.obj.shape == fourier.obj.shape - assert _relative_error(montage.obj, fourier.obj) < 1e-4 - - def test_matches_with_phase_flip(self, dataset4d): - """The phase-flip filter is BF-independent, so post-hoc application is exact.""" - defocus = integer_shift_defocus(1) - fourier, montage = _build_pair(dataset4d, defocus) - - fourier.reconstruct(deconvolution_kernel="prlx", parallax_flip_phase=True, verbose=False) - montage.reconstruct( - deconvolution_kernel="prlx", - parallax_flip_phase=True, - weight_normalize=False, - verbose=False, - ) - - assert _relative_error(montage.obj, fourier.obj) < 1e-4 - - def test_matches_with_butterworth_filters(self, dataset4d): - defocus = integer_shift_defocus(1) - fourier, montage = _build_pair(dataset4d, defocus) + def test_matches_parallax_kernel(self, fourier_montage): + fourier, montage = fourier_montage + for shared, montage_only in ( + (dict(upsampling_factor=1), dict(interpolation="bilinear")), + (dict(upsampling_factor=2), dict(interpolation="bilinear")), + # the phase-flip filter is BF-independent, so post-hoc application is exact + (dict(parallax_flip_phase=True), {}), + (dict(q_lowpass=0.2, q_highpass=0.02), {}), + # a quarter turn keeps every shift on a whole pixel + (dict(override_rotation_angle=90.0), {}), + ): + upsampling_factor = shared.get("upsampling_factor", 1) + defocus = integer_shift_defocus(1, upsampling_factor) + shifts, _ = montage._return_shifts_px( + shared.get("override_rotation_angle", 0.0), + {"C10": defocus}, + montage.bf_mask, + upsampling_factor, + ) + assert (shifts - shifts.round()).abs().max() < 1e-4, f"{shared=}" + assert shifts.abs().max() > 1.0, f"{shared=}" - recon_kwargs = dict( - deconvolution_kernel="prlx", - parallax_flip_phase=False, - q_lowpass=0.2, - q_highpass=0.02, - verbose=False, - ) - fourier.reconstruct(**recon_kwargs) - montage.reconstruct(weight_normalize=False, **recon_kwargs) + kwargs = dict( + deconvolution_kernel="prlx", + parallax_flip_phase=False, + override_aberration_coefs={"C10": defocus}, + verbose=False, + ) + kwargs.update(shared) + fourier.reconstruct(**kwargs) + montage.reconstruct(weight_normalize=False, **montage_only, **kwargs) - assert _relative_error(montage.obj, fourier.obj) < 1e-4 + assert montage.obj.shape == fourier.obj.shape, f"{shared=}" + assert _relative_error(montage.obj, fourier.obj) < 1e-4, f"{shared=}" - def test_weight_normalization_differs_only_by_a_constant(self, dataset4d): + def test_weight_normalization_differs_only_by_a_constant(self, fourier_montage): """On a full grid at U=1 the accumulated weight is exactly num_bf everywhere.""" - defocus = integer_shift_defocus(1) - _, montage = _build_pair(dataset4d, defocus) + _, montage = fourier_montage montage.reconstruct( deconvolution_kernel="prlx", @@ -329,15 +256,14 @@ def test_weight_normalization_differs_only_by_a_constant(self, dataset4d): rescaled = normalized * montage.num_bf / bf_weights assert _relative_error(rescaled, unnormalized) < 1e-4 - def test_residual_loss_tracks_the_fourier_one(self, dataset4d): + def test_residual_loss_tracks_the_fourier_one(self, fourier_montage): """Both measure the spread of the parallax-corrected images about their mean. At the integer-shift defocus the montage's nearest-pixel deposit is the exact shift, so the two agree up to the frequencies past the band limit that only the montage counts. ``band_limit=False`` because the two classes filter at different stages. """ - defocus = integer_shift_defocus(1) - fourier, montage = _build_pair(dataset4d, defocus) + fourier, montage = fourier_montage recon_kwargs = dict( deconvolution_kernel="prlx", @@ -350,10 +276,9 @@ def test_residual_loss_tracks_the_fourier_one(self, dataset4d): assert montage.residual_loss() == pytest.approx(fourier.residual_loss(), rel=5e-3) - def test_bf_mask_subsets_are_additive(self, dataset4d): + def test_bf_mask_subsets_are_additive(self, fourier_montage): """Checkerboard half-sets, rescaled by their BF weights, sum to the whole.""" - defocus = integer_shift_defocus(1) - _, montage = _build_pair(dataset4d, defocus) + _, montage = fourier_montage recon_kwargs = dict( deconvolution_kernel="prlx", @@ -371,18 +296,9 @@ def test_bf_mask_subsets_are_additive(self, dataset4d): assert _relative_error(halves[0] + halves[1], full) < 1e-4 - @pytest.mark.parametrize("kernel", ["ssb", "obf"]) - def test_accepts_the_deconvolution_kernels(self, dataset4d, kernel): - """Available as truncated real-space convolutions; see `TestRealSpaceKernels`.""" - _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) - - obj = montage.reconstruct(deconvolution_kernel=kernel, stencil_radius=6, verbose=False).obj - - assert np.isfinite(obj).all() - - def test_icom_matches_the_fourier_class(self, dataset4d): + def test_icom_matches_the_fourier_class(self, fourier_montage): """iCoM has no compact real-space form, but the FFT route never truncates.""" - fourier, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + fourier, montage = fourier_montage fourier.reconstruct(deconvolution_kernel="icom", verbose=False) montage.reconstruct(deconvolution_kernel="icom", convolution_mode="fft", verbose=False) @@ -392,10 +308,9 @@ def test_icom_matches_the_fourier_class(self, dataset4d): class TestRotationConvention: """`spatial_frequencies` rotates the k-grid; shifts must follow it, unflipped.""" - def test_ninety_degrees_rotates_the_shifts(self, dataset4d): - defocus = integer_shift_defocus(1) - _, montage = _build_pair(dataset4d, defocus) - coefs = {"C10": defocus} + def test_ninety_degrees_rotates_the_shifts(self, fourier_montage): + _, montage = fourier_montage + coefs = {"C10": integer_shift_defocus(1)} unrotated, _ = montage._return_shifts_px(0.0, coefs, montage.bf_mask, 1) rotated, _ = montage._return_shifts_px(90.0, coefs, montage.bf_mask, 1) @@ -405,63 +320,22 @@ def test_ninety_degrees_rotates_the_shifts(self, dataset4d): expected = torch.stack((-unrotated[:, 1], unrotated[:, 0]), dim=-1) assert torch.allclose(rotated, expected, atol=1e-4) - def test_rotation_changes_the_reconstruction(self, dataset4d): - defocus = integer_shift_defocus(1) - _, montage = _build_pair(dataset4d, defocus) - - montage.reconstruct(override_rotation_angle=0.0, verbose=False) - unrotated = montage.obj.copy() - montage.reconstruct(override_rotation_angle=30.0, verbose=False) - - assert not np.allclose(unrotated, montage.obj) - class TestReconstructDefaults: """The defaults are behaviour; pin them.""" - def test_interpolation_defaults_to_nearest(self, dataset4d): - _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) - - default = montage.reconstruct(upsampling_factor=2, verbose=False).obj.copy() - nearest = montage.reconstruct( - upsampling_factor=2, interpolation="nearest", verbose=False - ).obj - - assert np.array_equal(default, nearest) - - def test_nearest_is_a_roll_of_the_bright_field_images(self, dataset4d): - """On a raster scan, snapping moves every position of a BF image by one integer. - - `positions_px * U` is an exact integer, so `round(n + s) == n + round(s)`. - """ - _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) - upsampling_factor = 4 - shifts, _ = montage._return_shifts_px( - 0.0, montage.aberration_coefs, montage.bf_mask, upsampling_factor - ) - positions = montage.positions_px * upsampling_factor - coords = positions[None] + shifts[:, None] - - offsets = coords.round() - positions[None] - # every position of a given BF image moves by the same integer - assert torch.equal(offsets, shifts.round()[:, None].expand_as(offsets)) - - def test_gridded_constructors_flag_the_scan_as_gridded(self, dataset4d): - _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) - assert montage.gridded_scan is True - - def test_weight_normalize_defaults_off_for_a_raster_scan(self, dataset4d): + def test_weight_normalize_defaults_on_only_for_an_ungridded_scan( + self, dataset4d, fourier_montage + ): """Uniform density needs no correction, and normalizing amplifies edge noise.""" - _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) - + _, montage = fourier_montage + assert montage.gridded_scan is True default = montage.reconstruct(boundary="pad", verbose=False).obj.copy() unnormalized = montage.reconstruct( boundary="pad", weight_normalize=False, verbose=False ).obj - assert np.array_equal(default, unnormalized) - def test_weight_normalize_defaults_on_for_an_ungridded_scan(self, dataset4d): dataset3d, positions = TestNonGridScan._dataset3d_and_positions(dataset4d) recon = DirectPtychographyMontage.from_dataset3d( dataset3d, @@ -469,7 +343,6 @@ def test_weight_normalize_defaults_on_for_an_ungridded_scan(self, dataset4d): scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), **_common_kwargs(integer_shift_defocus(1)), ) - assert recon.gridded_scan is False default = recon.reconstruct(verbose=False).obj.copy() normalized = recon.reconstruct(weight_normalize=True, verbose=False).obj @@ -479,9 +352,9 @@ def test_weight_normalize_defaults_on_for_an_ungridded_scan(self, dataset4d): class TestPadBoundary: """`boundary="pad"` grows the canvas instead of wrapping.""" - def test_interior_matches_wrap(self, dataset4d): + def test_interior_matches_wrap(self, fourier_montage): defocus = integer_shift_defocus(1) - _, montage = _build_pair(dataset4d, defocus) + _, montage = fourier_montage # the filter is an FFT over the whole canvas, and "pad" makes that canvas larger, so # filtering then cropping is not the same operation as filtering the wrapped canvas. @@ -513,10 +386,13 @@ def test_interior_matches_wrap(self, dataset4d): assert interior_pad.shape == interior_wrap.shape assert _relative_error(interior_pad, interior_wrap) < 1e-5 - def test_canvas_covers_the_shifted_positions(self, dataset4d): - defocus = integer_shift_defocus(2) - _, montage = _build_pair(dataset4d, defocus) - montage.reconstruct(boundary="pad", verbose=False) + def test_canvas_covers_the_shifted_positions(self, fourier_montage): + _, montage = fourier_montage + montage.reconstruct( + boundary="pad", + override_aberration_coefs={"C10": integer_shift_defocus(2)}, + verbose=False, + ) # nothing was dropped: every (BF pixel, position) pair landed on the canvas assert float(montage.weights.sum()) == pytest.approx( @@ -524,8 +400,8 @@ def test_canvas_covers_the_shifted_positions(self, dataset4d): ) assert montage.obj.shape[0] > N and montage.obj.shape[1] > N - def test_pad_px_freezes_the_canvas(self, dataset4d): - _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + def test_pad_px_freezes_the_canvas(self, fourier_montage): + _, montage = fourier_montage shapes = set() for pixel_shift in (1, 2): @@ -555,13 +431,11 @@ def _dataset3d_and_positions(dataset4d): positions = scan_positions_px()[: n_scan * n_scan] * SCAN_SAMPLING return dataset3d, positions - def test_matches_from_dataset4d(self, dataset4d): + def test_matches_from_dataset4d(self, dataset4d, fourier_montage): defocus = integer_shift_defocus(1) dataset3d, positions = self._dataset3d_and_positions(dataset4d) - gridded = DirectPtychographyMontage.from_dataset4d( - dataset4d, edge_blend_pixels=0, boundary="wrap", **_common_kwargs(defocus) - ) + _, gridded = fourier_montage ungridded = DirectPtychographyMontage.from_dataset3d( dataset3d, positions, @@ -582,111 +456,6 @@ def test_matches_from_dataset4d(self, dataset4d): assert ungridded.obj.shape == gridded.obj.shape assert _relative_error(ungridded.obj, gridded.obj) < 1e-5 - def test_position_axis_order_is_row_col(self, dataset4d): - """Swapping the position columns must transpose the reconstruction.""" - defocus = integer_shift_defocus(1) - dataset3d, positions = self._dataset3d_and_positions(dataset4d) - - recon_kwargs = dict( - deconvolution_kernel="prlx", - parallax_flip_phase=False, - weight_normalize=False, - verbose=False, - ) - common = dict( - scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), - boundary="wrap", - **_common_kwargs(defocus), - ) - - straight = DirectPtychographyMontage.from_dataset3d(dataset3d, positions, **common) - swapped = DirectPtychographyMontage.from_dataset3d( - dataset3d, positions[:, ::-1].copy(), **common - ) - straight.reconstruct(**recon_kwargs) - swapped.reconstruct(**recon_kwargs) - - assert not np.allclose(straight.obj, swapped.obj) - - def test_scattered_positions_reconstruct(self, dataset4d): - """A jittered, shuffled scan still produces a supported montage.""" - defocus = integer_shift_defocus(1) - dataset3d, positions = self._dataset3d_and_positions(dataset4d) - - rng = np.random.default_rng(0) - order = rng.permutation(len(positions)) - jittered = positions[order] + rng.normal( - scale=0.25 * SCAN_SAMPLING, size=(len(positions), 2) - ) - shuffled = Dataset3d.from_array( - dataset3d.array[order], - name="shuffled", - units=dataset3d.units, - sampling=dataset3d.sampling, - ) - - recon = DirectPtychographyMontage.from_dataset3d( - shuffled, - jittered, - scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), - boundary="pad", - **_common_kwargs(defocus), - ) - recon.reconstruct(parallax_flip_phase=False, verbose=False) - - assert np.isfinite(recon.obj).all() - assert float(recon.weights.max()) > 0 - - def test_auto_scan_sampling_warns_and_infers(self, dataset4d): - dataset3d, positions = self._dataset3d_and_positions(dataset4d) - - with pytest.warns(UserWarning, match="Inferred scan_sampling"): - recon = DirectPtychographyMontage.from_dataset3d( - dataset3d, - positions, - scan_sampling="auto", - **_common_kwargs(integer_shift_defocus(1)), - ) - - assert recon.scan_sampling[0] == pytest.approx(SCAN_SAMPLING, rel=1e-6) - - def test_accepts_a_dataset2d_of_positions(self, dataset4d): - dataset3d, positions = self._dataset3d_and_positions(dataset4d) - positions_dataset = Dataset2d.from_array(positions, name="positions", units=("A", "A")) - - recon = DirectPtychographyMontage.from_dataset3d( - dataset3d, - positions_dataset, - scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), - **_common_kwargs(integer_shift_defocus(1)), - ) - assert recon.num_positions == len(positions) - - def test_rejects_positions_in_the_wrong_units(self, dataset4d): - dataset3d, positions = self._dataset3d_and_positions(dataset4d) - positions_dataset = Dataset2d.from_array( - positions, name="positions", units=("pixels", "pixels") - ) - - with pytest.raises(ValueError, match="must be given in 'A'"): - DirectPtychographyMontage.from_dataset3d( - dataset3d, - positions_dataset, - scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), - **_common_kwargs(integer_shift_defocus(1)), - ) - - def test_rejects_mismatched_position_count(self, dataset4d): - dataset3d, positions = self._dataset3d_and_positions(dataset4d) - - with pytest.raises(ValueError, match="rows but `dataset` has"): - DirectPtychographyMontage.from_dataset3d( - dataset3d, - positions[:-1], - scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), - **_common_kwargs(integer_shift_defocus(1)), - ) - def test_requires_an_explicit_rotation_angle(self, dataset4d): dataset3d, positions = self._dataset3d_and_positions(dataset4d) kwargs = _common_kwargs(integer_shift_defocus(1)) @@ -702,27 +471,9 @@ def test_requires_an_explicit_rotation_angle(self, dataset4d): class TestHyperparameterSearch: - def test_grid_search_recovers_the_seeded_defocus(self, dataset4d): + def test_residual_loss_is_minimized_at_the_true_defocus(self, fourier_montage): true_defocus = integer_shift_defocus(1) - _, montage = _build_pair(dataset4d, true_defocus) - - montage.grid_search_hyperparameters( - aberration_coefs={ - "C10": OptimizationParameter( - low=0.4 * true_defocus, high=1.6 * true_defocus, n_points=7 - ) - }, - parallax_flip_phase=False, - verbose=False, - ) - best = montage.hyperparameter_state.optimized_aberrations["C10"] - - step = (1.6 - 0.4) * true_defocus / 6 - assert abs(best - true_defocus) <= step - - def test_residual_loss_is_minimized_at_the_true_defocus(self, dataset4d): - true_defocus = integer_shift_defocus(1) - _, montage = _build_pair(dataset4d, true_defocus) + _, montage = fourier_montage losses = {} for scale in (0.5, 0.8, 1.0, 1.2, 1.5): @@ -775,7 +526,7 @@ class TestSparseEvents: """ @staticmethod - def _counts(dataset4d, scale=40.0): + def _counts(dataset4d, scale=4.0): """Integer counts, in the ``(N_pos, N, N)`` layout `from_dataset3d` takes.""" frames = to_numpy(dataset4d.array).reshape(-1, N, N) frames = frames - frames.min() @@ -869,24 +620,34 @@ def test_icom_matches_dense(self, pair, convolution_mode): ) assert _relative_error(b, a) < 1e-4 - @pytest.mark.parametrize("boundary, interpolation", [("wrap", "nearest"), ("pad", "bilinear")]) - def test_boundaries_and_interpolation_match_dense( - self, dataset4d, pair, boundary, interpolation - ): - if boundary == "wrap": - dense, sparse = pair - else: - dense, sparse = self._pair(dataset4d, integer_shift_defocus(1), boundary=boundary) + def test_pad_and_bilinear_match_dense(self, pair): + dense, sparse = pair a, b = self._both( dense, sparse, deconvolution_kernel="ssb", convolution_mode="stencil", stencil_radius=4, - interpolation=interpolation, + boundary="pad", + interpolation="bilinear", ) assert _relative_error(b, a) < 1e-4 + def test_complex_objects_match_dense(self, pair): + dense, sparse = pair + for kernel in ("obf", "icom"): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + a, b = self._both( + dense, + sparse, + deconvolution_kernel=kernel, + obj_type="complex", + convolution_mode="stencil", + stencil_radius=4, + ) + assert _relative_error(b, a) < 1e-4, f"{kernel=}" + def test_scattered_positions_match_dense(self, dataset4d): """An ungridded scan, where the coverage the scan mean rides on is not constant.""" counts = self._counts(dataset4d) @@ -1136,117 +897,114 @@ def test_rejects_a_vector_without_detector_fields(self): **_common_kwargs(0.0), ) - def test_dense_and_sparse_are_mutually_exclusive(self, pair): - dense, sparse = pair - assert sparse.sparse_frames - assert sparse.vbf_stack is None - assert sparse.events.num_positions == sparse.num_positions - - assert not dense.sparse_frames - assert dense.events is None + def test_survives_a_round_trip(self, pair, tmp_path): + _, sparse = pair + state = sparse.hyperparameter_state + state.optimized_aberrations = {"C10": 123.0} + state.optimized_rotation_angle = 7.5 + try: + sparse.reconstruct(deconvolution_kernel="prlx", verbose=False) + before = sparse.obj.copy() + + path = str(tmp_path / "sparse.zip") + sparse.save(path, mode="o") + restored = load(path) + + assert type(restored) is DirectPtychographyMontage + assert restored.sparse_frames + assert np.array_equal(restored.obj, before) + assert restored.hyperparameter_state.optimized_aberrations == {"C10": 123.0} + assert restored.hyperparameter_state.optimized_rotation_angle == 7.5 + assert restored.residual_loss() == pytest.approx(sparse.residual_loss()) + + restored.reconstruct(deconvolution_kernel="prlx", verbose=False) + assert np.allclose(restored.obj, before) + finally: + state.clear_optimized() class TestSerialization: - """Both classes must survive a save/load round-trip and stay usable afterwards.""" - - @pytest.mark.parametrize("cls_name", ["fourier", "montage", "sparse"]) - def test_round_trip(self, dataset4d, tmp_path, cls_name): - defocus = integer_shift_defocus(1) - if cls_name == "sparse": - recon = TestSparseEvents._pair(dataset4d, defocus)[1] - else: - fourier, montage = _build_pair(dataset4d, defocus) - recon = fourier if cls_name == "fourier" else montage - + def test_round_trip(self, dataset4d, tmp_path): + """Every setting a montage stores must survive save/load and reconstruct again.""" + recon = TestSharedCanvas._montage( + dataset4d, + integer_shift_defocus(1), + position_offset=(11.0, -3.0), + defocus_gradient=(7.0, -3.0), + ) recon.hyperparameter_state.optimized_aberrations = {"C10": 123.0} recon.hyperparameter_state.optimized_rotation_angle = 7.5 - recon.reconstruct(deconvolution_kernel="prlx", verbose=False) - before = recon.obj.copy() - - path = str(tmp_path / f"{cls_name}.zip") - recon.save(path, mode="o") - restored = load(path) - - assert type(restored) is type(recon) - assert np.array_equal(restored.obj, before) - assert restored.hyperparameter_state.optimized_aberrations == {"C10": 123.0} - assert restored.hyperparameter_state.optimized_rotation_angle == 7.5 - assert restored.gpts == recon.gpts - assert restored.residual_loss() == pytest.approx(recon.residual_loss()) - - # and it must still be able to reconstruct - restored.reconstruct(deconvolution_kernel="prlx", verbose=False) - assert np.allclose(restored.obj, before) - - def test_torch_size_attributes_round_trip_as_tuples(self, dataset4d): - """`torch.Size` is a tuple subclass; AutoSerialize used to choke on the subclass name.""" - _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) - - assert type(montage.gpts) is tuple - assert type(montage.scan_gpts) is tuple + _, empirical = TestFourierProbe._pair(dataset4d, DirectPtychographyMontage) + + for name, montage, kwargs in ( + ("ungridded", recon, dict(deconvolution_kernel="prlx")), + ("empirical", empirical, dict(deconvolution_kernel="ssb", stencil_radius=5)), + ): + montage.reconstruct(verbose=False, **kwargs) + before = montage.obj.copy() + + path = str(tmp_path / f"{name}.zip") + montage.save(path, mode="o") + restored = load(path) + + assert type(restored) is DirectPtychographyMontage, name + assert np.array_equal(restored.obj, before), name + assert restored.gpts == montage.gpts, name + restored.reconstruct(verbose=False, **kwargs) + assert np.allclose(restored.obj, before), name + + if name == "ungridded": + state = restored.hyperparameter_state + assert state.optimized_aberrations == {"C10": 123.0} + assert state.optimized_rotation_angle == 7.5 + assert restored.residual_loss() == pytest.approx(montage.residual_loss()) + assert restored.defocus_gradient == (7.0, -3.0) + assert restored.scan_origin == pytest.approx(montage.scan_origin) + assert restored.obj_origin == pytest.approx(montage.obj_origin) + else: + assert restored.fourier_probe is not None class TestVisualization: """The object sampling, and hence `visualize`'s scalebar, must follow the upsampling.""" - @pytest.mark.parametrize("cls_name", ["fourier", "montage"]) - @pytest.mark.parametrize("upsampling_factor", [1, 2, 3]) - def test_scalebar_follows_the_upsampling(self, dataset4d, cls_name, upsampling_factor): - fourier, montage = _build_pair(dataset4d, integer_shift_defocus(1, upsampling_factor)) - recon = fourier if cls_name == "fourier" else montage - recon.reconstruct( - deconvolution_kernel="prlx", upsampling_factor=upsampling_factor, verbose=False - ) - - expected = SCAN_SAMPLING / upsampling_factor - assert recon._obj_sampling[0] == pytest.approx(expected) - # the reported sampling must span the same field of view as the image itself - assert recon.obj.shape[0] * recon._obj_sampling[0] == pytest.approx(N * SCAN_SAMPLING) - - def test_sampling_resets_between_reconstructions(self, dataset4d): - _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) - - montage.reconstruct(upsampling_factor=3, verbose=False) - assert montage._obj_sampling[0] == pytest.approx(SCAN_SAMPLING / 3) - montage.reconstruct(upsampling_factor=1, verbose=False) - assert montage._obj_sampling[0] == pytest.approx(SCAN_SAMPLING) - - def test_sampling_is_defined_before_reconstructing(self, dataset4d): - fourier, montage = _build_pair(dataset4d, integer_shift_defocus(1)) - - for recon in (fourier, montage): - assert recon._obj_sampling[0] == pytest.approx(SCAN_SAMPLING) - - def test_fov_matches_the_scan_extent(self, dataset4d): - fourier, montage = _build_pair(dataset4d, integer_shift_defocus(1)) - - for recon in (fourier, montage): - assert recon.fov == pytest.approx((N * SCAN_SAMPLING, N * SCAN_SAMPLING)) - - @pytest.mark.parametrize("upsampling_factor", [1, 2]) - def test_padded_canvas_keeps_the_scan_sampling(self, dataset4d, upsampling_factor): + def test_scalebar_follows_the_upsampling(self, dataset4d): + for recon in _build_pair(dataset4d, integer_shift_defocus(1, 2)): + name = type(recon).__name__ + assert recon._obj_sampling[0] == pytest.approx(SCAN_SAMPLING), name + assert recon.fov == pytest.approx((N * SCAN_SAMPLING, N * SCAN_SAMPLING)), name + assert recon.scan_origin == (0.0, 0.0), name + + recon.reconstruct(deconvolution_kernel="prlx", upsampling_factor=2, verbose=False) + assert recon._obj_sampling[0] == pytest.approx(SCAN_SAMPLING / 2), name + # the reported sampling must span the same field of view as the image itself + assert recon.obj.shape[0] * recon._obj_sampling[0] == pytest.approx( + N * SCAN_SAMPLING + ), name + assert recon._obj_fov == pytest.approx(recon.fov), name + + def test_padded_canvas_keeps_the_scan_sampling(self, fourier_montage): """A padded canvas spans more than the scan, so its sampling is not fov/shape. Deriving from the *scan* field of view would under-report the pixel size by the padding fraction; `_obj_fov` reports the canvas extent instead. """ - _, montage = _build_pair(dataset4d, integer_shift_defocus(1, upsampling_factor)) - montage.reconstruct(boundary="pad", upsampling_factor=upsampling_factor, verbose=False) + _, montage = fourier_montage + montage.reconstruct( + boundary="pad", + upsampling_factor=2, + override_aberration_coefs={"C10": integer_shift_defocus(1, 2)}, + verbose=False, + ) - assert montage.obj.shape[0] > N * upsampling_factor # canvas really did grow - assert montage._obj_sampling[0] == pytest.approx(SCAN_SAMPLING / upsampling_factor) + assert montage.obj.shape[0] > N * 2 # canvas really did grow + assert montage._obj_sampling[0] == pytest.approx(SCAN_SAMPLING / 2) assert montage._obj_fov[0] > montage.fov[0] # and the reported extent still matches the image it describes assert montage.obj.shape[0] * montage._obj_sampling[0] == pytest.approx( montage._obj_fov[0] ) - def test_wrapped_canvas_spans_exactly_the_scan(self, dataset4d): - _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) - montage.reconstruct(boundary="wrap", upsampling_factor=2, verbose=False) - - assert montage._obj_fov == pytest.approx(montage.fov) - def test_visualize_before_reconstruct_raises(self, dataset4d): _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) with pytest.raises(RuntimeError, match="Run reconstruct"): @@ -1258,7 +1016,8 @@ class TestFourierProbe: DEFOCUS = staticmethod(lambda: integer_shift_defocus(1)) - def _pair(self, dataset4d, cls, **probe_kwargs): + @staticmethod + def _pair(dataset4d, cls, **probe_kwargs): """The same reconstruction twice: analytic, and with its own probe fed back in.""" defocus = integer_shift_defocus(1) kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0) @@ -1273,30 +1032,22 @@ def _pair(self, dataset4d, cls, **probe_kwargs): empirical = cls.from_dataset4d(dataset4d, fourier_probe=probe, **kwargs) return analytic, empirical - @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) - @pytest.mark.parametrize("kernel", ["ssb", "obf"]) - def test_empirical_probe_reproduces_the_analytic_one(self, dataset4d, cls, kernel): + def test_empirical_probe_reproduces_the_analytic_one(self, dataset4d): """The headline: the same probe, described two ways, must reconstruct the same. This is what makes the empirical path trustworthy -- there is no ground truth for a measured probe, so the only check available is that a probe we *can* write down - analytically goes through the array path unchanged. + analytically goes through the array path unchanged. `TestEmpiricalProbe` checks the + Fourier class. """ - analytic, empirical = self._pair(dataset4d, cls, normalize=False) - extra = {} if cls is DirectPtychography else {"stencil_radius": 5} - - analytic.reconstruct(deconvolution_kernel=kernel, verbose=False, **extra) - expected = analytic.obj.copy() - empirical.reconstruct(deconvolution_kernel=kernel, verbose=False, **extra) - - assert _relative_error(empirical.obj, expected) < 1e-5 + analytic, empirical = self._pair(dataset4d, DirectPtychographyMontage, normalize=False) - def test_is_analytic_reports_which_path_is_live(self, dataset4d): - analytic, empirical = self._pair(dataset4d, DirectPtychographyMontage) + for kernel in ("ssb", "obf"): + analytic.reconstruct(deconvolution_kernel=kernel, stencil_radius=5, verbose=False) + expected = analytic.obj.copy() + empirical.reconstruct(deconvolution_kernel=kernel, stencil_radius=5, verbose=False) - assert analytic.fourier_probe is None - assert empirical.fourier_probe is not None - assert empirical.fourier_probe.is_analytic is False + assert _relative_error(empirical.obj, expected) < 1e-5, f"{kernel=}" def test_zero_outside_the_detector(self, dataset4d): """Beyond the detector's Nyquist nothing was measured, so `psi` is zero there. @@ -1397,8 +1148,7 @@ def test_refined_sampling_converges_on_the_exact_shift(self): assert errors[0] > errors[1] > errors[2] assert errors[2] < 0.01 - @pytest.mark.parametrize("rotation_angle", [15.0, -30.0]) - def test_rotation_off_the_lattice_is_supported(self, dataset4d, rotation_angle): + def test_rotation_off_the_lattice_is_supported(self, dataset4d): """A rotation carries `q` off the probe's lattice, which is interpolated, not refused. It is a real experimental parameter. What it costs is a sampling error that the @@ -1413,7 +1163,7 @@ def test_rotation_off_the_lattice_is_supported(self, dataset4d, rotation_angle): kwargs = dict( _common_kwargs(defocus), edge_blend_pixels=0, boundary="wrap", intensity_threshold=0.5 ) - kwargs["rotation_angle"] = rotation_angle + kwargs["rotation_angle"] = 15.0 analytic = DirectPtychographyMontage.from_dataset4d(dataset4d, **kwargs) psi = analytic_probe_array(analytic, {"C10": defocus}) @@ -1506,43 +1256,22 @@ def test_aberration_only_features_raise(self, dataset4d, kwargs, match): with pytest.raises(NotImplementedError, match=match): empirical.reconstruct(verbose=False, override_aberration_coefs={"C10": 0.0}, **kwargs) - @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) - def test_prlx_runs_on_a_modelled_empirical_probe(self, dataset4d, cls): + def test_prlx_runs_on_a_modelled_empirical_probe(self, dataset4d): """Coefficients modelling a measured probe's phase are enough to run parallax. The probe here *is* the analytic one materialized as an array, so the model is exact and the two paths must agree. Whether a real measured probe is this well modelled is - what ``fit_aberrations_from_probe`` is for. + what ``fit_aberrations_from_probe`` is for; `test_fitted_coefficients_drive_parallax` + runs the montage on the fitted coefficients. """ - analytic, empirical = self._pair(dataset4d, cls, normalize=False) - extra = {} if cls is DirectPtychography else {"boundary": "wrap"} + analytic, empirical = self._pair(dataset4d, DirectPtychography, normalize=False) - analytic.reconstruct(deconvolution_kernel="prlx", verbose=False, **extra) + analytic.reconstruct(deconvolution_kernel="prlx", verbose=False) expected = analytic.obj.copy() - empirical.reconstruct(deconvolution_kernel="prlx", verbose=False, **extra) + empirical.reconstruct(deconvolution_kernel="prlx", verbose=False) assert correlation(to_numpy(empirical.obj), to_numpy(expected)) > 0.999 - def test_semiangle_cutoff_becomes_optional(self, dataset4d): - """The empirical probe carries its own aperture, whatever shape it is.""" - defocus = integer_shift_defocus(1) - kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0, boundary="wrap") - analytic = DirectPtychographyMontage.from_dataset4d(dataset4d, **kwargs) - psi = analytic_probe_array(analytic, {"C10": defocus}) - - kwargs["semiangle_cutoff"] = None - empirical = DirectPtychographyMontage.from_dataset4d( - dataset4d, - fourier_probe=FourierProbe.from_array( - psi, analytic.reciprocal_sampling, analytic.wavelength - ), - **kwargs, - ) - - assert empirical.semiangle_cutoff is None - empirical.reconstruct(deconvolution_kernel="ssb", stencil_radius=5, verbose=False) - assert np.isfinite(empirical.obj).all() - def test_normalize_gives_unit_intensity(self, dataset4d): _, empirical = self._pair(dataset4d, DirectPtychographyMontage, normalize=True) @@ -1568,25 +1297,11 @@ def test_shape_must_match_the_detector(self, dataset4d): with pytest.raises(ValueError, match="detector grid"): analytic.fourier_probe = wrong - def test_survives_a_round_trip(self, dataset4d, tmp_path): - _, empirical = self._pair(dataset4d, DirectPtychographyMontage) - empirical.reconstruct(deconvolution_kernel="ssb", stencil_radius=5, verbose=False) - before = empirical.obj.copy() - - path = str(tmp_path / "empirical.zip") - empirical.save(path, mode="o") - restored = load(path) - - assert restored.fourier_probe is not None - restored.reconstruct(deconvolution_kernel="ssb", stencil_radius=5, verbose=False) - assert np.allclose(restored.obj, before) - class TestWavelength: """`wavelength` given directly, for anything that is not an electron.""" - @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) - def test_energy_and_wavelength_agree(self, dataset4d, cls): + def test_energy_and_wavelength_agree(self, dataset4d): """The two routes must land on the same geometry, and hence the same image.""" common = dict( semiangle_cutoff=SEMIANGLE_CUTOFF, @@ -1596,8 +1311,10 @@ def test_energy_and_wavelength_agree(self, dataset4d, cls): edge_blend_pixels=0, verbose=False, ) - from_energy = cls.from_dataset4d(dataset4d, energy=PROBE_ENERGY, **common) - from_wavelength = cls.from_dataset4d( + from_energy = DirectPtychographyMontage.from_dataset4d( + dataset4d, energy=PROBE_ENERGY, **common + ) + from_wavelength = DirectPtychographyMontage.from_dataset4d( dataset4d, wavelength=from_energy.wavelength, **common ) @@ -1608,43 +1325,24 @@ def test_energy_and_wavelength_agree(self, dataset4d, cls): recon.reconstruct(deconvolution_kernel="prlx", verbose=False) assert np.array_equal(from_wavelength.obj, from_energy.obj) - @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) - @pytest.mark.parametrize( - "kwargs", - [ - {}, # neither - {"energy": PROBE_ENERGY, "wavelength": 0.02}, # both - ], - ) - def test_exactly_one_of_energy_or_wavelength(self, dataset4d, cls, kwargs): - with pytest.raises(ValueError, match="exactly one of `energy`"): - cls.from_dataset4d( - dataset4d, - semiangle_cutoff=SEMIANGLE_CUTOFF, - rotation_angle=0.0, - force_fitted_origin=ORIGIN, - verbose=False, - **kwargs, - ) - - def test_photon_wavelength_is_not_the_electron_one(self): - """The reason this exists: the de Broglie formula is wrong for photons. - - `electron_wavelength_angstrom` is relativistic and correct for electrons -- 0.019687 - Angstrom at 300 kV -- but it is the wrong physics for a 7.9 keV photon, where the - answer is `hc/E` = 1.5694 rather than 0.1375. - """ - from quantem.core.utils.utils import electron_wavelength_angstrom - - assert electron_wavelength_angstrom(300e3) == pytest.approx(0.019687, abs=1e-6) - - photon_energy_ev = 7900.08 - assert electron_wavelength_angstrom(photon_energy_ev) == pytest.approx(0.1375, abs=1e-3) - assert 12398.42 / photon_energy_ev == pytest.approx(1.5694, abs=1e-3) - - @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) - def test_wavelength_survives_a_round_trip(self, dataset4d, cls, tmp_path): - recon = cls.from_dataset4d( + def test_exactly_one_positive_value_is_given(self, dataset4d): + common = dict( + semiangle_cutoff=SEMIANGLE_CUTOFF, + rotation_angle=0.0, + force_fitted_origin=ORIGIN, + verbose=False, + ) + for kwargs, match in ( + ({}, "exactly one of `energy`"), + ({"energy": PROBE_ENERGY, "wavelength": 0.02}, "exactly one of `energy`"), + ({"wavelength": 0.0}, "wavelength must be greater than 0"), + ({"wavelength": -1.0}, "wavelength must be greater than 0"), + ): + with pytest.raises(ValueError, match=match): + DirectPtychographyMontage.from_dataset4d(dataset4d, **common, **kwargs) + + def test_wavelength_survives_a_round_trip(self, dataset4d, tmp_path): + recon = DirectPtychographyMontage.from_dataset4d( dataset4d, wavelength=1.5694, semiangle_cutoff=SEMIANGLE_CUTOFF, @@ -1655,37 +1353,24 @@ def test_wavelength_survives_a_round_trip(self, dataset4d, cls, tmp_path): ) recon.reconstruct(deconvolution_kernel="prlx", verbose=False) - path = str(tmp_path / f"{cls.__name__}.zip") + path = str(tmp_path / "wavelength.zip") recon.save(path, mode="o") restored = load(path) assert restored.wavelength == pytest.approx(1.5694) - @pytest.mark.parametrize("bad", [0.0, -1.0]) - def test_non_positive_wavelength_raises(self, dataset4d, bad): - with pytest.raises(ValueError): - DirectPtychographyMontage.from_dataset4d( - dataset4d, - wavelength=bad, - semiangle_cutoff=SEMIANGLE_CUTOFF, - rotation_angle=0.0, - force_fitted_origin=ORIGIN, - verbose=False, - ) - class TestSemiangleCutoff: """`semiangle_cutoff` sets the probe aperture and is never optional.""" - @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) - def test_from_virtual_bfs_requires_it(self, cls, dataset4d): + def test_is_required_and_positive(self, dataset4d, fourier_montage): """Enforced at runtime rather than by the signature. It used to be a required positional, but `wavelength` has to sit alongside `energy` with a default, and Python will not take a defaulted parameter before an undefaulted one. The guarantee is the error, so test the error. """ - _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + _, montage = fourier_montage vbf_dataset = Dataset3d.from_array( to_numpy(montage.vbf_stack).reshape(montage.num_bf, N, N), name="vBF stack", @@ -1698,9 +1383,8 @@ def test_from_virtual_bfs_requires_it(self, cls, dataset4d): units=("A^-1", "A^-1"), sampling=tuple(montage.reciprocal_sampling), ) - with pytest.raises(ValueError, match="`semiangle_cutoff` is required"): - cls.from_virtual_bfs( + DirectPtychographyMontage.from_virtual_bfs( vbf_dataset, bf_mask_dataset, energy=PROBE_ENERGY, @@ -1709,29 +1393,13 @@ def test_from_virtual_bfs_requires_it(self, cls, dataset4d): verbose=False, ) - @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) - def test_none_raises_a_clear_error(self, dataset4d, cls): + common = dict( + energy=PROBE_ENERGY, rotation_angle=0.0, force_fitted_origin=ORIGIN, verbose=False + ) with pytest.raises(ValueError, match="`semiangle_cutoff` is required"): - cls.from_dataset4d( - dataset4d, - energy=PROBE_ENERGY, - semiangle_cutoff=None, - rotation_angle=0.0, - force_fitted_origin=ORIGIN, - verbose=False, - ) - - @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) - def test_non_positive_raises(self, dataset4d, cls): - with pytest.raises(ValueError): - cls.from_dataset4d( - dataset4d, - energy=PROBE_ENERGY, - semiangle_cutoff=-1.0, - rotation_angle=0.0, - force_fitted_origin=ORIGIN, - verbose=False, - ) + DirectPtychographyMontage.from_dataset4d(dataset4d, semiangle_cutoff=None, **common) + with pytest.raises(ValueError, match="semiangle_cutoff must be greater than 0"): + DirectPtychographyMontage.from_dataset4d(dataset4d, semiangle_cutoff=-1.0, **common) class TestDefocusGradient: @@ -1757,11 +1425,10 @@ def test_none_and_zero_are_the_same_reconstruction(self, dataset4d): with_zero.reconstruct(verbose=False).obj, ) - def test_defocus_rate_is_the_analytic_lambda_k(self, dataset4d): - """`d shift / d C10 = wavelength * k`, independent of the other aberrations.""" - montage = DirectPtychographyMontage.from_dataset4d( - dataset4d, **_common_kwargs(integer_shift_defocus(1)) - ) + def test_defocus_rate_is_the_analytic_lambda_k(self, fourier_montage): + """`d shift / d C10 = wavelength * k`, independent of the other aberrations: chi is + linear in every magnitude, so the rate cannot depend on the rest.""" + _, montage = fourier_montage rate = montage._return_defocus_rate_px(0.0, montage.bf_mask, 1) kxa, kya = spatial_frequencies(montage.gpts, montage.sampling, device=montage.device) @@ -1773,37 +1440,21 @@ def test_defocus_rate_is_the_analytic_lambda_k(self, dataset4d): * montage.wavelength / scan_sampling ) - assert torch.allclose(rate, expected, atol=1e-9) - def test_defocus_rate_ignores_other_aberrations(self, dataset4d): - """chi is linear in every magnitude, so the rate cannot depend on the rest.""" - montage = DirectPtychographyMontage.from_dataset4d( - dataset4d, **_common_kwargs(integer_shift_defocus(1)) - ) - rate = montage._return_defocus_rate_px(0.0, montage.bf_mask, 1) - base = {"C10": 500.0, "C12": 40.0, "phi12": 0.7, "C30": 1.2e5} shifted = montage._return_shifts_px(0.0, {**base, "C10": 501.0}, montage.bf_mask, 1)[0] unshifted = montage._return_shifts_px(0.0, base, montage.bf_mask, 1)[0] - assert torch.allclose(shifted - unshifted, rate, atol=1e-6) - def test_delta_defocus_is_mean_zero(self, dataset4d): + def test_delta_defocus_is_mean_zero(self, fourier_montage): """Measuring from the centroid keeps the gradient orthogonal to the global C10.""" - montage = DirectPtychographyMontage.from_dataset4d( - dataset4d, **_common_kwargs(integer_shift_defocus(1)) - ) + _, montage = fourier_montage delta = montage._return_delta_c10((7.0, -3.0)) - assert delta is not None assert float(delta.mean().abs()) < 1e-9 assert float(delta.abs().max()) > 0 - - def test_zero_gradient_short_circuits(self, dataset4d): - montage = DirectPtychographyMontage.from_dataset4d( - dataset4d, **_common_kwargs(integer_shift_defocus(1)) - ) + # and no gradient short-circuits assert montage._return_delta_c10(None) is None assert montage._return_delta_c10((0.0, 0.0)) is None @@ -1823,11 +1474,9 @@ def test_padded_canvas_covers_the_gradient(self, dataset4d): assert tilted_shape[0] > flat_shape[0] assert tilted_shape[1] > flat_shape[1] - def test_shift_extrema_match_a_brute_force_scan(self, dataset4d): + def test_shift_extrema_match_a_brute_force_scan(self, fourier_montage): """The closed form must bound every (BF pixel, position) pair, with no slack.""" - montage = DirectPtychographyMontage.from_dataset4d( - dataset4d, **_common_kwargs(integer_shift_defocus(1)) - ) + _, montage = fourier_montage gradient = (30.0, -10.0) shifts = montage._return_shifts_px(0.0, montage.aberration_coefs, montage.bf_mask, 1)[0] rate = montage._return_defocus_rate_px(0.0, montage.bf_mask, 1) @@ -1861,8 +1510,8 @@ def corr(g): assert corr(gradient) > corr((-gradient[0], -gradient[1])) - @pytest.mark.parametrize("gradient", [(20.0, 0.0), (20.0, -10.0), (-15.0, 25.0)]) - def test_correcting_the_gradient_sharpens_the_reconstruction(self, gradient): + def test_correcting_the_gradient_sharpens_the_reconstruction(self): + gradient = (20.0, -10.0) montage, obj = _model_montage(gradient) reconstruct = dict(parallax_flip_phase=False, interpolation="bilinear", verbose=False) @@ -1903,41 +1552,38 @@ def test_defocus_map_is_flat_without_a_gradient(self): assert np.ptp(results["c10_best"]) < 0.05 * MODEL_DEFOCUS - @pytest.mark.parametrize("gradient", [(0.0, 0.0), (20.0, -10.0), (-15.0, 25.0)]) - def test_fit_defocus_gradient_recovers_the_seed(self, gradient): - montage, _ = _model_montage(gradient) - - montage.fit_defocus_gradient( - MODEL_C10_GRID, patch_grid=(3, 3), interpolation="bilinear", verbose=False - ) + def test_fit_defocus_gradient_recovers_the_seed(self): + # a zero gradient is the control + for gradient in ((0.0, 0.0), (20.0, -10.0)): + montage, _ = _model_montage(gradient) - assert montage.defocus_gradient is not None - scale = max(np.hypot(*gradient), 1.0) - assert np.hypot(*np.subtract(montage.defocus_gradient, gradient)) < 0.15 * scale - - def test_fit_defocus_gradient_updates_the_global_defocus(self): - montage, _ = _model_montage((20.0, 0.0)) - montage.hyperparameter_state.optimized_aberrations = {} - - montage.fit_defocus_gradient( - MODEL_C10_GRID, patch_grid=(3, 3), interpolation="bilinear", verbose=False - ) + montage.fit_defocus_gradient( + MODEL_C10_GRID, patch_grid=(3, 3), interpolation="bilinear", verbose=False + ) - assert "C10" in montage.hyperparameter_state.optimized_aberrations - assert "C10" in montage.hyperparameter_state.optimized_keys + assert montage.defocus_gradient is not None, f"{gradient=}" + scale = max(np.hypot(*gradient), 1.0) + error = np.hypot(*np.subtract(montage.defocus_gradient, gradient)) + assert error < 0.15 * scale, f"{gradient=}" - def test_fit_defocus_gradient_can_leave_the_defocus_alone(self): - montage, _ = _model_montage((20.0, 0.0)) + def test_fit_defocus_gradient_updates_the_global_defocus_on_request(self): + for update_defocus in (True, False): + montage, _ = _model_montage((20.0, 0.0)) - montage.fit_defocus_gradient( - MODEL_C10_GRID, - patch_grid=(3, 3), - interpolation="bilinear", - update_defocus=False, - verbose=False, - ) + montage.fit_defocus_gradient( + MODEL_C10_GRID, + patch_grid=(3, 3), + interpolation="bilinear", + update_defocus=update_defocus, + verbose=False, + ) - assert montage.hyperparameter_state.optimized_aberrations == {} + state = montage.hyperparameter_state + if update_defocus: + assert "C10" in state.optimized_aberrations + assert "C10" in state.optimized_keys + else: + assert state.optimized_aberrations == {} def test_endpoint_pinned_patches_are_invalid(self): """A grid that does not bracket the local defocus must be reported, not fitted.""" @@ -1964,25 +1610,17 @@ def test_fit_raises_when_too_few_patches_bracket(self): verbose=False, ) - def test_defocus_map_allows_a_one_dimensional_grid(self): - """A (P, 1) grid is a profile along one axis -- only the plane fit needs three.""" + def test_defocus_map_validates_its_grids(self): montage, _ = _model_montage((20.0, 0.0)) + # a (P, 1) grid is a profile along one axis -- only the plane fit needs three results = montage.defocus_map( MODEL_C10_GRID, patch_grid=(3, 1), interpolation="bilinear", verbose=False ) - assert results["c10_best"].shape == (3,) - def test_defocus_map_rejects_a_degenerate_grid(self): - montage, _ = _model_montage((0.0, 0.0)) - with pytest.raises(ValueError, match="must be positive"): montage.defocus_map(MODEL_C10_GRID, patch_grid=(0, 3), verbose=False) - - def test_defocus_map_needs_enough_trial_values(self): - montage, _ = _model_montage((0.0, 0.0)) - with pytest.raises(ValueError, match="at least 3 points"): montage.defocus_map([2000.0, 3000.0], patch_grid=(2, 2), verbose=False) @@ -2013,15 +1651,6 @@ def test_rejects_a_malformed_gradient(self, dataset4d): with pytest.raises(ValueError, match="must be a \\(row, col\\) pair"): montage.defocus_gradient = (1.0, 2.0, 3.0) - def test_survives_a_serialization_round_trip(self, dataset4d, tmp_path): - montage = DirectPtychographyMontage.from_dataset4d( - dataset4d, defocus_gradient=(7.0, -3.0), **_common_kwargs(integer_shift_defocus(1)) - ) - path = tmp_path / "montage.zip" - montage.save(path, mode="o") - - assert load(path).defocus_gradient == (7.0, -3.0) - class TestRealSpaceKernels: """SSB / OBF as truncated real-space convolutions. @@ -2146,6 +1775,10 @@ def test_warns_when_the_tolerance_cannot_be_met(self, dataset4d): def test_auto_respects_the_radius_cap(self, dataset4d): _, montage = self._pair(dataset4d) + montage.reconstruct(deconvolution_kernel="ssb", convolution_mode="stencil", verbose=False) + info = montage._stencil_info + assert info["stencil_radius"] >= 1 + assert 0.0 <= info["mean_error"] <= info["max_error"] with pytest.warns(UserWarning, match="truncation error"): montage.reconstruct( @@ -2158,40 +1791,16 @@ def test_auto_respects_the_radius_cap(self, dataset4d): assert montage._stencil_info["stencil_radius"] <= 3 - def test_auto_reports_what_it_chose(self, dataset4d): - _, montage = self._pair(dataset4d) - montage.reconstruct(deconvolution_kernel="ssb", convolution_mode="stencil", verbose=False) - - info = montage._stencil_info - assert info["stencil_radius"] >= 1 - assert 0.0 <= info["mean_error"] <= info["max_error"] - - @pytest.mark.parametrize("kernel", KERNELS) - def test_residual_loss_is_undefined(self, dataset4d, kernel): + def test_residual_loss_is_only_defined_for_parallax(self, dataset4d): + """`TestHyperparameterSearch` checks the same for `ssb`.""" _, montage = self._pair(dataset4d) - montage.reconstruct(deconvolution_kernel=kernel, stencil_radius=4, verbose=False) - + montage.reconstruct(deconvolution_kernel="obf", stencil_radius=4, verbose=False) with pytest.raises(NotImplementedError, match="only defined for the parallax"): montage.residual_loss() - def test_residual_loss_returns_after_a_parallax_reconstruction(self, dataset4d): - _, montage = self._pair(dataset4d) - montage.reconstruct(deconvolution_kernel="ssb", stencil_radius=4, verbose=False) montage.reconstruct(verbose=False) - assert montage.residual_loss() > 0 - def test_icom_runs_in_either_mode(self, dataset4d): - _, montage = self._pair(dataset4d) - for mode in ("fft", "stencil"): - montage.reconstruct( - deconvolution_kernel="icom", - convolution_mode=mode, - stencil_radius=8, - verbose=False, - ) - assert np.isfinite(montage.obj).all() - def test_icom_ignores_an_empirical_probe(self, dataset4d): """It never reads psi, so an empirical probe changes nothing -- not even a bit. @@ -2303,28 +1912,6 @@ def test_com_shift_is_the_k_weighted_first_moment(self, dataset4d): assert _relative_error(to_numpy(collapsed), to_numpy(expected)) < 1e-5 - def test_the_icom_kernel_is_the_ricom_kernel(self): - """`ifft2(-i q / |q|**2)` is `r / (2 pi |r|**2)`, the kernel riCOM writes down.""" - size = 128 - frequency = np.fft.fftfreq(size) - qx, qy = frequency[:, None], frequency[None, :] - q_square = qx**2 + qy**2 - q_square[0, 0] = 1.0 - operator = -1j * qx / q_square - operator[0, 0] = 0.0 - transformed = np.real(np.fft.ifft2(operator)) - - position = np.fft.fftfreq(size, 1 / size) - rx, ry = position[:, None], position[None, :] - r_square = rx**2 + ry**2 - r_square[0, 0] = 1.0 - analytic = rx / (2 * np.pi * r_square) - - # away from the singular origin and the periodic seam - distance = np.hypot(rx, ry) - inside = (distance > 3) & (distance < size / 4) - assert np.corrcoef(transformed[inside], analytic[inside])[0, 1] > 0.98 - def test_phase_flip_is_not_applied_to_deconvolution_kernels(self, dataset4d): """The kernels already invert the contrast transfer, as in `DirectPtychography`.""" _, montage = self._pair(dataset4d) @@ -2383,24 +1970,6 @@ def test_fft_mode_is_exact(self, dataset4d, kernel): assert _relative_error(montage.obj, fourier.obj) < 1e-5 - @pytest.mark.parametrize("kernel", KERNELS) - def test_fft_mode_beats_a_truncated_stencil(self, dataset4d, kernel): - fourier, montage = self._pair(dataset4d) - fourier.reconstruct(deconvolution_kernel=kernel, verbose=False) - reference = fourier.obj.copy() - - montage.reconstruct(deconvolution_kernel=kernel, convolution_mode="fft", verbose=False) - exact = _relative_error(montage.obj, reference) - montage.reconstruct( - deconvolution_kernel=kernel, - convolution_mode="stencil", - stencil_radius=5, - verbose=False, - ) - truncated = _relative_error(montage.obj, reference) - - assert exact < 0.01 * truncated - def test_pad_boundary_stays_linear(self, dataset4d): """A Fourier convolution wraps; `"pad"` doubles the canvas so that it does not. @@ -2422,9 +1991,7 @@ def test_pad_boundary_stays_linear(self, dataset4d): assert montage.obj.shape == stencil.shape assert _relative_error(montage.obj, stencil) < 0.05 - @pytest.mark.parametrize("interpolation", ["nearest", "bilinear"]) - @pytest.mark.parametrize("boundary", ["wrap", "pad"]) - def test_conv2d_stencil_matches_the_scatter(self, dataset4d, boundary, interpolation): + def test_conv2d_stencil_matches_the_scatter(self): """`splat_and_convolve` is a reorganization of `scatter_add_convolve`, not a change. The subtlety it has to reproduce is that the scatter tests the boundary at the @@ -2441,28 +2008,31 @@ def test_conv2d_stencil_matches_the_scatter(self, dataset4d, boundary, interpola coords = torch.rand(4, 30, 2) * torch.tensor([28.0, 24.0]) - torch.tensor([4.0, 4.0]) shape = (20, 16) - scattered = torch.zeros(shape[0] * shape[1], dtype=torch.complex64) - scatter_add_convolve( - values, - coords, - shape, - offsets, - weights, - boundary=boundary, - interpolation=interpolation, - out=scattered, - ) - convolved = splat_and_convolve( - values, - coords, - shape, - weights, - radius, - boundary=boundary, - interpolation=interpolation, - ).sum(0) - - assert _relative_error(to_numpy(convolved), to_numpy(scattered.view(shape))) < 1e-5 + for boundary in ("wrap", "pad"): + for interpolation in ("nearest", "bilinear"): + scattered = torch.zeros(shape[0] * shape[1], dtype=torch.complex64) + scatter_add_convolve( + values, + coords, + shape, + offsets, + weights, + boundary=boundary, + interpolation=interpolation, + out=scattered, + ) + convolved = splat_and_convolve( + values, + coords, + shape, + weights, + radius, + boundary=boundary, + interpolation=interpolation, + ).sum(0) + + error = _relative_error(to_numpy(convolved), to_numpy(scattered.view(shape))) + assert error < 1e-5, f"{boundary=} {interpolation=}" def test_splat_stack_matches_a_shared_canvas(self): """Summing the per-image canvases must give what the shared-canvas splat gives.""" @@ -2478,15 +2048,16 @@ def test_splat_stack_matches_a_shared_canvas(self): assert _relative_error(to_numpy(stack.sum(0)), to_numpy(shared.view(shape))) < 1e-5 - @pytest.mark.parametrize("stencil_radius, expected", [("auto", "fft"), (6, "stencil")]) - def test_auto_reads_the_stencil_radius(self, stencil_radius, expected): - """Naming a radius is a request to truncate; leaving it 'auto' takes the exact route.""" - resolve = DirectPtychographyMontage._resolve_convolution_mode - assert resolve("auto", "ssb", stencil_radius) == expected - - def test_parallax_ignores_the_mode(self): + def test_auto_reads_the_stencil_radius(self): + """Naming a radius is a request to truncate; leaving it 'auto' takes the exact route. + Parallax ignores the mode.""" resolve = DirectPtychographyMontage._resolve_convolution_mode - assert resolve("fft", "prlx", "auto") == "splat" + for args, expected in ( + (("auto", "ssb", "auto"), "fft"), + (("auto", "ssb", 6), "stencil"), + (("fft", "prlx", "auto"), "splat"), + ): + assert resolve(*args) == expected, f"{args=}" def test_unknown_mode_raises(self, dataset4d): _, montage = self._pair(dataset4d) @@ -2495,12 +2066,6 @@ def test_unknown_mode_raises(self, dataset4d): deconvolution_kernel="ssb", convolution_mode="direct", verbose=False ) - def test_fft_mode_leaves_no_stencil_info(self, dataset4d): - _, montage = self._pair(dataset4d) - montage.reconstruct(deconvolution_kernel="ssb", convolution_mode="fft", verbose=False) - - assert montage._stencil_info is None - class TestSharedCanvas: """`obj_origin` / `obj_fov` pin the canvas to a window of the *specimen*. @@ -2651,49 +2216,26 @@ def test_a_wrapped_canvas_smaller_than_the_scan_warns(self, dataset4d): verbose=False, ) - @pytest.mark.parametrize("boundary", ["wrap", "pad"]) - def test_a_canvas_covering_the_scan_is_silent(self, dataset4d, boundary, recwarn): + def test_a_canvas_covering_the_scan_is_silent(self, dataset4d, recwarn): montage = self._montage(dataset4d, integer_shift_defocus(1)) span = np.ptp(to_numpy(montage.positions_px), axis=0) * SCAN_SAMPLING - montage.reconstruct( - deconvolution_kernel="prlx", - boundary=boundary, - obj_origin=tuple(np.asarray(montage.scan_origin) - 4 * SCAN_SAMPLING), - obj_fov=tuple(span + 8 * SCAN_SAMPLING), - verbose=False, - ) - - assert not [w for w in recwarn if "wrap around" in str(w.message)] - - def test_wrap_defaults_to_the_scan_grid(self, dataset4d): - """The default is unchanged: no window given, the canvas is the scan.""" - montage = self._montage(dataset4d, integer_shift_defocus(1)) - montage.reconstruct(deconvolution_kernel="prlx", boundary="wrap", verbose=False) + for boundary in ("wrap", "pad"): + montage.reconstruct( + deconvolution_kernel="prlx", + boundary=boundary, + obj_origin=tuple(np.asarray(montage.scan_origin) - 4 * SCAN_SAMPLING), + obj_fov=tuple(span + 8 * SCAN_SAMPLING), + verbose=False, + ) - assert montage.obj.shape == tuple(montage.scan_gpts) + assert not [w for w in recwarn if "wrap around" in str(w.message)], f"{boundary=}" def test_obj_fov_must_be_positive(self, dataset4d): montage = self._montage(dataset4d, integer_shift_defocus(1)) with pytest.raises(ValueError, match="positive"): montage.reconstruct(deconvolution_kernel="prlx", obj_fov=(0.0, 10.0), verbose=False) - def test_scan_origin_survives_a_round_trip(self, dataset4d, tmp_path): - montage = self._montage(dataset4d, integer_shift_defocus(1), position_offset=(11.0, -3.0)) - montage.reconstruct(deconvolution_kernel="prlx", verbose=False) - - path = str(tmp_path / "montage.zip") - montage.save(path, mode="o") - restored = load(path) - - assert restored.scan_origin == pytest.approx(montage.scan_origin) - assert restored.obj_origin == pytest.approx(montage.obj_origin) - - def test_gridded_construction_has_a_zero_scan_origin(self, dataset4d): - """A raster acquisition's grid *defines* the coordinates, so its origin is zero.""" - _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) - assert montage.scan_origin == (0.0, 0.0) - class TestFrameDrift: """`estimate_frame_drift` on interleaved frames with a seeded, known displacement.""" @@ -2704,9 +2246,10 @@ class TestFrameDrift: FOV = (48 * CTF_SAMPLING, 48 * CTF_SAMPLING) DRIFT = np.array([[0.0, 0.0], [1.6, -0.8], [3.2, -1.6], [4.8, -2.4]]) * CTF_SAMPLING - def _frames(self, ctf_scene, drift=None): + @classmethod + def _frames(cls, ctf_scene, drift=None): complex_obj, _, probe, _, _ = ctf_scene - drift = self.DRIFT if drift is None else drift + drift = cls.DRIFT if drift is None else drift datasets, positions = ctf_interleaved_frames(complex_obj, probe, drift) montages = [] @@ -2720,27 +2263,28 @@ def _frames(self, ctf_scene, drift=None): ) montage.reconstruct( deconvolution_kernel="prlx", - obj_origin=self.ORIGIN, - obj_fov=self.FOV, + obj_origin=cls.ORIGIN, + obj_fov=cls.FOV, verbose=False, ) montages.append(montage) return montages, positions - def test_recovers_a_seeded_drift(self, ctf_scene): + @pytest.fixture(scope="class") + def measured(self, ctf_scene): + """The drift estimated from the default frames, which the tests only read.""" montages, _ = self._frames(ctf_scene) - measured = estimate_frame_drift(montages, verbose=False) + return estimate_frame_drift(montages, verbose=False) + def test_recovers_a_seeded_drift(self, measured): # the estimate is referred to the mean over frames, so compare mean-centered expected = self.DRIFT - self.DRIFT.mean(axis=0) assert measured == pytest.approx(expected, abs=0.35 * CTF_SCAN_SAMPLING) - - def test_drift_sums_to_zero(self, ctf_scene): - montages, _ = self._frames(ctf_scene) - measured = estimate_frame_drift(montages, verbose=False) assert measured.mean(axis=0) == pytest.approx((0.0, 0.0), abs=1e-9) - def test_correcting_the_positions_sharpens_the_combined_reconstruction(self, ctf_scene): + def test_correcting_the_positions_sharpens_the_combined_reconstruction( + self, ctf_scene, measured + ): """The point of the exercise: `positions - drift` must beat `positions`. Also pins the sign, which no amount of reasoning about correlation conventions @@ -2748,8 +2292,7 @@ def test_correcting_the_positions_sharpens_the_combined_reconstruction(self, ctf """ complex_obj, _, probe, ctf_full, _ = ctf_scene datasets, positions = ctf_interleaved_frames(complex_obj, probe, self.DRIFT) - montages, _ = self._frames(ctf_scene) - drift = estimate_frame_drift(montages, verbose=False) + drift = measured combined = Dataset3d.from_array( np.concatenate([d.array for d in datasets]), @@ -2788,8 +2331,11 @@ def test_zero_drift_is_recovered_as_zero(self, ctf_scene): measured = estimate_frame_drift(montages, verbose=False) assert np.abs(measured).max() < 0.35 * CTF_SCAN_SAMPLING - def test_rejects_frames_on_different_canvases(self, ctf_scene): + def test_rejects_frames_it_cannot_compare(self, ctf_scene): montages, _ = self._frames(ctf_scene) + with pytest.raises(ValueError, match="at least two"): + estimate_frame_drift(montages[:1], verbose=False) + montages[1].reconstruct( deconvolution_kernel="prlx", obj_origin=self.ORIGIN, @@ -2799,9 +2345,7 @@ def test_rejects_frames_on_different_canvases(self, ctf_scene): with pytest.raises(ValueError, match="share a canvas"): estimate_frame_drift(montages, verbose=False) - def test_rejects_frames_at_different_origins(self, ctf_scene): - """The subtle failure: same shape, wrong place, silently measuring the offset.""" - montages, _ = self._frames(ctf_scene) + # the subtle failure: same shape, wrong place, silently measuring the offset montages[1].reconstruct( deconvolution_kernel="prlx", obj_origin=(self.ORIGIN[0] + 3 * CTF_SAMPLING, self.ORIGIN[1]), @@ -2811,17 +2355,10 @@ def test_rejects_frames_at_different_origins(self, ctf_scene): with pytest.raises(ValueError, match="not a canvas origin"): estimate_frame_drift(montages, verbose=False) - def test_rejects_unreconstructed_frames(self, ctf_scene): - montages, _ = self._frames(ctf_scene) montages[1]._reset_reconstruction() with pytest.raises(ValueError, match="not been reconstructed"): estimate_frame_drift(montages, verbose=False) - def test_needs_at_least_two_frames(self, ctf_scene): - montages, _ = self._frames(ctf_scene) - with pytest.raises(ValueError, match="at least two"): - estimate_frame_drift(montages[:1], verbose=False) - def test_works_on_gridded_reconstructions(self, dataset4d): """Nothing about it is montage-specific -- two raster reconstructions also align.""" defocus = integer_shift_defocus(1) @@ -2838,129 +2375,6 @@ def test_works_on_gridded_reconstructions(self, dataset4d): assert np.abs(drift).max() < 1e-6 -@pytest.mark.skipif(not ACCELERATORS, reason="no accelerator available") -@pytest.mark.parametrize("device", ACCELERATORS) -class TestAccelerators: - """Both classes must run, and agree with CPU, on whatever accelerator is present. - - MPS has no float64 at all, so every positional tensor -- coordinates, shifts, canvas - origins -- has to take its dtype from the device rather than hardcode one. - """ - - @staticmethod - def _kwargs(dataset4d): - return dict(_common_kwargs(integer_shift_defocus(1)), edge_blend_pixels=0) - - @staticmethod - def _ungridded(dataset4d): - """A `Dataset3d` and positions, with the origin *fitted* rather than forced. - - Forcing the origin skips the centre-of-mass fit entirely, which is where the - device mismatch lived, so these deliberately let it run. - """ - dataset3d = Dataset3d.from_array( - np.asarray(dataset4d.array).reshape(-1, N, N), - name="ungridded", - sampling=(1.0, dataset4d.sampling[-2], dataset4d.sampling[-1]), - units=("index", "A^-1", "A^-1"), - ) - return dataset3d, scan_positions_px() * SCAN_SAMPLING - - @pytest.mark.parametrize("boundary", ["wrap", "pad"]) - def test_montage_matches_cpu(self, dataset4d, device, boundary): - def run(where): - montage = DirectPtychographyMontage.from_dataset4d( - dataset4d, device=where, boundary=boundary, **self._kwargs(dataset4d) - ) - return montage.reconstruct(verbose=False).obj - - on_device, on_cpu = run(device), run("cpu") - - assert on_device.shape == on_cpu.shape - assert np.abs(on_device - on_cpu).max() / np.abs(on_cpu).max() < 1e-4 - - def test_montage_upsampled_matches_cpu(self, dataset4d, device): - def run(where): - montage = DirectPtychographyMontage.from_dataset4d( - dataset4d, device=where, **self._kwargs(dataset4d) - ) - return montage.reconstruct( - upsampling_factor=2, interpolation="bilinear", verbose=False - ).obj - - assert np.abs(run(device) - run("cpu")).max() / np.abs(run("cpu")).max() < 1e-4 - - def test_defocus_gradient_matches_cpu(self, dataset4d, device): - def run(where): - montage = DirectPtychographyMontage.from_dataset4d( - dataset4d, - device=where, - boundary="pad", - defocus_gradient=(30.0, -10.0), - **self._kwargs(dataset4d), - ) - return montage.reconstruct(verbose=False).obj - - on_device, on_cpu = run(device), run("cpu") - - assert on_device.shape == on_cpu.shape - assert np.abs(on_device - on_cpu).max() / np.abs(on_cpu).max() < 1e-4 - - def test_real_space_kernel_runs(self, dataset4d, device): - montage = DirectPtychographyMontage.from_dataset4d( - dataset4d, device=device, **self._kwargs(dataset4d) - ) - obj = montage.reconstruct(deconvolution_kernel="ssb", stencil_radius=4, verbose=False).obj - - assert np.isfinite(obj).all() - - def test_residual_loss_and_search_run(self, dataset4d, device): - montage = DirectPtychographyMontage.from_dataset4d( - dataset4d, device=device, **self._kwargs(dataset4d) - ) - montage.reconstruct(verbose=False) - - assert np.isfinite(montage.residual_loss()) - - @pytest.mark.parametrize( - "cls", [DirectPtychography, DirectPtychographyMontage], ids=["fourier", "montage"] - ) - def test_ungridded_fits_the_origin_on_device(self, dataset4d, device, cls): - """Regression: probe positions arrive as numpy and met the measured origin on MPS.""" - dataset3d, positions = self._ungridded(dataset4d) - - reconstruction = cls.from_dataset3d( - dataset3d, - positions, - energy=PROBE_ENERGY, - semiangle_cutoff=SEMIANGLE_CUTOFF, - rotation_angle=0.0, - scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), - device=device, - verbose=False, - ) - obj = reconstruction.reconstruct(deconvolution_kernel="prlx", verbose=False).obj - - assert np.isfinite(obj).all() - - -@pytest.mark.skipif(not ACCELERATORS, reason="no accelerator available") -@pytest.mark.parametrize("device", ACCELERATORS) -def test_padded_canvas_shape_is_device_independent(dataset4d, device): - """Float noise must not push a canvas bound across an integer on one device only.""" - - def shape(where): - montage = DirectPtychographyMontage.from_dataset4d( - dataset4d, - device=where, - boundary="pad", - **dict(_common_kwargs(integer_shift_defocus(1)), edge_blend_pixels=0), - ) - return montage.reconstruct(verbose=False).obj.shape - - assert shape(device) == shape("cpu") - - class TestAberrationFitFromProbe: """Fitting ``chi(k)`` to a measured ``psi(k)``, and knowing when not to trust it. @@ -3097,22 +2511,6 @@ def _recon(dataset4d): kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0, boundary="wrap") return DirectPtychographyMontage.from_dataset4d(dataset4d, **kwargs) - def test_band_limit_is_twice_the_cutoff(self, dataset4d): - montage = self._recon(dataset4d) - expected = 2 * montage.semiangle_cutoff * 1e-3 / montage.wavelength - assert montage._return_band_limit() == pytest.approx(expected) - - @pytest.mark.parametrize("kernel", ["ssb", "obf"]) - def test_deconvolution_kernels_are_untouched(self, dataset4d, kernel): - """``gamma`` already vanishes past the band limit, so the clamp must be a no-op.""" - montage = self._recon(dataset4d) - montage.reconstruct(deconvolution_kernel=kernel, stencil_radius=5, verbose=False) - limited = to_numpy(montage.obj).copy() - montage.reconstruct( - deconvolution_kernel=kernel, stencil_radius=5, band_limit=False, verbose=False - ) - assert np.allclose(limited, to_numpy(montage.obj)) - def test_parallax_amplitude_now_scales_like_ssb(self, dataset4d): """Unfiltered ``prlx`` fell off as 1/U where every other kernel goes as 1/U**2, and that gap was the retained replicas.""" @@ -3136,8 +2534,11 @@ def spread(kernel, factor, **kwargs): assert unlimited == pytest.approx(2.0, rel=0.05) def test_an_explicit_cutoff_below_the_band_limit_still_wins(self, dataset4d): + """``gamma`` already vanishes past the band limit, so the kernels carrying it keep + whatever cutoff was asked for.""" montage = self._recon(dataset4d) band = montage._return_band_limit() + assert band == pytest.approx(2 * montage.semiangle_cutoff * 1e-3 / montage.wavelength) assert montage._resolve_q_lowpass(band / 4, "prlx") == pytest.approx(band / 4) assert montage._resolve_q_lowpass(band * 4, "prlx") == pytest.approx(band) @@ -3167,6 +2568,8 @@ def test_an_empirical_probe_reads_its_extent_off_the_bf_mask(self, dataset4d): assert empirical._return_band_limit() == pytest.approx( analytic._return_band_limit(), rel=0.25 ) + empirical.reconstruct(deconvolution_kernel="ssb", stencil_radius=5, verbose=False) + assert np.isfinite(empirical.obj).all() # --------------------------------------------------------------------------- @@ -3281,10 +2684,9 @@ def test_the_fft_route_matches_the_fourier_class(self, complex_pair, kernel): # energies, which agree to 3e-7, set how well it can match assert montage.residual_loss() == pytest.approx(fourier.residual_loss(), abs=1e-6) - @pytest.mark.parametrize("kernel", ["obf"]) - def test_the_fft_route_takes_the_ard_prior(self, complex_pair, kernel): + def test_the_fft_route_takes_the_ard_prior(self, complex_pair): fourier, montage = complex_pair - kwargs = dict(deconvolution_kernel=kernel, obj_type="complex", prior="ard", verbose=False) + kwargs = dict(deconvolution_kernel="obf", obj_type="complex", prior="ard", verbose=False) fourier.reconstruct(**kwargs) montage.reconstruct(convolution_mode="fft", **kwargs) @@ -3310,12 +2712,26 @@ def test_the_fft_route_takes_the_wiener_kernel(self, complex_pair, obj_type): expected = to_numpy(getattr(fourier, attribute)) assert _relative_error(to_numpy(getattr(montage, attribute)), expected) < 5e-3 - def test_the_stencil_route_refuses_the_wiener_kernel(self, complex_pair): + def test_the_stencil_route_refuses_the_nonlinear_and_multislice_solves(self, complex_pair): _, montage = complex_pair - with pytest.raises(ValueError, match="convolution_mode='fft'"): + with pytest.raises(ValueError, match="'wiener' kernel needs convolution_mode='fft'"): montage.reconstruct( deconvolution_kernel="wiener", convolution_mode="stencil", verbose=False ) + with pytest.raises(ValueError, match="prior='ard' needs convolution_mode='fft'"): + montage.reconstruct( + deconvolution_kernel="obf", convolution_mode="stencil", prior="ard", verbose=False + ) + with pytest.raises(NotImplementedError, match="num_slices=2 needs convolution_mode='fft'"): + montage.reconstruct( + deconvolution_kernel="obf", + num_slices=2, + slice_mode="joint", + slice_thicknesses=5.0, + convolution_mode="stencil", + stencil_radius=4, + verbose=False, + ) def test_wdd_is_refused(self, complex_pair): _, montage = complex_pair @@ -3338,13 +2754,6 @@ def test_the_focal_spread_reaches_the_fft_route(self, complex_pair): for recon in (fourier, montage): recon.focal_spread = 0.0 - def test_the_stencil_route_refuses_the_ard_prior(self, complex_pair): - _, montage = complex_pair - with pytest.raises(ValueError, match="convolution_mode='fft'"): - montage.reconstruct( - deconvolution_kernel="obf", convolution_mode="stencil", prior="ard", verbose=False - ) - def test_sections_are_single_plane_reconstructions(self, complex_pair): """Sectioning is dispatched before the convolution route is chosen, so one route covers both.""" @@ -3414,24 +2823,6 @@ def test_ricom_amplitude_is_the_zeroth_moment(self, complex_pair, convolution_mo expected = to_numpy(fourier.corrected_amplitude) assert _relative_error(to_numpy(montage.corrected_amplitude), expected) < 1e-4 - def test_several_slices_need_the_fft_route(self, complex_pair): - _, montage = complex_pair - with pytest.raises(NotImplementedError, match="convolution_mode='fft'"): - montage.reconstruct( - deconvolution_kernel="obf", - num_slices=2, - slice_mode="joint", - slice_thicknesses=5.0, - convolution_mode="stencil", - stencil_radius=4, - verbose=False, - ) - - def test_parallax_rejects_complex(self, complex_pair): - _, montage = complex_pair - with pytest.raises(ValueError, match="cannot separate amplitude from phase"): - montage.reconstruct(deconvolution_kernel="prlx", obj_type="complex", verbose=False) - def test_frame_mean_subtraction_rejects_complex(self, dataset4d): montage = DirectPtychographyMontage.from_dataset4d( dataset4d, @@ -3442,18 +2833,3 @@ def test_frame_mean_subtraction_rejects_complex(self, dataset4d): ) with pytest.raises(ValueError, match="zeroth moment"): montage.reconstruct(deconvolution_kernel="obf", obj_type="complex", verbose=False) - - @pytest.mark.parametrize("kernel", ["obf", "icom"]) - def test_sparse_frames_match_dense(self, dataset4d, kernel): - dense, sparse = TestSparseEvents._pair(dataset4d, integer_shift_defocus(1)) - kwargs = dict( - deconvolution_kernel=kernel, - obj_type="complex", - convolution_mode="stencil", - stencil_radius=4, - ) - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - a, b = TestSparseEvents._both(dense, sparse, **kwargs) - - assert _relative_error(b, a) < 1e-4 From 9459276cfe45e558851a817ae210f36d32c5607b Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 15:09:31 -0700 Subject: [PATCH 13/19] Consolidate the DirectPtychographyTiltSeries tests - test_one_untilted_acquisition_is_its_own_reconstruction[3]: each case simulated and built its own acquisition; the three cases now loop over one. - test_upsampling asserted only the shape and finiteness on the five-tilt series. It now runs on two noise-free acquisitions and checks the upsampled slices against the coarse ones on the coarse grid. - test_from_tilt_series_builds_each_acquisition: per-dataset keyword forwarding and the returned type are asserted by test_survives_a_round_trip. - test_rejects_bad_series[3], test_rejects_another_scan_grid, test_only_the_normal_equation_kernels_combine[4] (a single `kernel not in _KERNELS` branch) and test_a_tilt_series_cannot_be_sectioned -> test_rejects_what_it_cannot_combine. - The noise-free acquisitions are a module fixture shared by test_upsampling, the rejections and test_survives_a_round_trip_before_reconstructing; the layered dataset comes from conftest's tilted_layers_4dstem. Co-Authored-By: Claude Opus 5.5 --- .../test_direct_ptychography_tilt_series.py | 133 +++++++----------- 1 file changed, 52 insertions(+), 81 deletions(-) diff --git a/tests/diffractive_imaging/test_direct_ptychography_tilt_series.py b/tests/diffractive_imaging/test_direct_ptychography_tilt_series.py index 14f53d1b9..bac4705ef 100644 --- a/tests/diffractive_imaging/test_direct_ptychography_tilt_series.py +++ b/tests/diffractive_imaging/test_direct_ptychography_tilt_series.py @@ -7,11 +7,10 @@ from quantem.diffractive_imaging import DirectPtychography, DirectPtychographyTiltSeries from .conftest import ( - ML_ABERRATIONS, - ML_SAMPLING, + correlation, layer_purity, ml_kwargs, - thin_layer_4dstem, + tilted_layers_4dstem, twisted_square_cell, ) @@ -25,19 +24,9 @@ def _layers(): return [0.3 * layer for layer in layers] -def _shifted(image, shift): - q = np.fft.fftfreq(image.shape[0], ML_SAMPLING) - ramp = np.exp(-2j * np.pi * (q[:, None] * shift[0] + q[None, :] * shift[1])) - return np.fft.ifft2(np.fft.fft2(image) * ramp).real - - def _acquisition(tilt_mrad, dose, seed, drift=(0.0, 0.0)): """The two layers seen at a specimen tilt: each displaced by its depth times the tilt.""" - theta = np.asarray(tilt_mrad, dtype=float) * 1e-3 - logs = [ - 1j * _shifted(layer, z * theta + np.asarray(drift)) for layer, z in zip(_layers(), DEPTHS) - ] - return thin_layer_4dstem(logs, list(DEPTHS), ML_ABERRATIONS, dose=dose, seed=seed) + return tilted_layers_4dstem(tilt_mrad, SPACING, dose=dose, seed=seed, drift=drift) def _recon(dataset): @@ -56,24 +45,34 @@ def tilted(): return series, _recon(_acquisition((0, 0), total, seed=99)) +@pytest.fixture(scope="module") +def noiseless(): + """Noise-free acquisitions at no tilt and at 10 mrad.""" + return [_acquisition((0, 0), None, 0), _acquisition((10, 0), None, 0)] + + KWARGS = dict(obj_type="complex", num_slices=2, slice_mode="joint", slice_thicknesses=SPACING) class TestTiltSeries: - @pytest.mark.parametrize( - "kernel, obj_type", [("obf", "complex"), ("wiener", "pure_phase"), ("obf", "pure_phase")] - ) - def test_one_untilted_acquisition_is_its_own_reconstruction(self, kernel, obj_type): + def test_one_untilted_acquisition_is_its_own_reconstruction(self): recon = _recon(_acquisition((0, 0), 300.0, seed=0)) - kwargs = dict( - deconvolution_kernel=kernel, obj_type=obj_type, num_slices=2, slice_mode="joint" - ) - expected = recon.reconstruct(slice_thicknesses=SPACING, verbose=False, **kwargs).obj series = DirectPtychographyTiltSeries([recon], [(0, 0)]) - series.reconstruct(slice_thicknesses=SPACING, **kwargs) - - assert np.array_equal(series.obj, expected) - assert series.residual_loss() == pytest.approx(recon.residual_loss()) + for kernel, obj_type in ( + ("obf", "complex"), + ("wiener", "pure_phase"), + ("obf", "pure_phase"), + ): + kwargs = dict( + deconvolution_kernel=kernel, obj_type=obj_type, num_slices=2, slice_mode="joint" + ) + expected = recon.reconstruct(slice_thicknesses=SPACING, verbose=False, **kwargs).obj + series.reconstruct(slice_thicknesses=SPACING, **kwargs) + + assert np.array_equal(series.obj, expected), f"{kernel=} {obj_type=}" + assert series.residual_loss() == pytest.approx(recon.residual_loss()), ( + f"{kernel=} {obj_type=}" + ) def test_the_residual_prefers_the_true_tilts(self, tilted): series, _ = tilted @@ -116,53 +115,40 @@ def test_register_recovers_a_drift(self): assert series.shifts[0] == pytest.approx((0.0, 0.0)) assert series.shifts[1] == pytest.approx(drift, abs=0.03) - def test_upsampling(self, tilted): - series, _ = tilted + def test_upsampling(self, noiseless): + series = DirectPtychographyTiltSeries([_recon(d) for d in noiseless], [(0, 0), (10, 0)]) + coarse = to_numpy(series.reconstruct(deconvolution_kernel="obf", **KWARGS).corrected_bf) series.reconstruct(deconvolution_kernel="obf", upsampling_factor=2, **KWARGS) - assert series.corrected_bf.shape == (2, 128, 128) - assert np.isfinite(series.obj).all() - - def test_from_tilt_series_builds_each_acquisition(self): - datasets = [_acquisition((0, 0), None, 0), _acquisition((10, 0), None, 0)] - shared = ml_kwargs() - shared.pop("aberration_coefs") - series = DirectPtychography.from_tilt_series( - datasets, - [(0, 0), (10, 0)], - per_dataset_kwargs=[ - {"aberration_coefs": ML_ABERRATIONS}, - {"aberration_coefs": dict(ML_ABERRATIONS, C10=60.0)}, - ], - **shared, - ) + fine = to_numpy(series.corrected_bf) - assert isinstance(series, DirectPtychographyTiltSeries) - assert len(series) == 2 - assert series[1].hyperparameter_state.current_aberrations()["C10"] == pytest.approx(60.0) - - @pytest.mark.parametrize( - "make, match", - [ - (lambda r: DirectPtychographyTiltSeries([r], [(0, 0), (1, 0)]), "one row per"), - (lambda r: DirectPtychographyTiltSeries([r, "not one"], [(0, 0), (1, 0)]), "must be"), - (lambda r: DirectPtychographyTiltSeries([], []), "at least one"), - ], - ) - def test_rejects_bad_series(self, tilted, make, match): - series, _ = tilted - with pytest.raises((ValueError, TypeError), match=match): - make(series[0]) + assert fine.shape == (2, 128, 128) + for slice_index in range(2): + assert correlation(fine[slice_index, ::2, ::2], coarse[slice_index]) > 0.99 - def test_rejects_another_scan_grid(self, tilted): - series, _ = tilted + def test_rejects_what_it_cannot_combine(self, tilted, noiseless): from quantem.core.datastructures import Dataset4d - full = _acquisition((0, 0), None, 0) + series, _ = tilted + one = series[0] + with pytest.raises((ValueError, TypeError), match="one row per"): + DirectPtychographyTiltSeries([one], [(0, 0), (1, 0)]) + with pytest.raises((ValueError, TypeError), match="must be"): + DirectPtychographyTiltSeries([one, "not one"], [(0, 0), (1, 0)]) + with pytest.raises((ValueError, TypeError), match="at least one"): + DirectPtychographyTiltSeries([], []) + + full = noiseless[0] cropped = Dataset4d.from_array( np.asarray(full.array)[:32].copy(), sampling=full.sampling, units=full.units ) with pytest.raises(ValueError, match="scan grid"): - DirectPtychographyTiltSeries([series[0], _recon(cropped)], [(0, 0), (1, 0)]) + DirectPtychographyTiltSeries([one, _recon(cropped)], [(0, 0), (1, 0)]) + + # only the kernels solving the normal equations combine acquisitions + with pytest.raises(ValueError, match="normal"): + series.reconstruct(deconvolution_kernel="ssb", num_slices=1) + with pytest.raises(ValueError, match="solved jointly"): + series.reconstruct(num_slices=2, slice_thicknesses=SPACING, slice_mode="sectioned") def test_survives_a_round_trip(self, tmp_path): from quantem.core.io.serialize import load @@ -197,13 +183,10 @@ def test_survives_a_round_trip(self, tmp_path): restored.reconstruct(**kwargs) assert np.allclose(restored.obj, before, atol=1e-6) - def test_survives_a_round_trip_before_reconstructing(self, tmp_path): + def test_survives_a_round_trip_before_reconstructing(self, noiseless, tmp_path): from quantem.core.io.serialize import load - series = DirectPtychographyTiltSeries( - [_recon(_acquisition((0, 0), None, 0)), _recon(_acquisition((10, 0), None, 0))], - [(0, 0), (10, 0)], - ) + series = DirectPtychographyTiltSeries([_recon(d) for d in noiseless], [(0, 0), (10, 0)]) path = str(tmp_path / "fresh.zip") series.save(path, mode="o") restored = load(path) @@ -214,15 +197,3 @@ def test_survives_a_round_trip_before_reconstructing(self, tmp_path): restored.reconstruct(deconvolution_kernel="obf", **KWARGS) expected = series.reconstruct(deconvolution_kernel="obf", **KWARGS).obj assert np.allclose(restored.obj, expected, atol=1e-6) - - @pytest.mark.parametrize("kernel", ["ssb", "wdd", "prlx", "icom"]) - def test_only_the_normal_equation_kernels_combine(self, tilted, kernel): - series, _ = tilted - with pytest.raises(ValueError, match="normal"): - series.reconstruct(deconvolution_kernel=kernel, num_slices=1) - - -def test_a_tilt_series_cannot_be_sectioned(tilted): - series, _ = tilted - with pytest.raises(ValueError, match="solved jointly"): - series.reconstruct(num_slices=2, slice_thicknesses=SPACING, slice_mode="sectioned") From e8ab4857b40dbc566d6f0a6b97c38a125e281547 Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 15:19:29 -0700 Subject: [PATCH 14/19] Keep an odd upsampling factor in the scalebar test test_scalebar_follows_the_upsampling loops over U = 2 and 3 on one pair, so the odd factor the removed U=3 parametrizations covered is still checked for both classes. Co-Authored-By: Claude Opus 5.5 --- .../test_direct_ptychography_montage.py | 25 +++++++++++++------ 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/tests/diffractive_imaging/test_direct_ptychography_montage.py b/tests/diffractive_imaging/test_direct_ptychography_montage.py index 41182f669..f51f04693 100644 --- a/tests/diffractive_imaging/test_direct_ptychography_montage.py +++ b/tests/diffractive_imaging/test_direct_ptychography_montage.py @@ -969,19 +969,28 @@ class TestVisualization: """The object sampling, and hence `visualize`'s scalebar, must follow the upsampling.""" def test_scalebar_follows_the_upsampling(self, dataset4d): - for recon in _build_pair(dataset4d, integer_shift_defocus(1, 2)): + for recon in _build_pair(dataset4d, integer_shift_defocus(1)): name = type(recon).__name__ assert recon._obj_sampling[0] == pytest.approx(SCAN_SAMPLING), name assert recon.fov == pytest.approx((N * SCAN_SAMPLING, N * SCAN_SAMPLING)), name assert recon.scan_origin == (0.0, 0.0), name - recon.reconstruct(deconvolution_kernel="prlx", upsampling_factor=2, verbose=False) - assert recon._obj_sampling[0] == pytest.approx(SCAN_SAMPLING / 2), name - # the reported sampling must span the same field of view as the image itself - assert recon.obj.shape[0] * recon._obj_sampling[0] == pytest.approx( - N * SCAN_SAMPLING - ), name - assert recon._obj_fov == pytest.approx(recon.fov), name + for upsampling_factor in (2, 3): + case = f"{name} {upsampling_factor=}" + recon.reconstruct( + deconvolution_kernel="prlx", + upsampling_factor=upsampling_factor, + override_aberration_coefs={"C10": integer_shift_defocus(1, upsampling_factor)}, + verbose=False, + ) + assert recon._obj_sampling[0] == pytest.approx( + SCAN_SAMPLING / upsampling_factor + ), case + # the reported sampling must span the same field of view as the image itself + assert recon.obj.shape[0] * recon._obj_sampling[0] == pytest.approx( + N * SCAN_SAMPLING + ), case + assert recon._obj_fov == pytest.approx(recon.fov), case def test_padded_canvas_keeps_the_scan_sampling(self, fourier_montage): """A padded canvas spans more than the scan, so its sampling is not fov/shape. From 3ef8cc3169af0c047e1275af4d9a0e6f7fb1d74a Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 15:30:30 -0700 Subject: [PATCH 15/19] Add CLAUDE.md with rules for writing tests The test suite grew from approximately 2k lines in June to 22k lines in October, mostly from tests added during Claude Code sessions. CLAUDE.md asks future sessions to extend existing tests before adding new ones, 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. Co-Authored-By: Claude Opus 5.5 --- .gitignore | 2 +- CLAUDE.md | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) create mode 100644 CLAUDE.md diff --git a/.gitignore b/.gitignore index c5bc9776b..807f05aa1 100644 --- a/.gitignore +++ b/.gitignore @@ -191,6 +191,6 @@ ipynb-playground/ # cursor/CLI .cursor .claude -CLAUDE.md +CLAUDE.local.md AGENTS.md AGENT.md diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000..cd138bd36 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,23 @@ +# quantem + +## Tests + +The test suite protects the code against regressions for the next person who edits it. CI runs +every test on every push to every pull request, so each test adds to the CI time of every PR. + +- Before writing a new test, find the existing test for that code and extend it. Add a new + test file only for a new module. +- Add a test for a bug fix (a test that fails before the fix) or for a new code path. + Refactors, renames, argument forwarding, default values, getters, reprs and plotting do not + need new tests. +- Test each code path once. Do not parametrize over values that take the same branch, and do + not take cartesian products of parameters. When several inexpensive cases share one code + path, loop over them inside one test and name the case in the assert message. +- Every assert must be able to fail. Do not assert only `is not None`, a type, or a shape that + is fixed by construction, do not compare a result with itself, and do not use `try/except` + or `skip` in a way that hides a failure. +- Use the smallest problem that exercises the code path, such as small arrays, few iterations + and a coarse orientation plan, and build expensive inputs once in a module-scoped fixture. +- CI does not run tests marked `@pytest.mark.slow`, so mark a test as slow only when CI does + not need to cover it. +- When a test duplicates another test, delete it. From f8bb45007ddfe8f07882356cadfc026b3a37e2a3 Mon Sep 17 00:00:00 2001 From: Colin Ophus Date: Sun, 18 Oct 2026 15:49:11 -0700 Subject: [PATCH 16/19] ci: schedule tests by module and class With the default `--dist load`, each xdist worker rebuilds every module- and class-scoped fixture it uses. `--dist loadscope` assigns each test class (or each module's free functions) to a single worker, so the shared direct ptychography fixtures are built once. A local run with 4 workers dropped from approximately 118 s to 97 s. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/uv-pytests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/uv-pytests.yml b/.github/workflows/uv-pytests.yml index fa1c69b9c..04090c367 100644 --- a/.github/workflows/uv-pytests.yml +++ b/.github/workflows/uv-pytests.yml @@ -32,4 +32,4 @@ jobs: # -n logical, not auto: auto counts physical cores, 2 on GitHub's 4-vCPU runners. # tests/conftest.py gives each worker cores // workers torch and BLAS threads. - name: Run pytests - run: uv run pytest tests -n logical + run: uv run pytest tests -n logical --dist loadscope From 634d3d5efd85efb9d90a18552bf8daad9f9e85bf Mon Sep 17 00:00:00 2001 From: Georgios Varnavides Date: Mon, 19 Oct 2026 19:23:32 +0200 Subject: [PATCH 17/19] Read the tomography pose learning rate from its declared default TomographyDatasetBase declared DEFAULT_LRS = {"pose_lr": 5e-2}, while TomographyOpt._get_default_lr reads DEFAULT_LRS.get("pose", 5e-2), so the declared value was never used and the fallback, which happens to be equal, always was. The key is now "pose", the plain key every other DEFAULT_LRS uses. test_get_default_lr asserted 5e-2, which the fallback also gave, so it could not fail. It now compares each key with the model's own DEFAULT_LRS, and fails with a KeyError before this fix. dataset_models.py already differs from ruff format in lines this commit does not touch, so the formatter hook was skipped to keep the diff to the fix. Co-Authored-By: Claude Opus 5.5 --- src/quantem/tomography/dataset_models.py | 2 +- tests/tomography/test_tomography_opt.py | 8 ++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index 3d4ee028a..b49212ec1 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -160,7 +160,7 @@ class TomographyDatasetBase(AutoSerialize, OptimizerMixin, nn.Module): _token = object() DEFAULT_LRS = { - "pose_lr": 5e-2, + "pose": 5e-2, } def __init__( diff --git a/tests/tomography/test_tomography_opt.py b/tests/tomography/test_tomography_opt.py index c9898156d..ec5c773c1 100644 --- a/tests/tomography/test_tomography_opt.py +++ b/tests/tomography/test_tomography_opt.py @@ -125,9 +125,13 @@ def test_step_and_zero_grad_touch_each_optimizer_once(self, torch_device, monkey @requires_torch class TestOptHelpers: def test_get_default_lr(self, torch_device): + """Each key reads the model's own declared default rather than the fallback, which for + the pose happens to have the same value.""" tomo = _inr_tomography(torch_device) - assert tomo._get_default_lr("object") == pytest.approx(8e-6) - assert tomo._get_default_lr("pose") == pytest.approx(5e-2) + assert tomo._get_default_lr("object") == pytest.approx( + tomo.obj_model.DEFAULT_LRS["object"] + ) + assert tomo._get_default_lr("pose") == pytest.approx(tomo.dset.DEFAULT_LRS["pose"]) with pytest.raises(ValueError): tomo._get_default_lr("banana") From 64a6d31fa58a74ccd439975c7a68d17a6c837d89 Mon Sep 17 00:00:00 2001 From: Georgios Varnavides Date: Mon, 19 Oct 2026 19:23:32 +0200 Subject: [PATCH 18/19] Run the CPU tomography reconstruction test in CI test_tomography_cpu_recon.py exists to cover the learned reconstruction loop of Tomography.reconstruct on CPU-only CI, but it was marked slow, which CI skips. It takes under a second, so the slow marker is removed. Co-Authored-By: Claude Opus 5.5 --- tests/tomography/test_tomography_cpu_recon.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/tomography/test_tomography_cpu_recon.py b/tests/tomography/test_tomography_cpu_recon.py index 0ebf77522..5fb429577 100644 --- a/tests/tomography/test_tomography_cpu_recon.py +++ b/tests/tomography/test_tomography_cpu_recon.py @@ -8,8 +8,8 @@ It is skipped when CUDA is present: object models that go through ``setup_distributed`` must be built on CUDA when a CUDA device exists (the CPU path is only valid with no CUDA device), -and on GPU machines the ``requires_gpu`` tests already cover this path. Marked ``slow`` so it -runs under ``--runslow``. +and on GPU machines the ``requires_gpu`` tests already cover this path. It takes under a +second, so it is not marked ``slow``. """ import numpy as np @@ -19,7 +19,6 @@ from quantem.tomography.tomography_lite import TomographyLiteINR pytestmark = [ - pytest.mark.slow, pytest.mark.skipif( torch.cuda.is_available(), reason="CPU INR path is only valid without CUDA; GPU machines use the requires_gpu tests", From 04ac1e4608260bdfaee8360f5fd3dc3073a2d097 Mon Sep 17 00:00:00 2001 From: Georgios Varnavides Date: Mon, 19 Oct 2026 19:23:33 +0200 Subject: [PATCH 19/19] ci: run the slow tests on pull requests into main No CI job ran the tests marked slow. They now run with --runslow on pull requests into main, once per release, on every OS in the release matrix. Locally they add about 70 s with 4 workers (19 passed, 6 skipped without CUDA). CLAUDE.md's rule for the slow marker says so. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/uv-pytests.yml | 3 ++- CLAUDE.md | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/workflows/uv-pytests.yml b/.github/workflows/uv-pytests.yml index 04090c367..4772a0e9e 100644 --- a/.github/workflows/uv-pytests.yml +++ b/.github/workflows/uv-pytests.yml @@ -31,5 +31,6 @@ jobs: # -n logical, not auto: auto counts physical cores, 2 on GitHub's 4-vCPU runners. # tests/conftest.py gives each worker cores // workers torch and BLAS threads. + # Tests marked slow run only on pull requests into main, that is once per release. - name: Run pytests - run: uv run pytest tests -n logical --dist loadscope + run: uv run pytest tests -n logical --dist loadscope ${{ github.base_ref == 'main' && '--runslow' || '' }} diff --git a/CLAUDE.md b/CLAUDE.md index cd138bd36..24df760a2 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -18,6 +18,6 @@ every test on every push to every pull request, so each test adds to the CI time or `skip` in a way that hides a failure. - Use the smallest problem that exercises the code path, such as small arrays, few iterations and a coarse orientation plan, and build expensive inputs once in a module-scoped fixture. -- CI does not run tests marked `@pytest.mark.slow`, so mark a test as slow only when CI does - not need to cover it. +- CI runs tests marked `@pytest.mark.slow` only on pull requests into `main`, once per release, + so mark a test as slow only when checking it at each release is enough. - When a test duplicates another test, delete it.