Restructure loss coverage in reconstruction tests

- Revert test_Adam_gold_balls to default amplitude_mse
- Add new test_intensity_MSE test (gold balls + AdamReconstructor)
- Update near_field threshold to 17 (poisson_nll scale)
Thresholds for test_intensity_MSE to be tuned after GPU run.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
allevitan
2026-04-03 18:46:27 +02:00
co-authored by Claude Sonnet 4.6
parent 178e9a4e4e
commit 3ccdc24d3f
2 changed files with 34 additions and 5 deletions
+1 -1
View File
@@ -169,4 +169,4 @@ def test_near_field_ptycho(near_field_ptycho_cxi, reconstruction_device, show_pl
plt.close('all')
# If this fails, the reconstruction has gotten worse
assert model.loss_history[-1] < 3.9
assert model.loss_history[-1] < 17
+33 -4
View File
@@ -4,7 +4,6 @@ import cdtools
import torch as t
import numpy as np
import pickle
from matplotlib import pyplot as plt
from copy import deepcopy
@@ -38,11 +37,8 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
propagation_distance=2e-6,
units='um',
probe_fourier_crop=pad,
<<<<<<< HEAD
panel_plot_mode=False, # At least one check without panel plot mode
=======
loss='intensity_mse',#NOTE: Only to check that it works.
>>>>>>> f0f1577 (Add loss function coverage to slow reconstruction tests)
)
model.translation_offsets.data += 0.7 * \
@@ -129,6 +125,39 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
assert model_recon.loss_history[-1] < 0.09
@pytest.mark.slow
def test_intensity_MSE(gold_ball_cxi, reconstruction_device, show_plot):
"""
This test checks that the intensity_mse loss function works end-to-end
with the AdamReconstructor, using the Au particle dataset.
"""
print('\nTesting performance on the standard gold balls dataset ' +
'with intensity_mse loss')
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
model = cdtools.models.FancyPtycho.from_dataset(
dataset,
n_modes=1,
propagation_distance=-3e-6,
units='nm',
loss='intensity_mse',
)
model.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
recon = cdtools.reconstructors.AdamReconstructor(model=model,
dataset=dataset)
t.manual_seed(0)
for loss in recon.optimize(5, lr=.05, batch_size=10):
print(model.report())
# Threshold to be updated after running on a GPU machine
assert model.loss_history[-1] < 91
@pytest.mark.slow
def test_LBFGS_RPI(optical_data_ss_cxi,
optical_ptycho_incoherent_pickle,