chebax 0.1.0.dev0__tar.gz → 0.3.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- chebax-0.3.0/PKG-INFO +132 -0
- chebax-0.3.0/README.md +106 -0
- chebax-0.3.0/pyproject.toml +35 -0
- chebax-0.3.0/src/chebax/__init__.py +54 -0
- chebax-0.3.0/src/chebax/_src/generate.py +280 -0
- chebax-0.3.0/src/chebax/_src/pergroup.py +113 -0
- chebax-0.3.0/src/chebax/_src/pytree.py +87 -0
- chebax-0.3.0/src/chebax/_src/recipes/_common.py +134 -0
- chebax-0.3.0/src/chebax/_src/recipes/_gen_common.py +121 -0
- chebax-0.3.0/src/chebax/_src/recipes/besseli.py +376 -0
- chebax-0.3.0/src/chebax/_src/recipes/besseli_gen.py +84 -0
- chebax-0.3.0/src/chebax/_src/recipes/besseli_table.py +138 -0
- chebax-0.3.0/src/chebax/_src/recipes/besselj.py +251 -0
- chebax-0.3.0/src/chebax/_src/recipes/besselj_gen.py +146 -0
- chebax-0.3.0/src/chebax/_src/recipes/besselj_table.py +39 -0
- chebax-0.3.0/src/chebax/_src/recipes/besselj_table_ext.py +87 -0
- chebax-0.3.0/src/chebax/_src/recipes/besselk.py +209 -0
- chebax-0.3.0/src/chebax/_src/recipes/besselk_gen.py +107 -0
- chebax-0.3.0/src/chebax/_src/recipes/besselk_table.py +208 -0
- chebax-0.3.0/src/chebax/_src/recipes/bessely.py +185 -0
- chebax-0.3.0/src/chebax/_src/recipes/bessely_gen.py +103 -0
- chebax-0.3.0/src/chebax/_src/recipes/bessely_table.py +284 -0
- chebax-0.3.0/src/chebax/_src/recipes/betainc.py +255 -0
- chebax-0.3.0/src/chebax/_src/recipes/betainc_gen.py +78 -0
- chebax-0.3.0/src/chebax/_src/recipes/betainc_table.py +1791 -0
- chebax-0.3.0/src/chebax/_src/recipes/betainc_wide_gen.py +110 -0
- chebax-0.3.0/src/chebax/_src/recipes/betainc_wide_table.py +11390 -0
- chebax-0.3.0/src/chebax/_src/recipes/erf_family.py +75 -0
- chebax-0.3.0/src/chebax/_src/recipes/erf_gen.py +69 -0
- chebax-0.3.0/src/chebax/_src/recipes/erf_table.py +51 -0
- chebax-0.3.0/src/chebax/_src/recipes/gammainc.py +259 -0
- chebax-0.3.0/src/chebax/_src/recipes/gammainc_gen.py +99 -0
- chebax-0.3.0/src/chebax/_src/recipes/gammainc_large.py +178 -0
- chebax-0.3.0/src/chebax/_src/recipes/gammainc_large_gen.py +172 -0
- chebax-0.3.0/src/chebax/_src/recipes/gammainc_large_table.py +90 -0
- chebax-0.3.0/src/chebax/_src/recipes/gammainc_table.py +76 -0
- chebax-0.3.0/src/chebax/_src/recipes/hyp1f1.py +176 -0
- chebax-0.3.0/src/chebax/_src/recipes/hyp1f1_gen.py +139 -0
- chebax-0.3.0/src/chebax/_src/recipes/hyp1f1_table.py +6002 -0
- chebax-0.3.0/src/chebax/_src/recipes/lambertw.py +106 -0
- chebax-0.3.0/src/chebax/_src/recipes/matern.py +55 -0
- chebax-0.3.0/src/chebax/_src/recipes/quantiles.py +788 -0
- chebax-0.3.0/src/chebax/_src/recipes/spherical.py +100 -0
- chebax-0.3.0/src/chebax/_src/recipes/stdtr_gen.py +74 -0
- chebax-0.3.0/src/chebax/_src/recipes/stdtr_table.py +297 -0
- chebax-0.3.0/src/chebax/_src/recipes/vonmises.py +182 -0
- chebax-0.3.0/src/chebax/_src/recipes/vonmises_gen.py +68 -0
- chebax-0.3.0/src/chebax/_src/recipes/vonmises_table.py +127 -0
- chebax-0.3.0/src/chebax/_src/series.py +353 -0
- chebax-0.3.0/src/chebax/bake/__init__.py +319 -0
- chebax-0.3.0/src/chebax/bake/_jaxpr_emit.py +407 -0
- chebax-0.3.0/src/chebax/domains.py +59 -0
- chebax-0.3.0/src/chebax/numpyro.py +391 -0
- chebax-0.3.0/src/chebax/pytensor.py +196 -0
- chebax-0.3.0/src/chebax.egg-info/PKG-INFO +132 -0
- chebax-0.3.0/src/chebax.egg-info/SOURCES.txt +85 -0
- chebax-0.3.0/src/chebax.egg-info/requires.txt +19 -0
- chebax-0.3.0/tests/test_bake.py +582 -0
- chebax-0.3.0/tests/test_besseli.py +292 -0
- {chebax-0.1.0.dev0 → chebax-0.3.0}/tests/test_besselj.py +38 -7
- chebax-0.3.0/tests/test_besselj_m3.py +186 -0
- chebax-0.3.0/tests/test_besselk.py +308 -0
- chebax-0.3.0/tests/test_bessely.py +87 -0
- chebax-0.3.0/tests/test_betainc.py +234 -0
- chebax-0.3.0/tests/test_contracts.py +99 -0
- chebax-0.3.0/tests/test_erf_family.py +278 -0
- chebax-0.3.0/tests/test_gammainc.py +282 -0
- chebax-0.3.0/tests/test_gammainc_large.py +259 -0
- chebax-0.3.0/tests/test_hyp1f1.py +173 -0
- chebax-0.3.0/tests/test_int_params.py +65 -0
- chebax-0.3.0/tests/test_lambertw.py +80 -0
- chebax-0.3.0/tests/test_m1.py +465 -0
- chebax-0.3.0/tests/test_matern.py +108 -0
- chebax-0.3.0/tests/test_numpyro_module.py +502 -0
- chebax-0.3.0/tests/test_pergroup.py +199 -0
- chebax-0.3.0/tests/test_pytensor_plugin.py +142 -0
- chebax-0.3.0/tests/test_quantiles.py +691 -0
- chebax-0.3.0/tests/test_spherical.py +59 -0
- chebax-0.3.0/tests/test_vonmises.py +241 -0
- chebax-0.3.0/tests/test_x64_off.py +129 -0
- chebax-0.1.0.dev0/PKG-INFO +0 -35
- chebax-0.1.0.dev0/README.md +0 -17
- chebax-0.1.0.dev0/pyproject.toml +0 -26
- chebax-0.1.0.dev0/src/chebax/__init__.py +0 -11
- chebax-0.1.0.dev0/src/chebax/_src/generate.py +0 -106
- chebax-0.1.0.dev0/src/chebax/_src/recipes/besselj.py +0 -138
- chebax-0.1.0.dev0/src/chebax/_src/recipes/besselj_gen.py +0 -76
- chebax-0.1.0.dev0/src/chebax/_src/recipes/besselj_table.py +0 -37
- chebax-0.1.0.dev0/src/chebax/_src/series.py +0 -195
- chebax-0.1.0.dev0/src/chebax.egg-info/PKG-INFO +0 -35
- chebax-0.1.0.dev0/src/chebax.egg-info/SOURCES.txt +0 -19
- chebax-0.1.0.dev0/src/chebax.egg-info/requires.txt +0 -8
- chebax-0.1.0.dev0/tests/test_m1.py +0 -212
- {chebax-0.1.0.dev0 → chebax-0.3.0}/LICENSE +0 -0
- {chebax-0.1.0.dev0 → chebax-0.3.0}/setup.cfg +0 -0
- {chebax-0.1.0.dev0 → chebax-0.3.0}/src/chebax/_src/__init__.py +0 -0
- {chebax-0.1.0.dev0 → chebax-0.3.0}/src/chebax/_src/algorithms.py +0 -0
- {chebax-0.1.0.dev0 → chebax-0.3.0}/src/chebax/_src/recipes/__init__.py +0 -0
- {chebax-0.1.0.dev0 → chebax-0.3.0}/src/chebax.egg-info/dependency_links.txt +0 -0
- {chebax-0.1.0.dev0 → chebax-0.3.0}/src/chebax.egg-info/top_level.txt +0 -0
chebax-0.3.0/PKG-INFO
ADDED
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: chebax
|
|
3
|
+
Version: 0.3.0
|
|
4
|
+
Summary: Differentiable Chebyshev approximants for special functions, aimed at JAX and GPUs
|
|
5
|
+
Author-email: Andres Hernandez <andres.hernandez.deml@gmail.com>
|
|
6
|
+
License-Expression: BSD-3-Clause
|
|
7
|
+
Project-URL: Repository, https://github.com/andreshernandez-spec/chebax
|
|
8
|
+
Requires-Python: >=3.11
|
|
9
|
+
Description-Content-Type: text/markdown
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Requires-Dist: jax>=0.5.3
|
|
12
|
+
Requires-Dist: numpy>=1.26
|
|
13
|
+
Provides-Extra: gen
|
|
14
|
+
Requires-Dist: mpmath<1.5,>=1.4.1; extra == "gen"
|
|
15
|
+
Provides-Extra: test
|
|
16
|
+
Requires-Dist: pytest; extra == "test"
|
|
17
|
+
Requires-Dist: scipy; extra == "test"
|
|
18
|
+
Provides-Extra: examples
|
|
19
|
+
Requires-Dist: matplotlib; extra == "examples"
|
|
20
|
+
Requires-Dist: numpyro; extra == "examples"
|
|
21
|
+
Provides-Extra: numpyro
|
|
22
|
+
Requires-Dist: numpyro>=0.15; extra == "numpyro"
|
|
23
|
+
Provides-Extra: pytensor
|
|
24
|
+
Requires-Dist: pytensor>=3.2; extra == "pytensor"
|
|
25
|
+
Dynamic: license-file
|
|
26
|
+
|
|
27
|
+
# chebax
|
|
28
|
+
|
|
29
|
+
Differentiable special functions for JAX and GPUs, built from Chebyshev
|
|
30
|
+
approximants. A build-time generator (mpmath precision) produces fixed-degree
|
|
31
|
+
branchless polynomial kernels with exact derivative series — including gradients
|
|
32
|
+
with respect to function *parameters* (Bessel order, Beta shape, von Mises
|
|
33
|
+
concentration), which mainstream ML libraries lack. No mpmath at use time; the
|
|
34
|
+
runtime imports only jax and numpy.
|
|
35
|
+
|
|
36
|
+
```sh
|
|
37
|
+
pip install chebax
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
```python
|
|
41
|
+
import jax
|
|
42
|
+
jax.config.update("jax_enable_x64", True) # before building instances: tables
|
|
43
|
+
# materialize in the active precision
|
|
44
|
+
import jax.numpy as jnp
|
|
45
|
+
import chebax
|
|
46
|
+
|
|
47
|
+
jv = chebax.besselj(2.5) # any real order in [0, 10]; cached
|
|
48
|
+
x = jnp.linspace(0.1, 30.0, 200)
|
|
49
|
+
jv(x), jax.vmap(jax.grad(jv))(x) # values and dJ/dx, jit/vmap-safe
|
|
50
|
+
|
|
51
|
+
chebax.besselk_fn(1.7, x) # order as a traced scalar:
|
|
52
|
+
jax.grad(chebax.besselk_fn)(1.7, 5.0) # jax.grad w.r.t. the ORDER
|
|
53
|
+
chebax.betainc_fn(2.0, 3.0, 0.4) # dI/da, dI/db via jax.grad (jax#38610)
|
|
54
|
+
chebax.betaincinv(2.0, 3.0, 0.05) # differentiable Beta quantile (jax#2399)
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
## What's in
|
|
58
|
+
|
|
59
|
+
| family | factory | traced params | param derivative | notes |
|
|
60
|
+
|---|---|---|---|---|
|
|
61
|
+
| `besselj` | `besselj(v)`, `besselj(v, domain=(lo, hi))` | — | `besselj_dnu(v)` | x ≥ 0, validated to 1e4; `domain=` keeps only the regions covering the window and drops the rest of the evaluation, measured 2.5x f64 / 2.6x f32 (`experiments/04`), nan outside it |
|
|
62
|
+
| `bessely` | `bessely(v)` | — (structural) | `bessely_dnu(v)` | x ≥ 1e-6 |
|
|
63
|
+
| `besseli` | `besseli(v, scaled=)` | `besseli_fn`, `besseli_ratio` | `besseli_dnu(v, scaled=)` | `scaled` = scipy's ive; the ratio I_{v+1}/I_v is the circular-statistics workhorse |
|
|
64
|
+
| `besselk` | `besselk(v)` | `besselk_fn`, `log_besselk_fn` | `besselk_dnu(v)` | x ≥ 1e-6; the log form has no underflow ceiling |
|
|
65
|
+
| Matérn | — | `matern(nu, r, lengthscale)` | via `grad` | unit-variance correlation, LEARNABLE smoothness ν ∈ (0, 10] |
|
|
66
|
+
| `betainc` | `betainc(a, b)` | `betainc_fn`, `log_betainc_fn` | via `grad` of `_fn` | (a, b) ∈ [0.1, 100]² (panels, split at 10); the log form resolves the lower tail with no underflow floor |
|
|
67
|
+
| `gammainc` | `gammainc(a)`, `gammaincc(a)` | `gammainc_fn`, `gammaincc_fn`, `log_gammainc_fn`, `log_gammaincc_fn` | via `grad` of `_fn` | a ≥ 0.1 (above 10 via Temme-zone tables, whose v = 10/a axis reaches a = ∞), x ≥ 0; branchless (no while_loop), measured 10–27x vs jax's on GPU f64 at a ≤ 10 (`experiments/05`) |
|
|
68
|
+
| `hyp1f1` | `hyp1f1(a, b)` | `hyp1f1_fn`, `log_hyp1f1_fn` | via `grad` of `_fn` | Kummer M, (a, b) ∈ [0.1, 10]², x ≥ 0; jax's is documented-unstable (jax#21503); the log form has no overflow ceiling |
|
|
69
|
+
| spherical | `spherical_jn/yn(n)` | — | — | n ∈ [0, 9], via half-integer tables |
|
|
70
|
+
| quantiles | — | `betaincinv`, `gammaincinv`, `gammainccinv`, `chi2inv`, `stdtr`, `stdtrit` | via `grad` (IFT) | jax#2399/#5350/#20358; chi-squared at real dof to 2000; Student-t ν ∈ [0.2, 200] via dedicated slice tables |
|
|
71
|
+
| von Mises | — | `vonmises_cdf/icdf` | via `grad` | κ ∈ [0, 50] |
|
|
72
|
+
| erf family | — | `dawsn`, `erfcx` | — (no params) | recent jax ships both; kept as a worked parameter-free recipe (plain functions, so not `bake` inputs: bake takes Recipe instances) |
|
|
73
|
+
| Lambert W | — | `lambertw(x, k)` | — | both real branches, jax#13680 |
|
|
74
|
+
|
|
75
|
+
Plus the generic core (`fit`, `ChebSeries`, `PiecewiseCheb`) and bake emitters
|
|
76
|
+
(`chebax.bake.jax_module`, `chebax.bake.xsf_header`: self-contained pure-jax
|
|
77
|
+
modules and C++17 headers from an instance; `dtype="float"` emits an f32
|
|
78
|
+
CUDA-shaped kernel body, and `truncate_tol=1e-7` drops the converged f64
|
|
79
|
+
tail — measured 1.4–2.2x faster f32 evaluation, `experiments/12`).
|
|
80
|
+
|
|
81
|
+
For pymc users: `import chebax.pytensor` (installs with
|
|
82
|
+
`pip install chebax[pytensor]`) registers JAX-backend lowerings for the
|
|
83
|
+
pytensor ops that otherwise need tfp-nightly under
|
|
84
|
+
`pm.sample(nuts_sampler="numpyro"|"blackjax")` — betaincinv, gammaincinv,
|
|
85
|
+
gammainccinv, erfcx, erfcinv, ive, kve — and adds betainc gradients in its
|
|
86
|
+
shape parameters, so censored or truncated StudentT/Beta likelihoods with a
|
|
87
|
+
latent scalar shape sample instead of raising. Shape parameters must be
|
|
88
|
+
scalar (batched ones fall back to tfp or fail loudly); out-of-domain values
|
|
89
|
+
return nan, never silently wrong numbers. `gammainccinv` takes its
|
|
90
|
+
upper-tail probability directly (it used to be solved at `1 - p`, which
|
|
91
|
+
collapsed everything below ~1e-17 onto one point), so both tails resolve to
|
|
92
|
+
the same depth.
|
|
93
|
+
|
|
94
|
+
For numpyro users: `from chebax.numpyro import TruncatedGamma, TruncatedBeta,
|
|
95
|
+
TruncatedStudentT` (installs with `pip install chebax[numpyro]`) gives the
|
|
96
|
+
truncated distributions numpyro's location-scale machinery cannot cover
|
|
97
|
+
(numpyro#969, numpyro#1365): full Distribution classes with reparameterized
|
|
98
|
+
inverse-CDF sampling (`has_rsample`), truncation-normalized `log_prob`, and
|
|
99
|
+
gradients through every parameter including the shapes, so NUTS and SVI work
|
|
100
|
+
with a latent concentration or df. The truncation is normalized in log space,
|
|
101
|
+
so an interval whose two ends land in the same tail still works: `Gamma(3)`
|
|
102
|
+
truncated to [50, 51] has mass ~3e-20 and returns an ordinary density.
|
|
103
|
+
Shape parameters are uniform per call; domain boxes are in the module
|
|
104
|
+
docstring and checked at construction.
|
|
105
|
+
|
|
106
|
+
`notebooks/` holds themed, executed walkthroughs: why Chebyshev nodes work,
|
|
107
|
+
the Bessel family with a learnable Matérn kernel, differentiable quantiles,
|
|
108
|
+
truncated and circular distributions in numpyro, Gaussian tails / Lambert W /
|
|
109
|
+
binomial reliability, copulas, and baking. The numpyro and plotting
|
|
110
|
+
dependencies install with `pip install chebax[examples]`. `examples/` holds
|
|
111
|
+
plain scripts: a few of the notebook workflows in copy-paste form, plus
|
|
112
|
+
integration examples with no notebook counterpart (a GIG distribution for
|
|
113
|
+
efax, learned Matérn smoothness, truncated sampling, `chebax.numpyro` with
|
|
114
|
+
latent shape parameters under NUTS).
|
|
115
|
+
|
|
116
|
+
Parameters come first, evaluation point last (opposite of scipy). Orders/shapes
|
|
117
|
+
are uniform per call; per-element parameter arrays are out of scope by design.
|
|
118
|
+
The middle ground is `chebax.pergroup(fn, group_idx)`: a static integer array
|
|
119
|
+
assigns each element to a group, each group gets its own (traceable) parameter
|
|
120
|
+
set — one group per chain, mixture component, or plate level.
|
|
121
|
+
Accuracy is validated against mpmath at 40 dps — measured worst cases and the
|
|
122
|
+
per-family error metrics are in each test file's docstring; tables regenerate
|
|
123
|
+
bit-for-bit from checked-in generators.
|
|
124
|
+
|
|
125
|
+
## Docs
|
|
126
|
+
|
|
127
|
+
`docs/adding-a-recipe.md` is the workflow for new functions;
|
|
128
|
+
`docs/increments.md` the design log of every recipe (what was measured, which
|
|
129
|
+
traps were found); `PROJECT.md` the plan and evidence; `CLAUDE.md` how to work
|
|
130
|
+
in the repo. Grown out of a private research project on GPU Bessel functions;
|
|
131
|
+
its two load-bearing measurements are reproduced here in `experiments/`.
|
|
132
|
+
BSD-3-Clause.
|
chebax-0.3.0/README.md
ADDED
|
@@ -0,0 +1,106 @@
|
|
|
1
|
+
# chebax
|
|
2
|
+
|
|
3
|
+
Differentiable special functions for JAX and GPUs, built from Chebyshev
|
|
4
|
+
approximants. A build-time generator (mpmath precision) produces fixed-degree
|
|
5
|
+
branchless polynomial kernels with exact derivative series — including gradients
|
|
6
|
+
with respect to function *parameters* (Bessel order, Beta shape, von Mises
|
|
7
|
+
concentration), which mainstream ML libraries lack. No mpmath at use time; the
|
|
8
|
+
runtime imports only jax and numpy.
|
|
9
|
+
|
|
10
|
+
```sh
|
|
11
|
+
pip install chebax
|
|
12
|
+
```
|
|
13
|
+
|
|
14
|
+
```python
|
|
15
|
+
import jax
|
|
16
|
+
jax.config.update("jax_enable_x64", True) # before building instances: tables
|
|
17
|
+
# materialize in the active precision
|
|
18
|
+
import jax.numpy as jnp
|
|
19
|
+
import chebax
|
|
20
|
+
|
|
21
|
+
jv = chebax.besselj(2.5) # any real order in [0, 10]; cached
|
|
22
|
+
x = jnp.linspace(0.1, 30.0, 200)
|
|
23
|
+
jv(x), jax.vmap(jax.grad(jv))(x) # values and dJ/dx, jit/vmap-safe
|
|
24
|
+
|
|
25
|
+
chebax.besselk_fn(1.7, x) # order as a traced scalar:
|
|
26
|
+
jax.grad(chebax.besselk_fn)(1.7, 5.0) # jax.grad w.r.t. the ORDER
|
|
27
|
+
chebax.betainc_fn(2.0, 3.0, 0.4) # dI/da, dI/db via jax.grad (jax#38610)
|
|
28
|
+
chebax.betaincinv(2.0, 3.0, 0.05) # differentiable Beta quantile (jax#2399)
|
|
29
|
+
```
|
|
30
|
+
|
|
31
|
+
## What's in
|
|
32
|
+
|
|
33
|
+
| family | factory | traced params | param derivative | notes |
|
|
34
|
+
|---|---|---|---|---|
|
|
35
|
+
| `besselj` | `besselj(v)`, `besselj(v, domain=(lo, hi))` | — | `besselj_dnu(v)` | x ≥ 0, validated to 1e4; `domain=` keeps only the regions covering the window and drops the rest of the evaluation, measured 2.5x f64 / 2.6x f32 (`experiments/04`), nan outside it |
|
|
36
|
+
| `bessely` | `bessely(v)` | — (structural) | `bessely_dnu(v)` | x ≥ 1e-6 |
|
|
37
|
+
| `besseli` | `besseli(v, scaled=)` | `besseli_fn`, `besseli_ratio` | `besseli_dnu(v, scaled=)` | `scaled` = scipy's ive; the ratio I_{v+1}/I_v is the circular-statistics workhorse |
|
|
38
|
+
| `besselk` | `besselk(v)` | `besselk_fn`, `log_besselk_fn` | `besselk_dnu(v)` | x ≥ 1e-6; the log form has no underflow ceiling |
|
|
39
|
+
| Matérn | — | `matern(nu, r, lengthscale)` | via `grad` | unit-variance correlation, LEARNABLE smoothness ν ∈ (0, 10] |
|
|
40
|
+
| `betainc` | `betainc(a, b)` | `betainc_fn`, `log_betainc_fn` | via `grad` of `_fn` | (a, b) ∈ [0.1, 100]² (panels, split at 10); the log form resolves the lower tail with no underflow floor |
|
|
41
|
+
| `gammainc` | `gammainc(a)`, `gammaincc(a)` | `gammainc_fn`, `gammaincc_fn`, `log_gammainc_fn`, `log_gammaincc_fn` | via `grad` of `_fn` | a ≥ 0.1 (above 10 via Temme-zone tables, whose v = 10/a axis reaches a = ∞), x ≥ 0; branchless (no while_loop), measured 10–27x vs jax's on GPU f64 at a ≤ 10 (`experiments/05`) |
|
|
42
|
+
| `hyp1f1` | `hyp1f1(a, b)` | `hyp1f1_fn`, `log_hyp1f1_fn` | via `grad` of `_fn` | Kummer M, (a, b) ∈ [0.1, 10]², x ≥ 0; jax's is documented-unstable (jax#21503); the log form has no overflow ceiling |
|
|
43
|
+
| spherical | `spherical_jn/yn(n)` | — | — | n ∈ [0, 9], via half-integer tables |
|
|
44
|
+
| quantiles | — | `betaincinv`, `gammaincinv`, `gammainccinv`, `chi2inv`, `stdtr`, `stdtrit` | via `grad` (IFT) | jax#2399/#5350/#20358; chi-squared at real dof to 2000; Student-t ν ∈ [0.2, 200] via dedicated slice tables |
|
|
45
|
+
| von Mises | — | `vonmises_cdf/icdf` | via `grad` | κ ∈ [0, 50] |
|
|
46
|
+
| erf family | — | `dawsn`, `erfcx` | — (no params) | recent jax ships both; kept as a worked parameter-free recipe (plain functions, so not `bake` inputs: bake takes Recipe instances) |
|
|
47
|
+
| Lambert W | — | `lambertw(x, k)` | — | both real branches, jax#13680 |
|
|
48
|
+
|
|
49
|
+
Plus the generic core (`fit`, `ChebSeries`, `PiecewiseCheb`) and bake emitters
|
|
50
|
+
(`chebax.bake.jax_module`, `chebax.bake.xsf_header`: self-contained pure-jax
|
|
51
|
+
modules and C++17 headers from an instance; `dtype="float"` emits an f32
|
|
52
|
+
CUDA-shaped kernel body, and `truncate_tol=1e-7` drops the converged f64
|
|
53
|
+
tail — measured 1.4–2.2x faster f32 evaluation, `experiments/12`).
|
|
54
|
+
|
|
55
|
+
For pymc users: `import chebax.pytensor` (installs with
|
|
56
|
+
`pip install chebax[pytensor]`) registers JAX-backend lowerings for the
|
|
57
|
+
pytensor ops that otherwise need tfp-nightly under
|
|
58
|
+
`pm.sample(nuts_sampler="numpyro"|"blackjax")` — betaincinv, gammaincinv,
|
|
59
|
+
gammainccinv, erfcx, erfcinv, ive, kve — and adds betainc gradients in its
|
|
60
|
+
shape parameters, so censored or truncated StudentT/Beta likelihoods with a
|
|
61
|
+
latent scalar shape sample instead of raising. Shape parameters must be
|
|
62
|
+
scalar (batched ones fall back to tfp or fail loudly); out-of-domain values
|
|
63
|
+
return nan, never silently wrong numbers. `gammainccinv` takes its
|
|
64
|
+
upper-tail probability directly (it used to be solved at `1 - p`, which
|
|
65
|
+
collapsed everything below ~1e-17 onto one point), so both tails resolve to
|
|
66
|
+
the same depth.
|
|
67
|
+
|
|
68
|
+
For numpyro users: `from chebax.numpyro import TruncatedGamma, TruncatedBeta,
|
|
69
|
+
TruncatedStudentT` (installs with `pip install chebax[numpyro]`) gives the
|
|
70
|
+
truncated distributions numpyro's location-scale machinery cannot cover
|
|
71
|
+
(numpyro#969, numpyro#1365): full Distribution classes with reparameterized
|
|
72
|
+
inverse-CDF sampling (`has_rsample`), truncation-normalized `log_prob`, and
|
|
73
|
+
gradients through every parameter including the shapes, so NUTS and SVI work
|
|
74
|
+
with a latent concentration or df. The truncation is normalized in log space,
|
|
75
|
+
so an interval whose two ends land in the same tail still works: `Gamma(3)`
|
|
76
|
+
truncated to [50, 51] has mass ~3e-20 and returns an ordinary density.
|
|
77
|
+
Shape parameters are uniform per call; domain boxes are in the module
|
|
78
|
+
docstring and checked at construction.
|
|
79
|
+
|
|
80
|
+
`notebooks/` holds themed, executed walkthroughs: why Chebyshev nodes work,
|
|
81
|
+
the Bessel family with a learnable Matérn kernel, differentiable quantiles,
|
|
82
|
+
truncated and circular distributions in numpyro, Gaussian tails / Lambert W /
|
|
83
|
+
binomial reliability, copulas, and baking. The numpyro and plotting
|
|
84
|
+
dependencies install with `pip install chebax[examples]`. `examples/` holds
|
|
85
|
+
plain scripts: a few of the notebook workflows in copy-paste form, plus
|
|
86
|
+
integration examples with no notebook counterpart (a GIG distribution for
|
|
87
|
+
efax, learned Matérn smoothness, truncated sampling, `chebax.numpyro` with
|
|
88
|
+
latent shape parameters under NUTS).
|
|
89
|
+
|
|
90
|
+
Parameters come first, evaluation point last (opposite of scipy). Orders/shapes
|
|
91
|
+
are uniform per call; per-element parameter arrays are out of scope by design.
|
|
92
|
+
The middle ground is `chebax.pergroup(fn, group_idx)`: a static integer array
|
|
93
|
+
assigns each element to a group, each group gets its own (traceable) parameter
|
|
94
|
+
set — one group per chain, mixture component, or plate level.
|
|
95
|
+
Accuracy is validated against mpmath at 40 dps — measured worst cases and the
|
|
96
|
+
per-family error metrics are in each test file's docstring; tables regenerate
|
|
97
|
+
bit-for-bit from checked-in generators.
|
|
98
|
+
|
|
99
|
+
## Docs
|
|
100
|
+
|
|
101
|
+
`docs/adding-a-recipe.md` is the workflow for new functions;
|
|
102
|
+
`docs/increments.md` the design log of every recipe (what was measured, which
|
|
103
|
+
traps were found); `PROJECT.md` the plan and evidence; `CLAUDE.md` how to work
|
|
104
|
+
in the repo. Grown out of a private research project on GPU Bessel functions;
|
|
105
|
+
its two load-bearing measurements are reproduced here in `experiments/`.
|
|
106
|
+
BSD-3-Clause.
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=77"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "chebax"
|
|
7
|
+
version = "0.3.0"
|
|
8
|
+
description = "Differentiable Chebyshev approximants for special functions, aimed at JAX and GPUs"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "BSD-3-Clause"
|
|
11
|
+
authors = [{ name = "Andres Hernandez", email = "andres.hernandez.deml@gmail.com" }]
|
|
12
|
+
requires-python = ">=3.11"
|
|
13
|
+
# jax floor = the oldest version the full suite passes on (CI pins a job to it)
|
|
14
|
+
dependencies = ["jax>=0.5.3", "numpy>=1.26"]
|
|
15
|
+
|
|
16
|
+
[project.urls]
|
|
17
|
+
Repository = "https://github.com/andreshernandez-spec/chebax"
|
|
18
|
+
|
|
19
|
+
[project.optional-dependencies]
|
|
20
|
+
# tables record the mpmath version that built them; regeneration is verified
|
|
21
|
+
# byte for byte, so the build dependency is pinned to the recorded minor
|
|
22
|
+
gen = ["mpmath>=1.4.1,<1.5"]
|
|
23
|
+
test = ["pytest", "scipy"]
|
|
24
|
+
examples = ["matplotlib", "numpyro"]
|
|
25
|
+
numpyro = ["numpyro>=0.15"]
|
|
26
|
+
pytensor = ["pytensor>=3.2"]
|
|
27
|
+
|
|
28
|
+
[tool.setuptools.packages.find]
|
|
29
|
+
where = ["src"]
|
|
30
|
+
|
|
31
|
+
[tool.pytest.ini_options]
|
|
32
|
+
testpaths = ["tests"]
|
|
33
|
+
markers = [
|
|
34
|
+
"slow: table regeneration and other multi-minute tests",
|
|
35
|
+
]
|
|
@@ -0,0 +1,54 @@
|
|
|
1
|
+
"""chebax: differentiable Chebyshev approximants for special functions.
|
|
2
|
+
|
|
3
|
+
Full float64 accuracy needs jax's x64 mode (jax.config.update("jax_enable_x64", True)).
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
from chebax._src.generate import fit
|
|
7
|
+
from chebax._src.recipes.besseli import (besseli, besseli_dnu, besseli_fn,
|
|
8
|
+
besseli_ratio)
|
|
9
|
+
from chebax._src.recipes.erf_family import dawsn, erfcx
|
|
10
|
+
from chebax._src.recipes.lambertw import lambertw
|
|
11
|
+
from chebax._src.recipes.matern import matern
|
|
12
|
+
from chebax._src.recipes.spherical import spherical_jn, spherical_yn
|
|
13
|
+
from chebax._src.recipes.besselj import besselj, besselj_dnu
|
|
14
|
+
from chebax._src.recipes.besselk import (besselk, besselk_dnu, besselk_fn,
|
|
15
|
+
log_besselk_fn)
|
|
16
|
+
from chebax._src.recipes.bessely import bessely, bessely_dnu
|
|
17
|
+
from chebax._src.recipes.betainc import betainc, betainc_fn, log_betainc_fn
|
|
18
|
+
from chebax._src.recipes.gammainc import (gammainc, gammainc_fn, gammaincc,
|
|
19
|
+
gammaincc_fn, log_gammainc_fn,
|
|
20
|
+
log_gammaincc_fn)
|
|
21
|
+
from chebax._src.recipes.hyp1f1 import hyp1f1, hyp1f1_fn, log_hyp1f1_fn
|
|
22
|
+
from chebax._src.recipes.quantiles import (betaincinv, chi2inv,
|
|
23
|
+
gammainccinv, gammaincinv,
|
|
24
|
+
log_stdtr, log_stdtr_sf, stdtr,
|
|
25
|
+
stdtrit)
|
|
26
|
+
from chebax._src.recipes.vonmises import vonmises_cdf, vonmises_icdf
|
|
27
|
+
from chebax._src.pergroup import pergroup
|
|
28
|
+
from chebax._src.series import ChebSeries, PiecewiseCheb
|
|
29
|
+
|
|
30
|
+
# single-sourced from pyproject.toml via package metadata
|
|
31
|
+
try:
|
|
32
|
+
from importlib.metadata import version as _pkg_version
|
|
33
|
+
__version__ = _pkg_version("chebax")
|
|
34
|
+
except Exception: # not installed (e.g. running from a source checkout)
|
|
35
|
+
__version__ = "0+unknown"
|
|
36
|
+
__all__ = ["ChebSeries", "PiecewiseCheb", "besseli", "besseli_dnu", "besseli_fn",
|
|
37
|
+
"besseli_ratio", "besselj", "besselj_dnu", "besselk", "besselk_dnu", "besselk_fn",
|
|
38
|
+
"log_besselk_fn",
|
|
39
|
+
"bessely", "bessely_dnu", "betainc", "betainc_fn", "betaincinv",
|
|
40
|
+
"chi2inv", "dawsn", "erfcx", "fit", "gammainc", "gammainc_fn", "gammaincc",
|
|
41
|
+
"gammaincc_fn", "gammainccinv", "gammaincinv", "hyp1f1", "hyp1f1_fn", "lambertw",
|
|
42
|
+
"log_betainc_fn", "log_hyp1f1_fn",
|
|
43
|
+
"matern",
|
|
44
|
+
"log_gammainc_fn", "log_gammaincc_fn", "log_stdtr", "log_stdtr_sf",
|
|
45
|
+
"pergroup",
|
|
46
|
+
"spherical_jn",
|
|
47
|
+
"spherical_yn", "stdtr", "stdtrit", "vonmises_cdf", "vonmises_icdf",
|
|
48
|
+
"__version__"]
|
|
49
|
+
|
|
50
|
+
# submodule re-export so `import chebax; chebax.bake` works as documented
|
|
51
|
+
# (bake reads __version__, so this import must stay below its definition)
|
|
52
|
+
from chebax import bake # noqa: E402,F401
|
|
53
|
+
|
|
54
|
+
__all__.append("bake")
|
|
@@ -0,0 +1,280 @@
|
|
|
1
|
+
"""Build-time fitting: sample at Chebyshev points, DCT to coefficients.
|
|
2
|
+
|
|
3
|
+
The float64 path uses numpy. Passing dps switches sampling and the DCT to
|
|
4
|
+
mpmath at that precision; mpmath is imported inside that branch only (the
|
|
5
|
+
runtime-imports rule in CLAUDE.md).
|
|
6
|
+
|
|
7
|
+
Adaptive fitting doubles the node count until the coefficient tail sits at
|
|
8
|
+
the fit's own noise floor, then truncates there. A converged-looking tail
|
|
9
|
+
on one grid is never trusted by itself: on n first-kind points T_{2n-k}
|
|
10
|
+
samples exactly as -T_k, so an aliased fit can look converged (T_34 on 17
|
|
11
|
+
nodes reproduces -T_0 to the last bit). Every candidate the adaptive path
|
|
12
|
+
returns, the all-zero fit included, is first validated against fresh
|
|
13
|
+
samples on two finer grids (see _validated). A fixed deg= is not: there an
|
|
14
|
+
unresolved f is a legitimate request, so under-resolution cannot be told
|
|
15
|
+
from aliasing.
|
|
16
|
+
|
|
17
|
+
The noise floor is estimated from the trailing sixteenth of the fit's own
|
|
18
|
+
coefficients rather than assumed: float64 sampling noise scales with
|
|
19
|
+
the function's derivative (measured: ~1e-16 relative for exp, ~2.5e-15 for
|
|
20
|
+
sin(50x)), so a fixed tolerance either fails to converge on oscillatory
|
|
21
|
+
functions or keeps hundreds of noise coefficients.
|
|
22
|
+
"""
|
|
23
|
+
|
|
24
|
+
import numpy as np
|
|
25
|
+
|
|
26
|
+
from chebax._src import algorithms
|
|
27
|
+
from chebax._src.series import ChebSeries, PiecewiseCheb
|
|
28
|
+
|
|
29
|
+
# ~2*eps: the practical floor of the float64 build path. Use dps= to go below.
|
|
30
|
+
_DEFAULT_TOL = 5e-16
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def fit(f, domain=(-1.0, 1.0), *, deg=None, tol=None, breaks=None, dps=None, max_deg=256):
|
|
34
|
+
"""Chebyshev interpolant of f, as a ChebSeries (or PiecewiseCheb if breaks given).
|
|
35
|
+
|
|
36
|
+
f: callable on the domain. Vectorized over numpy arrays if possible;
|
|
37
|
+
with dps set it must accept mpmath floats instead.
|
|
38
|
+
domain: (a, b), finite with a < b. Ignored when breaks is given.
|
|
39
|
+
deg: fixed degree, at most max_deg; otherwise the degree adapts until
|
|
40
|
+
tol is met. A fixed degree is taken as an instruction: the
|
|
41
|
+
interpolant is returned as asked, with no convergence check and
|
|
42
|
+
no aliasing check (fit(T_29, deg=5) returns T_5, which is what
|
|
43
|
+
interpolating T_29 at 6 points gives). Leave deg unset for a
|
|
44
|
+
validated fit.
|
|
45
|
+
tol: target for the relative coefficient tail (default 5e-16), positive
|
|
46
|
+
and finite. The effective chop level is max(tol, the fit's
|
|
47
|
+
measured noise floor); float64 sampling noise can sit above tol
|
|
48
|
+
for functions with large derivatives. Build with dps to control
|
|
49
|
+
the floor.
|
|
50
|
+
breaks: strictly increasing finite knots, at least two, endpoints
|
|
51
|
+
included; one series per segment, each fit independently, padded
|
|
52
|
+
to a common degree.
|
|
53
|
+
dps: mpmath working precision for sampling and the DCT. Derivatives of
|
|
54
|
+
a float64-built fit sit at a ~deg^2 * eps floor (differentiation
|
|
55
|
+
amplifies sample noise); build with dps when that matters.
|
|
56
|
+
max_deg: hard cap on the returned degree. Reaching it without a validated
|
|
57
|
+
converged fit raises: f is not smooth enough on the interval
|
|
58
|
+
(split the domain) or genuinely needs a higher degree.
|
|
59
|
+
"""
|
|
60
|
+
max_deg = _check_int("max_deg", max_deg, 1)
|
|
61
|
+
if deg is not None:
|
|
62
|
+
deg = _check_int("deg", deg, 0)
|
|
63
|
+
if deg > max_deg:
|
|
64
|
+
raise ValueError(f"deg={deg} is above the max_deg={max_deg} cap")
|
|
65
|
+
if dps is not None:
|
|
66
|
+
dps = _check_int("dps", dps, 1)
|
|
67
|
+
if tol is None:
|
|
68
|
+
tol = _DEFAULT_TOL
|
|
69
|
+
else:
|
|
70
|
+
tol = float(tol)
|
|
71
|
+
if not np.isfinite(tol) or tol <= 0.0:
|
|
72
|
+
raise ValueError(f"tol must be positive and finite, got {tol}")
|
|
73
|
+
if breaks is not None:
|
|
74
|
+
br = [float(t) for t in breaks]
|
|
75
|
+
if len(br) < 2:
|
|
76
|
+
raise ValueError("breaks needs at least two knots")
|
|
77
|
+
if not all(np.isfinite(t) for t in br):
|
|
78
|
+
raise ValueError("breaks must be finite")
|
|
79
|
+
if not all(u < v for u, v in zip(br, br[1:])):
|
|
80
|
+
raise ValueError("breaks must be strictly increasing")
|
|
81
|
+
cs = [_fit_interval(f, br[i], br[i + 1], deg, tol, max_deg, dps)
|
|
82
|
+
for i in range(len(br) - 1)]
|
|
83
|
+
n = max(c.size for c in cs)
|
|
84
|
+
coef = np.zeros((len(cs), n))
|
|
85
|
+
for i, c in enumerate(cs):
|
|
86
|
+
coef[i, : c.size] = c
|
|
87
|
+
return PiecewiseCheb(coef, br)
|
|
88
|
+
try:
|
|
89
|
+
n_dom = len(domain)
|
|
90
|
+
except TypeError:
|
|
91
|
+
n_dom = -1
|
|
92
|
+
if n_dom != 2:
|
|
93
|
+
# (1,) used to reach an IndexError and a 3-tuple silently dropped
|
|
94
|
+
# its third element (review, 2026-08-02)
|
|
95
|
+
raise ValueError(f"domain must be a pair (a, b), got {domain!r}")
|
|
96
|
+
a, b = float(domain[0]), float(domain[1])
|
|
97
|
+
if not (np.isfinite(a) and np.isfinite(b) and a < b):
|
|
98
|
+
raise ValueError(f"domain must be finite with a < b, got ({a}, {b})")
|
|
99
|
+
return ChebSeries(_fit_interval(f, a, b, deg, tol, max_deg, dps), (a, b))
|
|
100
|
+
|
|
101
|
+
|
|
102
|
+
def _check_int(name, v, low):
|
|
103
|
+
kind = "a nonnegative integer" if low == 0 else "a positive integer"
|
|
104
|
+
if isinstance(v, bool) or not isinstance(v, (int, np.integer)) or v < low:
|
|
105
|
+
raise ValueError(f"{name} must be {kind}, got {v!r}")
|
|
106
|
+
return int(v)
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def _fit_interval(f, a, b, deg, tol, max_deg, dps):
|
|
110
|
+
dct, sample_unit = (_dct_mp(f, a, b, dps) if dps is not None
|
|
111
|
+
else _dct_np(f, a, b))
|
|
112
|
+
if deg is not None:
|
|
113
|
+
return dct(deg + 1)
|
|
114
|
+
return _adaptive(dct, sample_unit, tol, max_deg, (a, b))
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
_NOISE_CEIL = 1e-8 # a window this close to the function scale is signal
|
|
118
|
+
_PHI = 0.6180339887498949 # (sqrt(5) - 1)/2, irrational: see _offgrid
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
def _chop(c, tol):
|
|
122
|
+
"""Truncate at the noise level; returns (candidate, absolute level,
|
|
123
|
+
measured plateau relative to the scale, inf when there is none).
|
|
124
|
+
|
|
125
|
+
The level is max(tol, noise floor), where the noise floor is read off
|
|
126
|
+
the last sixteenth of the coefficients: for a still-decaying tail those
|
|
127
|
+
sit below tol and the requested tol governs (measured: exp's last
|
|
128
|
+
coefficients decay through 1e-16); for a converged fit they are the
|
|
129
|
+
sampling-noise plateau (sin(50x): ~5e-16 relative at n=129, edge-
|
|
130
|
+
singular sampling can push higher) and doubling the largest of them
|
|
131
|
+
clears the plateau's own scatter.
|
|
132
|
+
|
|
133
|
+
That window only holds noise if the fit has room to spare, and it does
|
|
134
|
+
not when the top mode lands inside it: fitting T20 on 21 points puts
|
|
135
|
+
c[20] = 1 in the last sixteenth, which read as a noise floor of 2 and
|
|
136
|
+
chopped the exact fit down to one coefficient (with max_deg = 20 there
|
|
137
|
+
is no doubling left, so `fit(T20, max_deg=20)` then raised). A real
|
|
138
|
+
plateau sits at the sampling floor, decades under the function scale
|
|
139
|
+
(the settled test below wants 1e-10 relative before it will believe
|
|
140
|
+
one), so a window within _NOISE_CEIL of the scale is signal and the
|
|
141
|
+
requested tol governs instead."""
|
|
142
|
+
scale = np.max(np.abs(c))
|
|
143
|
+
win = float(np.max(np.abs(c[-max(2, c.size // 16):])))
|
|
144
|
+
measured = win <= _NOISE_CEIL * scale
|
|
145
|
+
noise = 2.0 * win if measured else 0.0
|
|
146
|
+
level = max(tol * scale, noise)
|
|
147
|
+
keep = np.nonzero(np.abs(c) > level)[0]
|
|
148
|
+
return ((c[: keep[-1] + 1] if keep.size else c[:1]), level,
|
|
149
|
+
noise / scale if measured else np.inf)
|
|
150
|
+
|
|
151
|
+
|
|
152
|
+
def _offgrid(m):
|
|
153
|
+
"""m points cos(pi (i + PHI) / m), PHI irrational.
|
|
154
|
+
|
|
155
|
+
NOT the nodes of any Chebyshev grid, which is the whole point. On m
|
|
156
|
+
first-kind points T_K samples as +-T_j with (j, sign) fixed by
|
|
157
|
+
K mod 4m, so any set of harmonically related grids has whole families
|
|
158
|
+
of modes it cannot see: n, 2n and 4n all read T_{16n} as +1, which is
|
|
159
|
+
how `fit(T272)` used to come back as the constant 1 with sup error 2
|
|
160
|
+
(found in review, 2026-08-02). Here an alias would need
|
|
161
|
+
cos(K theta_i) = +-cos(j theta_i) at every point, i.e.
|
|
162
|
+
(K -+ j)(i + PHI)/m in 2Z for every i; differencing over i forces
|
|
163
|
+
K -+ j = 2 m q, and then 2 q PHI in Z, which for irrational PHI needs
|
|
164
|
+
q = 0. So the only mode that can hide is the one already in the
|
|
165
|
+
candidate."""
|
|
166
|
+
i = np.arange(m)
|
|
167
|
+
return np.cos(np.pi * (i + _PHI) / m)
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
def _validated(cand, level, sample_unit, n):
|
|
171
|
+
"""Check the candidate against fresh samples away from the fit grid.
|
|
172
|
+
|
|
173
|
+
A finite point set can never prove a fit: any f agreeing with cand at
|
|
174
|
+
the test points passes, and there are always such f. This is a
|
|
175
|
+
gross-error detector, nothing more. What it must not have is a
|
|
176
|
+
STRUCTURAL blind spot, and harmonically related Chebyshev grids have
|
|
177
|
+
one (see _offgrid), so the off-grid sets carry the argument and the 2n
|
|
178
|
+
grid is kept only because it is cheap and catches the ordinary
|
|
179
|
+
mistakes. Sizes are fixed by the fitting resolution n rather than by
|
|
180
|
+
the candidate's own size, which could otherwise land back on the
|
|
181
|
+
fitting grid.
|
|
182
|
+
|
|
183
|
+
An aliased fit is off by O(scale) here, ~1e15x the chop level, where
|
|
184
|
+
an honest one is at ~1x; the 1000x margin is the gap between those,
|
|
185
|
+
not an accuracy contract. The second term is the Clenshaw evaluation
|
|
186
|
+
floor, which grows with the degree and not just with sum|c|: a
|
|
187
|
+
degree-544 candidate lands 6e-13 off its own exact samples, over the
|
|
188
|
+
1000x chop margin, so an honest high-degree fit was being rejected
|
|
189
|
+
(measured on T544 once the off-grid sets started catching it)."""
|
|
190
|
+
allow = max(1000.0 * level,
|
|
191
|
+
100.0 * np.finfo(float).eps * cand.size
|
|
192
|
+
* float(np.sum(np.abs(cand))))
|
|
193
|
+
for t in (algorithms.chebpts(2 * n), _offgrid(2 * n + 1),
|
|
194
|
+
_offgrid(4 * n + 1)):
|
|
195
|
+
err = np.max(np.abs(algorithms.chebval(t, cand) - sample_unit(t)))
|
|
196
|
+
if err > allow:
|
|
197
|
+
return False
|
|
198
|
+
return True
|
|
199
|
+
|
|
200
|
+
|
|
201
|
+
def _adaptive(dct, sample_unit, tol, max_deg, interval):
|
|
202
|
+
n = min(17, max_deg + 1)
|
|
203
|
+
prev_rel = np.inf
|
|
204
|
+
while True:
|
|
205
|
+
c = dct(n)
|
|
206
|
+
scale = float(np.max(np.abs(c)))
|
|
207
|
+
if scale == 0.0:
|
|
208
|
+
# all-zero samples prove nothing off the grid, T_n vanishes on
|
|
209
|
+
# chebpts(n), so the zero series is validated like any other
|
|
210
|
+
cand, level, settled = np.zeros(1), 0.0, True
|
|
211
|
+
else:
|
|
212
|
+
cand, level, rel = _chop(c, tol)
|
|
213
|
+
# The tail has settled if the chop level reached the requested
|
|
214
|
+
# tol, or if the MEASURED plateau stopped improving under node
|
|
215
|
+
# doubling while being far below the function scale: that is the
|
|
216
|
+
# sampling-noise plateau, and more nodes cannot beat it. A tail
|
|
217
|
+
# that is small but still shrinking (a segment one doubling short
|
|
218
|
+
# of converged) must escalate, not chop. rel is inf when the
|
|
219
|
+
# window held signal rather than noise, so a resolution with no
|
|
220
|
+
# plateau to report cannot seed the comparison either way.
|
|
221
|
+
settled = level <= 2.0 * tol * scale or (rel <= 1e-10 and rel >= prev_rel / 4.0)
|
|
222
|
+
prev_rel = rel
|
|
223
|
+
at_cap = n >= max_deg + 1
|
|
224
|
+
# A candidate holding nearly every coefficient of its grid has not
|
|
225
|
+
# shown a resolved tail, so double instead. At the cap there is no
|
|
226
|
+
# doubling left and validation is the only evidence available.
|
|
227
|
+
chopped = cand.size <= c.size - max(2, c.size // 8)
|
|
228
|
+
if settled and (chopped or at_cap) and _validated(cand, level, sample_unit, n):
|
|
229
|
+
return cand
|
|
230
|
+
if at_cap:
|
|
231
|
+
raise ValueError(
|
|
232
|
+
f"no validated convergence to tol={tol:g} by degree {max_deg} on "
|
|
233
|
+
f"{interval}; f may not be smooth there (split the domain), may "
|
|
234
|
+
"need a higher max_deg, or float64 sampling noise may sit above "
|
|
235
|
+
"tol (build with dps)")
|
|
236
|
+
n = min(2 * n - 1, max_deg + 1)
|
|
237
|
+
|
|
238
|
+
|
|
239
|
+
def _dct_np(f, a, b):
|
|
240
|
+
def sample_unit(t):
|
|
241
|
+
xk = a + (b - a) * (t + 1) / 2
|
|
242
|
+
return _sample(f, xk)
|
|
243
|
+
|
|
244
|
+
def dct(n):
|
|
245
|
+
return algorithms.chebfit_dct(sample_unit(algorithms.chebpts(n)))
|
|
246
|
+
return dct, sample_unit
|
|
247
|
+
|
|
248
|
+
|
|
249
|
+
def _sample(f, xk):
|
|
250
|
+
# scalar-only callables (wrapped special functions) are common; fall back.
|
|
251
|
+
try:
|
|
252
|
+
y = np.asarray(f(xk), dtype=float)
|
|
253
|
+
if y.shape == xk.shape:
|
|
254
|
+
return y
|
|
255
|
+
except (TypeError, ValueError):
|
|
256
|
+
pass
|
|
257
|
+
return np.array([float(f(x)) for x in xk])
|
|
258
|
+
|
|
259
|
+
|
|
260
|
+
def _dct_mp(f, a, b, dps):
|
|
261
|
+
import mpmath as mp
|
|
262
|
+
|
|
263
|
+
def sample_unit(t):
|
|
264
|
+
with mp.workdps(dps):
|
|
265
|
+
mid = (mp.mpf(b) + mp.mpf(a)) / 2
|
|
266
|
+
half = (mp.mpf(b) - mp.mpf(a)) / 2
|
|
267
|
+
return np.array([float(f(mid + half * mp.mpf(float(tk)))) for tk in t])
|
|
268
|
+
|
|
269
|
+
def dct(n):
|
|
270
|
+
with mp.workdps(dps):
|
|
271
|
+
mid = (mp.mpf(b) + mp.mpf(a)) / 2
|
|
272
|
+
half = (mp.mpf(b) - mp.mpf(a)) / 2
|
|
273
|
+
fk = [f(mid + half * mp.cos(mp.pi * (2 * i + 1) / (2 * n))) for i in range(n)]
|
|
274
|
+
c = []
|
|
275
|
+
for j in range(n):
|
|
276
|
+
s = mp.fsum(fk[i] * mp.cos(mp.pi * j * (2 * i + 1) / (2 * n)) for i in range(n))
|
|
277
|
+
c.append(2 * s / n)
|
|
278
|
+
c[0] /= 2
|
|
279
|
+
return np.array([float(x) for x in c])
|
|
280
|
+
return dct, sample_unit
|