disco-eb 1.0.1__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 (25) hide show
  1. disco_eb-1.0.1/LICENSE +21 -0
  2. disco_eb-1.0.1/PKG-INFO +178 -0
  3. disco_eb-1.0.1/README.md +145 -0
  4. disco_eb-1.0.1/pyproject.toml +174 -0
  5. disco_eb-1.0.1/setup.cfg +4 -0
  6. disco_eb-1.0.1/src/disco_eb.egg-info/PKG-INFO +178 -0
  7. disco_eb-1.0.1/src/disco_eb.egg-info/SOURCES.txt +23 -0
  8. disco_eb-1.0.1/src/disco_eb.egg-info/dependency_links.txt +1 -0
  9. disco_eb-1.0.1/src/disco_eb.egg-info/requires.txt +12 -0
  10. disco_eb-1.0.1/src/disco_eb.egg-info/top_level.txt +1 -0
  11. disco_eb-1.0.1/src/discoeb/__init__.py +87 -0
  12. disco_eb-1.0.1/src/discoeb/component.py +1314 -0
  13. disco_eb-1.0.1/src/discoeb/eb_model.py +939 -0
  14. disco_eb-1.0.1/src/discoeb/perturbations.py +1620 -0
  15. disco_eb-1.0.1/src/discoeb/recombination.py +690 -0
  16. disco_eb-1.0.1/src/discoeb/spline_interpolation.py +246 -0
  17. disco_eb-1.0.1/src/discoeb/util.py +49 -0
  18. disco_eb-1.0.1/tests/test_boundary.py +32 -0
  19. disco_eb-1.0.1/tests/test_component.py +679 -0
  20. disco_eb-1.0.1/tests/test_eb_model.py +254 -0
  21. disco_eb-1.0.1/tests/test_eb_sparsity.py +286 -0
  22. disco_eb-1.0.1/tests/test_perturbations.py +428 -0
  23. disco_eb-1.0.1/tests/test_reionization.py +74 -0
  24. disco_eb-1.0.1/tests/test_split.py +82 -0
  25. disco_eb-1.0.1/tests/test_thermo.py +122 -0
