diff --git a/examples/fancy_ptycho.py b/examples/fancy_ptycho.py index 8e3a9d1..3ac8974 100644 --- a/examples/fancy_ptycho.py +++ b/examples/fancy_ptycho.py @@ -14,7 +14,6 @@ model = cdtools.models.FancyPtycho.from_dataset( propagation_distance=5e-3, # Propagate the initial probe guess by 5 mm units='mm', # Set the units for the live plots obj_view_crop=-50, # Expands the field of view in the object plot by 50 pix, - panel_plot_mode=True, # Organizes the live plots into panels ) if t.cuda.is_available(): @@ -34,7 +33,7 @@ recon = cdtools.reconstructors.AdamReconstructor(model, dataset) for loss in recon.optimize(50, lr=0.02, batch_size=10): print(model.report()) # Plotting is expensive, so we only do it every tenth epoch - if model.epoch % 2 == 0: + if model.epoch % 10 == 0: model.inspect(dataset) # It's common to chain several different reconstruction loops. Here, we diff --git a/examples/transmission_RPI.py b/examples/transmission_RPI.py index 7b8d65a..c8fd424 100644 --- a/examples/transmission_RPI.py +++ b/examples/transmission_RPI.py @@ -19,7 +19,7 @@ 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',panel_plot_mode=True) + initialization='random') # Let's do this reconstruction on the GPU, shall we? diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 9ec668b..44448b3 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -45,7 +45,7 @@ class FancyPtycho(CDIModel): near_field=False, angular_spectrum_propagator=None, inv_angular_spectrum_propagator=None, - panel_plot_mode=False, + panel_plot_mode=True, plot_level=2, ): @@ -255,7 +255,7 @@ class FancyPtycho(CDIModel): obj_view_crop=None, obj_padding=200, near_field=False, - panel_plot_mode=False, + panel_plot_mode=True, plot_level=2, ): @@ -846,6 +846,44 @@ class FancyPtycho(CDIModel): # We discard the U matrix and re-multiply S & Vh self.weights.data = S[:,:,None] * (Vh / probe_sqrt_intensities) + + def get_probe_intensities(self): + if not hasattr(self, 'weights'): + raise NotImplementedError( + "I don't know how to handle having no weights") + elif self.weights.ndim == 1: + probe_intensities = self.weights.detach().cpu().numpy()**2 + else: + # The big case, with OPRP + probe_matrix = np.zeros([self.probe.shape[0]]*2, + dtype=np.complex64) + np_probes = self.probe.detach().cpu().numpy() + for i in range(probe_matrix.shape[0]): + for j in range(probe_matrix.shape[0]): + probe_matrix[i,j] = np.sum(np_probes[i]*np_probes[j].conj()) + + weights = self.weights.detach().cpu().numpy() + + # The outer one is a sum, because the tensordot is what broadcasts + # the probe matrix along the shot dimension - the second one + # doesn't have to. + weighted_probe_matrices = np.sum( + np.tensordot(weights, probe_matrix, axes=1)[...,None] + * weights.conj().transpose((0,2,1))[...,None,:,:], + axis=-2 + ) + + basis_probe_intensities = np.trace( + probe_matrix, axis1=-2, axis2=-1) + probe_intensities = np.trace( + weighted_probe_matrices, axis1=-2, axis2=-1) + + # Imaginary part is already essentially zero up to rounding error + probe_intensities = np.real( + probe_intensities / basis_probe_intensities) + + return probe_intensities + def plot_wavefront_variation(self, dataset, fig=None, mode='amplitude', **kwargs): def get_probes(idx): @@ -862,22 +900,8 @@ class FancyPtycho(CDIModel): if mode.lower() == 'phase': return np.angle(ortho_probes.detach().cpu().numpy()) - probe_matrix = np.zeros([self.probe.shape[0]]*2, - dtype=np.complex64) - np_probes = self.probe.detach().cpu().numpy() - for i in range(probe_matrix.shape[0]): - for j in range(probe_matrix.shape[0]): - probe_matrix[i,j] = np.sum(np_probes[i]*np_probes[j].conj()) - - weights = self.weights.detach().cpu().numpy() - - probe_intensities = np.sum(np.tensordot(weights, probe_matrix, axes=1) - * weights.conj(), axis=2) - - # Imaginary part is already essentially zero up to rounding error - probe_intensities = np.real(probe_intensities) - - values = np.sum(probe_intensities, axis=1) + values = self.get_probe_intensities() + if mode.lower() == 'amplitude' or mode.lower() == 'root_sum_intensity': cmap = 'viridis' else: @@ -896,35 +920,9 @@ class FancyPtycho(CDIModel): def plot_illumination_intensity(self, fig, dataset): - if not hasattr(self, 'weights'): - raise NotImplementedError("I don't know how to handle having no weights") - elif self.weights.ndim == 1: - probe_intensities = self.weights.detach().cpu().numpy()**2 - else: - # The big case, with OPRP - probe_matrix = np.zeros([self.probe.shape[0]]*2, - dtype=np.complex64) - np_probes = self.probe.detach().cpu().numpy() - for i in range(probe_matrix.shape[0]): - for j in range(probe_matrix.shape[0]): - probe_matrix[i,j] = np.sum(np_probes[i]*np_probes[j].conj()) - - weights = self.weights.detach().cpu().numpy() - - # The outer one is a sum, because the tensordot is what broadcasts the - # probe matrix along the shot dimension - the second one doesn't have to. - weighted_probe_matrices = np.sum(np.tensordot(weights, probe_matrix, axes=1)[...,None] - * weights.conj().transpose((0,2,1))[...,None,:,:], axis=-2) - - basis_probe_intensities = np.trace(probe_matrix, axis1=-2, axis2=-1) - probe_intensities = np.trace(weighted_probe_matrices, axis1=-2, axis2=-1) - - # Imaginary part is already essentially zero up to rounding error - probe_intensities = np.real(probe_intensities / basis_probe_intensities) - p.plot_nanomap( self.corrected_translations(dataset), - probe_intensities, + self.get_probe_intensities(), fig=fig, cmap='magma', cmap_label='Intensity (a.u.)', @@ -1011,7 +1009,7 @@ class FancyPtycho(CDIModel): (self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe)), fig=fig, - title='Basis Probe', + title='Basis Probes', basis=self.probe_basis, units=self.units), }, @@ -1022,7 +1020,7 @@ class FancyPtycho(CDIModel): (self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe)), fig=fig, - title='Basis Probe', + title='Basis Probes', basis=self.probe_basis, units=self.units), }, @@ -1041,7 +1039,7 @@ class FancyPtycho(CDIModel): (self.probe if self.fourier_probe else tools.propagators.far_field(self.probe)), fig=fig, - title='Basis Probe, Fourier', + title='Basis Probes, Fourier', ), }, { @@ -1051,7 +1049,7 @@ class FancyPtycho(CDIModel): (self.probe if self.fourier_probe else tools.propagators.far_field(self.probe)), fig=fig, - title='Basis Probe, Fourier', + title='Basis Probes, Fourier', ), }, { diff --git a/src/cdtools/models/rpi.py b/src/cdtools/models/rpi.py index 4dadb16..0cf0b68 100644 --- a/src/cdtools/models/rpi.py +++ b/src/cdtools/models/rpi.py @@ -58,7 +58,7 @@ class RPI(CDIModel): propagation_distance=0, units='um', dtype=t.float32, - panel_plot_mode=False, + panel_plot_mode=True, plot_level=1, ): @@ -167,7 +167,7 @@ class RPI(CDIModel): phase_only=False, probe_threshold=0, dtype=t.float32, - panel_plot_mode=False, + panel_plot_mode=True, plot_level=1, ): complex_dtype = (t.ones([1], dtype=dtype) + @@ -291,7 +291,7 @@ class RPI(CDIModel): phase_only=False, initialization='random', dtype=t.float32, - panel_plot_mode=False, + panel_plot_mode=True, plot_level=1, ): diff --git a/src/cdtools/tools/plotting/plotting.py b/src/cdtools/tools/plotting/plotting.py index 5380871..362489c 100644 --- a/src/cdtools/tools/plotting/plotting.py +++ b/src/cdtools/tools/plotting/plotting.py @@ -202,7 +202,7 @@ def plot_image( fig.canvas.toolbar.home() fig.canvas.toolbar.update() - if len(im.shape) >= 3: + if num_images > 1: base = title if title is not None else '('.join(ax_title.split('(')[:-1])[:-1] fig.axes[0].set_title(base + f' ({fig.plot_idx+1} of {num_images})') @@ -295,7 +295,7 @@ def plot_image( if title is not None: ax.set_title(ax_title) - if len(im.shape) >= 3: + if num_images >= 3: ax.set_title(ax_title + f' ({fig.plot_idx+1} of {num_images})') if fig.canvas.toolbar is not None: