From c1663fa964d2af86496151ea0bade25301a24ecf Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Tue, 22 Feb 2022 20:36:50 -0500 Subject: [PATCH] Add some tools for datasets to store input intensity measurements --- CDTools/datasets/ptycho_2d_dataset.py | 24 +++++++++++++++++++----- CDTools/models/polarized_fancy_ptycho.py | 6 +++--- CDTools/tools/data/data.py | 2 +- 3 files changed, 23 insertions(+), 9 deletions(-) diff --git a/CDTools/datasets/ptycho_2d_dataset.py b/CDTools/datasets/ptycho_2d_dataset.py index 9cd0d71..e81e990 100644 --- a/CDTools/datasets/ptycho_2d_dataset.py +++ b/CDTools/datasets/ptycho_2d_dataset.py @@ -19,7 +19,8 @@ class Ptycho2DDataset(CDataset): It should save and load files compatible with most reconstruction programs, although it is only tested against SHARP. """ - def __init__(self, translations, patterns, axes=None, *args, **kwargs): + def __init__(self, translations, patterns, intensities=None, + axes=None, *args, **kwargs): """The __init__ function allows construction from python objects. The detector_geometry dictionary is defined to have the @@ -52,23 +53,27 @@ class Ptycho2DDataset(CDataset): background : array An initial guess for the not-previously-subtracted detector background + intensities : array + A list of measured shot-to-shot intensities """ - + super(Ptycho2DDataset,self).__init__(*args, **kwargs) self.axes = copy(axes) self.translations = t.tensor(translations, dtype=t.float32) self.patterns = t.as_tensor(patterns, dtype=t.float32) - if self.patterns.dtype == t.float64: - raise NotImplementedError('64-bit floats are not supported and precision will not be retained in reconstructions! Please explicitly convert your data to 32-bit or submit a pull request') 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) self.patterns.masked_fill_(t.isnan(self.patterns),0) - + if intensities is not None: + self.intensities = t.as_tensor(intensities, dtype=t.float32) + else: + self.intensities = None + def __len__(self): return self.patterns.shape[0] @@ -159,6 +164,12 @@ class Ptycho2DDataset(CDataset): if dataset.mask is None: dataset.mask = t.ones(dataset.patterns.shape[-2:]).to(dtype=t.bool) + try: + intensities = cdtdata.get_shot_to_shot_info('intensities') + dataset.intensities = t.as_tensor(intensities, dtype=t.float32) + except KeyError: + dataset.intensities = None + return dataset @@ -188,6 +199,9 @@ class Ptycho2DDataset(CDataset): cdtdata.add_data(cxi_file, self.patterns) cdtdata.add_ptycho_translations(cxi_file, self.translations) + if hasattr(self, 'intensities') and self.intensities is not None: + cdtdata.add_shot_to_shot_info(cxi_file, self.intensities, 'intensities') + def inspect(self, logarithmic=True, units='um'): """Launches an interactive plot for perusing the data diff --git a/CDTools/models/polarized_fancy_ptycho.py b/CDTools/models/polarized_fancy_ptycho.py index 6d307a2..bbb2bb2 100644 --- a/CDTools/models/polarized_fancy_ptycho.py +++ b/CDTools/models/polarized_fancy_ptycho.py @@ -294,13 +294,13 @@ class PolarizedFancyPtycho(FancyPtycho): else: raise NotImplementedError('Unstable Modes not Implemented for polarized light') - pol_probes = polarization.apply_linear_polarizer(prs, polarizer) - exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc( pol_probes, self.obj, pix_trans, shift_probe=True, multiple_modes=True, polarized=True) - + # We're losing some efficiency here, because we only need to keep + # around the scalar wavefield after analyzing the waves. + # But I think it's not a huge issue - Abe analyzed_exit_waves = polarization.apply_linear_polarizer(exit_waves, analyzer) return analyzed_exit_waves diff --git a/CDTools/tools/data/data.py b/CDTools/tools/data/data.py index bb97a0f..5fdd982 100644 --- a/CDTools/tools/data/data.py +++ b/CDTools/tools/data/data.py @@ -93,7 +93,7 @@ def get_sample_info(cxi_file): """ if 'entry_1/sample_1' not in cxi_file: return None - + s1 = cxi_file['entry_1/sample_1'] metadata_attrs = ['name','description','unit_cell_group']