disco_eb-1.0.1/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2023-2026 Oliver Hahn, Florian List, Natalia Porqueres and Lawrence Berry
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
disco_eb-1.0.1/PKG-INFO ADDED
@@ -0,0 +1,178 @@
1
+ Metadata-Version: 2.4
2
+ Name: disco-eb
3
+ Version: 1.0.1
4
+ Summary: A differentiable linear Einstein-Boltzmann solver for cosmology, on the GPU
5
+ Author: Florian List, Natalia Porqueres, Lawrence Berry
6
+ Author-email: Oliver Hahn <oliver.hahn@univie.ac.at>
7
+ License-Expression: MIT
8
+ Project-URL: Documentation, https://cosmo-sims.github.io/DISCO-EB/
9
+ Project-URL: Source, https://github.com/cosmo-sims/DISCO-EB
10
+ Project-URL: Issues, https://github.com/cosmo-sims/DISCO-EB/issues
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Environment :: GPU :: NVIDIA CUDA
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: Operating System :: POSIX :: Linux
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Classifier: Topic :: Scientific/Engineering :: Astronomy
20
+ Classifier: Topic :: Scientific/Engineering :: Physics
21
+ Requires-Python: >=3.10
22
+ Description-Content-Type: text/markdown
23
+ License-File: LICENSE
24
+ Requires-Dist: jax>=0.6.2
25
+ Requires-Dist: numpy>=2.2.6
26
+ Requires-Dist: modax-solvers>=0.2.2
27
+ Requires-Dist: numba-enzyme-cuda>=0.3.0
28
+ Provides-Extra: cuda12
29
+ Requires-Dist: jax[cuda12]>=0.6.2; extra == "cuda12"
30
+ Provides-Extra: cuda13
31
+ Requires-Dist: jax[cuda13]>=0.9.1; python_version >= "3.11" and extra == "cuda13"
32
+ Dynamic: license-file
33
+
34
+ <h1 align="center">DISCO-EB</h1>
35
+
36
+ <!--
37
+ The coverage badge reads a gist that the `tests` workflow rewrites on every
38
+ push to main. The id below is the repository variable COVERAGE_GIST_ID.
39
+ These badges sit above the `overview` snippet marker, so the documentation
40
+ site does not repeat them.
41
+ -->
42
+ <p align="center">
43
+ <a href="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/tests.yml"><img alt="tests" src="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/tests.yml/badge.svg?branch=main"></a>
44
+ <a href="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/tests.yml"><img alt="coverage" src="https://img.shields.io/endpoint?url=https%3A%2F%2Fgist.githubusercontent.com%2Flawrenceberry%2Fb451e14a9620fabb0d49c345f13b5414%2Fraw%2Fdiscoeb-coverage.json"></a>
45
+ <a href="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/lint.yml"><img alt="lint" src="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/lint.yml/badge.svg?branch=main"></a>
46
+ <a href="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/types.yml"><img alt="types" src="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/types.yml/badge.svg?branch=main"></a>
47
+ <a href="https://pypi.org/project/disco-eb/"><img alt="PyPI" src="https://img.shields.io/pypi/v/disco-eb"></a>
48
+ <a href="https://pypi.org/project/disco-eb/"><img alt="Downloads" src="https://img.shields.io/pypi/dm/disco-eb"></a>
49
+ <a href="https://pypi.org/project/disco-eb/"><img alt="Python" src="https://img.shields.io/pypi/pyversions/disco-eb"></a>
50
+ <a href="https://docs.jax.dev/"><img alt="JAX" src="https://img.shields.io/badge/JAX-compatible-3f51b5"></a>
51
+ <a href="https://arxiv.org/abs/2311.03291"><img alt="arXiv" src="https://img.shields.io/badge/arXiv-2311.03291-b31b1b.svg"></a>
52
+ <a href="https://www.conventionalcommits.org/en/v1.0.0/"><img alt="Conventional Commits" src="https://img.shields.io/badge/Conventional%20Commits-1.0.0-%23FE5196?logo=conventionalcommits&logoColor=white"></a>
53
+ <a href="LICENSE"><img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
54
+ </p>
55
+
56
+ <p align="center">
57
+ <a href="#installation">Installation</a> •
58
+ <a href="#quickstart">Quickstart</a> •
59
+ <a href="https://cosmo-sims.github.io/DISCO-EB/">Documentation</a> •
60
+ <a href="https://arxiv.org/abs/2311.03291">Paper</a> •
61
+ <a href="#citation">Citation</a>
62
+ </p>
63
+
64
+ <!-- start:overview -->
65
+
66
+ DISCO-EB is a **differentiable linear Einstein-Boltzmann solver** for cosmology, and the Einstein-Boltzmann module of the DISCO-DJ framework (DIfferentiable Simulations for COsmology, Done with Jax). It evolves the background, the recombination history and the linear perturbations of every wave mode of a **batch of cosmologies in one GPU launch**.
67
+
68
+ The solve runs as a [modax](https://github.com/lawrenceberry/modax) Rodas5P kernel. RECFAST and the optical depth step alongside the perturbation hierarchy in the same kernel, so the solve needs no tables. You choose which quantities to keep at which times, and the kernel evaluates them as it goes. The results are JAX arrays, differentiable in the cosmological parameters by forward sensitivities. DISCO-EB supports flat or curved ΛCDM, a CPL dark-energy fluid $(w_0, w_a)$, massive neutrinos and reionization.
69
+
70
+ <!-- end:overview -->
71
+
72
+ The [documentation](https://cosmo-sims.github.io/DISCO-EB/) covers the API and how to use it. The [paper](https://arxiv.org/abs/2311.03291) describes the method and its accuracy against CLASS.
73
+
74
+ ## Installation
75
+
76
+ <!-- start:install -->
77
+
78
+ ```bash
79
+ pip install "disco-eb[cuda13]" # or [cuda12], which Python 3.10 needs
80
+ ```
81
+
82
+ You need a GPU to run a solve, but not to import the package. modax brings the CUDA compiler stack with it. The CMB spectra live in a separate package, [DISCO-CMB](https://github.com/lawrenceberry/DISCO-CMB), which builds on this one.
83
+
84
+ <!-- end:install -->
85
+
86
+ ## Quickstart
87
+
88
+ <!-- start:quickstart -->
89
+
90
+ This example keeps the comoving curvature perturbation $\mathcal{R}$ and the CDM density contrast of 64 wave modes at 100 conformal times:
91
+
92
+ ```python
93
+ import jax
94
+
95
+ jax.config.update("jax_enable_x64", True)
96
+
97
+ import numpy as np
98
+ from discoeb import solve_eb
99
+
100
+ h = 0.6732
101
+ cosmology = {
102
+ "Omegam": 0.1424 / h**2, "Omegab": 0.02238 / h**2, "H0": 100 * h,
103
+ "w_DE_0": -1.0, "w_DE_a": 0.0, "cs2_DE": 1.0, "Omegak": 0.0,
104
+ "A_s": 2.1e-9, "n_s": 0.966, "k_p": 0.05,
105
+ "Tcmb": 2.7255, "YHe": 0.2454, "Neff": 3.046, "Nmnu": 0.0, "mnu": 0.0,
106
+ }
107
+ k = np.geomspace(1e-4, 1.0, 64) # 1/Mpc
108
+ tau = np.geomspace(1.0, 2e4, 100) # Mpc, the last times past today
109
+
110
+ solution = solve_eb(cosmology, k, tau, keep=("R", "clxc"))
111
+ R = solution["R"] # (1, 64, 100): cosmologies, wave modes, times
112
+ ```
113
+
114
+ <!-- end:quickstart -->
115
+
116
+ ## Development
117
+
118
+ <!-- start:development -->
119
+
120
+ The repository uses [uv](https://docs.astral.sh/uv/):
121
+
122
+ ```bash
123
+ git clone https://github.com/cosmo-sims/DISCO-EB.git && cd DISCO-EB
124
+ uv sync --extra cuda13 # or --extra cuda12
125
+
126
+ uv run pytest tests -rs # needs a GPU, and -rs shows skipped CLASS comparisons
127
+ uv run ruff format && uv run ruff check --fix
128
+ uv run ty check
129
+ uv run pre-commit install # the checks above, and a commit-message check
130
+ ```
131
+
132
+ The documentation (Sphinx with MyST and autodoc) builds in an environment of
133
+ its own, the package's CPU install plus the docs group:
134
+
135
+ ```bash
136
+ export UV_PROJECT_ENVIRONMENT=.venv-docs
137
+ uv sync --no-default-groups --group docs
138
+ uv run --no-sync sphinx-autobuild docs docs/_build/html # live preview
139
+ uv run --no-sync sphinx-build -W docs docs/_build/html # what CI runs
140
+ ```
141
+
142
+ <!-- end:development -->
143
+
144
+ ## Citation
145
+
146
+ <!-- start:citation -->
147
+
148
+ If you use DISCO-EB in your research, please cite the
149
+ [JCAP article](https://doi.org/10.1088/1475-7516/2024/06/063):
150
+
151
+ ```bibtex
152
+ @article{DISCO-EB:2024,
153
+ author = "Hahn, Oliver and List, Florian and Porqueres, Natalia",
154
+ title = "{DISCO-DJ I: a differentiable Einstein-Boltzmann solver for cosmology}",
155
+ eprint = "2311.03291",
156
+ archivePrefix = "arXiv",
157
+ primaryClass = "astro-ph.CO",
158
+ doi = "10.1088/1475-7516/2024/06/063",
159
+ journal = "JCAP",
160
+ volume = "06",
161
+ pages = "063",
162
+ year = "2024"
163
+ }
164
+ ```
165
+
166
+ <!-- end:citation -->
167
+
168
+ ## License
169
+
170
+ <!-- start:license -->
171
+
172
+ DISCO-EB uses the
173
+ [MIT License](https://github.com/cosmo-sims/DISCO-EB/blob/main/LICENSE). The
174
+ recombination module adapts code from
175
+ [Recfast](https://www.astro.ubc.ca/people/scott/recfast.html), whose UBC
176
+ license its file header reproduces.
177
+
178
+ <!-- end:license -->
disco_eb-1.0.1/README.md ADDED
@@ -0,0 +1,145 @@
1
+ <h1 align="center">DISCO-EB</h1>
2
+
3
+ <!--
4
+ The coverage badge reads a gist that the `tests` workflow rewrites on every
5
+ push to main. The id below is the repository variable COVERAGE_GIST_ID.
6
+ These badges sit above the `overview` snippet marker, so the documentation
7
+ site does not repeat them.
8
+ -->
9
+ <p align="center">
10
+ <a href="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/tests.yml"><img alt="tests" src="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/tests.yml/badge.svg?branch=main"></a>
11
+ <a href="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/tests.yml"><img alt="coverage" src="https://img.shields.io/endpoint?url=https%3A%2F%2Fgist.githubusercontent.com%2Flawrenceberry%2Fb451e14a9620fabb0d49c345f13b5414%2Fraw%2Fdiscoeb-coverage.json"></a>
12
+ <a href="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/lint.yml"><img alt="lint" src="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/lint.yml/badge.svg?branch=main"></a>
13
+ <a href="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/types.yml"><img alt="types" src="https://github.com/cosmo-sims/DISCO-EB/actions/workflows/types.yml/badge.svg?branch=main"></a>
14
+ <a href="https://pypi.org/project/disco-eb/"><img alt="PyPI" src="https://img.shields.io/pypi/v/disco-eb"></a>
15
+ <a href="https://pypi.org/project/disco-eb/"><img alt="Downloads" src="https://img.shields.io/pypi/dm/disco-eb"></a>
16
+ <a href="https://pypi.org/project/disco-eb/"><img alt="Python" src="https://img.shields.io/pypi/pyversions/disco-eb"></a>
17
+ <a href="https://docs.jax.dev/"><img alt="JAX" src="https://img.shields.io/badge/JAX-compatible-3f51b5"></a>
18
+ <a href="https://arxiv.org/abs/2311.03291"><img alt="arXiv" src="https://img.shields.io/badge/arXiv-2311.03291-b31b1b.svg"></a>
19
+ <a href="https://www.conventionalcommits.org/en/v1.0.0/"><img alt="Conventional Commits" src="https://img.shields.io/badge/Conventional%20Commits-1.0.0-%23FE5196?logo=conventionalcommits&logoColor=white"></a>
20
+ <a href="LICENSE"><img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
21
+ </p>
22
+
23
+ <p align="center">
24
+ <a href="#installation">Installation</a> •
25
+ <a href="#quickstart">Quickstart</a> •
26
+ <a href="https://cosmo-sims.github.io/DISCO-EB/">Documentation</a> •
27
+ <a href="https://arxiv.org/abs/2311.03291">Paper</a> •
28
+ <a href="#citation">Citation</a>
29
+ </p>
30
+
31
+ <!-- start:overview -->
32
+
33
+ DISCO-EB is a **differentiable linear Einstein-Boltzmann solver** for cosmology, and the Einstein-Boltzmann module of the DISCO-DJ framework (DIfferentiable Simulations for COsmology, Done with Jax). It evolves the background, the recombination history and the linear perturbations of every wave mode of a **batch of cosmologies in one GPU launch**.
34
+
35
+ The solve runs as a [modax](https://github.com/lawrenceberry/modax) Rodas5P kernel. RECFAST and the optical depth step alongside the perturbation hierarchy in the same kernel, so the solve needs no tables. You choose which quantities to keep at which times, and the kernel evaluates them as it goes. The results are JAX arrays, differentiable in the cosmological parameters by forward sensitivities. DISCO-EB supports flat or curved ΛCDM, a CPL dark-energy fluid $(w_0, w_a)$, massive neutrinos and reionization.
36
+
37
+ <!-- end:overview -->
38
+
39
+ The [documentation](https://cosmo-sims.github.io/DISCO-EB/) covers the API and how to use it. The [paper](https://arxiv.org/abs/2311.03291) describes the method and its accuracy against CLASS.
40
+
41
+ ## Installation
42
+
43
+ <!-- start:install -->
44
+
45
+ ```bash
46
+ pip install "disco-eb[cuda13]" # or [cuda12], which Python 3.10 needs
47
+ ```
48
+
49
+ You need a GPU to run a solve, but not to import the package. modax brings the CUDA compiler stack with it. The CMB spectra live in a separate package, [DISCO-CMB](https://github.com/lawrenceberry/DISCO-CMB), which builds on this one.
50
+
51
+ <!-- end:install -->
52
+
53
+ ## Quickstart
54
+
55
+ <!-- start:quickstart -->
56
+
57
+ This example keeps the comoving curvature perturbation $\mathcal{R}$ and the CDM density contrast of 64 wave modes at 100 conformal times:
58
+
59
+ ```python
60
+ import jax
61
+
62
+ jax.config.update("jax_enable_x64", True)
63
+
64
+ import numpy as np
65
+ from discoeb import solve_eb
66
+
67
+ h = 0.6732
68
+ cosmology = {
69
+ "Omegam": 0.1424 / h**2, "Omegab": 0.02238 / h**2, "H0": 100 * h,
70
+ "w_DE_0": -1.0, "w_DE_a": 0.0, "cs2_DE": 1.0, "Omegak": 0.0,
71
+ "A_s": 2.1e-9, "n_s": 0.966, "k_p": 0.05,
72
+ "Tcmb": 2.7255, "YHe": 0.2454, "Neff": 3.046, "Nmnu": 0.0, "mnu": 0.0,
73
+ }
74
+ k = np.geomspace(1e-4, 1.0, 64) # 1/Mpc
75
+ tau = np.geomspace(1.0, 2e4, 100) # Mpc, the last times past today
76
+
77
+ solution = solve_eb(cosmology, k, tau, keep=("R", "clxc"))
78
+ R = solution["R"] # (1, 64, 100): cosmologies, wave modes, times
79
+ ```
80
+
81
+ <!-- end:quickstart -->
82
+
83
+ ## Development
84
+
85
+ <!-- start:development -->
86
+
87
+ The repository uses [uv](https://docs.astral.sh/uv/):
88
+
89
+ ```bash
90
+ git clone https://github.com/cosmo-sims/DISCO-EB.git && cd DISCO-EB
91
+ uv sync --extra cuda13 # or --extra cuda12
92
+
93
+ uv run pytest tests -rs # needs a GPU, and -rs shows skipped CLASS comparisons
94
+ uv run ruff format && uv run ruff check --fix
95
+ uv run ty check
96
+ uv run pre-commit install # the checks above, and a commit-message check
97
+ ```
98
+
99
+ The documentation (Sphinx with MyST and autodoc) builds in an environment of
100
+ its own, the package's CPU install plus the docs group:
101
+
102
+ ```bash
103
+ export UV_PROJECT_ENVIRONMENT=.venv-docs
104
+ uv sync --no-default-groups --group docs
105
+ uv run --no-sync sphinx-autobuild docs docs/_build/html # live preview
106
+ uv run --no-sync sphinx-build -W docs docs/_build/html # what CI runs
107
+ ```
108
+
109
+ <!-- end:development -->
110
+
111
+ ## Citation
112
+
113
+ <!-- start:citation -->
114
+
115
+ If you use DISCO-EB in your research, please cite the
116
+ [JCAP article](https://doi.org/10.1088/1475-7516/2024/06/063):
117
+
118
+ ```bibtex
119
+ @article{DISCO-EB:2024,
120
+ author = "Hahn, Oliver and List, Florian and Porqueres, Natalia",
121
+ title = "{DISCO-DJ I: a differentiable Einstein-Boltzmann solver for cosmology}",
122
+ eprint = "2311.03291",
123
+ archivePrefix = "arXiv",
124
+ primaryClass = "astro-ph.CO",
125
+ doi = "10.1088/1475-7516/2024/06/063",
126
+ journal = "JCAP",
127
+ volume = "06",
128
+ pages = "063",
129
+ year = "2024"
130
+ }
131
+ ```
132
+
133
+ <!-- end:citation -->
134
+
135
+ ## License
136
+
137
+ <!-- start:license -->
138
+
139
+ DISCO-EB uses the
140
+ [MIT License](https://github.com/cosmo-sims/DISCO-EB/blob/main/LICENSE). The
141
+ recombination module adapts code from
142
+ [Recfast](https://www.astro.ubc.ca/people/scott/recfast.html), whose UBC
143
+ license its file header reproduces.
144
+
145
+ <!-- end:license -->
disco_eb-1.0.1/pyproject.toml ADDED
@@ -0,0 +1,174 @@
1
+ [project]
2
+ name = "disco-eb"
3
+ version = "1.0.1"
4
+ description = "A differentiable linear Einstein-Boltzmann solver for cosmology, on the GPU"
5
+ readme = "README.md"
6
+ requires-python = ">=3.10"
7
+ authors = [
8
+ { name = "Oliver Hahn", email = "oliver.hahn@univie.ac.at" },
9
+ { name = "Florian List" },
10
+ { name = "Natalia Porqueres" },
11
+ { name = "Lawrence Berry" },
12
+ ]
13
+ # PEP 639: an SPDX expression and the files themselves. recombination.py also
14
+ # carries the Recfast license, which its header reproduces.
15
+ license = "MIT"
16
+ license-files = ["LICENSE"]
17
+ classifiers = [
18
+ "Development Status :: 4 - Beta",
19
+ "Environment :: GPU :: NVIDIA CUDA",
20
+ "Intended Audience :: Science/Research",
21
+ "Operating System :: POSIX :: Linux",
22
+ "Programming Language :: Python :: 3.10",
23
+ "Programming Language :: Python :: 3.11",
24
+ "Programming Language :: Python :: 3.12",
25
+ "Programming Language :: Python :: 3.13",
26
+ "Topic :: Scientific/Engineering :: Astronomy",
27
+ "Topic :: Scientific/Engineering :: Physics",
28
+ ]
29
+ # Only what `src/discoeb` imports. The CUDA build of JAX is an extra, as in
30
+ # modax: the CPU install is enough to import the package, build the docs and
31
+ # type check it.
32
+ dependencies = [
33
+ "jax>=0.6.2",
34
+ "numpy>=2.2.6",
35
+ # The hierarchy of blocks the joint solve is built on (Block, Hierarchy),
36
+ # its forward sensitivities and quantity_hook, whose Enzyme tangents
37
+ # solve_eb's derivatives use. 0.2.0 for the Rodas5P method object that
38
+ # carries the LU precision and the PI gains. modax brings numba-cuda-mlir.
39
+ "modax-solvers>=0.2.2",
40
+ # 0.3.0 for the O3 prepasses the kernels need and Blackwell support.
41
+ "numba-enzyme-cuda>=0.3.0",
42
+ ]
43
+
44
+ [project.urls]
45
+ Documentation = "https://cosmo-sims.github.io/DISCO-EB/"
46
+ Source = "https://github.com/cosmo-sims/DISCO-EB"
47
+ Issues = "https://github.com/cosmo-sims/DISCO-EB/issues"
48
+
49
+ [project.optional-dependencies]
50
+ cuda12 = ["jax[cuda12]>=0.6.2"]
51
+ # JAX's CUDA 13 build starts after it dropped Python 3.10, so on 3.10 this
52
+ # extra installs nothing and only `cuda12` gives a GPU-enabled JAX.
53
+ cuda13 = ["jax[cuda13]>=0.9.1; python_version >= '3.11'"]
54
+
55
+ [dependency-groups]
56
+ # The docs build (Sphinx, MyST and autodoc) imports the package, so its
57
+ # environment is the project with its CPU dependencies plus this group:
58
+ # `uv sync --no-default-groups --group docs`. No GPU is needed.
59
+ docs = [
60
+ "sphinx>=8.1",
61
+ "myst-parser>=4.0",
62
+ "sphinx-book-theme>=1.1",
63
+ "sphinx-design>=0.6",
64
+ "sphinx-autobuild>=2024.10",
65
+ ]
66
+ dev = [
67
+ # The CLASS references the accuracy tests compare against. The tests skip
68
+ # without it, so a missing classy reads as a pass in `-q` output.
69
+ "classy>=3.3.4.0",
70
+ "ipykernel>=6.31.0",
71
+ "matplotlib>=3.9.4",
72
+ "pre-commit>=4.4.0",
73
+ "pytest>=9.0.2",
74
+ "pytest-benchmark>=5.2.3",
75
+ "pytest-cov>=7.0.0",
76
+ "pytest-regressions>=2.10.0",
77
+ "ruff>=0.15.5",
78
+ # The same type checker CI and the pre-commit hook run.
79
+ "ty>=0.0.83",
80
+ ]
81
+
82
+ [build-system]
83
+ requires = ["setuptools>=77"]
84
+ build-backend = "setuptools.build_meta"
85
+
86
+ [tool.setuptools.packages.find]
87
+ where = ["src"]
88
+
89
+ # The GPU stack must match modax's lock: these kernels are compiled by
90
+ # numba-cuda-mlir and linked by nvjitlink, and 0.5.2 with CUDA 13.4 emits a
91
+ # `.smax` that ptxas rejects when building the 50-variable Enzyme kernel. On
92
+ # Python 3.10 modax runs the whole stack on CUDA 12 instead, so the 13.x toolkit
93
+ # pin applies only above it.
94
+ [tool.uv]
95
+ constraint-dependencies = [
96
+ "numba-cuda-mlir==0.5.1",
97
+ "cuda-toolkit==13.2.1; python_version >= '3.11'",
98
+ ]
99
+ # numba-enzyme caps numba below 0.62 and llvmlite below 0.45 for its CPU
100
+ # pipeline, which hands llvmlite's IR to its vendored LLVM 15 tools. This
101
+ # project only uses its CUDA backend, where numba-cuda-mlir emits the IR, so
102
+ # the cap is lifted here rather than letting it drag numba and numpy back to
103
+ # the LLVM 15 releases (numba 0.61, numpy 2.2).
104
+ override-dependencies = ["numba>=0.61", "llvmlite>=0.44"]
105
+
106
+ [tool.pytest.ini_options]
107
+ testpaths = ["tests"]
108
+
109
+ [tool.semantic_release]
110
+ # python-semantic-release (the `publish` workflow) reads the Conventional
111
+ # Commits since the last `v*` tag, picks the next version, writes it here,
112
+ # updates CHANGELOG.md, commits, tags and builds.
113
+ version_toml = ["pyproject.toml:project.version"]
114
+ tag_format = "v{version}"
115
+ commit_parser = "conventional"
116
+ commit_message = "chore(release): v{version}"
117
+ allow_zero_version = true
118
+ major_on_zero = false
119
+ # uv.lock records the project's own version, so the build refreshes that one
120
+ # entry and stages it with the release commit. `set -e`: a failed step must
121
+ # stop the release.
122
+ build_command = """
123
+ set -e
124
+ python -m pip install "uv==0.13.7"
125
+ uv lock --upgrade-package "$PACKAGE_NAME"
126
+ git add uv.lock
127
+ uv build
128
+ """
129
+
130
+ [tool.semantic_release.branches.main]
131
+ match = "main"
132
+ prerelease = false
133
+
134
+ [tool.semantic_release.changelog]
135
+ mode = "update"
136
+ insertion_flag = "<!-- version list -->"
137
+ exclude_commit_patterns = ['^chore\(release\):']
138
+
139
+ [tool.semantic_release.changelog.default_templates]
140
+ changelog_file = "CHANGELOG.md"
141
+
142
+ [tool.ruff]
143
+ # The notebooks still call the removed diffrax `evolve_perturbations*` API and
144
+ # await porting to solve_eb; they are not linted until then.
145
+ extend-exclude = ["notebooks"]
146
+
147
+ [tool.ruff.lint]
148
+ # On top of the default `select`, enable import sorting.
149
+ extend-select = ["I"]
150
+
151
+ [tool.ruff.lint.per-file-ignores]
152
+ # The examples set environment variables and enable x64 before importing JAX
153
+ # and DISCO-EB, which must happen first.
154
+ "examples/*" = ["E402"]
155
+
156
+ [tool.coverage.run]
157
+ source = ["discoeb"]
158
+ branch = true
159
+
160
+ [tool.coverage.report]
161
+ exclude_also = [
162
+ "if TYPE_CHECKING:",
163
+ "raise NotImplementedError",
164
+ "if __name__ == .__main__.:",
165
+ ]
166
+
167
+ [tool.ty.src]
168
+ # The shipped package is what the type check gates.
169
+ include = ["src/discoeb"]
170
+
171
+ [tool.ty.rules]
172
+ # numba-cuda-mlir assembles `cuda.local` and friends onto the `cuda` namespace
173
+ # at import time rather than as real submodules.
174
+ possibly-missing-submodule = "ignore"
disco_eb-1.0.1/setup.cfg ADDED
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+