From bad75422717fd0eecfe1a136fddc2b96828a3321 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Wed, 27 Mar 2019 19:03:44 -0400 Subject: [PATCH] Finish writing the Dataset classes with loading, and write up all the tests --- CDTools/datasets.py | 93 ++++++++---- CDTools/tools/data.py | 20 +-- tests/conftest.py | 283 ++++++++++++++++++++++++++++++++++++ tests/test_datasets.py | 243 +++++++++++++++++++++++++++++++ tests/tools/test_data.py | 302 +-------------------------------------- 5 files changed, 609 insertions(+), 332 deletions(-) create mode 100644 tests/conftest.py create mode 100644 tests/test_datasets.py diff --git a/CDTools/datasets.py b/CDTools/datasets.py index 5bff38c..d4d364b 100644 --- a/CDTools/datasets.py +++ b/CDTools/datasets.py @@ -4,7 +4,7 @@ import torch as t from copy import copy # The naming overlap here is definitely going to get confusing. -from CDTools import tools +from CDTools.tools import data as cdtdata from torch.utils import data as torchdata __all__ = ['CDataset', 'Ptycho_2D_Dataset'] @@ -103,26 +103,33 @@ class CDataset(torchdata.Dataset): self.wavelength = wavelength self.detector_geometry = copy(detector_geometry) if mask is not None: - self.mask = t.Tensor(mask) + self.mask = t.tensor(mask) else: self.mask = None def to(self,*args,**kwargs): - if mask is not None: - self.mask.to(*args,**kwargs) + # The mask should always stay a uint8, but it should switch devices + mask_kwargs = copy(kwargs) + try: + mask_kwargs.pop('dtype') + except KeyError as r: + pass + + if self.mask is not None: + self.mask = self.mask.to(*args,**mask_kwargs) @classmethod def from_cxi(cls, cxi_file): - entry_info = tools.get_entry_info(cxi_file) - sample_info = tools.get_sample_info(cxi_file) - wavelength = tools.get_wavelength(cxi_file) - distance, basis, corner = tools.get_detector_geometry(cxi_file) + entry_info = cdtdata.get_entry_info(cxi_file) + sample_info = cdtdata.get_sample_info(cxi_file) + wavelength = cdtdata.get_wavelength(cxi_file) + distance, basis, corner = cdtdata.get_detector_geometry(cxi_file) detector_geometry = {'distance' : distance, 'basis' : basis, 'corner' : corner} - mask = tools.get_mask(cxi_file) + mask = cdtdata.get_mask(cxi_file) return cls(entry_info = entry_info, sample_info = sample_info, wavelength=wavelength, @@ -132,42 +139,78 @@ class CDataset(torchdata.Dataset): def to_cxi(self, cxi_file): if self.entry_info is not None: - tools.add_entry_info(cxi_file, self.entry_info) + cdtdata.add_entry_info(cxi_file, self.entry_info) if self.sample_info is not None: - tools.add_sample_info(cxi_file, self.sample_info) + cdtdata.add_sample_info(cxi_file, self.sample_info) if self.wavelength is not None: - tools.add_source(cxi_file, self.wavelength) + cdtdata.add_source(cxi_file, self.wavelength) if self.detector_geometry is not None: if 'corner' in self.detector_geometry: corner = self.detector_geometry['corner'] else: corner = None - tools.add_detector(cxi_file, - self.detector_info['wavelength'], - self.detector_info['basis'], + cdtdata.add_detector(cxi_file, + self.detector_geometry['distance'], + self.detector_geometry['basis'], corner = corner) if self.mask is not None: - tools.add_mask(cxi_file, mask) + cdtdata.add_mask(cxi_file, self.mask) +# +# This is the standard dataset for a 2D ptychography experiment, +# which saves and loads files compatible with most reconstruction +# programs (only tested against SHARP) +# + class Ptycho_2D_Dataset(CDataset): - def __init__(self,translations, patterns, **kwargs): + def __init__(self, translations, patterns, axes=None, *args, **kwargs): - super(CDataset,self).__init__(kwargs) + super(Ptycho_2D_Dataset,self).__init__(*args, **kwargs) + self.axes = copy(axes) self.translations = t.tensor(translations) self.patterns = t.tensor(patterns) - def to(self, *args, **kwargs): - super(CDataset,self).to(*args,**kwargs) - self.translations.to(*args, **kwargs) - self.patterns.to(*args, **kwargs) - - def __len__(self): return self.patterns.shape[0] - def __get__(self,index): + def __getitem__(self,index): return index, self.translations[index], self.patterns[index] + + + def to(self, *args, **kwargs): + super(Ptycho_2D_Dataset,self).to(*args,**kwargs) + self.translations = self.translations.to(*args, **kwargs) + self.patterns = self.patterns.to(*args, **kwargs) + + + # It sucks that I can't reuse the base factory method here, + # perhaps there is a way but I couldn't figure it out. + @classmethod + def from_cxi(cls, cxi_file): + entry_info = cdtdata.get_entry_info(cxi_file) + sample_info = cdtdata.get_sample_info(cxi_file) + wavelength = cdtdata.get_wavelength(cxi_file) + distance, basis, corner = cdtdata.get_detector_geometry(cxi_file) + detector_geometry = {'distance' : distance, + 'basis' : basis, + 'corner' : corner} + mask = cdtdata.get_mask(cxi_file) + patterns, axes = cdtdata.get_data(cxi_file) + + translations = cdtdata.get_ptycho_translations(cxi_file) + return cls(translations, patterns, axes=axes, + entry_info = entry_info, + sample_info = sample_info, + wavelength=wavelength, + detector_geometry=detector_geometry, + mask=mask) + + + def to_cxi(self, cxi_file): + super(Ptycho_2D_Dataset,self).to_cxi(cxi_file) + cdtdata.add_data(cxi_file, self.patterns, axes=self.axes) + cdtdata.add_ptycho_translations(cxi_file, self.translations) diff --git a/CDTools/tools/data.py b/CDTools/tools/data.py index 63bf9d6..7646391 100644 --- a/CDTools/tools/data.py +++ b/CDTools/tools/data.py @@ -143,6 +143,10 @@ def get_sample_info(cxi_file): # TODO: I should add the sample geometry as a valid metadata that can # be copied over + # Check if the metadata is empty + if metadata == {}: + metadata = None + return metadata @@ -315,12 +319,10 @@ def get_data(cxi_file): def get_ptycho_translations(cxi_file): """Gets an array of x,y,z translations, if such an array has been defined in the file - It applies two operations to the translations. First, it negates them, - because the CXI file format is designed to specify translations of the - samples and the CDTools code specifies translations of the optics. - Second, it transposes the array so that the first axis is translation - ID and the second axis is the (x,y,z) components of the translation - + It negates the translations, because the CXI file format is designed + to specify translations of the samples and the CDTools code specifies + translations of the optics. + Args: cxi_file (h5py.File) : a file object to be read @@ -338,7 +340,7 @@ def get_ptycho_translations(cxi_file): else: raise KeyError('Translations are not defined within cxi file') - translations = -np.array(cxi_file[pull_from]).astype(np.float32).transpose() + translations = -np.array(cxi_file[pull_from]).astype(np.float32) return translations @@ -537,7 +539,7 @@ def add_ptycho_translations(cxi_file, translations): It will add the translations to the file, negating them to conform to the standard in cxi files that the translations refer to the object's - translation, and also transposing them to match the cxi file specification. + translation. It will generally store them in 3 places: @@ -575,7 +577,7 @@ def add_ptycho_translations(cxi_file, translations): # accounting for the different definition between cxi files and # CDTools - translations = -translations.transpose() + translations = -translations g1.create_dataset('translation', data=translations) data1['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation') diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..a062cf6 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,283 @@ +from __future__ import division, print_function, absolute_import +import numpy as np +import h5py +import pytest +import datetime + + +# +# +# The following few fixtures define some standard data files +# for use to test the data loading capabilities, whether in the +# datasets directly or in the data tools file +# +# + +@pytest.fixture(scope='module') +def ptycho_cxi_1(): + """Creates an example file for CXI ptychography. This file is defined + to have everything done as correctly as possible with lots of attributes + defined. It will return both a dictionary describing what is expected + to be loaded and a file with the data stored in it. + """ + + expected = {} + f = h5py.File('ptycho_cxi_1',driver='core',backing_store=False) + + # Start by defining the basic structure + f.create_dataset('cxi_version', data=150) + f.create_dataset('number_of_entries',data=1) + + # Then define a bunch of metadata for entry_1 + e1f = f.create_group('entry_1') + expected['entry metadata'] = {} + e1e = expected['entry metadata'] + e1e['start_time'] = datetime.datetime.now() + e1f['start_time'] = np.string_(e1e['start_time'].isoformat()) + e1e['end_time'] = datetime.datetime.now() + e1f['end_time'] = np.string_(e1e['end_time'].isoformat()) + e1e['experiment_identifier'] = 'Fake Experiment 1' + e1f['experiment_identifier'] = np.string_(e1e['experiment_identifier']) + e1e['experiment_description'] = 'A fully defined ptychography experiment to test the data loading' + e1f['experiment_description'] = np.string_(e1e['experiment_description']) + e1e['program_name'] = 'CDTools' + e1f['program_name'] = np.string_(e1e['program_name']) + e1e['title'] = 'The one experiment we did' + e1f['title'] = np.string_(e1e['title']) + + # Set up the sample info + s1f = e1f.create_group('sample_1') + expected['sample info'] = {} + s1e = expected['sample info'] + s1e['name'] = 'Fake Sample' + s1f['name'] = np.string_(s1e['name']) + s1e['description'] = 'A sample that isn\'t real' + s1f['description'] = np.string_(s1e['description']) + s1e['unit_cell_group'] = 'P1' + s1f['unit_cell_group'] = np.string_(s1e['unit_cell_group']) + s1e['concentration'] = np.float32(np.random.rand()) + s1f['concentration'] = s1e['concentration'] + s1e['mass'] = np.float32(np.random.rand()) + s1f['mass'] = s1e['mass'] + s1e['temperature'] = np.float32(np.random.rand()*100) + s1f['temperature'] = s1e['temperature'] + s1e['thickness'] = np.float32(np.random.rand()*1e-7) + s1f['thickness'] = s1e['thickness'] + s1e['unit_cell_volume'] = np.float32(np.random.rand() * 1e-27) + s1f['unit_cell_volume'] = s1e['unit_cell_volume'] + s1e['unit_cell'] = np.array([1,1,1,90,90,90]).astype(np.float32) + s1f.create_dataset('unit_cell',data = s1e['unit_cell']) + + i1f = e1f.create_group('instrument_1') + source1f = i1f.create_group('source_1') + + energy = np.float32(1.3618e-16) #Joules, = 850 eV + source1f['energy'] = energy + expected['wavelength'] = np.float32(1.9864459e-25) / energy + source1f['wavelength'] = expected['wavelength'] + + d1f = i1f.create_group('detector_1') + expected['detector'] = {} + d1e = expected['detector'] + d1e['distance'] = np.float32(0.3) + d1f['distance'] = d1e['distance'] + d1e['basis'] = np.array([[0,-30e-6,0], + [-20e-6,0,0]]).astype(np.float32).transpose() + d1f.create_dataset('basis_vectors',data=d1e['basis']) + d1f['x_pixel_size'] = np.float32(20e-6) + d1f['y_pixel_size'] = np.float32(30e-6) + d1e['corner'] = np.array((2550e-6,3825e-6,0.3)).astype(np.float32) + d1f.create_dataset('corner_position', data=d1e['corner']) + + # Remember the format for the CXI file differs from the format used + # internally + mask = np.zeros((100,256,256)).astype(np.uint32) + expected['mask'] = np.ones((100,256,256)).astype(np.uint8) + d1f.create_dataset('mask',data=mask) + + data1f = e1f.create_group('data_1') + + data = np.random.rand(100,256,256).astype(np.float32) + expected['data'] = data + d1f.create_dataset('data',data=data) + data1f['data'] = h5py.SoftLink('/entry_1/instrument_1/detector_1/data') + + d1f['data'].attrs['axes'] = np.string_('translation:y:x') + expected['axes'] = ['translation','y','x'] + + g1f = s1f.create_group('geometry_1') + translations = np.arange(300).reshape((100,3)).astype(np.float32) + g1f.create_dataset('translation',data=translations) + data1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation') + d1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation') + expected['translations'] = -translations + + yield f, expected + + f.close() + + +@pytest.fixture(scope='module') +def ptycho_cxi_2(): + """Creates an example file for CXI ptychography. This file is defined + to have a subset of things missing. In particular, it: + + * Defines the wavelength but not the energy + * Defines the corner position but not the sample-detector distance + * Defines pixel sizes but no basis vectors + * Doesn't define a mask + * Only defines data in the relevant places, not under data_1 + * Doesn't explicitly define axes for the data arrays + * Is missing many allowed metadata attributes + + """ + + expected = {} + f = h5py.File('ptycho_cxi_2',driver='core',backing_store=False) + + # Start by defining the basic structure + f.create_dataset('cxi_version', data=150) + f.create_dataset('number_of_entries',data=1) + + # Then define a bunch of metadata for entry_1 + e1f = f.create_group('entry_1') + expected['entry metadata'] = {} + e1e = expected['entry metadata'] + e1e['title'] = 'The one experiment we did' + e1f['title'] = np.string_(e1e['title']) + + # Set up the sample info + s1f = e1f.create_group('sample_1') + expected['sample info'] = {} + s1e = expected['sample info'] + s1e['temperature'] = np.float32(np.random.rand()*100) + s1f['temperature'] = s1e['temperature'] + + i1f = e1f.create_group('instrument_1') + source1f = i1f.create_group('source_1') + + energy = np.float32(1.3618e-16) #Joules, = 850 eV + expected['wavelength'] = np.float32(1.9864459e-25) / energy + source1f['wavelength'] = expected['wavelength'] + + d1f = i1f.create_group('detector_1') + expected['detector'] = {} + d1e = expected['detector'] + d1e['distance'] = np.float32(0.3) + d1e['basis'] = np.array([[0,-30e-6,0], + [-20e-6,0,0]]).astype(np.float32).transpose() + d1f['x_pixel_size'] = np.float32(20e-6) + d1f['y_pixel_size'] = np.float32(30e-6) + d1e['corner'] = np.array((2550e-6,3825e-6,0.3)).astype(np.float32) + d1f.create_dataset('corner_position', data=d1e['corner']) + + # Remember the format for the CXI file differs from the format used + # internally + expected['mask'] = None + + data1f = e1f.create_group('data_1') + + data = np.random.rand(100,256,256).astype(np.float32) + expected['data'] = data + d1f.create_dataset('data',data=data) + + expected['axes'] = None + + g1f = s1f.create_group('geometry_1') + translations = np.arange(300).reshape((100,3)).astype(np.float32) + g1f.create_dataset('translation',data=translations) + expected['translations'] = -translations + + yield f, expected + + f.close() + + +@pytest.fixture(scope='module') +def ptycho_cxi_3(): + """Creates an example file for CXI ptychography. This file is defined + to have a different subset of information missing. In particular, it: + + * Has no sample_1 group + * Defines the energy but not the wavelength of light + * Has the data only defined under the data_1 group, not in the relevant places + * Defines the detector basis but no pixel sizes + * Defines a mask as all pixels flagged as "above the background" + * Defines the sample to detector distance but no corner location + * Is missing some of the allowed metadata + """ + + expected = {} + f = h5py.File('ptycho_cxi_3',driver='core',backing_store=False) + + # Start by defining the basic structure + f.create_dataset('cxi_version', data=150) + f.create_dataset('number_of_entries',data=1) + + # Then define a bunch of metadata for entry_1 + e1f = f.create_group('entry_1') + expected['entry metadata'] = {} + e1e = expected['entry metadata'] + e1e['start_time'] = datetime.datetime.now() + e1f['start_time'] = np.string_(e1e['start_time'].isoformat()) + e1e['end_time'] = datetime.datetime.now() + e1f['end_time'] = np.string_(e1e['end_time'].isoformat()) + + # Set up the sample info + expected['sample info'] = None + + i1f = e1f.create_group('instrument_1') + source1f = i1f.create_group('source_1') + + energy = np.float32(1.3618e-16) #Joules, = 850 eV + source1f['energy'] = energy + expected['wavelength'] = np.float32(1.9864459e-25) / energy + + d1f = i1f.create_group('detector_1') + expected['detector'] = {} + d1e = expected['detector'] + d1e['distance'] = np.float32(0.3) + d1f['distance'] = d1e['distance'] + d1e['basis'] = np.array([[0,-30e-6,0], + [-20e-6,0,0]]).astype(np.float32).transpose() + d1f.create_dataset('basis_vectors',data=d1e['basis']) + d1e['corner'] = None + + # Remember the format for the CXI file differs from the format used + # internally + mask = np.ones((100,256,256)).astype(np.uint32) * 0x00001000 + expected['mask'] = np.ones((100,256,256)).astype(np.uint8) + d1f.create_dataset('mask',data=mask) + + data1f = e1f.create_group('data_1') + + data = np.random.rand(100,256,256).astype(np.float32) + expected['data'] = data + data1f.create_dataset('data',data=data) + + data1f['data'].attrs['axes'] = np.string_('translation:y:x') + expected['axes'] = ['translation','y','x'] + + translations = np.arange(300).reshape((100,3)).astype(np.float32) + data1f.create_dataset('translation',data=translations) + expected['translations'] = -translations + + yield f, expected + + f.close() + + +# As specific issues start to crop up with loading CXI files from different +# beamlines, put a fixture here that replicates the issue so that we can +# ensure compatibility with many beamlines +# + + +@pytest.fixture(scope='module') +def test_ptycho_cxis(ptycho_cxi_1, ptycho_cxi_2, ptycho_cxi_3): + """Loads a list of tuples of ptychography CXI files and dictionaries, + describing the expected output from various functions on being called + on the cxi files. + """ + return [ptycho_cxi_1, ptycho_cxi_2, ptycho_cxi_3] + diff --git a/tests/test_datasets.py b/tests/test_datasets.py new file mode 100644 index 0000000..02dc473 --- /dev/null +++ b/tests/test_datasets.py @@ -0,0 +1,243 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.datasets import * +from CDTools.tools import data as cdtdata +import numpy as np +import torch as t +import h5py +import pytest +import datetime + + +# +# 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))) + 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)) + + +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)) + + + +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.uint8 + # 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') + + + +# +# And we then test the derived Ptychography class +# +# + + +def test_Ptycho_2D_Dataset_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 = Ptycho_2D_Dataset(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.tensor(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.tensor(patterns)) + assert t.allclose(dataset.translations, t.tensor(translations)) + + + +def test_Ptycho_2D_Dataset_from_cxi(test_ptycho_cxis): + for cxi, expected in test_ptycho_cxis: + dataset = Ptycho_2D_Dataset.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)) + + assert t.allclose(t.tensor(expected['data']),dataset.patterns) + assert t.allclose(t.tensor(expected['translations']),dataset.translations) + + + + +def test_Ptycho_2D_Dataset_to_cxi(test_ptycho_cxis, tmp_path): + for cxi, expected in test_ptycho_cxis: + dataset = Ptycho_2D_Dataset.from_cxi(cxi) + with cdtdata.create_cxi(tmp_path / 'test_Ptycho_2D_Dataset_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_Ptycho_2D_Dataset_to_cxi.cxi', 'r') as f: + read_dataset = Ptycho_2D_Dataset.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)) + + assert t.allclose(dataset.patterns, read_dataset.patterns) + assert t.allclose(dataset.translations, read_dataset.translations) + + +def test_Ptycho_2D_Dataset_to(ptycho_cxi_1): + dataset = Ptycho_2D_Dataset.from_cxi(ptycho_cxi_1[0]) + + dataset.to(dtype=t.float64) + assert dataset.mask.dtype == t.uint8 + 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.patterns.device == t.device('cuda:0') + assert dataset.translations.device == t.device('cuda:0') + + +def test_Ptycho_2D_Dataset_ops(ptycho_cxi_1): + cxi, expected = ptycho_cxi_1 + dataset = Ptycho_2D_Dataset.from_cxi(cxi) + + 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,:,:])) diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py index 78630c3..573a8e5 100644 --- a/tests/tools/test_data.py +++ b/tests/tools/test_data.py @@ -11,303 +11,9 @@ import numbers from pathlib import Path -# -# First, we write a few fixtures to generate specific data files we -# want with deliberate pathologies. -# -# * Energy defined but no wavelength -# * Wavelength defined but no energy -# * Distance and pixel pitch defined but no corner position or basis -# * Corner position and basis defined but no distance or pixel pitch -# * No mask defined -# * Mask defined -# -# Each file will get a fixture that loads the cxi file but also loads -# a dictionary with the relevant information for the cxi file -# -# - - -# This just grabs the directory whose name matches the file we're running -@pytest.fixture(scope='module') -def datadir(request): - filename = request.module.__file__ - test_dir, _ = os.path.splitext(filename) - return Path(test_dir) - - -@pytest.fixture(scope='module') -def ptycho_cxi_1(): - """Creates an example file for CXI ptychography. This file is defined - to have everything done as correctly as possible with lots of attributes - defined. It will return both a dictionary describing what is expected - to be loaded and a file with the data stored in it. - """ - - expected = {} - f = h5py.File('ptycho_cxi_1',driver='core',backing_store=False) - - # Start by defining the basic structure - f.create_dataset('cxi_version', data=150) - f.create_dataset('number_of_entries',data=1) - - # Then define a bunch of metadata for entry_1 - e1f = f.create_group('entry_1') - expected['entry metadata'] = {} - e1e = expected['entry metadata'] - e1e['start_time'] = datetime.datetime.now() - e1f['start_time'] = np.string_(e1e['start_time'].isoformat()) - e1e['end_time'] = datetime.datetime.now() - e1f['end_time'] = np.string_(e1e['end_time'].isoformat()) - e1e['experiment_identifier'] = 'Fake Experiment 1' - e1f['experiment_identifier'] = np.string_(e1e['experiment_identifier']) - e1e['experiment_description'] = 'A fully defined ptychography experiment to test the data loading' - e1f['experiment_description'] = np.string_(e1e['experiment_description']) - e1e['program_name'] = 'CDTools' - e1f['program_name'] = np.string_(e1e['program_name']) - e1e['title'] = 'The one experiment we did' - e1f['title'] = np.string_(e1e['title']) - - # Set up the sample info - s1f = e1f.create_group('sample_1') - expected['sample info'] = {} - s1e = expected['sample info'] - s1e['name'] = 'Fake Sample' - s1f['name'] = np.string_(s1e['name']) - s1e['description'] = 'A sample that isn\'t real' - s1f['description'] = np.string_(s1e['description']) - s1e['unit_cell_group'] = 'P1' - s1f['unit_cell_group'] = np.string_(s1e['unit_cell_group']) - s1e['concentration'] = np.float32(np.random.rand()) - s1f['concentration'] = s1e['concentration'] - s1e['mass'] = np.float32(np.random.rand()) - s1f['mass'] = s1e['mass'] - s1e['temperature'] = np.float32(np.random.rand()*100) - s1f['temperature'] = s1e['temperature'] - s1e['thickness'] = np.float32(np.random.rand()*1e-7) - s1f['thickness'] = s1e['thickness'] - s1e['unit_cell_volume'] = np.float32(np.random.rand() * 1e-27) - s1f['unit_cell_volume'] = s1e['unit_cell_volume'] - s1e['unit_cell'] = np.array([1,1,1,90,90,90]).astype(np.float32) - s1f.create_dataset('unit_cell',data = s1e['unit_cell']) - - i1f = e1f.create_group('instrument_1') - source1f = i1f.create_group('source_1') - - energy = np.float32(1.3618e-16) #Joules, = 850 eV - source1f['energy'] = energy - expected['wavelength'] = np.float32(1.9864459e-25) / energy - source1f['wavelength'] = expected['wavelength'] - - d1f = i1f.create_group('detector_1') - expected['detector'] = {} - d1e = expected['detector'] - d1e['distance'] = np.float32(0.3) - d1f['distance'] = d1e['distance'] - d1e['basis'] = np.array([[0,-30e-6,0], - [-20e-6,0,0]]).astype(np.float32).transpose() - d1f.create_dataset('basis_vectors',data=d1e['basis']) - d1f['x_pixel_size'] = np.float32(20e-6) - d1f['y_pixel_size'] = np.float32(30e-6) - d1e['corner'] = np.array((2550e-6,3825e-6,0.3)).astype(np.float32) - d1f.create_dataset('corner_position', data=d1e['corner']) - - # Remember the format for the CXI file differs from the format used - # internally - mask = np.zeros((100,256,256)).astype(np.uint32) - expected['mask'] = np.ones((100,256,256)).astype(np.uint8) - d1f.create_dataset('mask',data=mask) - - data1f = e1f.create_group('data_1') - - data = np.random.rand(100,256,256).astype(np.float32) - expected['data'] = data - d1f.create_dataset('data',data=data) - data1f['data'] = h5py.SoftLink('/entry_1/instrument_1/detector_1/data') - - d1f['data'].attrs['axes'] = np.string_('translation:y:x') - expected['axes'] = ['translation','y','x'] - - g1f = s1f.create_group('geometry_1') - translations = np.arange(300).reshape((100,3)).astype(np.float32) - g1f.create_dataset('translation',data=translations) - data1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation') - d1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation') - expected['translations'] = -translations.transpose() - - yield f, expected - - f.close() - - -@pytest.fixture(scope='module') -def ptycho_cxi_2(): - """Creates an example file for CXI ptychography. This file is defined - to have a subset of things missing. In particular, it: - - * Defines the wavelength but not the energy - * Defines the corner position but not the sample-detector distance - * Defines pixel sizes but no basis vectors - * Doesn't define a mask - * Only defines data in the relevant places, not under data_1 - * Doesn't explicitly define axes for the data arrays - * Is missing many allowed metadata attributes - - """ - - expected = {} - f = h5py.File('ptycho_cxi_2',driver='core',backing_store=False) - - # Start by defining the basic structure - f.create_dataset('cxi_version', data=150) - f.create_dataset('number_of_entries',data=1) - - # Then define a bunch of metadata for entry_1 - e1f = f.create_group('entry_1') - expected['entry metadata'] = {} - e1e = expected['entry metadata'] - e1e['title'] = 'The one experiment we did' - e1f['title'] = np.string_(e1e['title']) - - # Set up the sample info - s1f = e1f.create_group('sample_1') - expected['sample info'] = {} - s1e = expected['sample info'] - s1e['temperature'] = np.float32(np.random.rand()*100) - s1f['temperature'] = s1e['temperature'] - - i1f = e1f.create_group('instrument_1') - source1f = i1f.create_group('source_1') - - energy = np.float32(1.3618e-16) #Joules, = 850 eV - expected['wavelength'] = np.float32(1.9864459e-25) / energy - source1f['wavelength'] = expected['wavelength'] - - d1f = i1f.create_group('detector_1') - expected['detector'] = {} - d1e = expected['detector'] - d1e['distance'] = np.float32(0.3) - d1e['basis'] = np.array([[0,-30e-6,0], - [-20e-6,0,0]]).astype(np.float32).transpose() - d1f['x_pixel_size'] = np.float32(20e-6) - d1f['y_pixel_size'] = np.float32(30e-6) - d1e['corner'] = np.array((2550e-6,3825e-6,0.3)).astype(np.float32) - d1f.create_dataset('corner_position', data=d1e['corner']) - - # Remember the format for the CXI file differs from the format used - # internally - expected['mask'] = None - - data1f = e1f.create_group('data_1') - - data = np.random.rand(100,256,256).astype(np.float32) - expected['data'] = data - d1f.create_dataset('data',data=data) - - expected['axes'] = None - - g1f = s1f.create_group('geometry_1') - translations = np.arange(300).reshape((100,3)).astype(np.float32) - g1f.create_dataset('translation',data=translations) - expected['translations'] = -translations.transpose() - - yield f, expected - - f.close() - - -@pytest.fixture(scope='module') -def ptycho_cxi_3(): - """Creates an example file for CXI ptychography. This file is defined - to have a different subset of information missing. In particular, it: - - * Has no sample_1 group - * Defines the energy but not the wavelength of light - * Has the data only defined under the data_1 group, not in the relevant places - * Defines the detector basis but no pixel sizes - * Defines a mask as all pixels flagged as "above the background" - * Defines the sample to detector distance but no corner location - * Is missing some of the allowed metadata - """ - - expected = {} - f = h5py.File('ptycho_cxi_3',driver='core',backing_store=False) - - # Start by defining the basic structure - f.create_dataset('cxi_version', data=150) - f.create_dataset('number_of_entries',data=1) - - # Then define a bunch of metadata for entry_1 - e1f = f.create_group('entry_1') - expected['entry metadata'] = {} - e1e = expected['entry metadata'] - e1e['start_time'] = datetime.datetime.now() - e1f['start_time'] = np.string_(e1e['start_time'].isoformat()) - e1e['end_time'] = datetime.datetime.now() - e1f['end_time'] = np.string_(e1e['end_time'].isoformat()) - - # Set up the sample info - expected['sample info'] = None - - i1f = e1f.create_group('instrument_1') - source1f = i1f.create_group('source_1') - - energy = np.float32(1.3618e-16) #Joules, = 850 eV - source1f['energy'] = energy - expected['wavelength'] = np.float32(1.9864459e-25) / energy - - d1f = i1f.create_group('detector_1') - expected['detector'] = {} - d1e = expected['detector'] - d1e['distance'] = np.float32(0.3) - d1f['distance'] = d1e['distance'] - d1e['basis'] = np.array([[0,-30e-6,0], - [-20e-6,0,0]]).astype(np.float32).transpose() - d1f.create_dataset('basis_vectors',data=d1e['basis']) - d1e['corner'] = None - - # Remember the format for the CXI file differs from the format used - # internally - mask = np.ones((100,256,256)).astype(np.uint32) * 0x00001000 - expected['mask'] = np.ones((100,256,256)).astype(np.uint8) - d1f.create_dataset('mask',data=mask) - - data1f = e1f.create_group('data_1') - - data = np.random.rand(100,256,256).astype(np.float32) - expected['data'] = data - data1f.create_dataset('data',data=data) - - data1f['data'].attrs['axes'] = np.string_('translation:y:x') - expected['axes'] = ['translation','y','x'] - - translations = np.arange(300).reshape((100,3)).astype(np.float32) - data1f.create_dataset('translation',data=translations) - expected['translations'] = -translations.transpose() - - yield f, expected - - f.close() - - -# As specific issues start to crop up with loading CXI files from different -# beamlines, put a fixture here that replicates the issue so that we can -# ensure compatibility with many beamlines -# - - -@pytest.fixture(scope='module') -def test_ptycho_cxis(ptycho_cxi_1, ptycho_cxi_2, ptycho_cxi_3): - """Loads a list of tuples of ptychography CXI files and dictionaries, - describing the expected output from various functions on being called - on the cxi files. - """ - return [ptycho_cxi_1, ptycho_cxi_2, ptycho_cxi_3] - - # -# Now we have a bunch of tests of the data loading capabilities +# We start with a bunch of tests of the data loading capabilities # @@ -530,6 +236,6 @@ def test_add_ptycho_translations(tmp_path): read_translations_2 = np.array(f['entry_1/instrument_1/detector_1/translation']) read_translations_3 = np.array(f['entry_1/sample_1/geometry_1/translation']) - assert np.allclose(-translations.transpose(), read_translations_1) - assert np.allclose(-translations.transpose(), read_translations_2) - assert np.allclose(-translations.transpose(), read_translations_3) + assert np.allclose(-translations, read_translations_1) + assert np.allclose(-translations, read_translations_2) + assert np.allclose(-translations, read_translations_3)