diff --git a/docs/source/using_pty_chi/data_structures.rst b/docs/source/using_pty_chi/data_structures.rst index 1e3b0fe..2884480 100644 --- a/docs/source/using_pty_chi/data_structures.rst +++ b/docs/source/using_pty_chi/data_structures.rst @@ -37,6 +37,9 @@ parameters. - (1) Probe positions should follow row-major order, i.e., y-coordinates come first. (2) The aforementioned shape is of the internally stored position tensor. ``PtychographyTask`` takes x- and y-positions as separate data arguments. + (3) Positions are always expressed in object-pixel units, including when + the probe and object pixel sizes differ. Probe-grid coordinates are + temporary internal copies. * - Diffraction data - ``(n_positions, height, width)`` - diff --git a/docs/source/using_pty_chi/options.rst b/docs/source/using_pty_chi/options.rst index 9a014f8..99ca2e0 100644 --- a/docs/source/using_pty_chi/options.rst +++ b/docs/source/using_pty_chi/options.rst @@ -130,6 +130,37 @@ delay probe updates until later: options.probe_options.optimization_plan.start = 5 +Independent object and probe pixel sizes +---------------------------------------- + +The object and probe may use different real-space sampling. Object pixels are +configured by ``object_options.pixel_size_m`` (the x/width dimension) and +``object_options.pixel_size_aspect_ratio`` (width divided by height). The +corresponding optional probe settings are +``probe_options.pixel_size_m`` and +``probe_options.pixel_size_aspect_ratio``. Each omitted probe value inherits +the matching object value independently. For example, this changes only the +probe's y sampling: + +.. code-block:: python + + options.probe_options.pixel_size_m = None + options.probe_options.pixel_size_aspect_ratio = 1.25 + +Pixel height is ``pixel_size_m / pixel_size_aspect_ratio``. For each axis, a +native object dimension ``n`` is represented internally on the probe grid with + +.. math:: + + n_{probe} = \max(1, \operatorname{round}(n\,\Delta_{object}/\Delta_{probe})). + +Pty-Chi resamples the complete multislice object by centered Fourier +padding/cropping, using constant-amplitude normalization. Analytical object +gradients are mapped back with the matching Hermitian adjoint; direct object +projections use its normalized pseudoinverse. The equal-pixel-size case keeps +the original no-resampling path. + + OptimizationPlan ---------------- diff --git a/src/ptychi/api/options/base.py b/src/ptychi/api/options/base.py index a3c2ee0..2d280ee 100644 --- a/src/ptychi/api/options/base.py +++ b/src/ptychi/api/options/base.py @@ -716,6 +716,16 @@ class ProbeOptions(ParameterOptions): initial_guess: Optional[DataArray] = None """A (n_opr_modes, n_modes, h, w) complex tensor of the probe initial guess.""" + pixel_size_m: Optional[float] = PydanticField(default=None, gt=0) + """Probe pixel width in meters. If ``None``, inherit the object pixel width.""" + + pixel_size_aspect_ratio: Optional[float] = PydanticField(default=None, gt=0) + """Probe pixel width/height ratio. If ``None``, inherit the object ratio. + + The width and aspect ratio are inherited independently, so either value may be + overridden without specifying the other one. + """ + power_constraint: ProbePowerConstraintOptions = field( default_factory=ProbePowerConstraintOptions ) diff --git a/src/ptychi/api/task.py b/src/ptychi/api/task.py index 4391e0a..dcd0a0a 100644 --- a/src/ptychi/api/task.py +++ b/src/ptychi/api/task.py @@ -357,26 +357,49 @@ def _check_object_position_coverage( pos_x = self._as_numpy(position_x_px) obj_lateral_shape = object_shape[-2:] probe_lateral_shape = probe_shape[-2:] + object_pixel_width = self.object_options.pixel_size_m + object_pixel_height = ( + object_pixel_width / self.object_options.pixel_size_aspect_ratio + ) + probe_pixel_width = ( + object_pixel_width + if self.probe_options.pixel_size_m is None + else self.probe_options.pixel_size_m + ) + probe_aspect_ratio = ( + self.object_options.pixel_size_aspect_ratio + if self.probe_options.pixel_size_aspect_ratio is None + else self.probe_options.pixel_size_aspect_ratio + ) + probe_pixel_height = probe_pixel_width / probe_aspect_ratio + probe_grid_shape = ( + max(1, round(obj_lateral_shape[0] * object_pixel_height / probe_pixel_height)), + max(1, round(obj_lateral_shape[1] * object_pixel_width / probe_pixel_width)), + ) + scale = np.asarray(probe_grid_shape) / np.asarray(obj_lateral_shape) + pos_y_probe = pos_y * scale[0] + pos_x_probe = pos_x * scale[1] min_size = [ - int(np.ceil(pos_y.max() - pos_y.min() + probe_lateral_shape[-2])) + 2, - int(np.ceil(pos_x.max() - pos_x.min() + probe_lateral_shape[-1])) + 2, + int(np.ceil(pos_y_probe.max() - pos_y_probe.min() + probe_lateral_shape[-2])) + 2, + int(np.ceil(pos_x_probe.max() - pos_x_probe.min() + probe_lateral_shape[-1])) + 2, ] - if any(min_size[i] > obj_lateral_shape[i] for i in range(2)): + if any(min_size[i] > probe_grid_shape[i] for i in range(2)): logging.warning( - f"An object tensor with a lateral size of at least {min_size} is " + f"A probe-grid object tensor with a lateral size of at least {min_size} is " "required to avoid padding when extracting/placing patches, but the provided " - f"object size is {list(obj_lateral_shape)}." + f"object maps to {list(probe_grid_shape)} probe-grid pixels." ) if ( self.object_options.determine_position_origin_coords_by == api.ObjectPosOriginCoordsMethods.SUPPORT ): - buffer_center = np.array([np.round(x / 2) + 0.5 for x in obj_lateral_shape]) + native_center = np.array([np.round(x / 2) + 0.5 for x in obj_lateral_shape]) + buffer_center = native_center * scale if ( - pos_y.max() + buffer_center[0] + probe_lateral_shape[-2] // 2 > obj_lateral_shape[-2] - or pos_y.min() + buffer_center[0] - probe_lateral_shape[-2] // 2 < 0 - or pos_x.max() + buffer_center[1] + probe_lateral_shape[-1] // 2 > obj_lateral_shape[-1] - or pos_x.min() + buffer_center[1] - probe_lateral_shape[-1] // 2 < 0 + pos_y_probe.max() + buffer_center[0] + probe_lateral_shape[-2] // 2 > probe_grid_shape[-2] + or pos_y_probe.min() + buffer_center[0] - probe_lateral_shape[-2] // 2 < 0 + or pos_x_probe.max() + buffer_center[1] + probe_lateral_shape[-1] // 2 > probe_grid_shape[-1] + or pos_x_probe.min() + buffer_center[1] - probe_lateral_shape[-1] // 2 < 0 ): logging.warning( "`object_options.determine_center_coords_by` is set to `SUPPORT`. This assumes " @@ -739,6 +762,7 @@ def set_large_tensor_device( ) self.reconstructor.forward_model.move_intermediate_variables_to_device(device) + self.reconstructor.forward_model.clear_resampling_caches() if device.type == "cpu": AcceleratorModuleWrapper.get_module().empty_cache() diff --git a/src/ptychi/data_structures/object.py b/src/ptychi/data_structures/object.py index b4d5cdc..66cec20 100644 --- a/src/ptychi/data_structures/object.py +++ b/src/ptychi/data_structures/object.py @@ -279,6 +279,7 @@ def extract_patches( patch_shape: Tuple[int, int], integer_mode: bool = False, pad_for_shift: Optional[int] = 1, + object_array: Optional[Tensor] = None, ): """ Extract (n_patches, n_slices, h', w') patches from the object. @@ -296,6 +297,11 @@ def extract_patches( pad_for_shift : int, optional If given, patches larger than the intended size by this amount are cropped out from the object before shifting. + object_array : Tensor, optional + Explicit ``(n_slices, height, width)`` object array from which to extract + patches. If omitted, patches are extracted from ``self.data``. Positions + retain the same coordinate convention in either case: their origin is + ``self.pos_origin_coords``. Returns ------- @@ -304,16 +310,24 @@ def extract_patches( """ # Positions are provided with the origin in the center of the object support. # We shift the positions so that the origin is in the upper left corner. + if object_array is None: + object_array = self.data + if object_array.ndim != 3 or object_array.shape[0] != self.n_slices: + raise ValueError( + "`object_array` must have shape (n_slices, height, width) with " + f"n_slices={self.n_slices}." + ) + positions = positions + self.pos_origin_coords patches_all_slices = [] for i_slice in range(self.n_slices): if integer_mode: patches = ip.extract_patches_integer( - self.get_slice(i_slice), positions, patch_shape + object_array[i_slice], positions, patch_shape ) else: patches = self.extract_patches_function( - self.get_slice(i_slice), positions, patch_shape, pad=pad_for_shift + object_array[i_slice], positions, patch_shape, pad=pad_for_shift ) patches_all_slices.append(patches) patches_all_slices = torch.stack(patches_all_slices, dim=1) @@ -627,6 +641,25 @@ def calculate_illumination_map( else: probe_int = probe.get_mode_and_opr_mode(mode=0, opr_mode=0)[None, ...].abs() ** 2 + probe_pixel_width = ( + self.pixel_size_m + if probe.options.pixel_size_m is None + else probe.options.pixel_size_m + ) + probe_aspect_ratio = ( + self.options.pixel_size_aspect_ratio + if probe.options.pixel_size_aspect_ratio is None + else probe.options.pixel_size_aspect_ratio + ) + probe_pixel_height = probe_pixel_width / probe_aspect_ratio + object_pixel_height = self.pixel_size_m / self.options.pixel_size_aspect_ratio + physical_footprint_shape = ( + max(1, round(probe_int.shape[-2] * probe_pixel_height / object_pixel_height)), + max(1, round(probe_int.shape[-1] * probe_pixel_width / self.pixel_size_m)), + ) + if physical_footprint_shape != tuple(probe_int.shape[-2:]): + probe_int = ip.fourier_resize(probe_int, physical_footprint_shape).real.clamp_min_(0) + # Stitch probes of all positions on the object buffer # TODO: allow setting chunk size externally probe_sq_map = chunked_processing( diff --git a/src/ptychi/data_structures/probe_positions.py b/src/ptychi/data_structures/probe_positions.py index cdeda35..5407ea0 100644 --- a/src/ptychi/data_structures/probe_positions.py +++ b/src/ptychi/data_structures/probe_positions.py @@ -109,10 +109,31 @@ def update_position_weights( is done in batches to avoid memory issues. This can, but does not have to be the same as the batch size used for reconstruction. """ - half_probe_shape = [x // 2 for x in probe.shape[-2:]] - probe_intensity = ip.central_crop_or_pad( - probe.data[0, 0].abs() ** 2, half_probe_shape + probe_pixel_width = ( + object_.pixel_size_m + if probe.options.pixel_size_m is None + else probe.options.pixel_size_m ) + probe_aspect_ratio = ( + object_.options.pixel_size_aspect_ratio + if probe.options.pixel_size_aspect_ratio is None + else probe.options.pixel_size_aspect_ratio + ) + probe_pixel_height = probe_pixel_width / probe_aspect_ratio + object_pixel_height = ( + object_.pixel_size_m / object_.options.pixel_size_aspect_ratio + ) + probe_shape_on_object_grid = ( + max(1, round(probe.shape[-2] * probe_pixel_height / object_pixel_height)), + max(1, round(probe.shape[-1] * probe_pixel_width / object_.pixel_size_m)), + ) + half_probe_shape = [max(1, x // 2) for x in probe_shape_on_object_grid] + probe_intensity = probe.data[0, 0].abs() ** 2 + if probe_shape_on_object_grid != tuple(probe.shape[-2:]): + probe_intensity = ip.fourier_resize( + probe_intensity, probe_shape_on_object_grid + ).real.clamp_min_(0) + probe_intensity = ip.central_crop_or_pad(probe_intensity, half_probe_shape) total_variations = torch.zeros_like(self.data) obj_slice = object_.get_slice(self.get_slice_for_correction(object_.n_slices)) diff --git a/src/ptychi/forward_models.py b/src/ptychi/forward_models.py index 43b26a1..0761709 100644 --- a/src/ptychi/forward_models.py +++ b/src/ptychi/forward_models.py @@ -81,6 +81,10 @@ def forward(self, *args, **kwargs): def post_differentiation_hook(self, *args, **kwargs): pass + def clear_resampling_caches(self): + """Clear derived grid caches, if the forward model defines any.""" + return None + def record_intermediate_variable(self, name, var): if self.retain_intermediates: if isinstance(self.intermediate_variables[name], list): @@ -210,6 +214,17 @@ def __init__( self.probe_positions = parameter_group.probe_positions self.opr_mode_weights = parameter_group.opr_mode_weights + # Resampling caches are ordinary Python attributes on purpose: they are + # read-only derived values and must not appear in state_dict(). Each + # cache keeps only its newest entry on a given device. + self._resampling_caches = { + "object": {}, + "preconditioner": {}, + "positions": {}, + } + self._resampling_geometry_key = None + self._resampling_geometry = None + self.wavelength_m = wavelength_m self.free_space_propagation_distance_m = free_space_propagation_distance_m @@ -232,6 +247,286 @@ def __init__( self.check_inputs() + @property + def probe_pixel_size_m(self) -> float: + value = self.probe.options.pixel_size_m + return self.object.pixel_size_m if value is None else value + + @property + def probe_pixel_size_aspect_ratio(self) -> float: + value = self.probe.options.pixel_size_aspect_ratio + return self.object.options.pixel_size_aspect_ratio if value is None else value + + @property + def probe_pixel_height_m(self) -> float: + return self.probe_pixel_size_m / self.probe_pixel_size_aspect_ratio + + @property + def object_pixel_height_m(self) -> float: + return self.object.pixel_size_m / self.object.options.pixel_size_aspect_ratio + + @property + def probe_grid_object_shape(self) -> tuple[int, int]: + native_h, native_w = self.object.lateral_shape + return ( + max(1, round(native_h * self.object_pixel_height_m / self.probe_pixel_height_m)), + max(1, round(native_w * self.object.pixel_size_m / self.probe_pixel_size_m)), + ) + + @property + def probe_grid_scale(self) -> tuple[float, float]: + target_h, target_w = self.probe_grid_object_shape + native_h, native_w = self.object.lateral_shape + return target_h / native_h, target_w / native_w + + @property + def resampling_enabled(self) -> bool: + # Pixel geometry, rather than rounded shape alone, defines the identity + # fast path. This preserves exact legacy behavior for inherited grids. + return not ( + self.probe_pixel_size_m == self.object.pixel_size_m + and self.probe_pixel_height_m == self.object_pixel_height_m + ) + + def _current_geometry_key(self): + return ( + tuple(self.object.lateral_shape), + self.probe_grid_object_shape, + self.object.pixel_size_m, + self.object.options.pixel_size_aspect_ratio, + self.probe_pixel_size_m, + self.probe_pixel_size_aspect_ratio, + tuple(self.probe.get_spatial_shape()), + ) + + def clear_resampling_caches(self): + if hasattr(self, "_resampling_caches"): + for cache in self._resampling_caches.values(): + cache.clear() + self._resampling_geometry_key = self._current_geometry_key() + native_shape = tuple(self.object.lateral_shape) + target_shape = self.probe_grid_object_shape + source_slices = [] + target_slices = [] + for source_size, target_size in zip(native_shape, target_shape): + overlap = min(source_size, target_size) + source_start = source_size // 2 - overlap // 2 + target_start = target_size // 2 - overlap // 2 + source_slices.append(slice(source_start, source_start + overlap)) + target_slices.append(slice(target_start, target_start + overlap)) + native_numel = native_shape[0] * native_shape[1] + target_numel = target_shape[0] * target_shape[1] + self._resampling_geometry = { + "native_shape": native_shape, + "target_shape": target_shape, + "axis_scales": self.probe_grid_scale, + "forward_normalization": target_numel / native_numel, + "normalized_adjoint_normalization": native_numel / target_numel, + "source_slices": tuple(source_slices), + "target_slices": tuple(target_slices), + } + + @property + def resampling_geometry(self) -> dict: + """Read-only static geometry used by the probe-grid transforms.""" + self._ensure_current_geometry() + return self._resampling_geometry + + def _ensure_current_geometry(self): + geometry_key = self._current_geometry_key() + if self._resampling_geometry_key != geometry_key: + self.clear_resampling_caches() + # Propagators contain pixel-geometry-dependent state. + if hasattr(self, "free_space_propagator"): + self.build_free_space_propagator() + if hasattr(self, "in_object_propagator"): + self.build_in_object_propagator() + + def _apply(self, fn): + result = super()._apply(fn) + self.clear_resampling_caches() + return result + + @staticmethod + def _tensor_version_key(tensor: Tensor): + try: + storage_identity = tensor.untyped_storage().data_ptr() + except AttributeError: + storage_identity = tensor.storage().data_ptr() + return storage_identity, tensor._version + + def _cache_value(self, kind, source_tensors, value_factory, *, dtype, device): + """Return a version-keyed cached value, respecting autograd graph lifetime. + + Parameters + ---------- + kind : str + Cache namespace in ``self._resampling_caches``, such as ``"object"``, + ``"preconditioner"``, or ``"positions"``. + source_tensors : Iterable[Tensor] + Backing tensors on which the derived value depends. Their storage + identities and PyTorch mutation versions are included in the cache key. + value_factory : Callable[[], Any] + Zero-argument callable that computes the derived value after a cache miss. + dtype : torch.dtype + Data type of the derived value, included in the cache key. + device : torch.device or str + Device on which the derived value resides. Each cache namespace retains + only its latest entry for a given device. + + Returns + ------- + Any + The cached value, or the newly computed value after a cache miss. + + Notes + ----- + If gradient tracking is enabled and any source tensor requires gradients, + the value is recomputed and not cached so an earlier autograd graph is never + reused. + """ + self._ensure_current_geometry() + device = torch.device(device) + sources = tuple(source_tensors) + key = ( + tuple(self._tensor_version_key(source) for source in sources), + str(device), + dtype, + self._current_geometry_key(), + ) + graph_must_be_fresh = torch.is_grad_enabled() and any( + source.requires_grad for source in sources + ) + device_cache = self._resampling_caches[kind] + if not graph_must_be_fresh: + cached = device_cache.get(str(device)) + if cached is not None and cached[0] == key: + return cached[1] + value = value_factory() + if not graph_must_be_fresh: + device_cache[str(device)] = (key, value) + return value + + def get_probe_grid_object(self) -> Tensor: + """Return the complete multislice object represented on the probe grid.""" + self._ensure_current_geometry() + data = self.object.data + if not self.resampling_enabled: + return data + backing = self.object.tensor.data + return self._cache_value( + "object", + (backing,), + lambda: ip.fourier_resize(data, self.probe_grid_object_shape), + dtype=data.dtype, + device=data.device, + ) + + def get_probe_grid_preconditioner(self) -> Optional[Tensor]: + """Return a cached real, nonnegative copy of the native preconditioner.""" + preconditioner = self.object.preconditioner + if preconditioner is None or not self.resampling_enabled: + return preconditioner + return self._cache_value( + "preconditioner", + (preconditioner,), + lambda: ip.fourier_resize( + preconditioner, self.probe_grid_object_shape + ).real.clamp_min_(0), + dtype=preconditioner.dtype, + device=preconditioner.device, + ) + + def get_probe_grid_position_data(self): + """Return full transformed centers, snapped centers, and shift residuals.""" + positions = self.probe_positions.tensor + origin = self.object.pos_origin_coords + if not self.resampling_enabled: + absolute = positions + origin + snapped = positions.round() + origin + residual = positions - positions.round() + return absolute, snapped, residual + + def transform(): + scale = positions.new_tensor(self.probe_grid_scale) + absolute = (positions + origin) * scale + parity = positions.new_tensor( + [((size - 1) / 2) % 1 for size in self.probe.get_spatial_shape()] + ) + snapped = (absolute - parity).round() + parity + return absolute, snapped, absolute - snapped + + return self._cache_value( + "positions", + (positions, origin), + transform, + dtype=positions.dtype, + device=positions.device, + ) + + def transform_positions(self, positions: Tensor, *, absolute: bool = False) -> Tensor: + """Map object-pixel positions to probe-grid pixels without modifying inputs.""" + if not self.resampling_enabled: + return positions + self.object.pos_origin_coords if absolute else positions + scale = positions.new_tensor(self.probe_grid_scale) + if absolute: + return (positions + self.object.pos_origin_coords) * scale + return positions * scale + + def probe_displacements_to_object_pixels(self, displacements: Tensor) -> Tensor: + if not self.resampling_enabled: + return displacements + return displacements / displacements.new_tensor(self.probe_grid_scale) + + def _snap_probe_grid_centers(self, positions: Tensor, patch_shape) -> Tensor: + if not self.resampling_enabled: + return positions.round() + self.object.pos_origin_coords + centers = self.transform_positions(positions, absolute=True) + parity = centers.new_tensor([((size - 1) / 2) % 1 for size in patch_shape]) + return (centers - parity).round() + parity + + def place_object_patches_on_probe_grid( + self, positions: Tensor, patches: Tensor, *, integer_mode: bool = True + ) -> Tensor: + """Place patches on an empty probe-grid object buffer.""" + if integer_mode: + centers = self._snap_probe_grid_centers(positions, patches.shape[-2:]) + else: + centers = self.transform_positions(positions, absolute=True) + image = torch.zeros( + self.probe_grid_object_shape, dtype=patches.dtype, device=patches.device + ) + if integer_mode: + return ip.place_patches_integer(image, centers, patches, op="add") + return self.object.place_patches_function( + image, centers, patches, op="add", pad=self.pad_for_shift + ) + + def extract_probe_grid_patches( + self, image: Tensor, positions: Tensor, patch_shape=None, *, integer_mode: bool = True + ) -> Tensor: + """Extract patches from a probe-grid buffer at object-domain positions.""" + if patch_shape is None: + patch_shape = self.probe.get_spatial_shape() + if integer_mode: + centers = self._snap_probe_grid_centers(positions, patch_shape) + return ip.extract_patches_integer(image, centers, patch_shape) + centers = self.transform_positions(positions, absolute=True) + return self.object.extract_patches_function( + image, centers, patch_shape, pad=self.pad_for_shift + ) + + def object_update_to_native_grid(self, update: Tensor, *, normalized: bool = False) -> Tensor: + """Apply the resize adjoint (or normalized pseudoinverse) to an update.""" + if not self.resampling_enabled: + return update + return ip.fourier_resize( + update, + tuple(self.object.lateral_shape), + adjoint=True, + normalized_adjoint=normalized, + ) + def check_inputs(self): if self.probe.has_multiple_opr_modes: if self.opr_mode_weights is None: @@ -264,8 +559,8 @@ def build_in_object_propagator(self): wavelength_m=self.wavelength_m, width_px=self.probe.shape[-1], height_px=self.probe.shape[-2], - pixel_width_m=self.object.pixel_size_m, - pixel_height_m=self.object.pixel_size_m / self.object.options.pixel_size_aspect_ratio, + pixel_width_m=self.probe_pixel_size_m, + pixel_height_m=self.probe_pixel_height_m, propagation_distance_m=self.object.slice_spacings.data[0], ) self.in_object_propagator = AngularSpectrumPropagator(self.in_object_prop_params) @@ -278,8 +573,8 @@ def build_free_space_propagator(self): wavelength_m=self.wavelength_m, width_px=self.probe.shape[-1], height_px=self.probe.shape[-2], - pixel_width_m=self.object.pixel_size_m, - pixel_height_m=self.object.pixel_size_m, + pixel_width_m=self.probe_pixel_size_m, + pixel_height_m=self.probe_pixel_height_m, propagation_distance_m=self.free_space_propagation_distance_m, ) # TODO: AngularSpectrumPropagator uses analytical transfer function. Using the FFT @@ -290,19 +585,22 @@ def build_free_space_propagator(self): self.free_space_propagator = AngularSpectrumPropagator(params) def extract_object_patches(self, indices: Tensor) -> Tensor: - positions = self.probe_positions.data[indices] - if self.apply_subpixel_shifts_on_probe: - obj_patches = self.object.extract_patches( - positions.round().int(), self.probe.get_spatial_shape(), - integer_mode=True - ) - else: - obj_patches = self.object.extract_patches( - positions, self.probe.get_spatial_shape(), - pad_for_shift=self.pad_for_shift, - integer_mode=False - ) - return obj_patches + object_on_probe_grid = self.get_probe_grid_object() + absolute, snapped, _ = self.get_probe_grid_position_data() + absolute_positions = ( + snapped[indices] if self.apply_subpixel_shifts_on_probe else absolute[indices] + ) + # PlanarObject.extract_patches expects positions relative to its configured + # origin. Offset the probe-grid absolute centers so adding that origin in + # the object method recovers the intended centers in object_on_probe_grid. + positions = absolute_positions - self.object.pos_origin_coords + return self.object.extract_patches( + positions, + self.probe.get_spatial_shape(), + integer_mode=self.apply_subpixel_shifts_on_probe, + pad_for_shift=self.pad_for_shift, + object_array=object_on_probe_grid, + ) def get_unique_probes(self, indices: Tensor, always_return_probe_batch: bool = True) -> Tensor: """Get the unique probes for all positions in the batch. @@ -349,7 +647,7 @@ def shift_unique_probes(self, indices: Tensor, unique_probes: Tensor, first_mode If True, only the first mode is shifted. """ orig_shape = unique_probes.shape - fractional_shifts = self.probe_positions.data[indices] - self.probe_positions.data[indices].round() + fractional_shifts = self.get_probe_grid_position_data()[2][indices] if first_mode_only: unique_probe_to_shift = unique_probes[..., 0, :, :] @@ -477,8 +775,8 @@ def propagate_to_next_slice(self, psi: Tensor, slice_index: int): wavelength_m=self.wavelength_m, width_px=self.probe.shape[-1], height_px=self.probe.shape[-2], - pixel_width_m=self.object.pixel_size_m, - pixel_height_m=self.object.pixel_size_m / self.object.options.pixel_size_aspect_ratio, + pixel_width_m=self.probe_pixel_size_m, + pixel_height_m=self.probe_pixel_height_m, propagation_distance_m=self.object.slice_spacings.data[slice_index], ) self.in_object_propagator.update(self.in_object_prop_params) @@ -510,8 +808,8 @@ def propagate_to_previous_slice(self, psi: Tensor, slice_index: int): wavelength_m=self.wavelength_m, width_px=self.probe.shape[-1], height_px=self.probe.shape[-2], - pixel_width_m=self.object.pixel_size_m, - pixel_height_m=self.object.pixel_size_m / self.object.options.pixel_size_aspect_ratio, + pixel_width_m=self.probe_pixel_size_m, + pixel_height_m=self.probe_pixel_height_m, propagation_distance_m=self.object.slice_spacings.data[slice_index - 1], ) self.in_object_propagator.update(self.in_object_prop_params) diff --git a/src/ptychi/image_proc.py b/src/ptychi/image_proc.py index 1b97b41..c25c769 100644 --- a/src/ptychi/image_proc.py +++ b/src/ptychi/image_proc.py @@ -1844,3 +1844,72 @@ def central_crop_or_pad(img: Tensor, target_size: tuple[int, int]) -> Tensor: elif img.shape[-2 + i] < target_size[-2 + i]: img = central_pad(img, target_size_current_dim) return img + + +def fourier_resize( + image: Tensor, + target_shape: tuple[int, int], + *, + adjoint: bool = False, + normalized_adjoint: bool = False, +) -> Tensor: + """Resize the last two dimensions by centered Fourier cropping/padding. + + The forward operation preserves the amplitude of a constant input. With + ``adjoint=True`` this returns its exact Hermitian adjoint. Setting + ``normalized_adjoint=True`` additionally applies the normalization that + makes the reverse operation a pseudoinverse on retained Fourier modes. + + Parameters + ---------- + image : Tensor + A real or complex tensor with at least two dimensions. + target_shape : tuple[int, int] + Requested height and width. + adjoint : bool + Apply the Hermitian adjoint of a forward resize from ``target_shape`` + to ``image.shape[-2:]``. + normalized_adjoint : bool + Return the constant-amplitude reverse resize rather than the raw + Hermitian adjoint. This is only valid when ``adjoint`` is true. + """ + target_shape = tuple(int(v) for v in target_shape) + if len(target_shape) != 2 or any(v < 1 for v in target_shape): + raise ValueError("`target_shape` must contain two positive integers.") + if normalized_adjoint and not adjoint: + raise ValueError("`normalized_adjoint` requires `adjoint=True`.") + source_shape = tuple(image.shape[-2:]) + if source_shape == target_shape: + return image + + spectrum = torch.fft.fftshift(torch.fft.fft2(image, dim=(-2, -1)), dim=(-2, -1)) + # Align the zero-frequency samples explicitly. Symmetric ``F.pad`` puts + # the extra sample on the wrong side for some odd/even transitions. + source_slices = [] + target_slices = [] + for source_size, target_size in zip(source_shape, target_shape): + overlap = min(source_size, target_size) + source_start = source_size // 2 - overlap // 2 + target_start = target_size // 2 - overlap // 2 + source_slices.append(slice(source_start, source_start + overlap)) + target_slices.append(slice(target_start, target_start + overlap)) + resized_spectrum = torch.zeros( + (*spectrum.shape[:-2], *target_shape), dtype=spectrum.dtype, device=spectrum.device + ) + resized_spectrum[..., target_slices[0], target_slices[1]] = spectrum[ + ..., source_slices[0], source_slices[1] + ] + resized = torch.fft.ifft2( + torch.fft.ifftshift(resized_spectrum, dim=(-2, -1)), dim=(-2, -1) + ) + if not adjoint: + resized = resized * ( + (target_shape[0] * target_shape[1]) / (source_shape[0] * source_shape[1]) + ) + elif normalized_adjoint: + resized = resized * ( + (target_shape[0] * target_shape[1]) / (source_shape[0] * source_shape[1]) + ) + if not image.is_complex(): + resized = resized.real + return resized diff --git a/src/ptychi/reconstructors/base.py b/src/ptychi/reconstructors/base.py index d5bc456..e78e0bd 100644 --- a/src/ptychi/reconstructors/base.py +++ b/src/ptychi/reconstructors/base.py @@ -751,9 +751,8 @@ def adjoint_shift_probe_update_direction(self, indices, delta_p, first_mode_only Tensor A (batch_size, n_probe_modes, h, w) tensor giving the shifted probe update direction. """ - pos = self.parameter_group.probe_positions.data[indices] orig_shape = delta_p.shape - fractional_shifts = pos - pos.round() + fractional_shifts = self.forward_model.get_probe_grid_position_data()[2][indices] if first_mode_only: delta_p_to_shift = delta_p[..., 0, :, :] diff --git a/src/ptychi/reconstructors/bh.py b/src/ptychi/reconstructors/bh.py index 2ffc948..83505f7 100644 --- a/src/ptychi/reconstructors/bh.py +++ b/src/ptychi/reconstructors/bh.py @@ -160,7 +160,9 @@ def compute_updates( delta_o = delta_o.unsqueeze(0) delta_p_all_modes = delta_p[None, :, :] delta_pos = torch.zeros_like(probe_positions.data) - delta_pos[indices] = delta_pos0 + delta_pos[indices] = self.forward_model.probe_displacements_to_object_pixels( + delta_pos0 + ) elif o_opt and p_opt and (not pos_opt): (delta_o, delta_p) = self.compute_updates_object_probe(op, p, psi_far, gradF, d) @@ -321,19 +323,19 @@ def gradient_o(self, p, gradF): tmp = torch.conj(p) * gradF - obj = self.parameter_group.object - o = obj.place_patches_on_empty_buffer( - self.positions, - tmp[:, 0], - pad_for_shift=self.options.forward_model_options.pad_for_shift + o_probe_grid = self.forward_model.place_object_patches_on_probe_grid( + self.positions, + tmp[:, 0], + integer_mode=False, ) - patches = obj.extract_patches_function( - o, self.positions + obj.pos_origin_coords, - self.parameter_group.probe.get_spatial_shape(), - pad=self.options.forward_model_options.pad_for_shift + patches = self.forward_model.extract_probe_grid_patches( + o_probe_grid, + self.positions, + self.parameter_group.probe.get_spatial_shape(), + integer_mode=False, ) patches = patches[:, None] - return o, patches + return self.forward_model.object_update_to_native_grid(o_probe_grid), patches def gradient_p(self, op, gradF): """Gradient with respect to the probe""" diff --git a/src/ptychi/reconstructors/dm.py b/src/ptychi/reconstructors/dm.py index b7a211b..344b22e 100644 --- a/src/ptychi/reconstructors/dm.py +++ b/src/ptychi/reconstructors/dm.py @@ -17,7 +17,6 @@ ) from ptychi.api.options.dm import DMReconstructorOptions from ptychi.timing.timer_utils import timer -import ptychi.image_proc as ip if TYPE_CHECKING: import ptychi.data_structures.parameter_group as pg @@ -200,15 +199,12 @@ def update_probe(self, probe_update: Tensor): def calculate_exit_wave_chunk( self, start_pt: int, end_pt: int, return_obj_patches: bool = False ) -> Union[Tensor, Tuple[Tensor, Tensor]]: - object_ = self.parameter_group.object - probe = self.parameter_group.probe - positions = self.parameter_group.probe_positions.tensor - - obj_patches = object_.extract_patches( - positions[start_pt:end_pt].round().int(), probe.get_spatial_shape(), integer_mode=True - ) + indices = torch.arange( + start_pt, end_pt, device=self.parameter_group.object.data.device + ).long() + obj_patches = self.forward_model.extract_object_patches(indices) psi = self.forward_model.forward_real_space( - indices=torch.arange(start_pt, end_pt, device=obj_patches.device).long(), + indices=indices, obj_patches=obj_patches, ) @@ -293,31 +289,41 @@ def calculate_object_update(self, start_pts: list[int], end_pts: list[int]) -> T # Calculate object update # Iterating over the object numerator calculation is more memory efficient - object_numerator = torch.zeros_like(object_.get_slice(0)) - object_denominator = torch.zeros_like(object_.get_slice(0), dtype=positions.dtype) + object_numerator = torch.zeros( + self.forward_model.probe_grid_object_shape, + dtype=object_.data.dtype, + device=object_.data.device, + ) + object_denominator = torch.zeros( + self.forward_model.probe_grid_object_shape, + dtype=positions.dtype, + device=positions.device, + ) for i in range(len(start_pts)): indices = torch.arange(start_pts[i], end_pts[i], device=object_numerator.device).long() p = self.forward_model.get_unique_probes(indices, always_return_probe_batch=True) p = self.forward_model.shift_unique_probes(indices, p, first_mode_only=True) - object_numerator = ip.place_patches_integer( - object_numerator, - positions[start_pts[i] : end_pts[i]].round().int() + object_.pos_origin_coords, - patches=(p.conj() * self.psi[start_pts[i] : end_pts[i]]).sum(1), - op="add", + batch_positions = positions[start_pts[i] : end_pts[i]] + object_numerator += self.forward_model.place_object_patches_on_probe_grid( + batch_positions, + (p.conj() * self.psi[start_pts[i] : end_pts[i]]).sum(1), + integer_mode=True, ) - object_denominator = ip.place_patches_integer( - object_denominator, - positions[start_pts[i] : end_pts[i]].round().int() + object_.pos_origin_coords, - patches=(p.abs() ** 2).sum(1), - op="add", + object_denominator += self.forward_model.place_object_patches_on_probe_grid( + batch_positions, + (p.abs() ** 2).sum(1), + integer_mode=True, ) # Calculate DM object update object_update = object_numerator / torch.sqrt( object_denominator**2 + (0.05 * object_denominator.max()) ** 2 ) + object_update = self.forward_model.object_update_to_native_grid( + object_update, normalized=True + ) # Apply inertia object_update = object_.get_slice(0) * object_.options.inertia + object_update * ( 1 - object_.options.inertia @@ -352,5 +358,4 @@ def get_positions_update_chunk( probe, None, ) - - return delta_pos + return self.forward_model.probe_displacements_to_object_pixels(delta_pos) diff --git a/src/ptychi/reconstructors/lsqml.py b/src/ptychi/reconstructors/lsqml.py index e4fa797..6007f8b 100644 --- a/src/ptychi/reconstructors/lsqml.py +++ b/src/ptychi/reconstructors/lsqml.py @@ -18,7 +18,6 @@ import ptychi.maths as pmath import ptychi.api.enums as enums from ptychi.timing.timer_utils import timer -import ptychi.image_proc as ip from ptychi.parallel import MultiprocessMixin if TYPE_CHECKING: @@ -1233,14 +1232,19 @@ def _combine_object_patch_update_directions( # Stitch all delta O patches on the object buffer # Shape of delta_o_hat: (h_whole, w_whole) - delta_o_hat = self.parameter_group.object.place_patches_on_empty_buffer( - positions.round().int(), delta_o_patches, integer_mode=True + delta_o_hat = self.forward_model.place_object_patches_on_probe_grid( + positions, delta_o_patches, integer_mode=True ) delta_o_hat = delta_o_hat[None, ...] if onto_accumulated: - delta_o_hat = delta_o_hat + ( - -self.parameter_group.object.get_grad()[slice_index : slice_index + 1] - ) + if self.forward_model.resampling_enabled: + delta_o_hat = delta_o_hat + self._object_gradient_probe_grid[ + slice_index : slice_index + 1 + ] + else: + delta_o_hat = delta_o_hat + ( + -self.parameter_group.object.get_grad()[slice_index : slice_index + 1] + ) return delta_o_hat @timer() @@ -1260,17 +1264,18 @@ def _precondition_object_update_direction( """ delta_o_hat = delta_o_hat[slice_index] - preconditioner = self.parameter_group.object.preconditioner + preconditioner = self.forward_model.get_probe_grid_preconditioner() delta_o_hat = delta_o_hat / torch.sqrt( preconditioner**2 + (preconditioner.max() * alpha_mix) ** 2 ) # Re-extract delta O patches if positions is not None: - delta_o_patches = ip.extract_patches_integer( + delta_o_patches = self.forward_model.extract_probe_grid_patches( delta_o_hat, - positions.round().int() + self.parameter_group.object.pos_origin_coords, + positions, self.parameter_group.probe.shape[-2:], + integer_mode=True, ) return delta_o_hat[None, ...], delta_o_patches[:, None, :, :] @@ -1285,11 +1290,17 @@ def _precondition_accumulated_object_update_direction(self): """ delta_o_hat_full = [] for i_slice in range(self.parameter_group.object.n_slices): + accumulated = ( + self._object_gradient_probe_grid[i_slice : i_slice + 1] + if self.forward_model.resampling_enabled + else -self.parameter_group.object.get_grad()[i_slice : i_slice + 1] + ) delta_o_hat = self._precondition_object_update_direction( - -self.parameter_group.object.get_grad()[i_slice : i_slice + 1], + accumulated, positions=None, alpha_mix=self.options.preconditioning_damping_factor ) + delta_o_hat = self.forward_model.object_update_to_native_grid(delta_o_hat) delta_o_hat_full.append(delta_o_hat) delta_o_hat_full = torch.cat(delta_o_hat_full, dim=0) return delta_o_hat_full @@ -1304,9 +1315,21 @@ def _initialize_object_gradient(self): """ if self.options.batching_mode in [enums.BatchingModes.RANDOM, enums.BatchingModes.UNIFORM]: self.parameter_group.object.initialize_grad() + if self.forward_model.resampling_enabled: + self._object_gradient_probe_grid = torch.zeros( + (self.parameter_group.object.n_slices, *self.forward_model.probe_grid_object_shape), + dtype=self.parameter_group.object.data.dtype, + device=self.parameter_group.object.data.device, + ) else: if self.current_minibatch == 0: self.parameter_group.object.initialize_grad() + if self.forward_model.resampling_enabled: + self._object_gradient_probe_grid = torch.zeros( + (self.parameter_group.object.n_slices, *self.forward_model.probe_grid_object_shape), + dtype=self.parameter_group.object.data.dtype, + device=self.parameter_group.object.data.device, + ) def _initialize_object_step_size_buffer(self): if self.current_minibatch == 0: @@ -1347,7 +1370,15 @@ def _record_object_slice_gradient(self, i_slice, delta_o_hat, add_to_existing=Fa """ Record the gradient of one slice of a multislice object. """ - if not add_to_existing: + if self.forward_model.resampling_enabled: + if not add_to_existing: + self._object_gradient_probe_grid[i_slice].copy_(delta_o_hat[0]) + else: + self._object_gradient_probe_grid[i_slice].add_(delta_o_hat[0]) + if self.options.batching_mode != enums.BatchingModes.COMPACT: + native_delta = self.forward_model.object_update_to_native_grid(delta_o_hat)[0] + self.parameter_group.object.set_grad(-native_delta, slicer=i_slice) + elif not add_to_existing: self.parameter_group.object.set_grad(-delta_o_hat[0], slicer=i_slice) else: self.parameter_group.object.set_grad( @@ -1502,6 +1533,8 @@ def update_probe_positions( unique_probes, self.parameter_group.object.step_size, ) + if getattr(self.forward_model, "resampling_enabled", False): + delta_pos = self.forward_model.probe_displacements_to_object_pixels(delta_pos) if self.parameter_group.probe_positions.options.momentum_acceleration_gain > 0: delta_pos = self._clip_probe_position_update(delta_pos) delta_pos = self._apply_probe_position_momentum(indices, delta_pos) @@ -1697,21 +1730,34 @@ def _combine_object_patch_update_directions( # Stitch all delta O patches on the object buffer # Shape of delta_o_hat: (h_whole, w_whole) - delta_o_hat = self.parameter_group.object.place_patches_on_empty_buffer( - positions.round().int(), delta_o_patches, integer_mode=True + delta_o_hat = self.forward_model.place_object_patches_on_probe_grid( + positions, delta_o_patches, integer_mode=True ) - # Synchronize buffer across ranks. - delta_o_hat = self.sync_buffer( - delta_o_hat, - op=dist.ReduceOp.SUM, - ) + # Compact mode keeps a rank-local probe-grid accumulator. Its smaller + # native-grid adjoint is synchronized once at epoch end. + if self.options.batching_mode != enums.BatchingModes.COMPACT: + delta_o_hat = self.sync_buffer( + delta_o_hat, + op=dist.ReduceOp.SUM, + ) delta_o_hat = delta_o_hat[None, ...] if onto_accumulated: - delta_o_hat = delta_o_hat + ( - -self.parameter_group.object.get_grad()[slice_index : slice_index + 1] - ) + if self.forward_model.resampling_enabled: + delta_o_hat = delta_o_hat + self._object_gradient_probe_grid[ + slice_index : slice_index + 1 + ] + else: + delta_o_hat = delta_o_hat + ( + -self.parameter_group.object.get_grad()[slice_index : slice_index + 1] + ) + return delta_o_hat + + def _precondition_accumulated_object_update_direction(self): + delta_o_hat = super()._precondition_accumulated_object_update_direction() + if self.forward_model.resampling_enabled: + delta_o_hat = self.sync_buffer(delta_o_hat, op=dist.ReduceOp.SUM) return delta_o_hat diff --git a/src/ptychi/reconstructors/pie.py b/src/ptychi/reconstructors/pie.py index ab9a38f..752d035 100644 --- a/src/ptychi/reconstructors/pie.py +++ b/src/ptychi/reconstructors/pie.py @@ -12,7 +12,6 @@ ) from ptychi.metrics import MSELossOfSqrt from ptychi.timing.timer_utils import timer -import ptychi.image_proc as ip if TYPE_CHECKING: import ptychi.api as api @@ -108,14 +107,14 @@ def compute_updates( step_weight = self.calculate_object_step_weight(unique_probes[i_slice]) delta_o_patches = step_weight * delta_exwv_i delta_o_patches = delta_o_patches.sum(1, keepdim=True) - delta_o_i = ip.place_patches_integer( - torch.zeros_like(object_.get_slice(0)), - positions.round().int() + object_.pos_origin_coords, + delta_o_i = self.forward_model.place_object_patches_on_probe_grid( + positions, delta_o_patches[:, 0], - op="add", + integer_mode=True, + ) + delta_o[i_slice, ...] = self.forward_model.object_update_to_native_grid( + delta_o_i ) - - delta_o[i_slice, ...] = delta_o_i delta_pos = None if (probe_positions.optimization_enabled(self.current_epoch) @@ -123,13 +122,16 @@ def compute_updates( and i_slice == self.parameter_group.probe_positions.get_slice_for_correction(object_.n_slices) ): delta_pos = torch.zeros_like(probe_positions.data) - delta_pos[indices] = probe_positions.position_correction.get_update( + delta_pos_i = probe_positions.position_correction.get_update( delta_exwv_i, obj_patches[:, i_slice : i_slice + 1, ...], delta_o_patches, self.forward_model.intermediate_variables.shifted_unique_probes[i_slice], object_.step_size, ) + delta_pos[indices] = self.forward_model.probe_displacements_to_object_pixels( + delta_pos_i + ) delta_p_i = None if (i_slice == 0) and (probe.optimization_enabled(self.current_epoch)): diff --git a/src/ptychi/reconstructors/raar.py b/src/ptychi/reconstructors/raar.py index edf1c5c..881aad3 100644 --- a/src/ptychi/reconstructors/raar.py +++ b/src/ptychi/reconstructors/raar.py @@ -9,7 +9,6 @@ from torch.utils.data import DataLoader, Dataset import ptychi.api as api -import ptychi.image_proc as ip from ptychi.api.options.raar import RAARReconstructorOptions from ptychi.reconstructors.base import ( AnalyticalIterativePtychographyReconstructor, @@ -127,7 +126,7 @@ def compute_updates(self, y_true: Tensor) -> Tensor: q_object = self.calculate_object_projection(self.psi, start_pts, end_pts) pa_object = self.calculate_data_projected_object_projection(y_true, start_pts, end_pts) else: - q_object = object_.get_slice(0) + q_object = self.forward_model.get_probe_grid_object()[0] pa_object = q_object raar_error_squared, exit_wave_update = self.apply_raar_exit_wave_update( @@ -177,14 +176,10 @@ def initialize_exit_wave(self, start_pts: list[int], end_pts: list[int]) -> None @timer() def calculate_exit_wave_chunk(self, start_pt: int, end_pt: int) -> Tensor: - object_ = self.parameter_group.object - positions = self.parameter_group.probe_positions.tensor - obj_patches = object_.extract_patches( - positions[start_pt:end_pt].round().int(), - self.parameter_group.probe.get_spatial_shape(), - integer_mode=True, - ) - indices = torch.arange(start_pt, end_pt, device=obj_patches.device).long() + indices = torch.arange( + start_pt, end_pt, device=self.parameter_group.object.data.device + ).long() + obj_patches = self.forward_model.extract_object_patches(indices) return self.forward_model.forward_real_space(indices=indices, obj_patches=obj_patches) @timer() @@ -198,10 +193,17 @@ def apply_data_projection(self, exit_wave: Tensor, y_true: Tensor) -> Tensor: return self.forward_model.free_space_propagator.propagate_backward(propagated_exit_wave) def initialize_object_projection_terms(self) -> Tuple[Tensor, Tensor]: - object_ = self.parameter_group.object positions = self.parameter_group.probe_positions.tensor - numerator = torch.zeros_like(object_.get_slice(0)) - denominator = torch.zeros_like(object_.get_slice(0), dtype=positions.dtype) + numerator = torch.zeros( + self.forward_model.probe_grid_object_shape, + dtype=self.parameter_group.object.data.dtype, + device=positions.device, + ) + denominator = torch.zeros( + self.forward_model.probe_grid_object_shape, + dtype=positions.dtype, + device=positions.device, + ) return numerator, denominator @timer() @@ -213,7 +215,6 @@ def add_to_object_projection_terms( start_pt: int, end_pt: int, ) -> None: - object_ = self.parameter_group.object positions = self.parameter_group.probe_positions.tensor indices = torch.arange(start_pt, end_pt, device=numerator.device).long() shifted_probe = self.forward_model.get_unique_probes( @@ -222,32 +223,28 @@ def add_to_object_projection_terms( shifted_probe = self.forward_model.shift_unique_probes( indices, shifted_probe, first_mode_only=True ) - placement_positions = positions[start_pt:end_pt].round().int() + object_.pos_origin_coords numerator.copy_( - ip.place_patches_integer( - numerator, - placement_positions, - patches=(shifted_probe.conj() * exit_wave).sum(1), - op="add", + numerator + self.forward_model.place_object_patches_on_probe_grid( + positions[start_pt:end_pt], + (shifted_probe.conj() * exit_wave).sum(1), + integer_mode=True, ) ) denominator.copy_( - ip.place_patches_integer( - denominator, - placement_positions, - patches=(shifted_probe.abs() ** 2).sum(1), - op="add", + denominator + self.forward_model.place_object_patches_on_probe_grid( + positions[start_pt:end_pt], + (shifted_probe.abs() ** 2).sum(1), + integer_mode=True, ) ) def finish_object_projection(self, numerator: Tensor, denominator: Tensor) -> Tensor: - object_ = self.parameter_group.object tiny = torch.finfo(denominator.dtype).tiny projected_object = numerator / denominator.clamp_min(tiny) projected_object = torch.where( denominator > tiny, projected_object, - object_.get_slice(0), + self.forward_model.get_probe_grid_object()[0], ) return projected_object @@ -294,13 +291,13 @@ def calculate_data_projected_object_projection( def synthesize_exit_wave_chunk( self, projected_object: Tensor, start_pt: int, end_pt: int ) -> Tensor: - object_ = self.parameter_group.object positions = self.parameter_group.probe_positions.tensor indices = torch.arange(start_pt, end_pt, device=projected_object.device).long() - obj_patches = ip.extract_patches_integer( + obj_patches = self.forward_model.extract_probe_grid_patches( projected_object, - positions[start_pt:end_pt].round().int() + object_.pos_origin_coords, + positions[start_pt:end_pt], self.parameter_group.probe.get_spatial_shape(), + integer_mode=True, ) shifted_probe = self.forward_model.get_unique_probes( indices, always_return_probe_batch=True @@ -344,18 +341,13 @@ def apply_raar_exit_wave_update( @timer() def update_variable_probe(self, exit_wave_update: Tensor) -> None: """Update OPR eigenmodes and weights from the RAAR exit-wave update.""" - object_ = self.parameter_group.object probe = self.parameter_group.probe probe_positions = self.parameter_group.probe_positions opr_mode_weights = self.parameter_group.opr_mode_weights indices = torch.arange( probe_positions.n_scan_points, device=probe_positions.data.device ).long() - obj_patches = object_.extract_patches( - probe_positions.tensor.round().int(), - probe.get_spatial_shape(), - integer_mode=True, - ) + obj_patches = self.forward_model.extract_object_patches(indices) delta_p_i = obj_patches.conj() * exit_wave_update delta_p_i = self.adjoint_shift_probe_update_direction( @@ -376,6 +368,9 @@ def update_variable_probe(self, exit_wave_update: Tensor) -> None: @timer() def update_object(self, projected_object: Tensor) -> None: object_ = self.parameter_group.object + projected_object = self.forward_model.object_update_to_native_grid( + projected_object, normalized=True + ) object_update = ( object_.options.inertia * object_.get_slice(0) + (1 - object_.options.inertia) * projected_object @@ -390,19 +385,13 @@ def update_object(self, projected_object: Tensor) -> None: @timer() def calculate_probe_update(self, start_pts: list[int], end_pts: list[int]) -> Tensor: - object_ = self.parameter_group.object probe = self.parameter_group.probe - positions = self.parameter_group.probe_positions.tensor numerator = torch.zeros_like(probe.get_opr_mode(0)) denominator = torch.zeros_like(probe.get_opr_mode(0).abs()) for start_pt, end_pt in zip(start_pts, end_pts): indices = torch.arange(start_pt, end_pt, device=numerator.device).long() - obj_patches = object_.extract_patches( - positions[start_pt:end_pt].round().int(), - probe.get_spatial_shape(), - integer_mode=True, - ) + obj_patches = self.forward_model.extract_object_patches(indices) numerator_update = obj_patches.conj() * self.psi[start_pt:end_pt] denominator_update = obj_patches.abs() ** 2 numerator_update = self.adjoint_shift_probe_update_direction( @@ -428,19 +417,13 @@ def calculate_probe_update(self, start_pts: list[int], end_pts: list[int]) -> Te @timer() def update_probe_positions(self, start_pts: list[int], end_pts: list[int]) -> None: - object_ = self.parameter_group.object - probe = self.parameter_group.probe probe_positions = self.parameter_group.probe_positions positions = probe_positions.tensor delta_pos = torch.zeros_like(probe_positions.data) for start_pt, end_pt in zip(start_pts, end_pts): indices = torch.arange(start_pt, end_pt, device=positions.device).long() - obj_patches = object_.extract_patches( - positions[start_pt:end_pt].round().int(), - probe.get_spatial_shape(), - integer_mode=True, - ) + obj_patches = self.forward_model.extract_object_patches(indices) model_exit_wave = self.forward_model.forward_real_space( indices=indices, obj_patches=obj_patches ) @@ -460,10 +443,11 @@ def get_positions_update_chunk( probe_positions = self.parameter_group.probe_positions probe = self.forward_model.get_unique_probes(indices, always_return_probe_batch=True) probe = self.forward_model.shift_unique_probes(indices, probe, first_mode_only=True) - return probe_positions.position_correction.get_update( + delta_pos = probe_positions.position_correction.get_update( chi, obj_patches, None, probe, None, ) + return self.forward_model.probe_displacements_to_object_pixels(delta_pos) diff --git a/src/ptychi/workflows/progressive_resolution.py b/src/ptychi/workflows/progressive_resolution.py index e762e61..0f5240d 100644 --- a/src/ptychi/workflows/progressive_resolution.py +++ b/src/ptychi/workflows/progressive_resolution.py @@ -39,6 +39,10 @@ def run(self) -> None: level_options.object_options.pixel_size_m = ( self.task_options.object_options.pixel_size_m * factor ) + if self.task_options.probe_options.pixel_size_m is not None: + level_options.probe_options.pixel_size_m = ( + self.task_options.probe_options.pixel_size_m * factor + ) task = PtychographyTask( level_options, diff --git a/tests/test_probe_pixel_geometry.py b/tests/test_probe_pixel_geometry.py new file mode 100644 index 0000000..c6284d4 --- /dev/null +++ b/tests/test_probe_pixel_geometry.py @@ -0,0 +1,241 @@ +import json + +import pytest +import torch +from pydantic import ValidationError + +import ptychi.api as api +import ptychi.image_proc as image_proc +from ptychi.api.task import PtychographyTask + + +def _small_task(options): + options.reconstructor_options.default_device = api.Devices.CPU + options.reconstructor_options.num_epochs = 1 + options.reconstructor_options.allow_nondeterministic_algorithms = False + if hasattr(options.reconstructor_options, "batch_size"): + options.reconstructor_options.batch_size = 2 + if hasattr(options.reconstructor_options, "chunk_length"): + options.reconstructor_options.chunk_length = 1 + options.data_options.fft_shift = False + options.object_options.pixel_size_m = 1.0 + options.object_options.pixel_size_aspect_ratio = 1.0 + options.probe_options.pixel_size_m = 0.8 + options.probe_options.pixel_size_aspect_ratio = 1.2 + options.probe_options.optimizable = False + options.probe_position_options.optimizable = False + + diffraction = torch.rand(2, 5, 5, device="cpu") + 0.2 + object_data = torch.ones((1, 17, 19), dtype=torch.complex64, device="cpu") + probe_data = torch.ones((1, 1, 5, 5), dtype=torch.complex64, device="cpu") + position_x = torch.tensor([-2.0, 2.0], device="cpu") + position_y = torch.tensor([-1.0, 1.0], device="cpu") + original_x = position_x.clone() + original_y = position_y.clone() + + task = PtychographyTask( + options, + diffraction_data=diffraction, + object_data=object_data, + probe_data=probe_data, + probe_position_x_px=position_x, + probe_position_y_px=position_y, + ) + assert torch.equal(position_x, original_x) + assert torch.equal(position_y, original_y) + return task + + +def test_probe_pixel_options_validate_inherit_and_serialize(): + options = api.LSQMLOptions() + assert options.probe_options.pixel_size_m is None + assert options.probe_options.pixel_size_aspect_ratio is None + + options.probe_options.pixel_size_m = 2e-9 + options.probe_options.pixel_size_aspect_ratio = 1.5 + values = options.get_dict()["probe_options"] + assert values["pixel_size_m"] == 2e-9 + assert values["pixel_size_aspect_ratio"] == 1.5 + json.dumps(options.get_dict()) + + with pytest.raises(ValidationError): + options.probe_options.pixel_size_m = 0 + with pytest.raises(ValidationError): + options.probe_options.pixel_size_aspect_ratio = -1 + + +@pytest.mark.parametrize( + ("source_shape", "target_shape"), + [ + ((5, 6), (8, 3)), + ((4, 4), (7, 7)), + ((7, 8), (3, 10)), + ((5, 5), (4, 4)), + ((4, 5), (5, 4)), + ], +) +def test_fourier_resize_preserves_constants_and_has_exact_adjoint( + source_shape, target_shape +): + constant = torch.ones(source_shape) + resized_constant = image_proc.fourier_resize(constant, target_shape) + assert torch.allclose(resized_constant, torch.ones(target_shape), atol=1e-6) + + source = torch.randn(source_shape, dtype=torch.complex64) + target = torch.randn(target_shape, dtype=torch.complex64) + resized = image_proc.fourier_resize(source, target_shape) + adjoint = image_proc.fourier_resize(target, source_shape, adjoint=True) + lhs = torch.vdot(resized.reshape(-1), target.reshape(-1)) + rhs = torch.vdot(source.reshape(-1), adjoint.reshape(-1)) + assert torch.allclose(lhs, rhs, atol=2e-5, rtol=2e-5) + + +def test_anisotropic_geometry_and_versioned_caches(monkeypatch): + task = _small_task(api.EPIEOptions()) + forward_model = task.reconstructor.forward_model + expected_shape = ( + round(17 / (0.8 / 1.2)), + round(19 / 0.8), + ) + assert forward_model.probe_grid_object_shape == expected_shape + assert forward_model.probe_grid_scale == pytest.approx( + (expected_shape[0] / 17, expected_shape[1] / 19) + ) + + resize_calls = 0 + original_resize = image_proc.fourier_resize + + def counted_resize(*args, **kwargs): + nonlocal resize_calls + resize_calls += 1 + return original_resize(*args, **kwargs) + + monkeypatch.setattr(image_proc, "fourier_resize", counted_resize) + task.object.preconditioner = torch.ones(task.object.lateral_shape) + with torch.no_grad(): + first_object = forward_model.get_probe_grid_object() + assert forward_model.get_probe_grid_object() is first_object + first_preconditioner = forward_model.get_probe_grid_preconditioner() + assert forward_model.get_probe_grid_preconditioner() is first_preconditioner + assert resize_calls == 2 + + with torch.no_grad(): + task.object.set_data(task.object.data + 1) + forward_model.get_probe_grid_object() + task.object.preconditioner.add_(1) + forward_model.get_probe_grid_preconditioner() + assert resize_calls == 4 + assert torch.all(forward_model.get_probe_grid_preconditioner() >= 0) + assert not any("resampling" in key for key in forward_model.state_dict()) + + +def test_inherited_probe_geometry_uses_identity_path(): + task = _small_task(api.EPIEOptions()) + forward_model = task.reconstructor.forward_model + task.probe.options.pixel_size_m = None + task.probe.options.pixel_size_aspect_ratio = None + indices = torch.arange(task.probe_positions.n_scan_points) + + actual = forward_model.extract_object_patches(indices) + expected = task.object.extract_patches( + task.probe_positions.tensor.round().int(), + task.probe.get_spatial_shape(), + integer_mode=True, + ) + assert not forward_model.resampling_enabled + assert forward_model.probe_grid_object_shape == task.object.lateral_shape + assert torch.equal(actual, expected) + + +def test_object_patch_extraction_accepts_an_explicit_object_array(): + task = _small_task(api.EPIEOptions()) + replacement = torch.full_like(task.object.data, 3 + 2j) + positions = task.probe_positions.tensor.round().int() + + patches = task.object.extract_patches( + positions, + task.probe.get_spatial_shape(), + integer_mode=True, + object_array=replacement, + ) + + assert torch.all(patches == 3 + 2j) + assert torch.all(task.object.data == 1) + + +def test_autodiff_graph_cache_policy(): + task = _small_task(api.AutodiffPtychographyOptions()) + forward_model = task.reconstructor.forward_model + + first = forward_model.get_probe_grid_object() + second = forward_model.get_probe_grid_object() + assert first is not second + + task.object.set_optimizable(False) + third = forward_model.get_probe_grid_object() + fourth = forward_model.get_probe_grid_object() + assert third is fourth + + +@pytest.mark.parametrize( + ("options_factory", "integer_mode"), + [(api.EPIEOptions, True), (api.BHOptions, False)], +) +def test_resampled_patch_extraction_and_placement_are_adjoint( + options_factory, integer_mode +): + options = options_factory() + if not integer_mode: + options.reconstructor_options.forward_model_options.pad_for_shift = 1 + task = _small_task(options) + forward_model = task.reconstructor.forward_model + indices = torch.arange(task.probe_positions.n_scan_points) + + source = torch.randn_like(task.object.data) + with torch.no_grad(): + task.object.set_data(source) + patches = forward_model.extract_object_patches(indices)[:, 0] + cotangent = torch.randn_like(patches) + placed = forward_model.place_object_patches_on_probe_grid( + task.probe_positions.tensor, + cotangent, + integer_mode=integer_mode, + ) + native_adjoint = forward_model.object_update_to_native_grid(placed) + + lhs = torch.vdot(patches.reshape(-1), cotangent.reshape(-1)) + rhs = torch.vdot(source.reshape(-1), native_adjoint.reshape(-1)) + assert torch.allclose(lhs, rhs, atol=3e-4, rtol=3e-4) + + +@pytest.mark.parametrize( + "options_factory", + [ + api.AutodiffPtychographyOptions, + api.EPIEOptions, + api.BHOptions, + api.DMOptions, + api.RAAROptions, + ], +) +def test_lightweight_mismatched_grid_reconstruction(options_factory): + options = options_factory() + if isinstance(options, api.BHOptions): + options.reconstructor_options.forward_model_options.pad_for_shift = 1 + task = _small_task(options) + task.run() + assert torch.isfinite(task.get_data_to_cpu("object")).all() + + +@pytest.mark.parametrize( + "batching_mode", + [api.BatchingModes.RANDOM, api.BatchingModes.UNIFORM, api.BatchingModes.COMPACT], +) +def test_lightweight_mismatched_grid_lsqml(batching_mode): + options = api.LSQMLOptions() + options.reconstructor_options.rescale_probe_intensity_in_first_epoch = False + options.reconstructor_options.batching_mode = batching_mode + options.reconstructor_options.batch_size = 1 + task = _small_task(options) + task.run() + assert torch.isfinite(task.get_data_to_cpu("object")).all() diff --git a/tests/test_progressive_resolution_workflow.py b/tests/test_progressive_resolution_workflow.py index 4598bd8..7e13f6f 100644 --- a/tests/test_progressive_resolution_workflow.py +++ b/tests/test_progressive_resolution_workflow.py @@ -162,6 +162,8 @@ def test_workflow_runs_rounded_levels_and_transfers_results(monkeypatch): monkeypatch.setattr(progressive_resolution_module, "PtychographyTask", _FakeTask) data = _workflow_data() original_options = _task_options() + original_options.probe_options.pixel_size_m = 1.25 + original_options.probe_options.pixel_size_aspect_ratio = 1.4 workflow = ProgressiveResolutionWorkflow( original_options, workflow_options=api.ProgressiveResolutionWorkflowOptions( @@ -192,6 +194,16 @@ def test_workflow_runs_rounded_levels_and_transfers_results(monkeypatch): ] assert original_options.reconstructor_options.num_epochs == 100 assert original_options.object_options.pixel_size_m == 2.5 + assert [task.options.probe_options.pixel_size_m for task in workflow.tasks] == [ + 5.0, + 2.5, + 1.25, + ] + assert all( + task.options.probe_options.pixel_size_aspect_ratio == 1.4 + for task in workflow.tasks + ) + assert original_options.probe_options.pixel_size_m == 1.25 assert [tuple(task.input_data["diffraction_data"].shape[-2:]) for task in workflow.tasks] == [ (2, 3),