Add sim-to-dataset capability

CDTools/models/fancy_ptycho.py

@@ -2,10 +2,12 @@ from __future__ import division, print_function, absolute_import
 
 import torch as t
 from CDTools.models import CDIModel
+from CDTools.datasets import Ptycho_2D_Dataset
 from CDTools import tools
 from CDTools.tools import cmath
 from CDTools.tools import plotting as p
 from matplotlib import pyplot as plt
+from datetime import datetime
 import numpy as np
 from copy import copy
 
@@ -264,7 +266,32 @@ class FancyPtycho(CDIModel):
         
 
     def sim_to_dataset(self, args_list):
-        pass
+        # In the future, potentially add more control
+        # over what metadata is saved (names, etc.)
+        
+        # First, I need to gather all the relevant data
+        # that needs to be added to the dataset
+        entry_info = {'program_name': 'CDTools',
+                      'instrument_n': 'Simulated Data',
+                      'start_time': datetime.now()}
+
+        sample_info = {'description': 'A simulated sample'}
+        
+        detector_geometry = self.detector_geometry
+        mask = self.mask
+        wavelength = self.wavelength
+        indices, translations = args_list
+        
+        # Then we simulate the results
+        data = self.forward(indices, translations)
+
+        # And finally, we make the dataset
+        return Ptycho_2D_Dataset(translations, data,
+                                 entry_info = entry_info,
+                                 sample_info = sample_info,
+                                 wavelength=wavelength,
+                                 detector_geometry=detector_geometry,
+                                 mask=mask)
 
     
     def corrected_translations(self,dataset):

CDTools/models/simple_ptycho.py

@@ -2,11 +2,13 @@ from __future__ import division, print_function, absolute_import
 
 import torch as t
 from CDTools.models import CDIModel
+from CDTools.datasets import Ptycho_2D_Dataset
 from CDTools import tools
 from CDTools.tools import plotting as p
 from copy import copy
 from torch.utils import data as torchdata
 from matplotlib import pyplot as plt
+from datetime import datetime
 import numpy as np
 
 class SimplePtycho(CDIModel):
@@ -135,7 +137,33 @@ class SimplePtycho(CDIModel):
 
 
     def sim_to_dataset(self, args_list):
-        raise NotImplementedError()
+        # In the future, potentially add more control
+        # over what metadata is saved (names, etc.)
+        
+        # First, I need to gather all the relevant data
+        # that needs to be added to the dataset
+        entry_info = {'program_name': 'CDTools',
+                      'instrument_n': 'Simulated Data',
+                      'start_time': datetime.now()}
+
+        sample_info = {'description': 'A simulated sample'}
+        
+        detector_geometry = self.detector_geometry
+        mask = self.mask
+        wavelength = self.wavelength
+        indices, translations = args_list
+        
+        # Then we simulate the results
+        data = self.forward(indices, translations)
+
+        # And finally, we make the dataset
+        return Ptycho_2D_Dataset(translations, data,
+                                 entry_info = entry_info,
+                                 sample_info = sample_info,
+                                 wavelength=wavelength,
+                                 detector_geometry=detector_geometry,
+                                 mask=mask)
+
 
 
     plot_list = [

examples/gold_ball_ptycho.py

@@ -11,9 +11,11 @@ filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
 with h5py.File(filename,'r') as f:
     dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(f)
 
+
 model = CDTools.models.FancyPtycho.from_dataset(dataset,n_modes=3,randomize_ang=0.1*np.pi)
 
 
+
 # default is CPU with 32-bit floats
 model.to(device='cuda')
 dataset.get_as(device='cuda')
@@ -24,7 +26,8 @@ for i, loss in enumerate(model.Adam_optimize(3, dataset, batch_size=100)):
     print(i,loss)
 
 
-#with open('test_results.pickle', 'wb') as f:
+
+#with open('example_reconstructions/Au_balls.pickle', 'wb') as f:
 #    pickle.dump(model.save_results(dataset),f)
     
 model.inspect(dataset)