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

A first prototype from claude to test on Ra

Browse Source
This commit is contained in:
allevitan
2026-03-20 09:47:38 +01:00
parent ca3c48e83d
commit f3581b88db
4 changed files with 329 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/models/base.py
+215 -70
View File
@@ -58,12 +58,14 @@ class CDIModel(t.nn.Module):
functions.
"""
def __init__(self):
def __init__(self, panel_plot_mode=False, plot_level=np.inf):
super(CDIModel, self).__init__()
self.loss_history = []
self.training_history = ''
self.epoch = 0
self.panel_plot_mode = panel_plot_mode
self.plot_level = plot_level
def from_dataset(self, dataset):
raise NotImplementedError()
@@ -547,96 +549,238 @@ class CDIModel(t.nn.Module):
return msg
# By default, the plot_list is empty
# By default, the plot lists are empty
plot_panel_list = []
plot_list = []
def inspect(self, dataset=None, update=True):
"""Plots all the plots defined in the model's plot_list attribute
def inspect(self, dataset=None, replot_all=False):
"""Plots all the plots defined in the model's plot_panel_list and plot_list attributes
If update is set to True, it will update any previously plotted set
of plots, if one exists, and then redraw them. Otherwise, it will
plot a new set, and any subsequent updates will update the new set
Updates any previously plotted figures that are still open. Figures
that have been closed are left closed unless replot_all=True.
Optionally, a dataset can be passed, which will allow plotting of any
registered plots which need to incorporate some information from
the dataset (such as geometry or a comparison with measured data).
Plots can be registered in any subclass by defining the plot_list
attribute. This should be a list of tuples in the following format:
( 'Plot Title', function_to_generate_plot(self),
function_to_determine_whether_to_plot(self))
Plots can be registered in any subclass by defining plot_panel_list
and/or plot_list class attributes. See the CDIModel documentation for
the expected dict-based format of each.
Where the third element in the tuple (a function that returns
True if the plot is relevant) is not required.
When panel_plot_mode=True (set in __init__), plot_panel_list entries
are rendered as multi-subplot figures. When False (the default),
each subplot in plot_panel_list is rendered as its own figure,
prepended to any standalone plot_list entries.
The plot_level attribute (set in __init__, default np.inf) controls
which plots are shown: a panel or standalone plot is only shown when
its plot_level <= self.plot_level.
Parameters
----------
dataset : CDataset
Optional, a dataset matched to the model type
update : bool, default: True
Whether to update existing plots or plot new ones
replot_all : bool, default: False
If True, recreate figures that were previously closed by the user.
"""
# We find or create all the figures
first_update = False
if update and hasattr(self, 'figs') and self.figs:
figs = self.figs
elif update:
figs = None
self.figs = []
first_update = True
plot_panel_list = getattr(self, 'plot_panel_list', None) or []
plot_list = getattr(self, 'plot_list', None) or []
if self.panel_plot_mode and plot_panel_list:
self._inspect_panel(dataset=dataset, replot_all=replot_all)
else:
figs = None
self.figs = []
# Flatten plot_panel_list, assigning each subplot the panel's plot_level,
# then prepend to plot_list
flat = []
for panel in plot_panel_list:
panel_level = panel.get('plot_level', 0)
for plot in panel['plots']:
flat.append({**plot, 'plot_level': panel_level})
all_plots = flat + list(plot_list)
if not hasattr(self, '_flat_fig_map'):
self._flat_fig_map = {}
self.figs = self._do_inspect(all_plots, self._flat_fig_map,
dataset=dataset,
replot_all=replot_all)
plt.pause(0.05)
def _do_inspect(self, plot_list, fig_map, dataset=None, replot_all=False):
"""Core one-figure-per-plot rendering logic.
fig_map is a dict {title: figure} owned by the caller and updated
in-place. It tracks which figures are open across calls.
Behaviour:
replot_all=False — closed figures are skipped (left closed).
replot_all=True — closed figures are recreated.
Returns the list of figures that were rendered this call.
"""
if not plot_list:
return []
rendered = []
for plot in plot_list:
# Level filter
if plot.get('plot_level', 0) > self.plot_level:
continue
# Condition check
condition = plot.get('condition', None)
if condition is not None:
try:
if not condition(self):
continue
except TypeError:
if not condition(self, dataset):
continue
title = plot['title']
fig = fig_map.get(title)
if fig is not None and not plt.fignum_exists(fig.number):
# Figure was closed by the user
if replot_all:
fig = None
del fig_map[title]
else:
continue # leave it closed
if fig is None:
fig = plt.figure(num=title)
fig._panel_label = title
fig_map[title] = fig
idx = 0
for plots in self.plot_list:
# If a conditional is included in the plot, we check whether
# it is True
try:
if len(plots) >=3 and not plots[2](self):
continue
except TypeError as e:
if len(plots) >= 3 and not plots[2](self, dataset):
continue
name = plots[0]
plotter = plots[1]
if figs is None:
fig = plt.figure()
self.figs.append(fig)
else:
fig = figs[idx]
try: # We try just plotting using the simplest allowed signature
plotter(self,fig)
plt.title(name)
except TypeError as e:
# TypeError implies it wanted another argument, i.e. a dataset
plot['plot_func'](self, fig)
plt.title(title)
except TypeError:
if dataset is not None:
try:
plotter(self, fig, dataset)
plt.title(name)
except Exception as e: # Don't raise errors: it's just plots
plot['plot_func'](self, fig, dataset)
plt.title(title)
except Exception:
pass
except Exception as e: # Don't raise errors, it's just a plot
except Exception:
pass
idx += 1
rendered.append(fig)
try:
fig.canvas.draw_idle()
except Exception:
pass
if update:
# This seems to update the figure without blocking.
plt.draw()
fig.canvas.start_event_loop(0.001)
return rendered
def _inspect_panel(self, dataset=None, replot_all=False):
"""Multi-subplot panel rendering.
Creates one figure per plot_panel_list entry, placing each subplot's
plot_func output into the appropriate axes. Closed panels stay closed
on subsequent calls unless replot_all=True. Standalone plot_list
entries are then rendered via _do_inspect and appended to self.figs.
"""
plot_panel_list = getattr(self, 'plot_panel_list', None) or []
plot_list = getattr(self, 'plot_list', None) or []
n_panels = len(plot_panel_list)
# _panel_figs: list of figures (or None if never created / closed).
# _panel_axes: dict keyed by (panel_idx, row, col) → Axes.
# _standalone_fig_map: dict {title: figure} for standalone plot_list.
first_call = not hasattr(self, '_panel_figs')
if first_call:
self._panel_figs = [None] * n_panels
self._panel_axes = {}
self._standalone_fig_map = {}
if not hasattr(self, '_standalone_fig_map'):
self._standalone_fig_map = {}
for panel_idx, panel_def in enumerate(plot_panel_list):
panel_level = panel_def.get('plot_level', 0)
if panel_level > self.plot_level:
continue # skip entire panel
nrows, ncols = panel_def['grid']
figsize = panel_def.get('figure_size', None)
title = panel_def.get('title', '')
fig = self._panel_figs[panel_idx]
# Detect if a previously open figure was closed by the user.
if fig is not None and not plt.fignum_exists(fig.number):
self._panel_figs[panel_idx] = None
for k in [k for k in self._panel_axes if k[0] == panel_idx]:
del self._panel_axes[k]
fig = None
if fig is None:
if not first_call and not replot_all:
continue # was closed; leave it closed
fig = plt.figure(num=title, figsize=figsize)
fig._panel_label = title
self._panel_figs[panel_idx] = fig
else:
# Remove all axes and recreate them fresh each update.
# plt.colorbar() shrinks the parent axes to make room for
# itself, so clearing and recreating is simpler than trying
# to undo that resizing.
for ax in list(fig.axes):
ax.remove()
for k in [k for k in self._panel_axes if k[0] == panel_idx]:
del self._panel_axes[k]
for plot in panel_def['plots']:
condition = plot.get('condition', None)
if condition is not None:
try:
if not condition(self):
continue
except TypeError:
if not condition(self, dataset):
continue
row, col = plot['subplot']
position = row * ncols + col + 1 # 1-indexed for matplotlib
ax_key = (panel_idx, row, col)
ax = fig.add_subplot(nrows, ncols, position)
self._panel_axes[ax_key] = ax
try:
plot['plot_func'](self, ax)
ax.set_title(plot['title'])
except TypeError:
if dataset is not None:
try:
plot['plot_func'](self, ax, dataset)
ax.set_title(plot['title'])
except Exception:
pass
except Exception:
pass
try:
fig.canvas.draw_idle()
except Exception:
pass
# Rebuild self.figs from open panel figures + rendered standalone figures.
panel_figs = [f for f in self._panel_figs if f is not None]
standalone_rendered = self._do_inspect(
list(plot_list), self._standalone_fig_map,
dataset=dataset, replot_all=replot_all,
)
self.figs = panel_figs + standalone_rendered
if first_update:
# But this is needed the first time the figures update, or
# they won't get drawn at all
plt.pause(0.05 * len(self.figs))
def save_figures(self, prefix='', extension='.pdf'):
@@ -661,14 +805,15 @@ class CDIModel(t.nn.Module):
Default is .eps, the file extension to save with.
"""
if hasattr(self, 'figs') and self.figs:
figs = self.figs
else:
return # No figures to save
if not (hasattr(self, 'figs') and self.figs):
return # No figures to save
for fig in self.figs:
fig.savefig(prefix + fig.axes[0].get_title() + extension,
bbox_inches = 'tight')
if hasattr(fig, '_panel_label') and fig._panel_label:
label = fig._panel_label
else:
label = fig.axes[0].get_title() if fig.axes else 'figure'
fig.savefig(prefix + label + extension, bbox_inches='tight')
def compare(self, dataset, logarithmic=False):
src/cdtools/models/simple_ptycho.py
+64 -12
View File
@@ -15,10 +15,15 @@ class SimplePtycho(CDIModel):
probe_guess,
obj_guess,
min_translation = [0,0],
panel_plot_mode=False,
plot_level=0,
):
# We initialize the superclass
super().__init__()
super().__init__(
panel_plot_mode=panel_plot_mode,
plot_level=plot_level,
)
# We register all the constants, like wavelength, as buffers. This
# lets the model hook into some nice pytorch features, like using
@@ -43,7 +48,8 @@ class SimplePtycho(CDIModel):
@classmethod
def from_dataset(cls, dataset):
def from_dataset(cls, dataset,panel_plot_mode=False,
plot_level=0, ):
# We get the key geometry information from the dataset
wavelength = dataset.wavelength
@@ -76,7 +82,9 @@ class SimplePtycho(CDIModel):
probe_basis,
probe,
obj,
min_translation=min_translation
min_translation=min_translation,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level,
)
@@ -107,15 +115,59 @@ class SimplePtycho(CDIModel):
# This lists all the plots to display on a call to model.inspect()
plot_list = [
('Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)),
('Probe Phase',
lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis)),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
('Object Phase',
lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis))
plot_panel_list = [
{
# Title for window
'title' : 'Probe Results',
# (width, height) in inches
'figure_size': (7, 7),
# (nrows, ncols) for subplot grid
'grid': (2, 2),
# A setting to control how many plots are produced
'plot_level': 1,
# The list of plots to include
'plots' : [
{
'title': 'Probe Amplitude',
'subplot' : (0, 0),
'plot_func': lambda self, fig:
p.plot_amplitude(self.probe, fig,
basis=self.probe_basis)
},{
'title': 'Probe Phase',
'subplot' : (0, 1),
'plot_func': lambda self, fig:
p.plot_phase(self.probe, fig,
basis=self.probe_basis)
}
]
},
{
# Title for window
'title' : 'Object Results',
# (width, height) in inches
'figure_size': (7, 7),
# (nrows, ncols) for subplot grid
'grid': (2, 2),
# A setting to control how many plots are produced
'plot_level': 1,
# The list of plots to include
'plots' : [
{
'title': 'Object Amplitude',
'subplot' : (1, 0),
'plot_func': lambda self, fig:
p.plot_amplitude(self.obj, fig,
basis=self.probe_basis)
}, {
'title': 'Object Phase',
'subplot' : (1, 1),
'plot_func': lambda self, fig:
p.plot_phase(self.obj, fig,
basis=self.probe_basis)
},
]
},
]
def save_results(self, dataset):
src/cdtools/reconstructors/base.py
+12 -3
View File
@@ -16,6 +16,7 @@ from torch.utils import data as td
import threading
import queue
import time
from matplotlib import pyplot as plt
from typing import List, Union
if TYPE_CHECKING:
@@ -357,10 +358,18 @@ class Reconstructor:
try:
calc.start()
while calc.is_alive():
if hasattr(self.model, 'figs'):
self.model.figs[0].canvas.start_event_loop(0.01)
figs = getattr(self.model, 'figs', [])
open_fig = next(
(f for f in figs if plt.fignum_exists(f.number)),
None,
)
if open_fig is not None:
try:
open_fig.canvas.start_event_loop(0.01)
except Exception:
time.sleep(0.01)
else:
calc.join()
time.sleep(0.01)
except KeyboardInterrupt as e:
stop_event.set()
src/cdtools/tools/plotting/plotting.py
+38 -13
View File
@@ -164,7 +164,12 @@ def plot_image(
else:
im = im.detach().cpu().numpy()
if fig is None:
# Support passing an Axes object instead of a Figure
ax_mode = isinstance(fig, plt.Axes)
if ax_mode:
ax = fig
fig = ax.get_figure()
elif fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
@@ -173,8 +178,13 @@ def plot_image(
# given
def make_plot(idx):
plt.figure(fig.number)
title = plt.gca().get_title()
fig.clear()
if ax_mode:
title = ax.get_title()
ax.cla()
plt.sca(ax)
else:
title = plt.gca().get_title()
fig.clear()
# If im only has two dimensions, this reshape will add a leading
@@ -184,9 +194,10 @@ def plot_image(
s = im.shape
reshaped_im = im.reshape(-1,s[-2],s[-1])
num_images = reshaped_im.shape[0]
fig.plot_idx = idx % num_images
plot_holder = ax if ax_mode else fig
plot_holder.plot_idx = idx % num_images
to_plot = plot_func(reshaped_im[fig.plot_idx])
to_plot = plot_func(reshaped_im[plot_holder.plot_idx])
mpl_im = plt.imshow(
to_plot,
@@ -273,11 +284,13 @@ def plot_image(
plt.title(title)
if len(im.shape) >= 3:
plt.text(0.03, 0.03, str(fig.plot_idx), fontsize=14, transform=plt.gcf().transFigure)
text_transform = ax.transAxes if ax_mode else plt.gcf().transFigure
plt.text(0.03, 0.03, str(plot_holder.plot_idx), fontsize=14, transform=text_transform)
return fig
if hasattr(fig, 'plot_idx'):
result = make_plot(fig.plot_idx)
plot_holder = ax if ax_mode else fig
if hasattr(plot_holder, 'plot_idx'):
result = make_plot(plot_holder.plot_idx)
else:
result = make_plot(0)
@@ -285,15 +298,16 @@ def plot_image(
def on_action(event):
plot_holder = ax if ax_mode else fig
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)
update(plot_holder.plot_idx - 1)
elif event.key == 'down' or event.button == 'down':
update(fig.plot_idx + 1)
update(plot_holder.plot_idx + 1)
plt.draw()
if len(im.shape) >=3:
@@ -557,7 +571,11 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, inver
factor = get_units_factor(units)
if fig is None:
if isinstance(fig, plt.Axes):
ax = fig
fig = ax.get_figure()
plt.sca(ax)
elif fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
else:
@@ -614,7 +632,13 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
The figure object that was actually plotted to.
"""
if fig is None:
ax_mode = isinstance(fig, plt.Axes)
if ax_mode:
ax = fig
fig = ax.get_figure()
ax.cla()
plt.sca(ax)
elif fig is None:
fig = plt.figure()
else:
plt.figure(fig.number)
@@ -622,7 +646,8 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
factor = get_units_factor(units)
bbox = fig.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
plot_area = ax if ax_mode else fig
bbox = plot_area.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
if isinstance(translations, t.Tensor):
trans = translations.detach().cpu().numpy()
else: