mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-11 14:02:38 +02:00
395 lines
15 KiB
Python
395 lines
15 KiB
Python
from cdtools.datasets import *
|
|
from cdtools.tools import data as cdtdata
|
|
import numpy as np
|
|
import torch as t
|
|
import h5py
|
|
import datetime
|
|
from copy import deepcopy
|
|
import pytest
|
|
|
|
#
|
|
# We start by testing the CDataset base class
|
|
#
|
|
|
|
def test_CDataset_init():
|
|
entry_info = {'start_time': datetime.datetime.now(),
|
|
'title' : 'A simple test'}
|
|
sample_info = {'name': 'A test sample',
|
|
'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))
|
|
dataset = CDataset(entry_info, sample_info,
|
|
wavelength, detector_geometry, mask)
|
|
|
|
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
|
|
assert dataset.detector_geometry == detector_geometry
|
|
|
|
|
|
def test_CDataset_from_cxi(test_ptycho_cxis):
|
|
for cxi, expected in test_ptycho_cxis:
|
|
dataset = CDataset.from_cxi(cxi)
|
|
|
|
# The entry metadata loaded
|
|
for key in expected['entry metadata']:
|
|
assert dataset.entry_info[key] == expected['entry metadata'][key]
|
|
|
|
# Don't test for fidelity since this is tested in the data, just test
|
|
# that it is loaded
|
|
if expected['sample info'] is None:
|
|
assert dataset.sample_info is None
|
|
else:
|
|
assert dataset.sample_info is not None
|
|
|
|
assert np.isclose(dataset.wavelength,expected['wavelength'])
|
|
|
|
# Just check one of the loaded attributes
|
|
assert np.isclose(dataset.detector_geometry['distance'],
|
|
expected['detector']['distance'])
|
|
# Check that the other ones are loaded but not for fidelity
|
|
assert 'basis' in dataset.detector_geometry
|
|
if expected['detector']['corner'] is not None:
|
|
assert 'corner' in dataset.detector_geometry
|
|
|
|
if expected['mask'] is not None:
|
|
assert t.all(t.eq(t.tensor(expected['mask']),dataset.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)
|
|
with cdtdata.create_cxi(tmp_path / 'test_CDataset_to_cxi.cxi') as f:
|
|
dataset.to_cxi(f)
|
|
|
|
# Now we have to check that all the stuff was written
|
|
with h5py.File(tmp_path / 'test_CDataset_to_cxi.cxi', 'r') as f:
|
|
read_dataset = CDataset.from_cxi(f)
|
|
|
|
assert dataset.entry_info == read_dataset.entry_info
|
|
|
|
if dataset.sample_info is None:
|
|
assert read_dataset.sample_info is None
|
|
else:
|
|
assert read_dataset.sample_info is not None
|
|
|
|
assert np.isclose(dataset.wavelength, read_dataset.wavelength)
|
|
|
|
|
|
# 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))
|
|
|
|
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])
|
|
|
|
dataset.to(dtype=t.float32)
|
|
assert dataset.mask.dtype == t.bool
|
|
# If cuda is available, check that moving the mask to CUDA works.
|
|
if t.cuda.is_available():
|
|
dataset.to(device='cuda:0')
|
|
assert dataset.mask.device == t.device('cuda:0')
|
|
assert dataset.background.device == t.device('cuda:0')
|
|
|
|
|
|
#
|
|
# And we then test the derived Ptychography class
|
|
#
|
|
|
|
|
|
def test_Ptycho2DDataset_init():
|
|
entry_info = {'start_time': datetime.datetime.now(),
|
|
'title' : 'A simple test'}
|
|
sample_info = {'name': 'A test sample',
|
|
'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))
|
|
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)
|
|
|
|
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
|
|
assert dataset.detector_geometry == detector_geometry
|
|
assert t.allclose(dataset.patterns, t.as_tensor(patterns))
|
|
assert t.allclose(dataset.translations, t.as_tensor(translations))
|
|
|
|
|
|
def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
|
|
for cxi, expected in test_ptycho_cxis:
|
|
dataset = Ptycho2DDataset.from_cxi(cxi)
|
|
|
|
# The entry metadata loaded
|
|
for key in expected['entry metadata']:
|
|
assert dataset.entry_info[key] == expected['entry metadata'][key]
|
|
|
|
# Don't test for fidelity since this is tested in the data, just test
|
|
# that it is loaded
|
|
if expected['sample info'] is None:
|
|
assert dataset.sample_info is None
|
|
else:
|
|
assert dataset.sample_info is not None
|
|
|
|
assert np.isclose(dataset.wavelength,expected['wavelength'])
|
|
|
|
# Just check one of the loaded attributes
|
|
assert np.isclose(dataset.detector_geometry['distance'],
|
|
expected['detector']['distance'])
|
|
# Check that the other ones are loaded but not for fidelity
|
|
assert 'basis' in dataset.detector_geometry
|
|
if expected['detector']['corner'] is not None:
|
|
assert 'corner' in dataset.detector_geometry
|
|
|
|
if expected['mask'] is not None:
|
|
assert t.all(t.eq(t.tensor(expected['mask']),dataset.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)
|
|
|
|
|
|
def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
|
|
for cxi, expected in test_ptycho_cxis:
|
|
print('loading dataset')
|
|
dataset = Ptycho2DDataset.from_cxi(cxi)
|
|
print('dataset mask is type', dataset.mask.dtype)
|
|
with cdtdata.create_cxi(tmp_path / 'test_Ptycho2DDataset_to_cxi.cxi') as f:
|
|
dataset.to_cxi(f)
|
|
|
|
# Now we have to check that all the stuff was written
|
|
with h5py.File(tmp_path / 'test_Ptycho2DDataset_to_cxi.cxi', 'r') as f:
|
|
read_dataset = Ptycho2DDataset.from_cxi(f)
|
|
|
|
assert dataset.entry_info == read_dataset.entry_info
|
|
|
|
if dataset.sample_info is None:
|
|
assert read_dataset.sample_info is None
|
|
else:
|
|
assert read_dataset.sample_info is not None
|
|
|
|
assert np.isclose(dataset.wavelength, read_dataset.wavelength)
|
|
|
|
|
|
# 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))
|
|
|
|
if dataset.background is not None:
|
|
assert t.all(t.eq(dataset.background, read_dataset.background))
|
|
|
|
assert t.allclose(dataset.patterns, read_dataset.patterns)
|
|
assert t.allclose(dataset.translations, read_dataset.translations)
|
|
|
|
|
|
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.patterns.dtype == t.float64
|
|
assert dataset.translations.dtype == t.float64
|
|
# If cuda is available, check that moving the mask to CUDA works.
|
|
if t.cuda.is_available():
|
|
dataset.to(device='cuda:0')
|
|
assert dataset.mask.device == t.device('cuda:0')
|
|
assert dataset.background.device == t.device('cuda:0')
|
|
assert dataset.patterns.device == t.device('cuda:0')
|
|
assert dataset.translations.device == t.device('cuda:0')
|
|
|
|
|
|
def test_Ptycho2DDataset_ops(ptycho_cxi_1):
|
|
cxi, expected = ptycho_cxi_1
|
|
dataset = Ptycho2DDataset.from_cxi(cxi)
|
|
dataset.get_as('cpu')
|
|
|
|
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,:,:]))
|
|
|
|
|
|
def test_Ptycho2DDataset_get_as(ptycho_cxi_1):
|
|
cxi, expected = ptycho_cxi_1
|
|
dataset = Ptycho2DDataset.from_cxi(cxi)
|
|
if t.cuda.is_available():
|
|
dataset.get_as('cuda:0')
|
|
assert len(dataset) == expected['data'].shape[0]
|
|
|
|
(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,:]))
|
|
assert t.allclose(pattern.to(device='cpu'),
|
|
t.tensor(expected['data'][3,:,:]))
|
|
|
|
|
|
def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
|
|
for cxi, expected in test_ptycho_cxis:
|
|
dataset = Ptycho2DDataset.from_cxi(cxi)
|
|
|
|
# First we test the case of downsampling by 2 against some explicit
|
|
# calculations
|
|
copied_dataset = deepcopy(dataset)
|
|
copied_dataset.downsample(2)
|
|
|
|
# May start failing if the test datasets are changed to include
|
|
# a dataset with any dimension not even. That's a problem with the
|
|
# test, not the code. Sorry! -Abe
|
|
assert t.allclose(
|
|
copied_dataset.patterns,
|
|
dataset.patterns[:,::2,::2] +
|
|
dataset.patterns[:,1::2,::2] +
|
|
dataset.patterns[:,::2,1::2] +
|
|
dataset.patterns[:,1::2,1::2]
|
|
)
|
|
|
|
assert t.allclose(
|
|
copied_dataset.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]),
|
|
)
|
|
)
|
|
|
|
|
|
|
|
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]
|
|
)
|
|
|
|
# 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)
|
|
for factor in [1, 2, 3]:
|
|
copied_dataset = deepcopy(dataset)
|
|
copied_dataset.downsample(factor=factor)
|
|
|
|
expected_pattern_shape = np.concatenate(
|
|
[[dataset.patterns.shape[0]],
|
|
np.array(dataset.patterns.shape[-2:]) // factor]
|
|
)
|
|
|
|
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))
|
|
|
|
|
|
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,:])
|