chebax 0.1.0.dev0__tar.gz → 0.2.0__tar.gz

chebax-0.2.0/PKG-INFO

@@ -0,0 +1,106 @@
+Metadata-Version: 2.4
+Name: chebax
+Version: 0.2.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"
+Provides-Extra: pytensor
+Requires-Dist: pytensor>=3.2; extra == "pytensor"
+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).
+
+For pymc users: `import chebax.pytensor` (installs with
+`pip install chebax[pytensor]`) registers JAX-backend lowerings for the
+pytensor ops that otherwise need tfp-nightly under
+`pm.sample(nuts_sampler="numpyro"|"blackjax")` — betaincinv, gammaincinv,
+gammainccinv, erfcx, erfcinv, ive, kve — and adds betainc gradients in its
+shape parameters, so censored or truncated StudentT/Beta likelihoods with a
+latent scalar shape sample instead of raising. Shape parameters must be
+scalar (batched ones fall back to tfp or fail loudly); out-of-domain values
+return nan, never silently wrong numbers.
+
+`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.2.0/README.md

@@ -0,0 +1,82 @@
+# 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).
+
+For pymc users: `import chebax.pytensor` (installs with
+`pip install chebax[pytensor]`) registers JAX-backend lowerings for the
+pytensor ops that otherwise need tfp-nightly under
+`pm.sample(nuts_sampler="numpyro"|"blackjax")` — betaincinv, gammaincinv,
+gammainccinv, erfcx, erfcinv, ive, kve — and adds betainc gradients in its
+shape parameters, so censored or truncated StudentT/Beta likelihoods with a
+latent scalar shape sample instead of raising. Shape parameters must be
+scalar (batched ones fall back to tfp or fail loudly); out-of-domain values
+return nan, never silently wrong numbers.
+
+`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.2.0/pyproject.toml

@@ -0,0 +1,34 @@
+[build-system]
+requires = ["setuptools>=77"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "chebax"
+version = "0.2.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"]
+pytensor = ["pytensor>=3.2"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+markers = [
+    "slow: table regeneration and other multi-minute tests",
+]

chebax-0.2.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.2.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.2.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

chebax-0.2.0/src/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