diff --git a/.github/workflows/uv-pytests.yml b/.github/workflows/uv-pytests.yml index fa1c69b9c..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 + run: uv run pytest tests -n logical --dist loadscope ${{ github.base_ref == 'main' && '--runslow' || '' }} 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..24df760a2 --- /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 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. 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/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/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.""" diff --git a/tests/diffractive_imaging/conftest.py b/tests/diffractive_imaging/conftest.py index f9a12034e..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,26 +773,27 @@ 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. +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 - 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). + +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``. """ - 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))) + _, 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]: @@ -820,3 +810,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_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 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") 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..f51f04693 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,123 @@ 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)): + 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 + + 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. 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 +1025,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 +1041,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) + analytic, empirical = self._pair(dataset4d, DirectPtychographyMontage, normalize=False) - assert _relative_error(empirical.obj, expected) < 1e-5 + 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) - def test_is_analytic_reports_which_path_is_live(self, dataset4d): - analytic, empirical = self._pair(dataset4d, DirectPtychographyMontage) - - 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 +1157,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 +1172,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 +1265,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 +1306,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 +1320,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 +1334,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 +1362,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 +1392,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 +1402,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 +1434,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 +1449,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 +1483,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 +1519,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 +1561,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 +1619,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 +1660,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 +1784,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 +1800,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 +1921,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 +1979,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 +2000,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 +2017,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 +2057,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 +2075,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 +2225,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 +2255,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 +2272,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 +2301,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 +2340,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 +2354,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 +2364,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 +2384,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 +2520,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 +2543,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 +2577,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 +2693,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 +2721,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 +2763,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 +2832,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 +2842,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 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") 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.""" 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( 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, 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) 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) 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_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", diff --git a/tests/tomography/test_tomography_opt.py b/tests/tomography/test_tomography_opt.py index bae00ad9a..ec5c773c1 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,40 @@ 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): + """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( + 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") @@ -185,25 +147,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"} 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.""" 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)