diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index 5030ed7..a205412 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/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): diff --git a/CDTools/models/simple_ptycho.py b/CDTools/models/simple_ptycho.py index 4ba4f68..90d0ace 100644 --- a/CDTools/models/simple_ptycho.py +++ b/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 = [ diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index 3b0ae73..a590ddd 100644 --- a/examples/gold_ball_ptycho.py +++ b/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)