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

Finish writing the Dataset classes with loading, and write up all the tests

Browse Source
This commit is contained in:
Abe Levitan
2019-03-27 19:03:44 -04:00
parent 78d13a3f67
commit bad7542271
5 changed files with 609 additions and 332 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/datasets.py
+68 -25
View File
@@ -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)
CDTools/tools/data.py
+11 -9
View File
@@ -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')
tests/conftest.py
+283
View File
@@ -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]
tests/test_datasets.py
+243
View File
@@ -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,:,:]))
tests/tools/test_data.py
+4 -298
View File
@@ -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)