From e71c266fe4c5f574f4ba221f79a55c09952abbcc Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Sat, 5 Oct 2019 09:30:26 -0400 Subject: [PATCH] Update the ptycho dataset with the new, more sensible, approach --- CDTools/datasets/ptycho_2d_dataset.py | 42 +++++++++++++++------------ 1 file changed, 23 insertions(+), 19 deletions(-) diff --git a/CDTools/datasets/ptycho_2d_dataset.py b/CDTools/datasets/ptycho_2d_dataset.py index c4d417f..5bf16ab 100644 --- a/CDTools/datasets/ptycho_2d_dataset.py +++ b/CDTools/datasets/ptycho_2d_dataset.py @@ -42,6 +42,8 @@ class Ptycho2DDataset(CDataset): 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 list of names for the axes of the probe translations entry_info : dict A dictionary containing the entry_info metadata sample_info : dict @@ -70,7 +72,6 @@ class Ptycho2DDataset(CDataset): self.patterns.masked_fill_(t.isnan(self.patterns),0) - def __len__(self): return self.patterns.shape[0] @@ -82,6 +83,8 @@ class Ptycho2DDataset(CDataset): 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 @@ -139,26 +142,24 @@ class Ptycho2DDataset(CDataset): with h5py.File(cxi_file,'r') as f: return cls.from_cxi(f) - 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) - - dark = cdtdata.get_dark(cxi_file) + # Generate a base dataset + dataset = CDataset.from_cxi(cxi_file) + # Mutate the class to this subclass (BasicPtychoDataset) + dataset.__class__ = cls + + # Load the data that is only relevant for this class 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, background=dark) + # And now re-do the stuff from __init__ + dataset.translations = t.Tensor(translations).clone() + dataset.patterns = t.Tensor(patterns).clone() + dataset.axes = axes + if dataset.mask is None: + dataset.mask = t.ones(dataset.patterns.shape[-2:]).to(dtype=t.bool) + + return dataset + def to_cxi(self, cxi_file): """Saves out a Ptycho2DDataset as a .cxi file @@ -180,7 +181,10 @@ class Ptycho2DDataset(CDataset): return self.to_cxi(f) super(Ptycho2DDataset,self).to_cxi(cxi_file) - cdtdata.add_data(cxi_file, self.patterns, axes=self.axes) + if hasattr(self, 'axes'): + cdtdata.add_data(cxi_file, self.patterns, axes=self.axes) + else: + cdtdata.add_data(cxi_file, self.patterns) cdtdata.add_ptycho_translations(cxi_file, self.translations)