diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index 43f6360..04f4369 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -83,6 +83,7 @@ class Bragg2DPtycho(CDIModel): obj_view_crop=0, panel_plot_mode=False, plot_level=1, + translations=None, ): # We need the detector geometry @@ -176,7 +177,7 @@ class Bragg2DPtycho(CDIModel): shape = [s//oversampling for s in self.probe[0]] background = 1e-6 * t.ones(shape, dtype=t.float32) - self.background = t.nn.Parameter(background) + self.background = t.nn.Parameter(t.as_tensor(background, dtype=dtype)) if weights is None: self.weights = None @@ -255,7 +256,11 @@ class Bragg2DPtycho(CDIModel): self.loss_normalizer = tools.losses.IntensityMSENormalizer() else: raise KeyError('Specified loss function not supported') - + + if translations is not None: + self.register_buffer('original_translations', + t.as_tensor(translations, dtype=dtype)) + @classmethod def from_dataset( @@ -473,6 +478,7 @@ class Bragg2DPtycho(CDIModel): units=units, panel_plot_mode=panel_plot_mode, plot_level=plot_level, + translations=translations, ) @@ -590,11 +596,25 @@ class Bragg2DPtycho(CDIModel): mask=mask) - def corrected_translations(self,dataset): - translations = dataset.translations.to(dtype=self.probe.real.dtype, - device=self.probe.device) - t_offset = tools.interactions.pixel_to_translations(self.obj_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal) - return translations + t_offset + def corrected_translations(self, dataset=None): + if dataset is not None: + translations = dataset.translations.to( + dtype=self.probe.real.dtype, device=self.probe.device) + elif (hasattr(self, 'original_translations') and + self.original_translations is not None): + translations = self.original_translations.to( + dtype=self.probe.real.dtype, device=self.probe.device) + else: + raise ValueError( + 'Must provide a dataset or have original_translations stored ' + 'internally (via from_dataset or from_results_dict).') + if self.translation_offsets is not None: + t_offset = tools.interactions.pixel_to_translations( + self.obj_basis, + self.translation_offsets * self.translation_scale, + surface_normal=self.surface_normal) + return translations + t_offset + return translations plot_list = [ @@ -679,13 +699,48 @@ class Bragg2DPtycho(CDIModel): units=self.units, )}, {'title': 'Corrected Translations', - 'plot_func': lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)}, + 'plot_func': lambda self, fig: p.plot_translations(self.corrected_translations(), fig=fig, units=self.units)}, {'title': 'Background', 'plot_func': lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)}, ] - def save_results(self, dataset): + @classmethod + def from_results_dict(cls, results_dict, obj_view_crop=0, units='um'): + sd = results_dict['state_dict'] + translation_offsets = sd.get('translation_offsets') + model = cls( + wavelength=sd['wavelength'], + detector_geometry={ + 'basis': sd['det_basis'], + 'distance': sd.get('det_distance'), + 'corner': sd.get('det_corner'), + }, + obj_basis=sd['obj_basis'], + probe_guess=sd['probe'], + obj_guess=sd['obj'], + min_translation=sd.get('min_translation', np.array([0., 0.])), + probe_basis=sd.get('probe_basis'), + median_propagation=sd.get('median_propagation', 0.0), + background=sd['background'], + translation_offsets=translation_offsets, + mask=sd.get('mask'), + weights=sd.get('weights'), + translation_scale=float(sd.get('translation_scale', 1.0)), + saturation=sd.get('saturation'), + oversampling=int(sd.get('oversampling', 1)), + propagate_probe=bool(sd.get('propagate_probe', True)), + correct_tilt=bool(sd.get('correct_tilt', True)), + loss=results_dict.get('loss_function', 'amplitude mse'), + obj_view_crop=obj_view_crop, + units=units, + translations=sd.get('original_translations'), + ) + model._load_results_dict(results_dict) + return model + + + def save_results(self, dataset=None): # This will save out everything needed to recreate the object # in the same state, but it's not the best formatted. For example, # "background" stores the square root of the background, etc. @@ -694,8 +749,11 @@ class Bragg2DPtycho(CDIModel): # We also save out the main results in a more readable format obj_basis = self.obj_basis.detach().cpu().numpy() probe_basis = self.probe_basis.detach().cpu().numpy() - translations=self.corrected_translations(dataset).detach().cpu().numpy() - original_translations = dataset.translations.detach().cpu().numpy() + translations = self.corrected_translations(dataset).detach().cpu().numpy() + if dataset is not None: + original_translations = dataset.translations.detach().cpu().numpy() + else: + original_translations = self.original_translations.detach().cpu().numpy() probe = self.probe.detach().cpu().numpy() probe = probe * self.probe_norm.detach().cpu().numpy() obj = self.obj.detach().cpu().numpy() diff --git a/src/cdtools/models/multislice_ptycho.py b/src/cdtools/models/multislice_ptycho.py index 03b1b52..3547351 100644 --- a/src/cdtools/models/multislice_ptycho.py +++ b/src/cdtools/models/multislice_ptycho.py @@ -40,6 +40,7 @@ class MultislicePtycho(CDIModel): obj_view_crop=0, panel_plot_mode=False, plot_level=1, + translations=None, ): super(MultislicePtycho, self).__init__(panel_plot_mode=panel_plot_mode, @@ -113,7 +114,7 @@ class MultislicePtycho(CDIModel): shape = [s//oversampling for s in self.probe[0]] background = 1e-6 * t.ones(shape, dtype=t.float32) - self.background = t.nn.Parameter(background) + self.background = t.nn.Parameter(t.as_tensor(background, dtype=dtype)) if weights is None: self.weights = None @@ -182,6 +183,10 @@ class MultislicePtycho(CDIModel): else: raise KeyError('Specified loss function not supported') + if translations is not None: + self.register_buffer('original_translations', + t.as_tensor(translations, dtype=dtype)) + @classmethod def from_dataset(cls, @@ -420,6 +425,7 @@ class MultislicePtycho(CDIModel): obj_view_crop=obj_view_crop, panel_plot_mode=panel_plot_mode, plot_level=plot_level, + translations=translations, ) @@ -591,9 +597,18 @@ class MultislicePtycho(CDIModel): mask=mask) - def corrected_translations(self, dataset): - translations = dataset.translations.to( - dtype=t.float32, device=self.probe.device) + def corrected_translations(self, dataset=None): + if dataset is not None: + translations = dataset.translations.to( + dtype=t.float32, device=self.probe.device) + elif (hasattr(self, 'original_translations') and + self.original_translations is not None): + translations = self.original_translations.to( + dtype=t.float32, device=self.probe.device) + else: + raise ValueError( + 'Must provide a dataset or have original_translations stored ' + 'internally (via from_dataset or from_results_dict).') if (hasattr(self, 'translation_offsets') and self.translation_offsets is not None): t_offset = tools.interactions.pixel_to_translations( @@ -711,7 +726,7 @@ class MultislicePtycho(CDIModel): self.weights.data = S[:,:,None] * (Vh / probe_sqrt_intensities) - def plot_wavefront_variation(self, dataset, fig=None, mode='amplitude', **kwargs): + def plot_wavefront_variation(self, dataset=None, fig=None, mode='amplitude', **kwargs): def get_probes(idx): basis_prs = self.probe * self.probe_support[..., :, :] prs = t.sum(self.weights[idx, :, :, None, None] * basis_prs, @@ -761,24 +776,21 @@ class MultislicePtycho(CDIModel): plot_list = [ {'title': '', - 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( - dataset, + 'plot_func': lambda self, fig: self.plot_wavefront_variation( fig=fig, mode='root_sum_intensity', image_title='Root Summed Probe Intensities', image_colorbar_title='Square Root of Intensity'), 'condition': lambda self: len(self.weights.shape) >= 2}, {'title': '', - 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( - dataset, + 'plot_func': lambda self, fig: self.plot_wavefront_variation( fig=fig, mode='amplitude', image_title='Probe Amplitudes (scroll to view modes)', image_colorbar_title='Probe Amplitude'), 'condition': lambda self: len(self.weights.shape) >= 2}, {'title': '', - 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( - dataset, + 'plot_func': lambda self, fig: self.plot_wavefront_variation( fig=fig, mode='phase', image_title='Probe Phases (scroll to view modes)', @@ -814,8 +826,8 @@ class MultislicePtycho(CDIModel): fig=fig), 'condition': lambda self: len(self.weights.shape) >= 2}, {'title': '% of Power in Top Mode', - 'plot_func': lambda self, fig, dataset: p.plot_nanomap( - self.corrected_translations(dataset), + 'plot_func': lambda self, fig: p.plot_nanomap( + self.corrected_translations(), 100 * t.stack([ analysis.calc_mode_power_fractions( self.probe.data, @@ -884,13 +896,51 @@ class MultislicePtycho(CDIModel): cmap='cividis'), 'condition': lambda self: self.exponentiate_obj}, {'title': 'Corrected Translations', - 'plot_func': lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)}, + 'plot_func': lambda self, fig: p.plot_translations(self.corrected_translations(), fig=fig, units=self.units)}, {'title': 'Background', 'plot_func': lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)}, ] - def save_results(self, dataset): + @classmethod + def from_results_dict(cls, results_dict, obj_view_crop=0, units='um'): + sd = results_dict['state_dict'] + translation_offsets = sd.get('translation_offsets') + model = cls( + wavelength=sd['wavelength'], + detector_geometry={ + 'basis': sd['det_basis'], + 'distance': sd.get('det_distance'), + 'corner': sd.get('det_corner'), + }, + obj_basis=sd['obj_basis'], + probe_guess=sd['probe'], + obj_guess=sd['obj'], + interslice_propagator=sd['interslice_propagator'], + surface_normal=sd.get('surface_normal', np.array([0., 0., 1.])), + min_translation=sd.get('min_translation', np.array([0., 0.])), + background=sd['background'], + probe_basis=sd.get('probe_basis'), + translation_offsets=translation_offsets, + mask=sd.get('mask'), + weights=sd.get('weights'), + translation_scale=float(sd.get('translation_scale', 1.0)), + saturation=sd.get('saturation'), + oversampling=int(sd.get('oversampling', 1)), + fourier_probe=bool(sd.get('fourier_probe', False)), + simulate_probe_translation=bool(sd.get('simulate_probe_translation', False)), + simulate_finite_pixels=bool(sd.get('simulate_finite_pixels', False)), + exponentiate_obj=bool(sd.get('exponentiate_obj', False)), + loss=results_dict.get('loss_function', 'amplitude mse'), + obj_view_crop=obj_view_crop, + units=units, + translations=sd.get('original_translations'), + ) + model._load_results_dict(results_dict) + return model + + + def save_results(self, dataset=None): # This will save out everything needed to recreate the object # in the same state, but it's not the best formatted. For example, # "background" stores the square root of the background, etc. @@ -899,8 +949,11 @@ class MultislicePtycho(CDIModel): # We also save out the main results in a more readable format obj_basis = self.obj_basis.detach().cpu().numpy() probe_basis = self.probe_basis.detach().cpu().numpy() - translations=self.corrected_translations(dataset).detach().cpu().numpy() - original_translations = dataset.translations.detach().cpu().numpy() + translations = self.corrected_translations(dataset).detach().cpu().numpy() + if dataset is not None: + original_translations = dataset.translations.detach().cpu().numpy() + else: + original_translations = self.original_translations.detach().cpu().numpy() probe = self.probe.detach().cpu().numpy() probe = probe * self.probe_norm.detach().cpu().numpy() obj = self.obj.detach().cpu().numpy() diff --git a/src/cdtools/models/rpi.py b/src/cdtools/models/rpi.py index 4bce45a..6e94103 100644 --- a/src/cdtools/models/rpi.py +++ b/src/cdtools/models/rpi.py @@ -111,10 +111,11 @@ class RPI(CDIModel): # We always use multi-modes to store the object, so we convert it # if we just get a single 2D array as an input + obj_guess = t.as_tensor(obj_guess, dtype=complex_dtype) if obj_guess.dim() == 2: obj_guess = obj_guess[None, :, :] - - self.obj = t.nn.Parameter(t.as_tensor(obj_guess, dtype=complex_dtype)) + + self.obj = t.nn.Parameter(obj_guess) self.weights = t.nn.Parameter( t.eye(probe.shape[0], dtype=complex_dtype)) @@ -632,6 +633,32 @@ class RPI(CDIModel): ] + @classmethod + def from_results_dict(cls, results_dict, units='um'): + sd = results_dict['state_dict'] + model = cls( + wavelength=sd['wavelength'], + detector_geometry={ + 'basis': sd['det_basis'], + 'distance': sd.get('det_distance'), + 'corner': sd.get('det_corner'), + }, + probe_basis=sd['probe_basis'], + probe=sd['probe'], + obj_guess=sd['obj'], + background=sd.get('background'), + mask=sd.get('mask'), + saturation=sd.get('saturation'), + oversampling=int(sd.get('oversampling', 1)), + exponentiate_obj=bool(sd.get('exponentiate_obj', False)), + phase_only=bool(sd.get('phase_only', False)), + loss=results_dict.get('loss_function', 'amplitude mse'), + units=units, + ) + model._load_results_dict(results_dict) + return model + + def save_results(self, dataset=None): # dataset is set as a kwarg here because it isn't needed, but the # common pattern is to pass a dataset. This makes it okay if one