Metadata-Version: 2.4
Name: fftlog-lss
Version: 0.1.1
Summary: JAX-accelerated FFTLog implementation for Large Scale Structure cosmology
Author: Pierre Zhang
Author-email: Pierre Zhang <pierre.zhang@example.com>
Maintainer-email: Pierre Zhang <pierre.zhang@example.com>
License-Expression: MIT
Project-URL: Homepage, https://github.com/pierrexyz/fftlog
Project-URL: Repository, https://github.com/pierrexyz/fftlog
Project-URL: Documentation, https://github.com/pierrexyz/fftlog#readme
Project-URL: Bug Tracker, https://github.com/pierrexyz/fftlog/issues
Keywords: fft,logarithmic,transform,scientific,numerical,jax,gpu,acceleration
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Environment :: Console
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy>=1.20.0
Requires-Dist: scipy>=1.7.0
Provides-Extra: jax
Requires-Dist: jax>=0.4.0; extra == "jax"
Requires-Dist: jaxlib>=0.4.0; extra == "jax"
Requires-Dist: interpax>=0.3.0; extra == "jax"
Provides-Extra: dev
Requires-Dist: pytest>=6.0; extra == "dev"
Requires-Dist: pytest-cov>=2.0; extra == "dev"
Requires-Dist: black>=22.0; extra == "dev"
Requires-Dist: flake8>=4.0; extra == "dev"
Requires-Dist: mypy>=0.950; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx>=4.0; extra == "docs"
Requires-Dist: sphinx-rtd-theme>=1.0; extra == "docs"
Requires-Dist: nbsphinx>=0.8; extra == "docs"
Dynamic: author
Dynamic: license-file

# FFTLog

A Python implementation of FFTLog for fast logarithmic FFT transforms, developed for [PyBird](https://github.com/pierrexyz/pybird).

## Installation

### Basic installation (recommended)
```bash
pip install fftlog-lss
```

### With JAX acceleration
```bash
pip install fftlog-lss[jax]
```

### From source
```bash
git clone https://github.com/pierrexyz/fftlog.git
cd fftlog
pip install --editable .
```

### From source with JAX
```bash
git clone https://github.com/pierrexyz/fftlog.git
cd fftlog
pip install --editable .[jax]
```

## Dependencies

### Required
- Python >= 3.8
- numpy >= 1.20.0
- scipy >= 1.7.0

### Optional (for JAX acceleration)
- jax >= 0.4.0
- jaxlib >= 0.4.0
- interpax >= 0.3.0

## Quickstart

### Basic Usage (without JAX)

```python
import numpy as np
from scipy.stats import lognorm
from fftlog import FFTLog

# Create test function (power spectrum)
k = np.logspace(-4, 0, 200)
pk = lognorm.pdf(k, 2.1)

# Initialize FFTLog
fft = FFTLog(
    Nmax=512,       # Number of points
    xmin=1e-5,      # Minimum k value
    xmax=1e3,       # Maximum k value
    bias=-0.1,      # Bias parameter
    complex=False,  # Use real FFT
    window=0.1      # Anti-aliasing window
)

# FFTLog decomposition and reconstruction
pk_reconstructed = fft.rec(k, pk, k)

# Spherical Bessel Transform
s = np.arange(1., 1e3, 5.)
xi = fft.sbt(k, pk, s)  # Fast O(N log N) transform
```

### With JAX Acceleration

```python
import numpy as np
from scipy.stats import lognorm
from fftlog import FFTLog
from fftlog.config import set_jax_enabled
from jax import jit
import jax.numpy as jnp

# Enable JAX mode
set_jax_enabled(True)

# Create test function
k = np.logspace(-4, 0, 200)
pk = lognorm.pdf(k, 2.1)

# Initialize FFTLog (same as above)
fft = FFTLog(
    Nmax=512, xmin=1e-5, xmax=1e3, 
    bias=-0.1, complex=False, window=0.1
)

# Convert to JAX arrays and JIT compile
k_jax, pk_jax = jnp.array(k), jnp.array(pk)
get_coef_jit = jit(fft.Coef)

# Now much faster for repeated calls
coefficients = get_coef_jit(k_jax, pk_jax)
```

## Features

- Fast logarithmic FFT transforms
- Support for both real and complex transforms
- Spherical Bessel transforms
- Anti-aliasing windows
- Optional JAX acceleration for GPU/TPU support

## Documentation

For more detailed examples and documentation, see the notebooks in the `notebooks/` directory.

## License

This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

## Contributing

Contributions are welcome! Please feel free to submit a Pull Request.



