fkptjax

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.

Installation

Install from pip:

pip install fkptjax

Requirements

• Python 3.10+
• JAX 0.4.0+
• NumPy 1.24.0+
• SciPy 1.10.0+

Development

For development, install with dev dependencies:

git clone https://github.com/cosmodesi/fkptjax.git
cd fkptjax
pip install -e ".[dev]"

Run tests:

pytest

Run benchmarks (default: all calculators, 100 runs):

python tests/test.py

Run 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 --help

Run tests with coverage:

pytest --cov=fkptjax

Type checking:

mypy src/fkptjax

License

MIT License - see LICENSE file for details.

Authors

• David Kirkby dkirkby@uci.edu
• Matthew Dowicz mdowicz@uci.edu