Finish writing the Dataset classes with loading, and write up all the tests

CDTools/datasets.py

@@ -4,7 +4,7 @@ import torch as t
 from copy import copy
 
 # The naming overlap here is definitely going to get confusing.
-from CDTools import tools
+from CDTools.tools import data as cdtdata
 from torch.utils import data as torchdata
 
 __all__ = ['CDataset', 'Ptycho_2D_Dataset']
@@ -103,26 +103,33 @@ class CDataset(torchdata.Dataset):
         self.wavelength = wavelength
         self.detector_geometry = copy(detector_geometry)
         if mask is not None:
-            self.mask = t.Tensor(mask)
+            self.mask = t.tensor(mask)
         else:
             self.mask = None
 
             
     def to(self,*args,**kwargs):
-        if mask is not None:
-            self.mask.to(*args,**kwargs)          
+        # The mask should always stay a uint8, but it should switch devices
+        mask_kwargs = copy(kwargs)
+        try:
+            mask_kwargs.pop('dtype')
+        except KeyError as r:
+            pass
+        
+        if self.mask is not None:
+            self.mask = self.mask.to(*args,**mask_kwargs)          
 
     
     @classmethod
     def from_cxi(cls, cxi_file):
-        entry_info = tools.get_entry_info(cxi_file)
-        sample_info = tools.get_sample_info(cxi_file)
-        wavelength = tools.get_wavelength(cxi_file)
-        distance, basis, corner = tools.get_detector_geometry(cxi_file)
+        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 = tools.get_mask(cxi_file)
+        mask = cdtdata.get_mask(cxi_file)
         return cls(entry_info = entry_info,
                    sample_info = sample_info,
                    wavelength=wavelength,
@@ -132,42 +139,78 @@ class CDataset(torchdata.Dataset):
     
     def to_cxi(self, cxi_file):
         if self.entry_info is not None:
-            tools.add_entry_info(cxi_file, self.entry_info)
+            cdtdata.add_entry_info(cxi_file, self.entry_info)
         if self.sample_info is not None:
-            tools.add_sample_info(cxi_file, self.sample_info)
+            cdtdata.add_sample_info(cxi_file, self.sample_info)
         if self.wavelength is not None:
-            tools.add_source(cxi_file, self.wavelength)
+            cdtdata.add_source(cxi_file, self.wavelength)
         if self.detector_geometry is not None:
             if 'corner' in self.detector_geometry:
                 corner = self.detector_geometry['corner']
             else:
                 corner = None
-            tools.add_detector(cxi_file,
-                               self.detector_info['wavelength'],
-                               self.detector_info['basis'],
+            cdtdata.add_detector(cxi_file,
+                               self.detector_geometry['distance'],
+                               self.detector_geometry['basis'],
                                corner = corner)
         if self.mask is not None:
-            tools.add_mask(cxi_file, mask)
+            cdtdata.add_mask(cxi_file, self.mask)
         
 
 
+#
+# This is the standard dataset for a 2D ptychography experiment,
+# which saves and loads files compatible with most reconstruction
+# programs (only tested against SHARP)
+#
+
 class Ptycho_2D_Dataset(CDataset):
 
-    def __init__(self,translations, patterns, **kwargs):
+    def __init__(self, translations, patterns, axes=None, *args, **kwargs):
 
-        super(CDataset,self).__init__(kwargs)
+        super(Ptycho_2D_Dataset,self).__init__(*args, **kwargs)
+        self.axes = copy(axes)
         self.translations = t.tensor(translations)
         self.patterns = t.tensor(patterns)
         
 
-    def to(self, *args, **kwargs):
-        super(CDataset,self).to(*args,**kwargs)
-        self.translations.to(*args, **kwargs)
-        self.patterns.to(*args, **kwargs)
-        
-        
     def __len__(self):
         return self.patterns.shape[0]
 
-    def __get__(self,index):
+    def __getitem__(self,index):
         return index, self.translations[index], self.patterns[index]
+
+
+    def to(self, *args, **kwargs):
+        super(Ptycho_2D_Dataset,self).to(*args,**kwargs)
+        self.translations = self.translations.to(*args, **kwargs)
+        self.patterns = self.patterns.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):
+        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)
+        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)
+    
+
+    def to_cxi(self, cxi_file):
+        super(Ptycho_2D_Dataset,self).to_cxi(cxi_file)
+        cdtdata.add_data(cxi_file, self.patterns, axes=self.axes)
+        cdtdata.add_ptycho_translations(cxi_file, self.translations)

CDTools/tools/data.py

@@ -143,6 +143,10 @@ def get_sample_info(cxi_file):
     # TODO: I should add the sample geometry as a valid metadata that can
     # be copied over
     
+    # Check if the metadata is empty
+    if metadata == {}:
+        metadata = None
+    
     return metadata
 
 
@@ -315,12 +319,10 @@ def get_data(cxi_file):
 def get_ptycho_translations(cxi_file):
     """Gets an array of x,y,z translations, if such an array has been defined in the file
 
-    It applies two operations to the translations. First, it negates them,
-    because the CXI file format is designed to specify translations of the
-    samples and the CDTools code specifies translations of the optics.
-    Second, it transposes the array so that the first axis is translation
-    ID and the second axis is the (x,y,z) components of the translation
-
+    It negates the translations, because the CXI file format is designed
+    to specify translations of the samples and the CDTools code specifies
+    translations of the optics.
+    
     Args:
         cxi_file (h5py.File) : a file object to be read
     
@@ -338,7 +340,7 @@ def get_ptycho_translations(cxi_file):
     else:
         raise KeyError('Translations are not defined within cxi file')
 
-    translations = -np.array(cxi_file[pull_from]).astype(np.float32).transpose()
+    translations = -np.array(cxi_file[pull_from]).astype(np.float32)
     return translations
     
 
@@ -537,7 +539,7 @@ def add_ptycho_translations(cxi_file, translations):
     
     It will add the translations to the file, negating them to conform to
     the standard in cxi files that the translations refer to the object's
-    translation, and also transposing them to match the cxi file specification.
+    translation.
     
     It will generally store them in 3 places:
 
@@ -575,7 +577,7 @@ def add_ptycho_translations(cxi_file, translations):
 
     # accounting for the different definition between cxi files and
     # CDTools
-    translations = -translations.transpose()
+    translations = -translations
     
     g1.create_dataset('translation', data=translations)
     data1['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation')