diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index ea1ddd7..3c70f86 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -398,6 +398,7 @@ def from_dataset(cls, dataset.patterns.shape[-2:], dtype=t.float32) if probe_fourier_crop is not None: + initial_probe_shape = np.array(probe.shape) probe = tools.propagators.far_field(probe) probe = probe[probe_fourier_crop : probe.shape[-2] - probe_fourier_crop, @@ -405,7 +406,7 @@ def from_dataset(cls, - probe_fourier_crop] probe = tools.propagators.inverse_far_field(probe) - scale_factor = np.array(det_shape) / np.array(probe.shape) + scale_factor = initial_probe_shape / np.array(probe.shape) probe_basis = obj_basis * scale_factor[None,:] else: probe_basis = obj_basis.clone()