disco-eb 1.0.1__py3-none-any.whl

disco_eb-1.0.1.dist-info/METADATA

@@ -0,0 +1,178 @@
+Metadata-Version: 2.4
+Name: disco-eb
+Version: 1.0.1
+Summary: A differentiable linear Einstein-Boltzmann solver for cosmology, on the GPU
+Author: Florian List, Natalia Porqueres, Lawrence Berry
+Author-email: Oliver Hahn <oliver.hahn@univie.ac.at>
+License-Expression: MIT
+Project-URL: Documentation, https://cosmo-sims.github.io/DISCO-EB/
+Project-URL: Source, https://github.com/cosmo-sims/DISCO-EB
+Project-URL: Issues, https://github.com/cosmo-sims/DISCO-EB/issues
+Classifier: Development Status :: 4 - Beta
+Classifier: Environment :: GPU :: NVIDIA CUDA
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Astronomy
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.6.2
+Requires-Dist: numpy>=2.2.6
+Requires-Dist: modax-solvers>=0.2.2
+Requires-Dist: numba-enzyme-cuda>=0.3.0
+Provides-Extra: cuda12
+Requires-Dist: jax[cuda12]>=0.6.2; extra == "cuda12"
+Provides-Extra: cuda13
+Requires-Dist: jax[cuda13]>=0.9.1; python_version >= "3.11" and extra == "cuda13"
+Dynamic: license-file
+
+<h1 align="center">DISCO-EB</h1>
+
+<!--
+The coverage badge reads a gist that the `tests` workflow rewrites on every
+push to main. The id below is the repository variable COVERAGE_GIST_ID.
+These badges sit above the `overview` snippet marker, so the documentation
+site does not repeat them.
+-->
+<p align="center">
+  <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>
+  <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>
+  <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>
+  <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>
+  <a href="https://pypi.org/project/disco-eb/"><img alt="PyPI" src="https://img.shields.io/pypi/v/disco-eb"></a>
+  <a href="https://pypi.org/project/disco-eb/"><img alt="Downloads" src="https://img.shields.io/pypi/dm/disco-eb"></a>
+  <a href="https://pypi.org/project/disco-eb/"><img alt="Python" src="https://img.shields.io/pypi/pyversions/disco-eb"></a>
+  <a href="https://docs.jax.dev/"><img alt="JAX" src="https://img.shields.io/badge/JAX-compatible-3f51b5"></a>
+  <a href="https://arxiv.org/abs/2311.03291"><img alt="arXiv" src="https://img.shields.io/badge/arXiv-2311.03291-b31b1b.svg"></a>
+  <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>
+  <a href="LICENSE"><img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
+</p>
+
+<p align="center">
+  <a href="#installation">Installation</a> •
+  <a href="#quickstart">Quickstart</a> •
+  <a href="https://cosmo-sims.github.io/DISCO-EB/">Documentation</a> •
+  <a href="https://arxiv.org/abs/2311.03291">Paper</a> •
+  <a href="#citation">Citation</a>
+</p>
+
+<!-- start:overview -->
+
+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**.
+
+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.
+
+<!-- end:overview -->
+
+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.
+
+## Installation
+
+<!-- start:install -->
+
+```bash
+pip install "disco-eb[cuda13]"   # or [cuda12], which Python 3.10 needs
+```
+
+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.
+
+<!-- end:install -->
+
+## Quickstart
+
+<!-- start:quickstart -->
+
+This example keeps the comoving curvature perturbation $\mathcal{R}$ and the CDM density contrast of 64 wave modes at 100 conformal times:
+
+```python
+import jax
+
+jax.config.update("jax_enable_x64", True)
+
+import numpy as np
+from discoeb import solve_eb
+
+h = 0.6732
+cosmology = {
+    "Omegam": 0.1424 / h**2, "Omegab": 0.02238 / h**2, "H0": 100 * h,
+    "w_DE_0": -1.0, "w_DE_a": 0.0, "cs2_DE": 1.0, "Omegak": 0.0,
+    "A_s": 2.1e-9, "n_s": 0.966, "k_p": 0.05,
+    "Tcmb": 2.7255, "YHe": 0.2454, "Neff": 3.046, "Nmnu": 0.0, "mnu": 0.0,
+}
+k = np.geomspace(1e-4, 1.0, 64)      # 1/Mpc
+tau = np.geomspace(1.0, 2e4, 100)    # Mpc, the last times past today
+
+solution = solve_eb(cosmology, k, tau, keep=("R", "clxc"))
+R = solution["R"]  # (1, 64, 100): cosmologies, wave modes, times
+```
+
+<!-- end:quickstart -->
+
+## Development
+
+<!-- start:development -->
+
+The repository uses [uv](https://docs.astral.sh/uv/):
+
+```bash
+git clone https://github.com/cosmo-sims/DISCO-EB.git && cd DISCO-EB
+uv sync --extra cuda13  # or --extra cuda12
+
+uv run pytest tests -rs   # needs a GPU, and -rs shows skipped CLASS comparisons
+uv run ruff format && uv run ruff check --fix
+uv run ty check
+uv run pre-commit install   # the checks above, and a commit-message check
+```
+
+The documentation (Sphinx with MyST and autodoc) builds in an environment of
+its own, the package's CPU install plus the docs group:
+
+```bash
+export UV_PROJECT_ENVIRONMENT=.venv-docs
+uv sync --no-default-groups --group docs
+uv run --no-sync sphinx-autobuild docs docs/_build/html   # live preview
+uv run --no-sync sphinx-build -W docs docs/_build/html    # what CI runs
+```
+
+<!-- end:development -->
+
+## Citation
+
+<!-- start:citation -->
+
+If you use DISCO-EB in your research, please cite the
+[JCAP article](https://doi.org/10.1088/1475-7516/2024/06/063):
+
+```bibtex
+@article{DISCO-EB:2024,
+    author = "Hahn, Oliver and List, Florian and Porqueres, Natalia",
+    title = "{DISCO-DJ I: a differentiable Einstein-Boltzmann solver for cosmology}",
+    eprint = "2311.03291",
+    archivePrefix = "arXiv",
+    primaryClass = "astro-ph.CO",
+    doi = "10.1088/1475-7516/2024/06/063",
+    journal = "JCAP",
+    volume = "06",
+    pages = "063",
+    year = "2024"
+}
+```
+
+<!-- end:citation -->
+
+## License
+
+<!-- start:license -->
+
+DISCO-EB uses the
+[MIT License](https://github.com/cosmo-sims/DISCO-EB/blob/main/LICENSE). The
+recombination module adapts code from
+[Recfast](https://www.astro.ubc.ca/people/scott/recfast.html), whose UBC
+license its file header reproduces.
+
+<!-- end:license -->

disco_eb-1.0.1.dist-info/RECORD

@@ -0,0 +1,12 @@
+disco_eb-1.0.1.dist-info/licenses/LICENSE,sha256=mJVm-1igGT7a9zg2Z9minrQczT_8jkQDjNp3W3K6LwI,1125
+discoeb/__init__.py,sha256=oPN2uLjo1r2wTNDo-NP-aKl3bgCE_vP_JBD4-98eZhU,3524
+discoeb/component.py,sha256=_nBLK6gVwxmSprLdeMymp7iZSXS-Vn-BXthTxmdzEHo,53723
+discoeb/eb_model.py,sha256=5FSMsvJ45HaALb8cAIQNCPzL9YI0_r9WC0egO3QjFBw,36638
+discoeb/perturbations.py,sha256=jN0K_5rX51oSqXlZTAg-QGGyEPGNEMcQ7xESJ98pD54,60148
+discoeb/recombination.py,sha256=MP_vP9SCdN82wVn_LuWPd7xiRPv8DriZYi5YW-qlTzY,27234
+discoeb/spline_interpolation.py,sha256=lv2NST_8COkXY3-1Lu4RLTZdFxSqjvnygfwPytA6WuI,8927
+discoeb/util.py,sha256=7KKf8Vswt3kYwQrMn753SJhtj37joAvEjzQ6ws7xHso,1508
+disco_eb-1.0.1.dist-info/METADATA,sha256=Fn-r0P8O6_cHVVcEcJKxVMn1LjhGw_kASkRczBetnVA,7827
+disco_eb-1.0.1.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+disco_eb-1.0.1.dist-info/top_level.txt,sha256=4O3aR-FoF7U-FCR143Ev0aHmnBQk--AD0DjSyOxL38g,8
+disco_eb-1.0.1.dist-info/RECORD,,

disco_eb-1.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

disco_eb-1.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023-2026 Oliver Hahn, Florian List, Natalia Porqueres and Lawrence Berry
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

disco_eb-1.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+discoeb

discoeb/__init__.py

@@ -0,0 +1,87 @@
+"""DISCO-EB, and its interface: what the CMB code -- DISCO-CMB -- may use.
+
+DISCO-EB solves the background, the thermal history (RECFAST with
+reionization) and the linear perturbation hierarchy of each wave mode on the
+GPU: :mod:`discoeb.component` (the component model), :mod:`discoeb.eb_model`
+(the species and their equations), :mod:`discoeb.recombination` (the
+thermal history) and :mod:`discoeb.perturbations` (layouts, batching, the modax launches). The CMB
+code, the separate package ``discocmb`` (DISCO-CMB), turns its output into
+angular power spectra -- line-of-sight sources, grids, transfer functions,
+projection and lensing. It imports from the package itself alone (and the generic
+numerics of :mod:`discoeb.spline_interpolation`); DISCO-EB does not depend
+on it.
+
+**The contract.** A caller asks for *named quantities at chosen conformal
+times* of a batch of cosmologies and wave modes, and gets them back
+evaluated inside the kernel launch from the solver's dense output -- never the
+state history:
+
+* :func:`solve_eb` ``(cosmologies, k, tau, keep=names, extra=components)``
+  returns an :class:`EBSolution` whose ``values`` are ``(n_cosmologies,
+  len(keep), n_k, n_tau)``, with each cosmology's conformal age ``tau0``;
+  past its own ``tau0`` a cosmology's trajectories hold their state today.
+* A name is a state element (:func:`state_names`; the thermal block's
+  ``lna``, ``thermo_0..2`` and the optical depth ``kappa``) or any quantity
+  the model provides (``adotoa``, ``opacity``, ``sigma``, ``dgrho``,
+  ``clxg``, ``polter``, ``R`` ...). A caller extends the vocabulary with
+  ``extra``: stateless :class:`Component` subclasses whose providers read the
+  model's quantities by name and add their own, compiled into the same
+  kernel -- this is how the CMB's line-of-sight sources are defined, so the
+  equations of both parts live in one component model and nothing of the
+  CMB enters the hierarchy's right-hand side.
+* Error weights are given by state name (``{"theta_2": f(k)}``), so the
+  caller needs neither the state's index order nor the launch order.
+* Differentiation: :data:`DIFFERENTIABLE_PARAMETERS` given as JAX values, or
+  the nine :data:`DIRECTED_NAMES` inputs (:func:`directed_inputs`) passed as
+  ``inputs``, make ``values`` differentiable by forward sensitivities formed
+  in the kernel (``jax.jvp``, ``jax.grad``), at fixed ``tau0`` and
+  reionization redshift.
+
+The background for building grids comes from :func:`background_history`,
+:func:`density_coefficients` and :func:`conformal_ages`; cosmologies are
+parameter dictionaries or the objects :func:`as_cosmology` returns.
+"""
+
+from .component import Component
+from .perturbations import (
+    DIFFERENTIABLE_PARAMETERS,
+    DIRECTED_NAMES,
+    PERTURB_ATOL,
+    PERTURB_FIRST_STEP,
+    PERTURB_MAX_STEPS,
+    PERTURB_RTOL,
+    TAU_START,
+    EBSolution,
+    PerturbationLayout,
+    as_cosmology,
+    background_history,
+    conformal_ages,
+    density_coefficients,
+    directed_inputs,
+    joint_model,
+    solve_eb,
+    state_names,
+)
+from .recombination import THERMO_ERROR_WEIGHT
+
+__all__ = [
+    "Component",
+    "DIFFERENTIABLE_PARAMETERS",
+    "DIRECTED_NAMES",
+    "EBSolution",
+    "PERTURB_ATOL",
+    "PERTURB_FIRST_STEP",
+    "PERTURB_MAX_STEPS",
+    "PERTURB_RTOL",
+    "PerturbationLayout",
+    "TAU_START",
+    "THERMO_ERROR_WEIGHT",
+    "as_cosmology",
+    "background_history",
+    "conformal_ages",
+    "density_coefficients",
+    "directed_inputs",
+    "joint_model",
+    "solve_eb",
+    "state_names",
+]