Allow user to specify output domain (fourier) for pipeline

.gitignore

@@ -134,3 +134,8 @@ _version_save.py
 
 # uv package manager
 uv.lock
+
+# ai related
+README_ai.md
+.claude/
+CLAUDE.md

CHANGELOG

@@ -1,4 +1,5 @@
-0.6.2
+0.7.0
+ - feat: allow user to specify (output) domain to skip ifft (#28, #29)
  - docs: clarify `filter_size` in `filter.get_filter_array` (#6)
  - enh: add 'physical radius' as disk filter option (#3, #26)
  - ref: use get_available_interfaces in get_best_interface (#14, #25)

docs/sec_code_reference.rst

@@ -47,6 +47,15 @@ Cupy
    :inherited-members:
 
 
+.. _sec_code_fourier_field_data:
+
+Fourier Field Data
+------------------
+.. automodule:: qpretrieve.fourier.fourier_field_data
+   :members:
+   :inherited-members:
+
+
 .. _sec_code_ifer:
 
 Interference image analysis

docs/sec_fourier_domain_pipeline.rst

@@ -0,0 +1,69 @@
+.. _sec_fourier_domain_pipeline:
+
+=======================
+Fourier Domain Pipeline
+=======================
+
+Since version 0.7.0, :meth:`.OffAxisHologram.run_pipeline` accepts an
+``output_domain`` keyword argument. By default (``output_domain="spatial"``) the
+pipeline returns the reconstructed complex field as usual. When
+``output_domain="fourier"`` is set, the inverse FFT is skipped and a
+:class:`~qpretrieve.fourier.fourier_field_data.FourierFieldData` is returned instead.
+
+The :class:`~qpretrieve.fourier.fourier_field_data.FourierFieldData` holds the
+filtered Fourier data and all the metadata needed to reconstruct the spatial field.
+Calling its :meth:`~qpretrieve.fourier.fourier_field_data.FourierFieldData.finalize`
+method performs the inverse FFT and returns the spatial field identically to the
+default path.
+
+Spatial Output (Default)
+------------------------
+
+.. code-block:: python
+
+    import numpy as np
+    import qpretrieve
+
+    edata = np.load("examples/data/hologram_cell.npz")
+    oah = qpretrieve.OffAxisHologram(edata["data"])
+    field = oah.run_pipeline()      # returns complex spatial field
+    print(type(field))              # numpy.ndarray
+    print(field.shape)              # (1, H, W)
+
+Fourier Output (skipping the inverse FFT)
+-----------------------------------------
+
+This is useful when combined with field propagation, see the Note above.
+
+.. code-block:: python
+
+    import numpy as np
+    import qpretrieve
+
+    edata = np.load("examples/data/hologram_cell.npz")
+    oah = qpretrieve.OffAxisHologram(edata["data"])
+    fourier_data = oah.run_pipeline(output_domain="fourier")
+
+    # The fourier_data carries the filtered Fourier data.
+    # Call finalize() to recover the spatial field when needed.
+    field = fourier_data.finalize()
+    print(field.shape)              # (1, H, W) — identical to the default path
+
+
+.. admonition:: Combining `qpretrieve` and `nrefocus` pipelines
+
+    The Fourier output is most useful when the result is passed directly to a
+    wave propagation library such as `nrefocus
+    <https://nrefocus.readthedocs.io>`_, which can consume the
+    :class:`~qpretrieve.fourier.fourier_field_data.FourierFieldData` object and
+    skip its own forward FFT. This avoids a redundant iFFT + FFT pair at the
+    qpretrieve/nrefocus boundary.
+    For an nrefocus-integrated working example see the :ref:`sec_examples`.
+
+    *Comparing Spatial vs. Fourier Pipeline*
+
+    For unpadded, square spatial input data, the default spatial domain
+    pipeline and the fourier domain pipeline are identical. There is only
+    floating point imprecision.
+    For padded pipelines, the pipelines are not identical due to padding and
+    unpadding causing inconsistencies at the boundary of the images.

docs/sec_getting_started.rst

@@ -9,4 +9,5 @@ Getting started
   sec_basic_use
   sec_array_layout
   sec_ndarray_backend
+  sec_fourier_domain_pipeline
   sec_userapi

qpretrieve/__init__.py

@@ -2,5 +2,6 @@
 from ._version import version as __version__
 from ._ndarray_backend import get_ndarray_backend, set_ndarray_backend
 from .interfere import OffAxisHologram, QLSInterferogram
+from .fourier import FourierFieldData
 from . import filter
 from . import fourier

qpretrieve/fourier/__init__.py

@@ -3,6 +3,7 @@
 from typing import Type
 
 from .base import FFTFilter
+from .fourier_field_data import FourierFieldData, finalize_fourier_field
 from .ff_numpy import FFTFilterNumpy
 
 try:

qpretrieve/fourier/base.py

@@ -8,6 +8,7 @@
 from .. import filter
 from ..utils import padding_3d, mean_3d
 from ..data_array_layout import convert_data_to_3d_array_layout
+from .fourier_field_data import FourierFieldData, finalize_fourier_field
 
 
 class FFTCache:
@@ -215,7 +216,9 @@ def backend_check(self):
 
     def filter(self, filter_name: str, filter_size: float,
                freq_pos: (float, float),
-               scale_to_filter: bool | float = False) -> xp.ndarray:
+               scale_to_filter: bool | float = False,
+               output_domain: str = "spatial"
+               ) -> xp.ndarray | FourierFieldData:
         """
         Parameters
         ----------
@@ -250,6 +253,12 @@ def filter(self, filter_name: str, filter_size: float,
             a boolean array but a floating-point array), the higher you
             set `scale_to_filter`, the more information will be included
             in the scaled image.
+        output_domain: str
+            Either ``"spatial"`` or ``"fourier"``. Spatial returns the
+            inverse-transformed field, Fourier returns a
+            :class:`~qpretrieve.fourier.fourier_field_data.FourierFieldData`.
+
+            .. versionadded:: 0.7.0
 
         Notes
         -----
@@ -271,6 +280,9 @@ def filter(self, filter_name: str, filter_size: float,
                                 str(self.dtype_conversion),
                                 ])
 
+        if output_domain not in ("spatial", "fourier"):
+            raise ValueError("`output_domain` must be 'spatial' or 'fourier'.")
+
         inv_data = FFTCache.get_item(weakref_key)
 
         if inv_data is not None:
@@ -302,30 +314,39 @@ def filter(self, filter_name: str, filter_size: float,
                 # We now have the interesting peak already shifted to
                 # the first entry of our array in `shifted`.
                 fft_used = fft_used[:, cslice, cslice]
-
-            field = self._ifft(xp.fft.ifftshift(fft_used, axes=(-2, -1)))
-
-            if self.padding:
-                # revert padding
-                sx, sy = self.origin.shape[-2:]
-                if scale_to_filter:
-                    sx = int(xp.ceil(sx * 2 * crad / osize))
-                    sy = int(xp.ceil(sy * 2 * crad / osize))
-
-                field = field[:, :sx, :sy]
-
-                if scale_to_filter:
-                    # Scale the absolute value of the field. This does not
-                    # have any influence on the phase, but on the amplitude.
-                    field *= (2 * crad / osize) ** 2
-            # Add FFT to cache
-            # (The cache will only be cleared if this instance is deleted)
-            FFTCache.add_item(weakref_key, self.fft_origin,
-                              (filt_array, fft_used, field))
+            field = None
 
         self.fft_filtered[:] = fft_filtered
         self.fft_used = fft_used
-        return field
+        if output_domain == "spatial":
+            if field is None:
+                field = finalize_fourier_field(
+                    fft_in=fft_used,
+                    ifft_fn=self._ifft,
+                    input_shape=self.origin.shape[-2:],
+                    fft_shape=self.fft_origin.shape[-2:],
+                    padding=self.padding,
+                    scale_to_filter=scale_to_filter,
+                    crop_radius=fft_used.shape[-2] // 2 if scale_to_filter else None,  # noqa: E501
+                )
+                FFTCache.add_item(weakref_key, self.fft_origin,
+                                  (filt_array, fft_used, field))
+            return field
+
+        # Preserve the FFT intermediates so the field can be
+        # compute later without recomputing the filter.
+        FFTCache.add_item(weakref_key, self.fft_origin,
+                          (filt_array, fft_used, None))
+        crop_radius = fft_used.shape[-2] // 2 if scale_to_filter else None
+        return FourierFieldData(
+            fft_used=fft_used,
+            ifft_fn=self._ifft,
+            input_shape=self.origin.shape[-2:],
+            fft_shape=self.fft_origin.shape[-2:],
+            padding=self.padding,
+            scale_to_filter=scale_to_filter,
+            crop_radius=crop_radius,
+        )
 
     def _result_type(self, dtype_in) -> xp.dtype:
         """Wrapper on `np.result_type` to provide correct fft dtype"""