xpunwrap 0.0.0__py3-none-any.whl

xpunwrap/__init__.py

@@ -0,0 +1,29 @@
+from collections.abc import Callable
+from typing import Any
+
+from ._ndarray_backend import get_ndarray_backend, set_ndarray_backend
+
+from .algorithms import (
+    algo_ls_poisson,
+    algo_ls_poisson_pg,
+    algo_ls_weighted,
+    algo_skimage_unwrap,
+)
+
+
+def algos_available() -> dict[str, Callable[..., Any]]:
+    """Return all available phase unwrapping algorithms.
+
+    Returns
+    -------
+        Mapping from algorithm name to callable, in the order:
+        ``algo_ls_poisson``, ``algo_ls_poisson_pg``, ``algo_ls_weighted``,
+        ``algo_skimage_unwrap``.
+    """
+    algos = {}
+    for algo in [algo_ls_poisson,
+                 algo_ls_poisson_pg,
+                 algo_ls_weighted,
+                 algo_skimage_unwrap]:
+        algos[algo.__name__] = algo
+    return algos

xpunwrap/_dtype_utils.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import types
+from typing import Any
+
+import numpy as np
+
+
+def real_pi(xp: types.ModuleType, dtype: Any) -> Any:
+    """Return pi cast to the requested real dtype.
+
+    Parameters
+    ----------
+    xp : module
+        Active ndarray backend (``numpy`` or ``cupy``).
+    dtype : dtype-like
+        Target real floating-point dtype.
+
+    Returns
+    -------
+    scalar
+        ``pi`` as a scalar of ``dtype``.
+    """
+    dt = xp.dtype(dtype)
+    return dt.type(xp.pi)
+
+
+def complex_dtype_for_real(xp: types.ModuleType, dtype: Any) -> np.dtype:
+    """Map a real dtype to the corresponding complex dtype.
+
+    ``float16`` and ``float32`` map to ``complex64``; ``float64`` maps to
+    ``complex128``. Complex inputs are returned unchanged. Uncommon real
+    dtypes fall back to ``complex128``.
+
+    Parameters
+    ----------
+    xp : module
+        Active ndarray backend (``numpy`` or ``cupy``).
+    dtype : dtype-like
+        Input dtype, real or complex.
+
+    Returns
+    -------
+    dtype
+        Corresponding complex dtype.
+    """
+    dt = xp.dtype(dtype)
+    if dt.kind == "c":
+        return dt
+    if dt == xp.float16 or dt == xp.float32:
+        return xp.complex64
+    if dt == xp.float64:
+        return xp.complex128
+    return xp.complex128

xpunwrap/_ndarray_backend.py

@@ -0,0 +1,96 @@
+"""
+Module that controls and exposes the active ndarray backend.
+NumPy is used for CPU-only. CuPy is used for GPU.
+"""
+
+import importlib
+import types
+from typing import Any
+
+_default_backend = "numpy"
+_xp = importlib.import_module(_default_backend)
+
+
+class NDArrayBackend:
+    """Proxy object exposing the current ndarray backend.
+
+    All attribute access is forwarded to the active backend module, so
+    ``xp.array``, ``xp.zeros``, etc. resolve to the correct implementation
+    without any conditional imports in calling code.
+    """
+
+    def __init__(self) -> None:
+        self._xp = _xp
+
+    def get(self) -> types.ModuleType:
+        """Return the currently active backend module.
+
+        Returns
+        -------
+        module
+            The active ndarray backend (``numpy`` or ``cupy``).
+        """
+        return self._xp
+
+    def set(self, backend_name: str = "numpy") -> None:
+        """Switch the active ndarray backend.
+
+        Parameters
+        ----------
+        backend_name : str
+            Name of the backend to activate. ``"numpy"`` selects CPU,
+            ``"cupy"`` selects GPU. Default ``"numpy"``.
+
+        Raises
+        ------
+        ImportError
+            If the requested backend is not installed.
+        """
+        global _xp
+        try:
+            self._xp = importlib.import_module(backend_name)
+            _xp = self._xp  # keep global in sync
+        except ModuleNotFoundError as err:
+            raise ImportError(f"The backend '{backend_name}' is not "
+                              f"installed. Either install it or use the "
+                              f"default backend: 'numpy'.") from err
+
+    def __getattr__(self, name: str) -> Any:
+        """Delegate attribute lookup to the active backend module.
+
+        Parameters
+        ----------
+        name : str
+            Attribute name.
+
+        Returns
+        -------
+        object
+            The corresponding attribute from the active backend.
+        """
+        return getattr(self._xp, name)
+
+    def is_numpy(self) -> bool:
+        """Return ``True`` if the active backend is NumPy (CPU).
+
+        Returns
+        -------
+        bool
+        """
+        return self._xp.__name__.startswith("numpy")
+
+    def is_cupy(self) -> bool:
+        """Return ``True`` if the active backend is CuPy (GPU).
+
+        Returns
+        -------
+        bool
+        """
+        return self._xp.__name__.startswith("cupy")
+
+
+# Export a single global proxy instance
+xp = NDArrayBackend()
+# This is what is imported by the user
+get_ndarray_backend = xp.get
+set_ndarray_backend = xp.set

