Write the model and dataset tutorials, update the dependencies, and various small bug fixes

examples/basic_ptycho_dataset.py

@@ -0,0 +1,69 @@
+from __future__ import division, print_function, absolute_import
+import numpy as np
+import torch as t
+from matplotlib import pyplot as plt
+from CDTools.datasets import CDataset
+from CDTools.tools import data as cdtdata
+
+
+__all__ = ['BasicPtychoDataset']
+
+
+class BasicPtychoDataset(CDataset):
+    """The standard dataset for a 2D ptychography scan"""
+
+    def __init__(self, translations, patterns, *args, **kwargs):
+        """Initialize the dataset from python objects"""
+
+        super(BasicPtychoDataset,self).__init__(*args, **kwargs)
+        self.translations = t.Tensor(translations).clone()
+        self.patterns = t.Tensor(patterns).clone()
+
+    def __len__(self):
+        return self.patterns.shape[0]
+
+    def _load(self, index):
+        return (index, self.translations[index]), self.patterns[index]
+
+
+    def to(self, *args, **kwargs):
+        """Sends the relevant data to the given device and dtype"""
+        super(BasicPtychoDataset,self).to(*args,**kwargs)
+        self.translations = self.translations.to(*args, **kwargs)
+        self.patterns = self.patterns.to(*args, **kwargs)
+
+
+    @classmethod
+    def from_cxi(cls, cxi_file):
+        """Generates a new CDataset from a .cxi file directly"""
+
+        # 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)
+
+        # And now re-add it
+        dataset.translations = t.Tensor(translations).clone()
+        dataset.patterns = t.Tensor(patterns).clone()
+
+        return dataset
+
+    
+    def to_cxi(self, cxi_file):
+        """Saves out a BasicPtychoDataset as a .cxi file"""
+
+        super(BasicPtychoDataset,self).to_cxi(cxi_file)
+        cdtdata.add_data(cxi_file, self.patterns, axes=self.axes)
+        cdtdata.add_ptycho_translations(cxi_file, self.translations)
+
+
+    def inspect(self):
+        """Plots a random diffraction pattern"""
+
+        index = np.random.randint(len(self))
+        plt.figure()
+        plt.imshow(self.patterns[index,:,:].cpu().numpy())

examples/simple_ptycho.py

@@ -4,7 +4,7 @@ import CDTools
 from matplotlib import pyplot as plt
 
 # First, we load an example dataset from a .cxi file
-filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
+filename = 'example_data/lab_ptycho_data.cxi'
 dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
 
 # Next, we create a ptychography model from the dataset

examples/simple_ptycho_model.py

@@ -0,0 +1,107 @@
+from __future__ import division, print_function, absolute_import
+import numpy as np
+import torch as t
+from CDTools.models import CDIModel
+from CDTools import tools
+
+__all__ = ['SimplePtycho']
+
+class SimplePtycho(CDIModel):
+    """A simple ptychography model for exploring ideas and extensions"""
+    
+    def __init__(self, probe_basis, probe_guess, obj_guess,
+                 min_translation = [0,0]):
+        # We have to initialize the Module
+        super(SimplePtycho,self).__init__()
+
+        # We first save the relevant information
+        self.min_translation = t.Tensor(min_translation)
+        self.probe_basis = t.Tensor(probe_basis)
+
+        # We rescale the probe so it learns at the same rate as the object
+        self.probe_norm = t.max(tools.cmath.cabs(probe_guess.to(t.float32)))
+        self.probe = t.nn.Parameter(probe_guess.to(t.float32)
+                                    / self.probe_norm)
+        self.obj = t.nn.Parameter(obj_guess.to(t.float32))
+
+
+
+    @classmethod
+    def from_dataset(cls, dataset):
+        # First, load the information from the dataset
+        wavelength = dataset.wavelength
+        det_basis = dataset.detector_geometry['basis']
+        det_shape = dataset[0][1].shape
+        distance = dataset.detector_geometry['distance']
+        (indices, translations), patterns = dataset[:]
+        
+        # Then, generate the probe geometry
+        ewg = tools.initializers.exit_wave_geometry
+        probe_basis, probe_shape, det_slice =  ewg(det_basis,
+                                                   det_shape,
+                                                   wavelength,
+                                                   distance,
+                                                   opt_for_fft=False)
+
+        # Next generate the object geometry
+        pix_translations = tools.interactions.translations_to_pixel(
+            probe_basis, translations)
+        obj_size, min_translation = tools.initializers.calc_object_setup(
+            probe_shape, pix_translations)
+
+        # Finally, initialize the probe and  object using this information
+        probe = tools.initializers.SHARP_style_probe(dataset,
+                                                     probe_shape,
+                                                     det_slice)
+
+        obj = t.ones(obj_size+(2,))
+
+        return cls(probe_basis, probe, obj, min_translation=min_translation)
+    
+
+    def interaction(self, index, translations):
+        pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
+                                                             translations)
+        pix_trans -= self.min_translation
+        return tools.interactions.ptycho_2D_round(self.probe_norm * self.probe,
+                                                  self.obj,
+                                                  pix_trans)
+
+    def forward_propagator(self, wavefields):
+        return tools.propagators.far_field(wavefields)
+
+    def measurement(self, wavefields):
+        return tools.measurements.intensity(wavefields)
+
+    def loss(self, sim_data, real_data):
+        return tools.losses.amplitude_mse(real_data, sim_data)
+
+
+    def to(self, *args, **kwargs):
+        super(SimplePtycho, self).to(*args, **kwargs)
+        self.min_translation = self.min_translation.to(*args,**kwargs)
+        self.probe_basis = self.probe_basis.to(*args,**kwargs)
+        self.probe_norm = self.probe_norm.to(*args,**kwargs)
+
+        
+   
+    plot_list = [
+        ('Probe Amplitude',
+         lambda self, fig: tools.plotting.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)),
+        ('Probe Phase',
+         lambda self, fig: tools.plotting.plot_phase(self.probe, fig=fig, basis=self.probe_basis)),
+        ('Object Amplitude',
+         lambda self, fig: tools.plotting.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
+        ('Object Phase',
+         lambda self, fig: tools.plotting.plot_phase(self.obj, fig=fig, basis=self.probe_basis))
+    ]
+    
+    def save_results(self):
+        probe = tools.cmath.torch_to_complex(self.probe.detach().cpu())
+        probe = probe * self.probe_norm.detach().cpu().numpy()
+        obj = tools.cmath.torch_to_complex(self.obj.detach().cpu())
+        return {'probe':probe,'obj':obj}
+
+    
+if __name__ == '__main__':
+