Perturbation theory calculations for LCDM and Modified Gravity theories using "fk"-Kernels implemented in Python with JAX.
Based on the C code at https://github.com/alejandroaviles/fkpt, which is based on the paper Rodriguez-Meza, M. A. et al, "fkPT: Constraining scale-dependent modified gravity with the full-shape galaxy power spectrum", JCAP03(2024)049.
See KFUNCTIONS.md for details on what quantities this code calculates. See the included notebook for examples of how to perform the calculations. See BENCHMARKS.md for timing benchmarks on different platforms.
Install from pip:
pip install fkptjax- Python 3.10+
- JAX 0.4.0+
- NumPy 1.24.0+
- SciPy 1.10.0+
For development, install with dev dependencies:
git clone https://github.com/cosmodesi/fkptjax.git
cd fkptjax
pip install -e ".[dev]"Run tests:
pytestRun benchmarks (default: all calculators, 100 runs):
python tests/test.pyRun benchmarks with options:
# Run only NumPy calculator
python tests/test.py -c numpy
# Run with custom data file and 50 runs
python tests/test.py -d mydata.npz -n 50
# Run both NumPy and JAX CPU calculators
python tests/test.py -c numpy jax-cpu
# Show all options
python tests/test.py --helpRun tests with coverage:
pytest --cov=fkptjaxType checking:
mypy src/fkptjaxMIT License - see LICENSE file for details.
- David Kirkby [email protected]
- Matthew Dowicz [email protected]