PyCosmoMMF 1.0.2__tar.gz → 2.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (79) hide show
  1. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/.github/workflows/cd.yml +1 -1
  2. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/.github/workflows/ci.yml +2 -2
  3. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/.gitignore +4 -0
  4. pycosmommf-2.0.0/CHANGELOG.md +224 -0
  5. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/PKG-INFO +8 -5
  6. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/basic_usage.md +66 -16
  7. pycosmommf-2.0.0/docs/changelog.md +3 -0
  8. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/index.md +3 -1
  9. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/pyproject.toml +16 -5
  10. pycosmommf-2.0.0/src/pycosmommf/_cupy_tagging.py +343 -0
  11. pycosmommf-2.0.0/src/pycosmommf/_jax_backend.py +577 -0
  12. pycosmommf-2.0.0/src/pycosmommf/_version.py +24 -0
  13. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/src/pycosmommf/filter.py +42 -37
  14. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/src/pycosmommf/hessian.py +11 -9
  15. pycosmommf-2.0.0/src/pycosmommf/signatures.py +273 -0
  16. pycosmommf-2.0.0/src/pycosmommf/tagging.py +456 -0
  17. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/src/pycosmommf/utils.py +6 -6
  18. pycosmommf-2.0.0/tests/test_cupy_tagging.py +168 -0
  19. pycosmommf-2.0.0/tests/test_integration.py +97 -0
  20. pycosmommf-2.0.0/tests/test_jax_backend.py +213 -0
  21. pycosmommf-2.0.0/tests/test_tagging.py +170 -0
  22. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/tests/test_unit.py +4 -1
  23. pycosmommf-1.0.2/src/pycosmommf/_version.py +0 -21
  24. pycosmommf-1.0.2/src/pycosmommf/signatures.py +0 -138
  25. pycosmommf-1.0.2/src/pycosmommf/tagging.py +0 -307
  26. pycosmommf-1.0.2/tests/test_integration.py +0 -73
  27. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/.copier-answers.yml +0 -0
  28. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/.git_archival.txt +0 -0
  29. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/.gitattributes +0 -0
  30. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/.github/CONTRIBUTING.md +0 -0
  31. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/.github/dependabot.yml +0 -0
  32. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/.github/release.yml +0 -0
  33. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/.pre-commit-config.yaml +0 -0
  34. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/.readthedocs.yaml +0 -0
  35. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/LICENSE +0 -0
  36. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/README.md +0 -0
  37. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/CosmoMMF_Dark.png +0 -0
  38. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/CosmoMMF_light.png +0 -0
  39. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/basic.css +0 -0
  40. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/check-solid.svg +0 -0
  41. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/clipboard.min.js +0 -0
  42. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/copy-button.svg +0 -0
  43. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/copybutton.css +0 -0
  44. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/copybutton.js +0 -0
  45. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/copybutton_funcs.js +0 -0
  46. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/debug.css +0 -0
  47. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/doctools.js +0 -0
  48. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/documentation_options.js +0 -0
  49. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/file.png +0 -0
  50. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/language_data.js +0 -0
  51. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/minus.png +0 -0
  52. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/plus.png +0 -0
  53. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/pygments.css +0 -0
  54. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/scripts/furo-extensions.js +0 -0
  55. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/scripts/furo.js +0 -0
  56. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/scripts/furo.js.LICENSE.txt +0 -0
  57. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/scripts/furo.js.map +0 -0
  58. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/searchtools.js +0 -0
  59. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/skeleton.css +0 -0
  60. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/sphinx_highlight.js +0 -0
  61. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/styles/furo-extensions.css +0 -0
  62. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/styles/furo-extensions.css.map +0 -0
  63. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/styles/furo.css +0 -0
  64. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/_static/styles/furo.css.map +0 -0
  65. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/api/pycosmommf.rst +0 -0
  66. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/api_tab.md +0 -0
  67. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/conf.py +0 -0
  68. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/images/CosmoMMF_Dark.png +0 -0
  69. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/images/CosmoMMF_light.png +0 -0
  70. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/images/final_NEXUSPLUS_Signatures_hydro.png +0 -0
  71. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/images/final_NEXUSPLUS_Signatures_hydro_dark.png +0 -0
  72. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/images/final_tagging_figure.png +0 -0
  73. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/docs/images/final_tagging_figure_dark.png +0 -0
  74. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/noxfile.py +0 -0
  75. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/src/pycosmommf/__init__.py +0 -0
  76. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/src/pycosmommf/_version.pyi +0 -0
  77. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/src/pycosmommf/py.typed +0 -0
  78. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/tests/__init__.py +0 -0
  79. {pycosmommf-1.0.2 → pycosmommf-2.0.0}/tests/test_package.py +0 -0
