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

Update the tests to work better when checking the model plotting, and make sure to cover panel_plot_mode=False

Browse Source
This commit is contained in:
allevitan
2026-03-24 16:19:26 +01:00
parent 29e9fddd72
commit ad2b09f2f0
4 changed files with 64 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/tools/plotting/plotting.py
+13 -3
View File
@@ -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
tests/models/test_fancy_ptycho.py
+18 -10
View File
@@ -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
tests/models/test_simple_ptycho.py
+6 -2
View File
@@ -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
tests/test_reconstructors.py
+27 -13
View File
@@ -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])