Merging a whole buncha work back into master

CDTools/datasets/polarized_ptycho_2d_dataset.py

@@ -115,7 +115,7 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset):
     # 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):
+    def from_cxi(cls, cxi_file, cut_zeros=True):
         """Generates a new PolarizedPtycho2DDataset from a .cxi file directly
 
         This generates a new PolarizedPtycho2DDataset from a .cxi file storing
@@ -125,19 +125,22 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset):
         ----------
         file : str, pathlib.Path, or h5py.File
             The .cxi file to load from
+        cut_zeros : bool
+            Default True, whether to set all negative data to zero
 
         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)
+                return cls.from_cxi(f, cut_zeros=cut_zeros)
 
         # Generate a base dataset
-        dataset = Ptycho2DDataset.from_cxi(cxi_file)
+        dataset = Ptycho2DDataset.from_cxi(cxi_file, cut_zeros=cut_zeros)
 
         # Mutate the class to this subclass (PolarizedPtycho2DDataset)
         dataset.__class__ = cls

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.translations = t.tensor(translations)
         
-        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')
+        self.patterns = t.as_tensor(patterns)
 
         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]
 
@@ -118,7 +123,7 @@ class Ptycho2DDataset(CDataset):
     # 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):
+    def from_cxi(cls, cxi_file, cut_zeros=True):
         """Generates a new Ptycho2DDataset from a .cxi file directly
 
         This generates a new Ptycho2DDataset from a .cxi file storing
@@ -128,6 +133,8 @@ class Ptycho2DDataset(CDataset):
         ----------
         file : str, pathlib.Path, or h5py.File
             The .cxi file to load from
+        cut_zeros : bool
+            Default True, whether to set all negative data to zero
 
         Returns
         -------
@@ -145,11 +152,10 @@ class Ptycho2DDataset(CDataset):
         dataset.__class__ = cls
 
         # Load the data that is only relevant for this class
-        patterns, axes = cdtdata.get_data(cxi_file)
+        patterns, axes = cdtdata.get_data(cxi_file, cut_zeros=cut_zeros)
         translations = cdtdata.get_ptycho_translations(cxi_file)
         # And now re-do the stuff from __init__
         dataset.translations = t.tensor(translations, dtype=t.float32)
-
         dataset.patterns = t.as_tensor(patterns)
         if dataset.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')
@@ -158,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(cxi_file, 'intensities')
+            dataset.intensities = t.as_tensor(intensities, dtype=t.float32)
+        except KeyError:
+            dataset.intensities = None
+            
         return dataset
 
 
@@ -187,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', log_offset=1):
         """Launches an interactive plot for perusing the data
@@ -234,5 +249,5 @@ class Ptycho2DDataset(CDataset):
         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)
+        return 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)