From e7d7323e5057113ac42117dc356db2d435167f5d Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Mon, 26 Jul 2021 14:46:02 -0400 Subject: [PATCH] Added a polarization dependent ptychography dataset class so we can save and load data! --- CDTools/datasets/__init__.py | 3 +- CDTools/datasets/base.py | 13 +- .../datasets/polarized_ptycho_2d_dataset.py | 219 ++++++++++++++++++ CDTools/datasets/ptycho_2d_dataset.py | 9 +- CDTools/tools/analysis/analysis.py | 2 +- CDTools/tools/data/data.py | 156 +++++++++---- tests/conftest.py | 13 ++ tests/test_datasets.py | 188 ++++++++++++++- tests/tools/test_data.py | 25 ++ tests/tools/test_losses.py | 1 + 10 files changed, 560 insertions(+), 69 deletions(-) create mode 100644 CDTools/datasets/polarized_ptycho_2d_dataset.py diff --git a/CDTools/datasets/__init__.py b/CDTools/datasets/__init__.py index 4a2494c..40f267f 100644 --- a/CDTools/datasets/__init__.py +++ b/CDTools/datasets/__init__.py @@ -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 diff --git a/CDTools/datasets/base.py b/CDTools/datasets/base.py index f4d3067..ef253ba 100644 --- a/CDTools/datasets/base.py +++ b/CDTools/datasets/base.py @@ -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 diff --git a/CDTools/datasets/polarized_ptycho_2d_dataset.py b/CDTools/datasets/polarized_ptycho_2d_dataset.py new file mode 100644 index 0000000..95c426c --- /dev/null +++ b/CDTools/datasets/polarized_ptycho_2d_dataset.py @@ -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) + diff --git a/CDTools/datasets/ptycho_2d_dataset.py b/CDTools/datasets/ptycho_2d_dataset.py index 0b84d54..c5bc416 100644 --- a/CDTools/datasets/ptycho_2d_dataset.py +++ b/CDTools/datasets/ptycho_2d_dataset.py @@ -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) diff --git a/CDTools/tools/analysis/analysis.py b/CDTools/tools/analysis/analysis.py index 2bc67da..bef729b 100644 --- a/CDTools/tools/analysis/analysis.py +++ b/CDTools/tools/analysis/analysis.py @@ -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 diff --git a/CDTools/tools/data/data.py b/CDTools/tools/data/data.py index b3f05c4..2462cbd 100644 --- a/CDTools/tools/data/data.py +++ b/CDTools/tools/data/data.py @@ -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/ + 2) entry_1/sample_1/geometry_1/ + 3) entry_1/instrument_1/detector_1/ + + 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/ path + 2) A softlink at entry_1/data_1/ + 3) A softlink at entry_1/instrument_1/detector_1/ + + 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') diff --git a/tests/conftest.py b/tests/conftest.py index bbf4ecb..ec51da2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 30404f3..7a825f0 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -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])) diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py index a39b15d..cf8b5dd 100644 --- a/tests/tools/test_data.py +++ b/tests/tools/test_data.py @@ -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) diff --git a/tests/tools/test_losses.py b/tests/tools/test_losses.py index 0c7df38..e8934ca 100644 --- a/tests/tools/test_losses.py +++ b/tests/tools/test_losses.py @@ -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