mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-11 14:02:38 +02:00
269 lines
8.8 KiB
Python
269 lines
8.8 KiB
Python
import pytest
|
|
import time
|
|
import torch as t
|
|
from matplotlib import pyplot as plt
|
|
|
|
import cdtools
|
|
|
|
# Force all reconstructions to use the same RNG seed
|
|
t.manual_seed(0)
|
|
|
|
|
|
def test_center_probe(lab_ptycho_cxi):
|
|
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
|
|
model = cdtools.models.FancyPtycho.from_dataset(
|
|
dataset,
|
|
n_modes=3,
|
|
fourier_probe=False
|
|
)
|
|
base_probe = model.probe.detach().clone()
|
|
model.center_probes()
|
|
centered_probe = model.probe.detach().clone()
|
|
|
|
fourier_model = cdtools.models.FancyPtycho.from_dataset(
|
|
dataset,
|
|
n_modes=3,
|
|
fourier_probe=True,
|
|
)
|
|
|
|
fourier_model.probe.data = cdtools.tools.propagators.far_field(
|
|
base_probe
|
|
)
|
|
|
|
fourier_model.probe.detach().clone()
|
|
fourier_model.center_probes()
|
|
fourier_centered_probe = fourier_model.probe.detach().clone()
|
|
ifft_fourier_centered_probe = cdtools.tools.propagators.inverse_far_field(
|
|
fourier_centered_probe)
|
|
|
|
# So we know the code had to do something
|
|
assert not t.allclose(base_probe, centered_probe)
|
|
# And checking that they both do the same thing, whether or not
|
|
# fourier_probe was set to True
|
|
assert t.allclose(
|
|
centered_probe,
|
|
ifft_fourier_centered_probe,
|
|
atol=1e-4,
|
|
rtol=1e-3
|
|
)
|
|
|
|
def test_lab_ptycho_data_loading(lab_ptycho_cxi):
|
|
|
|
print('\nTesting a few unusual data loading scenarios.')
|
|
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
|
|
|
|
# Test that it will properly load an initialization for the weights
|
|
# from the intensities with OPRP on
|
|
dataset.intensities = t.rand(len(dataset))
|
|
|
|
model = cdtools.models.FancyPtycho.from_dataset(
|
|
dataset,
|
|
n_modes=4,
|
|
dm_rank=1,
|
|
)
|
|
|
|
# And test the case without OPRP
|
|
model = cdtools.models.FancyPtycho.from_dataset(
|
|
dataset,
|
|
n_modes=2,
|
|
)
|
|
|
|
|
|
|
|
@pytest.mark.slow
|
|
def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
|
|
|
|
print('\nTesting performance on the standard transmission ptycho dataset')
|
|
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
|
|
|
|
# Test the masking system
|
|
dataset.mask[110:115,65:70] = 0
|
|
dataset.patterns[...,~dataset.mask] = t.max(dataset.patterns)
|
|
|
|
model = cdtools.models.FancyPtycho.from_dataset(
|
|
dataset,
|
|
n_modes=3,
|
|
oversampling=2,
|
|
dm_rank=2,
|
|
exponentiate_obj=True,
|
|
probe_support_radius=120,
|
|
propagation_distance=5e-3,
|
|
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,
|
|
plot_level=4, # test with all plots
|
|
)
|
|
|
|
print('Running reconstruction on provided reconstruction_device,',
|
|
reconstruction_device)
|
|
model.to(device=reconstruction_device)
|
|
dataset.get_as(device=reconstruction_device)
|
|
|
|
for loss in model.Adam_optimize(50, dataset, lr=0.02, batch_size=10):
|
|
print(model.report())
|
|
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:
|
|
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:
|
|
model.inspect(dataset, min_interval=10)
|
|
|
|
model.tidy_probes()
|
|
|
|
if show_plot:
|
|
model.inspect(dataset)
|
|
model.compare(dataset)
|
|
time.sleep(3)
|
|
plt.close('all')
|
|
|
|
# Simply test that this does not fail
|
|
results = model.save_results(dataset)
|
|
|
|
# If this fails, the reconstruction has gotten worse
|
|
assert model.loss_history[-1] < 0.38
|
|
|
|
|
|
@pytest.mark.slow
|
|
def test_near_field_ptycho(near_field_ptycho_cxi, reconstruction_device, show_plot):
|
|
|
|
print('\nTesting performance on the standard transmission ptycho dataset')
|
|
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(near_field_ptycho_cxi)
|
|
|
|
model = cdtools.models.FancyPtycho.from_dataset(
|
|
dataset,
|
|
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
|
|
loss='poisson_nll',
|
|
)
|
|
|
|
print('Running reconstruction on provided reconstruction_device,',
|
|
reconstruction_device)
|
|
model.to(device=reconstruction_device)
|
|
dataset.get_as(device=reconstruction_device)
|
|
|
|
for loss in model.Adam_optimize(100, dataset, lr=0.04, batch_size=10):
|
|
print(model.report())
|
|
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:
|
|
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] < 18
|
|
|
|
|
|
def test_fancy_ptycho_from_results_dict(lab_ptycho_cxi, tmp_path):
|
|
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
|
|
|
|
t.manual_seed(42)
|
|
model = cdtools.models.FancyPtycho.from_dataset(
|
|
dataset,
|
|
n_modes=2,
|
|
)
|
|
|
|
# Verify original_translations is stored after from_dataset
|
|
assert hasattr(model, 'original_translations')
|
|
assert model.original_translations is not None
|
|
|
|
# Run a few epochs to get non-trivial state
|
|
for loss in model.Adam_optimize(5, dataset, batch_size=10):
|
|
pass
|
|
|
|
# Test from_results_dict with in-memory dict (no dataset argument needed)
|
|
results_dict = model.save_results()
|
|
loaded_model = cdtools.models.FancyPtycho.from_results_dict(results_dict)
|
|
|
|
# Test from_results_h5 via a temporary file
|
|
h5_path = str(tmp_path / 'fancy_ptycho_test.h5')
|
|
model.save_to_h5(h5_path)
|
|
loaded_model_h5 = cdtools.models.FancyPtycho.from_results_h5(h5_path)
|
|
|
|
# Verify training metadata is restored
|
|
assert loaded_model.epoch == model.epoch
|
|
assert loaded_model.loss_history == model.loss_history
|
|
|
|
# Verify original_translations round-trips correctly
|
|
assert t.allclose(
|
|
loaded_model.original_translations,
|
|
model.original_translations,
|
|
)
|
|
|
|
# Verify all parameters and buffers are restored exactly
|
|
original_sd = model.state_dict()
|
|
loaded_sd = loaded_model.state_dict()
|
|
loaded_h5_sd = loaded_model_h5.state_dict()
|
|
for key in original_sd:
|
|
assert t.allclose(original_sd[key].float(), loaded_sd[key].float()), \
|
|
f'from_results_dict: state_dict mismatch for key {key}'
|
|
assert t.allclose(original_sd[key].float(), loaded_h5_sd[key].float()), \
|
|
f'from_results_h5: state_dict mismatch for key {key}'
|
|
|
|
# Verify forward pass produces identical output
|
|
(indices, translations), patterns = dataset[:5]
|
|
with t.no_grad():
|
|
original_out = model(indices, translations)
|
|
loaded_out = loaded_model(indices, translations)
|
|
loaded_h5_out = loaded_model_h5(indices, translations)
|
|
|
|
assert t.allclose(original_out, loaded_out), \
|
|
'from_results_dict: forward pass output mismatch'
|
|
assert t.allclose(original_out, loaded_h5_out), \
|
|
'from_results_h5: forward pass output mismatch'
|
|
|
|
|
|
def test_fancy_ptycho_from_results_dict_with_missing_keys(lab_ptycho_cxi):
|
|
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
|
|
|
|
t.manual_seed(42)
|
|
model = cdtools.models.FancyPtycho.from_dataset(dataset, n_modes=1)
|
|
results_dict = model.save_results()
|
|
|
|
# Strip top-level training metadata
|
|
for key in ('loss_history', 'epoch', 'training_history'):
|
|
results_dict.pop(key, None)
|
|
|
|
# Strip defaultable state_dict keys
|
|
sd_keys_to_strip = (
|
|
'exponentiate_obj', 'phase_only', 'near_field', 'fourier_probe',
|
|
'simulate_probe_translation', 'simulate_finite_pixels',
|
|
'translation_scale', 'oversampling', 'surface_normal', 'min_translation',
|
|
)
|
|
for key in sd_keys_to_strip:
|
|
results_dict['state_dict'].pop(key, None)
|
|
|
|
loaded = cdtools.models.FancyPtycho.from_results_dict(results_dict)
|
|
|
|
assert loaded.loss_history == []
|
|
assert loaded.epoch == 0
|
|
assert loaded.training_history == ''
|
|
assert bool(loaded.exponentiate_obj) == False
|
|
assert bool(loaded.phase_only) == False
|
|
assert bool(loaded.near_field) == False
|
|
assert bool(loaded.fourier_probe) == False
|
|
assert bool(loaded.simulate_probe_translation) == False
|
|
assert bool(loaded.simulate_finite_pixels) == False
|
|
assert float(loaded.translation_scale) == 1.0
|
|
assert int(loaded.oversampling) == 1
|
|
assert t.allclose(loaded.surface_normal, t.tensor([0., 0., 1.]))
|
|
assert t.allclose(loaded.min_translation, t.tensor([0., 0.]))
|