@@ -29,7 +29,7 @@ jobs:
29
29
  with:
30
30
  fetch-depth: 0
31
31
 
32
- - uses: hynek/build-and-inspect-python-package@v2
32
+ - uses: hynek/build-and-inspect-python-package@v3.0.2
33
33
 
34
34
  publish:
35
35
  needs: [dist]
@@ -38,8 +38,8 @@ jobs:
38
38
  strategy:
39
39
  fail-fast: false
40
40
  matrix:
41
- python-version: ["3.8", "3.12"]
42
- runs-on: [ubuntu-latest, windows-latest, macos-latest, macos-13]
41
+ python-version: ["3.9", "3.12"]
42
+ runs-on: [ubuntu-latest, windows-latest, macos-latest]
43
43
 
44
44
  steps:
45
45
  - uses: actions/checkout@v4
@@ -156,3 +156,7 @@ Thumbs.db
156
156
  # Common editor files
157
157
  *~
158
158
  *.swp
159
+
160
+ # AGENTS
161
+ CLAUDE.md
162
+ AGENTS.md
@@ -0,0 +1,224 @@
1
+ # Changelog
2
+
3
+ All notable changes to `pycosmommf` are documented here. Version numbers come
4
+ from git tags (`hatch-vcs`); releases are published on
5
+ [PyPI](https://pypi.org/project/pycosmommf/).
6
+
7
+ ## Version 2.0.0
8
+
9
+ Changes since [1.0.3](https://pypi.org/project/pycosmommf/1.0.3/), the latest
10
+ release on PyPI. This release changes the return value of
11
+ `calc_structure_bools()` (see Breaking changes below).
12
+
13
+ ### Highlights
14
+
15
+ - **GPU backend.** `maximum_signature()` and `calc_structure_bools()` take a new
16
+ `backend` argument. `backend="gpu"` runs the signatures with JAX and the
17
+ structure tagging with CuPy. Install it with `pip install 'pycosmommf[gpu]'`.
18
+ The GPU backend is capable of computing the cosmic web structure tags on a
19
+ 1024^3 grid in just a couple minutes.
20
+ - **Faster CPU signatures.** The Hessian eigenvalues are now computed with an
21
+ analytic closed-form solver instead of a per-voxel `np.linalg.eigvals` call,
22
+ and the FFTs run in complex64. `maximum_signature()` on the CPU is about 4×
23
+ faster at 128³.
24
+ - **One structure filter.** `calc_structure_bools()` returns a single `uint8`
25
+ array that labels every voxel, instead of four boolean arrays.
26
+ - **Tagging bug fixes.** A race condition in the cluster finder dropped voxels
27
+ from large clumps, and a threshold with no clumps could be selected as the
28
+ cluster threshold. Both are fixed, so cluster filters can differ from 1.0.3.
29
+
30
+ ### Breaking changes
31
+
32
+ - `calc_structure_bools()` returns one `uint8` array, `structure_filter`, with
33
+ the same shape as the density field. Each voxel holds its structure type: `1`
34
+ = cluster, `2` = filament, `3` = wall, `4` = void. Every voxel has exactly one
35
+ type, so the array has no zeros. With `verbose_flag=True` the function returns
36
+ `(structure_filter, summary_data)`.
37
+
38
+ To migrate code written for 1.0.3:
39
+
40
+ ```python
41
+ # 1.0.3
42
+ clusbool, filbool, wallbool, voidbool = calc_structure_bools(rho, sigs, False)
43
+
44
+ # now
45
+ structure_filter = calc_structure_bools(rho, sigs, False)
46
+ clusbool = structure_filter == 1
47
+ filbool = structure_filter == 2
48
+ wallbool = structure_filter == 3
49
+ voidbool = structure_filter == 4
50
+ ```
51
+
52
+ The `uint8` filter uses one byte per voxel. The four boolean arrays used four.
53
+
54
+ - `maximum_signature(..., algorithm="NEXUSPLUS")` raises `ValueError` if the
55
+ density field has a negative voxel. NEXUS+ smooths `log10(density)`, so a
56
+ negative input (usually an overdensity δ passed instead of δ+1 = ρ/⟨ρ⟩) used
57
+ to fill the whole output with NaN without an error. The error message tells
58
+ you to pass `density + 1`. `algorithm="NEXUS"` still accepts any values.
59
+
60
+ - All computation and all returned numbers are float32 on both backends. Inputs
61
+ are converted to float32. The thresholds, virialized fractions, mass curves,
62
+ and volume and mass fractions in `summary_data` are float32. In 1.0.3 some
63
+ were float64.
64
+
65
+ ### Added
66
+
67
+ #### GPU backend
68
+
69
+ - New `backend` argument for `maximum_signature()` and `calc_structure_bools()`.
70
+ It takes `"cpu"` (default) or `"gpu"`. Any other value raises `ValueError`.
71
+ - `backend="cpu"` uses only NumPy, SciPy, Numba and scikit-image. It never
72
+ imports JAX or CuPy.
73
+ - `backend="gpu"` uses JAX for `maximum_signature()` and CuPy for
74
+ `calc_structure_bools()`.
75
+ - New install extra `gpu`: `pip install 'pycosmommf[gpu]'` installs
76
+ `jax[cuda12]>=0.4.30` and `cupy-cuda12x`. JAX and CuPy are imported only when
77
+ `backend="gpu"` is used. If they are missing, you get an `ImportError` that
78
+ gives the install command.
79
+ - New `maximum_signature(..., backend="gpu", return_numpy=False)` keeps the
80
+ signatures on the GPU as a JAX array. Pass that array directly to
81
+ `calc_structure_bools(..., backend="gpu")`: CuPy reads it through DLPack
82
+ without a copy. With the default `return_numpy=True`, the result is a NumPy
83
+ array, the same as the CPU backend.
84
+ - `calc_structure_bools(..., backend="gpu")` accepts NumPy, CuPy or JAX GPU
85
+ arrays for `density_cube`, `max_sigs` and `clusbool`. It always returns NumPy
86
+ arrays.
87
+
88
+ Typical GPU use:
89
+
90
+ ```python
91
+ import pycosmommf
92
+
93
+ sigs = pycosmommf.maximum_signature(Rs, rho, backend="gpu", return_numpy=False)
94
+ structure_filter = pycosmommf.calc_structure_bools(rho, sigs, False, backend="gpu")
95
+ ```
96
+
97
+ How the GPU signatures (`_jax_backend.py`) work:
98
+
99
+ - The whole per-scale step (smoothing, Hessian, eigenvalues, signatures and
100
+ running maximum) is one `jax.jit` graph. The running maximum buffer is donated
101
+ (`donate_argnums`), so XLA updates it in place.
102
+ - The eigenvalues use the same analytic solver as the CPU backend (see
103
+ Performance below). There is no `(N, N, N, 3, 3)` tensor and no call to
104
+ `jnp.linalg.eigh`.
105
+ - The six Hessian components stay as six separate arrays, and the signature
106
+ stage is tiled along the first axis for grids larger than 128³.
107
+ - Peak GPU memory is about 9–10× the size of one float32 field. A 512³ NEXUS+
108
+ run needs 11.8 GiB and takes 5.4 s on a V100S; a naive JAX port ran out of
109
+ memory on a 32 GB GPU. A 1024³ run, the largest supported grid, needs about 65
110
+ GiB, so it needs an 80 GB GPU. A 40 GB GPU is not enough at 1024³. The backend
111
+ runs on a single GPU.
112
+ - On Apple `jax-metal`, which has no FFT or complex support, the backend warns
113
+ and runs JAX on the CPU.
114
+
115
+ How the GPU tagging (`_cupy_tagging.py`) works:
116
+
117
+ - It follows the CPU tagging step for step and calls the same helper functions
118
+ for the threshold choices, so both backends make the same decisions.
119
+ - Clumps are labelled with `cupyx.scipy.ndimage.label` and 26-connectivity, the
120
+ same connectivity as the CPU backend. Per-clump sums use GPU atomics over the
121
+ foreground voxels only.
122
+ - Each point of a mass change curve is one fused float32 reduction kernel, with
123
+ no masks or compacted copies.
124
+ - CuPy allocates from CUDA's stream-ordered memory pool while it tags. This lets
125
+ it reuse memory that JAX freed in the same process. Peak GPU memory is about
126
+ 7× one float32 field (about 29 GiB at 1024³).
127
+ - At 1024³, tagging takes 3.8 s on an A100, against 334 s on the CPU.
128
+ - On the 128³, 256³ and 1024³ validation cubes, GPU and CPU tagging gave the
129
+ same thresholds and the same structure filters. In general the two backends
130
+ agree up to float32 summation order, not bit for bit.
131
+
132
+ Recommended settings for large grids: set `XLA_PYTHON_CLIENT_PREALLOCATE=false`
133
+ and `XLA_PYTHON_CLIENT_ALLOCATOR=cuda_async` before you import JAX. The default
134
+ JAX allocator fragments memory and runs out of memory at 1024³.
135
+
136
+ GPU and CPU signatures differ by about 10⁻⁶ (relative) after smoothing and by up
137
+ to a few 10⁻³ (absolute) in the final signatures, because XLA and SciPy use
138
+ different float32 FFTs. These differences do not change the structure tagging.
139
+
140
+ #### Other additions
141
+
142
+ - If no clump is virialized at any cluster-signature threshold, the cluster
143
+ filter is empty and a warning says that the grid resolution is too low to
144
+ resolve virialized clusters. `make_the_clusbool()` then returns `None` as the
145
+ threshold. In 1.0.3 the function picked a threshold anyway: the lowest one, or
146
+ one with no clumps (see Fixed below).
147
+
148
+ ### Performance
149
+
150
+ - **Analytic eigenvalues on the CPU.** `signatures_from_hessian()` used to build
151
+ a 3×3 array for every voxel and call `np.linalg.eigvals` and `np.sort` on it.
152
+ It now calls `_eigvalsh_3x3_sym`, a scalar Numba function that uses the
153
+ closed-form trigonometric solution for real symmetric 3×3 matrices (Smith
154
+ 1961, _Communications of the ACM_ 4, 168). The function allocates nothing per
155
+ voxel and returns the eigenvalues already sorted. It rescales each matrix by
156
+ its largest absolute entry before it applies the formula, so it stays accurate
157
+ for near-zero Hessians in voids. The CPU and GPU backends use the same solver.
158
+ - **complex64 FFTs on the CPU.** Smoothing and the Hessian use `scipy.fft`
159
+ instead of `np.fft`. `scipy.fft` keeps float32 input in complex64. With numpy
160
+ < 2, `np.fft` always computes in complex128.
161
+ - **Fewer float64 promotions.** Numba kernels use float32 constants, so they no
162
+ longer promote float32 arithmetic to float64.
163
+
164
+ CPU timings, 1.0.3 against the current code, on a synthetic log-normal 128³
165
+ field with NEXUS+, on one node with 128 CPU cores. Each step was timed on its
166
+ second call, so the Numba compilation of module-level functions is excluded:
167
+
168
+ | Step | 1.0.3 | Now | Speedup |
169
+ | --------------------------------------- | ------ | ------ | ------- |
170
+ | `smooth_loggauss()`, 1 scale | 1.00 s | 0.23 s | 4.3× |
171
+ | `fast_hessian_from_smoothed()`, 1 scale | 4.68 s | 1.02 s | 4.6× |
172
+ | `signatures_from_hessian()`, 1 scale | 0.49 s | 0.17 s | 2.9× |
173
+ | `maximum_signature()`, 1 scale | 6.47 s | 1.60 s | 4.0× |
174
+ | `maximum_signature()`, 10 scales | 58.9 s | 13.8 s | 4.3× |
175
+
176
+ At 256³ the current code takes 53 s for 10 scales. The 1.0.3 run at 256³ did not
177
+ finish (the process was killed). On an A100 GPU, the same 256³ field takes 1.7 s
178
+ for 10 scales with `backend="gpu"` (9.7 s on the first call, which includes JIT
179
+ compilation), and the GPU tagging takes 0.08 s against 5.9 s on the CPU, with an
180
+ identical structure filter.
181
+
182
+ ### Fixed
183
+
184
+ - **Race condition in the cluster finder.** The per-clump overdensity sums ran
185
+ in a Numba `prange` loop that added into shared arrays. Threads overwrote each
186
+ other's additions and dropped voxels from large clumps, so the mean
187
+ overdensity of a clump, and therefore the cluster threshold, could be wrong
188
+ and change from run to run. The sums are now a serial float32 loop
189
+ (`_clump_sums`).
190
+ - **Thresholds without clumps.** If a cluster-signature threshold had no clumps,
191
+ its virialized fraction was 0/0 = NaN. `np.argmin` returns the first NaN, so
192
+ that empty threshold could be selected as the cluster threshold. These
193
+ fractions are now NaN on purpose, and `np.nanargmin` skips them.
194
+ - **Threshold comparisons.** The thresholds were float64 and the signatures were
195
+ float32. Thresholds are now float32, so the comparisons are exact and do not
196
+ depend on NumPy's scalar promotion rules.
197
+
198
+ ### Dependencies and packaging
199
+
200
+ - `scipy` is now a declared dependency. The CPU FFTs use `scipy.fft`.
201
+ - New optional extra `gpu` (`jax[cuda12]>=0.4.30`, `cupy-cuda12x`).
202
+ - The `requires-python` upper bound (`<3.13`) is removed. The requirement is now
203
+ `>=3.8`.
204
+
205
+ ### Internal changes
206
+
207
+ - `calc_structure_bools()` separates the computation from the printed report.
208
+ The CPU and GPU backends share the report and the `summary_data` code.
209
+ - `wavevectors3D()` builds the wavevectors directly in float32. The values are
210
+ unchanged.
211
+ - `sphere()`, `cylinder()` and `wall()` in `utils.py` return float32 arrays.
212
+
213
+ ### Documentation and tests
214
+
215
+ - `docs/basic_usage.md` describes the new `structure_filter` return value and
216
+ the GPU backend.
217
+ - New `tests/test_tagging.py` (CPU tagging), `tests/test_jax_backend.py` (JAX vs
218
+ CPU agreement for each stage) and `tests/test_cupy_tagging.py` (CuPy vs CPU
219
+ tagging with NumPy, CuPy and JAX inputs). The GPU tests skip when JAX, CuPy or
220
+ a CUDA GPU is not available, so CI runs only the CPU tests.
221
+
222
+ ## 1.0.3 and earlier
223
+
224
+ See the [GitHub releases](https://github.com/James11222/pycosmommf/releases).
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.4
1
+ Metadata-Version: 2.5
2
2
  Name: pycosmommf
3
- Version: 1.0.2
3
+ Version: 2.0.0
4
4
  Summary: A package for identifying structures in the cosmic web.
5
5
  Project-URL: Homepage, https://github.com/James11222/pycosmommf
6
6
  Project-URL: Bug Tracker, https://github.com/James11222/pycosmommf/issues
@@ -27,7 +27,7 @@ License: Copyright 2024 James Sunseri
27
27
  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
28
28
  SOFTWARE.
29
29
  License-File: LICENSE
30
- Classifier: Development Status :: 1 - Planning
30
+ Classifier: Development Status :: 5 - Production/Stable
31
31
  Classifier: Intended Audience :: Developers
32
32
  Classifier: Intended Audience :: Science/Research
33
33
  Classifier: License :: OSI Approved :: MIT License
@@ -35,17 +35,17 @@ Classifier: Operating System :: OS Independent
35
35
  Classifier: Programming Language :: Python
36
36
  Classifier: Programming Language :: Python :: 3
37
37
  Classifier: Programming Language :: Python :: 3 :: Only
38
- Classifier: Programming Language :: Python :: 3.8
39
38
  Classifier: Programming Language :: Python :: 3.9
40
39
  Classifier: Programming Language :: Python :: 3.10
41
40
  Classifier: Programming Language :: Python :: 3.11
42
41
  Classifier: Programming Language :: Python :: 3.12
43
42
  Classifier: Topic :: Scientific/Engineering
44
43
  Classifier: Typing :: Typed
45
- Requires-Python: <3.13,>=3.8
44
+ Requires-Python: >=3.9
46
45
  Requires-Dist: numba
47
46
  Requires-Dist: numpy>=1.22
48
47
  Requires-Dist: scikit-image
48
+ Requires-Dist: scipy
49
49
  Provides-Extra: dev
50
50
  Requires-Dist: pytest-cov>=3; extra == 'dev'
51
51
  Requires-Dist: pytest>=6; extra == 'dev'
@@ -55,6 +55,9 @@ Requires-Dist: myst-parser>=0.13; extra == 'docs'
55
55
  Requires-Dist: sphinx-autodoc-typehints; extra == 'docs'
56
56
  Requires-Dist: sphinx-copybutton; extra == 'docs'
57
57
  Requires-Dist: sphinx>=7.0; extra == 'docs'
58
+ Provides-Extra: gpu
59
+ Requires-Dist: cupy-cuda12x; extra == 'gpu'
60
+ Requires-Dist: jax[cuda12]>=0.4.30; extra == 'gpu'
58
61
  Provides-Extra: test
59
62
  Requires-Dist: pytest-cov>=3; extra == 'test'
60
63
  Requires-Dist: pytest>=6; extra == 'test'
@@ -13,16 +13,16 @@ to analyze the effects of baryonic matter on the Cosmic Web.
13
13
  The NEXUS+ algorithm contains several steps as described in our paper
14
14
  ([Sunseri et al. 2022](https://ui.adsabs.harvard.edu/abs/2023PhRvD.107b3514S/abstract)).
15
15
  In general, we start with a density field (note: we more specifically mean a 1 +
16
- δ field), smooth it with a logarithmic Gaussian smoothing filter, then compute
17
- the hessian of the smoothed density field, use the eigenvalues of the hessian
18
- matrix to calculate the structure type signatures, find the maximum signatures
19
- over a range of smoothing scales, and apply physically based threshold criterion
20
- to categorize structures within the Cosmic Web. The entire package is
21
- implemented in `python` and all of these steps are summarized inside of two
22
- functions. The first function `maximum_signature()` does the first several steps
23
- of the NEXUS+ algorithm to compute the maximum structure signatures, the second
24
- function is `calc_structure_bools()` which uses physical criteria to tag
25
- structures into 4 categories: clusters, filaments, walls, and voids.
16
+ δ field), smooth it with a log-Gaussian smoothing filter, then compute the
17
+ hessian of the smoothed density field, use the eigenvalues of the hessian matrix
18
+ to calculate the structure type signatures, find the maximum signatures over a
19
+ range of smoothing scales, and apply physically based threshold criterion to
20
+ categorize structures within the Cosmic Web. The entire package is implemented
21
+ in `python` and all of these steps are summarized inside of two functions. The
22
+ first function `maximum_signature()` does the first several steps of the NEXUS+
23
+ algorithm to compute the maximum structure signatures, the second function is
24
+ `calc_structure_bools()` which uses physical criteria to tag structures into 4
25
+ categories: clusters, filaments, walls, and voids.
26
26
 
27
27
  We also make the data products from our paper
28
28
  ([Sunseri et al. 2022](https://ui.adsabs.harvard.edu/abs/2023PhRvD.107b3514S/abstract))
@@ -81,7 +81,7 @@ API Docs tab for a detailed summary of what is inside `summary_data`.
81
81
  verbose_flag = True # or False
82
82
 
83
83
 
84
- clusbool, filbool, wallbool, voidbool, summary_data = PyCosmoMMF.calc_structure_bools(
84
+ structure_filter, summary_data = PyCosmoMMF.calc_structure_bools(
85
85
  density_field, max_signatures, verbose_flag
86
86
  ) # tag structures
87
87
  ```
@@ -92,11 +92,23 @@ clusbool, filbool, wallbool, voidbool, summary_data = PyCosmoMMF.calc_structure_
92
92
  verbose_flag = False
93
93
 
94
94
 
95
- clusbool, filbool, wallbool, voidbool = PyCosmoMMF.calc_structure_bools(
95
+ structure_filter = PyCosmoMMF.calc_structure_bools(
96
96
  density_field, max_signatures, verbose_flag
97
97
  ) # tag structures
98
98
  ```
99
99
 
100
+ `structure_filter` is a `uint8` array with the same shape as the density field
101
+ that gives the structure type of every voxel: `1` = cluster, `2` = filament, `3`
102
+ = wall and `4` = void. Every voxel has exactly one type. Boolean filters for a
103
+ single structure type are one comparison away:
104
+
105
+ ```python
106
+ clusbool = structure_filter == 1
107
+ filbool = structure_filter == 2
108
+ wallbool = structure_filter == 3
109
+ voidbool = structure_filter == 4
110
+ ```
111
+
100
112
  We also note in the `calc_structure_bools()` function, one can use their own
101
113
  cluster boolean filter instead of the one generated by the NEXUS+ formalism
102
114
  (using virialization of clusters as a tool for determining spurious detections).
@@ -114,7 +126,7 @@ clusbool_ext = np.load(
114
126
  verbose_flag = False
115
127
 
116
128
 
117
- clusbool, filbool, wallbool, voidbool = PyCosmoMMF.calc_structure_bools(
129
+ structure_filter = PyCosmoMMF.calc_structure_bools(
118
130
  density_field, max_signatures, verbose_flag, clusbool_ext
119
131
  ) # tag structures
120
132
  ```
@@ -128,9 +140,8 @@ but other values can be 200 or 500 corresponding to `R_200` or `R_500`.
128
140
  average overdensity greater than this value, they are thought to be
129
141
  virialized/collapsed.
130
142
 
131
- The boolean filters for each structure type produced by `calc_structure_bools()`
132
- can be used to tag structures within a density field, the results of this can be
133
- seen below
143
+ The structure filter produced by `calc_structure_bools()` can be used to tag
144
+ structures within a density field, the results of this can be seen below
134
145
 
135
146
  ```{image} images/final_tagging_figure.png
136
147
  :class: only-light
@@ -141,6 +152,45 @@ seen below
141
152
  :class: only-dark
142
153
  ```
143
154
 
155
+ #### Running on a GPU
156
+
157
+ Both `maximum_signature()` and `calc_structure_bools()` take a `backend`
158
+ argument. The default, `backend="cpu"`, uses NumPy, SciPy, Numba and
159
+ scikit-image. `backend="gpu"` computes the signatures with JAX and tags the
160
+ structures with CuPy. It needs a CUDA GPU and the `gpu` extra:
161
+
162
+ ```bash
163
+ pip install 'pycosmommf[gpu]'
164
+ ```
165
+
166
+ To keep the signatures on the GPU between the two steps, pass
167
+ `return_numpy=False` to `maximum_signature()`. It then returns a JAX array, and
168
+ `calc_structure_bools()` reads that array without copying it. The returned
169
+ `structure_filter` is always a NumPy array.
170
+
171
+ ```python
172
+ import os
173
+
174
+ # set these before JAX is imported; they are needed for large grids
175
+ os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
176
+ os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "cuda_async"
177
+
178
+ import PyCosmoMMF
179
+
180
+ max_signatures = PyCosmoMMF.maximum_signature(
181
+ Rs, density_field, algorithm="NEXUSPLUS", backend="gpu", return_numpy=False
182
+ )
183
+ structure_filter = PyCosmoMMF.calc_structure_bools(
184
+ density_field, max_signatures, verbose_flag, backend="gpu"
185
+ )
186
+ ```
187
+
188
+ The GPU backend runs on one GPU. The GPU signatures need about 9–10× the memory
189
+ of one float32 density field, and the GPU tagging needs about 7×. A 1024³ grid,
190
+ the largest grid PyCosmoMMF supports, needs about 65 GiB, so it needs a GPU with
191
+ 80 GB of memory. A 40 GB GPU is not enough at 1024³. The CPU and GPU backends
192
+ agree to float32 precision and give the same structure tagging.
193
+
144
194
  #### Additional Code Information
145
195
 
146
196
  - Note: The NEXUS+ implementation of tagging clusters is highly dependent on the
@@ -0,0 +1,3 @@
1
+ ```{include} ../CHANGELOG.md
2
+
3
+ ```
@@ -4,6 +4,7 @@
4
4
 
5
5
  basic_usage
6
6
  api_tab
7
+ changelog
7
8
 
8
9
  ```
9
10
 
@@ -24,7 +25,8 @@ api_tab
24
25
  # PyCosmoMMF
25
26
 
26
27
  PyCosmoMMF is a package for identifying structures in the Cosmic Web from a
27
- given overdensity field.
28
+ given overdensity field. Version 2.0 includes a gpu backend for running the code
29
+ with NVIDIA GPUs.
28
30
 
29
31
  ### Authors
30
32
 
@@ -11,9 +11,9 @@ authors = [
11
11
  description = "A package for identifying structures in the cosmic web."
12
12
  readme = "README.md"
13
13
  license.file = "LICENSE"
14
- requires-python = ">=3.8, <3.13"
14
+ requires-python = ">=3.9"
15
15
  classifiers = [
16
- "Development Status :: 1 - Planning",
16
+ "Development Status :: 5 - Production/Stable",
17
17
  "Intended Audience :: Science/Research",
18
18
  "Intended Audience :: Developers",
19
19
  "License :: OSI Approved :: MIT License",
@@ -21,7 +21,6 @@ classifiers = [
21
21
  "Programming Language :: Python",
22
22
  "Programming Language :: Python :: 3",
23
23
  "Programming Language :: Python :: 3 :: Only",
24
- "Programming Language :: Python :: 3.8",
25
24
  "Programming Language :: Python :: 3.9",
26
25
  "Programming Language :: Python :: 3.10",
27
26
  "Programming Language :: Python :: 3.11",
@@ -33,7 +32,8 @@ dynamic = ["version"]
33
32
  dependencies = [
34
33
  "numba",
35
34
  "numpy>=1.22",
36
- "scikit-image"
35
+ "scikit-image",
36
+ "scipy",
37
37
  ]
38
38
 
39
39
  [project.optional-dependencies]
@@ -52,6 +52,10 @@ docs = [
52
52
  "sphinx_autodoc_typehints",
53
53
  "furo>=2023.08.17",
54
54
  ]
55
+ gpu = [
56
+ "jax[cuda12]>=0.4.30",
57
+ "cupy-cuda12x",
58
+ ]
55
59
 
56
60
  [project.urls]
57
61
  Homepage = "https://github.com/James11222/pycosmommf"
@@ -75,6 +79,7 @@ addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"]
75
79
  xfail_strict = true
76
80
  filterwarnings = [
77
81
  "error",
82
+ "ignore:jax-metal does not support FFT:UserWarning",
78
83
  ]
79
84
  log_cli_level = "INFO"
80
85
  testpaths = [
@@ -84,9 +89,16 @@ testpaths = [
84
89
 
85
90
  [tool.coverage]
86
91
  run.source = ["pycosmommf"]
92
+ # GitHub Actions has no free GPUs, so the GPU backend is left out of coverage:
93
+ # its two modules and the branches that dispatch to them.
94
+ run.omit = [
95
+ "*/pycosmommf/_jax_backend.py",
96
+ "*/pycosmommf/_cupy_tagging.py",
97
+ ]
87
98
  report.exclude_also = [
88
99
  '\.\.\.',
89
100
  'if typing.TYPE_CHECKING:',
101
+ 'if backend == "gpu":',
90
102
  ]
91
103
 
92
104
  [tool.ruff]
@@ -113,7 +125,6 @@ extend-select = [
113
125
  "YTT", # flake8-2020
114
126
  "EXE", # flake8-executable
115
127
  "NPY", # NumPy specific rules
116
- "PD", # pandas-vet
117
128
  ]
118
129
  ignore = [
119
130
  "F403", # star imports