chebax 0.1.0__tar.gz

chebax-0.1.0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, Andres Hernandez
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

chebax-0.1.0/PKG-INFO

@@ -0,0 +1,94 @@
+Metadata-Version: 2.4
+Name: chebax
+Version: 0.1.0
+Summary: Differentiable Chebyshev approximants for special functions, aimed at JAX and GPUs
+Author-email: Andres Hernandez <andres.hernandez.deml@gmail.com>
+License-Expression: BSD-3-Clause
+Project-URL: Repository, https://github.com/andreshernandez-spec/chebax
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.5.3
+Requires-Dist: numpy>=1.26
+Provides-Extra: gen
+Requires-Dist: mpmath<1.5,>=1.4.1; extra == "gen"
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: scipy; extra == "test"
+Provides-Extra: examples
+Requires-Dist: matplotlib; extra == "examples"
+Requires-Dist: numpyro; extra == "examples"
+Dynamic: license-file
+
+# chebax
+
+Differentiable special functions for JAX and GPUs, built from Chebyshev
+approximants. A build-time generator (mpmath precision) produces fixed-degree
+branchless polynomial kernels with exact derivative series — including gradients
+with respect to function *parameters* (Bessel order, Beta shape, von Mises
+concentration), which mainstream ML libraries lack. No mpmath at use time; the
+runtime imports only jax and numpy.
+
+```sh
+pip install chebax
+```
+
+```python
+import jax
+jax.config.update("jax_enable_x64", True)   # before building instances: tables
+                                            # materialize in the active precision
+import jax.numpy as jnp
+import chebax
+
+jv = chebax.besselj(2.5)                    # any real order in [0, 10]; cached
+x = jnp.linspace(0.1, 30.0, 200)
+jv(x), jax.vmap(jax.grad(jv))(x)            # values and dJ/dx, jit/vmap-safe
+
+chebax.besselk_fn(1.7, x)                   # order as a traced scalar:
+jax.grad(chebax.besselk_fn)(1.7, 5.0)       #   jax.grad w.r.t. the ORDER
+chebax.betainc_fn(2.0, 3.0, 0.4)            # dI/da, dI/db via jax.grad (jax#38610)
+chebax.betaincinv(2.0, 3.0, 0.05)           # differentiable Beta quantile (jax#2399)
+```
+
+## What's in
+
+| family | factory | traced params | param derivative | notes |
+|---|---|---|---|---|
+| `besselj` | `besselj(v)` | — | `besselj_dnu(v)` | x ≥ 0, validated to 1e4 |
+| `bessely` | `bessely(v)` | — (structural) | `bessely_dnu(v)` | x ≥ 1e-6 |
+| `besseli` | `besseli(v, scaled=)` | `besseli_fn` | `besseli_dnu(v, scaled=)` | `scaled` = scipy's ive |
+| `besselk` | `besselk(v)` | `besselk_fn`, `log_besselk_fn` | `besselk_dnu(v)` | x ≥ 1e-6; Matérn demo; the log form has no underflow ceiling |
+| `betainc` | `betainc(a, b)` | `betainc_fn` | via `grad` of `_fn` | (a, b) ∈ [0.1, 10]² |
+| spherical | `spherical_jn/yn(n)` | — | — | n ∈ [0, 9], via half-integer tables |
+| quantiles | — | `betaincinv`, `gammaincinv`, `stdtr`, `stdtrit` | via `grad` (IFT) | jax#2399/#5350/#20358 |
+| von Mises | — | `vonmises_cdf/icdf` | via `grad` | κ ∈ [0, 50] |
+| erf family | — | `dawsn`, `erfcx` | — (no params) | recent jax ships both; kept for the C++ bake path |
+| Lambert W | — | `lambertw(x, k)` | — | both real branches, jax#13680 |
+
+Plus the generic core (`fit`, `ChebSeries`, `PiecewiseCheb`) and bake emitters
+(`chebax.bake.jax_module`, `chebax.bake.xsf_header`: self-contained pure-jax
+modules and C++17 headers from an instance).
+
+`notebooks/` holds themed, executed walkthroughs: why Chebyshev nodes work,
+the Bessel family with a learnable Matérn kernel, differentiable quantiles,
+truncated and circular distributions in numpyro, Gaussian tails / Lambert W /
+binomial reliability, copulas, and baking. The numpyro and plotting
+dependencies install with `pip install chebax[examples]`. `examples/` holds
+plain scripts: a few of the notebook workflows in copy-paste form, plus
+integration examples with no notebook counterpart (a GIG distribution for
+efax, learned Matérn smoothness, truncated sampling).
+
+Parameters come first, evaluation point last (opposite of scipy). Orders/shapes
+are uniform per call; per-element parameter arrays are out of scope by design.
+Accuracy is validated against mpmath at 40 dps — measured worst cases and the
+per-family error metrics are in each test file's docstring; tables regenerate
+bit-for-bit from checked-in generators.
+
+## Docs
+
+`docs/adding-a-recipe.md` is the workflow for new functions;
+`docs/increments.md` the design log of every recipe (what was measured, which
+traps were found); `PROJECT.md` the plan and evidence; `CLAUDE.md` how to work
+in the repo. Grown out of a private research project on GPU Bessel functions;
+its two load-bearing measurements are reproduced here in `experiments/`.
+BSD-3-Clause.

chebax-0.1.0/README.md

@@ -0,0 +1,72 @@
+# chebax
+
+Differentiable special functions for JAX and GPUs, built from Chebyshev
+approximants. A build-time generator (mpmath precision) produces fixed-degree
+branchless polynomial kernels with exact derivative series — including gradients
+with respect to function *parameters* (Bessel order, Beta shape, von Mises
+concentration), which mainstream ML libraries lack. No mpmath at use time; the
+runtime imports only jax and numpy.
+
+```sh
+pip install chebax
+```
+
+```python
+import jax
+jax.config.update("jax_enable_x64", True)   # before building instances: tables
+                                            # materialize in the active precision
+import jax.numpy as jnp
+import chebax
+
+jv = chebax.besselj(2.5)                    # any real order in [0, 10]; cached
+x = jnp.linspace(0.1, 30.0, 200)
+jv(x), jax.vmap(jax.grad(jv))(x)            # values and dJ/dx, jit/vmap-safe
+
+chebax.besselk_fn(1.7, x)                   # order as a traced scalar:
+jax.grad(chebax.besselk_fn)(1.7, 5.0)       #   jax.grad w.r.t. the ORDER
+chebax.betainc_fn(2.0, 3.0, 0.4)            # dI/da, dI/db via jax.grad (jax#38610)
+chebax.betaincinv(2.0, 3.0, 0.05)           # differentiable Beta quantile (jax#2399)
+```
+
+## What's in
+
+| family | factory | traced params | param derivative | notes |
+|---|---|---|---|---|
+| `besselj` | `besselj(v)` | — | `besselj_dnu(v)` | x ≥ 0, validated to 1e4 |
+| `bessely` | `bessely(v)` | — (structural) | `bessely_dnu(v)` | x ≥ 1e-6 |
+| `besseli` | `besseli(v, scaled=)` | `besseli_fn` | `besseli_dnu(v, scaled=)` | `scaled` = scipy's ive |
+| `besselk` | `besselk(v)` | `besselk_fn`, `log_besselk_fn` | `besselk_dnu(v)` | x ≥ 1e-6; Matérn demo; the log form has no underflow ceiling |
+| `betainc` | `betainc(a, b)` | `betainc_fn` | via `grad` of `_fn` | (a, b) ∈ [0.1, 10]² |
+| spherical | `spherical_jn/yn(n)` | — | — | n ∈ [0, 9], via half-integer tables |
+| quantiles | — | `betaincinv`, `gammaincinv`, `stdtr`, `stdtrit` | via `grad` (IFT) | jax#2399/#5350/#20358 |
+| von Mises | — | `vonmises_cdf/icdf` | via `grad` | κ ∈ [0, 50] |
+| erf family | — | `dawsn`, `erfcx` | — (no params) | recent jax ships both; kept for the C++ bake path |
+| Lambert W | — | `lambertw(x, k)` | — | both real branches, jax#13680 |
+
+Plus the generic core (`fit`, `ChebSeries`, `PiecewiseCheb`) and bake emitters
+(`chebax.bake.jax_module`, `chebax.bake.xsf_header`: self-contained pure-jax
+modules and C++17 headers from an instance).
+
+`notebooks/` holds themed, executed walkthroughs: why Chebyshev nodes work,
+the Bessel family with a learnable Matérn kernel, differentiable quantiles,
+truncated and circular distributions in numpyro, Gaussian tails / Lambert W /
+binomial reliability, copulas, and baking. The numpyro and plotting
+dependencies install with `pip install chebax[examples]`. `examples/` holds
+plain scripts: a few of the notebook workflows in copy-paste form, plus
+integration examples with no notebook counterpart (a GIG distribution for
+efax, learned Matérn smoothness, truncated sampling).
+
+Parameters come first, evaluation point last (opposite of scipy). Orders/shapes
+are uniform per call; per-element parameter arrays are out of scope by design.
+Accuracy is validated against mpmath at 40 dps — measured worst cases and the
+per-family error metrics are in each test file's docstring; tables regenerate
+bit-for-bit from checked-in generators.
+
+## Docs
+
+`docs/adding-a-recipe.md` is the workflow for new functions;
+`docs/increments.md` the design log of every recipe (what was measured, which
+traps were found); `PROJECT.md` the plan and evidence; `CLAUDE.md` how to work
+in the repo. Grown out of a private research project on GPU Bessel functions;
+its two load-bearing measurements are reproduced here in `experiments/`.
+BSD-3-Clause.