xpunwrap/algorithms/__init__.py

@@ -0,0 +1,4 @@
+from .ls_poisson import algo_ls_poisson
+from .ls_poisson_pg import algo_ls_poisson_pg
+from .ls_weighted import algo_ls_weighted
+from .skimage_unwrap import algo_skimage_unwrap

xpunwrap/algorithms/_ls_common.py

@@ -0,0 +1,92 @@
+"""
+Shared building blocks for the least-squares Poisson unwrap algorithms.
+
+These helpers are used by both ``ls_poisson`` and ``ls_poisson_pg``.
+"""
+
+from .._dtype_utils import complex_dtype_for_real, real_pi
+from .._ndarray_backend import xp
+from ..fourier import get_fft_engine
+
+
+def wrap_phase(x: xp.ndarray) -> xp.ndarray:
+    """
+    Wrap phase to [-pi, pi).
+
+    Parameters
+    ----------
+    x : xp.ndarray
+        Input phase values.
+
+    Returns
+    -------
+    xp.ndarray
+        Wrapped values in [-pi, pi).
+    """
+    dtype = x.dtype
+    pi = real_pi(xp, dtype)
+    two_pi = dtype.type(2) * pi
+    return (x + pi) % two_pi - pi
+
+
+def divergence_stack(gx: xp.ndarray, gy: xp.ndarray) -> xp.ndarray:
+    """
+    Compute the periodic divergence of a wrapped-gradient stack.
+
+    Parameters
+    ----------
+    gx, gy : xp.ndarray
+        Gradient stacks, shape (N, H, W).
+
+    Returns
+    -------
+    xp.ndarray
+        Divergence, shape (N, H, W).
+    """
+    # Periodic backward-difference divergence used by the FFT Poisson solve.
+    # div g = (g_x - g_x shifted right) + (g_y - g_y shifted down)
+    return (
+        gx - xp.roll(gx, 1, axis=2)
+        + gy - xp.roll(gy, 1, axis=1)
+    )
+
+
+def poisson_solve_fft_stack(rhs: xp.ndarray) -> xp.ndarray:
+    """
+    Solve the periodic Poisson equation in the Fourier domain.
+
+    Returns the positive-denominator solution; callers multiply by -1 to
+    match the discrete Laplacian sign convention.
+
+    Parameters
+    ----------
+    rhs : xp.ndarray
+        Right-hand side, shape (N, H, W).
+
+    Returns
+    -------
+    xp.ndarray
+        Solution, shape (N, H, W).
+    """
+    _, H, W = rhs.shape
+    dtype = rhs.dtype
+    pi = real_pi(xp, dtype)
+    two = dtype.type(2)
+
+    ky = xp.fft.fftfreq(H).astype(dtype, copy=False).reshape(1, H, 1)
+    kx = xp.fft.fftfreq(W).astype(dtype, copy=False).reshape(1, 1, W)
+
+    denom = (two - two * xp.cos(two * pi * kx)) + \
+            (two - two * xp.cos(two * pi * ky))
+
+    denom[:, 0, 0] = dtype.type(1)
+
+    complex_dtype = complex_dtype_for_real(xp, dtype)
+    complex_dt = xp.dtype(complex_dtype)
+    fft = get_fft_engine()
+    rhs_hat = fft.fft2(rhs.astype(complex_dtype, copy=False))
+    phi_hat = rhs_hat / denom
+    phi_hat[:, 0, 0] = complex_dt.type(0)
+
+    out = xp.real(fft.ifft2(phi_hat))
+    return out.astype(dtype, copy=False)

xpunwrap/algorithms/_plane_utils.py

