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

Added a polarization dependent ptychography dataset class so we can save and load data!

Browse Source
This commit is contained in:
Abe Levitan
2021-07-26 14:46:02 -04:00
parent a36c6c8893
commit e7d7323e50
10 changed files with 560 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/datasets/__init__.py
+2 -1
View File
@@ -29,7 +29,8 @@ from __future__ import division, print_function, absolute_import
# I don't believe that __all__ really needed, but it's nice to define it
# to be explicit that import * is safe
__all__ = ['CDataset','Ptycho2DDataset']
__all__ = ['CDataset','Ptycho2DDataset','PolarizedPtycho2DDataset']
from CDTools.datasets.base import CDataset
from CDTools.datasets.ptycho_2d_dataset import Ptycho2DDataset
from CDTools.datasets.polarized_ptycho_2d_dataset import PolarizedPtycho2DDataset
CDTools/datasets/base.py
+3 -10
View File
@@ -19,11 +19,7 @@ import numpy as np
import torch as t
from copy import copy
import h5py
try:
import pathlib
except ImportError:
import pathlib2 as pathlib
import pathlib
from CDTools.tools import data as cdtdata
from CDTools.tools import plotting
from torch.utils import data as torchdata
@@ -88,14 +84,11 @@ class CDataset(torchdata.Dataset):
self.wavelength = wavelength
self.detector_geometry = copy(detector_geometry)
if mask is not None:
if isinstance(mask, t.Tensor):
self.mask = mask.detach().to(dtype=t.bool)
else:
self.mask = t.BoolTensor(mask)
self.mask = t.tensor(mask, dtype=t.bool)
else:
self.mask = None
if background is not None:
self.background = t.Tensor(background)
self.background = t.tensor(background, dtype=t.float32)
else:
self.background = None
CDTools/datasets/polarized_ptycho_2d_dataset.py
+219
View File
@@ -0,0 +1,219 @@
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
from copy import copy
import h5py
import pathlib
from CDTools.datasets import CDataset, Ptycho2DDataset
from CDTools.tools import data as cdtdata
from CDTools.tools import plotting
from torch.utils import data as torchdata
from matplotlib import pyplot as plt
from matplotlib.widgets import Slider
from matplotlib import ticker
__all__ = ['PolarizedPtycho2DDataset']
class PolarizedPtycho2DDataset(Ptycho2DDataset):
"""The standard dataset for a 2D ptychography scan
Subclasses datasets.CDataset
This class loads and saves 2D ptychography scan data from .cxi files.
It should save and load files compatible with most reconstruction
programs, although it is only tested against SHARP.
"""
def __init__(self, translations, polarizer_angles, analyzer_angles, patterns, axes=None, *args, **kwargs):
"""The __init__ function allows construction from python objects.
The detector_geometry dictionary is defined to have the
entries defined by the outputs of data.get_detector_geometry.
Parameters
----------
translations : array (nx3 t.tensor when polarized=True)
An nx3 array containing the probe translations at each scan point
patterns : array
An nxmxl array containing the full stack of measured diffraction patterns
axes : list(str)
A
of names for the axes of the probe translations
entry_info : dict
A dictionary containing the entry_info metadata
sample_info : dict
A dictionary containing the sample_info metadata
wavelength : float
The wavelength of light used in the experiment
detector_geometry : dict
A dictionary containing the various detector geometry
parameters
mask : array
A mask for the detector, defined as 1 for live pixels, 0
for dead
background : array
An initial guess for the not-previously-subtracted
detector background
"""
super(PolarizedPtycho2DDataset,self).__init__(translations, patterns,
*args, **kwargs)
self.polarizer = t.tensor(polarizer_angles, dtype=t.float32)
self.analyzer = t.tensor(analyzer_angles, dtype=t.float32)
def _load(self, index):
""" Internal function to load data
This function is used internally by the global __getitem__ function
defined in the base class, which handles moving data around when
the dataset is (for example) storing the data on the CPU but
getting data as GPU tensors.
It loads data in the format (inputs, output)
The inputs for a 2D ptychogaphy data set are:
1) The indices of the patterns to use
2) The recorded probe positions associated with those points
3) The angles of the polarizers if polarized=True
Parameters
----------
index (polarized=False): int or slice
The index or indices of the scan points to use
index (polarized=True):
tuple ((phi1, phi2), ind)
ind - index or indices of the scan points to use (in a (phi1, phi2) polarization state), int or slice
(phi1, phi2) - angles of the 1st and 2nd polarizers, ints
Returns
-------
inputs : tuple
A tuple of the inputs to the related forward models
if polarized: inputs = ((phi1, phi2), ind, transl[(phi1, phi2)][ind])
outputs : tuple
The output pattern or stack of output patterns
"""
return ((index, self.translations[index],
self.polarizer[index], self.analyzer[index]),
self.patterns[index])
def to(self, *args, **kwargs):
"""Sends the relevant data to the given device and dtype
This function sends the stored translations, patterns,
mask and background to the specified device and dtype
Accepts the same parameters as torch.Tensor.to
"""
super(PolarizedPtycho2DDataset, self).to(*args, **kwargs)
self.polarizer = self.polarizer.to(*args, **kwargs)
self.analyzer = self.analyzer.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):
"""Generates a new CDataset from a .cxi file directly
This generates a new PolarizedPtycho2DDataset from a .cxi file storing
a 2D ptychography scan.
Parameters
----------
file : str, pathlib.Path, or h5py.File
The .cxi file to load from
Returns
-------
dataset : PolarizedPtycho2DDataset
The constructed dataset object
"""
# If a bare string is passed
if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
with h5py.File(cxi_file, 'r') as f:
return cls.from_cxi(f, polarized)
# Generate a base dataset
dataset = Ptycho2DDataset.from_cxi(cxi_file)
# Mutate the class to this subclass (PolarizedPtycho2DDataset)
dataset.__class__ = cls
# Now, we save out the polarizer and analyzer states
polarizer = cdtdata.get_shot_to_shot_info(cxi_file, 'polarizer_angle')
analyzer = cdtdata.get_shot_to_shot_info(cxi_file, 'analyzer_angle')
dataset.analyzer = t.tensor(analyzer, dtype=t.float32)
dataset.polarizer = t.tensor(polarizer, dtype=t.float32)
return dataset
def to_cxi(self, cxi_file, polarized=False):
"""Saves out a PolarizedPtycho2DDataset as a .cxi file
This function saves all the compatible information in a
PolarizedPtycho2DDataset object into a .cxi file. This saved .cxi file
should be compatible with any standard .cxi file based
reconstruction tool, such as SHARP.
Parameters
----------
cxi_file : str, pathlib.Path, or h5py.File
The .cxi file to write to
"""
# If a bare string is passed
if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
with cdtdata.create_cxi(cxi_file) as f:
return self.to_cxi(f, polarized)
# This saves the translations, patterns, etc.
super(PolarizedPtycho2DDataset, self).to_cxi(cxi_file)
# Now, we save out the polarizer and analyzer states
cdtdata.add_shot_to_shot_info(cxi_file, self.polarizer,
'polarizer_angle')
cdtdata.add_shot_to_shot_info(cxi_file, self.analyzer,
'analyzer_angle')
def inspect(self, logarithmic=True, units='um'):
"""Launches an interactive plot for perusing the data
This launches an interactive plotting tool in matplotlib that
shows the spatial map constructed from the integrated intensity
at each position on the left, next to a panel on the right that
can display a base-10 log plot of the detector readout at each
position.
"""
def get_images(idx):
inputs, output = self[idx]
meas_data = output.detach().cpu().numpy()
if hasattr(self, 'mask') and self.mask is not None:
mask = self.mask.detach().cpu().numpy()
else:
mask = 1
if logarithmic:
return np.log(meas_data) / np.log(10) * mask
else:
return meas_data * mask
translations = self.translations.detach().cpu().numpy()
nanomap_values = (self.mask.to(t.float32) * self.patterns).sum(dim=(1,2)).detach().cpu().numpy()
if logarithmic:
cbar_title='Log Base 10 of Diffraction Intensity'
else:
cbar_title='Diffraction Intensity'
plotting.plot_nanomap_with_images(self.translations.detach().cpu(), get_images, values=nanomap_values, nanomap_units=units, image_title='Diffraction Pattern', image_colorbar_title=cbar_title)
CDTools/datasets/ptycho_2d_dataset.py
+3 -6
View File
@@ -3,10 +3,7 @@ import numpy as np
import torch as t
from copy import copy
import h5py
try:
import pathlib
except ImportError:
import pathlib2 as pathlib
import pathlib
from CDTools.datasets import CDataset
from CDTools.tools import data as cdtdata
@@ -64,8 +61,8 @@ class Ptycho2DDataset(CDataset):
super(Ptycho2DDataset,self).__init__(*args, **kwargs)
self.axes = copy(axes)
self.translations = t.Tensor(translations).clone()
self.patterns = t.Tensor(patterns).clone()
self.translations = t.tensor(translations, dtype=t.float32)
self.patterns = t.tensor(patterns, dtype=t.float32)
if self.mask is None:
self.mask = t.ones(self.patterns.shape[-2:]).to(dtype=t.bool)
self.mask.masked_fill_(t.isnan(t.sum(self.patterns,dim=(0,))),0)
CDTools/tools/analysis/analysis.py
+1 -1
View File
@@ -183,7 +183,7 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
else:
single_probe = False
normalization = t.sqrt(t.sum(t.abs(probe[0])**2) / (len(probe[0].view(-1))))
normalization = t.sqrt(t.mean(t.abs(probe[0])**2))
probe = probe / normalization
obj = obj * normalization
CDTools/tools/data/data.py
+113 -43
View File
@@ -21,6 +21,7 @@ __all__ = ['get_entry_info',
'get_mask',
'get_dark',
'get_data',
'get_shot_to_shot_info',
'get_ptycho_translations',
'create_cxi',
'add_entry_info',
@@ -30,8 +31,10 @@ __all__ = ['get_entry_info',
'add_mask',
'add_dark',
'add_data',
'add_shot_to_shot_info',
'add_ptycho_translations']
#
# Functions to inspect the basic attributes of a cxi file represented as an
# h5 file object
@@ -365,6 +368,48 @@ def get_data(cxi_file, cut_zeroes = True):
return data, axes
def get_shot_to_shot_info(cxi_file, field_name):
"""Gets a specified dataset of shot-to-shot information from the cxi file
The data is assumed to be in the form of an array, with one dimension
being the number of patterns being stored in the dataset. This is
helpful for storing additional readback data on the shot-to-shot
level that may be important but doens't have a clearly defined
place to be stored in the .cxi file specification. Such data includes
shot-to-shot probe intensity measurements, polarizer positions, etc.
It will look for this data in 3 places (in the following order):
1) entry_1/data_1/<field_name>
2) entry_1/sample_1/geometry_1/<field_name>
3) entry_1/instrument_1/detector_1/<field_name>
This function is also used internally to read out the translations
associated with a ptychography experiment
Parameters
----------
cxi_file : h5py.File
A file object to be read
field_name : str
The name of the field to be read from
Returns
-------
data : np.array
An array storing the translations defined in the cxi file
"""
if 'entry_1/data_1/' + field_name in cxi_file:
pull_from = 'entry_1/data_1/' + field_name
elif 'entry_1/sample_1/geometry_1/' + field_name in cxi_file:
pull_from = 'entry_1/sample_1/geometry_1/' + field_name
elif 'entry_1/instrument_1/detector_1/' in cxi_file:
pull_from = 'entry_1/instrument_1/detector_1/' + field_name
else:
raise KeyError('Data is not defined within cxi file')
return np.array(cxi_file[pull_from]).astype(np.float32)
def get_ptycho_translations(cxi_file):
"""Gets an array of x,y,z translations, if such an array has been defined in the file
@@ -382,27 +427,14 @@ def get_ptycho_translations(cxi_file):
-------
translations : np.array
An array storing the translations defined in the cxi file
axes : list(str)
A list of the axes defined in the axes attribute, if any
"""
if 'entry_1/data_1/translation' in cxi_file:
pull_from = 'entry_1/data_1/translation'
elif 'entry_1/sample_1/geometry_1/translation' in cxi_file:
pull_from = 'entry_1/sample_1/geometry_1/translation'
elif 'entry_1/instrument_1/detector_1/translation' in cxi_file:
pull_from = 'entry_1/instrument_1/detector_1/translation'
else:
raise KeyError('Translations are not defined within cxi file')
translations = -np.array(cxi_file[pull_from]).astype(np.float32)
return translations
translations = get_shot_to_shot_info(cxi_file, 'translation')
return -translations
#
# It might be useful to make some helper functions to help write cxi files
# Now we move on to the helper functions for writing CXI files
#
@@ -650,6 +682,70 @@ def add_data(cxi_file, data, axes=None):
axes_str = str(axes)
det1['data'].attrs['axes'] = np.string_(axes_str)
def add_shot_to_shot_info(cxi_file, data, field_name):
"""Adds a specified dataset of shot-to-shot information to the cxi file
The data is assumed to be in the form of an array, with one dimension
being the number of patterns being stored in the dataset. This is
helpful for storing additional readback data on the shot-to-shot
level that may be important but doens't have a clearly defined
place to be stored in the .cxi file specification. Such data includes
shot-to-shot probe intensity measurements, polarizer positions, etc.
This function is also used internally to store the translations
associated with a ptychography experiment
It will store this data in 3 places:
1) The entry_1/sample_1/geometry_1/<field_name> path
2) A softlink at entry_1/data_1/<field_name>
3) A softlink at entry_1/instrument_1/detector_1/<field_name>
The geometry and detector paths may not always be relevant, but this
ensures that the data is always available in any of the places that
an eventual reader may go to look for, e.g., the translations.
Parameters
----------
cxi_file : h5py.File
The file to add the translations to
data : array
The data to be saved
field_name : str
The field name to save the data under
"""
if 'entry_1/sample_1' not in cxi_file:
cxi_file['entry_1'].create_group('sample_1')
s1 = cxi_file['entry_1/sample_1']
if 'geometry_1' not in s1:
s1.create_group('geometry_1')
g1 = s1['geometry_1']
if 'entry_1/data_1' not in cxi_file:
cxi_file['entry_1'].create_group('data_1')
data1 = cxi_file['entry_1/data_1']
if 'entry_1/instrument_1' not in cxi_file:
cxi_file['entry_1'].create_group('instrument_1')
i1 = cxi_file['entry_1/instrument_1']
if 'detector_1' not in i1:
i1.create_group('detector_1')
det1 = i1['detector_1']
if isinstance(data, t.Tensor):
data = data.detach().cpu().numpy()
g1.create_dataset(field_name, data=data)
data1[field_name] = h5py.SoftLink('/entry_1/sample_1/geometry_1/'
+ field_name)
det1[field_name] = h5py.SoftLink('/entry_1/sample_1/geometry_1/'
+ field_name)
def add_ptycho_translations(cxi_file, translations):
"""Adds the specified translations to the cxi file
@@ -671,34 +767,8 @@ def add_ptycho_translations(cxi_file, translations):
translations : array
The translations to be saved
"""
if 'entry_1/sample_1' not in cxi_file:
cxi_file['entry_1'].create_group('sample_1')
s1 = cxi_file['entry_1/sample_1']
if 'geometry_1' not in s1:
s1.create_group('geometry_1')
g1 = s1['geometry_1']
if 'entry_1/data_1' not in cxi_file:
cxi_file['entry_1'].create_group('data_1')
data1 = cxi_file['entry_1/data_1']
if 'entry_1/instrument_1' not in cxi_file:
cxi_file['entry_1'].create_group('instrument_1')
i1 = cxi_file['entry_1/instrument_1']
if 'detector_1' not in i1:
i1.create_group('detector_1')
det1 = i1['detector_1']
if isinstance(translations, t.Tensor):
translations = translations.detach().cpu().numpy()
# accounting for the different definition between cxi files and
# CDTools
translations = -translations
g1.create_dataset('translation', data=translations)
data1['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation')
det1['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation')
add_shot_to_shot_info(cxi_file, translations, 'translation')
tests/conftest.py
+13
View File
@@ -292,6 +292,19 @@ def ptycho_cxi_3():
f.close()
@pytest.fixture(scope='module')
def polarized_ptycho_cxi(ptycho_cxi_1):
f, expected = ptycho_cxi_1
data1f = f['entry_1/data_1']
expected['analyzer_angle'] = np.random.rand(100).astype(np.float32)
data1f.create_dataset('analyzer_angle', data=expected['analyzer_angle'])
expected['polarizer_angle'] = np.random.rand(100).astype(np.float32)
data1f.create_dataset('polarizer_angle', data=expected['polarizer_angle'])
yield f, expected
# 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
tests/test_datasets.py
+180 -8
View File
@@ -118,11 +118,9 @@ def test_CDataset_to(ptycho_cxi_1):
assert dataset.background.device == t.device('cuda:0')
#
# And we then test the derived Ptychography class
#
#
def test_Ptycho2DDataset_init():
@@ -156,7 +154,6 @@ def test_Ptycho2DDataset_init():
assert t.allclose(dataset.translations, t.Tensor(translations))
def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
dataset = Ptycho2DDataset.from_cxi(cxi)
@@ -191,10 +188,8 @@ def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
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')
@@ -236,7 +231,7 @@ def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
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])
@@ -253,7 +248,6 @@ def test_Ptycho2DDataset_to(ptycho_cxi_1):
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)
@@ -282,5 +276,183 @@ def test_Ptycho2DDataset_get_as(ptycho_cxi_1):
t.tensor(expected['translations'][3,:]))
assert t.allclose(pattern.to(device='cpu'),
t.tensor(expected['data'][3,:,:]))
#
# And we now test the derived class for polarization-dependent ptycho
#
def test_PolarizedPtycho2DDataset_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)
analyzer = np.random.rand(20)
polarizer = np.random.rand(20)
dataset = PolarizedPtycho2DDataset(translations, polarizer,
analyzer, 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.Tensor(patterns))
assert t.allclose(dataset.translations, t.Tensor(translations))
assert t.allclose(dataset.analyzer, t.Tensor(analyzer))
assert t.allclose(dataset.polarizer, t.Tensor(polarizer))
def test_PolarizedPtycho2DDataset_from_cxi(polarized_ptycho_cxi):
cxi, expected = polarized_ptycho_cxi
dataset = PolarizedPtycho2DDataset.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.Tensor(expected['dark']),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['analyzer_angle']),dataset.analyzer)
assert t.allclose(t.tensor(expected['polarizer_angle']),dataset.polarizer)
def test_PolarizedPtycho2DDataset_to_cxi(polarized_ptycho_cxi, tmp_path):
cxi, expected = polarized_ptycho_cxi
print('loading dataset')
dataset = PolarizedPtycho2DDataset.from_cxi(cxi)
print('dataset mask is type', dataset.mask.dtype)
with cdtdata.create_cxi(tmp_path / 'test_PolarizedPtycho2DDataset_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_PolarizedPtycho2DDataset_to_cxi.cxi', 'r') as f:
read_dataset = PolarizedPtycho2DDataset.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)
assert t.allclose(dataset.analyzer, read_dataset.analyzer)
assert t.allclose(dataset.analyzer, read_dataset.analyzer)
def test_PolarizedPtycho2DDataset_to(polarized_ptycho_cxi):
dataset = PolarizedPtycho2DDataset.from_cxi(polarized_ptycho_cxi[0])
dataset.to(dtype=t.float64)
assert dataset.mask.dtype == t.bool
assert dataset.patterns.dtype == t.float64
assert dataset.translations.dtype == t.float64
assert dataset.analyzer.dtype == t.float64
assert dataset.polarizer.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')
assert dataset.analyzer.device == t.device('cuda:0')
assert dataset.polarizer.device == t.device('cuda:0')
def test_PolarizedPtycho2DDataset_ops(polarized_ptycho_cxi):
cxi, expected = polarized_ptycho_cxi
dataset = PolarizedPtycho2DDataset.from_cxi(cxi)
dataset.get_as('cpu')
assert len(dataset) == expected['data'].shape[0]
(idx, translation, polarizer, analyzer), 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(analyzer, t.tensor(expected['analyzer_angle'][3]))
assert t.allclose(polarizer, t.tensor(expected['polarizer_angle'][3]))
def test_PolarizedPtycho2DDataset_get_as(polarized_ptycho_cxi):
cxi, expected = polarized_ptycho_cxi
dataset = PolarizedPtycho2DDataset.from_cxi(cxi)
if t.cuda.is_available():
dataset.get_as('cuda:0')
assert len(dataset) == expected['data'].shape[0]
(idx, translation, polarizer, analyzer), 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,:,:]))
assert t.allclose(polarizer.to(device='cpu'),
t.tensor(expected['polarizer_angle'][3]))
assert t.allclose(analyzer.to(device='cpu'),
t.tensor(expected['analyzer_angle'][3]))
tests/tools/test_data.py
+25
View File
@@ -79,6 +79,13 @@ def test_get_data(test_ptycho_cxis):
assert axes == expected['axes']
def test_get_shot_to_shot_info(polarized_ptycho_cxi):
cxi, expected = polarized_ptycho_cxi
for key in ('analyzer_angle', 'polarizer_angle'):
assert np.allclose(data.get_shot_to_shot_info(cxi, key),
expected[key])
def test_get_ptycho_translations(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
assert np.allclose(data.get_ptycho_translations(cxi),
@@ -241,7 +248,25 @@ def test_add_data(tmp_path):
assert np.allclose(fake_data.numpy(),read_data)
def test_add_shot_to_shot_info(tmp_path):
analyzer = np.random.rand(100)
with data.create_cxi(tmp_path / 'test_add_shot_to_shot_info.cxi') as f:
data.add_shot_to_shot_info(f, analyzer, 'analyzer_angle')
with h5py.File(tmp_path / 'test_add_shot_to_shot_info.cxi') as f:
# Check this directly since we want to make sure it saved
# it in all the places it should have
read_analyzer_1 = np.array(f['entry_1/data_1/analyzer_angle'])
read_analyzer_2 = np.array(f['entry_1/instrument_1/detector_1/analyzer_angle'])
read_analyzer_3 = np.array(f['entry_1/sample_1/geometry_1/analyzer_angle'])
assert np.allclose(analyzer, read_analyzer_1)
assert np.allclose(analyzer, read_analyzer_2)
assert np.allclose(analyzer, read_analyzer_3)
def test_add_ptycho_translations(tmp_path):
translations = np.random.rand(3,100)
tests/tools/test_losses.py
+1
View File
@@ -10,6 +10,7 @@ import torch as t
# that any optimizations in the future don't change the results
def test_amplitude_mse():
# Make some fake data
data = np.random.rand(10,100,100)
# And add some noise to it