cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-10-13 11:30:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

honestly, not sure

Browse Source
This commit is contained in:
Abe Levitan
2024-02-12 09:49:26 -05:00
parent dd30f58d2d
commit 1502d567bf
4 changed files with 84 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/models/bragg_2d_ptycho.py
+26 -23
View File
@@ -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))
]
src/cdtools/tools/plotting/plotting.py
+49 -8
View File
@@ -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)