Metadata-Version: 2.5
Name: spexial
Version: 0.1.0
Summary: scipy.special in JAX.
Project-URL: Bug Tracker, https://github.com/JAXtronomy/spexial/issues
Project-URL: Changelog, https://github.com/JAXtronomy/spexial/releases
Project-URL: Documentation, https://jaxtronomy.github.io/spexial
Project-URL: Homepage, https://github.com/JAXtronomy/spexial
Author: Spexial Maintainers
Author-email: Nathaniel Starkman <nstarman@users.noreply.github.com>
License-Expression: MIT
License-File: LICENSE
Keywords: jax,scipy,special-functions
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Typing :: Typed
Requires-Python: >=3.11
Requires-Dist: jax>=0.7.2
Requires-Dist: jaxlib>=0.7.2
Requires-Dist: jaxtyping>=0.3.9
Provides-Extra: cpu
Requires-Dist: jax[cpu]>=0.7.2; extra == 'cpu'
Provides-Extra: cuda12
Requires-Dist: jax[cuda12]>=0.7.2; extra == 'cuda12'
Provides-Extra: cuda12-local
Requires-Dist: jax[cuda12-local]>=0.7.2; extra == 'cuda12-local'
Description-Content-Type: text/markdown

<h1 align='center'> spexial </h1>
<h3 align="center"><code>scipy.special</code> in JAX</h3>

<p align="center">
<a href="https://pypi.org/project/spexial/"><img alt="PyPI version" src="https://img.shields.io/pypi/v/spexial"></a>
<a href="https://pypi.org/project/spexial/"><img alt="PyPI platforms" src="https://img.shields.io/pypi/pyversions/spexial"></a>
<a href="https://github.com/JAXtronomy/spexial/actions"><img alt="Actions Status" src="https://github.com/JAXtronomy/spexial/workflows/CI/badge.svg"></a>
<a href="https://codecov.io/gh/JAXtronomy/spexial"><img alt="codecov" src="https://codecov.io/gh/JAXtronomy/spexial/graph/badge.svg"></a>
<a href="https://jaxtronomy.github.io/spexial"><img alt="Documentation" src="https://img.shields.io/badge/docs-jaxtronomy.github.io-blue"></a>
</p>

`spexial` provides special functions for JAX, following the `scipy.special` API. The implementations are written in terms of JAX primitives, so they compose with `jit`, `grad` and `vmap`, and run on CPU, GPU and TPU.

It exists to fill gaps in `jax.scipy.special`: `zeta` accepts negative integers where JAX returns `nan`, `spence` accepts complex arguments where JAX is real-only, and the modified Bessel `K`, Gegenbauer and polylogarithm functions have no JAX counterpart at any version. Where SciPy already covers a case, `spexial` matches it rather than claiming to exceed it.

## Installation

```bash
pip install spexial
```

or

```bash
uv add spexial
```

## Example

```pycon
>>> import jax
>>> jax.config.update("jax_enable_x64", True)

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> # Gegenbauer polynomial C_n^alpha(x), matching scipy.special.eval_gegenbauer
>>> sp.eval_gegenbauer(3, 0.5, 0.25)
Array(-0.3359375, dtype=float64, weak_type=True)

```

Everything is vectorisable and differentiable in the usual way:

```pycon
>>> xs = jnp.linspace(-1.0, 1.0, 5)
>>> jax.vmap(lambda x: sp.eval_gegenbauer(3, 0.5, x))(xs)
Array([-1.    ,  0.4375, -0.    , -0.4375,  1.    ], dtype=float64)

```

## What is here

| Function            | `scipy.special` counterpart                           |
| ------------------- | ----------------------------------------------------- |
| `comb`              | `comb` (the `exact=False` variant)                    |
| `gamma`             | `gamma` — JAX's value, plus an analytic derivative    |
| `eval_gegenbauer`   | `eval_gegenbauer`                                     |
| `eval_gegenbauers`  | -- returns every order up to `n`                      |
| `K0`, `K1`, `K2`    | `k0`, `k1`, `kn`                                      |
| `K0e`, `K1e`, `K2e` | `k0e`, `k1e`, `kve` — scaled by `e^z`, no upper limit |
| `Li`                | -- the polylogarithm                                  |
| `spence`            | `spence` — complex too, which JAX rejects             |
| `zeta`              | `zeta` — negative integers, which JAX gives as `nan`  |

**Read [Accuracy and domains](https://jaxtronomy.github.io/spexial/reference/accuracy-and-domains/) before relying on any of these.** It records, per function, the domain each is tested over and the tolerance it actually meets. Some are not machine-precision — the modified Bessel functions are accurate to about `1e-7`, and `zeta` does not implement the critical strip.

## Documentation

<https://jaxtronomy.github.io/spexial>

## Development

We welcome contributions! See [CONTRIBUTING.md](CONTRIBUTING.md).

```bash
uv sync --group dev
uv run nox -s all      # lint -> test -> docs
```

## Citation

If you use `spexial` in work you publish, please cite it — see [CITATION.cff](CITATION.cff). Several routines originate in the [LINX](https://github.com/cgiovanetti/LINX) code.

## License

MIT. See [LICENSE](LICENSE).
