Add CuPy GPU backend and phase congruence post-processing - #20
Merged
Conversation
Add CuPy as an alternate GPU backend alongside JAX. CuPy is a better fit for this use case since we only need GPU-accelerated FFTs (via cupyx.scipy.fft.dctn/idctn) without autodiff. The backend is selected via `backend="cupy"` in the API or `--backend cupy` on the CLI. Add a `--congruent` flag that post-processes the unwrapped phase so it differs from the wrapped input by an integer multiple of 2*pi at every pixel. The ADMM solver doesn't guarantee this property, so `make_congruent` computes k = round((unwrapped - wrapped) / 2pi) and returns wrapped + 2*pi*k. https://claude.ai/code/session_01XAyGBPgci1PkE5UpqMDMm9
Replace the duplicated _*_jax and _*_cupy functions with a single set of generic helpers (_apply_gradient_x, _apply_gradient_y, _apply_divergence, _est_wrapped_gradient, _p_shrink, _make_laplace_kernel) that take an xp array module parameter (np, jnp, or cp). Both the JAX JIT-compiled step and the CuPy unwrap loop now call the same generic helpers—jnp is just a Python module constant to JAX's tracer, so it works transparently with @jax.jit. The JAX public API functions (est_wrapped_gradient_jax, p_shrink_jax, make_laplace_kernel_jax) are kept as thin wrappers for backward compat. https://claude.ai/code/session_01XAyGBPgci1PkE5UpqMDMm9
Drop separate unwrap_jax/unwrap_cupy/JAX wrappers and the numpy sparse-matrix code path. Now there's one ADMM loop parameterized by xp (the array module) and a pair of DCT/IDCT callables resolved at the top of unwrap(). All helper functions (est_wrapped_gradient, p_shrink, make_laplace_kernel) take xp=np by default. core.py goes from ~560 lines to ~277. https://claude.ai/code/session_01XAyGBPgci1PkE5UpqMDMm9
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds GPU acceleration support via CuPy and introduces a phase congruence post-processing function to improve unwrapping results.
Key Changes
CuPy GPU Backend: Implemented full CuPy-accelerated phase unwrapping with:
cupyx.scipy.fftfor efficient Fourier-domain solvingPhase Congruence Post-Processing: Added
make_congruent()function that:congruent=Trueflag inunwrap()API Enhancements:
unwrap()function withbackendparameter supporting 'numpy', 'jax', and 'cupy'congruentparameter to enable phase congruence correction--backendand--congruentoptionsTesting: Added comprehensive test coverage for:
Code Quality:
__all__export list for readabilitypyproject.tomlhttps://claude.ai/code/session_01XAyGBPgci1PkE5UpqMDMm9