From ad2b09f2f0cdedfde81e89f76c7923c2d7069d6b Mon Sep 17 00:00:00 2001 From: allevitan Date: Tue, 24 Mar 2026 16:19:26 +0100 Subject: [PATCH] Update the tests to work better when checking the model plotting, and make sure to cover panel_plot_mode=False --- src/cdtools/tools/plotting/plotting.py | 16 +++++++++-- tests/models/test_fancy_ptycho.py | 28 +++++++++++------- tests/models/test_simple_ptycho.py | 8 ++++-- tests/test_reconstructors.py | 40 +++++++++++++++++--------- 4 files changed, 64 insertions(+), 28 deletions(-) diff --git a/src/cdtools/tools/plotting/plotting.py b/src/cdtools/tools/plotting/plotting.py index 30fbbac..a4847b7 100644 --- a/src/cdtools/tools/plotting/plotting.py +++ b/src/cdtools/tools/plotting/plotting.py @@ -204,7 +204,7 @@ def plot_image( suffix = {1: 'st', 2: 'nd', 3: 'rd'} def get_suffix(n): if n % 100 not in (11, 12, 13): - suffix.get(n % 10, 'th') + return suffix.get(n % 10, 'th') else: return 'th' return f"{n}{get_suffix(n)}" @@ -1074,7 +1074,11 @@ def plot_nanomap_with_images( # 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 isinstance(translations, t.Tensor): + translations = translations.detach().cpu().numpy() + + if isinstance(values, t.Tensor): + values = values.detach().cpu().numpy() if convention.lower() != 'probe': translations = translations * -1 @@ -1082,9 +1086,13 @@ def plot_nanomap_with_images( s = calculate_sizes(0) nanomap_units_factor = get_units_factor(nanomap_units) + + # Suppresses a warning from ax.scatter() + if values is None: + cmap = None nanomap = axes[0].scatter(nanomap_units_factor * translations[:,0], nanomap_units_factor * translations[:,1], - s=s,c=values, picker=True, cmap=cmap) + s=s, c=values, picker=True, cmap=cmap) axes[0].invert_xaxis() axes[0].set_facecolor('k') @@ -1101,7 +1109,9 @@ def plot_nanomap_with_images( # This seems to do a good job of leaving the appropriate space # where the colorbar should have been to avoid stretching the # nanomap plot, while still not showing the (now useless) colorbar. + pos = axes[0].get_position() cb1.remove() + axes[0].set_position(pos) # Now we set up the second plot, which shows the individual # diffraction patterns diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py index db05a75..4de9f4a 100644 --- a/tests/models/test_fancy_ptycho.py +++ b/tests/models/test_fancy_ptycho.py @@ -1,5 +1,7 @@ import pytest +import time import torch as t +from matplotlib import pyplot as plt import cdtools @@ -67,6 +69,7 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot): units='mm', obj_view_crop=-50, use_qe_mask=True, # test this in the case where no qe mask is defined + panel_plot_mode=True, # test with panel plot mode ) print('Running reconstruction on provided reconstruction_device,', @@ -76,24 +79,26 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot): for loss in model.Adam_optimize(50, dataset, lr=0.02, batch_size=10): print(model.report()) - if show_plot and model.epoch % 10 == 0: - model.inspect(dataset) + if show_plot: + model.inspect(dataset, min_interval=10) for loss in model.Adam_optimize(50, dataset, lr=0.005, batch_size=50): print(model.report()) - if show_plot and model.epoch % 10 == 0: - model.inspect(dataset) + if show_plot: + model.inspect(dataset, min_interval=10) for loss in model.Adam_optimize(25, dataset, lr=0.001, batch_size=50): print(model.report()) - if show_plot and model.epoch % 10 == 0: - model.inspect(dataset) + if show_plot: + model.inspect(dataset, min_interval=10) model.tidy_probes() if show_plot: model.inspect(dataset) model.compare(dataset) + time.sleep(3) + plt.close('all') # If this fails, the reconstruction has gotten worse assert model.loss_history[-1] < 0.0013 @@ -110,6 +115,7 @@ def test_near_field_ptycho(near_field_ptycho_cxi, reconstruction_device, show_pl n_modes=1, near_field=True, propagation_distance=3.65e-3, # 3.65 downstream from focus + panel_plot_mode=False, # test without panel plot mode ) print('Running reconstruction on provided reconstruction_device,', @@ -119,19 +125,21 @@ def test_near_field_ptycho(near_field_ptycho_cxi, reconstruction_device, show_pl for loss in model.Adam_optimize(100, dataset, lr=0.04, batch_size=10): print(model.report()) - if show_plot and model.epoch % 10 == 0: - model.inspect(dataset) + if show_plot: + model.inspect(dataset, min_interval=10) for loss in model.Adam_optimize(50, dataset, lr=0.005, batch_size=50): print(model.report()) - if show_plot and model.epoch % 10 == 0: - model.inspect(dataset) + if show_plot: + model.inspect(dataset, min_interval=10) model.tidy_probes() if show_plot: model.inspect(dataset) model.compare(dataset) + time.sleep(3) + plt.close('all') # If this fails, the reconstruction has gotten worse assert model.loss_history[-1] < 0.005 diff --git a/tests/models/test_simple_ptycho.py b/tests/models/test_simple_ptycho.py index b6b1868..225e463 100644 --- a/tests/models/test_simple_ptycho.py +++ b/tests/models/test_simple_ptycho.py @@ -1,5 +1,7 @@ import pytest +import time import torch as t +from matplotlib import pyplot as plt import cdtools @@ -18,12 +20,14 @@ def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot): for loss in model.Adam_optimize(100, dataset, batch_size=10): print(model.report()) - if show_plot and model.epoch % 10 == 0: - model.inspect(dataset) + if show_plot: + model.inspect(dataset, min_interval=10) if show_plot: model.inspect(dataset) model.compare(dataset) + time.sleep(3) + plt.close('all') # If this fails, the reconstruction got worse assert model.loss_history[-1] < 0.013 diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py index 132a8c2..3f9245b 100644 --- a/tests/test_reconstructors.py +++ b/tests/test_reconstructors.py @@ -1,4 +1,5 @@ import pytest +import time import cdtools import torch as t import numpy as np @@ -36,7 +37,8 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): probe_support_radius=50, propagation_distance=2e-6, units='um', - probe_fourier_crop=pad + probe_fourier_crop=pad, + panel_plot_mode=False, # At least one check without panel plot mode ) model.translation_offsets.data += 0.7 * \ @@ -67,8 +69,8 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): lr=lr_tup[i], batch_size=batch_size_tup[i]): print(model_recon.report()) - if show_plot and model_recon.epoch % 10 == 0: - model_recon.inspect(dataset) + if show_plot: + model_recon.inspect(dataset, min_interval=10) # Check hyperparameter update assert recon.optimizer.param_groups[0]['lr'] == lr_tup[i] @@ -86,6 +88,8 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): if show_plot: model_recon.inspect(dataset) model_recon.compare(dataset) + time.sleep(3) + plt.close('all') # ******* Reconstructions with CDIModel.Adam_optimize ******* print('Running reconstruction using CDIModel.Adam_optimize on provided' + @@ -99,14 +103,16 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): lr=lr_tup[i], batch_size=batch_size_tup[i]): print(model.report()) - if show_plot and model.epoch % 10 == 0: - model.inspect(dataset) + if show_plot: + model.inspect(dataset, min_interval=10) model.tidy_probes() if show_plot: model.inspect(dataset) model.compare(dataset) + time.sleep(3) + plt.close('all') # Ensure equivalency between the model reconstructions during the first # pass, where they should be identical @@ -170,8 +176,8 @@ def test_LBFGS_RPI(optical_data_ss_cxi, for loss in recon.optimize(iterations, lr=0.4, regularization_factor=reg_factor_tup[i]): - if show_plot and i == 0: - model_recon.inspect(dataset) + if show_plot: + model_recon.inspect(dataset, min_interval=10) print(model_recon.report()) # Check hyperparameter update (or lack thereof) @@ -180,6 +186,8 @@ def test_LBFGS_RPI(optical_data_ss_cxi, if show_plot: model_recon.inspect(dataset) model_recon.compare(dataset) + time.sleep(3) + plt.close('all') # Check model pointing assert id(model_recon) == id(recon.model) @@ -193,13 +201,15 @@ def test_LBFGS_RPI(optical_data_ss_cxi, dataset, lr=0.4, regularization_factor=reg_factor_tup[i]): # noqa - if show_plot and i == 0: - model.inspect(dataset) + if show_plot: + model.inspect(dataset, min_interval=10) print(model.report()) if show_plot: model.inspect(dataset) model.compare(dataset) + time.sleep(3) + plt.close('all') # Check loss equivalency between the two reconstructions assert np.allclose(model.loss_history[:epoch_tup[0]], model_recon.loss_history[:epoch_tup[0]]) @@ -271,8 +281,8 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): lr=lr, batch_size=batch_size): print(model_recon.report()) - if show_plot and model_recon.epoch % 10 == 0: - model_recon.inspect(dataset) + if show_plot: + model_recon.inspect(dataset, min_interval=10) # Check hyperparameter update assert recon.optimizer.param_groups[0]['lr'] == lr @@ -290,6 +300,8 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): if show_plot: model_recon.inspect(dataset) model_recon.compare(dataset) + time.sleep(3) + plt.close('all') # ******* Reconstructions with cdtools.CDIModel.SGD_optimize ******* print('Running reconstruction using CDIModel.SGD_optimize on provided' + @@ -301,14 +313,16 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): lr=lr, batch_size=batch_size): print(model.report()) - if show_plot and model.epoch % 10 == 0: - model.inspect(dataset) + if show_plot: + model.inspect(dataset, min_interval=10) model.tidy_probes() if show_plot: model.inspect(dataset) model.compare(dataset) + time.sleep(3) + plt.close('all') # Ensure equivalency between the model reconstructions assert np.allclose(model_recon.loss_history[-1], model.loss_history[-1])