FEAT: allow probe pixel size to be different from object; in that case, object is Fourier-rescaled to match the probe sampling in the forward model and the adjoint is calculated accordingly

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)``
      -

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
 ----------------
 

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
     )

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()

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(

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))