From 920f17bc66ce3eb0314e9cc7574fea8e6300c700 Mon Sep 17 00:00:00 2001 From: yoshikisd Date: Mon, 23 Dec 2024 04:12:09 +0000 Subject: [PATCH] Added a test for crop_2D_translations in Ptycho2DDataset. --- tests/test_datasets.py | 54 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 53 insertions(+), 1 deletion(-) diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 1b135a1..92f5980 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -5,7 +5,7 @@ import torch as t import h5py import datetime from copy import deepcopy - +import pytest # # We start by testing the CDataset base class @@ -340,3 +340,55 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis): assert np.allclose(np.array(dataset.background.shape) // factor, np.array(copied_dataset.background.shape)) + +def test_Ptycho2DDataset_crop_2D_translations(ptycho_cxi_1): + # Grab dataset + cxi, expected = ptycho_cxi_1 + dataset = Ptycho2DDataset.from_cxi(cxi) + copied_dataset = deepcopy(dataset) + + # Test 1: Complain when the the bounds of an ROI are correctly defined, + # but it does not contain any sample positions inside of it. The + # translations in ptycho_cxi_1 ranges from 0m to -300m in both x and y + # (it looks like a line scan). We select an ROI at (0, -300) which should + # not contain any translation positions. + with pytest.raises(ValueError) as excinfo: + copied_dataset.crop_2D_translations(roi=(0,-5,-295,-300)) + assert('(i.e., patterns and translations will be empty)') in str(excinfo.value) + + # Test 2: Draw an ROI that's centered in the middle of the x/y translation range, + # but complain because the bounds of the ROI are incorrectly defined + x_left, y_top = dataset.translations[10,:2] + x_right, y_bottom = dataset.translations[-11,:2] + + with pytest.raises(ValueError) as excinfo: + copied_dataset.crop_2D_translations(roi=(x_right,x_left,y_bottom,y_top)) + assert('(i.e., patterns and translations will be empty)') in str(excinfo.value) + + # Test 3: Basically Test 2, but don't complain when the bounds are set correctly + copied_dataset.crop_2D_translations(roi=(x_left,x_right,y_top,y_bottom)) + + # Test 4: Make sure that the first and last x/y elements in dataset.translate + # match x_left, x_right, y_top, and y_bottom + assert t.allclose(t.tensor([x_left, y_top]), copied_dataset.translations[0,:2]) + + assert t.allclose(t.tensor([x_right, y_bottom]), copied_dataset.translations[-1,:2]) + + # Test 5: Check if the shape of dataset.patterns and dataset.translate is correct + # (designed to be 20 fewer rows here) + expected_patterns_shape = np.concatenate([[dataset.patterns.shape[0] - 20], + dataset.patterns.shape[-2:]]) + + expected_translations_shape = np.concatenate([[dataset.translations.shape[0] - 20], + dataset.translations.shape[-1:]]) + + assert np.allclose(np.array(copied_dataset.patterns.shape), + expected_patterns_shape) + + assert np.allclose(np.array(copied_dataset.translations.shape), + expected_translations_shape) + + # Test 6: Check if the contents of dataset.patterns and dataset.translate is correct + assert t.allclose(copied_dataset.patterns, dataset.patterns[10:-10,:]) + + assert t.allclose(copied_dataset.translations, dataset.translations[10:-10,:])