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

Add title option to plot_image and wrappers; misc fixes

Browse Source 
- Add title parameter to plot_image, plot_real, plot_imag, plot_amplitude,
  plot_phase, and plot_colorized
- Various fixes to base.py and fancy_ptycho.py

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
allevitan
2026-03-20 17:20:52 +01:00
co-authored by Claude Sonnet 4.6
parent dd744a00a4
commit f4260837bc
4 changed files with 145 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/fancy_ptycho.py
+1 -4
View File
@@ -14,9 +14,7 @@ model = cdtools.models.FancyPtycho.from_dataset(
propagation_distance=5e-3, # Propagate the initial probe guess by 5 mm
units='mm', # Set the units for the live plots
obj_view_crop=-50, # Expands the field of view in the object plot by 50 pix,
exponentiate_obj=False,
panel_plot_mode=True,
plot_level=2,
panel_plot_mode=True, # Organizes the live plots into panels
)
if t.cuda.is_available():
@@ -30,7 +28,6 @@ if t.cuda.is_available():
# e.g. estimates of the moments of individual parameters
recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
# The learning rate parameter sets the alpha for Adam.
# The beta parameters are (0.9, 0.999) by default
# The batch size sets the minibatch size
src/cdtools/models/base.py
+31 -7
View File
@@ -675,6 +675,7 @@ class CDIModel(t.nn.Module):
if not condition(self, dataset):
continue
figsize = plot.get('figure_size', None)
if self.has_inspect_been_called and \
not replot_all and \
not plt.fignum_exists(plot['title']):
@@ -682,22 +683,30 @@ class CDIModel(t.nn.Module):
if not self.has_inspect_been_called:
fig = plt.figure(plot['title'],
figsize=figsize,
constrained_layout=True)
else:
with plt.rc_context({'figure.raise_window': False}):
fig = plt.figure(plot['title'],
figsize = panel_def.get('figure_size', None)
constrained_layout=True)
try:
plot['plot_func'](self, fig)
plt.title(plot['title'])
if plt.gca().get_title().strip() == '':
plt.title(plot['title'])
except TypeError:
if dataset is not None:
try:
plot['plot_func'](self, fig, dataset)
plt.title(plot['title'])
if plt.gca().get_title().strip() == '':
plt.title(plot['title'])
except KeyboardInterrupt:
raise
except Exception:
pass
except KeyboardInterrupt:
raise
except Exception:
pass
@@ -723,6 +732,15 @@ class CDIModel(t.nn.Module):
panel_level = panel_def.get('plot_level', 1)
if panel_level > self.plot_level:
continue # skip entire panel
panel_condition = panel_def.get('condition', None)
if panel_condition is not None:
try:
if not panel_condition(self):
continue
except TypeError:
if not panel_condition(self, dataset):
continue
nrows, ncols = panel_def['grid']
figsize = panel_def.get('figure_size', None)
@@ -767,16 +785,22 @@ class CDIModel(t.nn.Module):
try:
plot['plot_func'](self, subfig)
plt.gca().set_title(plot['title'])
if plt.gca().get_title().strip() == '':
plt.title(plot['title'])
except TypeError:
if dataset is not None:
try:
plot['plot_func'](self, subfig, dataset)
plt.gca().set_title(plot['title'])
except TypeError:#Exception:
if plt.gca().get_title().strip() == '':
plt.title(plot['title'])
except KeyboardInterrupt:
raise
except Exception:
pass
#except Exception:
# pass
except KeyboardInterrupt:
raise
except Exception:
pass
rendered.append(fig)
src/cdtools/models/fancy_ptycho.py
+75 -30
View File
@@ -896,11 +896,35 @@ class FancyPtycho(CDIModel):
def plot_illumination_intensity(self, fig, dataset):
if not hasattr(self, 'weights') or self.weights.ndim != 1:
raise NotImplementedError('Not yet implemented for OPRP')
if not hasattr(self, 'weights'):
raise NotImplementedError("I don't know how to handle having no weights")
elif self.weights.ndim == 1:
probe_intensities = self.weights.detach().cpu().numpy()**2
else:
# The big case, with OPRP
probe_matrix = np.zeros([self.probe.shape[0]]*2,
dtype=np.complex64)
np_probes = self.probe.detach().cpu().numpy()
for i in range(probe_matrix.shape[0]):
for j in range(probe_matrix.shape[0]):
probe_matrix[i,j] = np.sum(np_probes[i]*np_probes[j].conj())
weights = self.weights.detach().cpu().numpy()
# The outer one is a sum, because the tensordot is what broadcasts the
# probe matrix along the shot dimension - the second one doesn't have to.
weighted_probe_matrices = np.sum(np.tensordot(weights, probe_matrix, axes=1)[...,None]
* weights.conj().transpose((0,2,1))[...,None,:,:], axis=-2)
basis_probe_intensities = np.trace(probe_matrix, axis1=-2, axis2=-1)
probe_intensities = np.trace(weighted_probe_matrices, axis1=-2, axis2=-1)
# Imaginary part is already essentially zero up to rounding error
probe_intensities = np.real(probe_intensities / basis_probe_intensities)
p.plot_nanomap(
self.corrected_translations(dataset),
self.weights**2,
probe_intensities,
fig=fig,
cmap='magma',
cmap_label='Intensity (a.u.)',
@@ -908,7 +932,7 @@ class FancyPtycho(CDIModel):
convention='probe',
invert_xaxis=True
)
def plot_translations_and_originals(self, fig, dataset):
"""Only used to make a plot for the plot list."""
@@ -987,6 +1011,7 @@ class FancyPtycho(CDIModel):
(self.probe if not self.fourier_probe
else tools.propagators.inverse_far_field(self.probe)),
fig=fig,
title='Basis Probe',
basis=self.probe_basis,
units=self.units),
},
@@ -997,6 +1022,7 @@ class FancyPtycho(CDIModel):
(self.probe if not self.fourier_probe
else tools.propagators.inverse_far_field(self.probe)),
fig=fig,
title='Basis Probe',
basis=self.probe_basis,
units=self.units),
},
@@ -1014,7 +1040,9 @@ class FancyPtycho(CDIModel):
'plot_func': lambda self, fig: p.plot_colorized(
(self.probe if self.fourier_probe
else tools.propagators.far_field(self.probe)),
fig=fig),
fig=fig,
title='Basis Probe, Fourier',
),
},
{
'title': 'Basis Probes, Fourier Amplitude',
@@ -1022,13 +1050,14 @@ class FancyPtycho(CDIModel):
'plot_func': lambda self, fig: p.plot_amplitude(
(self.probe if self.fourier_probe
else tools.propagators.far_field(self.probe)),
fig=fig),
fig=fig,
title='Basis Probe, Fourier',
),
},
{
'title': 'Illumination Intensity',
'subplot': (0,1),
'plot_func': lambda self, fig, dataset: self.plot_illumination_intensity(fig, dataset),
'condition': lambda self: hasattr(self, 'weights') and self.weights.ndim == 1
},
{
'title': 'Detector Background',
@@ -1047,11 +1076,48 @@ class FancyPtycho(CDIModel):
},
],
},
{
'title': 'Unstable Probe Refinement Details',
'plot_level': 2,
'figure_size': (9,3.5),
'grid': (1,2),
'condition': lambda self: len(self.weights.shape) >= 2,
'plots': [
{
'title': '% of Power in Top Mode',
'subplot': (0,0),
'plot_func': lambda self, fig, dataset: p.plot_nanomap(
self.corrected_translations(dataset),
100 * t.stack([
analysis.calc_mode_power_fractions(
self.probe.data,
weight_matrix=self.weights.data[i])[0]
for i in range(self.weights.shape[0])
], dim=0),
fig=fig,
units=self.units),
'condition': lambda self: len(self.weights.shape) >= 2
},
{
'title': 'Average Weight Matrix Amplitudes',
'subplot': (0,1),
'plot_func': lambda self, fig: p.plot_amplitude(
np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0),
fig=fig),
'condition': lambda self: len(self.weights.shape) >= 2
},
]
}
]
plot_list = [
{'title': 'Quantum Efficiency Mask',
'plot_level': 2,
'plot_func': lambda self, fig: p.plot_amplitude(self.qe_mask, fig=fig),
'condition': lambda self: (hasattr(self, 'qe_mask') and self.qe_mask is not None)},
{'title': 'Per-Exposure Probe Intensity',
'plot_level': 3,
'figure_size': (8,5.3),
'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation(
dataset,
fig=fig,
@@ -1061,6 +1127,7 @@ class FancyPtycho(CDIModel):
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': 'Per-Exposure Probe Amplitudes',
'plot_level': 3,
'figure_size': (8,5.3),
'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation(
dataset,
fig=fig,
@@ -1070,6 +1137,7 @@ class FancyPtycho(CDIModel):
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': 'Per-Exposure Probe Phases',
'plot_level': 3,
'figure_size': (8,5.3),
'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation(
dataset,
fig=fig,
@@ -1077,29 +1145,6 @@ class FancyPtycho(CDIModel):
image_title='Probe Phases (scroll to view modes)',
image_colorbar_title='Probe Phase'),
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': 'Average Weight Matrix Amplitudes',
'plot_level': 1,
'plot_func': lambda self, fig: p.plot_amplitude(
np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0),
fig=fig),
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': '% of Power in Top Mode',
'plot_level': 3,
'plot_func': lambda self, fig, dataset: p.plot_nanomap(
self.corrected_translations(dataset),
100 * t.stack([
analysis.calc_mode_power_fractions(
self.probe.data,
weight_matrix=self.weights.data[i])[0]
for i in range(self.weights.shape[0])
], dim=0),
fig=fig,
units=self.units),
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': 'Quantum Efficiency Mask',
'plot_level': 3,
'plot_func': lambda self, fig: p.plot_amplitude(self.qe_mask, fig=fig),
'condition': lambda self: (hasattr(self, 'qe_mask') and self.qe_mask is not None)},
]
src/cdtools/tools/plotting/plotting.py
+38 -35
View File
@@ -105,6 +105,7 @@ def plot_image(
vmin=None,
vmax=None,
interpolation=None,
title=None,
**kwargs
):
"""Plots an image with a colorbar and on an appropriate spatial grid
@@ -169,12 +170,13 @@ def plot_image(
# stack of images, or the only image if only a single image has been
# given
def make_plot(idx):
#plt.figure(fig.number)
#title = plt.gca().get_title()
try:
title = fig.axes[0].get_title()
except IndexError:
title = ''
if title is not None:
ax_title = title
else:
try:
ax_title = fig.axes[0].get_title()
except IndexError:
ax_title = ''
# 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
@@ -190,7 +192,6 @@ def plot_image(
# By only updating the data, and not redrawing the fig, we
# don't "reset" the home positions of the other
if hasattr(fig, '_current_im'):
print('Just changing data')
fig._current_im.set_data(to_plot)
fig._current_im.autoscale()
# We need to go to the "home" position before updating it
@@ -200,10 +201,11 @@ def plot_image(
if fig.canvas.toolbar is not None:
fig.canvas.toolbar.home()
fig.canvas.toolbar.update()
# Replace existing mode number
for artist in fig.texts:
artist.set_text(f'Mode {fig.plot_idx}')
if len(im.shape) >= 3:
base = title if title is not None else '('.join(ax_title.split('(')[:-1])[:-1]
fig.axes[0].set_title(base + f' ({fig.plot_idx+1} of {num_images})')
return fig
fig.clear()
@@ -291,10 +293,10 @@ def plot_image(
ax.set_xlabel('j (pixels)')
ax.set_ylabel('i (pixels)')
ax.set_title(title)
if title is not None:
ax.set_title(ax_title)
if len(im.shape) >= 3:
fig.text(0.03, 0.03, f'Mode {fig.plot_idx}', fontsize=14)
ax.set_title(ax_title + f' ({fig.plot_idx+1} of {num_images})')
if fig.canvas.toolbar is not None:
fig.canvas.toolbar.update()
@@ -335,7 +337,7 @@ def plot_image(
return result_fig
def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', **kwargs):
def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', title=None, **kwargs):
"""Plots the real part of a complex array with dimensions NxM
If a figure is given explicitly, it will clear that existing figure and
@@ -369,11 +371,11 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_
plot_func = lambda x: np.real(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
title=title, **kwargs)
def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', **kwargs):
def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', title=None, **kwargs):
"""Plots the imaginary part of a complex array with dimensions NxM
If a figure is given explicitly, it will clear that existing figure and
@@ -407,10 +409,10 @@ def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_
plot_func = lambda x: np.imag(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
title=title, **kwargs)
def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Amplitude (a.u.)', **kwargs):
def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Amplitude (a.u.)', title=None, **kwargs):
"""Plots the amplitude of a complex array with dimensions NxM
If a figure is given explicitly, it will clear that existing figure and
@@ -444,7 +446,7 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis',
plot_func = lambda x: np.absolute(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
title=title, **kwargs)
def plot_phase(
@@ -456,6 +458,7 @@ def plot_phase(
cmap_label='Phase (rad)',
vmin=None,
vmax=None,
title=None,
**kwargs
):
""" Plots the phase of a complex array with dimensions NxM
@@ -506,14 +509,14 @@ def plot_phase(
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
vmin=vmin,vmax=vmax,
vmin=vmin, vmax=vmax, title=title,
**kwargs)
def plot_amplitude_surfacenorm():
pass
def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', title=None, **kwargs):
""" Plots the colorized version of a complex array with dimensions NxM
The darkness corresponds to the intensity of the image, and the color
@@ -545,7 +548,7 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
"""
plot_func = lambda x: colorize(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, show_cbar=False, **kwargs)
units=units, show_cbar=False, title=title, **kwargs)
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, clear_fig=True, label=None, color=None, marker='.', **kwargs):
@@ -713,20 +716,19 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# mode, i.e. on a figure that already has this thing showing.
if fig is None:
fig = plt.figure(figsize=(8,5.3))
fig = plt.figure(figsize=(20,4.5), constrained_layout=True)
else:
plt.figure(fig.number)
plt.gcf().clear()
fig = plt.figure(fig.number, figsize=(20,4.5), constrained_layout=True)
fig.clear()
if hasattr(fig, 'nanomap_cids'):
for cid in fig.nanomap_cids:
fig.canvas.mpl_disconnect(cid)
# Does figsize work with the fig.subplots, or just for plt.subplots?
axes = fig.subplots(1,2)
gs = fig.add_gridspec(2, 2, height_ratios=[0.9,0.1], width_ratios=[1,1])
fig.tight_layout(rect=[0.04, 0.09, 0.98, 0.96])
plt.subplots_adjust(wspace=0.25) #avoids overlap of labels with plots
axslider = plt.axes([0.15,0.06,0.75,0.03])
axes = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1])]
axslider = fig.add_subplot(gs[1, :]) # full width
# This gets the set of sizes for the points in the nanomap
def calculate_sizes(idx):
@@ -779,11 +781,12 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
axes[0].set_facecolor('k')
axes[0].set_xlabel('Translation x ('+nanomap_units+')', labelpad=1)
axes[0].set_ylabel('Translation y ('+nanomap_units+')', labelpad=1)
axes[0].set_aspect('equal')
cb1 = plt.colorbar(nanomap, ax=axes[0], orientation='horizontal',
format='%.2e',
ticks=ticker.LinearLocator(numticks=5),
pad=0.17,fraction=0.1)
cb1.ax.set_title(nanomap_colorbar_title, size="medium", pad=5)
ticks=ticker.LinearLocator(numticks=5))#,
#pad=0.17,fraction=0.1)
cb1.ax.set_title(nanomap_colorbar_title, size="medium")#, pad=5)
cb1.ax.tick_params(labelrotation=20)
if values is None:
# This seems to do a good job of leaving the appropriate space
@@ -836,8 +839,8 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',
format='%.2e',
ticks=ticker.LinearLocator(numticks=5),
pad=0.17,fraction=0.1)
ticks=ticker.LinearLocator(numticks=5))#,
#pad=-0.17)#,fraction=0.1)
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))