Metadata-Version: 2.4
Name: fkptjax
Version: 0.2.1
Summary: Perturbation theory calculations for LCDM and Modified Gravity theories using fk-Kernels implemented in Python with JAX
Author-email: David Kirkby <dkirkby@uci.edu>, Matthew Dowicz <mdowicz@uci.edu>
License: MIT
Project-URL: Homepage, https://github.com/cosmodesi/fkptjax
Project-URL: Repository, https://github.com/cosmodesi/fkptjax
Keywords: cosmology,perturbation-theory,modified-gravity,jax
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Physics
Classifier: Topic :: Scientific/Engineering :: Astronomy
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.4.0
Requires-Dist: jaxlib>=0.4.0
Requires-Dist: numpy>=1.24.0
Requires-Dist: scipy>=1.10.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: pytest-cov>=4.0; extra == "dev"
Requires-Dist: mypy>=1.0; extra == "dev"
Dynamic: license-file

# fkptjax

[![PyPI version](https://badge.fury.io/py/fkptjax.svg)](https://badge.fury.io/py/fkptjax)
[![Python](https://img.shields.io/pypi/pyversions/fkptjax.svg)](https://pypi.org/project/fkptjax/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![Tests](https://github.com/cosmodesi/fkptjax/actions/workflows/tests.yml/badge.svg)](https://github.com/cosmodesi/fkptjax/actions/workflows/tests.yml)

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](https://doi.org/10.1088/1475-7516/2024/03/049).

See [KFUNCTIONS.md](KFUNCTIONS.md) for details on what quantities this code calculates. See the [included notebook](examples/kfunctions.ipynb) for examples of how to perform the calculations. See [BENCHMARKS.md](BENCHMARKS.md) for timing benchmarks on different platforms.

## Installation

Install from pip:

```bash
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:

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

Run tests:

```bash
pytest
```

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

```bash
python tests/test.py
```

Run benchmarks with options:

```bash
# 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:

```bash
pytest --cov=fkptjax
```

Type checking:

```bash
mypy src/fkptjax
```

## License

MIT License - see LICENSE file for details.

## Authors

- David Kirkby <dkirkby@uci.edu>
- Matthew Dowicz <mdowicz@uci.edu>
