cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-10-12 11:00:26 +02:00
Code Issues Packages Projects Releases Wiki Activity

linting test_datasets.py

Browse Source
This commit is contained in:
gnzng
2025-07-07 13:49:46 -07:00
parent 4d812db68a
commit f463144347
1 changed files with 89 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_datasets.py
+89 -108
View File
@@ -18,20 +18,20 @@ from cdtools.tools import data as cdtdata
def test_CDataset_init():
entry_info = {'start_time': datetime.datetime.now(),
'title' : 'A simple test'}
'title': 'A simple test'}
sample_info = {'name': 'A test sample',
'mass' : 3.4,
'unit_cell' : np.array([1,1,1,87,84.5,90])}
'mass': 3.4,
'unit_cell': np.array([1, 1, 1, 87, 84.5, 90])}
wavelength = 1e-9
detector_geometry = {'distance': 0.7,
'basis': np.array([[0,-30e-6,0],
[-20e-6,0,0]]).transpose(),
'corner': np.array((2550e-6,3825e-6,0.3))}
mask = np.ones((256,256))
'basis': np.array([[0, -30e-6, 0],
[-20e-6, 0, 0]]).transpose(),
'corner': np.array((2550e-6, 3825e-6, 0.3))}
mask = np.ones((256, 256))
dataset = CDataset(entry_info, sample_info,
wavelength, detector_geometry, mask)
assert t.all(t.eq(dataset.mask,t.tensor(mask.astype(bool))))
assert t.all(t.eq(dataset.mask, t.tensor(mask.astype(bool))))
assert dataset.entry_info == entry_info
assert dataset.sample_info == sample_info
assert dataset.wavelength == wavelength
@@ -53,7 +53,7 @@ def test_CDataset_from_cxi(test_ptycho_cxis):
else:
assert dataset.sample_info is not None
assert np.isclose(dataset.wavelength,expected['wavelength'])
assert np.isclose(dataset.wavelength, expected['wavelength'])
# Just check one of the loaded attributes
assert np.isclose(dataset.detector_geometry['distance'],
@@ -64,18 +64,16 @@ def test_CDataset_from_cxi(test_ptycho_cxis):
assert 'corner' in dataset.detector_geometry
if expected['mask'] is not None:
assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask))
assert t.all(t.eq(t.tensor(expected['mask']), dataset.mask))
if expected['qe_mask'] is not None:
assert t.all(t.eq(t.tensor(expected['qe_mask']),dataset.qe_mask))
assert t.all(t.eq(t.tensor(expected['qe_mask']), dataset.qe_mask))
if expected['dark'] is not None:
assert t.all(t.eq(t.as_tensor(expected['dark'], dtype=t.float32),
dataset.background))
def test_CDataset_to_cxi(test_ptycho_cxis, tmp_path):
for cxi, expected in test_ptycho_cxis:
dataset = CDataset.from_cxi(cxi)
@@ -95,27 +93,24 @@ def test_CDataset_to_cxi(test_ptycho_cxis, tmp_path):
assert np.isclose(dataset.wavelength, read_dataset.wavelength)
# Just check one of the loaded attributes
# Just check one of the loaded attributes
assert np.isclose(dataset.detector_geometry['distance'],
read_dataset.detector_geometry['distance'])
# Check that the other ones are loaded but not for fidelity
assert 'basis' in read_dataset.detector_geometry
if dataset.detector_geometry['corner'] is not None:
assert 'corner' in read_dataset.detector_geometry
if dataset.mask is not None:
assert t.all(t.eq(dataset.mask,read_dataset.mask))
assert t.all(t.eq(dataset.mask, read_dataset.mask))
if dataset.qe_mask is not None:
assert t.all(t.eq(dataset.qe_mask,read_dataset.qe_mask))
assert t.all(t.eq(dataset.qe_mask, read_dataset.qe_mask))
if dataset.background is not None:
assert t.all(t.eq(dataset.background, read_dataset.background))
def test_CDataset_to(ptycho_cxi_1):
dataset = CDataset.from_cxi(ptycho_cxi_1[0])
@@ -136,28 +131,28 @@ def test_CDataset_to(ptycho_cxi_1):
def test_Ptycho2DDataset_init():
entry_info = {'start_time': datetime.datetime.now(),
'title' : 'A simple test'}
'title': 'A simple test'}
sample_info = {'name': 'A test sample',
'mass' : 3.4,
'unit_cell' : np.array([1,1,1,87,84.5,90])}
'mass': 3.4,
'unit_cell': np.array([1, 1, 1, 87, 84.5, 90])}
wavelength = 1e-9
detector_geometry = {'distance': 0.7,
'basis': np.array([[0,-30e-6,0],
[-20e-6,0,0]]).transpose(),
'corner': np.array((2550e-6,3825e-6,0.3))}
mask = np.ones((256,256))
qe_mask = 1.2*np.ones((256,256), dtype=np.float32)
patterns = np.random.rand(20,256,256)
translations = np.random.rand(20,3)
dataset = Ptycho2DDataset(translations, patterns,
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
'basis': np.array([[0, -30e-6, 0],
[-20e-6, 0, 0]]).transpose(),
'corner': np.array((2550e-6, 3825e-6, 0.3))}
mask = np.ones((256, 256))
qe_mask = 1.2 * np.ones((256, 256), dtype=np.float32)
patterns = np.random.rand(20, 256, 256)
translations = np.random.rand(20, 3)
assert t.all(t.eq(dataset.mask,t.BoolTensor(mask)))
dataset = Ptycho2DDataset(translations, patterns,
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
assert t.all(t.eq(dataset.mask, t.BoolTensor(mask)))
assert dataset.entry_info == entry_info
assert dataset.sample_info == sample_info
assert dataset.wavelength == wavelength
@@ -167,15 +162,15 @@ def test_Ptycho2DDataset_init():
# Also test one with a qe_mask
dataset = Ptycho2DDataset(translations, patterns,
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask,
qe_mask=qe_mask)
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask,
qe_mask=qe_mask)
assert t.all(t.eq(dataset.mask,t.BoolTensor(mask)))
assert t.all(t.eq(dataset.qe_mask,t.as_tensor(qe_mask)))
assert t.all(t.eq(dataset.mask, t.BoolTensor(mask)))
assert t.all(t.eq(dataset.qe_mask, t.as_tensor(qe_mask)))
assert dataset.entry_info == entry_info
assert dataset.sample_info == sample_info
assert dataset.wavelength == wavelength
@@ -199,7 +194,7 @@ def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
else:
assert dataset.sample_info is not None
assert np.isclose(dataset.wavelength,expected['wavelength'])
assert np.isclose(dataset.wavelength, expected['wavelength'])
# Just check one of the loaded attributes
assert np.isclose(dataset.detector_geometry['distance'],
@@ -210,18 +205,17 @@ def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
assert 'corner' in dataset.detector_geometry
if expected['mask'] is not None:
assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask))
assert t.all(t.eq(t.tensor(expected['mask']), dataset.mask))
if expected['qe_mask'] is not None:
assert t.all(t.eq(t.tensor(expected['qe_mask']),dataset.qe_mask))
assert t.all(t.eq(t.tensor(expected['qe_mask']), dataset.qe_mask))
if expected['dark'] is not None:
assert t.all(t.eq(t.as_tensor(expected['dark'], dtype=t.float32),
dataset.background))
assert t.allclose(t.tensor(expected['data']),dataset.patterns)
assert t.allclose(t.tensor(expected['translations']),dataset.translations)
assert t.allclose(t.tensor(expected['data']), dataset.patterns)
assert t.allclose(t.tensor(expected['translations']), dataset.translations)
def test_Ptycho2DDataset_from_cxi_64bit(test_ptycho_cxis):
@@ -279,20 +273,19 @@ def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
assert np.isclose(dataset.wavelength, read_dataset.wavelength)
# Just check one of the loaded attributes
# Just check one of the loaded attributes
assert np.isclose(dataset.detector_geometry['distance'],
read_dataset.detector_geometry['distance'])
# Check that the other ones are loaded but not for fidelity
assert 'basis' in read_dataset.detector_geometry
if dataset.detector_geometry['corner'] is not None:
assert 'corner' in read_dataset.detector_geometry
if dataset.mask is not None:
assert t.all(t.eq(dataset.mask,read_dataset.mask))
assert t.all(t.eq(dataset.mask, read_dataset.mask))
if dataset.qe_mask is not None:
assert t.all(t.eq(dataset.qe_mask,read_dataset.qe_mask))
assert t.all(t.eq(dataset.qe_mask, read_dataset.qe_mask))
if dataset.background is not None:
assert t.all(t.eq(dataset.background, read_dataset.background))
@@ -303,7 +296,6 @@ def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
def test_Ptycho2DDataset_to(ptycho_cxi_1):
dataset = Ptycho2DDataset.from_cxi(ptycho_cxi_1[0])
dataset.to(dtype=t.float64)
assert dataset.mask.dtype == t.bool
assert dataset.qe_mask.dtype == t.float64
@@ -327,8 +319,8 @@ def test_Ptycho2DDataset_ops(ptycho_cxi_1):
assert len(dataset) == expected['data'].shape[0]
(idx, translation), pattern = dataset[3]
assert idx == 3
assert t.allclose(translation, t.tensor(expected['translations'][3,:]))
assert t.allclose(pattern, t.tensor(expected['data'][3,:,:]))
assert t.allclose(translation, t.tensor(expected['translations'][3, :]))
assert t.allclose(pattern, t.tensor(expected['data'][3, :, :]))
def test_Ptycho2DDataset_get_as(ptycho_cxi_1):
@@ -341,12 +333,12 @@ def test_Ptycho2DDataset_get_as(ptycho_cxi_1):
(idx, translation), pattern = dataset[3]
assert str(translation.device) == 'cuda:0'
assert str(pattern.device) == 'cuda:0'
assert idx == 3
assert t.allclose(translation.to(device='cpu'),
t.tensor(expected['translations'][3,:]))
t.tensor(expected['translations'][3, :]))
assert t.allclose(pattern.to(device='cpu'),
t.tensor(expected['data'][3,:,:]))
t.tensor(expected['data'][3, :, :]))
def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
@@ -365,18 +357,15 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
masked_patterns = dataset.mask * dataset.patterns
assert t.allclose(
copied_dataset.patterns,
masked_patterns[:,::2,::2] +
masked_patterns[:,1::2,::2] +
masked_patterns[:,::2,1::2] +
masked_patterns[:,1::2,1::2]
masked_patterns[:, ::2, ::2] + masked_patterns[:, 1::2, ::2] + masked_patterns[:, ::2, 1::2] + masked_patterns[:, 1::2, 1::2]
)
if dataset.qe_mask is None:
manually_downsampled_mask = t.logical_and(
t.logical_and(dataset.mask[::2,::2],
dataset.mask[1::2,::2]),
t.logical_and(dataset.mask[::2,1::2],
dataset.mask[1::2,1::2])
t.logical_and(dataset.mask[::2, ::2],
dataset.mask[1::2, ::2]),
t.logical_and(dataset.mask[::2, 1::2],
dataset.mask[1::2, 1::2])
)
assert t.allclose(
copied_dataset.mask,
@@ -384,10 +373,10 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
)
else:
manually_downsampled_mask = t.logical_or(
t.logical_or(dataset.mask[::2,::2],
dataset.mask[1::2,::2]),
t.logical_or(dataset.mask[::2,1::2],
dataset.mask[1::2,1::2])
t.logical_or(dataset.mask[::2, ::2],
dataset.mask[1::2, ::2]),
t.logical_or(dataset.mask[::2, 1::2],
dataset.mask[1::2, 1::2])
)
assert t.allclose(
copied_dataset.mask,
@@ -396,8 +385,7 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
masked_qe_mask = dataset.mask * dataset.qe_mask
manually_downsampled_qe_mask = (
masked_qe_mask[::2,::2] + masked_qe_mask[1::2,::2]
+ masked_qe_mask[::2,1::2] + masked_qe_mask[1::2,1::2]
masked_qe_mask[::2, ::2] + masked_qe_mask[1::2, ::2] + masked_qe_mask[::2, 1::2] + masked_qe_mask[1::2, 1::2]
) / 4
assert t.allclose(
@@ -405,15 +393,11 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
manually_downsampled_qe_mask
)
if dataset.background is not None:
assert t.allclose(
copied_dataset.background,
dataset.background[::2,::2] +
dataset.background[1::2,::2] +
dataset.background[::2,1::2] +
dataset.background[1::2,1::2]
)
dataset.background[::2, 1::2] + dataset.background[1::2, 1::2]
)
# And then we just test the shape for a few factors, and check that
# it doesn't fail on edge cases (e.g. factor=1)
@@ -428,10 +412,10 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
assert np.allclose(expected_pattern_shape,
np.array(copied_dataset.patterns.shape))
assert np.allclose(np.array(dataset.mask.shape) // factor,
np.array(copied_dataset.mask.shape))
if dataset.background is not None:
assert np.allclose(np.array(dataset.background.shape) // factor,
np.array(copied_dataset.background.shape))
@@ -468,13 +452,13 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
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
# 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_translations(roi=(0,-5,-295,-300))
assert('(i.e., patterns and translations will be empty)') in str(excinfo.value)
copied_dataset.crop_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
# and make sure that the first and last x/y elements in dataset.translate
@@ -482,40 +466,37 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
# Make tuples that will store the x and y positions. This will be used for
# making permutations of the x/y positions in the ROI later.
x_left, y_top = dataset.translations[10,:2]
x_right, y_bottom = dataset.translations[-11,:2]
x_left, y_top = dataset.translations[10, :2]
x_right, y_bottom = dataset.translations[-11, :2]
x_permutations = ((x_left, x_right), (x_right, x_left))
y_permutations = ((y_top, y_bottom), (y_bottom, y_top))
roi_permutations = tuple((x1, x2, y1, y2) for (x1, x2), (y1, y2) in
itertools.product(x_permutations, y_permutations))
roi_permutations = tuple((x1, x2, y1, y2) for (x1, x2), (y1, y2) in itertools.product(x_permutations, y_permutations))
# Get the dataset
copied_dataset.crop_translations(roi=roi_permutations[0])
# Execute the actual test
assert (copied_dataset.translations[0,0] in x_permutations[0]) and \
(copied_dataset.translations[-1,0] in x_permutations[0])
assert (copied_dataset.translations[0,1] in y_permutations[0]) and \
(copied_dataset.translations[-1,1] in y_permutations[0])
assert (copied_dataset.translations[0, 0] in x_permutations[0]) and \
(copied_dataset.translations[-1, 0] in x_permutations[0])
assert (copied_dataset.translations[0, 1] in y_permutations[0]) and \
(copied_dataset.translations[-1, 1] in y_permutations[0])
# Test 3: Check if the shape of dataset.patterns and dataset.translate is correct
# (designed to be 20 fewer rows here)
# In the future, this should include a check for dataset.intensities once
# an appropriate cxi file is set up for conftest.
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:]])
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 4: Make sure that we always get the same result no matter what order we
# define the left/right and bottom/top values in roi, provided that roi[:2]
# define the left/right and bottom/top values in roi, provided that roi[:2]
# and roi[2:] correspond with the x and y coordinates, respectively.
# Check each permutation
@@ -526,6 +507,6 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
copied_dataset.crop_translations(roi=roi)
# 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.patterns, dataset.patterns[10:-10, :])
assert t.allclose(copied_dataset.translations, dataset.translations[10:-10,:])
assert t.allclose(copied_dataset.translations, dataset.translations[10:-10, :])