chebax 0.1.0__py3-none-any.whl

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/_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/_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/_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

chebax/_src/recipes/_common.py

@@ -0,0 +1,95 @@
+"""Shared runtime machinery for recipes: parameter-table reconstruction,
+domain checks, the safeguarded root solver, and a float64 digamma.
+
+Reconstruction convention (matches _gen_common's emitted axis note): for a
+2-D table, TABLE[k, j] is the j-th parameter-direction coefficient of the
+k-th argument-direction coefficient, the parameter mapped affinely from
+[lo, hi] onto [-1, 1]. `param_coefs` is the eager numpy path used by the
+cached factories; `traced_coefs` is the jax path used by the *_fn traced
+variants (differentiable in the parameter through the series custom_jvp).
+"""
+
+import math
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from chebax._src import algorithms
+from chebax._src.series import _chebval
+
+
+def canon_tag():
+    """Cache-key tag for the canonical float dtype: factory instances built
+    under different x64 settings must not share a cache slot."""
+    return str(jnp.empty(()).dtype)
+
+
+def canon_float(x):
+    """Promote to the canonical float dtype (float64 under x64).
+
+    Solvers mix inputs with reconstructed float64 tables inside a fori_loop
+    carry; an explicit float32 input under x64 raised a carry dtype
+    mismatch instead of computing. One deliberate dtype for everything."""
+    return jnp.asarray(x, dtype=jnp.empty(()).dtype)
+
+
+def param_coefs(table, lo, hi, p):
+    """Coefficient vector at parameter p: Clenshaw across table rows. numpy."""
+    t = 2.0 * (p - lo) / (hi - lo) - 1.0
+    return np.array([algorithms.chebval(t, row) for row in table])
+
+
+def param_coefs_der(table, lo, hi, p):
+    """d/dp of the coefficient vector, via chebder along the parameter axis."""
+    t = 2.0 * (p - lo) / (hi - lo) - 1.0
+    d = np.array([algorithms.chebval(t, algorithms.chebder(row)) for row in table])
+    return d * (2.0 / (hi - lo))
+
+
+def traced_coefs(table, lo, hi, p):
+    """jax version of param_coefs: p may be a tracer (uniform per call)."""
+    t = 2.0 * (p - lo) / (hi - lo) - 1.0
+    return jax.vmap(lambda row: _chebval(t, row, (-1.0, 1.0)))(jnp.asarray(table))
+
+
+def check_range(owner, name, v, lo, hi):
+    v = float(v)
+    if not lo <= v <= hi:
+        raise ValueError(f"{owner} table covers {name} in [{lo}, {hi}], got {v}")
+    return v
+
+
+def newton_bisect(f_and_df, x0, lo0, hi0, iters):
+    """Fixed-count safeguarded Newton for increasing f: keep [lo, hi]
+    bracketing the root, fall back to the midpoint (or doubling while hi is
+    inf) when the Newton step leaves the bracket. The bounds test is
+    INCLUSIVE: a Newton step that stalls exactly at the current iterate
+    (converged) must be accepted, or a one-sided approach whose far bound
+    never tightened catapults a converged iterate to the midpoint of the
+    half-open bracket (measured, the quantile-toolkit bug)."""
+
+    def step(_, carry):
+        x, lo, hi = carry
+        f, df = f_and_df(x)
+        lo = jnp.where(f < 0, x, lo)
+        hi = jnp.where(f < 0, hi, x)
+        xn = x - f / df
+        mid = jnp.where(jnp.isfinite(hi), 0.5 * (lo + hi), 2.0 * x + 1.0)
+        ok = (xn >= lo) & (xn <= hi) & jnp.isfinite(xn)
+        return jnp.where(ok, xn, mid), lo, hi
+
+    x, _, _ = jax.lax.fori_loop(0, iters, step, (x0, lo0, hi0))
+    return x
+
+
+def digamma64(z):
+    """float64 psi(z) for z >= 1: recurrence up to z >= 15, then the
+    Bernoulli series through z^-12 (A&S 6.3.5 / 6.3.18). Truncation < 1e-16."""
+    acc = 0.0
+    while z < 15.0:
+        acc -= 1.0 / z
+        z += 1.0
+    zi = 1.0 / (z * z)
+    tail = 1 / 12 - zi * (1 / 120 - zi * (1 / 252 - zi * (1 / 240 - zi * (1 / 132 - zi * (691 / 32760)))))
+    return acc + math.log(z) - 0.5 / z - zi * tail

chebax/_src/recipes/_gen_common.py