@@ -0,0 +1,65 @@
+from .._ndarray_backend import xp
+from ._ls_common import wrap_phase
+
+
+def restore_mean_plane(
+    phase_unwrapped: xp.ndarray,
+    phase_wrapped: xp.ndarray,
+) -> xp.ndarray:
+    """Reintroduce the mean wrapped gradient plane removed by Poisson solvers.
+
+    Least-squares Poisson solvers discard the null-space component (a linear
+    ramp). This function estimates that ramp from the mean wrapped gradients of
+    the input and adds it back, anchored so that the value at pixel ``(0, 0)``
+    matches the original wrapped phase.
+
+    Parameters
+    ----------
+    phase_unwrapped : xp.ndarray
+        Solver output, shape (H, W) or (N, H, W).
+    phase_wrapped : xp.ndarray
+        Original wrapped phase, same shape as ``phase_unwrapped``.
+
+    Returns
+    -------
+    xp.ndarray
+        Phase with the mean gradient plane restored, same shape as input.
+
+    Raises
+    ------
+    ValueError
+        If ``phase_unwrapped`` is not 2-D or 3-D.
+    """
+    input_2d = False
+    if phase_unwrapped.ndim == 2:
+        input_2d = True
+        phase_unwrapped = xp.expand_dims(phase_unwrapped, axis=0)
+        phase_wrapped = xp.expand_dims(phase_wrapped, axis=0)
+    elif phase_unwrapped.ndim != 3:
+        raise ValueError(
+            "phase_unwrapped must have shape (H, W) or (N, H, W)."
+        )
+
+    dtype = phase_unwrapped.dtype
+
+    gx = wrap_phase(xp.diff(phase_wrapped, axis=2,
+                            append=phase_wrapped[:, :, -1:]))
+    gy = wrap_phase(xp.diff(phase_wrapped, axis=1,
+                            append=phase_wrapped[:, -1:, :]))
+
+    gx_mean = gx.mean(axis=(1, 2), keepdims=True)
+    gy_mean = gy.mean(axis=(1, 2), keepdims=True)
+
+    _, H, W = phase_unwrapped.shape
+    x_idx = xp.arange(W, dtype=dtype).reshape(1, 1, W)
+    y_idx = xp.arange(H, dtype=dtype).reshape(1, H, 1)
+    plane = gx_mean * x_idx + gy_mean * y_idx
+    anchor = (
+        phase_wrapped[:, 0, 0]
+        - (phase_unwrapped[:, 0, 0] + plane[:, 0, 0])
+    )
+    out = phase_unwrapped + plane + anchor[:, None, None]
+
+    if input_2d:
+        out = out[0]
+    return out

xpunwrap/algorithms/ls_poisson.py

@@ -0,0 +1,68 @@
+from .._ndarray_backend import xp
+from ._ls_common import divergence_stack, poisson_solve_fft_stack, wrap_phase
+from ._plane_utils import restore_mean_plane
+
+
+def algo_ls_poisson(
+    phase_wrapped: xp.ndarray,
+    restore_plane: bool = False,
+) -> xp.ndarray:
+    """
+    Batched 2D phase unwrapping using a least-squares Poisson solver.
+
+    Parameters
+    ----------
+    phase_wrapped : xp.ndarray
+        Wrapped phase, shape (H, W) or (N, H, W), values in [-pi, pi).
+    restore_plane : bool, optional
+        If True, add back the mean wrapped gradient plane. Default False.
+
+    Returns
+    -------
+    phase_unwrapped : xp.ndarray
+        Unwrapped phase, same shape as input.
+
+    Notes
+    -----
+    The FFT Poisson solver assumes periodic boundary conditions: the input
+    domain is treated as if its left/right and top/bottom edges connect. No
+    zero-padding is applied before the FFT. If the wrapped phase is not
+    periodic at the boundaries, the solver may produce artifacts (ringing or
+    slope errors) near the domain edges. For non-periodic data consider
+    cropping the region of interest away from the edges or using
+    :func:`algo_ls_weighted`, which can suppress discontinuous boundary regions
+    via its border-weight mask.
+
+    References
+    ----------
+    .. [1] D. C. Ghiglia and M. D. Pritt, "Two-Dimensional Phase Unwrapping:
+       Theory, Algorithms, and Software," Wiley, 1998.
+
+    """
+    input_2d = False
+    if phase_wrapped.ndim == 2:
+        input_2d = True
+        phase_wrapped = xp.expand_dims(phase_wrapped, axis=0)
+    elif phase_wrapped.ndim != 3:
+        raise ValueError("phase_wrapped must have shape (H, W) or (N, H, W).")
+
+    dtype = phase_wrapped.dtype
+    # Wrapped gradients with periodic forward differences.
+    # This is the discrete operator used by the FFT Poisson solver.
+    gx = wrap_phase(xp.roll(phase_wrapped, -1, axis=2) - phase_wrapped)
+    gy = wrap_phase(xp.roll(phase_wrapped, -1, axis=1) - phase_wrapped)
+
+    rhs = divergence_stack(gx, gy)
+
+    # Solve the discrete Poisson equation in the Fourier domain.
+    # The FFT helper returns the positive-denominator form, so apply -1 to
+    # match the Laplacian sign convention here.
+    phase_unwrapped = poisson_solve_fft_stack(rhs)
+    phase_unwrapped *= dtype.type(-1)
+
+    if restore_plane:
+        phase_unwrapped = restore_mean_plane(phase_unwrapped, phase_wrapped)
+
+    if input_2d:
+        phase_unwrapped = phase_unwrapped[0]
+    return phase_unwrapped

