diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index 705a3ed..e3ef2da 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -7,6 +7,8 @@ import torch as t filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi' dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename) +dataset.inspect() + # Next, we create a ptychography model from the dataset # Note that we explicitly ask for two incoherent probe modes model = cdtools.models.FancyPtycho.from_dataset( @@ -20,10 +22,10 @@ model = cdtools.models.FancyPtycho.from_dataset( with model.save_on_exit('example_reconstructions/gold_balls.h5', dataset): # Now, we run a short reconstruction from the dataset - for loss in model.Adam_optimize(1, dataset, batch_size=50): + for loss in model.Adam_optimize(10, dataset, batch_size=50): # And we liveplot the updates to the model as they happen print(model.report()) - model.inspect(dataset) + #model.inspect(dataset) # This orthogonalizes the incoherent probe modes model.tidy_probes() diff --git a/examples/lab_bragg_2d_ptycho.py b/examples/lab_bragg_2d_ptycho.py index 331c3b5..801ee5a 100644 --- a/examples/lab_bragg_2d_ptycho.py +++ b/examples/lab_bragg_2d_ptycho.py @@ -18,13 +18,16 @@ dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename) model = cdtools.models.Bragg2DPtycho.from_dataset( dataset, probe_support_radius=60, - correct_tilt=False + correct_tilt=False, + units='mm', ) #model.to(device='cuda') #dataset.get_as(device='cuda') - +model.inspect(dataset) +plt.show() +exit() model.translation_offsets.requires_grad = False for loss in model.Adam_optimize(100, dataset): model.inspect(dataset) diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index e27ea1e..47e1917 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -223,6 +223,7 @@ class Bragg2DPtycho(CDIModel): correct_tilt=True, lens=False, obj_padding=200, + units='um', ): wavelength = dataset.wavelength @@ -249,7 +250,7 @@ class Bragg2DPtycho(CDIModel): wavelength, distance, oversampling=oversampling) - + # now we grab the sample surface normal if hasattr(dataset, 'sample_info') and \ dataset.sample_info is not None and \ @@ -288,7 +289,7 @@ class Bragg2DPtycho(CDIModel): projector = np.linalg.pinv(mat)[:, :3] obj_basis = t.Tensor(np.dot(projector, ew_basis)) - + # Now we need a much better way to handle the translations here # than translations_to_pixel @@ -385,8 +386,7 @@ class Bragg2DPtycho(CDIModel): if not(correct_tilt is True or correct_tilt is False): raise NotImplementedError('No auto option implemented yet') - - + return cls(wavelength, det_geo, obj_basis, probe, obj, min_translation=min_translation, probe_basis=probe_basis, @@ -399,7 +399,9 @@ class Bragg2DPtycho(CDIModel): oversampling=oversampling, propagate_probe=propagate_probe, correct_tilt=correct_tilt, - lens=lens) + lens=lens, + units=units, + ) def interaction(self, index, translations): @@ -427,17 +429,7 @@ class Bragg2DPtycho(CDIModel): padding = [self.oversampling * self.background.shape[-2] - prs.shape[-2], self.oversampling * self.background.shape[-1] - prs.shape[-1]] - if any([p != 0 for p in padding]): # For probe_fourier_crop != 0. - padding = [padding[-1]//2, padding[-1]-padding[-1]//2, - padding[-2]//2, padding[-2]-padding[-2]//2] - prs = tools.propagators.far_field(prs) - prs = t.nn.functional.pad(prs, padding) - prs = tools.propagators.inverse_far_field(prs) - - # we do the upscaling before the propagation because the propagator - # is currently calculated based on the object pixel pitch, not the - # probe one - + # probe propagation if self.propagate_probe: for j in range(prs.shape[0]): # I believe this -1 sign is in error, but I need a dataset with @@ -447,6 +439,15 @@ class Bragg2DPtycho(CDIModel): * self.universal_propagator) prs[j] = tools.propagators.near_field(prs[j], propagator) + # we do the upscaling after the propagation because the propagator + # is calculated based on the probe basis (before upsampling) + if any([p != 0 for p in padding]): # For probe_fourier_crop != 0. + padding = [padding[-1]//2, padding[-1]-padding[-1]//2, + padding[-2]//2, padding[-2]-padding[-2]//2] + prs = tools.propagators.far_field(prs) + prs = t.nn.functional.pad(prs, padding) + prs = tools.propagators.inverse_far_field(prs) + exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc( prs, self.obj,pix_trans, shift_probe=True, multiple_modes=True) @@ -566,29 +567,31 @@ class Bragg2DPtycho(CDIModel): self.probe, fig=fig, basis=self.probe_basis, - units=self.units + units=self.units, )), ('Basis Probe Real Space Phases', lambda self, fig: p.plot_phase( self.probe, fig=fig, basis=self.probe_basis, - units=self.units + units=self.units, )), - ('Object Amplitude', + ('Object Amplitude, Surface Normal View', lambda self, fig: p.plot_amplitude( self.obj, fig=fig, - basis=self.obj_basis + basis=self.obj_basis, + units=self.units, )), - ('Object Phase', + ('Object Phase, Surface Normal View', lambda self, fig: p.plot_phase( self.obj, fig=fig, - basis=self.obj_basis + basis=self.obj_basis, + units=self.units, )), ('Corrected Translations', - lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)), + 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)) ] diff --git a/src/cdtools/tools/plotting/plotting.py b/src/cdtools/tools/plotting/plotting.py index 07177fb..58dc999 100644 --- a/src/cdtools/tools/plotting/plotting.py +++ b/src/cdtools/tools/plotting/plotting.py @@ -12,6 +12,7 @@ import matplotlib.pyplot as plt from matplotlib.colors import hsv_to_rgb from matplotlib.widgets import Slider from matplotlib import ticker, patheffects +from matplotlib import transforms as mtransforms __all__ = ['colorize', 'plot_amplitude', 'plot_phase', @@ -82,14 +83,15 @@ def get_units_factor(units): return factor -def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, interpolation=None, **kwargs): +def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, view_basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, interpolation=None, **kwargs): """Plots an image with a colorbar and on an appropriate spatial grid If a figure is given explicitly, it will clear that existing figure and plot over it. Otherwise, it will generate a new figure. If a basis is explicitly passed, the image will be plotted in real-space - coordinates + coordinates. If the optional view_basis is passed as well, then the + image will be projected into the plane of the view basis. Finally, if a function is passed to the plot_func argument, this function will be called on each slice of data before it is plotted. This is used @@ -162,16 +164,55 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', np_basis = basis.detach().cpu().numpy() else: np_basis = basis - # This fails if the basis is not rectangular - basis_norm = np.linalg.norm(np_basis, axis = 0) - basis_norm = basis_norm * get_units_factor(units) - extent = [0, im.shape[-1]*basis_norm[1], 0, - im.shape[-2]*basis_norm[0]] + + np_basis = np_basis * get_units_factor(units) + basis_norm = np.linalg.norm(np_basis, axis = 0) + + normed_basis = np_basis / basis_norm + normed_z = np.cross(normed_basis[:,1], normed_basis[:,0]) + normed_z /= np.linalg.norm(normed_z) + normed_yprime = np.cross(normed_z, normed_basis[:,1]) + normed_yprime /= np.linalg.norm(normed_yprime) + + extent = [0, im.shape[-1], 0, im.shape[-2]] + else: extent=None - plt.imshow(to_plot, cmap = cmap, extent = extent, interpolation=interpolation) + mpl_im = plt.imshow( + to_plot, + cmap = cmap, + extent = extent, + interpolation = interpolation + ) + plt.gca().set_facecolor('k') + if basis is not None: + normed_ortho_basis = np.stack( + [normed_basis[:,1], normed_yprime], axis=1) + + coeffs = np.matmul(normed_ortho_basis.transpose(), + np_basis[:,::-1]) + [[a,c],[b,d]] = coeffs + + transform = mtransforms.Affine2D.from_values(a,b,c,d,0,0) + + trans_data = transform + plt.gca().transData + + mpl_im.set_transform(trans_data) + corners = np.array([[0,0], + [im.shape[-1],0], + [0, im.shape[-2]], + [im.shape[-1], im.shape[-2]]]) + corners = np.matmul( + normed_ortho_basis.transpose(), + np.matmul(np_basis[:,::-1], corners.transpose())) + mins = np.min(corners, axis=1) + maxes = np.max(corners, axis=1) + plt.gca().set_xlim([mins[0], maxes[0]]) + plt.gca().set_ylim([mins[1], maxes[1]]) + + cbar = plt.colorbar() if cmap_label is not None: cbar.set_label(cmap_label)