Metadata-Version: 2.4
Name: mlx-airy
Version: 0.1.1
Summary: GPU-accelerated Airy functions for Apple Silicon using MLX
Author-email: Sheng-Kai Huang <akai@fawstudio.com>
License-Expression: MIT
Project-URL: Homepage, https://github.com/akaiHuang/mlx-airy
Keywords: airy,special-functions,gpu,mlx,apple-silicon,quantum,optics
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: MacOS
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 :: Mathematics
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy
Requires-Dist: mlx>=0.20.0
Provides-Extra: benchmark
Requires-Dist: scipy; extra == "benchmark"
Provides-Extra: test
Requires-Dist: scipy; extra == "test"
Requires-Dist: pytest; extra == "test"
Dynamic: license-file

# mlx-airy

GPU-accelerated [Airy functions](https://dlmf.nist.gov/9) for Apple Silicon using [MLX](https://github.com/ml-explore/mlx).

Neither MLX nor JAX currently provide Airy functions
(see [JAX issue #25244](https://github.com/jax-ml/jax/issues/25244)).
This package fills that gap for the MLX ecosystem.

## What are Airy functions?

Airy functions Ai(x) and Bi(x) are the two linearly independent solutions to
the **Airy differential equation**:

```
y'' - x y = 0
```

They appear throughout physics and engineering:
- **Quantum mechanics** -- WKB connection formulas, tunneling through linear potentials
- **Optics** -- diffraction near caustics, rainbow scattering
- **Fluid dynamics** -- boundary layers, Stokes phenomena
- **Astrophysics** -- gravitational lensing near fold caustics

## Installation

```bash
pip install -e ".[test,benchmark]"
```

Requires Python >= 3.10 and Apple Silicon (MLX).

## Usage

```python
import mlx.core as mx
from mlx_airy import airy

x = mx.linspace(-15.0, 15.0, 10000)
ai, aip, bi, bip = airy(x)  # Ai(x), Ai'(x), Bi(x), Bi'(x)
```

All four functions are returned in a single call, matching the
`scipy.special.airy` convention.

## Implementation

| Region   | Method               | Reference          |
|----------|----------------------|--------------------|
| \|x\| < 5  | Taylor series (50 terms, Horner) | DLMF 9.4  |
| x >= 5   | Asymptotic expansion | DLMF 9.7.5--9.7.8 |
| x <= -5  | Oscillatory asymptotic | DLMF 9.7.9--9.7.12 |

Computation is float32 throughout, running entirely on the Apple GPU.

## Tests

```bash
pytest tests/ -v
```

## Benchmark

```bash
python -m mlx_airy.benchmark
```

Compares accuracy and throughput against `scipy.special.airy`.

Measured on Apple M1 Max (MLX 0.31.1, SciPy 1.16.2, float32),
`x in [-15, 15]`, all four outputs `(Ai, Ai', Bi, Bi')`:

| N         | scipy (ms) | mlx (ms) | Speedup |
|----------:|-----------:|---------:|--------:|
| 1,000     |      3.30  |   32.07  |   0.1x  |
| 10,000    |     15.76  |   31.65  |   0.5x  |
| 100,000   |    147.69  |   35.55  |   4.2x  |
| 1,000,000 |   1393.08  |  209.09  |   6.7x  |

mlx-airy wins above ~30-50k points, where the Metal dispatch cost is
amortised; below that, scipy's vectorised C path is faster. Worst-case
float32 relative error is ~3e-4. Full numbers in
[`benchmark_results.md`](benchmark_results.md).

## Author

**Sheng-Kai Huang** (akai@fawstudio.com)

## License

MIT
