fkptjax 0.1.0

Perturbation theory calculations for LCDM and Modified Gravity theories using fk-Kernels implemented in Python with JAX

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.

Installation

Install from source:

pip install .

For development, install with dev dependencies:

pip install -e ".[dev]"

Requirements

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

Usage

import fkptjax

# Example usage will be added as the package develops

See the examples/ directory for more detailed usage examples.

Development

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 [e-mail]
• Matthew Dowicz [e-mail]