From 2e6db5517dcc008b736e68c0e01888bb2ddaa2ce Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 22 Feb 2024 15:57:35 +0100 Subject: [PATCH] get the projected views for bragg_2d_Ptycho working --- .gitignore | 3 +- examples/gold_ball_ptycho.py | 2 +- src/cdtools/models/bragg_2d_ptycho.py | 56 +++++++++++++++++- src/cdtools/tools/plotting/plotting.py | 81 +++++++++++++++----------- 4 files changed, 103 insertions(+), 39 deletions(-) diff --git a/.gitignore b/.gitignore index 8032c1f..dcb9f36 100644 --- a/.gitignore +++ b/.gitignore @@ -7,4 +7,5 @@ docs/build build/* dist -*/example_data/* \ No newline at end of file +*/example_data/* +*.h5 \ No newline at end of file diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index e3ef2da..4ac5333 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -25,7 +25,7 @@ with model.save_on_exit('example_reconstructions/gold_balls.h5', dataset): 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/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index 463568b..84cdd13 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -52,6 +52,10 @@ __all__ = ['Bragg2DPtycho'] # in ggasp. So I believe this should not be a problem # +beam_basis = np.array([[0,-1], + [-1,0], + [0,0]]) + class Bragg2DPtycho(CDIModel): def __init__( @@ -537,20 +541,36 @@ class Bragg2DPtycho(CDIModel): 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', + ('Basis Probe Real Space Amplitudes, Surface Normal View', lambda self, fig: p.plot_amplitude( self.probe, fig=fig, basis=self.probe_basis, units=self.units, )), - ('Basis Probe Real Space Phases', + ('Basis Probe Real Space Phases, Surface Normal View', 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( + 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( + 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( self.obj, @@ -565,6 +585,38 @@ class Bragg2DPtycho(CDIModel): basis=self.obj_basis, units=self.units, )), + ('Object Amplitude, Beam View', + lambda self, fig: p.plot_amplitude( + self.obj, + fig=fig, + basis=self.obj_basis, + view_basis=beam_basis, + units=self.units, + )), + ('Object Phase, Beam View', + lambda self, fig: p.plot_phase( + self.obj, + fig=fig, + basis=self.obj_basis, + view_basis=beam_basis, + units=self.units, + )), + ('Object Amplitude, Detector View', + lambda self, fig: p.plot_amplitude( + self.obj, + fig=fig, + basis=self.obj_basis, + view_basis=self.det_basis, + units=self.units, + )), + ('Object Phase, Detector View', + lambda self, fig: p.plot_phase( + self.obj, + 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', diff --git a/src/cdtools/tools/plotting/plotting.py b/src/cdtools/tools/plotting/plotting.py index 58dc999..3043778 100644 --- a/src/cdtools/tools/plotting/plotting.py +++ b/src/cdtools/tools/plotting/plotting.py @@ -83,7 +83,7 @@ def get_units_factor(units): return factor -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): +def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, view_basis='ortho', 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 @@ -158,60 +158,71 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, view_basis=None, to_plot = plot_func(reshaped_im[fig.plot_idx]) - #Plot in a basis if it exists, otherwise dont + mpl_im = plt.imshow( + to_plot, + cmap = cmap, + interpolation = interpolation + ) + plt.gca().set_facecolor('k') + if basis is not None: + # we've closed over basis, so we can't edit it if isinstance(basis,t.Tensor): np_basis = basis.detach().cpu().numpy() else: np_basis = basis - 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 + if isinstance(view_basis, str) and view_basis.lower() == 'ortho': - 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) + # In this case, we construct a basis whose x-axis is + # parallel with the x-axis of the image basis, and whose + # y-axis lies in the x-y plane of the basis, perpendicular + # to the x-axis - coeffs = np.matmul(normed_ortho_basis.transpose(), - np_basis[:,::-1]) - [[a,c],[b,d]] = coeffs + 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) + + np_view_basis = np.stack( + [normed_yprime, normed_basis[:,1]], axis=1) + + else: + # We've also closed over view_basis, so we can't update it + if isinstance(view_basis,t.Tensor): + np_view_basis = view_basis.detach().cpu().numpy() + else: + np_view_basis = view_basis + + # We always normalize the view basis + view_basis_norm = np.linalg.norm(np_view_basis, axis = 0) + np_view_basis = np_view_basis / view_basis_norm + + # Holy cow, this works! + transform_matrix = \ + np.linalg.lstsq(np_view_basis[:,::-1], np_basis[:,::-1])[0] + [[a,c],[b,d]] = transform_matrix 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())) + corners = np.array([[-0.5,-0.5], + [im.shape[-1]-0.5,-0.5], + [-0.5, im.shape[-2]-0.5], + [im.shape[-1]-0.5, im.shape[-2]-0.5]]) + corners = np.matmul(transform_matrix,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]]) - + plt.gca().invert_yaxis() cbar = plt.colorbar() if cmap_label is not None: