Metadata-Version: 2.5
Name: spexial
Version: 0.2.1
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
Requires-Dist: typing-extensions>=4.5
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` |
| `incomplete_beta` | -- the _unregularized_ `B(a, b, z)`; `betainc` is the regularized form |
| `k0`, `k1`, `k2` | `k0`, `k1`, `kn` |
| `k0e`, `k1e`, `k2e` | `k0e`, `k1e`, `kve` — scaled by `e^z`, no upper limit |
| `polylog` | -- the polylogarithm |
| `spence` | `spence` — complex too, which JAX rejects |
| `sph_legendre_p` | `sph_legendre_p` — absent from JAX at any version |
| `sph_harm_y` | `sph_harm_y` — JAX's returns incorrect values for array degrees |
| `sph_harm_y_cart` | -- from a Cartesian direction; correct gradients on the z-axis |
| `sph_harm_y_cart_all` | `sph_harm_y_all` (layout only) — the whole `(l, m)` table in one pass |
| `sph_harm_y_cart_all_terms` | -- as above, returned as separate arrays so reductions stay fused |
| `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).