xpunwrap/algorithms/ls_poisson_pg.py

@@ -0,0 +1,117 @@
+from .._ndarray_backend import xp
+from ._ls_common import divergence_stack, poisson_solve_fft_stack, wrap_phase
+from ._plane_utils import restore_mean_plane
+
+
+def algo_ls_poisson_pg(
+    phase_wrapped: xp.ndarray,
+    restore_plane: bool = False,
+) -> xp.ndarray:
+    """
+    Least-squares unwrapping with periodic gradient enforcement.
+
+    Parameters
+    ----------
+    phase_wrapped : xp.ndarray
+        Wrapped phase, shape (H, W) or (N, H, W), values in [-pi, pi).
+    restore_plane : bool, optional
+        If True, add back the mean wrapped gradient plane. Default False.
+
+    Returns
+    -------
+    phase_unwrapped : xp.ndarray
+        Unwrapped phase, same shape as input.
+
+    Notes
+    -----
+    The FFT Poisson solver assumes periodic boundary conditions: the input
+    domain is treated as if its left/right and top/bottom edges connect. No
+    zero-padding is applied before the FFT. This algorithm explicitly enforces
+    periodicity on the wrapped gradients (via
+    :func:`enforce_periodic_gradients_stack`) before solving, which reduces
+    gradient discontinuities at domain edges compared to
+    :func:`algo_ls_poisson`. If the wrapped phase is not periodic at the
+    boundaries, artifacts near the edges may still occur.
+
+    References
+    ----------
+    .. [1] D. C. Ghiglia and M. D. Pritt, "Two-Dimensional Phase Unwrapping:
+       Theory, Algorithms, and Software," Wiley, 1998.
+
+    """
+    input_2d = False
+    if phase_wrapped.ndim == 2:
+        input_2d = True
+        phase_wrapped = xp.expand_dims(phase_wrapped, axis=0)
+    elif phase_wrapped.ndim != 3:
+        raise ValueError("phase_wrapped must have shape (H, W) or (N, H, W).")
+
+    # Periodic gradients.
+    gx, gy = wrapped_gradients_stack(phase_wrapped)
+    gx, gy = enforce_periodic_gradients_stack(gx, gy)
+    rhs = divergence_stack(gx, gy)
+
+    # The FFT helper returns the positive-denominator form, so apply -1 to
+    # match the Laplacian sign convention here.
+    phi = poisson_solve_fft_stack(rhs)
+    phi *= phi.dtype.type(-1)
+
+    if restore_plane:
+        phi = restore_mean_plane(phi, phase_wrapped)
+    if input_2d:
+        phi = phi[0]
+    return phi
+
+
+def wrapped_gradients_stack(
+        phi: xp.ndarray,
+) -> tuple[xp.ndarray, xp.ndarray]:
+    """
+    Wrapped forward gradients for a stack of phases.
+
+    Parameters
+    ----------
+    phi : xp.ndarray
+        Wrapped phase, shape (N, H, W).
+
+    Returns
+    -------
+    gx, gy : xp.ndarray
+        Wrapped gradients with shape (N, H, W).
+    """
+    gx = wrap_phase(phi[:, :, 1:] - phi[:, :, :-1])
+    gy = wrap_phase(phi[:, 1:, :] - phi[:, :-1, :])
+
+    # Pad with zeros to restore the full (N, H, W) shape lost by diff.
+    # The padded edge values are immediately overwritten with the correct
+    # periodic boundary values in enforce_periodic_gradients_stack.
+    gx = xp.pad(gx, ((0, 0), (0, 0), (0, 1)))
+    gy = xp.pad(gy, ((0, 0), (0, 1), (0, 0)))
+
+    return gx, gy
+
+
+def enforce_periodic_gradients_stack(
+        gx: xp.ndarray,
+        gy: xp.ndarray,
+) -> tuple[xp.ndarray, xp.ndarray]:
+    """
+    Enforce periodic boundary conditions on wrapped gradients.
+
+    Parameters
+    ----------
+    gx, gy : xp.ndarray
+        Gradient stacks, shape (N, H, W).
+
+    Returns
+    -------
+    gx, gy : xp.ndarray
+        Periodic gradients, shape (N, H, W).
+    """
+    # x-direction periodicity
+    gx[:, :, -1] = gx[:, :, 0]
+
+    # y-direction periodicity
+    gy[:, -1, :] = gy[:, 0, :]
+
+    return gx, gy