From 30a989fe3eac36e5c10812dc988c4f05537dfd9d Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Fri, 30 Aug 2024 16:10:26 +0200 Subject: [PATCH] Add first tests for models, and replace t.tensor with t.as_tensor in models --- src/cdtools/models/base.py | 6 +- src/cdtools/models/bragg_2d_ptycho.py | 42 +++---- src/cdtools/models/fancy_ptycho.py | 46 ++++---- src/cdtools/models/multislice_ptycho.py | 46 ++++---- src/cdtools/models/rpi.py | 44 +++---- src/cdtools/models/simple_ptycho.py | 7 +- .../tools/interactions/interactions.py | 4 +- tests/conftest.py | 49 +++++++- tests/models/test_fancy_ptycho.py | 108 ++++++++++++++++++ tests/models/test_simple_ptycho.py | 22 ++++ 10 files changed, 269 insertions(+), 105 deletions(-) create mode 100644 tests/models/test_fancy_ptycho.py create mode 100644 tests/models/test_simple_ptycho.py diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index efbcb30..d2de6f7 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -117,18 +117,18 @@ class CDIModel(t.nn.Module): The datatype to convert the values to before registering """ self.register_buffer('det_basis', - t.tensor(detector_geometry['basis'], + t.as_tensor(detector_geometry['basis'], dtype=dtype)) if 'distance' in detector_geometry \ and detector_geometry['distance'] is not None: self.register_buffer('det_distance', - t.tensor(detector_geometry['distance'], + t.as_tensor(detector_geometry['distance'], dtype=dtype)) if 'corner' in detector_geometry \ and detector_geometry['corner'] is not None: self.register_buffer('det_corner', - t.tensor(detector_geometry['corner'], + t.as_tensor(detector_geometry['corner'], dtype=dtype)) def get_detector_geometry(self): diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index 38a3eca..473c35f 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -92,23 +92,23 @@ class Bragg2DPtycho(CDIModel): super(Bragg2DPtycho, self).__init__() self.register_buffer('wavelength', - t.tensor(wavelength, dtype=dtype)) + t.as_tensor(wavelength, dtype=dtype)) self.store_detector_geometry(detector_geometry, dtype=dtype) self.register_buffer('min_translation', - t.tensor(min_translation, dtype=dtype)) + t.as_tensor(min_translation, dtype=dtype)) self.register_buffer('median_propagation', - t.tensor(median_propagation, dtype=dtype)) + t.as_tensor(median_propagation, dtype=dtype)) self.register_buffer('obj_basis', - t.tensor(obj_basis, dtype=dtype)) + t.as_tensor(obj_basis, dtype=dtype)) if probe_basis is None: self.register_buffer('probe_basis', - t.tensor(obj_basis, dtype=dtype)) + t.as_tensor(obj_basis, dtype=dtype)) else: self.register_buffer('probe_basis', - t.tensor(probe_basis, dtype=dtype)) + t.as_tensor(probe_basis, dtype=dtype)) self.units = units @@ -119,22 +119,22 @@ class Bragg2DPtycho(CDIModel): np.array(obj_basis)[:,0]) surface_normal /= np.linalg.norm(surface_normal) self.register_buffer('surface_normal', - t.tensor(surface_normal, dtype=dtype)) + t.as_tensor(surface_normal, dtype=dtype)) if saturation is None: self.saturation = None else: self.register_buffer('saturation', - t.tensor(saturation, dtype=dtype)) + t.as_tensor(saturation, dtype=dtype)) if mask is None: self.mask = None else: self.register_buffer('mask', - t.tensor(mask, dtype=t.bool)) + t.as_tensor(mask, dtype=t.bool)) - probe_guess = t.tensor(probe_guess, dtype=t.complex64) - obj_guess = t.tensor(obj_guess, dtype=t.complex64) + probe_guess = t.as_tensor(probe_guess, dtype=t.complex64) + obj_guess = t.as_tensor(obj_guess, dtype=t.complex64) # We rescale the probe here so it learns at the same rate as the # object @@ -150,7 +150,7 @@ class Bragg2DPtycho(CDIModel): if probe_support is None: probe_support = t.ones_like(self.probe[0], dtype=t.bool) self.register_buffer('probe_support', - t.tensor(probe_support, dtype=t.bool)) + t.as_tensor(probe_support, dtype=t.bool)) self.probe.data *= self.probe_support if background is None: @@ -164,26 +164,26 @@ class Bragg2DPtycho(CDIModel): self.weights = None else: # No incoherent + unstable here yet - self.weights = t.nn.Parameter(t.tensor(weights, + self.weights = t.nn.Parameter(t.as_tensor(weights, dtype=t.float32)) if translation_offsets is None: self.translation_offsets = None else: - t_o = t.tensor(translation_offsets, dtype=t.float32) + t_o = t.as_tensor(translation_offsets, dtype=t.float32) t_o = t_o / translation_scale self.translation_offsets = t.nn.Parameter(t_o) self.register_buffer('translation_scale', - t.tensor(translation_scale, dtype=dtype)) + t.as_tensor(translation_scale, dtype=dtype)) self.register_buffer('oversampling', - t.tensor(oversampling, dtype=int)) + t.as_tensor(oversampling, dtype=int)) self.register_buffer('propagate_probe', - t.tensor(propagate_probe, dtype=bool)) + t.as_tensor(propagate_probe, dtype=bool)) self.register_buffer('correct_tilt', - t.tensor(correct_tilt, dtype=bool)) + t.as_tensor(correct_tilt, dtype=bool)) if correct_tilt: k_map, intensity_map = \ @@ -195,9 +195,9 @@ class Bragg2DPtycho(CDIModel): self.wavelength,dtype=t.float32, lens=lens) self.register_buffer('k_map', - t.tensor(k_map, dtype=dtype)) + t.as_tensor(k_map, dtype=dtype)) self.register_buffer('intensity_map', - t.tensor(intensity_map, dtype=dtype)) + t.as_tensor(intensity_map, dtype=dtype)) else: self.k_map = None @@ -205,7 +205,7 @@ class Bragg2DPtycho(CDIModel): # The propagation direction of the probe self.register_buffer('prop_dir', - t.tensor([0, 0, 1], dtype=dtype)) + t.as_tensor([0, 0, 1], dtype=dtype)) # This propagator should be able to be multiplied by the propagation # distance each time to get a propagator diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 7b4a540..611fc72 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -44,39 +44,39 @@ class FancyPtycho(CDIModel): super(FancyPtycho, self).__init__() self.register_buffer('wavelength', - t.tensor(wavelength, dtype=dtype)) + t.as_tensor(wavelength, dtype=dtype)) self.store_detector_geometry(detector_geometry, dtype=dtype) self.register_buffer('min_translation', - t.tensor(min_translation, dtype=dtype)) + t.as_tensor(min_translation, dtype=dtype)) self.register_buffer('obj_basis', - t.tensor(obj_basis, dtype=dtype)) + t.as_tensor(obj_basis, dtype=dtype)) if probe_basis is None: self.register_buffer('probe_basis', - t.tensor(obj_basis, dtype=dtype)) + t.as_tensor(obj_basis, dtype=dtype)) else: self.register_buffer('probe_basis', - t.tensor(probe_basis, dtype=dtype)) + t.as_tensor(probe_basis, dtype=dtype)) self.register_buffer('surface_normal', - t.tensor(surface_normal, dtype=dtype)) + t.as_tensor(surface_normal, dtype=dtype)) if saturation is None: self.saturation = None else: self.register_buffer('saturation', - t.tensor(saturation, dtype=dtype)) + t.as_tensor(saturation, dtype=dtype)) self.register_buffer('fourier_probe', - t.tensor(fourier_probe, dtype=bool)) + t.as_tensor(fourier_probe, dtype=bool)) self.register_buffer('exponentiate_obj', - t.tensor(exponentiate_obj, dtype=bool)) + t.as_tensor(exponentiate_obj, dtype=bool)) self.register_buffer('phase_only', - t.tensor(phase_only, dtype=bool)) + t.as_tensor(phase_only, dtype=bool)) # Not sure how to make this a buffer... self.units = units @@ -85,10 +85,10 @@ class FancyPtycho(CDIModel): self.mask = None else: self.register_buffer('mask', - t.tensor(mask, dtype=t.bool)) + t.as_tensor(mask, dtype=t.bool)) - probe_guess = t.tensor(probe_guess, dtype=t.complex64) - obj_guess = t.tensor(obj_guess, dtype=t.complex64) + probe_guess = t.as_tensor(probe_guess, dtype=t.complex64) + obj_guess = t.as_tensor(obj_guess, dtype=t.complex64) # We rescale the probe here so it learns at the same rate as the @@ -121,34 +121,36 @@ class FancyPtycho(CDIModel): # weights and complex-valued per-mode weight matrices if len(weights.shape) == 1: # This is if it's just a list of numbers - self.weights = t.nn.Parameter(t.tensor(weights, + self.weights = t.nn.Parameter(t.as_tensor(weights, dtype=t.float32)) else: # Now this is a matrix of weights, so it needs to be complex - self.weights = t.nn.Parameter(t.tensor(weights, + self.weights = t.nn.Parameter(t.as_tensor(weights, dtype=t.complex64)) if translation_offsets is None: self.translation_offsets = None else: - t_o = t.tensor(translation_offsets, dtype=t.float32) + t_o = t.as_tensor(translation_offsets, dtype=t.float32) t_o = t_o / translation_scale self.translation_offsets = t.nn.Parameter(t_o) self.register_buffer('translation_scale', - t.tensor(translation_scale, dtype=dtype)) + t.as_tensor(translation_scale, dtype=dtype)) if probe_support is None: probe_support = t.ones_like(self.probe[0], dtype=t.bool) self.register_buffer('probe_support', - t.tensor(probe_support, dtype=t.bool)) + t.as_tensor(probe_support, dtype=t.bool)) self.probe.data *= self.probe_support self.register_buffer('oversampling', - t.tensor(oversampling, dtype=int)) + t.as_tensor(oversampling, dtype=int)) - self.register_buffer('simulate_probe_translation', - t.tensor(simulate_probe_translation, dtype=bool)) + self.register_buffer( + 'simulate_probe_translation', + t.as_tensor(simulate_probe_translation, dtype=bool) + ) if simulate_probe_translation: Is = t.arange(self.probe.shape[-2], dtype=dtype) @@ -162,7 +164,7 @@ class FancyPtycho(CDIModel): self.register_buffer('simulate_finite_pixels', - t.tensor(simulate_finite_pixels, dtype=bool)) + t.as_tensor(simulate_finite_pixels, dtype=bool)) # Here we set the appropriate loss function if (loss.lower().strip() == 'amplitude mse' diff --git a/src/cdtools/models/multislice_ptycho.py b/src/cdtools/models/multislice_ptycho.py index f9eab3e..23df19a 100644 --- a/src/cdtools/models/multislice_ptycho.py +++ b/src/cdtools/models/multislice_ptycho.py @@ -44,40 +44,40 @@ class MultislicePtycho(CDIModel): super(MultislicePtycho, self).__init__() self.register_buffer('wavelength', - t.tensor(wavelength, dtype=dtype)) + t.as_tensor(wavelength, dtype=dtype)) self.store_detector_geometry(detector_geometry, dtype=dtype) self.register_buffer('min_translation', - t.tensor(min_translation, dtype=dtype)) + t.as_tensor(min_translation, dtype=dtype)) self.register_buffer('obj_basis', - t.tensor(obj_basis, dtype=dtype)) + t.as_tensor(obj_basis, dtype=dtype)) self.register_buffer('exponentiate_obj', - t.tensor(exponentiate_obj, dtype=bool)) + t.as_tensor(exponentiate_obj, dtype=bool)) self.register_buffer('interslice_propagator', - t.tensor(interslice_propagator, dtype=t.complex64)) + t.as_tensor(interslice_propagator, dtype=t.complex64)) if probe_basis is None: self.register_buffer('probe_basis', - t.tensor(obj_basis, dtype=dtype)) + t.as_tensor(obj_basis, dtype=dtype)) else: self.register_buffer('probe_basis', - t.tensor(probe_basis, dtype=dtype)) + t.as_tensor(probe_basis, dtype=dtype)) self.register_buffer('surface_normal', - t.tensor(surface_normal, dtype=dtype)) + t.as_tensor(surface_normal, dtype=dtype)) if saturation is None: self.saturation = None else: self.register_buffer('saturation', - t.tensor(saturation, dtype=dtype)) + t.as_tensor(saturation, dtype=dtype)) self.register_buffer('fourier_probe', - t.tensor(fourier_probe, dtype=bool)) + t.as_tensor(fourier_probe, dtype=bool)) # Not sure how to make this a buffer... self.units = units @@ -86,10 +86,10 @@ class MultislicePtycho(CDIModel): self.mask = None else: self.register_buffer('mask', - t.tensor(mask, dtype=t.bool)) + t.as_tensor(mask, dtype=t.bool)) - probe_guess = t.tensor(probe_guess, dtype=t.complex64) - obj_guess = t.tensor(obj_guess, dtype=t.complex64) + probe_guess = t.as_tensor(probe_guess, dtype=t.complex64) + obj_guess = t.as_tensor(obj_guess, dtype=t.complex64) # We rescale the probe here so it learns at the same rate as the # object @@ -121,34 +121,36 @@ class MultislicePtycho(CDIModel): # weights and complex-valued per-mode weight matrices if len(weights.shape) == 1: # This is if it's just a list of numbers - self.weights = t.nn.Parameter(t.tensor(weights, + self.weights = t.nn.Parameter(t.as_tensor(weights, dtype=t.float32)) else: # Now this is a matrix of weights, so it needs to be complex - self.weights = t.nn.Parameter(t.tensor(weights, + self.weights = t.nn.Parameter(t.as_tensor(weights, dtype=t.complex64)) if translation_offsets is None: self.translation_offsets = None else: - t_o = t.tensor(translation_offsets, dtype=t.float32) + t_o = t.as_tensor(translation_offsets, dtype=t.float32) t_o = t_o / translation_scale self.translation_offsets = t.nn.Parameter(t_o) self.register_buffer('translation_scale', - t.tensor(translation_scale, dtype=dtype)) + t.as_tensor(translation_scale, dtype=dtype)) if probe_support is None: probe_support = t.ones_like(self.probe[0], dtype=t.bool) self.register_buffer('probe_support', - t.tensor(probe_support, dtype=t.bool)) + t.as_tensor(probe_support, dtype=t.bool)) self.probe.data *= self.probe_support self.register_buffer('oversampling', - t.tensor(oversampling, dtype=int)) + t.as_tensor(oversampling, dtype=int)) - self.register_buffer('simulate_probe_translation', - t.tensor(simulate_probe_translation, dtype=bool)) + self.register_buffer( + 'simulate_probe_translation', + t.as_tensor(simulate_probe_translation, dtype=bool) + ) if simulate_probe_translation: Is = t.arange(self.probe.shape[-2], dtype=dtype) @@ -161,7 +163,7 @@ class MultislicePtycho(CDIModel): self.register_buffer('J_phase', J_phase) self.register_buffer('simulate_finite_pixels', - t.tensor(simulate_finite_pixels, dtype=bool)) + t.as_tensor(simulate_finite_pixels, dtype=bool)) # Here we set the appropriate loss function if (loss.lower().strip() == 'amplitude mse' diff --git a/src/cdtools/models/rpi.py b/src/cdtools/models/rpi.py index 9d8a9ff..339b623 100644 --- a/src/cdtools/models/rpi.py +++ b/src/cdtools/models/rpi.py @@ -66,7 +66,7 @@ class RPI(CDIModel): 1j * t.ones([1], dtype=dtype)).dtype self.register_buffer('wavelength', - t.tensor(wavelength, dtype=dtype)) + t.as_tensor(wavelength, dtype=dtype)) self.store_detector_geometry(detector_geometry, dtype=dtype) @@ -75,9 +75,9 @@ class RPI(CDIModel): # used a bandlimiting constraint and had a larger basis, the user is # expected to upsample it explicitly before doing RPI. self.register_buffer('probe_basis', - t.tensor(probe_basis, dtype=dtype)) + t.as_tensor(probe_basis, dtype=dtype)) - scale_factor = t.tensor([probe.shape[-1]/obj_guess.shape[-1], + scale_factor = t.as_tensor([probe.shape[-1]/obj_guess.shape[-1], probe.shape[-2]/obj_guess.shape[-2]]) self.register_buffer('obj_basis', (self.probe_basis * scale_factor).to(dtype=dtype)) @@ -86,7 +86,7 @@ class RPI(CDIModel): self.saturation = None else: self.register_buffer('saturation', - t.tensor(saturation, dtype=dtype)) + t.as_tensor(saturation, dtype=dtype)) # not sure how to make this a buffer, or if I have to self.units = units @@ -95,23 +95,23 @@ class RPI(CDIModel): self.mask = None else: self.register_buffer('mask', - t.tensor(mask, dtype=t.bool)) + t.as_tensor(mask, dtype=t.bool)) - self.register_buffer('probe', t.tensor(probe, dtype=complex_dtype)) + self.register_buffer('probe', t.as_tensor(probe, dtype=complex_dtype)) self.register_buffer('exponentiate_obj', - t.tensor(exponentiate_obj, dtype=bool)) + t.as_tensor(exponentiate_obj, dtype=bool)) self.register_buffer('phase_only', - t.tensor(phase_only, dtype=bool)) + t.as_tensor(phase_only, dtype=bool)) # We always use multi-modes to store the object, so we convert it # if we just get a single 2D array as an input if obj_guess.dim() == 2: obj_guess = obj_guess[None, :, :] - self.obj = t.nn.Parameter(t.tensor(obj_guess, dtype=complex_dtype)) + self.obj = t.nn.Parameter(t.as_tensor(obj_guess, dtype=complex_dtype)) self.weights = t.nn.Parameter( t.eye(probe.shape[0], dtype=complex_dtype)) @@ -124,40 +124,26 @@ class RPI(CDIModel): dtype=dtype) self.register_buffer('background', - t.tensor(background, dtype=t.float32)) + t.as_tensor(background, dtype=t.float32)) if obj_support is None: obj_support = t.ones_like(self.obj[0, ...], dtype=int) self.register_buffer('obj_support', - t.tensor(obj_support, dtype=int)) + t.as_tensor(obj_support, dtype=int)) self.obj.data = self.obj * self.obj_support[None, ...] self.register_buffer('oversampling', - t.tensor(oversampling, dtype=int)) + t.as_tensor(oversampling, dtype=int)) self.register_buffer('propagation_distance', - t.tensor(propagation_distance, dtype=dtype)) + t.as_tensor(propagation_distance, dtype=dtype)) # The propagation direction of the probe. For now it's fixed, # but perhaps it would need to be updated in the future self.register_buffer('prop_dir', - t.tensor([0, 0, 1], dtype=dtype)) - - # This propagator should be able to be multiplied by the propagation - # distance each time to get a propagator - #universal_propagator = t.angle(ggasp( - # self.probe.shape[-2:], - # self.probe_basis, self.wavelength, - # t.tensor([0, 0, self.wavelength/(2*np.pi)], dtype=dtype), - # propagation_vector=self.prop_dir, - # dtype=complex_dtype, - # propagate_along_offset=True)) - - # TODO: probably doesn't support non-float-32 dtypes - #self.register_buffer('universal_propagator', - # universal_propagator) + t.as_tensor([0, 0, 1], dtype=dtype)) @classmethod @@ -229,7 +215,7 @@ class RPI(CDIModel): and dataset.background is not None: background = t.sqrt(dataset.background) elif background is not None: - background = t.sqrt(t.Tensor(background).to(dtype=t.float32)) + background = t.sqrt(t.as_tensor(background).to(dtype=t.float32)) det_geo = dataset.detector_geometry diff --git a/src/cdtools/models/simple_ptycho.py b/src/cdtools/models/simple_ptycho.py index f189c5c..13dde61 100644 --- a/src/cdtools/models/simple_ptycho.py +++ b/src/cdtools/models/simple_ptycho.py @@ -13,9 +13,6 @@ __all__ = ['SimplePtycho'] class SimplePtycho(CDIModel): """A simple ptychography model for exploring ideas and extensions - - - """ def __init__( self, @@ -40,8 +37,8 @@ class SimplePtycho(CDIModel): self.register_buffer('min_translation', t.as_tensor(min_translation)) self.register_buffer('probe_basis', t.as_tensor(probe_basis)) - probe_guess = t.tensor(probe_guess, dtype=t.complex64) - obj_guess = t.tensor(obj_guess, dtype=t.complex64) + probe_guess = t.as_tensor(probe_guess, dtype=t.complex64) + obj_guess = t.as_tensor(obj_guess, dtype=t.complex64) # We rescale the probe here so it learns at the same rate as the # object diff --git a/src/cdtools/tools/interactions/interactions.py b/src/cdtools/tools/interactions/interactions.py index 2d8642e..db330bb 100644 --- a/src/cdtools/tools/interactions/interactions.py +++ b/src/cdtools/tools/interactions/interactions.py @@ -464,8 +464,8 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe, dim=(-1,-2))) - if probe_support is not None: - shifted_probe = shifted_probe * probe_support[..., :, :] + # Note: resist the temptation to remultiply by the probe support here, + # it will fail if you have a probe which is restricted in Fourier space # TODO This is a kludge, I will fix this. if multiple_modes and len(selections.shape) == 3: # Multi-mode probe diff --git a/tests/conftest.py b/tests/conftest.py index 92a094c..9eceb0c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,8 +14,41 @@ import datetime def pytest_addoption(parser): parser.addoption( - "--plot", action="store", default=False, help="plot: True to show test plots" + "--plot", + action="store_true", + default=False, + help="when set, shows the test plots" ) + parser.addoption( + "--reconstruction_device", + action="store", + default="cuda", + help="What device to run reconstructions on, if they are being run" + ) + parser.addoption( + "--runslow", + action="store_true", + default=False, + help="run slow tests, primarily full reconstruction tests." + ) + + +def pytest_configure(config): + config.addinivalue_line("markers", "slow: mark test as slow to run") + + +def pytest_collection_modifyitems(config, items): + if config.getoption("--runslow"): + # --runslow given in cli: do not skip slow tests + return + skip_slow = pytest.mark.skip(reason="need --runslow option to run") + for item in items: + if "slow" in item.keywords: + item.add_marker(skip_slow) + +@pytest.fixture +def reconstruction_device(request): + return request.config.getoption("--reconstruction_device") @pytest.fixture @@ -319,3 +352,17 @@ def test_ptycho_cxis(ptycho_cxi_1, ptycho_cxi_2, ptycho_cxi_3): on the cxi files. """ return [ptycho_cxi_1, ptycho_cxi_2, ptycho_cxi_3] + + +@pytest.fixture(scope='module') +def gold_ball_cxi(pytestconfig): + return str(pytestconfig.rootpath) + \ + '/examples/example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi' + +@pytest.fixture(scope='module') +def lab_ptycho_cxi(pytestconfig): + return str(pytestconfig.rootpath) + \ + '/examples/example_data/lab_ptycho_data.cxi' + + + diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py new file mode 100644 index 0000000..c9eaa1b --- /dev/null +++ b/tests/models/test_fancy_ptycho.py @@ -0,0 +1,108 @@ +import pytest +import cdtools +import torch as t + +import cdtools +from matplotlib import pyplot as plt + +@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, + ) + + 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 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 + + diff --git a/tests/models/test_simple_ptycho.py b/tests/models/test_simple_ptycho.py new file mode 100644 index 0000000..770b1a0 --- /dev/null +++ b/tests/models/test_simple_ptycho.py @@ -0,0 +1,22 @@ +import cdtools +from matplotlib import pyplot as plt + +def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot): + dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi) + + model = cdtools.models.SimplePtycho.from_dataset(dataset) + + model.to(device=reconstruction_device) + dataset.get_as(device=reconstruction_device) + + 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) + model.compare(dataset) + + # If this fails, the reconstruction got worse + assert model.loss_history[-1] < 0.013