From b9dcb03a9316603ede47ec611fe56cdb06e65211 Mon Sep 17 00:00:00 2001 From: allevitan Date: Mon, 13 Apr 2026 17:10:58 +0200 Subject: [PATCH] Remove the need to use dataset for model.inspect() and model.save_results(), and update the examples accordingly. Docs and ipython notebooks not yet updated --- examples/fancy_ptycho.py | 6 +++--- examples/gold_ball_ptycho.py | 10 +++++----- examples/gold_ball_split.py | 2 +- examples/near_field_ptycho.py | 8 ++++---- examples/simple_ptycho.py | 2 +- examples/transmission_RPI.py | 8 ++++---- examples/tutorial_finale.py | 2 +- examples/tutorial_simple_ptycho.py | 2 +- src/cdtools/models/fancy_ptycho.py | 29 +++++++++++++++-------------- 9 files changed, 35 insertions(+), 34 deletions(-) diff --git a/examples/fancy_ptycho.py b/examples/fancy_ptycho.py index d6dfc74..14d75f8 100644 --- a/examples/fancy_ptycho.py +++ b/examples/fancy_ptycho.py @@ -34,7 +34,7 @@ for loss in recon.optimize(50, lr=0.02, batch_size=10): print(model.report()) # Because plotting can be expensive, setting a minimum plotting interval # (in seconds) can avoid excessive replots. - model.inspect(dataset, min_interval=10) + model.inspect(min_interval=10) # It's common to chain several different reconstruction loops. Here, we # started with an aggressive refinement to find the probe in the previous @@ -42,12 +42,12 @@ for loss in recon.optimize(50, lr=0.02, batch_size=10): # and larger minibatch for loss in recon.optimize(50, lr=0.005, batch_size=50): print(model.report()) - model.inspect(dataset, min_interval=10) + model.inspect(min_interval=10) # This orthogonalizes the recovered probe modes model.tidy_probes() # Setting replot_all will reopen any windows which were closed earlier -model.inspect(dataset, replot_all=True) +model.inspect(replot_all=True) model.compare(dataset) plt.show() diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index aeb510c..32a05a1 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -54,11 +54,11 @@ with model.save_on_exception( for loss in recon.optimize(20, lr=0.005, batch_size=50): print(model.report()) - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) for loss in recon.optimize(50, lr=0.002, batch_size=100): print(model.report()) - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) # We can often reset our guess of the probe positions once we have a # good guess of probe and object, but in this case it causes the @@ -69,14 +69,14 @@ with model.save_on_exception( # the loss fails to improve after 10 epochs for loss in recon.optimize(100, lr=0.001, batch_size=100, schedule=True): print(model.report()) - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) model.tidy_probes() # This saves the final result -model.save_to_h5('example_reconstructions/gold_balls.h5', dataset) +model.save_to_h5('example_reconstructions/gold_balls.h5') -model.inspect(dataset, replot_all=True) +model.inspect(replot_all=True) model.compare(dataset) plt.show() diff --git a/examples/gold_ball_split.py b/examples/gold_ball_split.py index fde19c6..b73b508 100644 --- a/examples/gold_ball_split.py +++ b/examples/gold_ball_split.py @@ -54,4 +54,4 @@ for label, dataset in zip(labels, datasets): model.tidy_probes() - model.save_to_h5(f'example_reconstructions/gold_balls_{label}.h5', dataset) + model.save_to_h5(f'example_reconstructions/gold_balls_{label}.h5') diff --git a/examples/near_field_ptycho.py b/examples/near_field_ptycho.py index af0012d..e446554 100644 --- a/examples/near_field_ptycho.py +++ b/examples/near_field_ptycho.py @@ -34,21 +34,21 @@ if t.cuda.is_available(): model.to(device='cuda') dataset.get_as(device='cuda') -model.inspect(dataset) +model.inspect() recon = cdtools.reconstructors.AdamReconstructor(model, dataset) for loss in recon.optimize(100, lr=0.04, batch_size=10): print(model.report()) - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) for loss in recon.optimize(50, lr=0.005, batch_size=50): print(model.report()) - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) # This orthogonalizes the recovered probe modes model.tidy_probes() -model.inspect(dataset, replot_all=True) +model.inspect(replot_all=True) model.compare(dataset) plt.show() diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index 41c0c6c..f2aec3b 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -30,7 +30,7 @@ for loss in model.Adam_optimize(100, dataset, batch_size=10): # We print a quick report of the optimization status print(model.report()) # And liveplot the updates to the model as they happen - model.inspect(dataset) + model.inspect() # We open a comparison of the simulated and measured data model.compare(dataset) diff --git a/examples/transmission_RPI.py b/examples/transmission_RPI.py index 32c02d5..c972db2 100644 --- a/examples/transmission_RPI.py +++ b/examples/transmission_RPI.py @@ -31,19 +31,19 @@ if t.cuda.is_available(): # The regularization is an L2 regularizer that empirically helps accelerate # convergence for loss in model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05]): - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) print(model.report()) # Now we use the regularizer to damp all but the top modes for loss in model.LBFGS_optimize(50, dataset, lr=0.4, regularization_factor=[0.001,0.1]): - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) print(model.report()) # Save results to an h5 file -model.save_to_h5('example_reconstructions/transmission_RPI.h5', dataset) +model.save_to_h5('example_reconstructions/transmission_RPI.h5') # Finally, we plot the results -model.inspect(dataset, replot_all=True) +model.inspect(replot_all=True) model.compare(dataset) plt.show() diff --git a/examples/tutorial_finale.py b/examples/tutorial_finale.py index 4e95ec6..b92155e 100644 --- a/examples/tutorial_finale.py +++ b/examples/tutorial_finale.py @@ -17,7 +17,7 @@ if t.cuda.is_available(): dataset.get_as(device='cuda') for loss in model.Adam_optimize(10, dataset): - model.inspect(dataset) + model.inspect() print(model.report()) model.inspect(dataset) diff --git a/examples/tutorial_simple_ptycho.py b/examples/tutorial_simple_ptycho.py index 67544d7..ef87366 100644 --- a/examples/tutorial_simple_ptycho.py +++ b/examples/tutorial_simple_ptycho.py @@ -130,7 +130,7 @@ class SimplePtycho(CDIModel): }, ] - def save_results(self, dataset): + def save_results(self, dataset=None): # This will save out everything needed to recreate the object # in the same state, but it's not the best formatted. base_results = super().save_results() diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index ebb966d..325860c 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -907,7 +907,7 @@ class FancyPtycho(CDIModel): return probe_intensities - def plot_wavefront_variation(self, dataset, fig=None, mode='amplitude', **kwargs): + def plot_wavefront_variation(self, dataset=None, fig=None, mode='amplitude', **kwargs): def get_probes(idx): basis_prs = self.probe * self.probe_support[..., :, :] prs = t.sum(self.weights[idx, :, :, None, None] * basis_prs, @@ -941,7 +941,7 @@ class FancyPtycho(CDIModel): **kwargs), - def plot_illumination_intensity(self, fig, dataset): + def plot_illumination_intensity(self, fig, dataset=None): """Plots the probe intensity nanomap. Only used to make a plot for the plot list.""" p.plot_nanomap( self.corrected_translations(dataset), @@ -956,10 +956,14 @@ class FancyPtycho(CDIModel): plt.gca().set_aspect('equal') - def plot_translations_and_originals(self, fig, dataset): + def plot_translations_and_originals(self, fig, dataset=None): """Only used to make a plot for the plot list.""" + if dataset is not None: + original_translations = dataset.translations + else: + original_translations = self.original_translations p.plot_translations( - dataset.translations, + original_translations, fig=fig, units=self.units, label='original translations', @@ -1093,7 +1097,7 @@ class FancyPtycho(CDIModel): { 'title': 'Illumination Intensity', 'subplot': (0,1), - 'plot_func': lambda self, fig, dataset: self.plot_illumination_intensity(fig, dataset), + 'plot_func': lambda self, fig: self.plot_illumination_intensity(fig), }, { 'title': 'Detector Background', @@ -1103,7 +1107,7 @@ class FancyPtycho(CDIModel): { 'title': 'Corrected Translations', 'subplot': (0,2), - 'plot_func': lambda self, fig, dataset: self.plot_translations_and_originals(fig, dataset), + 'plot_func': lambda self, fig: self.plot_translations_and_originals(fig), }, { 'title': 'Loss History', @@ -1122,8 +1126,8 @@ class FancyPtycho(CDIModel): { 'title': '% of Power in Top Mode', 'subplot': (0,0), - 'plot_func': lambda self, fig, dataset: p.plot_nanomap( - self.corrected_translations(dataset), + 'plot_func': lambda self, fig: p.plot_nanomap( + self.corrected_translations(), 100 * t.stack([ analysis.calc_mode_power_fractions( self.probe.data, @@ -1154,8 +1158,7 @@ class FancyPtycho(CDIModel): {'title': 'Per-Exposure Probe Intensity', 'plot_level': 3, 'figure_size': (8,5.3), - 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( - dataset, + 'plot_func': lambda self, fig: self.plot_wavefront_variation( fig=fig, mode='root_sum_intensity', image_title='Root Summed Probe Intensities', @@ -1164,8 +1167,7 @@ class FancyPtycho(CDIModel): {'title': 'Per-Exposure Probe Amplitudes', 'plot_level': 3, 'figure_size': (8,5.3), - 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( - dataset, + 'plot_func': lambda self, fig: self.plot_wavefront_variation( fig=fig, mode='amplitude', image_title='Probe Amplitudes (scroll to view modes)', @@ -1174,8 +1176,7 @@ class FancyPtycho(CDIModel): {'title': 'Per-Exposure Probe Phases', 'plot_level': 3, 'figure_size': (8,5.3), - 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( - dataset, + 'plot_func': lambda self, fig: self.plot_wavefront_variation( fig=fig, mode='phase', image_title='Probe Phases (scroll to view modes)',