@@ -0,0 +1,121 @@
+"""Shared build-time machinery for all recipe generators.
+
+Every generator uses the same two-stage (or three-stage) construction:
+Chebyshev-Gauss nodes, a plain DCT to coefficients, a parameter-direction
+fit per coefficient, and a writer that emits the baked table as a diffable
+python module. All fit arithmetic stays in mpmath (the M1 floor rule:
+rounding between stages puts eps-level sample noise under later fits).
+
+Table axis convention, written into every emitted header: for a 2-D table,
+TABLE[k, j] is the j-th parameter-direction coefficient of the k-th
+argument-direction coefficient; parameters map affinely onto [-1, 1] from
+the domain recorded in META. Emission is deterministic (repr of python
+floats round-trips; no timestamps), so identical inputs give identical
+bytes regardless of numpy version.
+"""
+
+import hashlib
+import importlib
+import pathlib
+
+import numpy as np
+
+DPS = 40  # mpmath working precision for every generator
+
+
+def table_path(out_dir, generator_file, name):
+    """Output location for an emitted table: the package directory by
+    default, or out_dir (the regeneration tests emit to a temp dir and
+    compare full files against the packaged ones)."""
+    base = pathlib.Path(out_dir) if out_dir is not None else pathlib.Path(generator_file).parent
+    return base / name
+
+
+def _generator_sha(generator_module):
+    """sha256 over the generator's own source plus this shared module: the
+    emitted file records exactly the code that built it, so any edit to
+    either shows up as a hash mismatch until the table is regenerated."""
+    h = hashlib.sha256()
+    for mod_name in (generator_module, __name__):
+        m = importlib.import_module(mod_name)
+        h.update(pathlib.Path(m.__file__).read_bytes())
+    return h.hexdigest()[:16]
+
+
+def nodes(mp, n):
+    """First-kind Chebyshev points, the sampling nodes of every fit."""
+    return [mp.cos(mp.pi * (2 * i + 1) / (2 * n)) for i in range(n)]
+
+
+def dct(mp, samples):
+    """Plain Chebyshev coefficients from samples at nodes(mp, len(samples))."""
+    n = len(samples)
+    c = []
+    for j in range(n):
+        s = mp.fsum(samples[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 c
+
+
+def param_fit(mp, rows):
+    """Second stage: fit each coefficient across the parameter nodes, in mp."""
+    nc = len(rows[0])
+    return [dct(mp, [rows[i][k] for i in range(len(rows))]) for k in range(nc)]
+
+
+def to_f64(rows):
+    return np.array([[float(x) for x in row] for row in rows])
+
+
+def _fmt_array(name, arr):
+    arr = np.asarray(arr)
+    lines = []
+    if arr.ndim == 1:
+        lines.append(f"{name} = np.array([")
+        for i in range(0, arr.size, 4):
+            lines.append("    " + ", ".join(repr(float(v)) for v in arr[i:i + 4]) + ",")
+        lines.append("])")
+    elif arr.ndim == 2:
+        lines.append(f"{name} = np.array([")
+        for row in arr:
+            lines.append("    [" + ", ".join(repr(float(v)) for v in row) + "],")
+        lines.append("])")
+    elif arr.ndim == 3:
+        lines.append(f"{name} = np.array([")
+        for block in arr:
+            lines.append("  [")
+            for row in block:
+                lines.append("    [" + ", ".join(repr(float(v)) for v in row) + "],")
+            lines.append("  ],")
+        lines.append("])")
+    else:
+        raise ValueError(f"unsupported table rank {arr.ndim}")
+    return lines
+
+
+def write_table_module(path, generator_module, axis_note, meta, scalars, arrays):
+    """Emit a baked table as a python module. Deterministic bytes."""
+    import mpmath
+
+    meta = dict(meta, dps=DPS, mpmath=mpmath.__version__,
+                generator_sha256=_generator_sha(generator_module))
+    lines = [
+        f'"""Baked table. Generated by {generator_module}; do not edit.',
+        f"Regenerate with:  python -m {generator_module}",
+        "",
+        f"Axes: {axis_note}",
+        '"""',
+        "",
+        "import numpy as np",
+        "",
+        f"META = {meta!r}",
+    ]
+    for k, v in scalars.items():
+        lines.append(f"{k} = {v!r}")
+    for name, arr in arrays.items():
+        lines.append("")
+        lines.extend(_fmt_array(name, arr))
+    path.write_text("\n".join(lines) + "\n")
+    shapes = ", ".join(f"{k} {np.asarray(v).shape}" for k, v in arrays.items())
+    print(f"wrote {path.name}: {shapes}")