diff --git a/.gitignore b/.gitignore index 33d3fee..45f0c47 100644 --- a/.gitignore +++ b/.gitignore @@ -134,3 +134,8 @@ _version_save.py # uv package manager uv.lock + +# ai related +README_ai.md +.claude/ +CLAUDE.md diff --git a/CHANGELOG b/CHANGELOG index 70fb10c..d2d4332 100644 --- a/CHANGELOG +++ b/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) diff --git a/docs/sec_code_reference.rst b/docs/sec_code_reference.rst index c5178d3..cf6aa23 100644 --- a/docs/sec_code_reference.rst +++ b/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 diff --git a/docs/sec_fourier_domain_pipeline.rst b/docs/sec_fourier_domain_pipeline.rst new file mode 100644 index 0000000..c9f1c63 --- /dev/null +++ b/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 + `_, 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. diff --git a/docs/sec_getting_started.rst b/docs/sec_getting_started.rst index bcdfacf..026b01e 100644 --- a/docs/sec_getting_started.rst +++ b/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 diff --git a/qpretrieve/__init__.py b/qpretrieve/__init__.py index b099735..c77d07b 100644 --- a/qpretrieve/__init__.py +++ b/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 diff --git a/qpretrieve/fourier/__init__.py b/qpretrieve/fourier/__init__.py index 630e05d..931d1fb 100644 --- a/qpretrieve/fourier/__init__.py +++ b/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: diff --git a/qpretrieve/fourier/base.py b/qpretrieve/fourier/base.py index 5f3b689..c195b54 100644 --- a/qpretrieve/fourier/base.py +++ b/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""" diff --git a/qpretrieve/fourier/fourier_field_data.py b/qpretrieve/fourier/fourier_field_data.py new file mode 100644 index 0000000..3d08d97 --- /dev/null +++ b/qpretrieve/fourier/fourier_field_data.py @@ -0,0 +1,187 @@ +""" + .. versionadded:: 0.7.0 +""" +from __future__ import annotations + +from dataclasses import dataclass, field as dataclass_field +from typing import Callable + +from .._ndarray_backend import xp + + +def finalize_fourier_field( + fft_in: xp.ndarray, + ifft_fn: Callable[[xp.ndarray], xp.ndarray], + input_shape: tuple[int, int], + fft_shape: tuple[int, int], + padding: bool | int, + scale_to_filter: bool | float, + crop_radius: int | None = None, +) -> xp.ndarray: + """Convert a propagated Fourier field back to the spatial domain. + + Applies ``ifftshift`` to move DC from the centre to index 0, then + calls the inverse FFT. If padding or Fourier-space cropping was used + during filtering the result is trimmed and rescaled to match the + original hologram dimensions. + + Parameters + ---------- + fft_in : ndarray + Fourier-domain data in the **fftshifted** convention (DC at centre), + shaped ``(..., fy, fx)``. + ifft_fn : callable + Inverse FFT function that accepts an ndarray and returns an ndarray + of the same shape. Must match the ndarray backend (numpy, cupy, …) + used when the data was created. + input_shape : tuple of int + Spatial shape ``(sy, sx)`` of the original (unpadded) hologram. + fft_shape : tuple of int + Shape of the FFT array ``(..., fy, fx)`` *before* any Fourier-space + cropping was applied. + padding : bool or int + Whether boundary padding was applied during filtering. If truthy, + the spatial field is cropped back to ``input_shape`` after the iFFT. + scale_to_filter : bool or float + Whether (or by how much) the Fourier array was cropped to the filter + support. If truthy, the crop window and amplitude scale factor are + recomputed from ``crop_radius`` and ``fft_shape``. + crop_radius : int or None + Radius in Fourier pixels of the crop window used when + ``scale_to_filter`` is active. Required when ``scale_to_filter`` is + truthy; ignored otherwise. + + Returns + ------- + field : ndarray + Complex spatial field, shaped ``(..., sy, sx)``. + """ + field = ifft_fn(xp.fft.ifftshift(fft_in, axes=(-2, -1))) + + if padding: + sx, sy = input_shape + if scale_to_filter: + if crop_radius is None: + raise ValueError( + "crop_radius is required when scale_to_filter is set") + osize = fft_shape[-1] + sx = int(xp.ceil(sx * 2 * crop_radius / osize)) + sy = int(xp.ceil(sy * 2 * crop_radius / osize)) + + field = field[:, :sx, :sy] + + if scale_to_filter: + osize = fft_shape[-1] + field *= (2 * crop_radius / osize) ** 2 + + return field + + +@dataclass(slots=True) +class FourierFieldData: + """Fourier-domain field data returned by ``output_domain="fourier"``. + + Produced by :meth:`.OffAxisHologram.run_pipeline` (and the underlying + :meth:`.FFTFilter.filter`) when ``output_domain="fourier"`` is requested. + Instead of performing the inverse FFT, qpretrieve packages the filtered + Fourier data together with the reconstruction metadata needed to recover + the spatial field on demand. + + Call :meth:`finalize` to apply the inverse FFT and obtain the spatial + field. Downstream libraries such as `nrefocus` can consume this object + directly via duck-typing (they detect :attr:`fft_used` and skip their own + forward FFT, avoiding a redundant iFFT + FFT pair at the pipeline + boundary). + + .. versionadded:: 0.7.0 + + Attributes + ---------- + fft_used : ndarray + Filtered FFT data in the **fftshifted** convention (DC at the centre + of the array). Shape is ``(..., fy, fx)``, where ``fy`` and ``fx`` + are reduced relative to the full FFT when ``scale_to_filter`` is set. + ifft_fn : callable + Inverse FFT callable that matches the ndarray backend (numpy, cupy, + …) used during filtering. + input_shape : tuple of int + Spatial shape ``(sy, sx)`` of the original unpadded hologram. + fft_shape : tuple of int + Shape of the full FFT array ``(..., fy0, fx0)`` before any + Fourier-space cropping was applied. + padding : bool or int + Whether boundary padding was applied during filtering. + scale_to_filter : bool or float + Whether (or by what factor) the Fourier array was cropped to the + filter support. Controls the crop and amplitude rescaling in + :func:`finalize_fourier_field`. + crop_radius : int or None + Radius in Fourier pixels of the crop window, set when + ``scale_to_filter`` is active. + output_domain : str + Domain hint used by nrefocus for duck-typing. Always ``"spatial"`` + — indicates that :meth:`finalize` should be called to convert to the + spatial domain. + """ + + fft_used: xp.ndarray + ifft_fn: Callable[[xp.ndarray], xp.ndarray] + input_shape: tuple[int, int] + fft_shape: tuple[int, int] + padding: bool | int + scale_to_filter: bool | float + crop_radius: int | None = None + output_domain: str = "spatial" + _field: xp.ndarray | None = dataclass_field( + default=None, init=False, repr=False) + + def finalize(self, propagated_fft: xp.ndarray | None = None) -> xp.ndarray: + """Return the reconstructed spatial field. + + Applies ``ifftshift`` + inverse FFT to convert from the Fourier + domain, then crops and rescales to match the original hologram + dimensions (reversing any padding or Fourier-space cropping applied + during filtering). + + The result is cached in :attr:`field` after the first call. + + Parameters + ---------- + propagated_fft : ndarray or None + Optional Fourier-domain array in the fftshifted convention, + typically the output of a wave propagation step (e.g. from + nrefocus). When provided, this is inverse-transformed instead + of :attr:`fft_used`, allowing the propagated field to be + reconstructed with the same cropping/scaling as the original. + If ``None``, :attr:`fft_used` is used. + + Returns + ------- + field : ndarray + Complex spatial field shaped ``(..., sy, sx)``, where ``sy`` and + ``sx`` are the spatial dimensions of the original hologram. + """ + fft_in = self.fft_used if propagated_fft is None else propagated_fft + field = finalize_fourier_field( + fft_in=fft_in, + ifft_fn=self.ifft_fn, + input_shape=self.input_shape, + fft_shape=self.fft_shape, + padding=self.padding, + scale_to_filter=self.scale_to_filter, + crop_radius=self.crop_radius, + ) + self._field = field + return field + + @property + def field(self) -> xp.ndarray: + """Cached spatial field. + + Returns the result of the last :meth:`finalize` call. If + :meth:`finalize` has not been called yet, it is invoked with no + arguments (i.e. using :attr:`fft_used`). + """ + if self._field is None: + return self.finalize() + return self._field diff --git a/qpretrieve/interfere/if_oah.py b/qpretrieve/interfere/if_oah.py index 367670b..237886d 100644 --- a/qpretrieve/interfere/if_oah.py +++ b/qpretrieve/interfere/if_oah.py @@ -1,4 +1,7 @@ +from __future__ import annotations + from .._ndarray_backend import xp +from ..fourier import FourierFieldData from .base import BaseInterferogram @@ -12,13 +15,31 @@ class OffAxisHologram(BaseInterferogram): "scale_to_filter": False, "sideband_freq": None, "invert_phase": False, + "output_domain": "spatial", } + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._fourier_field_data = None + + @property + def field(self) -> xp.ndarray: + """Retrieved complex field information.""" + if self._field is None: + if self._fourier_field_data is not None: + self.compute_field() + else: + self.run_pipeline(output_domain="spatial") + return self._field + @property def phase(self) -> xp.ndarray: """Retrieved phase information""" if self._field is None: - self.run_pipeline() + if self._fourier_field_data is not None: + self.compute_field() + else: + self.run_pipeline(output_domain="spatial") if self._phase is None: self._phase = xp.angle(self._field) return self._phase @@ -27,12 +48,39 @@ def phase(self) -> xp.ndarray: def amplitude(self) -> xp.ndarray: """Retrieved amplitude information""" if self._field is None: - self.run_pipeline() + if self._fourier_field_data is not None: + self.compute_field() + else: + self.run_pipeline(output_domain="spatial") if self._amplitude is None: self._amplitude = xp.abs(self._field) return self._amplitude - def run_pipeline(self, **pipeline_kws) -> xp.ndarray: + def compute_field(self, + propagated_fft: xp.ndarray | None = None) -> xp.ndarray: + """Compute the field using the current pipeline settings. + + If the field was skipped previously with ``output_domain='fourier'``, + this will reuse the cached Fourier intermediates instead of + recomputing the full filter path. + + Parameters + ---------- + propagated_fft: ndarray, optional + Fourier-domain field after external propagation. If omitted, + the stored qpretrieve Fourier artifact is computed directly. + + """ + if self._field is None: + if self._fourier_field_data is None: + self.run_pipeline(output_domain="spatial") + else: + self._field = self._fourier_field_data.finalize( + propagated_fft=propagated_fft) + return self._field + + def run_pipeline(self, output_domain: str = "spatial", + **pipeline_kws) -> xp.ndarray | FourierFieldData: r"""Run OAH analysis pipeline Parameters @@ -78,7 +126,13 @@ def run_pipeline(self, **pipeline_kws) -> xp.ndarray: Illumination wavelength in meters for physical-radius mode. invert_phase: bool Invert the phase data. + output_domain: str + Either ``"spatial"`` or ``"fourier"``. Spatial returns the + field, Fourier returns a :class:`FourierFieldData`. + + .. versionadded:: 0.7.0 """ + pipeline_kws["output_domain"] = output_domain for key in self.default_pipeline_kws: if key not in pipeline_kws: pipeline_kws[key] = self.get_pipeline_kw(key) @@ -101,21 +155,39 @@ def run_pipeline(self, **pipeline_kws) -> xp.ndarray: filter_size = float(fsize) freq_pos = tuple(float(x) for x in pipeline_kws["sideband_freq"]) - field = self.fft.filter( - filter_name=pipeline_kws["filter_name"], - filter_size=filter_size, - freq_pos=freq_pos, - scale_to_filter=pipeline_kws["scale_to_filter"]) - - if pipeline_kws["invert_phase"]: - field.imag *= -1 + if pipeline_kws["output_domain"] == "spatial": + # spatial pipeline + field = self.fft.filter( + filter_name=pipeline_kws["filter_name"], + filter_size=filter_size, + freq_pos=freq_pos, + scale_to_filter=pipeline_kws["scale_to_filter"], + output_domain="spatial") + + if pipeline_kws["invert_phase"]: + field.imag *= -1 + + self._field = field + self._fourier_field_data = None + else: + # direct fourier pipeline + artifact = self.fft.filter( + filter_name=pipeline_kws["filter_name"], + filter_size=filter_size, + freq_pos=freq_pos, + scale_to_filter=pipeline_kws["scale_to_filter"], + output_domain="fourier") + self._field = None + self._fourier_field_data = artifact - self._field = field self._phase = None self._amplitude = None self.pipeline_kws.update(pipeline_kws) - return self.field + if pipeline_kws["output_domain"] == "spatial": + return self._field + else: + return self._fourier_field_data def find_peak_cosine( diff --git a/tests/conftest.py b/tests/conftest.py index ac30996..3b6832c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -43,9 +43,9 @@ def set_ndarray_backend_to_cupy(): qpretrieve.set_ndarray_backend('numpy') -@pytest.fixture(params=[64]) # default param for size +@pytest.fixture() def hologram(request): - size = request.param + size = getattr(request, "param", 64) x = np.arange(size).reshape(-1, 1) - size / 2 y = np.arange(size).reshape(1, -1) - size / 2 diff --git a/tests/requirements.txt b/tests/requirements.txt index e29c7b2..b07624e 100644 --- a/tests/requirements.txt +++ b/tests/requirements.txt @@ -1,2 +1,4 @@ h5py pytest +pytest-benchmark +nrefocus>=0.8.0 diff --git a/tests/test_bm/test_bm_full_pipeline.py b/tests/test_bm/test_bm_full_pipeline.py new file mode 100644 index 0000000..b125f1d --- /dev/null +++ b/tests/test_bm/test_bm_full_pipeline.py @@ -0,0 +1,49 @@ +import pytest +import qpretrieve + +pytest.importorskip("nrefocus") + +import nrefocus # noqa: E402 + +_PROPAGATION_KWARGS = dict(d=1.5, nm=1.533, res=8.25, + method="fresnel", padding=False) + + +def _run_spatial(hologram): + """Original legacy pipeline with all Fourier transforms""" + holo = qpretrieve.OffAxisHologram(hologram, padding=False) + field = holo.run_pipeline(output_domain="spatial") + return nrefocus.refocus(field=field, **_PROPAGATION_KWARGS) + + +def _run_fourier(hologram): + """Let nrefocus do the qpretrieve iFFT, nrefocus knows what to do.""" + holo = qpretrieve.OffAxisHologram(hologram, padding=False) + artifact = holo.run_pipeline(output_domain="fourier") + return nrefocus.refocus(field=artifact, **_PROPAGATION_KWARGS) + + +def _run_fourier_no_finalize(hologram): + """Output the data in Fourier domain""" + holo = qpretrieve.OffAxisHologram(hologram, padding=False) + artifact = holo.run_pipeline(output_domain="fourier") + return nrefocus.refocus(field=artifact, output_domain="fourier", + **_PROPAGATION_KWARGS) + + +def test_bm_full_pipeline_spatial(benchmark, hologram): + """Original legacy pipeline with all Fourier transforms""" + result = benchmark(_run_spatial, hologram) + assert result is not None + + +def test_bm_full_pipeline_fourier(benchmark, hologram): + """Don't do the qpretrieve iFFT, nrefocus knows what to do with FFT""" + result = benchmark(_run_fourier, hologram) + assert result is not None + + +def test_bm_full_pipeline_fourier_no_finalize(benchmark, hologram): + """Output the data in Fourier domain""" + result = benchmark(_run_fourier_no_finalize, hologram) + assert result is not None diff --git a/tests/test_bm/test_bm_oah_output_domain.py b/tests/test_bm/test_bm_oah_output_domain.py new file mode 100644 index 0000000..0d26c9c --- /dev/null +++ b/tests/test_bm/test_bm_oah_output_domain.py @@ -0,0 +1,34 @@ +import qpretrieve + + +def _make_holo(hologram): + return qpretrieve.OffAxisHologram(hologram) + + +def test_bm_oah_output_domain_spatial(benchmark, hologram): + """Original legacy pipeline with iFFT""" + field = benchmark( + lambda: _make_holo(hologram).run_pipeline(output_domain="spatial")) + + assert field.shape == (1, hologram.shape[0], hologram.shape[1]) + + +def test_bm_oah_output_domain_fourier(benchmark, hologram): + """Pipeline without iFFT, output in Fourier domain""" + result = benchmark( + lambda: _make_holo(hologram).run_pipeline(output_domain="fourier")) + + assert result is not None + + +def test_bm_oah_output_domain_fourier_then_spatial(benchmark, hologram): + """Pipeline initially without iFFT, then iFFT (for use with nrefocus)""" + + def run(): + holo = _make_holo(hologram) + artifact = holo.run_pipeline(output_domain="fourier") + return artifact.finalize() # or holo.compute_field() + + field = benchmark(run) + + assert field.shape == (1, hologram.shape[0], hologram.shape[1]) diff --git a/tests/test_oah.py b/tests/test_oah.py index 27953ad..f8be148 100644 --- a/tests/test_oah.py +++ b/tests/test_oah.py @@ -164,9 +164,9 @@ def test_get_field_filter_names(hologram): @pytest.mark.parametrize("hologram", [62, 63, 64], indirect=True) def test_get_field_interpretation_fourier_index(hologram): - """Filter size in Fourier space using Fourier index new in 0.7.0""" + """Filter size in Fourier space using Fourier index""" data = hologram - shape_expected = (1, hologram.shape[-2], hologram.shape[-1]) + shape_expected = (1, data.shape[-2], data.shape[-1]) holo = qpretrieve.OffAxisHologram(data) ft_data = holo.fft_origin @@ -192,7 +192,7 @@ def test_get_field_interpretation_fourier_index(hologram): @pytest.mark.parametrize("hologram", [62, 63, 64], indirect=["hologram"]) def test_get_field_interpretation_fourier_index_control(hologram): - """Filter size in Fourier space using Fourier index new in 0.7.0""" + """Filter size in Fourier space using Fourier index""" data = hologram holo = qpretrieve.OffAxisHologram(data) @@ -241,7 +241,7 @@ def test_get_field_interpretation_fourier_index_mask_1(hologram, filter_size): @pytest.mark.parametrize("hologram", [62, 63, 64, 134, 135], indirect=["hologram"]) def test_get_field_interpretation_fourier_index_mask_2(hologram): - """Filter size in Fourier space using Fourier index new in 0.7.0""" + """Filter size in Fourier space using Fourier index""" data = hologram holo = qpretrieve.OffAxisHologram(data) diff --git a/tests/test_output_domain_option.py b/tests/test_output_domain_option.py new file mode 100644 index 0000000..a9b6c37 --- /dev/null +++ b/tests/test_output_domain_option.py @@ -0,0 +1,73 @@ +import numpy as np + +import qpretrieve +from qpretrieve import fourier +from qpretrieve.fourier import FourierFieldData + + +def test_fftfilter_filter_can_return_fourier_domain(): + x = np.linspace(-100, 100, 100) + xx, yy = np.meshgrid(x, -x, indexing="ij") + gauss = np.exp(-(xx ** 2 + yy ** 2) / 625) + + ft = fourier.FFTFilterNumpy(gauss, subtract_mean=False) + field = ft.filter( + filter_name="disk", + filter_size=0.25, + freq_pos=(0, 0), + scale_to_filter=False, + output_domain="fourier", + ) + + assert isinstance(field, FourierFieldData) + assert field.output_domain == "spatial" + assert ft.fft_used is not None + + +def test_oah_run_pipeline_can_return_fourier_domain(hologram): + holo = qpretrieve.OffAxisHologram(hologram) + + field = holo.run_pipeline(output_domain="fourier") + + assert isinstance(field, FourierFieldData) + assert field.output_domain == "spatial" + assert holo._field is None + assert holo._fourier_field_data is field + assert holo._phase is None + assert holo._amplitude is None + + +def test_oah_compute_field_reuses_cached_intermediates(hologram): + holo = qpretrieve.OffAxisHologram(hologram) + + artifact = holo.run_pipeline(output_domain="fourier") + field = holo.compute_field() + + assert field is holo._field + assert field.shape == (1, hologram.shape[0], hologram.shape[1]) + assert holo._phase is None + assert holo._amplitude is None + assert artifact.field is field + + +def test_oah_compute_field_matches_output_domain_spatial(hologram): + holo_spatial = qpretrieve.OffAxisHologram(hologram) + field_spatial = holo_spatial.run_pipeline(output_domain="spatial") + + holo_fourier = qpretrieve.OffAxisHologram(hologram) + artifact_fourier = holo_fourier.run_pipeline(output_domain="fourier") + field_fourier = artifact_fourier.finalize() + + assert np.allclose(field_spatial, field_fourier) + + +def test_oah_compute_field_accepts_propagated_fft(hologram): + holo = qpretrieve.OffAxisHologram(hologram) + artifact = holo.run_pipeline(output_domain="fourier") + + # here the propagated_fft could be a fft from nrefocus + field = holo.compute_field(propagated_fft=artifact.fft_used) + + assert field is holo._field + assert holo.phase.shape == field.shape + assert holo.amplitude.shape == field.shape