cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-09 21:12:42 +02:00
Code Issues Packages Projects Releases Wiki Activity

get the projected views for bragg_2d_Ptycho working

Browse Source
This commit is contained in:
levita_a
2024-02-22 15:57:35 +01:00
parent 30f9b142b6
commit 2e6db5517d
4 changed files with 103 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2 -1
View File
@@ -7,4 +7,5 @@
docs/build
build/*
dist
*/example_data/*
*/example_data/*
*.h5
examples/gold_ball_ptycho.py
+1 -1
View File
@@ -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()
src/cdtools/models/bragg_2d_ptycho.py
+54 -2
View File
@@ -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',
src/cdtools/tools/plotting/plotting.py
+46 -35
View File
@@ -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: