Examples

Field Retrieval & Unwrapping

Field retrieval (qpretrieve) and phase unwrapping (xpunwrap) on GPU.

gpu_field_retr_phase_unwrapping.py

import matplotlib.pyplot as plt
import numpy as np
import qpretrieve
import xpunwrap

# Force GPU backend for both libraries.
qpretrieve.set_ndarray_backend("cupy")
xpunwrap.set_ndarray_backend("cupy")
xp = xpunwrap.get_ndarray_backend()

fft_interface = qpretrieve.fourier.FFTFilterCupy
edata = np.load("./data/hologram_cell.npz")
holo = qpretrieve.OffAxisHologram(edata["data"], fft_interface)
bg = qpretrieve.OffAxisHologram(edata["bg_data"], fft_interface)

holo.run_pipeline(filter_name="disk", filter_size=1 / 2, scale_to_filter=True)
bg.process_like(holo)
phase_wrapped = xp.asarray(holo.phase - bg.phase).astype(xp.float32)

# Unwrap the phase with all available algorithms
outputs = {}
for algo_name, algo in xpunwrap.algos_available().items():
    outputs[algo_name] = algo(phase_wrapped)
skimage_out = outputs.get("algo_skimage_unwrap", None)
if skimage_out is not None:
    outputs_no_skimage = {k: v for k, v in outputs.items() if k != "algo_skimage_unwrap"}
else:
    outputs_no_skimage = outputs

# plot the wrapped and unwrapped phases
plt.style.use("dark_background")
fig, axes = plt.subplots(2, 3, figsize=(8, 6))
fig.suptitle("Field Retrieval + Phase Unwrapping (GPU)", fontsize=18)
axes = axes.flatten(order="F")

axes[0].imshow(phase_wrapped.get()[0])
axes[0].set_title("Wrapped Phase")

# With column-major flattening, bottom-right is index 5.
skimage_ax = axes[5]

# Fill remaining non-skimage algorithms first (slots 1..4).
plot_slots = [2, 3, 4]
for slot, (algo_name, arr) in zip(plot_slots, outputs_no_skimage.items()):
    ax = axes[slot]
    ax.imshow(arr.get()[0])
    ax.set_title(f"Unwrapped\n{algo_name}")

if skimage_out is not None:
    skimage_ax.imshow(skimage_out.get()[0])
    skimage_ax.set_title("Unwrapped (CPU-only)\nalgo_skimage_unwrap")
else:
    skimage_ax.text(0.5, 0.5, "skimage missing", ha="center", va="center")
    skimage_ax.set_title("Unwrapped\nalgo_skimage_unwrap")

for ax in axes:
    ax.set_axis_off()

plt.tight_layout(w_pad=4.5)
plt.savefig("gpu_field_retr_phase_unwrapping.png")
# plt.show()

CPU vs GPU phase unwrapping speed comparison

Compares four pipelines on hologram_cell data:

• CPU skimage + NumPy (qpretrieve NumPy FFT + skimage unwrap)

• CPU skimage + pyFFTW (qpretrieve pyFFTW FFT + skimage unwrap)

• CPU ls_poisson + pyFFTW (qpretrieve pyFFTW FFT + xpunwrap Poisson)

• GPU ls_poisson + CuPy (full CuPy pipeline)

Run:

python examples/compare_cpu_gpu.py

compare_cpu_gpu.py

from __future__ import annotations

import time
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

import qpretrieve
import xpunwrap
import xpunwrap.fourier as xp_fourier

DATA_PATH = Path(__file__).parent.parent / "tests" / "data" / "hologram_cell.npz"
OUTPUT_PNG = Path(__file__).parent / "compare_cpu_gpu.png"
STACK_DEPTH = 256
TILE_SIZE = 256


def _load_raw_numpy() -> dict[str, np.ndarray]:
    edata = np.load(DATA_PATH)
    data = np.asarray(edata["data"], dtype=np.float32)
    bg = np.asarray(edata["bg_data"], dtype=np.float32)
    tile_h = int(np.ceil(TILE_SIZE / data.shape[-2]))
    tile_w = int(np.ceil(TILE_SIZE / data.shape[-1]))
    data = np.tile(data, (tile_h, tile_w))[:TILE_SIZE, :TILE_SIZE]
    bg = np.tile(bg, (tile_h, tile_w))[:TILE_SIZE, :TILE_SIZE]
    return {"data": data, "bg": bg}


def _stack_phase(arr: np.ndarray) -> np.ndarray:
    return np.repeat(arr[None, ...], repeats=STACK_DEPTH, axis=0)


def _time_skimage_pipeline(
        raw: dict[str, np.ndarray],
        fft_interface,
        repeats: int = 3,
):
    """Shared timing loop for skimage-based pipelines."""
    qpretrieve.set_ndarray_backend("numpy")
    try:
        algo = xpunwrap.algos_available()["algo_skimage_unwrap"]
    except Exception:
        return None

    def _run():
        holo = qpretrieve.OffAxisHologram(raw["data"], fft_interface)
        bg = qpretrieve.OffAxisHologram(raw["bg"], fft_interface)
        holo.run_pipeline(
            filter_name="disk", filter_size=1 / 2, scale_to_filter=True)
        bg.process_like(holo)
        algo(np.asarray(holo.phase - bg.phase, dtype=np.float32))

    _run()  # warmup (important for pyFFTW plan compilation)
    times = []
    for _ in range(repeats):
        start = time.perf_counter()
        _run()
        times.append((time.perf_counter() - start) * 1000.0)
    return float(np.mean(times)), float(np.std(times))


def time_cpu_skimage_numpy(raw: dict[str, np.ndarray], repeats: int = 3):
    """CPU pipeline: qpretrieve (NumPy FFT) + skimage unwrap."""
    return _time_skimage_pipeline(
        raw, qpretrieve.fourier.FFTFilterNumpy, repeats)


def time_cpu_skimage_pyfftw(raw: dict[str, np.ndarray], repeats: int = 3):
    """CPU pipeline: qpretrieve (pyFFTW) + skimage unwrap."""
    return _time_skimage_pipeline(
        raw, qpretrieve.fourier.FFTFilterPyFFTW, repeats)


def time_cpu_ls_poisson_pyfftw(raw: dict[str, np.ndarray], repeats: int = 3):
    """CPU pipeline: qpretrieve (pyFFTW) + xpunwrap ls_poisson (pyFFTW)."""
    from xpunwrap.fourier import FFTEnginePyFFTW
    if FFTEnginePyFFTW is None:
        return None

    qpretrieve.set_ndarray_backend("numpy")
    xpunwrap.set_ndarray_backend("numpy")
    fft_interface = qpretrieve.fourier.FFTFilterPyFFTW
    xp_fourier.PREFERRED_ENGINE = "FFTEnginePyFFTW"
    algo = xpunwrap.algos_available()["algo_ls_poisson"]

    try:
        # Warmup
        holo = qpretrieve.OffAxisHologram(raw["data"], fft_interface)
        bg = qpretrieve.OffAxisHologram(raw["bg"], fft_interface)
        holo.run_pipeline(
            filter_name="disk", filter_size=1 / 2, scale_to_filter=True)
        bg.process_like(holo)
        algo(np.asarray(holo.phase - bg.phase, dtype=np.float32))

        times = []
        for _ in range(repeats):
            start = time.perf_counter()
            holo = qpretrieve.OffAxisHologram(raw["data"], fft_interface)
            bg = qpretrieve.OffAxisHologram(raw["bg"], fft_interface)
            holo.run_pipeline(
                filter_name="disk", filter_size=1 / 2, scale_to_filter=True)
            bg.process_like(holo)
            algo(np.asarray(holo.phase - bg.phase, dtype=np.float32))
            times.append((time.perf_counter() - start) * 1000.0)
    finally:
        xp_fourier.PREFERRED_ENGINE = None

    return float(np.mean(times)), float(np.std(times))


