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

Merging a whole buncha work back into master

Browse Source
This commit is contained in:
Abe Levitan
2022-04-18 15:51:45 -07:00
parent 1b5ef2271c 20a1479859
commit f3cee949e0
22 changed files with 3382 additions and 771 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/tools/plotting/plotting.py
+54 -43
View File
@@ -84,7 +84,7 @@ def get_units_factor(units):
def plot_image(im, plot_func=lambda x: x, fig=None, 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.
@@ -94,7 +94,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
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
internally to enable the plot_real, plot_image, plot_phase, etc. functions.
Parameters
----------
@@ -122,7 +122,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
# convert to numpy
if isinstance(im, t.Tensor):
# If final dimension is 2, assume it is a complex array. If not,
@@ -144,7 +144,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
title = plt.gca().get_title()
fig.clear()
# If im only has two dimensions, this reshape will add a leading
# dimension, and update will be called on index 0. If it has 3 or more
# dimensions, then all the leading dimensions will be compressed into
@@ -153,9 +153,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
reshaped_im = im.reshape(-1,s[-2],s[-1])
num_images = reshaped_im.shape[0]
fig.plot_idx = idx % num_images
to_plot = plot_func(reshaped_im[fig.plot_idx])
#Plot in a basis if it exists, otherwise dont
if basis is not None:
if isinstance(basis,t.Tensor):
@@ -183,7 +183,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
plt.title(title)
if len(im.shape) >= 3:
@@ -196,14 +196,14 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
result = make_plot(0)
update = make_plot
def on_action(event):
if not hasattr(event, 'button'):
event.button = None
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up':
update(fig.plot_idx - 1)
elif event.key == 'down' or event.button == 'down':
@@ -219,9 +219,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
fig.my_callbacks = []
fig.my_callbacks.append(fig.canvas.mpl_connect('key_press_event',on_action))
fig.my_callbacks.append(fig.canvas.mpl_connect('scroll_event',on_action))
return result
def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', **kwargs):
"""Plots the real part of a complex array with dimensions NxM
@@ -258,7 +258,7 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', **kwargs):
@@ -306,7 +306,7 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis',
If a basis is explicitly passed, the image will be plotted in real-space
coordinates.
Parameters
----------
im : array
@@ -418,7 +418,7 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
units=units, **kwargs)
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwargs):
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, **kwargs):
"""Plots a set of probe translations in a nicely formatted way
Parameters
@@ -431,6 +431,8 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
Default is um, units to report in (assuming input in m)
lines : bool
Whether to plot lines indicating the path taken
invert_xaxis : bool
Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -455,6 +457,9 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
translations = translations * factor
plt.plot(translations[:,0], translations[:,1],'k.')
if invert_xaxis:
plt.gca().invert_xaxis()
if lines:
plt.plot(translations[:,0], translations[:,1],'b-', linewidth=0.5)
plt.xlabel('X (' + units + ')')
@@ -463,7 +468,7 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
return fig
def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='probe'):
def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='probe', invert_xaxis=True):
"""Plots a set of nanomap data in a flexible way
Parameters
@@ -478,6 +483,8 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
Default is um, units to report in (assuming input in m)
convention : str
Default is 'probe', alternative is 'obj'. Whether the translations refer to the probe or object.
invert_xaxis : bool
Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective
Returns
-------
@@ -511,7 +518,9 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
s /= 4 # A rough value to make the size work out
plt.scatter(factor * trans[:,0],factor * trans[:,1],s=s,c=values)
if invert_xaxis:
plt.gca().invert_xaxis()
plt.gca().set_facecolor('k')
plt.xlabel('Translation x (' + units + ')')
plt.ylabel('Translation y (' + units + ')')
@@ -522,12 +531,12 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
def plot_nanomap_with_images(translations, get_image_func, values=None, mask=None, basis=None, fig=None, nanomap_units='$\\mu$m', image_units='$\\mu$m', convention='probe', image_title='Image', image_colorbar_title='Image Amplitude', nanomap_colorbar_title='Integrated Intensity', cmap='viridis', **kwargs):
"""Plots a nanomap, with an image or stack of images for each point
In many situations, ptychography data or the output of ptychography
reconstructions is formatted as a set of images associated with various
points in real space. This function is designed to allow for browsing
through this kind of data, by making it possible to visualize a
"""
# This should pull heavily from the dataset.inspect function
@@ -560,11 +569,11 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
s0 /= 4 # A rough value to make the size work out
s = np.ones(translations.shape[0]) * s0
s[idx] *= 4
return s
def update_colorbar(im):
#
# This solves the problem of the colorbar being changed
@@ -573,29 +582,29 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
if hasattr(im, 'norecurse') and im.norecurse:
im.norecurse=False
return
im.norecurse=True
# This is needed to update the colorbar
# only change limits if array contains multiple values
if np.min(im.get_array()) != np.max(im.get_array()):
im.set_clim(vmin=np.min(im.get_array()),
vmax=np.max(im.get_array()))
#
# The meatiest part of this program, here we just go through and
# set up the plot how we want it
#
# First we set up the left-hand plot, which shows an overview map
axes[0].set_title('Relative Displacement Map')
translations = translations.detach().cpu().numpy()
if convention.lower() != 'probe':
translations = translations * -1
s = calculate_sizes(0)
nanomap_units_factor = get_units_factor(nanomap_units)
nanomap = axes[0].scatter(nanomap_units_factor * translations[:,0],
nanomap_units_factor * translations[:,1],
@@ -616,7 +625,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# where the colorbar should have been to avoid stretching the
# nanomap plot, while still not showing the (now useless) colorbar.
cb1.remove()
# Now we set up the second plot, which shows the individual
# diffraction patterns
axes[1].set_title(image_title)
@@ -636,7 +645,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# This fails if the basis is not rectangular
basis_norm = np.linalg.norm(np_basis, axis = 0)
basis_norm = basis_norm * get_units_factor(image_units)
extent = [0, example_im.shape[-1]*basis_norm[1], 0,
example_im.shape[-2]*basis_norm[0]]
else:
@@ -657,9 +666,9 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
axes[1].text_box.set_path_effects(
[patheffects.Stroke(linewidth=2, foreground='black'),
patheffects.Normal()])
meas = axes[1].imshow(im, extent=extent, cmap=cmap)
cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',
format='%.2e',
ticks=ticker.LinearLocator(numticks=5),
@@ -667,7 +676,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
cb2.ax.tick_params(labelrotation=20)
cb2.ax.set_title(image_colorbar_title, size="medium", pad=5)
cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas))
# This function handles all the updating, except for moving the
# slider value. This is done because the slider widget is
# ultimately responsible for triggering an update, so all other
@@ -677,7 +686,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# We have to explicitly make it an integer because the slider will
# output floats (even if they are still integer-valued)
idx = int(idx)
# Get the new data for this index
im = get_image_func(idx)
if len(im.shape) >= 3:
@@ -688,22 +697,22 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
axes[1].image_idx = im_idx
axes[1].text_box.set_text(str(im_idx))
im = im.reshape(-1,im.shape[-2],im.shape[-1])[im_idx]
# Now we resize the nanomap to show the new selection
axes[0].collections[0].set_sizes(calculate_sizes(idx))
# And we update the data in the image as well
ax_im = axes[1].images[-1]
ax_im.set_data(im)
update_colorbar(ax_im)
#
# Now we define the functions to handle various kinds of events
# that can be thrown our way
#
# We start by creating the slider here, so it can be used
# by the update hooks.
slider = Slider(axslider, 'Image #', 0, translations.shape[0]-1, valstep=1, valfmt="%d")
@@ -717,7 +726,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# while the mouse is within the image display
im = im.reshape(-1,im.shape[-2],im.shape[-1])
im_idx = axes[1].image_idx
if event.key == 'up' or event.button == 'up' \
or event.key == 'left':
im_idx = (im_idx - 1) % im.shape[0]
@@ -735,7 +744,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
event.button = None
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up' or event.key == 'left':
idx = slider.val - 1
elif event.key == 'down' or event.button == 'down' or event.key == 'right':
@@ -744,7 +753,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# This prevents errors from being thrown on irrelevant key
# or mouse input
return
# Handle the wraparound and trigger the update
idx = int(idx) % translations.shape[0]
slider.set_val(idx)
@@ -755,17 +764,19 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# for example, scroll events that happen over the nanomap
if event.mouseevent.button == 1:
slider.set_val(event.ind[0])
# Here we connect the various update functions
cid1 = fig.canvas.mpl_connect('pick_event',on_pick)
cid2 = fig.canvas.mpl_connect('key_press_event',on_action)
cid3 = fig.canvas.mpl_connect('scroll_event',on_action)
# It's so dumb that matplotlib doesn't automatically track this for you
fig.nanomap_cids = [cid1,cid2,cid3]
fig.nanomap_cids = [cid1,cid2,cid3]
slider.on_changed(update)
# Throw an extra update into the mix just to get rid of any things
# (like the nanomap dot sizes) that otherwise would change on the
# first update
update(0)
return fig