Repository navigation
Added GPU backend for pycosmommf - #9
Merged
Merged
Conversation
Adds backend='jax' parameter to maximum_signature(). The JAX path replaces Numba JIT loops with fully-vectorised XLA operations: k-space filter construction via broadcasting, all six Hessian components in a single jnp.stack of IFFT calls, and batched jnp.linalg.eigh for eigenvalue computation instead of a per-voxel loop. Output is always returned as a numpy ndarray for API compatibility. Includes test_jax_backend.py verifying shape and value agreement with the CPU backend (rtol≤5e-4 for individual pipeline stages, atol≤5e-3 end-to-end, reflecting inherent float32 XLA vs FFTW differences). JAX tests auto-skip when jax is not installed. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
The previous JAX backend OOMed on 512^3 inputs at the eigh step because
signatures_from_hessian_jax materialized a (N,N,N,3,3) tensor and called
jnp.linalg.eigh, which also allocates an eigenvector workspace and a
transposed copy. Combined with jnp.stack of 6 Hessian components and an
un-jitted per-scale Python loop, peak live memory was ~25-30x the scalar
field size.
Rewrite the hot path:
* Replace eigh with an analytic Smith-1961 closed-form solver for the
real symmetric 3x3 eigenvalue problem, operating directly on the 6
component arrays. No (...,3,3) tensor, no eigenvectors, no LAPACK
workspace. Rescale by max-abs before the trig formula so the
det/p^3 quotient stays well-conditioned on near-zero background
Hessians.
* Drop the jnp.stack of Hessian components in the hot path; they flow
into the eigenvalue solver as 6 separate (N,N,N) arrays. The public
fast_hessian_from_smoothed_jax keeps the (N,N,N,6) layout for
backward compatibility.
* JIT the whole per-scale step with donate_argnums=(sigmax,) so XLA
can reuse the running-max buffer and fuse eigenvalue temporaries
into single kernel passes.
* Tile the signature stage along axis 0 (auto-picked 8 tiles for
grids >= 128^3); the tile loop unrolls at trace time inside the
jit, giving XLA the full graph to plan slab-by-slab in-place
updates of sigmax.
* Upload wavevectors to the device once, not per scale.
Peak GPU memory on a real 512^3 NEXUSPLUS run drops from >32 GiB (OOM)
to 11.8 GiB; runtime is 5.4 s on a V100S.
Also add input validation for NEXUSPLUS in both the dispatcher and the
JAX entry point: the log-Gauss smoothing requires a strictly positive
field, but passing an overdensity delta (which is negative in voids)
previously produced NaN silently through the whole pipeline. Now it
raises ValueError with a message pointing at the delta -> delta+1 fix.
Loosen the direct-eigensolver-comparison test atol from 1e-5 to 1e-4:
the analytic Smith formula and NumPy's iterative eigvals differ by a
few ULP on near-degenerate Hessians, which the signature step's
Heaviside thresholds can amplify to ~5e-5 absolute. All such voxels
sit far below the structure-tagging cutoff.
All 31 tests pass.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…oth the CPU and GPU versions of the code
- src/pycosmommf/tagging.py: the race condition, threshold comparisons, float64 sums, NaN handling. - tests/test_tagging.py 2. CuPy backend: - src/pycosmommf/_cupy_tagging.py (new) - tagging.py: the backend= option, plus splitting the computation from the printed report. - signatures.py and _jax_backend.py: the return_numpy option. - pyproject.toml: the [cupy] extra. - CLAUDE.md - tests/test_cupy_tagging.py and tests/test_jax_backend.py 3. Speedup: already folded into _cupy_tagging.py; it no longer stands as a separate change.
…hroughout Backends - maximum_signature() and calc_structure_bools() now take backend=cpu | gpu. cpu uses only NumPy, SciPy, Numba and scikit-image; gpu uses JAX for the signatures and CuPy for the tagging. The old jax name is removed and raises ValueError. - New _cupy_tagging.py runs calc_structure_bools on the GPU: 26-connected clump labelling with cupyx.scipy.ndimage.label, per-clump sums with atomics over foreground voxels only, and fused reduction kernels for the mass curves. JAX arrays are accepted without a copy via DLPack, and CuPy allocates from CUDA's stream-ordered pool so it can reuse memory JAX freed. At 1024^3 tagging takes 3.8 s on an A100 vs 334 s on CPU. - maximum_signature(..., backend=gpu, return_numpy=False) keeps the result on the GPU for the tagging step. - Install extra is now pycosmommf[gpu] (jax[cuda12] + cupy-cuda12x); scipy is listed as a dependency. float32 everywhere - Inputs are converted to float32 and every array, FFT, sum and returned value stays float32 on both backends. CPU FFTs use scipy.fft (complex64) instead of np.fft, which computes in complex128 on numpy < 2. Numba kernels use float32 constants to avoid promotion. Tagging fixes - Per-clump overdensity sums no longer race: the Numba prange loop with += on shared arrays dropped voxels from large clumps. It is now a serial float32 loop. - A threshold with no clumps gets a NaN virialized fraction instead of dividing 0/0 and possibly being selected by argmin. - If no clump is virialized at any threshold, the cluster filter is empty and a warning says the grid resolution is too low to resolve virialized clusters. - Computation is split from the printed report so both backends share the report and summary_data code. Validation on real 128^3, 256^3 and 1024^3 cubes: GPU and CPU tagging give identical thresholds and identical masks (0 mismatched voxels).
(1 cluster, 2 filament, 3 wall, 4 void; no zeros) instead of four boolean arrays; with verbose_flag it returns (structure_filter, summary_data).
There was a problem hiding this comment.
🔵 Needs a closer look
It introduces ~1,900 lines across two new GPU backends, replaces the core eigenvalue computation, migrates everything to float32, and makes a breaking public API change, with the GPU paths not exercised in CI.
1 open issue
What changed in this PR
This PR adds a GPU execution path to pycosmommf and migrates the whole pipeline to float32. Despite the description's framing around documentation, the substantive work is new numerical code: a JAX backend for maximum_signature() and a CuPy backend for calc_structure_bools(), both reached via a new backend="gpu" parameter. It also replaces the per-voxel np.linalg.eigvals call with an analytic Smith-1961 3×3 symmetric eigensolver (shared, in sync, between the CPU Numba path and the JAX path), switches FFTs to scipy.fft (keeping float32 in complex64), and changes calc_structure_bools() to return a single uint8 structure_filter (1=cluster, 2=filament, 3=wall, 4=void) instead of four boolean masks.
Changes:
- New GPU backends (
_jax_backend.py,_cupy_tagging.py), lazily imported and not star-exported; old backend names"jax"/"cupy"now raiseValueError. - Project-wide float32 migration, new analytic eigensolver, and refactor of
tagging.pyinto reusable helpers shared by both backends; breaking change tocalc_structure_bools()return type. - Supporting updates:
scipydependency +[gpu]extra, relaxedrequires-python, new/updated tests, docs, and a contributor guide (CLAUDE.md).
| File | Description |
|---|---|
src/pycosmommf/_jax_backend.py |
New JAX backend: memory-efficient hessian/eigenvalues/signatures, tiling, device fallback. |
src/pycosmommf/_cupy_tagging.py |
New CuPy tagging backend mirroring CPU logic with fused kernels and async memory pool. |
src/pycosmommf/tagging.py |
Split compute/report, new structure_filter output, shared helpers, cluster-threshold logic. |
src/pycosmommf/signatures.py |
Analytic Smith-1961 eigensolver, density validation, backend/return_numpy dispatch. |
src/pycosmommf/filter.py |
scipy.fft, float32 constants, new _wavevector replacing fftfreq. |
src/pycosmommf/hessian.py |
scipy.fft, float32 -R² prefactor. |
src/pycosmommf/utils.py |
Synthetic geometry generators now return float32. |
tests/* |
New/updated unit, integration, tagging, JAX and CuPy tests (GPU tests skip without hardware). |
pyproject.toml |
scipy dep, [gpu] extra, relaxed Python cap, jax-metal warning filter. |
docs/basic_usage.md |
Usage updated for the single structure_filter return value. |
CLAUDE.md |
New contributor/architecture guide. |
.gitignore |
Ignores CLAUDE.md/AGENTS.md (trailing whitespace + contradicts committing CLAUDE.md). |
Review effort: Balanced
Comment on lines
+1
to
+10
| """ | ||
| JAX backend for pycosmommf. Requires jax to be installed: | ||
|
|
||
| pip install 'pycosmommf[gpu]' | ||
|
|
||
| Reached through ``maximum_signature(..., backend="gpu")``. All public | ||
| functions mirror the CPU API but operate on JAX arrays, in float32 (JAX's | ||
| default precision) like the CPU backend. maximum_signature_jax() returns a | ||
| numpy ndarray unless ``return_numpy=False``, in which case the JAX array | ||
| stays on the GPU for ``calc_structure_bools(..., backend="gpu")``. |
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
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.
This pull request introduces several key updates to the documentation and package configuration for
pycosmommf, focusing on clarifying usage, improving dependency management, and providing comprehensive guidance for contributors and users. The most significant changes include the addition of an extensive contributor guide (CLAUDE.md), updates to usage instructions in the documentation to reflect the current API, and enhancements to optional dependency handling for GPU support.Documentation and Usage Updates:
CLAUDE.mdfile outlining the project's architecture, backend design, key conventions, and testing structure to guide contributors and advanced users.docs/basic_usage.mdto reflect thatcalc_structure_bools()now returns a singlestructure_filterarray (with optionalsummary_data), replacing the previous multiple boolean mask outputs. Added clear instructions on how to derive boolean masks for each structure type fromstructure_filter. [1] [2] [3] [4]Dependency and Configuration Improvements:
scipyas a required dependency and introduced an optional[gpu]extra inpyproject.tomlfor installing JAX and CuPy with CUDA support, facilitating GPU backend usage. [1] [2]Testing and CI Adjustments:
jax-metallacking FFT support, ensuring smoother test runs on Apple hardware.These changes improve the clarity and maintainability of the project, making it easier for users to install, test, and contribute, while ensuring the documentation accurately reflects the current API and architecture.