From 91db5b3c6497cddb403900f780cd60fcb90d7d3d Mon Sep 17 00:00:00 2001 From: allevitan Date: Fri, 20 Mar 2026 19:01:22 +0100 Subject: [PATCH] Update remaining models to new plot registration system --- examples/transmission_RPI.py | 10 +- src/cdtools/models/bragg_2d_ptycho.py | 91 ++++++++------- src/cdtools/models/multislice_2d_ptycho.py | 99 ++++++++-------- src/cdtools/models/multislice_ptycho.py | 123 ++++++++++---------- src/cdtools/models/rpi.py | 127 ++++++++++++++------- 5 files changed, 257 insertions(+), 193 deletions(-) diff --git a/examples/transmission_RPI.py b/examples/transmission_RPI.py index feff19e..7b8d65a 100644 --- a/examples/transmission_RPI.py +++ b/examples/transmission_RPI.py @@ -1,5 +1,6 @@ import cdtools import pickle +import torch as t from matplotlib import pyplot as plt # First, we load an example dataset from a .cxi file @@ -18,17 +19,18 @@ dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(ss_filename) # Note that we explicitly as for two incoherent probe modes model = cdtools.models.RPI.from_dataset(dataset, probe, [500,500], background=background, n_modes=2, - initialization='random') + initialization='random',panel_plot_mode=True) # Let's do this reconstruction on the GPU, shall we? -model.to(device='cuda') -dataset.get_as(device='cuda') +if t.cuda.is_available(): + model.to(device='cuda') + dataset.get_as(device='cuda') # Note that the inspect step takes the vast majority of the time # The regularization is an L2 regularizer that empirically helps accelerate # convergence -for loss in model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05]): +for loss in model.Adam_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05]): model.inspect(dataset) print(model.report()) diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index 507323d..bd8b580 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -80,6 +80,8 @@ class Bragg2DPtycho(CDIModel): units='um', dtype=t.float32, obj_view_crop=0, + panel_plot_mode=False, + plot_level=1, ): # We need the detector geometry @@ -91,7 +93,8 @@ class Bragg2DPtycho(CDIModel): # translation_offsets can stay 2D for now # propagate_probe and correct_tilt are important! - super(Bragg2DPtycho, self).__init__() + super(Bragg2DPtycho, self).__init__(panel_plot_mode=panel_plot_mode, + plot_level=plot_level) self.register_buffer('wavelength', t.as_tensor(wavelength, dtype=dtype)) self.store_detector_geometry(detector_geometry, @@ -258,7 +261,9 @@ class Bragg2DPtycho(CDIModel): obj_padding=200, obj_view_crop=None, units='um', - surface_normal=None + surface_normal=None, + panel_plot_mode=False, + plot_level=1, ): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -436,8 +441,8 @@ class Bragg2DPtycho(CDIModel): return cls(wavelength, det_geo, obj_basis, probe, obj, min_translation=min_translation, probe_basis=probe_basis, - median_propagation =median_propagation, - translation_offsets = translation_offsets, + median_propagation=median_propagation, + translation_offsets=translation_offsets, weights=weights, mask=mask, background=background, translation_scale=translation_scale, saturation=saturation, @@ -448,6 +453,8 @@ class Bragg2DPtycho(CDIModel): lens=lens, obj_view_crop=obj_view_crop, units=units, + panel_plot_mode=panel_plot_mode, + plot_level=plot_level, ) @@ -578,90 +585,90 @@ class Bragg2DPtycho(CDIModel): plot_list = [ - ('Basis Probe Fourier Space Amplitudes', - lambda self, fig: p.plot_amplitude(tools.propagators.inverse_far_field(self.probe), fig=fig)), - ('Basis Probe Fourier Space Phases', - lambda self, fig: p.plot_phase(tools.propagators.inverse_far_field(self.probe), fig=fig)), - ('Basis Probe Real Space Amplitudes, Surface Normal View', - lambda self, fig: p.plot_amplitude( + {'title': 'Basis Probe Fourier Space Amplitudes', + 'plot_func': lambda self, fig: p.plot_amplitude(tools.propagators.inverse_far_field(self.probe), fig=fig)}, + {'title': 'Basis Probe Fourier Space Phases', + 'plot_func': lambda self, fig: p.plot_phase(tools.propagators.inverse_far_field(self.probe), fig=fig)}, + {'title': 'Basis Probe Real Space Amplitudes, Surface Normal View', + 'plot_func': lambda self, fig: p.plot_amplitude( self.probe, fig=fig, basis=self.probe_basis, units=self.units, - )), - ('Basis Probe Real Space Phases, Surface Normal View', - lambda self, fig: p.plot_phase( + )}, + {'title': 'Basis Probe Real Space Phases, Surface Normal View', + 'plot_func': lambda self, fig: p.plot_phase( self.probe, fig=fig, basis=self.probe_basis, units=self.units, - )), - ('Basis Probe Real Space Amplitudes, Beam View', - lambda self, fig: p.plot_amplitude( + )}, + {'title': 'Basis Probe Real Space Amplitudes, Beam View', + 'plot_func': lambda self, fig: p.plot_amplitude( self.probe, fig=fig, basis=self.probe_basis, view_basis=beam_basis, units=self.units, - )), - ('Basis Probe Real Space Phases, Beam View', - lambda self, fig: p.plot_phase( + )}, + {'title': 'Basis Probe Real Space Phases, Beam View', + 'plot_func': lambda self, fig: p.plot_phase( self.probe, fig=fig, basis=self.probe_basis, view_basis=beam_basis, units=self.units, - )), - ('Object Amplitude, Surface Normal View', - lambda self, fig: p.plot_amplitude( + )}, + {'title': 'Object Amplitude, Surface Normal View', + 'plot_func': lambda self, fig: p.plot_amplitude( self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, units=self.units, - )), - ('Object Phase, Surface Normal View', - lambda self, fig: p.plot_phase( + )}, + {'title': 'Object Phase, Surface Normal View', + 'plot_func': lambda self, fig: p.plot_phase( self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, units=self.units, - )), - ('Object Amplitude, Beam View', - lambda self, fig: p.plot_amplitude( + )}, + {'title': 'Object Amplitude, Beam View', + 'plot_func': lambda self, fig: p.plot_amplitude( self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, view_basis=beam_basis, units=self.units, - )), - ('Object Phase, Beam View', - lambda self, fig: p.plot_phase( + )}, + {'title': 'Object Phase, Beam View', + 'plot_func': lambda self, fig: p.plot_phase( self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, view_basis=beam_basis, units=self.units, - )), - ('Object Amplitude, Detector View', - lambda self, fig: p.plot_amplitude( + )}, + {'title': 'Object Amplitude, Detector View', + 'plot_func': lambda self, fig: p.plot_amplitude( self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, view_basis=self.det_basis, units=self.units, - )), - ('Object Phase, Detector View', - lambda self, fig: p.plot_phase( + )}, + {'title': 'Object Phase, Detector View', + 'plot_func': lambda self, fig: p.plot_phase( self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, view_basis=self.det_basis, units=self.units, - )), - ('Corrected Translations', - lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)), - ('Background', - lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)) + )}, + {'title': 'Corrected Translations', + 'plot_func': lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), 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)}, ] diff --git a/src/cdtools/models/multislice_2d_ptycho.py b/src/cdtools/models/multislice_2d_ptycho.py index 1d6b3cd..d6663be 100644 --- a/src/cdtools/models/multislice_2d_ptycho.py +++ b/src/cdtools/models/multislice_2d_ptycho.py @@ -47,9 +47,12 @@ class Multislice2DPtycho(CDIModel): prevent_aliasing=True, phase_only=False, units='um', + panel_plot_mode=False, + plot_level=1, ): - super(Multislice2DPtycho, self).__init__() + super(Multislice2DPtycho, self).__init__(panel_plot_mode=panel_plot_mode, + plot_level=plot_level) self.wavelength = t.tensor(wavelength) self.detector_geometry = copy(detector_geometry) self.dz = dz @@ -154,7 +157,7 @@ class Multislice2DPtycho(CDIModel): @classmethod - def from_dataset(cls, dataset, dz, nz, probe_convergence_semiangle, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True, probe_support_radius=None): + def from_dataset(cls, dataset, dz, nz, probe_convergence_semiangle, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True, probe_support_radius=None, panel_plot_mode=False, plot_level=1): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -286,7 +289,7 @@ class Multislice2DPtycho(CDIModel): surface_normal=surface_normal, probe_support=probe_support, min_translation=min_translation, - translation_offsets = translation_offsets, + translation_offsets=translation_offsets, weights=Ws, mask=mask, background=background, translation_scale=translation_scale, saturation=saturation, @@ -296,7 +299,9 @@ class Multislice2DPtycho(CDIModel): exponentiate_obj=exponentiate_obj, units=units, fourier_probe=fourier_probe, phase_only=phase_only, - prevent_aliasing=prevent_aliasing) + prevent_aliasing=prevent_aliasing, + panel_plot_mode=panel_plot_mode, + plot_level=plot_level) def interaction(self, index, translations): @@ -581,22 +586,22 @@ class Multislice2DPtycho(CDIModel): # Needs to be updated to allow for plotting to an existing figure - plot_list = [ - ('Probe Fourier Space Amplitude', - lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)), - ('Probe Fourier Space Phase', - lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)), - ('Probe Real Space Amplitude', - lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)), - ('Probe Real Space Phase', - lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)), - ('Average Weight Matrix Amplitudes', - lambda self, fig: p.plot_amplitude( + plot_list = [ + {'title': 'Probe Fourier Space Amplitude', + 'plot_func': lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)}, + {'title': 'Probe Fourier Space Phase', + 'plot_func': lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)}, + {'title': 'Probe Real Space Amplitude', + 'plot_func': lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)}, + {'title': 'Probe Real Space Phase', + 'plot_func': lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)}, + {'title': 'Average Weight Matrix Amplitudes', + 'plot_func': lambda self, fig: p.plot_amplitude( np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0), fig=fig), - lambda self: len(self.weights.shape) >= 2), - ('% of Power in Top Mode', - lambda self, fig, dataset: p.plot_nanomap( + '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), 100 * t.stack([ analysis.calc_mode_power_fractions( @@ -606,35 +611,35 @@ class Multislice2DPtycho(CDIModel): ], dim=0), fig=fig, units=self.units), - lambda self: len(self.weights.shape) >= 2), - ('Slice by Slice Real Part of T', - lambda self, fig: p.plot_real(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'), - lambda self: self.exponentiate_obj), - ('Slice by Slice Imaginary Part of T', - lambda self, fig: p.plot_imag(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units), - lambda self: self.exponentiate_obj), - ('Integrated Real Part of T', - lambda self, fig: p.plot_real(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'), - lambda self: (self.exponentiate_obj) and self.obj.dim() >= 3), - ('Integrated Imaginary Part of T', - lambda self, fig: p.plot_imag(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units), - lambda self: (self.exponentiate_obj) and self.obj.dim() >= 3), - ('Slice by Slice Amplitude of Object Function', - lambda self, fig: p.plot_amplitude(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units), - lambda self: not self.exponentiate_obj), - ('Slice by Slice Phase of Object Function', - lambda self, fig: p.plot_phase(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units,cmap='cividis'), - lambda self: not self.exponentiate_obj), - ('Amplitude of Stacked Object Function', - lambda self, fig: p.plot_amplitude(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units), - lambda self: (not self.exponentiate_obj) and self.obj.dim() >=3), - ('Phase of Stacked Object Function', - lambda self, fig: p.plot_phase(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'), - lambda self: (not self.exponentiate_obj) and self.obj.dim() >= 3), - ('Corrected Translations', - lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)), - ('Background', - lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)) + 'condition': lambda self: len(self.weights.shape) >= 2}, + {'title': 'Slice by Slice Real Part of T', + 'plot_func': lambda self, fig: p.plot_real(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'), + 'condition': lambda self: self.exponentiate_obj}, + {'title': 'Slice by Slice Imaginary Part of T', + 'plot_func': lambda self, fig: p.plot_imag(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units), + 'condition': lambda self: self.exponentiate_obj}, + {'title': 'Integrated Real Part of T', + 'plot_func': lambda self, fig: p.plot_real(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'), + 'condition': lambda self: self.exponentiate_obj and self.obj.dim() >= 3}, + {'title': 'Integrated Imaginary Part of T', + 'plot_func': lambda self, fig: p.plot_imag(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units), + 'condition': lambda self: self.exponentiate_obj and self.obj.dim() >= 3}, + {'title': 'Slice by Slice Amplitude of Object Function', + 'plot_func': lambda self, fig: p.plot_amplitude(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units), + 'condition': lambda self: not self.exponentiate_obj}, + {'title': 'Slice by Slice Phase of Object Function', + 'plot_func': lambda self, fig: p.plot_phase(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'), + 'condition': lambda self: not self.exponentiate_obj}, + {'title': 'Amplitude of Stacked Object Function', + 'plot_func': lambda self, fig: p.plot_amplitude(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units), + 'condition': lambda self: (not self.exponentiate_obj) and self.obj.dim() >= 3}, + {'title': 'Phase of Stacked Object Function', + 'plot_func': lambda self, fig: p.plot_phase(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'), + 'condition': lambda self: (not self.exponentiate_obj) and self.obj.dim() >= 3}, + {'title': 'Corrected Translations', + 'plot_func': lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), 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)}, ] diff --git a/src/cdtools/models/multislice_ptycho.py b/src/cdtools/models/multislice_ptycho.py index afcd83e..cec773c 100644 --- a/src/cdtools/models/multislice_ptycho.py +++ b/src/cdtools/models/multislice_ptycho.py @@ -39,10 +39,13 @@ class MultislicePtycho(CDIModel): simulate_finite_pixels=False, dtype=t.float32, exponentiate_obj=False, - obj_view_crop=0 + obj_view_crop=0, + panel_plot_mode=False, + plot_level=1, ): - super(MultislicePtycho, self).__init__() + super(MultislicePtycho, self).__init__(panel_plot_mode=panel_plot_mode, + plot_level=plot_level) self.register_buffer('wavelength', t.as_tensor(wavelength, dtype=dtype)) self.store_detector_geometry(detector_geometry, @@ -202,6 +205,8 @@ class MultislicePtycho(CDIModel): obj_view_crop=None, obj_padding=200, exponentiate_obj=False, + panel_plot_mode=False, + plot_level=1, ): wavelength = dataset.wavelength @@ -409,6 +414,8 @@ class MultislicePtycho(CDIModel): simulate_finite_pixels=simulate_finite_pixels, exponentiate_obj=exponentiate_obj, obj_view_crop=obj_view_crop, + panel_plot_mode=panel_plot_mode, + plot_level=plot_level, ) @@ -750,61 +757,61 @@ class MultislicePtycho(CDIModel): plot_list = [ - ('', - lambda self, fig, dataset: self.plot_wavefront_variation( + {'title': '', + 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( dataset, fig=fig, mode='root_sum_intensity', image_title='Root Summed Probe Intensities', image_colorbar_title='Square Root of Intensity'), - lambda self: len(self.weights.shape) >= 2), - ('', - lambda self, fig, dataset: self.plot_wavefront_variation( + 'condition': lambda self: len(self.weights.shape) >= 2}, + {'title': '', + 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( dataset, fig=fig, mode='amplitude', image_title='Probe Amplitudes (scroll to view modes)', image_colorbar_title='Probe Amplitude'), - lambda self: len(self.weights.shape) >= 2), - ('', - lambda self, fig, dataset: self.plot_wavefront_variation( + 'condition': lambda self: len(self.weights.shape) >= 2}, + {'title': '', + 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( dataset, fig=fig, mode='phase', image_title='Probe Phases (scroll to view modes)', image_colorbar_title='Probe Phase'), - lambda self: len(self.weights.shape) >= 2), - ('Basis Probe Fourier Space Amplitudes', - lambda self, fig: p.plot_amplitude( + 'condition': lambda self: len(self.weights.shape) >= 2}, + {'title': 'Basis Probe Fourier Space Amplitudes', + 'plot_func': lambda self, fig: p.plot_amplitude( (self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe)), - fig=fig)), - ('Basis Probe Fourier Space Phases', - lambda self, fig: p.plot_phase( + fig=fig)}, + {'title': 'Basis Probe Fourier Space Phases', + 'plot_func': lambda self, fig: p.plot_phase( (self.probe if self.fourier_probe - else tools.propagators.inverse_far_field(self.probe)) - , fig=fig)), - ('Basis Probe Real Space Amplitudes', - lambda self, fig: p.plot_amplitude( + else tools.propagators.inverse_far_field(self.probe)), + fig=fig)}, + {'title': 'Basis Probe Real Space Amplitudes', + 'plot_func': lambda self, fig: p.plot_amplitude( (self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe)), fig=fig, basis=self.probe_basis, - units=self.units)), - ('Basis Probe Real Space Phases', - lambda self, fig: p.plot_phase( + units=self.units)}, + {'title': 'Basis Probe Real Space Phases', + 'plot_func': lambda self, fig: p.plot_phase( (self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe)), fig=fig, basis=self.probe_basis, - units=self.units)), - ('Average Weight Matrix Amplitudes', - lambda self, fig: p.plot_amplitude( + units=self.units)}, + {'title': 'Average Weight Matrix Amplitudes', + 'plot_func': lambda self, fig: p.plot_amplitude( np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0), fig=fig), - lambda self: len(self.weights.shape) >= 2), - ('% of Power in Top Mode', - lambda self, fig, dataset: p.plot_nanomap( + '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), 100 * t.stack([ analysis.calc_mode_power_fractions( @@ -814,69 +821,69 @@ class MultislicePtycho(CDIModel): ], dim=0), fig=fig, units=self.units), - lambda self: len(self.weights.shape) >= 2), - ('Object Amplitude', - lambda self, fig: p.plot_amplitude( + 'condition': lambda self: len(self.weights.shape) >= 2}, + {'title': 'Object Amplitude', + 'plot_func': lambda self, fig: p.plot_amplitude( self.obj[(np.s_[:],) + self.obj_view_slice], fig=fig, basis=self.obj_basis, units=self.units), - lambda self: not self.exponentiate_obj), - ('Object (T) Imaginary Part', - lambda self, fig: p.plot_imag( + 'condition': lambda self: not self.exponentiate_obj}, + {'title': 'Object (T) Imaginary Part', + 'plot_func': lambda self, fig: p.plot_imag( self.obj[(np.s_[:],) + self.obj_view_slice], fig=fig, basis=self.obj_basis, units=self.units), - lambda self: self.exponentiate_obj), - ('Object Phase', - lambda self, fig: p.plot_phase( + 'condition': lambda self: self.exponentiate_obj}, + {'title': 'Object Phase', + 'plot_func': lambda self, fig: p.plot_phase( self.obj[(np.s_[:],) + self.obj_view_slice], fig=fig, basis=self.obj_basis, units=self.units), - lambda self: not self.exponentiate_obj), - ('Object (T) Real Part', - lambda self, fig: p.plot_real( + 'condition': lambda self: not self.exponentiate_obj}, + {'title': 'Object (T) Real Part', + 'plot_func': lambda self, fig: p.plot_real( self.obj[(np.s_[:],) + self.obj_view_slice], fig=fig, basis=self.obj_basis, units=self.units, cmap='cividis'), - lambda self: self.exponentiate_obj), - ('Object Product Amplitude', - lambda self, fig: p.plot_amplitude( + 'condition': lambda self: self.exponentiate_obj}, + {'title': 'Object Product Amplitude', + 'plot_func': lambda self, fig: p.plot_amplitude( t.prod(self.obj, dim=0)[self.obj_view_slice], fig=fig, basis=self.obj_basis, units=self.units), - lambda self: not self.exponentiate_obj), - ('Object (T) Sum Imaginary Part', - lambda self, fig: p.plot_imag( + 'condition': lambda self: not self.exponentiate_obj}, + {'title': 'Object (T) Sum Imaginary Part', + 'plot_func': lambda self, fig: p.plot_imag( t.sum(self.obj, dim=0)[self.obj_view_slice], fig=fig, basis=self.obj_basis, units=self.units), - lambda self: self.exponentiate_obj), - ('Object Product Phase', - lambda self, fig: p.plot_phase( + 'condition': lambda self: self.exponentiate_obj}, + {'title': 'Object Product Phase', + 'plot_func': lambda self, fig: p.plot_phase( t.prod(self.obj, dim=0)[self.obj_view_slice], fig=fig, basis=self.obj_basis, units=self.units), - lambda self: not self.exponentiate_obj), - ('Object (T) Sum Real Part', - lambda self, fig: p.plot_real( + 'condition': lambda self: not self.exponentiate_obj}, + {'title': 'Object (T) Sum Real Part', + 'plot_func': lambda self, fig: p.plot_real( t.sum(self.obj, dim=0)[self.obj_view_slice], fig=fig, basis=self.obj_basis, units=self.units, cmap='cividis'), - lambda self: self.exponentiate_obj), - ('Corrected Translations', - lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)), - ('Background', - lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)) + '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)}, + {'title': 'Background', + 'plot_func': lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)}, ] diff --git a/src/cdtools/models/rpi.py b/src/cdtools/models/rpi.py index 339b623..4dadb16 100644 --- a/src/cdtools/models/rpi.py +++ b/src/cdtools/models/rpi.py @@ -58,9 +58,12 @@ class RPI(CDIModel): propagation_distance=0, units='um', dtype=t.float32, + panel_plot_mode=False, + plot_level=1, ): - - super(RPI, self).__init__() + + super(RPI, self).__init__(panel_plot_mode=panel_plot_mode, + plot_level=plot_level) complex_dtype = (t.ones([1], dtype=dtype) + 1j * t.ones([1], dtype=dtype)).dtype @@ -164,6 +167,8 @@ class RPI(CDIModel): phase_only=False, probe_threshold=0, dtype=t.float32, + panel_plot_mode=False, + plot_level=1, ): complex_dtype = (t.ones([1], dtype=dtype) + 1j * t.ones([1], dtype=dtype)).dtype @@ -227,7 +232,7 @@ class RPI(CDIModel): # This will be superceded later by a call to init_obj, but it sets # the shape if obj_size is None: - obj_size = (np.array(self.probe.shape[-2:]) // 2).astype(int) + obj_size = (np.array(probe.shape[-2:]) // 2).astype(int) dummy_init_obj = t.ones([n_modes, obj_size[0], obj_size[1]], dtype=complex_dtype) @@ -247,14 +252,16 @@ class RPI(CDIModel): obj_support = t.as_tensor(binary_dilation(obj_support)) rpi_object = cls(wavelength, det_geo, ew_basis, - probe, dummy_init_obj, + probe, dummy_init_obj, background=background, mask=mask, saturation=saturation, obj_support=obj_support, oversampling=oversampling, exponentiate_obj=exponentiate_obj, phase_only=phase_only, - weight_matrix=weight_matrix) + weight_matrix=weight_matrix, + panel_plot_mode=panel_plot_mode, + plot_level=plot_level) # I don't love this pattern, where I do the "real" obj initialization # after creating the rpi object. But, I chose this so that I could @@ -283,7 +290,9 @@ class RPI(CDIModel): exponentiate_obj=False, phase_only=False, initialization='random', - dtype=t.float32 + dtype=t.float32, + panel_plot_mode=False, + plot_level=1, ): complex_dtype = (t.ones([1], dtype=dtype) + @@ -308,7 +317,7 @@ class RPI(CDIModel): # This will be superceded later by a call to init_obj, but it sets # the shape if obj_size is None: - obj_size = (np.array(self.probe.shape[-2:]) // 2).astype(int) + obj_size = (np.array(probe.shape[-2:]) // 2).astype(int) dummy_init_obj = t.ones([n_modes, obj_size[0], obj_size[1]], dtype=complex_dtype) @@ -327,6 +336,8 @@ class RPI(CDIModel): mask=mask, exponentiate_obj=exponentiate_obj, phase_only=phase_only, + panel_plot_mode=panel_plot_mode, + plot_level=plot_level, ) rpi_object.init_obj(initialization) @@ -365,7 +376,7 @@ class RPI(CDIModel): obj_shape=obj_shape, n_modes=n_modes) else: - raise KeyError('Initialization "' + str(initialization) + \ + raise KeyError('Initialization "' + str(initialization_type) + \ '" invalid - use "spectral", "uniform", or "random"') @@ -531,40 +542,72 @@ class RPI(CDIModel): def sim_to_dataset(self, args_list): raise NotImplementedError('No sim to dataset yet, sorry!') - plot_list = [ - ('Root Sum Squared Amplitude of all Probes', - lambda self, fig: p.plot_amplitude( - np.sqrt(np.sum((t.abs(t.sum(self.weights[..., None, None].detach() * self.probe, axis=-3))**2).cpu().numpy(),axis=0)), - fig=fig, basis=self.probe_basis)), - ('Object Amplitude', - lambda self, fig: p.plot_amplitude( - self.obj, - fig=fig, - basis=self.obj_basis, - units=self.units), - lambda self: not self.exponentiate_obj), - ('Object Phase', - lambda self, fig: p.plot_phase( - self.obj, - fig=fig, - basis=self.obj_basis, - units=self.units), - lambda self: not self.exponentiate_obj), - ('Real Part of T', - lambda self, fig: p.plot_real( - self.obj, - fig=fig, - basis=self.obj_basis, - units=self.units, - cmap='cividis'), - lambda self: self.exponentiate_obj), - ('Imaginary Part of T', - lambda self, fig: p.plot_imag( - self.obj, - fig=fig, - basis=self.obj_basis, - units=self.units), - lambda self: self.exponentiate_obj), + plot_panel_list = [ + { + 'title': 'RPI Results', + 'plot_level': 1, + 'figure_size': (12, 3.5), + 'grid': (1, 3), + 'plots': [ + { + 'title': 'Object Phase', + 'subplot': (0, 0), + 'plot_func': lambda self, fig: p.plot_phase( + self.obj, + fig=fig, + title='Object Phase', + basis=self.obj_basis, + units=self.units), + 'condition': lambda self: not self.exponentiate_obj, + }, + { + 'title': 'Real Part of T', + 'subplot': (0, 0), + 'plot_func': lambda self, fig: p.plot_real( + self.obj, + fig=fig, + title='Real Part of T', + basis=self.obj_basis, + units=self.units, + cmap='cividis'), + 'condition': lambda self: self.exponentiate_obj, + }, + { + 'title': 'Object Amplitude', + 'subplot': (0, 1), + 'plot_func': lambda self, fig: p.plot_amplitude( + self.obj, + fig=fig, + title='Object Amplitude', + basis=self.obj_basis, + units=self.units), + 'condition': lambda self: not self.exponentiate_obj, + }, + { + 'title': 'Imaginary Part of T', + 'subplot': (0, 1), + 'plot_func': lambda self, fig: p.plot_imag( + self.obj, + fig=fig, + title='Imaginary Part of T', + basis=self.obj_basis, + units=self.units), + 'condition': lambda self: self.exponentiate_obj, + }, + { + 'title': 'Root Sum Squared Amplitude of all Probes', + 'subplot': (0, 2), + 'plot_func': lambda self, fig: p.plot_amplitude( + np.sqrt(np.sum( + (t.abs(t.sum(self.weights[..., None, None].detach() + * self.probe, axis=-3))**2 + ).cpu().numpy(), axis=0)), + fig=fig, + basis=self.probe_basis, + units=self.units), + }, + ], + }, ]