chebax-0.1.0/pyproject.toml

@@ -0,0 +1,33 @@
+[build-system]
+requires = ["setuptools>=77"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "chebax"
+version = "0.1.0"
+description = "Differentiable Chebyshev approximants for special functions, aimed at JAX and GPUs"
+readme = "README.md"
+license = "BSD-3-Clause"
+authors = [{ name = "Andres Hernandez", email = "andres.hernandez.deml@gmail.com" }]
+requires-python = ">=3.11"
+# jax floor = the oldest version the full suite passes on (CI pins a job to it)
+dependencies = ["jax>=0.5.3", "numpy>=1.26"]
+
+[project.urls]
+Repository = "https://github.com/andreshernandez-spec/chebax"
+
+[project.optional-dependencies]
+# tables record the mpmath version that built them; regeneration is verified
+# byte for byte, so the build dependency is pinned to the recorded minor
+gen = ["mpmath>=1.4.1,<1.5"]
+test = ["pytest", "scipy"]
+examples = ["matplotlib", "numpyro"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+markers = [
+    "slow: table regeneration and other multi-minute tests",
+]

chebax-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

chebax-0.1.0/src/chebax/__init__.py

@@ -0,0 +1,38 @@
+"""chebax: differentiable Chebyshev approximants for special functions.
+
+Full float64 accuracy needs jax's x64 mode (jax.config.update("jax_enable_x64", True)).
+"""
+
+from chebax._src.generate import fit
+from chebax._src.recipes.besseli import besseli, besseli_dnu, besseli_fn
+from chebax._src.recipes.erf_family import dawsn, erfcx
+from chebax._src.recipes.lambertw import lambertw
+from chebax._src.recipes.spherical import spherical_jn, spherical_yn
+from chebax._src.recipes.besselj import besselj, besselj_dnu
+from chebax._src.recipes.besselk import (besselk, besselk_dnu, besselk_fn,
+                                         log_besselk_fn)
+from chebax._src.recipes.bessely import bessely, bessely_dnu
+from chebax._src.recipes.betainc import betainc, betainc_fn
+from chebax._src.recipes.quantiles import betaincinv, gammaincinv, stdtr, stdtrit
+from chebax._src.recipes.vonmises import vonmises_cdf, vonmises_icdf
+from chebax._src.series import ChebSeries, PiecewiseCheb
+
+# single-sourced from pyproject.toml via package metadata
+try:
+    from importlib.metadata import version as _pkg_version
+    __version__ = _pkg_version("chebax")
+except Exception:  # not installed (e.g. running from a source checkout)
+    __version__ = "0+unknown"
+__all__ = ["ChebSeries", "PiecewiseCheb", "besseli", "besseli_dnu", "besseli_fn",
+           "besselj", "besselj_dnu", "besselk", "besselk_dnu", "besselk_fn",
+           "log_besselk_fn",
+           "bessely", "bessely_dnu", "betainc", "betainc_fn", "betaincinv",
+           "dawsn", "erfcx", "fit", "gammaincinv", "lambertw", "spherical_jn",
+           "spherical_yn", "stdtr", "stdtrit", "vonmises_cdf", "vonmises_icdf",
+           "__version__"]
+
+# submodule re-export so `import chebax; chebax.bake` works as documented
+# (bake reads __version__, so this import must stay below its definition)
+from chebax import bake  # noqa: E402,F401
+
+__all__.append("bake")

chebax-0.1.0/src/chebax/_src/algorithms.py

@@ -0,0 +1,74 @@
+"""The textbook algorithms, in plain numpy, on [-1, 1].
+
+Coefficient convention is numpy.polynomial's: plain c[0], lowest order first,
+p(t) = sum(c[k] * T_k(t)). These are the reference implementations; the jax
+runtime in series.py must agree with them (tested against numpy.polynomial).
+Domain handling lives in the series classes, not here.
+"""
+
+import numpy as np
+
+
+def chebpts(n):
+    """First-kind Chebyshev points cos(pi*(2i+1)/(2n)), decreasing from ~1 to ~-1."""
+    i = np.arange(n)
+    return np.cos(np.pi * (2 * i + 1) / (2 * n))
+
+
+def chebfit_dct(samples):
+    """Interpolation coefficients from samples taken at chebpts(len(samples)).
+
+    Discrete orthogonality of cosines on the first-kind points; O(n^2), which
+    is fine at generator degrees.
+    """
+    n = len(samples)
+    i = np.arange(n)
+    j = np.arange(n)[:, None]
+    # reduce the angle with exact integer arithmetic; the raw product reaches
+    # hundreds of radians at large n and costs ~angle*eps of cosine accuracy
+    weights = np.cos(np.pi * ((j * (2 * i + 1)) % (4 * n)) / (2 * n))
+    c = (2.0 / n) * (weights @ np.asarray(samples, dtype=float))
+    c[0] /= 2
+    return c
+
+
+def chebval(t, c):
+    """Clenshaw evaluation of a plain Chebyshev series."""
+    t = np.asarray(t, dtype=float)
+    b1 = np.zeros_like(t)
+    b2 = np.zeros_like(t)
+    for ck in np.asarray(c, dtype=float)[:0:-1]:
+        b1, b2 = 2 * t * b1 - b2 + ck, b1
+    return t * b1 - b2 + c[0]
+
+
+def chebder(c):
+    """Coefficients of the derivative (numpy.polynomial's exact recurrence)."""
+    c = np.array(c, dtype=float)
+    n = c.size - 1
+    if n == 0:
+        return np.zeros(1)
+    der = np.zeros(n)
+    for j in range(n, 2, -1):
+        der[j - 1] = 2 * j * c[j]
+        c[j - 2] += j * c[j] / (j - 2)
+    if n > 1:
+        der[1] = 4 * c[2]
+    der[0] = c[1]
+    return der
+
+
+def chebint(c):
+    """Coefficients of the antiderivative that vanishes at t=0
+    (numpy.polynomial's k=0, lbnd=0 convention)."""
+    c = np.asarray(c, dtype=float)
+    n = c.size
+    out = np.zeros(n + 1)
+    out[1] = c[0]
+    if n > 1:
+        out[2] = c[1] / 4
+    for j in range(2, n):
+        out[j + 1] = c[j] / (2 * (j + 1))
+        out[j - 1] -= c[j] / (2 * (j - 1))
+    out[0] -= chebval(0.0, out)
+    return out

chebax-0.1.0/src/chebax/_src/generate.py

@@ -0,0 +1,177 @@
+"""Build-time fitting: sample at Chebyshev points, DCT to coefficients.
+
+The float64 path uses numpy. Passing dps switches sampling and the DCT to
+mpmath at that precision; mpmath is imported inside that branch only (the
+runtime-imports rule in CLAUDE.md).
+
+Adaptive fitting doubles the node count until the coefficient tail sits at
+the fit's own noise floor, then truncates there. A converged-looking tail
+on one grid is never trusted by itself: on n first-kind points T_{2n-k}
+samples exactly as -T_k, so an aliased fit can look converged (T_34 on 17
+nodes reproduces -T_0 to the last bit). Every candidate is therefore
+validated against fresh samples at off-grid points before it is returned.
+
+The noise floor is estimated from the trailing quarter of the candidate's
+own coefficients rather than assumed: float64 sampling noise scales with
+the function's derivative (measured: ~1e-16 relative for exp, ~2.5e-15 for
+sin(50x)), so a fixed tolerance either fails to converge on oscillatory
+functions or keeps hundreds of noise coefficients.
+"""
+
+import numpy as np
+
+from chebax._src import algorithms
+from chebax._src.series import ChebSeries, PiecewiseCheb
+
+# ~2*eps: the practical floor of the float64 build path. Use dps= to go below.
+_DEFAULT_TOL = 5e-16
+
+
+def fit(f, domain=(-1.0, 1.0), *, deg=None, tol=None, breaks=None, dps=None, max_deg=256):
+    """Chebyshev interpolant of f, as a ChebSeries (or PiecewiseCheb if breaks given).
+
+    f:       callable on the domain. Vectorized over numpy arrays if possible;
+             with dps set it must accept mpmath floats instead.
+    domain:  (a, b). Ignored when breaks is given.
+    deg:     fixed degree; otherwise the degree adapts until tol is met.
+    tol:     target for the relative coefficient tail (default 5e-16). The
+             effective chop level is max(tol, the fit's measured noise floor);
+             float64 sampling noise can sit above tol for functions with
+             large derivatives. Build with dps to control the floor.
+    breaks:  strictly increasing knots, endpoints included; one series per
+             segment, each fit independently, padded to a common degree.
+    dps:     mpmath working precision for sampling and the DCT. Derivatives of
+             a float64-built fit sit at a ~deg^2 * eps floor (differentiation
+             amplifies sample noise); build with dps when that matters.
+    max_deg: hard cap on the returned degree. Reaching it without a validated
+             converged fit raises: f is not smooth enough on the interval
+             (split the domain) or genuinely needs a higher degree.
+    """
+    if tol is None:
+        tol = _DEFAULT_TOL
+    if breaks is not None:
+        br = [float(t) for t in breaks]
+        cs = [_fit_interval(f, br[i], br[i + 1], deg, tol, max_deg, dps)
+              for i in range(len(br) - 1)]
+        n = max(c.size for c in cs)
+        coef = np.zeros((len(cs), n))
+        for i, c in enumerate(cs):
+            coef[i, : c.size] = c
+        return PiecewiseCheb(coef, br)
+    a, b = float(domain[0]), float(domain[1])
+    return ChebSeries(_fit_interval(f, a, b, deg, tol, max_deg, dps), (a, b))
+
+
+def _fit_interval(f, a, b, deg, tol, max_deg, dps):
+    dct, sample_unit = (_dct_mp(f, a, b, dps) if dps is not None
+                        else _dct_np(f, a, b))
+    if deg is not None:
+        return dct(deg + 1)
+    return _adaptive(dct, sample_unit, tol, max_deg, (a, b))
+
+
+def _chop(c, tol):
+    """Truncate at the noise level; returns (candidate, absolute level).
+
+    The level is max(tol, noise floor), where the noise floor is read off
+    the last eighth of the coefficients: for a still-decaying tail those
+    sit below tol and the requested tol governs (measured: exp's last
+    coefficients decay through 1e-16); for a converged fit they are the
+    sampling-noise plateau (sin(50x): ~5e-16 relative at n=129, edge-
+    singular sampling can push higher) and doubling the largest of them
+    clears the plateau's own scatter."""
+    scale = np.max(np.abs(c))
+    noise = 2.0 * float(np.max(np.abs(c[-max(2, c.size // 16):])))
+    level = max(tol * scale, noise)
+    keep = np.nonzero(np.abs(c) > level)[0]
+    return (c[: keep[-1] + 1] if keep.size else c[:1]), level
+
+
+def _validated(cand, level, sample_unit):
+    """Check the candidate against fresh samples at off-grid points.
+
+    chebpts(m) for m != n shares no points with the fitting grid, so an
+    aliased fit cannot pass: its off-grid error is O(scale), ~1e15x the
+    chop level, while an honest fit's is ~1x. The 1000x margin makes this
+    a gross-error detector, not the accuracy contract.
+    """
+    m = cand.size + 11
+    t = algorithms.chebpts(m)
+    y = sample_unit(t)
+    err = np.max(np.abs(algorithms.chebval(t, cand) - y))
+    allow = max(1000.0 * level,
+                100.0 * np.finfo(float).eps * float(np.sum(np.abs(cand))))
+    return err <= allow
+
+
+def _adaptive(dct, sample_unit, tol, max_deg, interval):
+    n = min(17, max_deg + 1)
+    prev_rel = np.inf
+    while True:
+        c = dct(n)
+        scale = float(np.max(np.abs(c)))
+        if scale == 0.0:
+            return np.zeros(1)
+        cand, level = _chop(c, tol)
+        rel = level / scale
+        # The tail has settled if the chop level reached the requested tol,
+        # or if it stopped improving under node doubling while being far
+        # below the function scale: that is the sampling-noise plateau, and
+        # more nodes cannot beat it. A small-but-still-shrinking tail (a
+        # segment one doubling short of converged) must escalate, not chop.
+        settled = level <= 2.0 * tol * scale or (rel <= 1e-10 and rel >= prev_rel / 4.0)
+        if settled and cand.size <= c.size - max(2, c.size // 8):
+            if _validated(cand, level, sample_unit):
+                return cand
+        prev_rel = rel
+        if n >= max_deg + 1:
+            raise ValueError(
+                f"no validated convergence to tol={tol:g} by degree {max_deg} on "
+                f"{interval}; f may not be smooth there (split the domain), may "
+                "need a higher max_deg, or float64 sampling noise may sit above "
+                "tol (build with dps)")
+        n = min(2 * n - 1, max_deg + 1)
+
+
+def _dct_np(f, a, b):
+    def sample_unit(t):
+        xk = a + (b - a) * (t + 1) / 2
+        return _sample(f, xk)
+
+    def dct(n):
+        return algorithms.chebfit_dct(sample_unit(algorithms.chebpts(n)))
+    return dct, sample_unit
+
+
+def _sample(f, xk):
+    # scalar-only callables (wrapped special functions) are common; fall back.
+    try:
+        y = np.asarray(f(xk), dtype=float)
+        if y.shape == xk.shape:
+            return y
+    except (TypeError, ValueError):
+        pass
+    return np.array([float(f(x)) for x in xk])
+
+
+def _dct_mp(f, a, b, dps):
+    import mpmath as mp
+
+    def sample_unit(t):
+        with mp.workdps(dps):
+            mid = (mp.mpf(b) + mp.mpf(a)) / 2
+            half = (mp.mpf(b) - mp.mpf(a)) / 2
+            return np.array([float(f(mid + half * mp.mpf(float(tk)))) for tk in t])
+
+    def dct(n):
+        with mp.workdps(dps):
+            mid = (mp.mpf(b) + mp.mpf(a)) / 2
+            half = (mp.mpf(b) - mp.mpf(a)) / 2
+            fk = [f(mid + half * mp.cos(mp.pi * (2 * i + 1) / (2 * n))) for i in range(n)]
+            c = []
+            for j in range(n):
+                s = mp.fsum(fk[i] * mp.cos(mp.pi * j * (2 * i + 1) / (2 * n)) for i in range(n))
+                c.append(2 * s / n)
+            c[0] /= 2
+            return np.array([float(x) for x in c])
+    return dct, sample_unit

chebax-0.1.0/src/chebax/_src/pytree.py

@@ -0,0 +1,59 @@
+"""Base class for recipe callables, retiring per-class pytree boilerplate.
+
+A recipe instance is a pytree whose children are ChebSeries (listed in
+_series_fields) and whose auxiliary data are static scalars (listed in
+_static_fields). Constructors take statics first, series after, in field
+order; _post_init derives cached constants (called by both __init__ and
+tree_unflatten). astype casts every series; __repr__ shows the statics and
+the coefficient dtype.
+
+Subclasses register themselves with the standard
+@jax.tree_util.register_pytree_node_class decorator at the class site.
+"""
+
+class Recipe:
+    _static_fields = ()
+    _series_fields = ()
+
+    def __init__(self, *args):
+        names = self._static_fields + self._series_fields
+        if len(args) != len(names):
+            raise TypeError(f"{type(self).__name__} expects {names}, got {len(args)} args")
+        for name, val in zip(names, args):
+            setattr(self, name, val)
+        self._post_init()
+
+    def _post_init(self):
+        pass
+
+    def astype(self, dtype):
+        return type(self)(*[getattr(self, f) for f in self._static_fields],
+                          *[getattr(self, f).astype(dtype) for f in self._series_fields])
+
+    def __repr__(self):
+        statics = ", ".join(f"{f}={getattr(self, f)!r}" for f in self._static_fields)
+        dtype = "-"
+        # first series child with a coefficient array; children may be
+        # nested recipes (spherical wraps a cylindrical instance)
+        node = self
+        while getattr(node, "_series_fields", ()):
+            node = getattr(node, node._series_fields[0])
+            coef = getattr(node, "coef", None)
+            if coef is not None:
+                dtype = coef.dtype
+                break
+        return f"{type(self).__name__}({statics}, dtype={dtype})"
+
+    def tree_flatten(self):
+        return (tuple(getattr(self, f) for f in self._series_fields),
+                tuple(getattr(self, f) for f in self._static_fields))
+
+    @classmethod
+    def tree_unflatten(cls, aux, children):
+        obj = object.__new__(cls)
+        for name, val in zip(cls._static_fields, aux):
+            setattr(obj, name, val)
+        for name, val in zip(cls._series_fields, children):
+            setattr(obj, name, val)
+        obj._post_init()
+        return obj