From 25c7b66c7c5f56de7221510ec36f34fa71e75331 Mon Sep 17 00:00:00 2001 From: allevitan Date: Tue, 24 Mar 2026 11:10:23 +0100 Subject: [PATCH] Do a final review of all the examples to ensure they work and are well structured --- .flake8 | 2 +- examples/gold_ball_split.py | 6 +++--- examples/gold_ball_synthesize.py | 6 +++--- examples/near_field_ptycho.py | 1 - examples/simple_ptycho.py | 2 +- examples/tutorial_finale.py | 6 ++++-- examples/tutorial_simple_ptycho.py | 28 ++++++++++++++++++++-------- pyproject.toml | 3 ++- src/cdtools/models/base.py | 11 +++++++++++ src/cdtools/models/fancy_ptycho.py | 10 ++++++++++ src/cdtools/models/simple_ptycho.py | 25 ++++++++++--------------- 11 files changed, 65 insertions(+), 35 deletions(-) diff --git a/.flake8 b/.flake8 index 9214cb9..3e08f47 100644 --- a/.flake8 +++ b/.flake8 @@ -1,2 +1,2 @@ [flake8] -ignore = E501, W503 \ No newline at end of file +ignore = E501, W503, E731 \ No newline at end of file diff --git a/examples/gold_ball_split.py b/examples/gold_ball_split.py index 9fc5b08..fde19c6 100644 --- a/examples/gold_ball_split.py +++ b/examples/gold_ball_split.py @@ -32,9 +32,9 @@ for label, dataset in zip(labels, datasets): model.weights.requires_grad = False - device = 'cuda' - model.to(device=device) - dataset.get_as(device=device) + if t.cuda.is_available(): + model.to(device='cuda') + dataset.get_as(device='cuda') # Create the reconstructor recon = cdtools.reconstructors.AdamReconstructor(model, dataset) diff --git a/examples/gold_ball_synthesize.py b/examples/gold_ball_synthesize.py index 3a2c68a..4131905 100644 --- a/examples/gold_ball_synthesize.py +++ b/examples/gold_ball_synthesize.py @@ -5,11 +5,11 @@ import numpy as np # We load all three reconstructions half_1 = cdtools.tools.data.h5_to_nested_dict( - f'example_reconstructions/gold_balls_half_1.h5') + 'example_reconstructions/gold_balls_half_1.h5') half_2 = cdtools.tools.data.h5_to_nested_dict( - f'example_reconstructions/gold_balls_half_2.h5') + 'example_reconstructions/gold_balls_half_2.h5') full = cdtools.tools.data.h5_to_nested_dict( - f'example_reconstructions/gold_balls_full.h5') + 'example_reconstructions/gold_balls_full.h5') # This defines the region of recovered object to use for the analysis. pad = 260 diff --git a/examples/near_field_ptycho.py b/examples/near_field_ptycho.py index b574813..4534ee5 100644 --- a/examples/near_field_ptycho.py +++ b/examples/near_field_ptycho.py @@ -6,7 +6,6 @@ filename = 'example_data/PETRAIII_P25_Near_Field_Ptycho.cxi' dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename) dataset.inspect() -plt.show() # Setting near_field equal to True uses an angular spectrum propagator in # lieu of the default Fourier-transform propagator for far-field ptychography. diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index af8aeb1..8218a22 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -24,7 +24,7 @@ if t.cuda.is_available(): dataset.get_as(device='cuda') model.inspect(dataset) -print('hi') + # We run the reconstruction for loss in model.Adam_optimize(30, dataset, batch_size=10): # We print a quick report of the optimization status diff --git a/examples/tutorial_finale.py b/examples/tutorial_finale.py index 9f698b7..4e95ec6 100644 --- a/examples/tutorial_finale.py +++ b/examples/tutorial_finale.py @@ -1,5 +1,6 @@ from tutorial_basic_ptycho_dataset import BasicPtychoDataset from tutorial_simple_ptycho import SimplePtycho +import torch as t from h5py import File from matplotlib import pyplot as plt @@ -11,8 +12,9 @@ dataset.inspect() model = SimplePtycho.from_dataset(dataset) -model.to(device='cuda') -dataset.get_as(device='cuda') +if t.cuda.is_available(): + model.to(device='cuda') + dataset.get_as(device='cuda') for loss in model.Adam_optimize(10, dataset): model.inspect(dataset) diff --git a/examples/tutorial_simple_ptycho.py b/examples/tutorial_simple_ptycho.py index 7e8a3e5..67544d7 100644 --- a/examples/tutorial_simple_ptycho.py +++ b/examples/tutorial_simple_ptycho.py @@ -108,14 +108,26 @@ 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)) + { + 'title': 'Probe Amplitude', + 'plot_func': lambda self, fig: + p.plot_amplitude(self.probe, fig, basis=self.probe_basis), + }, + { + 'title': 'Probe Phase', + 'plot_func': lambda self, fig: + p.plot_phase(self.probe, fig, basis=self.probe_basis) + }, + { + 'title': 'Object Amplitude', + 'plot_func': lambda self, fig: + p.plot_amplitude(self.obj, fig, basis=self.probe_basis) + }, + { + 'title': 'Object Phase', + 'plot_func': lambda self, fig: + p.plot_phase(self.obj, fig, basis=self.probe_basis) + }, ] def save_results(self, dataset): diff --git a/pyproject.toml b/pyproject.toml index 9ba1438..aacaecf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,4 +50,5 @@ docs = [ ] [tool.ruff] -line-length = 79 \ No newline at end of file +line-length = 79 +ignore = ["E501", "E731"] \ No newline at end of file diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index d3cb122..e751659 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -61,6 +61,16 @@ class CDIModel(t.nn.Module): """ def __init__(self, panel_plot_mode=False, plot_level=np.inf): + """Initializes the CDIModel base class. + + Parameters + ---------- + panel_plot_mode : bool, default: False + If True, plot_panel_list entries are rendered as multi-subplot + figures. If False, each subplot is rendered as its own figure. + plot_level : float, default: np.inf + Only plots whose plot_level <= this value are shown. + """ super(CDIModel, self).__init__() self.loss_history = [] @@ -640,6 +650,7 @@ class CDIModel(t.nn.Module): def _is_backend_interactive( self ): + """Returns True if the current matplotlib backend is interactive.""" backend = matplotlib.get_backend().lower() try: # matplotlib >= 3.9 diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 3edb9ac..b363591 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -848,6 +848,15 @@ class FancyPtycho(CDIModel): def get_probe_intensities(self): + """Returns the effective probe intensity at each scan position. + + Handles both the simple (1D weights) and OPRP (2D weights) cases. + + Returns + ------- + probe_intensities : np.ndarray + Array of probe intensities, one per scan position. + """ if not hasattr(self, 'weights'): raise NotImplementedError( "I don't know how to handle having no weights") @@ -920,6 +929,7 @@ class FancyPtycho(CDIModel): def plot_illumination_intensity(self, fig, dataset): + """Plots the probe intensity nanomap. Only used to make a plot for the plot list.""" p.plot_nanomap( self.corrected_translations(dataset), self.get_probe_intensities(), diff --git a/src/cdtools/models/simple_ptycho.py b/src/cdtools/models/simple_ptycho.py index 234ed5d..734ba20 100644 --- a/src/cdtools/models/simple_ptycho.py +++ b/src/cdtools/models/simple_ptycho.py @@ -110,28 +110,23 @@ class SimplePtycho(CDIModel): plot_list = [ { 'title': 'Probe Amplitude', - 'subplot' : (0, 0), 'plot_func': lambda self, fig: - p.plot_amplitude(self.probe, fig, - basis=self.probe_basis), - }, { + 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) - }, { + p.plot_phase(self.probe, fig, basis=self.probe_basis) + }, + { 'title': 'Object Amplitude', - 'subplot' : (1, 0), 'plot_func': lambda self, fig: - p.plot_amplitude(self.obj, fig, - basis=self.probe_basis) - }, { + 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) + p.plot_phase(self.obj, fig, basis=self.probe_basis) }, ]