mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 21:42:39 +02:00
149 lines
4.5 KiB
Python
149 lines
4.5 KiB
Python
import pytest
|
|
import cdtools
|
|
import torch as t
|
|
|
|
import cdtools
|
|
from matplotlib import pyplot as plt
|
|
|
|
|
|
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_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
|
|
)
|
|
|
|
@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)
|
|
|
|
model = cdtools.models.FancyPtycho.from_dataset(
|
|
dataset,
|
|
n_modes=3,
|
|
oversampling=2,
|
|
exponentiate_obj=True,
|
|
dm_rank=2,
|
|
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
|
|
)
|
|
|
|
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(70, dataset, lr=0.02, batch_size=10):
|
|
print(model.report())
|
|
if show_plot and model.epoch % 10 == 0:
|
|
model.inspect(dataset)
|
|
|
|
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)
|
|
|
|
model.tidy_probes()
|
|
|
|
if show_plot:
|
|
model.inspect(dataset)
|
|
model.compare(dataset)
|
|
|
|
# If this fails, the reconstruction has gotten worse
|
|
assert model.loss_history[-1] < 0.001
|
|
|
|
|
|
@pytest.mark.slow
|
|
def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
|
|
|
|
print('\nTesting performance on the standard gold balls dataset')
|
|
|
|
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
|
|
|
|
pad = 10
|
|
dataset.pad(pad)
|
|
|
|
model = cdtools.models.FancyPtycho.from_dataset(
|
|
dataset,
|
|
n_modes=3,
|
|
probe_support_radius=50,
|
|
propagation_distance=2e-6,
|
|
units='um',
|
|
probe_fourier_crop=pad
|
|
)
|
|
|
|
model.translation_offsets.data += \
|
|
0.7 * t.randn_like(model.translation_offsets)
|
|
|
|
# Not much probe intensity instability in this dataset, no need for this
|
|
model.weights.requires_grad = False
|
|
|
|
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(20, dataset, lr=0.005, batch_size=50):
|
|
print(model.report())
|
|
if show_plot and model.epoch % 10 == 0:
|
|
model.inspect(dataset)
|
|
|
|
for loss in model.Adam_optimize(50, dataset, lr=0.002, batch_size=100):
|
|
print(model.report())
|
|
if show_plot and model.epoch % 10 == 0:
|
|
model.inspect(dataset)
|
|
|
|
for loss in model.Adam_optimize(100, dataset, lr=0.001, batch_size=100,
|
|
schedule=True):
|
|
print(model.report())
|
|
if show_plot and model.epoch % 10 == 0:
|
|
model.inspect(dataset)
|
|
|
|
model.tidy_probes()
|
|
|
|
if show_plot:
|
|
model.inspect(dataset)
|
|
model.compare(dataset)
|
|
|
|
# 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
|
|
|
|
|