Get the tests passing again after the change of the optimization functions in the model classes

This commit is contained in:
2025-10-16 15:58:28 +02:00
parent 8328b852d3
commit f022b5e2c4
2 changed files with 16 additions and 9 deletions
+7 -2
View File
@@ -56,8 +56,8 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
dataset,
n_modes=3,
oversampling=2,
exponentiate_obj=True,
dm_rank=2,
exponentiate_obj=True,
probe_support_radius=120,
propagation_distance=5e-3,
units='mm',
@@ -70,7 +70,7 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
model.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
for loss in model.Adam_optimize(70, dataset, lr=0.02, batch_size=10):
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)
@@ -79,6 +79,11 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
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)
model.tidy_probes()
+9 -7
View File
@@ -91,7 +91,8 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
' reconstruction_device,', reconstruction_device)
t.manual_seed(0)
for i, iterations in enumerate(epoch_tup):
# We only need to test the first loop to ensure it's identical
for i, iterations in enumerate(epoch_tup[:1]):
for loss in model.Adam_optimize(iterations,
dataset,
lr=lr_tup[i],
@@ -106,14 +107,15 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
model.inspect(dataset)
model.compare(dataset)
# Ensure equivalency between the model reconstructions
assert np.allclose(model_recon.loss_history[-1], model.loss_history[-1])
# Ensure equivalency between the model reconstructions during the first
# pass, where they should be identical
assert np.allclose(model_recon.loss_history[:epoch_tup[0]], model.loss_history[:epoch_tup[0]])
# Ensure reconstructions have reached a certain loss tolerance. This just
# comes from running a reconstruction when it was working well and
# choosing a rough value. If it triggers this assertion error, something
# changed to make the final quality worse!
assert model.loss_history[-1] < 0.0001
assert model_recon.loss_history[-1] < 0.0001
@pytest.mark.slow
@@ -184,7 +186,7 @@ def test_LBFGS_RPI(optical_data_ss_cxi,
print('Running reconstruction using CDIModel.LBFGS_optimize.' +
'optimize on provided reconstruction_device,', reconstruction_device)
t.manual_seed(0)
for i, iterations in enumerate(epoch_tup):
for i, iterations in enumerate(epoch_tup[:1]):
for loss in model.LBFGS_optimize(iterations,
dataset,
lr=0.4,
@@ -198,12 +200,12 @@ def test_LBFGS_RPI(optical_data_ss_cxi,
model.compare(dataset)
# Check loss equivalency between the two reconstructions
assert np.allclose(model.loss_history[-1], model_recon.loss_history[-1])
assert np.allclose(model.loss_history[:epoch_tup[0]], model_recon.loss_history[:epoch_tup[0]])
# The final loss when testing this was 2.28607e-3. Based on this, we set
# a threshold of 2.3e-3 for the tested loss. If this value has been
# exceeded, the reconstructions have gotten worse.
assert model.loss_history[-1] < 0.0023
assert model_recon.loss_history[-1] < 0.0023
@pytest.mark.slow