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.
Files changed (100) hide show
  1. chebax-0.3.0/PKG-INFO +132 -0
  2. chebax-0.3.0/README.md +106 -0
  3. chebax-0.3.0/pyproject.toml +35 -0
  4. chebax-0.3.0/src/chebax/__init__.py +54 -0
  5. chebax-0.3.0/src/chebax/_src/generate.py +280 -0
  6. chebax-0.3.0/src/chebax/_src/pergroup.py +113 -0
  7. chebax-0.3.0/src/chebax/_src/pytree.py +87 -0
  8. chebax-0.3.0/src/chebax/_src/recipes/_common.py +134 -0
  9. chebax-0.3.0/src/chebax/_src/recipes/_gen_common.py +121 -0
  10. chebax-0.3.0/src/chebax/_src/recipes/besseli.py +376 -0
  11. chebax-0.3.0/src/chebax/_src/recipes/besseli_gen.py +84 -0
  12. chebax-0.3.0/src/chebax/_src/recipes/besseli_table.py +138 -0
  13. chebax-0.3.0/src/chebax/_src/recipes/besselj.py +251 -0
  14. chebax-0.3.0/src/chebax/_src/recipes/besselj_gen.py +146 -0
  15. chebax-0.3.0/src/chebax/_src/recipes/besselj_table.py +39 -0
  16. chebax-0.3.0/src/chebax/_src/recipes/besselj_table_ext.py +87 -0
  17. chebax-0.3.0/src/chebax/_src/recipes/besselk.py +209 -0
  18. chebax-0.3.0/src/chebax/_src/recipes/besselk_gen.py +107 -0
  19. chebax-0.3.0/src/chebax/_src/recipes/besselk_table.py +208 -0
  20. chebax-0.3.0/src/chebax/_src/recipes/bessely.py +185 -0
  21. chebax-0.3.0/src/chebax/_src/recipes/bessely_gen.py +103 -0
  22. chebax-0.3.0/src/chebax/_src/recipes/bessely_table.py +284 -0
  23. chebax-0.3.0/src/chebax/_src/recipes/betainc.py +255 -0
  24. chebax-0.3.0/src/chebax/_src/recipes/betainc_gen.py +78 -0
  25. chebax-0.3.0/src/chebax/_src/recipes/betainc_table.py +1791 -0
  26. chebax-0.3.0/src/chebax/_src/recipes/betainc_wide_gen.py +110 -0
  27. chebax-0.3.0/src/chebax/_src/recipes/betainc_wide_table.py +11390 -0
  28. chebax-0.3.0/src/chebax/_src/recipes/erf_family.py +75 -0
  29. chebax-0.3.0/src/chebax/_src/recipes/erf_gen.py +69 -0
  30. chebax-0.3.0/src/chebax/_src/recipes/erf_table.py +51 -0
  31. chebax-0.3.0/src/chebax/_src/recipes/gammainc.py +259 -0
  32. chebax-0.3.0/src/chebax/_src/recipes/gammainc_gen.py +99 -0
  33. chebax-0.3.0/src/chebax/_src/recipes/gammainc_large.py +178 -0
  34. chebax-0.3.0/src/chebax/_src/recipes/gammainc_large_gen.py +172 -0
  35. chebax-0.3.0/src/chebax/_src/recipes/gammainc_large_table.py +90 -0
  36. chebax-0.3.0/src/chebax/_src/recipes/gammainc_table.py +76 -0
  37. chebax-0.3.0/src/chebax/_src/recipes/hyp1f1.py +176 -0
  38. chebax-0.3.0/src/chebax/_src/recipes/hyp1f1_gen.py +139 -0
  39. chebax-0.3.0/src/chebax/_src/recipes/hyp1f1_table.py +6002 -0
  40. chebax-0.3.0/src/chebax/_src/recipes/lambertw.py +106 -0
  41. chebax-0.3.0/src/chebax/_src/recipes/matern.py +55 -0
  42. chebax-0.3.0/src/chebax/_src/recipes/quantiles.py +788 -0
  43. chebax-0.3.0/src/chebax/_src/recipes/spherical.py +100 -0
  44. chebax-0.3.0/src/chebax/_src/recipes/stdtr_gen.py +74 -0
  45. chebax-0.3.0/src/chebax/_src/recipes/stdtr_table.py +297 -0
  46. chebax-0.3.0/src/chebax/_src/recipes/vonmises.py +182 -0
  47. chebax-0.3.0/src/chebax/_src/recipes/vonmises_gen.py +68 -0
  48. chebax-0.3.0/src/chebax/_src/recipes/vonmises_table.py +127 -0
  49. chebax-0.3.0/src/chebax/_src/series.py +353 -0
  50. chebax-0.3.0/src/chebax/bake/__init__.py +319 -0
  51. chebax-0.3.0/src/chebax/bake/_jaxpr_emit.py +407 -0
  52. chebax-0.3.0/src/chebax/domains.py +59 -0
  53. chebax-0.3.0/src/chebax/numpyro.py +391 -0
  54. chebax-0.3.0/src/chebax/pytensor.py +196 -0
  55. chebax-0.3.0/src/chebax.egg-info/PKG-INFO +132 -0
  56. chebax-0.3.0/src/chebax.egg-info/SOURCES.txt +85 -0
  57. chebax-0.3.0/src/chebax.egg-info/requires.txt +19 -0
  58. chebax-0.3.0/tests/test_bake.py +582 -0
  59. chebax-0.3.0/tests/test_besseli.py +292 -0
  60. {chebax-0.1.0.dev0 → chebax-0.3.0}/tests/test_besselj.py +38 -7
  61. chebax-0.3.0/tests/test_besselj_m3.py +186 -0
  62. chebax-0.3.0/tests/test_besselk.py +308 -0
  63. chebax-0.3.0/tests/test_bessely.py +87 -0
  64. chebax-0.3.0/tests/test_betainc.py +234 -0
  65. chebax-0.3.0/tests/test_contracts.py +99 -0
  66. chebax-0.3.0/tests/test_erf_family.py +278 -0
  67. chebax-0.3.0/tests/test_gammainc.py +282 -0
  68. chebax-0.3.0/tests/test_gammainc_large.py +259 -0
  69. chebax-0.3.0/tests/test_hyp1f1.py +173 -0
  70. chebax-0.3.0/tests/test_int_params.py +65 -0
  71. chebax-0.3.0/tests/test_lambertw.py +80 -0
  72. chebax-0.3.0/tests/test_m1.py +465 -0
  73. chebax-0.3.0/tests/test_matern.py +108 -0
  74. chebax-0.3.0/tests/test_numpyro_module.py +502 -0
  75. chebax-0.3.0/tests/test_pergroup.py +199 -0
  76. chebax-0.3.0/tests/test_pytensor_plugin.py +142 -0
  77. chebax-0.3.0/tests/test_quantiles.py +691 -0
  78. chebax-0.3.0/tests/test_spherical.py +59 -0
  79. chebax-0.3.0/tests/test_vonmises.py +241 -0
  80. chebax-0.3.0/tests/test_x64_off.py +129 -0
  81. chebax-0.1.0.dev0/PKG-INFO +0 -35
  82. chebax-0.1.0.dev0/README.md +0 -17
  83. chebax-0.1.0.dev0/pyproject.toml +0 -26
  84. chebax-0.1.0.dev0/src/chebax/__init__.py +0 -11
  85. chebax-0.1.0.dev0/src/chebax/_src/generate.py +0 -106
  86. chebax-0.1.0.dev0/src/chebax/_src/recipes/besselj.py +0 -138
  87. chebax-0.1.0.dev0/src/chebax/_src/recipes/besselj_gen.py +0 -76
  88. chebax-0.1.0.dev0/src/chebax/_src/recipes/besselj_table.py +0 -37
  89. chebax-0.1.0.dev0/src/chebax/_src/series.py +0 -195
  90. chebax-0.1.0.dev0/src/chebax.egg-info/PKG-INFO +0 -35
  91. chebax-0.1.0.dev0/src/chebax.egg-info/SOURCES.txt +0 -19
  92. chebax-0.1.0.dev0/src/chebax.egg-info/requires.txt +0 -8
  93. chebax-0.1.0.dev0/tests/test_m1.py +0 -212
  94. {chebax-0.1.0.dev0 → chebax-0.3.0}/LICENSE +0 -0
  95. {chebax-0.1.0.dev0 → chebax-0.3.0}/setup.cfg +0 -0
  96. {chebax-0.1.0.dev0 → chebax-0.3.0}/src/chebax/_src/__init__.py +0 -0
  97. {chebax-0.1.0.dev0 → chebax-0.3.0}/src/chebax/_src/algorithms.py +0 -0
  98. {chebax-0.1.0.dev0 → chebax-0.3.0}/src/chebax/_src/recipes/__init__.py +0 -0
  99. {chebax-0.1.0.dev0 → chebax-0.3.0}/src/chebax.egg-info/dependency_links.txt +0 -0
  100. {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