def time_gpu_pipeline(raw: dict[str, np.ndarray], repeats: int = 3):
    """GPU pipeline: qpretrieve (CuPy FFT) + xpunwrap ls_poisson (CuPy FFT)."""
    try:
        import cupy as cp  # noqa: F401
    except Exception:
        return None

    qpretrieve.set_ndarray_backend("cupy")
    fft_interface = qpretrieve.fourier.FFTFilterCupy
    xpunwrap.set_ndarray_backend("cupy")
    xp = xpunwrap.get_ndarray_backend()
    algo = xpunwrap.algos_available()["algo_ls_poisson"]

    def _sync():
        if hasattr(xp, "cuda"):
            xp.cuda.Stream.null.synchronize()

    # Warmup
    _sync()
    holo_gpu = xp.asarray(raw["data"])
    bg_gpu = xp.asarray(raw["bg"])
    holo_obj = qpretrieve.OffAxisHologram(holo_gpu, fft_interface)
    bg_obj = qpretrieve.OffAxisHologram(bg_gpu, fft_interface)
    holo_obj.run_pipeline(
        filter_name="disk", filter_size=1 / 2, scale_to_filter=True)
    bg_obj.process_like(holo_obj)
    algo(xp.asarray(holo_obj.phase - bg_obj.phase, dtype=xp.float32))
    _sync()

    times = []
    for _ in range(repeats):
        _sync()
        start = time.perf_counter()
        holo_gpu = xp.asarray(raw["data"])
        bg_gpu = xp.asarray(raw["bg"])
        holo_obj = qpretrieve.OffAxisHologram(holo_gpu, fft_interface)
        bg_obj = qpretrieve.OffAxisHologram(bg_gpu, fft_interface)
        holo_obj.run_pipeline(
            filter_name="disk", filter_size=1 / 2, scale_to_filter=True)
        bg_obj.process_like(holo_obj)
        algo(xp.asarray(holo_obj.phase - bg_obj.phase, dtype=xp.float32))
        _sync()
        times.append((time.perf_counter() - start) * 1000.0)
    return float(np.mean(times)), float(np.std(times))


def main():
    raw = _load_raw_numpy()

    skimage_numpy = time_cpu_skimage_numpy(raw)
    skimage_pyfftw = time_cpu_skimage_pyfftw(raw)
    ls_pyfftw = time_cpu_ls_poisson_pyfftw(raw)
    ls_cupy = time_gpu_pipeline(raw)

    # Wong (2011) / Nature Methods colorblind-safe palette
    clr_blue, clr_green = "#56B4E9", "#009E73"
    clr_orange, clr_purple = "#D55E00", "#CC79A7"
    entries = [
        ("skimage\n+ NumPy\n(CPU)", skimage_numpy, clr_blue),
        ("skimage\n+ pyFFTW\n(CPU)", skimage_pyfftw, clr_green),
        ("ls_poisson\n+ pyFFTW\n(CPU)", ls_pyfftw, clr_orange),
        ("ls_poisson\n+ CuPy\n(GPU)", ls_cupy, clr_purple),
    ]
    entries = [(lbl, res, col) for lbl, res, col in entries if res is not None]

    labels = [e[0] for e in entries]
    means = [e[1][0] for e in entries]
    stds = [e[1][1] for e in entries]
    colors = [e[2] for e in entries]

    plt.style.use("dark_background")
    fig, ax = plt.subplots(figsize=(6, 4), facecolor="#181C24")
    x = np.arange(len(labels))
    ax.bar(x, means, yerr=stds, color=colors, edgecolor="white",
           error_kw={"ecolor": "white", "elinewidth": 1.5,
                     "capsize": 5, "capthick": 1.5})
    ax.set_xticks(x)
    ax.set_xticklabels(labels, ha="center")
    ax.set_ylabel("Time (ms)")
    ax.set_title("Phase unwrapping speed comparison")
    ax.grid(axis="y", linestyle="--", alpha=0.3)
    ax.set_facecolor("#181C24")

    for xi, mean, std in zip(x, means, stds):
        ax.text(xi, mean + std, f"{mean:.1f} ms",
                ha="center", va="bottom", fontsize=8)

    fig.tight_layout()
    fig.savefig(OUTPUT_PNG, dpi=200)
    print(f"Saved {OUTPUT_PNG}")


if __name__ == "__main__":
    main()

Poisson LS Algorithm Walkthrough

Step-by-step walkthrough of the Least-squares Poisson Periodic Gradient algorithm.

The aim here is to visualise the steps of the algorithm with matplotlib with example datasets, showing edge cases where phase unwrapping fails. Where possible the scikit-image 2D phase unwrapping algorithm will be used for comparison.

The example data is the data used by M. Gdeisat and F. Lilley in their excellent “Two-Dimensional Phase Unwrapping Problem” guide.

walkthrough-ls_poisson_pg.py

import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple

import xpunwrap
from xpunwrap.algorithms.ls_poisson_pg import (
    wrapped_gradients_stack,
    enforce_periodic_gradients_stack,
    divergence_stack,
    poisson_solve_fft_stack,
)
from xpunwrap.algorithms._plane_utils import restore_mean_plane

# use numpy for this example for simplicity
xpunwrap.set_ndarray_backend('numpy')
xp = xpunwrap.get_ndarray_backend()


# load/create the example data
def _generate_phase(nrx: int = 512, nry: int = 512) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Create the continuous phase image f(x, y) = 20*exp(-0.25*(x^2+y^2)) + 2x + y."""
    tx = np.linspace(-3.0, 3.0, nrx)
    ty = np.linspace(-3.0, 3.0, nry)
    x, y = np.meshgrid(tx, ty)
    image = 20.0 * np.exp(-0.25 * (x ** 2 + y ** 2)) + 2.0 * x + y
    return x, y, image


def _wrap_phase(image: np.ndarray) -> np.ndarray:
    """Wrap phase to [-pi, pi] using angle of complex exponential."""
    return np.angle(np.exp(1j * image))


x, y, image = _generate_phase()
arr = _wrap_phase(image)
phase_wrapped = np.repeat(arr[None, ...], repeats=10, axis=0)

# the `algo_ls_poisson_pg` is broken up into the following steps:
#    initial dimensionality check
#    gx, gy = wrapped_gradients_stack(phase_wrapped)
#    gx, gy = enforce_periodic_gradients_stack(gx, gy)
#    rhs = divergence_stack(gx, gy)
#    phi = poisson_solve_fft_stack(rhs)
#    phase inversion
#    restore_mean_plane

# Compute forward finite differences of the wrapped phase.
# Then wrap those gradients back into [-pi, pi) so they remain a valid
# representation of wrapped phase changes.
gx1, gy1 = wrapped_gradients_stack(phase_wrapped)
# Enforce periodic boundary conditions on the wrapped gradients so that the
# edges match up. This makes the subsequent FFT-based Poisson solve consistent
# with a periodic domain.
gx2, gy2 = enforce_periodic_gradients_stack(gx1, gy1)
# Take the divergence of those periodic wrapped gradients; this produces the
# right-hand side for the Poisson equation whose solution is the unwrapped phase.
rhs = divergence_stack(gx2, gy2)
# Solve the Poisson equation in the Fourier domain (batched over the stack);
# this yields an unwrapped phase up to a linear ramp.
phi = poisson_solve_fft_stack(rhs)
# Negate to match the sign convention used by the wrapped-gradient computation
# so that gradients of the solution align with the original wrapped gradients.
phi *= -1
phi_restored = restore_mean_plane(phi, phase_wrapped)

# plotting
fontsize = 12
nrows, ncols = 4, 2
fig = plt.figure(figsize=(8, 12))

ax1 = fig.add_subplot(nrows, ncols, 1)
ax1.set_title("Wrapped Phase Example", fontsize=fontsize)
im1 = ax1.imshow(phase_wrapped[1])
fig.colorbar(im1, ax=ax1, fraction=0.046, pad=0.04)

ax2 = fig.add_subplot(nrows, ncols, 2, projection='3d')
ax2.set_title("Wrapped Phase Example 3D proj", fontsize=fontsize)
surf2 = ax2.plot_surface(x, y, phase_wrapped[1], cmap="viridis", linewidth=0, antialiased=True)
fig.colorbar(surf2, ax=ax2, fraction=0.046, pad=0.04, shrink=0.6)

ax3 = fig.add_subplot(nrows, ncols, 3)
ax3.set_title("Wrapped Gradients + Per. Grad. (x)", fontsize=fontsize)
im3 = ax3.imshow(gx2[1])
fig.colorbar(im3, ax=ax3, fraction=0.046, pad=0.04)

ax4 = fig.add_subplot(nrows, ncols, 4)
ax4.set_title("Wrapped Gradients + Per. Grad. (y)", fontsize=fontsize)
im4 = ax4.imshow(gy2[1])
fig.colorbar(im4, ax=ax4, fraction=0.046, pad=0.04)

ax5 = fig.add_subplot(nrows, ncols, 5)
ax5.set_title("Poisson RHS", fontsize=fontsize)
im5 = ax5.imshow(rhs[1])
fig.colorbar(im5, ax=ax5, fraction=0.046, pad=0.04)

ax6 = fig.add_subplot(nrows, ncols, 6)
ax6.set_title("Unwrapped phase (with inversion)", fontsize=fontsize)
im6 = ax6.imshow(phi[1])
fig.colorbar(im6, ax=ax6, fraction=0.046, pad=0.04)

ax7 = fig.add_subplot(nrows, ncols, 7)
ax7.set_title("Unwrapped phase (restore)", fontsize=fontsize)
im7 = ax7.imshow(phi_restored[1])
fig.colorbar(im7, ax=ax7, fraction=0.046, pad=0.04)

ax8 = fig.add_subplot(nrows, ncols, 8, projection='3d')
ax8.set_title("Unwrapped phase (restore) 3D proj", fontsize=fontsize)
surf8 = ax8.plot_surface(x, y, phi_restored[1], cmap="viridis", linewidth=0, antialiased=True)
fig.colorbar(surf8, ax=ax8, fraction=0.046, pad=0.04, shrink=0.6)

for ax in [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8]:
    ax.set_xticks([])
    ax.set_yticks([])

plt.tight_layout()
plt.savefig("walkthrough-ls_poisson_pg.png")
plt.show()

Download base simulation script: unwrapping_simulation.py

LS-Poisson variants

Generate only the LS-Poisson variants figure from unwrapping_simulation. Outputs: unwrapping_comparison.png

unwrapping_simulation_poisson.py

from unwrapping_simulation import main


if __name__ == "__main__":
    main(plot_itoh=False, plot_poisson=True, plot_2d_comparison=False)

Itoh and Skimage unwrapping

Generate only the Itoh/skimage comparison figure from unwrapping_simulation. This graph is a reproduction of the graph in “Two-Dimensional Phase Unwrapping Problem By Dr. Munther Gdeisat and Dr. Francis Lilley”.

Outputs: unwrapping_itoh.png

unwrapping_simulation_itoh.py

from unwrapping_simulation import main


if __name__ == "__main__":
    main(plot_itoh=True, plot_poisson=False, plot_2d_comparison=False)

Phase Unwrapping Algorithm comparison

Generate only the 2D unwrapped comparison/difference maps vs skimage. Outputs: unwrapping_2d_comparison.png

unwrapping_simulation_comparison.py

from unwrapping_simulation import main


if __name__ == "__main__":
    main(plot_itoh=False, plot_poisson=False, plot_2d_comparison=True)

NDArray backend and FFT Engine Selection

Demonstrates how to choose between the NumPy (CPU) and CuPy (GPU) array backends and how the FFT engine is selected automatically to match.

The FFT engine is only relevant for the least-squares Poisson solvers (algo_ls_poisson and algo_ls_poisson_pg); the weighted solver (algo_ls_weighted) is Jacobi-based and uses no FFT.

Run:

python examples/using_backends_and_fft-engines.py

using_backends_and_fft-engines.py

from __future__ import annotations

import numpy as np

import xpunwrap
import xpunwrap.fourier as fourier

# 1. Dumm wrapped phase data (N=1, H=64, W=64)
rng = np.random.default_rng(0)
_tx = np.linspace(-3.0, 3.0, 64)
_x, _y = np.meshgrid(_tx, _tx)
_phase = (20.0 * np.exp(-0.25 * (_x ** 2 + _y ** 2)) +
          rng.normal(0, 0.3, (64, 64)))
wrapped_np = np.angle(np.exp(1j * _phase))[None, ...].astype(np.float32)

# 2. NumPy array backend → FFT engine auto-selects pyFFTW (if installed)
xpunwrap.set_ndarray_backend("numpy")
xp = xpunwrap.get_ndarray_backend()

fft_cls = fourier.get_best_engine()
print(f"Array backend : {xp.__name__} (CPU)")
print(f"FFT engine    : {fft_cls.__name__}")

wrapped = xp.asarray(wrapped_np)
result_auto = xpunwrap.algo_ls_poisson(wrapped, restore_plane=False)
print(f"Result shape  : {result_auto.shape}, dtype: {result_auto.dtype}")

# 3. Force the NumPy FFT engine via PREFERRED_ENGINE
fourier.PREFERRED_ENGINE = "FFTEngineNumpy"
fft_cls_forced = fourier.get_best_engine()
print(f"\nArray backend : {xp.__name__} (CPU)")
print(f"Preferred FFT : {fourier.PREFERRED_ENGINE}")
print(f"FFT engine    : {fft_cls_forced.__name__}")

result_numpy_fft = xpunwrap.algo_ls_poisson(wrapped, restore_plane=False)
print(f"Result shape  : {result_numpy_fft.shape}, "
      f"dtype: {result_numpy_fft.dtype}")

# Both FFT engines produce the same answer (within float32 rounding).
max_diff = float(np.max(np.abs(
    np.asarray(result_auto) - np.asarray(result_numpy_fft)
)))
print(f"Max |auto - numpy| : {max_diff:.2e}  (should be near zero)")

# Reset so the rest of this script uses auto-selection again.
fourier.PREFERRED_ENGINE = None

# 4. CuPy array backend (skipped when CuPy is not installed)
try:
    import cupy  # noqa: F401

    _cupy_available = True
except ImportError:
    _cupy_available = False

if _cupy_available:
    xpunwrap.set_ndarray_backend("cupy")
    xp_gpu = xpunwrap.get_ndarray_backend()
    fft_cls_gpu = fourier.get_best_engine()
    print(f"\nArray backend : {xp_gpu.__name__} (GPU)")
    print(f"FFT engine    : {fft_cls_gpu.__name__}")

    wrapped_gpu = xp_gpu.asarray(wrapped_np)
    result_gpu = xpunwrap.algo_ls_poisson(wrapped_gpu, restore_plane=False)
    print(f"Result shape  : {result_gpu.shape}, dtype: {result_gpu.dtype}")

    # Compare GPU result against CPU result.
    max_diff_gpu = float(np.max(np.abs(
        np.asarray(result_auto) - result_gpu.get()
    )))
    print(f"Max |cpu - gpu| : {max_diff_gpu:.2e}")
else:
    print("\nCuPy not installed — GPU section skipped.")

# 5. Use xp directly for array operations alongside the unwrapper
xpunwrap.set_ndarray_backend("numpy")
xp = xpunwrap.get_ndarray_backend()

wrapped = xp.asarray(wrapped_np)
result = xpunwrap.algo_ls_poisson(wrapped, restore_plane=False)

# All array ops go through xp, so this code is backend-agnostic.
residual = xp.angle(xp.exp(1j * result)) - xp.angle(xp.exp(1j * wrapped))
print(f"\nMean abs residual (wrapped): {float(
    xp.mean(xp.abs(residual))):.4f} rad")

"""
Example output:

Array backend : numpy (CPU)
FFT engine    : FFTEnginePyFFTW
Result shape  : (1, 64, 64), dtype: float32

Array backend : numpy (CPU)
Preferred FFT : FFTEngineNumpy
FFT engine    : FFTEngineNumpy
Result shape  : (1, 64, 64), dtype: float32
Max |auto - numpy| : 4.77e-06  (should be near zero)

Array backend : cupy (GPU)
FFT engine    : FFTEngineCupy
Result shape  : (1, 64, 64), dtype: float32
Max |cpu - gpu| : 4.29e-06

Mean abs residual (wrapped): 0.1254 rad
"""

Processing Pipeline with GPU and Zarr

Example processing pipeline for phase data on the GPU: - Load with Zarr in GPU - Field retrieval - Phase unwrapping - Write to Zarr

processing_pipeline_gpu.py

import zarr
import qpretrieve
import xpunwrap

# set cupy backend
xpunwrap.set_ndarray_backend("cupy")
qpretrieve.set_ndarray_backend("cupy")
xp = xpunwrap.get_ndarray_backend()

# load data with zarr via GPU memory
zarr.config.enable_gpu()

# load to cpu (mightn't need this)
edata = xp.load("./data/hologram_cell.npz")

# load to gpu
# caveat: Zarr loads the data into GPU, but encoding and decoding happens
# on the CPU. It is a work in progress.
# One can use `kvikio` to directly do GPU loading, but
# `pip install kvikio-cu12` only works on linux, or use conda/mamba to install
# on Windows, I assume. See https://docs.rapids.ai/api/kvikio/stable/zarr/
store_holo = zarr.storage.MemoryStore()
store_bg = zarr.storage.MemoryStore()
data_holo = zarr.create_array(store=store_holo, data=edata["data"])
data_bg = zarr.create_array(store=store_bg, data=edata["bg_data"])

# field retrieval on the GPU
fft_interface = qpretrieve.fourier.FFTFilterCupy
holo = qpretrieve.OffAxisHologram(data_holo[:], fft_interface)
bg = qpretrieve.OffAxisHologram(data_bg[:], fft_interface)
holo.run_pipeline(filter_name="disk", filter_size=1 / 2, scale_to_filter=True)
bg.process_like(holo)
phase_wrapped = xp.asarray(holo.phase - bg.phase).astype(xp.float32)

# unwrap the phase on the GPU
phase_unwrapped = xpunwrap.algo_ls_poisson_pg(phase_wrapped)

# write from GPU to Zarr file
path = "temp_unwrapped_cell_phase.zip"
with zarr.storage.ZipStore(path, mode='w') as disk:
    data_phase = zarr.create_array(store=disk, data=phase_unwrapped)

# reopen the file to ensure it is correct
with zarr.storage.ZipStore(path, mode='r') as disk_reload:
    data_reload = zarr.open_array(store=disk_reload, mode='r')
    assert data_phase.shape == data_reload.shape
    assert xp.allclose(phase_unwrapped, data_reload[:])