diff --git a/.claude/skills/jaxsr/SKILL.md b/.claude/skills/jaxsr/SKILL.md index 4bbd7eb..5b1c908 100644 --- a/.claude/skills/jaxsr/SKILL.md +++ b/.claude/skills/jaxsr/SKILL.md @@ -227,6 +227,67 @@ next_pts = study.suggest_next(n_points=5, strategy="uncertainty") study.save("catalyst.jaxsr") ``` +### Additive (Boosting-Style) Symbolic Regression + +Use when the signal is a sum of several simple effects and you want many small +interpretable terms instead of one large expression. Fits `f(x) = c + Σ_k η_k · +g_k(x)` by stagewise residual fitting (analogous to gradient boosting with +symbolic weak learners). Reuses `fit_symbolic` for each term. + +```python +from jaxsr.additive import StagewiseSymbolicRegressor + +model = StagewiseSymbolicRegressor( + n_terms=10, # number of boosting stages (terms) + learning_rate=0.2, # shrinkage (used when refit_coefficients=False) + max_complexity=4, # max basis terms per stage — keep small + refit_coefficients=True, # re-solve all linear weights by OLS each stage + early_stopping=False, # stop on a validation split when it stops improving + validation_fraction=0.2, +) +model.fit(X, y) + +print(model) # pretty structural summary +model.expressions_ # per-term expression strings +model.intercept_, model.coefficients_ +model.predict(X_new) +model.to_expression() # single combined SymPy expression +model.save("additive.json") # JSON round-trip (models are NOT picklable) +loaded = StagewiseSymbolicRegressor.load("additive.json") +``` + +Notes: +- Prefer `refit_coefficients=True` for accuracy with squared error; keep + `max_complexity` small (2–4) to favour many simple terms. +- **Robust / quantile regression:** set `loss` to `"absolute_error"`, + `"huber"`, or `"quantile"` (or an instance like `QuantileLoss(0.9)` / + `HuberLoss(delta=2.0)`). These are fit by gradient boosting with a per-stage + line search; use `refit_coefficients=False` (OLS refit only applies to + squared error and is auto-disabled with a warning otherwise). Fit several + quantiles to build prediction intervals. +- **Structural uncertainty:** `bootstrap_additive(model, X, y, n_bootstrap=...)` + refits on resamples and returns `["inclusion_probabilities"]` (how often each + basis is selected — a posterior-inclusion-probability proxy) and `["models"]`; + pass those to `bootstrap_predict_additive(models, X_new)` for an ensemble + prediction interval. Diffuse probabilities (~0.5) flag that the data don't + determine one expression (common with collinear features). +- `include_transcendental`/`include_ratios` are off by default; if enabled, a + stage that would produce non-finite predictions falls back to a finite basis. +- **Compositional discovery (experimental):** `RecursiveSymbolicRegressor` + grows the basis library along the residual (composing unary funcs, products, + ratios of discovered terms) to reach compositions a flat library misses + (e.g. `exp(x0*x1)`). Deterministic FFX-style search; competitive with GP on + simple targets but won't match PySR/Operon on hard ones. Not serialisable + (composed bases are closures). +- `BackfittingSymbolicRegressor` (GAM-style: a fixed set of terms revised + across sweeps, warm-started from stagewise; squared error only) is available. + It is never worse than stagewise+refit on training and helps specifically + when `max_complexity` is small (single-basis terms) and features are + collinear — where greedy forward selection gets stuck and re-discovery + escapes it. With larger per-term budgets it matches stagewise+refit, so + prefer `StagewiseSymbolicRegressor` there. A Bayesian (BART/iBART) variant is + future work. + ## Quick Reference: CLI ```bash @@ -325,6 +386,11 @@ See `guides/rsm.md` for RSM designs, canonical analysis, and optimization. See `guides/active-learning.md` for acquisition functions and adaptive sampling. +### "One expression isn't enough / the signal is a sum of many effects" + +See `guides/additive.md` for boosting-style additive symbolic regression +(`StagewiseSymbolicRegressor`): fit residuals stagewise into many small terms. + ## Templates Ready-to-use scripts and notebook starters are in `templates/`: @@ -332,6 +398,7 @@ Ready-to-use scripts and notebook starters are in `templates/`: | Template | Use Case | |----------|----------| | `basic-regression.py` | Discover an equation from X, y data | +| `additive-regression.py` | Boosting-style additive SR: robust/quantile losses, bootstrap UQ, backfitting, recursive expansion | | `constrained-model.py` | Add physical constraints to model | | `doe-study.py` | Full DOE workflow from design to report | | `uncertainty-analysis.py` | Compare all UQ methods | diff --git a/.claude/skills/jaxsr/guides/additive.md b/.claude/skills/jaxsr/guides/additive.md new file mode 100644 index 0000000..838f582 --- /dev/null +++ b/.claude/skills/jaxsr/guides/additive.md @@ -0,0 +1,332 @@ +# Additive Symbolic Regression + +Additive symbolic regression fits a model as a **sum of small symbolic +expressions**: + +``` +f(x) = c + eta_1 * g_1(x) + eta_2 * g_2(x) + ... + eta_K * g_K(x) +``` + +where each `g_k(x)` is a small, interpretable symbolic expression discovered by +the existing JAXSR machinery. This is analogous to **gradient boosting**, except +each weak learner is a symbolic expression rather than a decision tree. + +The submodule lives in `jaxsr.additive`. + +## Three flavours of symbolic regression + +| Approach | What it does | Status | +|----------|--------------|--------| +| **Single-expression** (`jaxsr.SymbolicRegressor`) | Fits one sparse expression over a fixed basis library. | Available | +| **Stagewise additive** (`jaxsr.additive.StagewiseSymbolicRegressor`) | Repeatedly fits a small expression to the *residual* and adds it to the ensemble. Old terms are **frozen**. | Available | +| **Backfitting additive** (`jaxsr.additive.BackfittingSymbolicRegressor`) | Maintains a fixed set of terms and **revises** each one in place across sweeps (GAM-style). | Available (squared error); Bayesian variant planned | + +The key distinction between the two additive variants: + +- **Stagewise**: once a term is discovered it never changes; only its linear + weight may be re-estimated. +- **Backfitting**: terms are revised repeatedly, each conditioned on the current + fit of all the others. + +## Scope: what this is (and isn't) good for + +**In one line:** JAXSR is a *linear method over a fixed feature space* — it +selects a sparse combination of basis functions you supply. It is not a +free-composition equation discoverer. + +That distinction decides whether it is the right tool: + +- **Good fit:** the right building blocks are on the menu (or you can add them + with `BasisLibrary.add_custom`), and you want an interpretable, robust, + uncertainty-aware additive model. On targets that live in the library it is + fast and accurate, and the additive layer adds robust/quantile losses and + structural-uncertainty bootstrapping that genetic-programming tools don't + offer out of the box. +- **Wrong fit:** you want to *discover* an unknown compositional law such as + `exp(x0*x1)`, `x0 / (1 + x1**2)`, or `sin(2*x0)`. These are not single basis + functions, and the space of such compositions is infinite and continuously + parameterized, so no fixed library enumerates them in advance. For that, + reach for a genetic-programming or neural symbolic-regression tool (PySR, + Operon, AI-Feynman), which *search* the space of expressions instead of + selecting from a fixed dictionary — or try the experimental + [`RecursiveSymbolicRegressor`](#recursive-basis-expansion-experimental), which + grows compositions along the residual and partially lifts this ceiling. + +The limit is one of **discovery, not representation**: the linear-in-basis model +fits any of those targets perfectly the moment the exact term is in the library +(e.g. `library.add_custom("exp(x0*x1)", lambda X: jnp.exp(X[:, 0] * X[:, 1]))`) — +it simply cannot figure out *which* composition it needs without being told. +JAXSR's parametric bases can additionally fit a few constants *inside* a +pre-specified nonlinearity (e.g. `sin(a*x0)` with `a` optimized), but that still +requires you to name the functional form. The additive extensions in this guide +raise the *statistical* sophistication (boosting, robust/quantile losses, +backfitting, structural UQ); they do not change this expressiveness boundary. + +## Quick start + +```python +import numpy as np +from jaxsr.additive import StagewiseSymbolicRegressor + +rng = np.random.default_rng(0) +X = rng.uniform(-2, 2, size=(200, 2)) +y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + 0.1 * rng.normal(size=200) + +model = StagewiseSymbolicRegressor( + n_terms=5, + learning_rate=0.2, + max_complexity=6, + refit_coefficients=True, +) +model.fit(X, y) + +print(model) # pretty structural summary +print(model.expressions_) # per-term expression strings +print(model.coefficients_) # per-term weights +print(model.intercept_) # additive intercept +y_pred = model.predict(X) +``` + +The `print(model)` output looks like: + +``` +StagewiseSymbolicRegressor( + intercept = 1.07 + terms = + + 1 * (y = 2*x0 - 1.07 + 0.5*x1^2) + ... +) +``` + +## The stagewise algorithm + +1. Initialise the intercept to `mean(y)` and the prediction to that constant. +2. Compute the residual `y - prediction`. +3. Fit a small symbolic expression `g_k` to the residual (via + `jaxsr.fit_symbolic`). +4. Append `g_k` to the ensemble. +5. If `refit_coefficients=True`, rebuild the design matrix `Phi[:, j] = g_j(X)` + and re-solve `y ~= intercept + Phi @ coefficients` by least squares. + Otherwise, update `prediction += learning_rate * g_k(X)`. +6. Record train (and optional validation) loss. +7. Repeat until `n_terms` terms are added or early stopping triggers. + +## Key parameters + +| Parameter | Meaning | +|-----------|---------| +| `n_terms` | Maximum number of boosting stages (terms). | +| `learning_rate` | Shrinkage on each stage when `refit_coefficients=False`. | +| `max_complexity` | Complexity budget per term (max basis terms). Keep small to favour many simple terms. | +| `refit_coefficients` | Re-solve all linear weights by OLS after each stage. | +| `loss` | `"squared_error"` (default), `"absolute_error"`, `"huber"`, `"quantile"`, or a `Loss` instance. See [Losses](#losses-robust-and-quantile-regression). | +| `early_stopping` | Hold out a validation split and stop when it stops improving. | +| `validation_fraction`, `patience`, `min_delta` | Early-stopping controls. | +| `max_poly_degree`, `include_transcendental`, `include_ratios` | Which basis functions each term may use. | +| `information_criterion` | Complexity control within each term (`"aic"`, `"aicc"`, `"bic"`). | + +## Coefficient refitting + +- `refit_coefficients=False`: the weights are the learning-rate-scaled stagewise + weights (`coefficients_[k] == learning_rate`). +- `refit_coefficients=True`: after each new term, the intercept and *all* per-term + weights are re-solved by ordinary least squares over the discovered symbolic + features. This decouples term discovery (nonlinear, greedy) from term weighting + (linear, global) and typically improves accuracy. + +The refit uses `jnp.linalg.lstsq` (SVD-based, minimum-norm), so the highly +correlated columns produced by later boosting stages do not cause instability. + +## Combined expression + +`model.to_expression()` returns a single simplified SymPy expression combining +all terms: + +```python +expr = model.to_expression() # requires sympy +``` + +## Saving and loading + +Fitted models serialize to JSON (each term is stored via the underlying +`SymbolicRegressor` state), mirroring the rest of jaxsr. Note that the models +are **not picklable** — the basis-function closures cannot be pickled — so use +`save`/`load` rather than `pickle`: + +```python +model.save("additive_model.json") +loaded = StagewiseSymbolicRegressor.load("additive_model.json") +``` + +## Structural uncertainty (bootstrap) + +A single fitted expression can hide the fact that the *structure* itself is +uncertain — several different basis sets may explain the data about equally +well (this is common with collinear features). `bootstrap_additive` refits the +model on bootstrap resamples and reports, for each basis function, how often it +is selected — a cheap approximation to a posterior inclusion probability — +together with a predictive ensemble: + +```python +from jaxsr.additive import ( + StagewiseSymbolicRegressor, + bootstrap_additive, + bootstrap_predict_additive, +) + +est = StagewiseSymbolicRegressor(n_terms=3, max_complexity=2) +res = bootstrap_additive(est, X, y, n_bootstrap=100, random_state=0) + +# How stable is the discovered structure? +for name, prob in res["inclusion_probabilities"].items(): + print(f"{name:10s} selected in {prob:.0%} of resamples") + +# Prediction intervals that reflect *structural* variability, not just noise +pi = bootstrap_predict_additive(res["models"], X_new, alpha=0.1) +pi["mean"], pi["lower"], pi["upper"] +``` + +**How to read it.** Inclusion probabilities near 0 or 1 mean the structure is +identifiable and the single fitted expression is trustworthy. **Diffuse** values +(e.g. a basis selected 50–60% of the time) mean the data do not determine one +expression — no single symbolic model should be over-trusted, and the bootstrap +intervals are the honest summary. This also works as a decision gate for heavier +Bayesian modelling: if the probabilities are already crisp, there is little +structural uncertainty left to quantify. It works for both the stagewise and +backfitting regressors. + +## Early stopping + +With `early_stopping=True`, a validation split (`validation_fraction`) is held +out. After each stage the validation loss is recorded; training stops once it +fails to improve by at least `min_delta` for `patience` consecutive stages, and +the model rolls back to the best iteration. + +## Losses: robust and quantile regression + +This is where additive symbolic regression goes beyond ordinary least-squares +symbolic regression. Each weak learner fits the negative gradient `-dL/dy_pred` +(gradient boosting), so you can target losses that OLS selection cannot: + +| `loss` | Class | Use when | +|--------|-------|----------| +| `"squared_error"` (default) | `SquaredError` | Standard regression | +| `"absolute_error"` | `AbsoluteError` | Outliers present (fits the median) | +| `"huber"` | `HuberLoss(delta=1.35)` | Outliers, but keep efficiency near zero | +| `"quantile"` | `QuantileLoss(quantile=0.5)` | Quantiles / prediction intervals / asymmetric cost | + +Pass a name for defaults, or an instance to customise: + +```python +from jaxsr.additive import StagewiseSymbolicRegressor, QuantileLoss, HuberLoss + +# Robust regression: heavy outliers barely move the fit +robust = StagewiseSymbolicRegressor(loss="huber", learning_rate=0.5).fit(X, y) + +# 90th-percentile regression (build intervals by fitting several quantiles) +q90 = StagewiseSymbolicRegressor(loss=QuantileLoss(0.9), learning_rate=0.5).fit(X, y) +``` + +**How non-squared losses are fit.** Each stage fits a symbolic term to the +negative gradient, then a **line search** picks the step size that minimises the +loss (`learning_rate` shrinks that step). Because the ordinary least-squares +coefficient refit targets squared error, `refit_coefficients=True` is ignored for +non-squared losses (a warning is issued) and gradient boosting is used instead — +so set `refit_coefficients=False` explicitly for robust/quantile models. + +The optimal constant initialisation adapts to the loss: mean for squared error, +median for absolute/Huber, and the empirical quantile for quantile loss. + +Add further losses (Poisson, logistic, ...) by subclassing `Loss` and +registering them in `jaxsr.additive.losses._LOSSES`. + +## Backfitting (GAM-style) + +`BackfittingSymbolicRegressor` maintains a **fixed** number of terms and +*revises* each one across sweeps, instead of freezing them. Each sweep removes a +term, re-discovers its expression on the partial residual, and puts it back: + +```python +from jaxsr.additive import BackfittingSymbolicRegressor + +model = BackfittingSymbolicRegressor(n_terms=4, n_sweeps=6, max_complexity=3) +model.fit(X, y) # warm-started from a stagewise fit, then refined by sweeps +``` + +``` +for sweep in 1..n_sweeps: + for term j: + partial_residual = y - intercept - sum_{i != j} coef_i * g_i(X) + g_j = fit_symbolic(X, partial_residual, ...) # re-discover structure + intercept, coef = OLS refit over all terms +(stop when the training loss stops improving by `tol`) +``` + +It is warm-started from a stagewise fit and currently supports **squared error +only**. Structure re-discovery makes the sweep a heuristic (no monotonicity +guarantee), so the best-loss iterate is kept. + +**When does it actually help?** Backfitting starts from the stagewise+refit fit +and keeps the best-loss iterate, so **it is never worse than +`StagewiseSymbolicRegressor(refit_coefficients=True)` on the training data** — +the only question is whether the sweeps improve on it. That hinges entirely on +whether re-discovery changes the *set* of selected basis functions: + +- **Generous per-term budget** (`max_complexity` ≥ 2–3): greedy usually already + finds a sufficient basis set, so the joint least-squares refit makes the two + essentially identical. Backfitting adds nothing here — prefer the stagewise + regressor. +- **Small per-term budget and collinear features** (`max_complexity=1`, the + GAM-style single-basis regime): greedy forward selection can lock into a + *suboptimal* basis set that a single forward pass cannot undo. Backfitting's + coordinate-descent re-discovery escapes it, changing the basis union and + improving the fit — we have measured up to roughly **+0.04 train / +0.06 test + R²** in this regime, with no downside in the cases where it does not help. + +So reach for backfitting when you want **small, revisable single-basis terms +over correlated features** (or a fixed-size GAM-style decomposition); use the +stagewise regressor for larger per-term expressions. Its other forward-looking +value is as the foundation for a future **Bayesian backfitting** variant +(BART/iBART-style), which would sample a *posterior over symbolic structure* — +genuinely beyond point-estimate SR — using the same partial-residual sweep with +conjugate marginal likelihoods. + +## Recursive basis expansion (experimental) + +The [Scope](#scope-what-this-is-and-isnt-good-for) section notes that a fixed +library cannot *discover* compositional forms like `x0*sin(x1)` or `exp(x0*x1)`. +`RecursiveSymbolicRegressor` is an experimental step past that ceiling. Instead +of enumerating a huge composition space up front (which explodes +combinatorially), it **grows the library lazily along the residual**: + +1. Fit a sparse model over the current library; take the residual. +2. Compose the currently useful terms (selected terms + features) with a small + operator set (unary functions, products, ratios) — one new layer. +3. Screen: drop non-finite candidates, deduplicate, keep the top few by + correlation with the residual. +4. Add the survivors and refit. Repeat. The effective composition depth is the + number of rounds, because a term found in one round feeds the next. + +```python +from jaxsr.additive import RecursiveSymbolicRegressor + +model = RecursiveSymbolicRegressor(n_expansions=3, max_terms=6, beam_width=25) +model.fit(X, y) +print(model.expression_) # e.g. recovers "exp((x0)*(x1))" exactly +print(model.history_) # library size / n_terms / train R^2 per round +``` + +This is essentially Fast Function Extraction (FFX) / symbolic feature +construction — a deterministic, bounded cousin of genetic programming. On simple +compositional targets it substantially beats a flat library (e.g. `exp(x0*x1)`: +R² 1.00 vs 0.70) and is competitive with a strong GP engine (matched Operon on +`x0*sin(x1)` in our tests). + +**Caveats.** It re-enters search-based territory: cost grows with `beam_width`, +`n_expansions`, and feature count, and it will not match a mature GP (PySR, +Operon) on hard, high-dimensional, or deeply nested targets. The result is an +ordinary `SymbolicRegressor` over the grown library (so predict/`expression_`/ +scoring and the base regressor's non-finite-basis guard and negligible-term +pruning all apply), but the composed bases are Python closures, so the fitted +model is **not** serialisable via `save`/`load`, and `to_sympy` may not parse +deeply nested term names. diff --git a/.claude/skills/jaxsr/templates/additive-regression.py b/.claude/skills/jaxsr/templates/additive-regression.py new file mode 100644 index 0000000..f1c6d2b --- /dev/null +++ b/.claude/skills/jaxsr/templates/additive-regression.py @@ -0,0 +1,99 @@ +""" +Additive Symbolic Regression — boosting-style ensembles of small expressions. + +This template shows the ``jaxsr.additive`` workflows: +1. Stagewise additive regression (gradient boosting with symbolic weak learners) +2. Robust regression under outliers (Huber / absolute-error losses) +3. Quantile regression (prediction intervals via the pinball loss) +4. Structural uncertainty (bootstrap basis-inclusion probabilities) +5. Backfitting (GAM-style: revise terms instead of freezing them) +6. Recursive expansion (reach compositions a flat library misses) + +Pick the section you need; each block is self-contained after the imports. +""" + +import numpy as np + +from jaxsr import fit_symbolic +from jaxsr.additive import ( + BackfittingSymbolicRegressor, + QuantileLoss, + RecursiveSymbolicRegressor, + StagewiseSymbolicRegressor, + bootstrap_additive, + bootstrap_predict_additive, +) + +# Replace with your own data. X shape (n_samples, n_features), y shape (n_samples,). +rng = np.random.default_rng(0) +X = rng.uniform(-2, 2, size=(300, 2)) +y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + 0.1 * rng.normal(size=300) + +# ============================================================================= +# 1. Stagewise additive regression +# Many small interpretable terms; keep max_complexity small. +# refit_coefficients=True re-solves all weights by least squares each stage. +# ============================================================================= +model = StagewiseSymbolicRegressor( + n_terms=10, + learning_rate=0.2, + max_complexity=4, + refit_coefficients=True, + early_stopping=False, # set True + validation_fraction to guard overfitting +) +model.fit(X, y) +print(model) # pretty structural summary +print(model.expressions_) # per-term expression strings +print(model.intercept_, model.coefficients_) +print("combined:", model.to_expression()) +model.save("additive_model.json") # JSON round-trip (models are NOT picklable) + +# ============================================================================= +# 2. Robust regression (outliers) — use refit_coefficients=False for any +# non-squared loss (OLS refit only applies to squared error). +# ============================================================================= +robust = StagewiseSymbolicRegressor( + loss="huber", # or "absolute_error" + n_terms=8, + max_complexity=3, + learning_rate=0.5, + refit_coefficients=False, +).fit(X, y) + +# ============================================================================= +# 3. Quantile regression — fit several quantiles to build a prediction band. +# ============================================================================= +q_models = { + q: StagewiseSymbolicRegressor( + loss=QuantileLoss(q), + n_terms=10, + max_complexity=3, + learning_rate=0.5, + refit_coefficients=False, + ).fit(X, y) + for q in (0.1, 0.5, 0.9) +} +lower, median, upper = (q_models[q].predict(X) for q in (0.1, 0.5, 0.9)) + +# ============================================================================= +# 4. Structural uncertainty — how stable is the discovered structure? +# ============================================================================= +res = bootstrap_additive(model, X, y, n_bootstrap=100, random_state=0) +print(res["inclusion_probabilities"]) # {basis: fraction selected} +pi = bootstrap_predict_additive(res["models"], X) # mean / lower / upper / ... + +# ============================================================================= +# 5. Backfitting (GAM-style) — fixed set of terms, revised across sweeps. +# ============================================================================= +bf = BackfittingSymbolicRegressor(n_terms=4, n_sweeps=6, max_complexity=3).fit(X, y) + +# ============================================================================= +# 6. Recursive expansion (experimental) — discover compositions like exp(x0*x1) +# that a flat library cannot reach. +# ============================================================================= +Xc = rng.uniform(-2, 2, size=(400, 2)) +yc = np.exp(Xc[:, 0] * Xc[:, 1]) +rec = RecursiveSymbolicRegressor(n_expansions=3, max_terms=6, beam_width=25).fit(Xc, yc) +print(rec.expression_) # e.g. "exp((x0)*(x1))" +# Compare with a flat library (which misses it): +flat = fit_symbolic(Xc, yc, max_terms=6, include_transcendental=True) diff --git a/CLAUDE.md b/CLAUDE.md index d10b577..1d1cacc 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -184,6 +184,7 @@ Map which JAXSR features have guide/template/notebook coverage. - Known-model fitting (`guides/known-model-fitting.md`) - Scikit-learn integration (`guides/sklearn-integration.md`) - CLI reference (`guides/cli.md`) +- Additive symbolic regression (`guides/additive.md`) **Gaps to fill:** - Metrics comparison guide (when to use R² vs AIC vs cross-validation) diff --git a/README.md b/README.md index 5c8979b..c510c0c 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ JAXSR is a fully open-source symbolic regression library built on JAX that disco - **Adaptive Sampling**: Intelligently suggest new data points to improve model quality - **JAX-Accelerated**: JIT compilation and GPU support for fast computation - **Symbolic Classification**: Discover interpretable logistic models for binary and multiclass problems via IRLS + sparse selection +- **Additive Symbolic Regression**: Boosting-style ensembles of small symbolic expressions — fit residuals stagewise for many simple, interpretable terms (`jaxsr.additive`) - **Scikit-learn Compatible**: Full estimator protocol (`get_params`/`set_params`/`clone`) — works with `cross_val_score`, `GridSearchCV`, `Pipeline` - **Symbolic Export**: Export to SymPy, LaTeX, or pure Python/NumPy functions @@ -253,6 +254,97 @@ y_pred = clf.predict(jnp.array(X)) Multiclass problems are handled automatically via one-vs-rest (OVR), giving each class its own interpretable expression. The classifier also supports coefficient intervals, conformal prediction sets, SymPy/LaTeX export, and save/load. +## Additive Symbolic Regression + +`jaxsr.additive` fits a model as a sum of small symbolic expressions, +`f(x) = c + Σ_k η_k · g_k(x)` — analogous to gradient boosting, but each weak +learner is an interpretable symbolic expression rather than a decision tree. +`StagewiseSymbolicRegressor` repeatedly fits a small expression to the current +residual and (optionally) refits all linear coefficients by least squares: + +> **Scope:** JAXSR is a *linear method over a fixed basis library* — it selects a +> sparse combination of basis functions you supply. It excels at interpretable, +> robust, uncertainty-aware modeling when the right building blocks are on the +> menu, but it does **not** discover unknown compositional forms like +> `exp(x0*x1)` or `sin(2*x0)` (the limit is discovery, not representation). For +> free-composition equation discovery, use a genetic-programming or neural +> symbolic-regression tool (PySR, Operon, AI-Feynman). + +```python +import numpy as np +from jaxsr.additive import StagewiseSymbolicRegressor + +# Additive target: y = 2*x0 + 0.5*x1^2 + noise +rng = np.random.default_rng(0) +X = rng.uniform(-2, 2, size=(200, 2)) +y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + 0.1 * rng.normal(size=200) + +model = StagewiseSymbolicRegressor( + n_terms=5, + learning_rate=0.2, + max_complexity=6, + refit_coefficients=True, + early_stopping=False, +) +model.fit(X, y) + +print(model) # pretty structural summary of the ensemble +print(model.expressions_) # per-term expression strings +print(model.intercept_, model.coefficients_) +y_pred = model.predict(X) +combined = model.to_expression() # single combined SymPy expression +model.save("additive_model.json") # JSON round-trip (models are not picklable) +``` + +Prefer this over a single large expression when the signal is a sum of several +simple effects: keep `max_complexity` small and let the ensemble accumulate many +interpretable terms. Early stopping on a validation split guards against +overfitting on noisy data. + +Because each term is fit by gradient boosting (fitting the negative gradient), +you can also target losses ordinary least-squares selection cannot — **robust** +and **quantile** symbolic regression: + +```python +from jaxsr.additive import StagewiseSymbolicRegressor, QuantileLoss + +# Robust to outliers (Huber); use refit_coefficients=False for non-squared losses +robust = StagewiseSymbolicRegressor(loss="huber", learning_rate=0.5, + refit_coefficients=False).fit(X, y) + +# 90th-percentile regression (fit several quantiles to build intervals) +q90 = StagewiseSymbolicRegressor(loss=QuantileLoss(0.9), learning_rate=0.5, + refit_coefficients=False).fit(X, y) +``` + +An experimental `RecursiveSymbolicRegressor` grows the basis library along the +residual (composing unary functions, products, and ratios of discovered terms), +partially lifting the fixed-library ceiling — it can recover compositional +targets like `exp(x0*x1)` that a flat library misses. + +Available losses: `"squared_error"` (default), `"absolute_error"`, `"huber"`, +`"quantile"`. A `BackfittingSymbolicRegressor` (GAM-style, where a fixed set of +terms is revised rather than frozen) is also available for squared error. It is +never worse than stagewise+refit on training and genuinely helps when per-term +complexity is small and features are collinear (where greedy selection gets +stuck); it is also the foundation for a future Bayesian (BART/iBART-style) +variant. + +`bootstrap_additive` quantifies **structural uncertainty** — how often each +basis function is selected across bootstrap resamples (a proxy for posterior +inclusion probability), plus a predictive ensemble whose intervals reflect +structural variability, not just coefficient noise: + +```python +from jaxsr.additive import bootstrap_additive, bootstrap_predict_additive + +res = bootstrap_additive(model, X, y, n_bootstrap=100, random_state=0) +res["inclusion_probabilities"] # {basis: fraction selected} +pi = bootstrap_predict_additive(res["models"], X_new) # mean / lower / upper +``` + +See `docs/guides/additive-symbolic-regression.md`. + ## Visualization ```python @@ -335,6 +427,9 @@ See the `docs/examples/` directory for complete worked examples: - `chemical_kinetics.py`: Discovering rate laws from kinetic data - `heat_transfer.py`: Heat transfer correlations +The `examples/` directory also has a standalone script, +`additive_symbolic_regression.py`, for boosting-style additive models. + ## API Reference See the [documentation](docs/) for full API details. diff --git a/docs/_toc.yml b/docs/_toc.yml index a01eae3..90ba8ff 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -8,6 +8,7 @@ parts: - caption: Guides chapters: + - file: guides/additive-symbolic-regression - file: guides/doe_guide - file: guides/acquisition - file: guides/cli_guide @@ -18,6 +19,7 @@ parts: - caption: Examples chapters: - file: examples/basic_usage + - file: examples/additive_symbolic_regression - file: examples/chemical_kinetics - file: examples/heat_transfer - file: examples/uncertainty_quantification diff --git a/docs/api/additive.rst b/docs/api/additive.rst new file mode 100644 index 0000000..1f19671 --- /dev/null +++ b/docs/api/additive.rst @@ -0,0 +1,50 @@ +jaxsr.additive +============== + +Additive (boosting-style) symbolic regression: fit a model as a sum of small +symbolic expressions, ``f(x) = c + sum_k eta_k * g_k(x)``. + +.. automodule:: jaxsr.additive + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: jaxsr.additive.stagewise + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: jaxsr.additive.ensemble + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: jaxsr.additive.losses + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: jaxsr.additive.coefficient_refit + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: jaxsr.additive.backfitting + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: jaxsr.additive.recursive + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: jaxsr.additive.base + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: jaxsr.additive.uncertainty + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/index.rst b/docs/api/index.rst index b7f5513..ba6a6f7 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -17,3 +17,4 @@ Full API documentation generated from source docstrings. rsm plotting classifier + additive diff --git a/docs/examples/additive_symbolic_regression.ipynb b/docs/examples/additive_symbolic_regression.ipynb new file mode 100644 index 0000000..35b0473 --- /dev/null +++ b/docs/examples/additive_symbolic_regression.ipynb @@ -0,0 +1,691 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "c1ef880d", + "metadata": {}, + "source": [ + "# Additive Symbolic Regression\n", + "\n", + "Additive symbolic regression fits a model as a **sum of small symbolic expressions**,\n", + "$f(x) = c + \\sum_k \\eta_k\\, g_k(x)$, where each $g_k$ is discovered by the existing JAXSR machinery.\n", + "This is analogous to **gradient boosting**, except each weak learner is an interpretable symbolic expression rather than a decision tree.\n", + "\n", + "This notebook shows the `StagewiseSymbolicRegressor`: it repeatedly fits a small expression to the current residual, then (optionally) refits all linear coefficients by least squares." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "65964650", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T21:42:08.032766Z", + "iopub.status.busy": "2026-07-02T21:42:08.032590Z", + "iopub.status.idle": "2026-07-02T21:42:08.855904Z", + "shell.execute_reply": "2026-07-02T21:42:08.854882Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "((200, 2), (100, 2))" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import numpy as np\n", + "\n", + "from jaxsr.additive import StagewiseSymbolicRegressor\n", + "\n", + "# Additive target: y = 2*x0 + 0.5*x1^2 + noise\n", + "rng = np.random.default_rng(0)\n", + "X = rng.uniform(-2, 2, size=(300, 2))\n", + "y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + 0.1 * rng.normal(size=300)\n", + "\n", + "# hold out a test set\n", + "n_train = 200\n", + "X_train, y_train = X[:n_train], y[:n_train]\n", + "X_test, y_test = X[n_train:], y[n_train:]\n", + "X_train.shape, X_test.shape" + ] + }, + { + "cell_type": "markdown", + "id": "be5a3483", + "metadata": {}, + "source": [ + "## Fit a stagewise additive model\n", + "\n", + "Keep `max_complexity` small so each stage is a simple expression and the ensemble accumulates many interpretable terms. With `refit_coefficients=True`, all linear weights are re-solved by least squares after each stage." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "df161198", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T21:42:08.858040Z", + "iopub.status.busy": "2026-07-02T21:42:08.857770Z", + "iopub.status.idle": "2026-07-02T21:42:14.519612Z", + "shell.execute_reply": "2026-07-02T21:42:14.518339Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "StagewiseSymbolicRegressor(\n", + " intercept = 1.075\n", + " terms =\n", + " + 1.001 * (y = 1.998*x0 - 1.07 + 0.5015*x1^2)\n", + " + 2.31 * (y = - 0.003903*x0^2)\n", + " + 1.865 * (y = 0.002939*x1)\n", + " + 1.124 * (y = 0.001106*x0*x1)\n", + " + 6.931 * (y = - 0.0001612*x1^3)\n", + ")\n" + ] + } + ], + "source": [ + "model = StagewiseSymbolicRegressor(\n", + " n_terms=5,\n", + " learning_rate=0.2,\n", + " max_complexity=4,\n", + " refit_coefficients=True,\n", + ")\n", + "model.fit(X_train, y_train)\n", + "print(model)" + ] + }, + { + "cell_type": "markdown", + "id": "58068d48", + "metadata": {}, + "source": [ + "## Inspect the learned ensemble" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "7ea9c7a0", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T21:42:14.522041Z", + "iopub.status.busy": "2026-07-02T21:42:14.521822Z", + "iopub.status.idle": "2026-07-02T21:42:14.526913Z", + "shell.execute_reply": "2026-07-02T21:42:14.525993Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "intercept: 1.0750973224639893\n", + "coefficients: [1.0005701780319214, 2.3100991249084473, 1.8648916482925415, 1.1235284805297852, 6.930931568145752]\n", + "\n", + "term 0: y = 1.998*x0 - 1.07 + 0.5015*x1^2\n", + "term 1: y = - 0.003903*x0^2\n", + "term 2: y = 0.002939*x1\n", + "term 3: y = 0.001106*x0*x1\n", + "term 4: y = - 0.0001612*x1^3\n" + ] + } + ], + "source": [ + "print('intercept:', model.intercept_)\n", + "print('coefficients:', model.coefficients_)\n", + "print()\n", + "for i, expr in enumerate(model.expressions_):\n", + " print(f'term {i}: {expr}')" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "78847b7b", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T21:42:14.528845Z", + "iopub.status.busy": "2026-07-02T21:42:14.528660Z", + "iopub.status.idle": "2026-07-02T21:42:15.407397Z", + "shell.execute_reply": "2026-07-02T21:42:15.406417Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "train R^2: 0.9982670545578003\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "test R^2: 0.9988469481468201\n" + ] + } + ], + "source": [ + "print('train R^2:', model.score(X_train, y_train))\n", + "print('test R^2:', model.score(X_test, y_test))" + ] + }, + { + "cell_type": "markdown", + "id": "b986e721", + "metadata": {}, + "source": [ + "## Combined symbolic expression\n", + "\n", + "`to_expression()` collapses the ensemble into a single simplified SymPy expression." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "ab64ae2e", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T21:42:15.409685Z", + "iopub.status.busy": "2026-07-02T21:42:15.409503Z", + "iopub.status.idle": "2026-07-02T21:42:15.891981Z", + "shell.execute_reply": "2026-07-02T21:42:15.891107Z" + } + }, + "outputs": [ + { + "data": { + "text/latex": [ + "$\\displaystyle 0.00124225317639448 x_{0} x_{1} + 1.99923764521122 x_{0} - 0.0090172041649631 x_{0}^{2.0} + 0.00548121186382547 x_{1} + 0.501793767439942 x_{1}^{2.0} - 0.00111730957318264 x_{1}^{3.0} + 0.00453644099624739$" + ], + "text/plain": [ + "0.00124225317639448*x0*x1 + 1.99923764521122*x0 - 0.0090172041649631*x0**2.0 + 0.00548121186382547*x1 + 0.501793767439942*x1**2.0 - 0.00111730957318264*x1**3.0 + 0.00453644099624739" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "model.to_expression()" + ] + }, + { + "cell_type": "markdown", + "id": "21de3fe8", + "metadata": {}, + "source": [ + "## Why additive? Many small terms vs one big expression\n", + "\n", + "When the signal is a sum of several simple effects, a tiny per-stage budget still recovers it because the ensemble accumulates terms across stages — where a single expression with the same small budget cannot." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "c891e663", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T21:42:15.894396Z", + "iopub.status.busy": "2026-07-02T21:42:15.894116Z", + "iopub.status.idle": "2026-07-02T21:42:21.413019Z", + "shell.execute_reply": "2026-07-02T21:42:21.411741Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "stagewise (max_complexity=1, 10 stages) test R^2: 0.9931\n", + "single expression (max_terms=1) test R^2: 0.3119\n" + ] + } + ], + "source": [ + "from jaxsr import fit_symbolic\n", + "\n", + "rng = np.random.default_rng(2)\n", + "Xm = rng.uniform(-1.5, 1.5, size=(500, 4))\n", + "ym = (\n", + " 1.5 * Xm[:, 0]\n", + " + 2.0 * Xm[:, 1] ** 2\n", + " - 1.0 * Xm[:, 2] ** 3\n", + " + 0.8 * Xm[:, 3]\n", + " + 0.2 * rng.normal(size=500)\n", + ")\n", + "Xtr, ytr, Xte, yte = Xm[:300], ym[:300], Xm[300:], ym[300:]\n", + "\n", + "stagewise = StagewiseSymbolicRegressor(n_terms=10, max_complexity=1).fit(Xtr, ytr)\n", + "single = fit_symbolic(Xtr, ytr, max_terms=1, include_transcendental=False)\n", + "\n", + "print('stagewise (max_complexity=1, 10 stages) test R^2:', round(float(stagewise.score(Xte, yte)), 4))\n", + "print('single expression (max_terms=1) test R^2:', round(float(single.score(Xte, yte)), 4))" + ] + }, + { + "cell_type": "markdown", + "id": "4bf5887e", + "metadata": {}, + "source": [ + "## Early stopping and saving\n", + "\n", + "Early stopping holds out a validation split and stops once it stops improving. Models serialize to JSON (they are not picklable because the basis functions use closures)." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "18ebecfc", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T21:42:21.415204Z", + "iopub.status.busy": "2026-07-02T21:42:21.414992Z", + "iopub.status.idle": "2026-07-02T21:42:25.216306Z", + "shell.execute_reply": "2026-07-02T21:42:25.215094Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "terms used after early stopping: 2\n", + "reloaded test R^2: 0.9989\n" + ] + } + ], + "source": [ + "import tempfile, os\n", + "\n", + "es = StagewiseSymbolicRegressor(\n", + " n_terms=20, max_complexity=3,\n", + " early_stopping=True, validation_fraction=0.25, patience=3, random_state=0,\n", + ").fit(X_train, y_train)\n", + "print('terms used after early stopping:', es.n_terms_)\n", + "\n", + "path = os.path.join(tempfile.mkdtemp(), 'additive_model.json')\n", + "es.save(path)\n", + "loaded = StagewiseSymbolicRegressor.load(path)\n", + "print('reloaded test R^2:', round(float(loaded.score(X_test, y_test)), 4))\n", + "os.remove(path)" + ] + }, + { + "cell_type": "markdown", + "id": "ae85f56d", + "metadata": {}, + "source": [ + "## Robust and quantile regression (beyond ordinary SR)\n", + "\n", + "Because each term is fit by **gradient boosting** (fitting the negative gradient), additive symbolic regression can target losses that least-squares selection cannot. Use `refit_coefficients=False` for non-squared losses.\n", + "\n", + "First, robustness to outliers with the Huber and absolute-error losses:" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "780a1673", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T21:42:25.218481Z", + "iopub.status.busy": "2026-07-02T21:42:25.218229Z", + "iopub.status.idle": "2026-07-02T21:42:28.761636Z", + "shell.execute_reply": "2026-07-02T21:42:28.760085Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "squared_error MAE vs clean signal = 2.764\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "huber MAE vs clean signal = 0.144\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "absolute_error MAE vs clean signal = 0.049\n" + ] + } + ], + "source": [ + "# Contaminate 8% of the targets with heavy outliers\n", + "rng = np.random.default_rng(1)\n", + "Xr = rng.uniform(-2, 2, size=(500, 2))\n", + "clean = 2.0 * Xr[:, 0] + 0.5 * Xr[:, 1] ** 2\n", + "yr = clean + rng.normal(0, 0.1, 500)\n", + "yr[rng.choice(500, 40, replace=False)] += 30.0\n", + "\n", + "for loss in ['squared_error', 'huber', 'absolute_error']:\n", + " m = StagewiseSymbolicRegressor(\n", + " n_terms=8, max_complexity=4, learning_rate=0.5,\n", + " loss=loss, refit_coefficients=False,\n", + " ).fit(Xr, yr)\n", + " mae = float(np.mean(np.abs(np.array(m.predict(Xr)) - clean)))\n", + " print(f'{loss:15s} MAE vs clean signal = {mae:.3f}')" + ] + }, + { + "cell_type": "markdown", + "id": "79fbd672", + "metadata": {}, + "source": [ + "The squared-error fit is dragged toward the outliers; Huber and absolute-error stay close to the true signal.\n", + "\n", + "Next, **quantile regression** — fit several quantiles to form a prediction band:" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "02f0c39f", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T21:42:28.763568Z", + "iopub.status.busy": "2026-07-02T21:42:28.763361Z", + "iopub.status.idle": "2026-07-02T21:42:30.943932Z", + "shell.execute_reply": "2026-07-02T21:42:30.942719Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "q=0.1: fraction of points below prediction = 0.100\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "q=0.5: fraction of points below prediction = 0.514\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "q=0.9: fraction of points below prediction = 0.890\n" + ] + } + ], + "source": [ + "from jaxsr.additive import QuantileLoss\n", + "\n", + "rng = np.random.default_rng(2)\n", + "Xq = rng.uniform(-2, 2, size=(500, 2))\n", + "yq = 2.0 * Xq[:, 0] + 0.5 * Xq[:, 1] ** 2 + rng.normal(0, 1.0, 500)\n", + "\n", + "for q in [0.1, 0.5, 0.9]:\n", + " m = StagewiseSymbolicRegressor(\n", + " n_terms=10, max_complexity=3, learning_rate=0.5,\n", + " loss=QuantileLoss(q), refit_coefficients=False,\n", + " ).fit(Xq, yq)\n", + " coverage = float(np.mean(yq <= np.array(m.predict(Xq))))\n", + " print(f'q={q}: fraction of points below prediction = {coverage:.3f}')" + ] + }, + { + "cell_type": "markdown", + "id": "f7a9d141", + "metadata": {}, + "source": [ + "## Structural uncertainty: is the discovered expression stable?\n", + "\n", + "A single fitted expression can hide that the *structure* itself is uncertain — several basis sets may explain the data about equally well (common with collinear features). `bootstrap_additive` refits on bootstrap resamples and reports how often each basis function is selected (a proxy for a posterior inclusion probability), plus a predictive ensemble." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "2983f9ba", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T21:42:30.946122Z", + "iopub.status.busy": "2026-07-02T21:42:30.945902Z", + "iopub.status.idle": "2026-07-02T21:42:43.245525Z", + "shell.execute_reply": "2026-07-02T21:42:43.244487Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "x2 selected in 100% of resamples\n", + "x1 selected in 62% of resamples\n", + "x0 selected in 50% of resamples\n", + "x0^3 selected in 45% of resamples\n", + "x1^3 selected in 14% of resamples\n", + "x2^2 selected in 11% of resamples\n" + ] + } + ], + "source": [ + "from jaxsr.additive import bootstrap_additive, bootstrap_predict_additive\n", + "\n", + "# Collinear features: x2 is nearly a copy of x0, so the structure is not\n", + "# well determined even though the fit is good.\n", + "rng = np.random.default_rng(0)\n", + "a = rng.normal(0, 1, 400)\n", + "c0, c1 = a, 0.9 * a + 0.1 * rng.normal(0, 1, 400)\n", + "c2 = 0.8 * a + 0.2 * rng.normal(0, 1, 400)\n", + "Xc = np.column_stack([c0, c1, c2])\n", + "yc = c0 - c1 + 0.5 * c2 + rng.normal(0, 0.1, 400)\n", + "\n", + "est = StagewiseSymbolicRegressor(n_terms=3, max_complexity=1)\n", + "res = bootstrap_additive(est, Xc, yc, n_bootstrap=80, random_state=0)\n", + "for name, prob in list(res['inclusion_probabilities'].items())[:6]:\n", + " print(f'{name:8s} selected in {prob:.0%} of resamples')" + ] + }, + { + "cell_type": "markdown", + "id": "ed335fd8", + "metadata": {}, + "source": [ + "Inclusion probabilities near 0 or 1 mean the structure is identifiable; **diffuse** values (~0.5) mean the data don't pin down one expression. The bootstrap predictive interval is then the honest summary:" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "a5c3fa5a", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T21:42:43.247815Z", + "iopub.status.busy": "2026-07-02T21:42:43.247625Z", + "iopub.status.idle": "2026-07-02T21:42:43.832166Z", + "shell.execute_reply": "2026-07-02T21:42:43.830953Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "x[0] mean=-0.00 90% interval=[-0.04, +0.04]\n", + "x[1] mean=-0.01 90% interval=[-0.06, +0.04]\n", + "x[2] mean=+0.37 90% interval=[+0.30, +0.47]\n", + "x[3] mean=+0.03 90% interval=[-0.11, +0.11]\n", + "x[4] mean=-0.24 90% interval=[-0.29, -0.16]\n" + ] + } + ], + "source": [ + "pi = bootstrap_predict_additive(res['models'], Xc[:5], alpha=0.1)\n", + "for i in range(5):\n", + " lo, hi, mid = pi['lower'][i], pi['upper'][i], pi['mean'][i]\n", + " print(f'x[{i}] mean={mid:+.2f} 90% interval=[{lo:+.2f}, {hi:+.2f}]')" + ] + }, + { + "cell_type": "markdown", + "id": "04c28885", + "metadata": {}, + "source": [ + "## Backfitting: revising terms instead of freezing them\n", + "\n", + "`BackfittingSymbolicRegressor` keeps a **fixed** set of terms and re-discovers each one on the partial residual across sweeps (GAM-style), warm-started from a stagewise fit. For squared error it typically *matches* stagewise+refit rather than beating it (the joint least-squares refit already optimises the linear combination), so prefer it when you specifically want a fixed-size, revisable decomposition." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "3c72570f", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T21:42:43.833986Z", + "iopub.status.busy": "2026-07-02T21:42:43.833794Z", + "iopub.status.idle": "2026-07-02T21:42:44.310569Z", + "shell.execute_reply": "2026-07-02T21:42:44.309337Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "backfitting test R^2: 0.9988\n", + "sweeps run: 1\n" + ] + } + ], + "source": [ + "from jaxsr.additive import BackfittingSymbolicRegressor\n", + "\n", + "bf = BackfittingSymbolicRegressor(n_terms=4, n_sweeps=6, max_complexity=3).fit(\n", + " X_train, y_train\n", + ")\n", + "print('backfitting test R^2:', round(float(bf.score(X_test, y_test)), 4))\n", + "print('sweeps run:', len(bf.training_history_) - 1)" + ] + }, + { + "cell_type": "markdown", + "id": "8723fb53", + "metadata": {}, + "source": [ + "## Recursive basis expansion: reaching compositions (experimental)\n", + "\n", + "A fixed library can only *select* from basis functions you supply — it cannot *discover* a composition like `exp(x0*x1)` that isn't in it. `RecursiveSymbolicRegressor` grows the library along the residual, composing the useful terms with unary functions, products, and ratios each round. On this target it recovers `exp(x0*x1)` exactly, where a flat library with transcendentals cannot." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "034c95a1", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-02T21:42:44.312876Z", + "iopub.status.busy": "2026-07-02T21:42:44.312670Z", + "iopub.status.idle": "2026-07-02T21:42:54.502732Z", + "shell.execute_reply": "2026-07-02T21:42:54.501056Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/user/jaxsr/src/jaxsr/regressor.py:1816: UserWarning: Excluding 4 basis function(s) with non-finite values on the training data; they will not be selected.\n", + " return model.fit(X, y)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "recursive test R^2: 1.0 -> y = 0.9998*exp((x0)*(x1))\n", + "flat lib test R^2: 0.7172\n" + ] + } + ], + "source": [ + "from jaxsr import fit_symbolic\n", + "from jaxsr.additive import RecursiveSymbolicRegressor\n", + "\n", + "rng = np.random.default_rng(0)\n", + "Xc = rng.uniform(-2, 2, size=(400, 2))\n", + "yc = np.exp(Xc[:, 0] * Xc[:, 1]) + 0.02 * rng.normal(size=400)\n", + "Xtr2, ytr2, Xte2, yte2 = Xc[:200], yc[:200], Xc[200:], yc[200:]\n", + "\n", + "def r2(y, p):\n", + " y, p = np.asarray(y, float), np.asarray(p, float)\n", + " return 1 - np.sum((y - p) ** 2) / np.sum((y - y.mean()) ** 2)\n", + "\n", + "rec = RecursiveSymbolicRegressor(n_expansions=3, max_terms=6, beam_width=25).fit(Xtr2, ytr2)\n", + "flat = fit_symbolic(Xtr2, ytr2, max_terms=6, include_transcendental=True)\n", + "print('recursive test R^2:', round(r2(yte2, rec.predict(Xte2)), 4), '->', rec.expression_[:40])\n", + "print('flat lib test R^2:', round(r2(yte2, flat.predict(Xte2)), 4))" + ] + }, + { + "cell_type": "markdown", + "id": "6bbf170d", + "metadata": {}, + "source": [ + "It is an experimental, deterministic FFX-style search: competitive with a strong genetic-programming engine on simple compositional targets, but its cost grows with `beam_width`/`n_expansions`/feature count and it will not match PySR/Operon on hard problems. The composed bases are closures, so the model is not serialisable." + ] + }, + { + "cell_type": "markdown", + "id": "a638aac5", + "metadata": {}, + "source": [ + "## Difference from single-expression and backfitting\n", + "\n", + "- **Single-expression** (`jaxsr.SymbolicRegressor`): one sparse expression over a fixed basis library.\n", + "- **Stagewise additive**: terms are discovered one at a time on the residual and then **frozen**.\n", + "- **Backfitting additive** (`BackfittingSymbolicRegressor`): a fixed set of terms is **revised** repeatedly, each conditioned on the others (GAM-style; squared error only). A Bayesian (BART/iBART) variant is future work.\n", + "\n", + "See `docs/guides/additive-symbolic-regression.md` for the full guide." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.15" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/guides/additive-symbolic-regression.md b/docs/guides/additive-symbolic-regression.md new file mode 100644 index 0000000..838f582 --- /dev/null +++ b/docs/guides/additive-symbolic-regression.md @@ -0,0 +1,332 @@ +# Additive Symbolic Regression + +Additive symbolic regression fits a model as a **sum of small symbolic +expressions**: + +``` +f(x) = c + eta_1 * g_1(x) + eta_2 * g_2(x) + ... + eta_K * g_K(x) +``` + +where each `g_k(x)` is a small, interpretable symbolic expression discovered by +the existing JAXSR machinery. This is analogous to **gradient boosting**, except +each weak learner is a symbolic expression rather than a decision tree. + +The submodule lives in `jaxsr.additive`. + +## Three flavours of symbolic regression + +| Approach | What it does | Status | +|----------|--------------|--------| +| **Single-expression** (`jaxsr.SymbolicRegressor`) | Fits one sparse expression over a fixed basis library. | Available | +| **Stagewise additive** (`jaxsr.additive.StagewiseSymbolicRegressor`) | Repeatedly fits a small expression to the *residual* and adds it to the ensemble. Old terms are **frozen**. | Available | +| **Backfitting additive** (`jaxsr.additive.BackfittingSymbolicRegressor`) | Maintains a fixed set of terms and **revises** each one in place across sweeps (GAM-style). | Available (squared error); Bayesian variant planned | + +The key distinction between the two additive variants: + +- **Stagewise**: once a term is discovered it never changes; only its linear + weight may be re-estimated. +- **Backfitting**: terms are revised repeatedly, each conditioned on the current + fit of all the others. + +## Scope: what this is (and isn't) good for + +**In one line:** JAXSR is a *linear method over a fixed feature space* — it +selects a sparse combination of basis functions you supply. It is not a +free-composition equation discoverer. + +That distinction decides whether it is the right tool: + +- **Good fit:** the right building blocks are on the menu (or you can add them + with `BasisLibrary.add_custom`), and you want an interpretable, robust, + uncertainty-aware additive model. On targets that live in the library it is + fast and accurate, and the additive layer adds robust/quantile losses and + structural-uncertainty bootstrapping that genetic-programming tools don't + offer out of the box. +- **Wrong fit:** you want to *discover* an unknown compositional law such as + `exp(x0*x1)`, `x0 / (1 + x1**2)`, or `sin(2*x0)`. These are not single basis + functions, and the space of such compositions is infinite and continuously + parameterized, so no fixed library enumerates them in advance. For that, + reach for a genetic-programming or neural symbolic-regression tool (PySR, + Operon, AI-Feynman), which *search* the space of expressions instead of + selecting from a fixed dictionary — or try the experimental + [`RecursiveSymbolicRegressor`](#recursive-basis-expansion-experimental), which + grows compositions along the residual and partially lifts this ceiling. + +The limit is one of **discovery, not representation**: the linear-in-basis model +fits any of those targets perfectly the moment the exact term is in the library +(e.g. `library.add_custom("exp(x0*x1)", lambda X: jnp.exp(X[:, 0] * X[:, 1]))`) — +it simply cannot figure out *which* composition it needs without being told. +JAXSR's parametric bases can additionally fit a few constants *inside* a +pre-specified nonlinearity (e.g. `sin(a*x0)` with `a` optimized), but that still +requires you to name the functional form. The additive extensions in this guide +raise the *statistical* sophistication (boosting, robust/quantile losses, +backfitting, structural UQ); they do not change this expressiveness boundary. + +## Quick start + +```python +import numpy as np +from jaxsr.additive import StagewiseSymbolicRegressor + +rng = np.random.default_rng(0) +X = rng.uniform(-2, 2, size=(200, 2)) +y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + 0.1 * rng.normal(size=200) + +model = StagewiseSymbolicRegressor( + n_terms=5, + learning_rate=0.2, + max_complexity=6, + refit_coefficients=True, +) +model.fit(X, y) + +print(model) # pretty structural summary +print(model.expressions_) # per-term expression strings +print(model.coefficients_) # per-term weights +print(model.intercept_) # additive intercept +y_pred = model.predict(X) +``` + +The `print(model)` output looks like: + +``` +StagewiseSymbolicRegressor( + intercept = 1.07 + terms = + + 1 * (y = 2*x0 - 1.07 + 0.5*x1^2) + ... +) +``` + +## The stagewise algorithm + +1. Initialise the intercept to `mean(y)` and the prediction to that constant. +2. Compute the residual `y - prediction`. +3. Fit a small symbolic expression `g_k` to the residual (via + `jaxsr.fit_symbolic`). +4. Append `g_k` to the ensemble. +5. If `refit_coefficients=True`, rebuild the design matrix `Phi[:, j] = g_j(X)` + and re-solve `y ~= intercept + Phi @ coefficients` by least squares. + Otherwise, update `prediction += learning_rate * g_k(X)`. +6. Record train (and optional validation) loss. +7. Repeat until `n_terms` terms are added or early stopping triggers. + +## Key parameters + +| Parameter | Meaning | +|-----------|---------| +| `n_terms` | Maximum number of boosting stages (terms). | +| `learning_rate` | Shrinkage on each stage when `refit_coefficients=False`. | +| `max_complexity` | Complexity budget per term (max basis terms). Keep small to favour many simple terms. | +| `refit_coefficients` | Re-solve all linear weights by OLS after each stage. | +| `loss` | `"squared_error"` (default), `"absolute_error"`, `"huber"`, `"quantile"`, or a `Loss` instance. See [Losses](#losses-robust-and-quantile-regression). | +| `early_stopping` | Hold out a validation split and stop when it stops improving. | +| `validation_fraction`, `patience`, `min_delta` | Early-stopping controls. | +| `max_poly_degree`, `include_transcendental`, `include_ratios` | Which basis functions each term may use. | +| `information_criterion` | Complexity control within each term (`"aic"`, `"aicc"`, `"bic"`). | + +## Coefficient refitting + +- `refit_coefficients=False`: the weights are the learning-rate-scaled stagewise + weights (`coefficients_[k] == learning_rate`). +- `refit_coefficients=True`: after each new term, the intercept and *all* per-term + weights are re-solved by ordinary least squares over the discovered symbolic + features. This decouples term discovery (nonlinear, greedy) from term weighting + (linear, global) and typically improves accuracy. + +The refit uses `jnp.linalg.lstsq` (SVD-based, minimum-norm), so the highly +correlated columns produced by later boosting stages do not cause instability. + +## Combined expression + +`model.to_expression()` returns a single simplified SymPy expression combining +all terms: + +```python +expr = model.to_expression() # requires sympy +``` + +## Saving and loading + +Fitted models serialize to JSON (each term is stored via the underlying +`SymbolicRegressor` state), mirroring the rest of jaxsr. Note that the models +are **not picklable** — the basis-function closures cannot be pickled — so use +`save`/`load` rather than `pickle`: + +```python +model.save("additive_model.json") +loaded = StagewiseSymbolicRegressor.load("additive_model.json") +``` + +## Structural uncertainty (bootstrap) + +A single fitted expression can hide the fact that the *structure* itself is +uncertain — several different basis sets may explain the data about equally +well (this is common with collinear features). `bootstrap_additive` refits the +model on bootstrap resamples and reports, for each basis function, how often it +is selected — a cheap approximation to a posterior inclusion probability — +together with a predictive ensemble: + +```python +from jaxsr.additive import ( + StagewiseSymbolicRegressor, + bootstrap_additive, + bootstrap_predict_additive, +) + +est = StagewiseSymbolicRegressor(n_terms=3, max_complexity=2) +res = bootstrap_additive(est, X, y, n_bootstrap=100, random_state=0) + +# How stable is the discovered structure? +for name, prob in res["inclusion_probabilities"].items(): + print(f"{name:10s} selected in {prob:.0%} of resamples") + +# Prediction intervals that reflect *structural* variability, not just noise +pi = bootstrap_predict_additive(res["models"], X_new, alpha=0.1) +pi["mean"], pi["lower"], pi["upper"] +``` + +**How to read it.** Inclusion probabilities near 0 or 1 mean the structure is +identifiable and the single fitted expression is trustworthy. **Diffuse** values +(e.g. a basis selected 50–60% of the time) mean the data do not determine one +expression — no single symbolic model should be over-trusted, and the bootstrap +intervals are the honest summary. This also works as a decision gate for heavier +Bayesian modelling: if the probabilities are already crisp, there is little +structural uncertainty left to quantify. It works for both the stagewise and +backfitting regressors. + +## Early stopping + +With `early_stopping=True`, a validation split (`validation_fraction`) is held +out. After each stage the validation loss is recorded; training stops once it +fails to improve by at least `min_delta` for `patience` consecutive stages, and +the model rolls back to the best iteration. + +## Losses: robust and quantile regression + +This is where additive symbolic regression goes beyond ordinary least-squares +symbolic regression. Each weak learner fits the negative gradient `-dL/dy_pred` +(gradient boosting), so you can target losses that OLS selection cannot: + +| `loss` | Class | Use when | +|--------|-------|----------| +| `"squared_error"` (default) | `SquaredError` | Standard regression | +| `"absolute_error"` | `AbsoluteError` | Outliers present (fits the median) | +| `"huber"` | `HuberLoss(delta=1.35)` | Outliers, but keep efficiency near zero | +| `"quantile"` | `QuantileLoss(quantile=0.5)` | Quantiles / prediction intervals / asymmetric cost | + +Pass a name for defaults, or an instance to customise: + +```python +from jaxsr.additive import StagewiseSymbolicRegressor, QuantileLoss, HuberLoss + +# Robust regression: heavy outliers barely move the fit +robust = StagewiseSymbolicRegressor(loss="huber", learning_rate=0.5).fit(X, y) + +# 90th-percentile regression (build intervals by fitting several quantiles) +q90 = StagewiseSymbolicRegressor(loss=QuantileLoss(0.9), learning_rate=0.5).fit(X, y) +``` + +**How non-squared losses are fit.** Each stage fits a symbolic term to the +negative gradient, then a **line search** picks the step size that minimises the +loss (`learning_rate` shrinks that step). Because the ordinary least-squares +coefficient refit targets squared error, `refit_coefficients=True` is ignored for +non-squared losses (a warning is issued) and gradient boosting is used instead — +so set `refit_coefficients=False` explicitly for robust/quantile models. + +The optimal constant initialisation adapts to the loss: mean for squared error, +median for absolute/Huber, and the empirical quantile for quantile loss. + +Add further losses (Poisson, logistic, ...) by subclassing `Loss` and +registering them in `jaxsr.additive.losses._LOSSES`. + +## Backfitting (GAM-style) + +`BackfittingSymbolicRegressor` maintains a **fixed** number of terms and +*revises* each one across sweeps, instead of freezing them. Each sweep removes a +term, re-discovers its expression on the partial residual, and puts it back: + +```python +from jaxsr.additive import BackfittingSymbolicRegressor + +model = BackfittingSymbolicRegressor(n_terms=4, n_sweeps=6, max_complexity=3) +model.fit(X, y) # warm-started from a stagewise fit, then refined by sweeps +``` + +``` +for sweep in 1..n_sweeps: + for term j: + partial_residual = y - intercept - sum_{i != j} coef_i * g_i(X) + g_j = fit_symbolic(X, partial_residual, ...) # re-discover structure + intercept, coef = OLS refit over all terms +(stop when the training loss stops improving by `tol`) +``` + +It is warm-started from a stagewise fit and currently supports **squared error +only**. Structure re-discovery makes the sweep a heuristic (no monotonicity +guarantee), so the best-loss iterate is kept. + +**When does it actually help?** Backfitting starts from the stagewise+refit fit +and keeps the best-loss iterate, so **it is never worse than +`StagewiseSymbolicRegressor(refit_coefficients=True)` on the training data** — +the only question is whether the sweeps improve on it. That hinges entirely on +whether re-discovery changes the *set* of selected basis functions: + +- **Generous per-term budget** (`max_complexity` ≥ 2–3): greedy usually already + finds a sufficient basis set, so the joint least-squares refit makes the two + essentially identical. Backfitting adds nothing here — prefer the stagewise + regressor. +- **Small per-term budget and collinear features** (`max_complexity=1`, the + GAM-style single-basis regime): greedy forward selection can lock into a + *suboptimal* basis set that a single forward pass cannot undo. Backfitting's + coordinate-descent re-discovery escapes it, changing the basis union and + improving the fit — we have measured up to roughly **+0.04 train / +0.06 test + R²** in this regime, with no downside in the cases where it does not help. + +So reach for backfitting when you want **small, revisable single-basis terms +over correlated features** (or a fixed-size GAM-style decomposition); use the +stagewise regressor for larger per-term expressions. Its other forward-looking +value is as the foundation for a future **Bayesian backfitting** variant +(BART/iBART-style), which would sample a *posterior over symbolic structure* — +genuinely beyond point-estimate SR — using the same partial-residual sweep with +conjugate marginal likelihoods. + +## Recursive basis expansion (experimental) + +The [Scope](#scope-what-this-is-and-isnt-good-for) section notes that a fixed +library cannot *discover* compositional forms like `x0*sin(x1)` or `exp(x0*x1)`. +`RecursiveSymbolicRegressor` is an experimental step past that ceiling. Instead +of enumerating a huge composition space up front (which explodes +combinatorially), it **grows the library lazily along the residual**: + +1. Fit a sparse model over the current library; take the residual. +2. Compose the currently useful terms (selected terms + features) with a small + operator set (unary functions, products, ratios) — one new layer. +3. Screen: drop non-finite candidates, deduplicate, keep the top few by + correlation with the residual. +4. Add the survivors and refit. Repeat. The effective composition depth is the + number of rounds, because a term found in one round feeds the next. + +```python +from jaxsr.additive import RecursiveSymbolicRegressor + +model = RecursiveSymbolicRegressor(n_expansions=3, max_terms=6, beam_width=25) +model.fit(X, y) +print(model.expression_) # e.g. recovers "exp((x0)*(x1))" exactly +print(model.history_) # library size / n_terms / train R^2 per round +``` + +This is essentially Fast Function Extraction (FFX) / symbolic feature +construction — a deterministic, bounded cousin of genetic programming. On simple +compositional targets it substantially beats a flat library (e.g. `exp(x0*x1)`: +R² 1.00 vs 0.70) and is competitive with a strong GP engine (matched Operon on +`x0*sin(x1)` in our tests). + +**Caveats.** It re-enters search-based territory: cost grows with `beam_width`, +`n_expansions`, and feature count, and it will not match a mature GP (PySR, +Operon) on hard, high-dimensional, or deeply nested targets. The result is an +ordinary `SymbolicRegressor` over the grown library (so predict/`expression_`/ +scoring and the base regressor's non-finite-basis guard and negligible-term +pruning all apply), but the composed bases are Python closures, so the fitted +model is **not** serialisable via `save`/`load`, and `to_sympy` may not parse +deeply nested term names. diff --git a/docs/intro.md b/docs/intro.md index a90094a..7b807ec 100644 --- a/docs/intro.md +++ b/docs/intro.md @@ -14,6 +14,7 @@ Key features: - **Multiple Selection Strategies**: Choose from greedy, exhaustive, or LASSO-based methods - **Uncertainty Quantification**: Prediction intervals, Bayesian Model Averaging, conformal prediction, and bootstrap methods - **Physical Constraints**: Incorporate domain knowledge through constraints +- **Additive Symbolic Regression**: Boosting-style ensembles of small symbolic expressions (`jaxsr.additive`) - **JAX-Powered**: GPU acceleration, JIT compilation, automatic differentiation - **Scikit-learn Compatible**: Familiar fit/predict interface diff --git a/examples/additive_symbolic_regression.py b/examples/additive_symbolic_regression.py new file mode 100644 index 0000000..1a6a5ae --- /dev/null +++ b/examples/additive_symbolic_regression.py @@ -0,0 +1,143 @@ +"""Worked examples for the ``jaxsr.additive`` submodule. + +Demonstrates, end to end: + +1. Stagewise additive regression (gradient boosting with symbolic weak learners) +2. Robust regression under outliers (Huber / absolute-error losses) +3. Quantile regression (pinball loss -> calibrated coverage) +4. Structural uncertainty via bootstrap (basis inclusion probabilities) +5. Backfitting (GAM-style: revise terms instead of freezing them) +6. Recursive basis expansion (reach compositions a flat library misses) + +Run with:: + + python examples/additive_symbolic_regression.py +""" + +import numpy as np + +from jaxsr import fit_symbolic +from jaxsr.additive import ( + BackfittingSymbolicRegressor, + QuantileLoss, + RecursiveSymbolicRegressor, + StagewiseSymbolicRegressor, + bootstrap_additive, + bootstrap_predict_additive, +) + + +def r2(y, p): + y, p = np.asarray(y, float), np.asarray(p, float) + return 1 - np.sum((y - p) ** 2) / (np.sum((y - y.mean()) ** 2) + 1e-12) + + +def section(title): + print("\n" + "=" * 68) + print(title) + print("=" * 68) + + +def stagewise_example(): + section("1. Stagewise additive regression: y = 2*x0 + 0.5*x1^2") + rng = np.random.default_rng(0) + X = rng.uniform(-2, 2, size=(200, 2)) + y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + 0.1 * rng.normal(size=200) + + model = StagewiseSymbolicRegressor( + n_terms=5, learning_rate=0.2, max_complexity=6, refit_coefficients=True + ).fit(X, y) + print(model) + print("R^2:", round(model.score(X, y), 4)) + print("combined:", model.to_expression()) + + +def robust_example(): + section("2. Robust regression: 8% heavy outliers, evaluate on clean signal") + rng = np.random.default_rng(1) + X = rng.uniform(-2, 2, size=(500, 2)) + clean = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + y = clean + rng.normal(0, 0.1, 500) + y[rng.choice(500, 40, replace=False)] += 25.0 # outliers + + for loss in ["squared_error", "huber", "absolute_error"]: + m = StagewiseSymbolicRegressor( + n_terms=8, max_complexity=3, learning_rate=0.5, loss=loss, refit_coefficients=False + ).fit(X, y) + mae = float(np.mean(np.abs(np.asarray(m.predict(X)) - clean))) + print(f" {loss:15s} MAE vs clean signal = {mae:.3f}") + + +def quantile_example(): + section("3. Quantile regression: empirical coverage should match q") + rng = np.random.default_rng(2) + X = rng.uniform(-2, 2, size=(500, 2)) + y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + rng.normal(0, 1.0, 500) + + for q in [0.1, 0.5, 0.9]: + m = StagewiseSymbolicRegressor( + n_terms=10, + max_complexity=3, + learning_rate=0.5, + loss=QuantileLoss(q), + refit_coefficients=False, + ).fit(X, y) + coverage = float(np.mean(y <= np.asarray(m.predict(X)))) + print(f" q={q}: fraction below prediction = {coverage:.3f}") + + +def uncertainty_example(): + section("4. Structural uncertainty: bootstrap inclusion probabilities") + # Collinear features -> the structure is not well determined. + rng = np.random.default_rng(0) + x0 = rng.normal(0, 1, 400) + x1 = 0.9 * x0 + 0.1 * rng.normal(0, 1, 400) + x2 = 0.8 * x0 + 0.2 * rng.normal(0, 1, 400) + X = np.column_stack([x0, x1, x2]) + y = x0 - x1 + 0.5 * x2 + rng.normal(0, 0.1, 400) + + est = StagewiseSymbolicRegressor(n_terms=3, max_complexity=1) + res = bootstrap_additive(est, X, y, n_bootstrap=80, random_state=0) + for name, prob in list(res["inclusion_probabilities"].items())[:5]: + print(f" {name:8s} selected in {prob:.0%} of resamples") + pi = bootstrap_predict_additive(res["models"], X[:3], alpha=0.1) + print(" 90% prediction intervals (first 3):") + for i in range(3): + print(f" x[{i}]: [{pi['lower'][i]:+.2f}, {pi['upper'][i]:+.2f}]") + + +def backfitting_example(): + section("5. Backfitting (GAM-style): revise terms instead of freezing") + rng = np.random.default_rng(3) + X = rng.uniform(-2, 2, size=(300, 2)) + y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + rng.normal(0, 0.1, 300) + + bf = BackfittingSymbolicRegressor(n_terms=3, n_sweeps=6, max_complexity=3).fit(X, y) + print("R^2:", round(bf.score(X, y), 4), "| sweeps run:", len(bf.training_history_) - 1) + + +def recursive_example(): + section("6. Recursive expansion: reach compositions a flat library misses") + rng = np.random.default_rng(0) + X = rng.uniform(-2, 2, size=(400, 2)) + y = np.exp(X[:, 0] * X[:, 1]) + 0.02 * rng.normal(size=400) + Xtr, ytr, Xte, yte = X[:200], y[:200], X[200:], y[200:] + + rec = RecursiveSymbolicRegressor(n_expansions=3, max_terms=6, beam_width=25).fit(Xtr, ytr) + flat = fit_symbolic(Xtr, ytr, max_terms=6, include_transcendental=True) + print(" target: exp(x0*x1)") + print(f" recursive test R^2 = {r2(yte, rec.predict(Xte)):.4f} -> {rec.expression_[:48]}") + print(f" flat-library test R^2 = {r2(yte, flat.predict(Xte)):.4f}") + + +def main() -> None: + stagewise_example() + robust_example() + quantile_example() + uncertainty_example() + backfitting_example() + recursive_example() + + +if __name__ == "__main__": + main() diff --git a/src/jaxsr/additive/__init__.py b/src/jaxsr/additive/__init__.py new file mode 100644 index 0000000..817c918 --- /dev/null +++ b/src/jaxsr/additive/__init__.py @@ -0,0 +1,64 @@ +""" +Additive symbolic regression for JAXSR. + +Fits models of the form ``f(x) = c + sum_k eta_k * g_k(x)``, where each +``g_k`` is a small symbolic expression discovered by the existing JAXSR +machinery. This is analogous to gradient boosting, except each weak learner +is an interpretable symbolic expression rather than a decision tree. + +Public API +---------- +StagewiseSymbolicRegressor + Boosting-style regressor that fits each new symbolic term to the current + residual and freezes it. This is the first-milestone workhorse. +BackfittingSymbolicRegressor + GAM-style regressor that revises terms in place across sweeps (warm-started + from a stagewise fit); squared-error only. +AdditiveSymbolicModel + Plain container for a fitted additive model. +Loss, SquaredError, AbsoluteError, HuberLoss, QuantileLoss, get_loss + Loss abstraction and registry. Squared error is the default; absolute + error, Huber, and quantile (pinball) losses enable robust and quantile + regression via gradient boosting. +refit_ols + Least-squares refit of intercept and per-term coefficients. +bootstrap_additive, bootstrap_predict_additive + Bootstrap structural uncertainty: basis-function inclusion probabilities and + a predictive ensemble that reflects structural variability. +""" + +from __future__ import annotations + +from .backfitting import BackfittingSymbolicRegressor +from .coefficient_refit import refit_ols +from .ensemble import AdditiveSymbolicModel, additive_predict +from .losses import ( + AbsoluteError, + HuberLoss, + Loss, + QuantileLoss, + SquaredError, + get_loss, + loss_from_config, +) +from .recursive import RecursiveSymbolicRegressor +from .stagewise import StagewiseSymbolicRegressor +from .uncertainty import bootstrap_additive, bootstrap_predict_additive + +__all__ = [ + "AbsoluteError", + "AdditiveSymbolicModel", + "BackfittingSymbolicRegressor", + "HuberLoss", + "Loss", + "QuantileLoss", + "RecursiveSymbolicRegressor", + "SquaredError", + "StagewiseSymbolicRegressor", + "additive_predict", + "bootstrap_additive", + "bootstrap_predict_additive", + "get_loss", + "loss_from_config", + "refit_ols", +] diff --git a/src/jaxsr/additive/backfitting.py b/src/jaxsr/additive/backfitting.py new file mode 100644 index 0000000..ff6d1e5 --- /dev/null +++ b/src/jaxsr/additive/backfitting.py @@ -0,0 +1,296 @@ +""" +Backfitting additive symbolic regression (GAM-style). + +Unlike :class:`~jaxsr.additive.stagewise.StagewiseSymbolicRegressor`, where each +discovered term is *frozen*, the backfitting regressor maintains a fixed number +of terms and repeatedly *revises* each one. A "sweep" visits every term in +turn, removes it from the ensemble, and re-discovers its symbolic expression on +the *partial residual* (the target minus every other term's contribution):: + + for sweep in 1..n_sweeps: + for term j: + partial_residual = y - intercept - sum_{i != j} coef_i * g_i(X) + g_j = fit_symbolic(X, partial_residual, ...) # re-discover structure + intercept, coef = OLS refit over all terms + +This lets early terms -- originally fit against a residual still polluted by +effects that had not yet been discovered -- clean themselves up once the other +terms are in place. It is the classic backfitting algorithm behind generalized +additive models, with symbolic expressions as the smoothers. + +The regressor is warm-started from a stagewise fit, so a single ``fit`` call +first runs stagewise boosting and then refines it by backfitting. + +Scope +----- +This is the deterministic, squared-error version. A future Bayesian variant +(BART/iBART-style, sampling a posterior over symbolic structures) can build on +the same partial-residual sweep; see the project notes. +""" + +from __future__ import annotations + +import jax.numpy as jnp + +from ..regressor import fit_symbolic +from .base import _BaseAdditiveRegressor +from .coefficient_refit import refit_ols +from .ensemble import AdditiveSymbolicModel +from .losses import SquaredError, get_loss +from .stagewise import StagewiseSymbolicRegressor + + +class BackfittingSymbolicRegressor(_BaseAdditiveRegressor): + """ + Backfitting additive symbolic regression (GAM-style). + + Maintains ``n_terms`` symbolic components and revises each one in place + across repeated sweeps, conditioning on the current fit of all other terms. + Contrast with :class:`~jaxsr.additive.stagewise.StagewiseSymbolicRegressor`, + which freezes terms once discovered. + + The model is warm-started with a stagewise fit and then refined by + backfitting. Only squared-error loss is supported. + + Parameters + ---------- + n_terms : int + Number of symbolic components to maintain and revise. + n_sweeps : int + Maximum number of backfitting sweeps over the terms. + max_complexity : int + Complexity budget (max basis terms) for each component. + loss : str + Loss function. Only ``"squared_error"`` is supported. + tol : float + Convergence tolerance: stop when the training MSE improves by less than + ``tol`` between consecutive sweeps. + max_poly_degree : int + Maximum polynomial degree available to each component. + include_transcendental : bool + Whether to allow transcendental terms in each component. + include_ratios : bool + Whether to allow ratio terms in each component. + strategy : str + Selection strategy for each component. + information_criterion : str + Information criterion for complexity control. + feature_names : list of str, optional + Names for the input features. + random_state : int, optional + Seed forwarded to the stagewise warm start. + + Attributes + ---------- + model_ : AdditiveSymbolicModel + The fitted additive model. + intercept_ : float + Fitted intercept. + coefficients_ : list of float + Fitted per-term coefficients. + expressions_ : list of str + Human-readable expression for each term. + terms_ : list of SymbolicRegressor + The fitted symbolic terms. + training_history_ : list of dict + Per-sweep diagnostics (``sweep`` index and ``train_loss``). + n_terms_ : int + Number of terms in the fitted model. + + Examples + -------- + >>> from jaxsr.additive import BackfittingSymbolicRegressor + >>> model = BackfittingSymbolicRegressor(n_terms=3, n_sweeps=5) + >>> model.fit(X, y) # doctest: +SKIP + >>> print(model) # doctest: +SKIP + """ + + def __init__( + self, + n_terms: int = 5, + n_sweeps: int = 10, + max_complexity: int = 4, + loss: str = "squared_error", + tol: float = 1e-6, + max_poly_degree: int = 3, + include_transcendental: bool = False, + include_ratios: bool = False, + strategy: str = "greedy_forward", + information_criterion: str = "bic", + feature_names: list[str] | None = None, + random_state: int | None = None, + ): + self.n_terms = n_terms + self.n_sweeps = n_sweeps + self.max_complexity = max_complexity + self.loss = loss + self.tol = tol + self.max_poly_degree = max_poly_degree + self.include_transcendental = include_transcendental + self.include_ratios = include_ratios + self.strategy = strategy + self.information_criterion = information_criterion + self.feature_names = feature_names + self.random_state = random_state + + self.model_: AdditiveSymbolicModel | None = None + self._is_fitted = False + + def _validate_params(self) -> None: + """Validate constructor parameters.""" + if self.n_terms < 1: + raise ValueError(f"n_terms must be >= 1, got {self.n_terms}.") + if self.n_sweeps < 1: + raise ValueError(f"n_sweeps must be >= 1, got {self.n_sweeps}.") + if self.max_complexity < 1: + raise ValueError(f"max_complexity must be >= 1, got {self.max_complexity}.") + + def _fit_term( + self, + X: jnp.ndarray, + target: jnp.ndarray, + feature_names: list[str], + ): + """Discover a single symbolic term for ``target`` (reuses fit_symbolic).""" + return fit_symbolic( + X, + target, + feature_names=feature_names, + max_terms=self.max_complexity, + max_poly_degree=self.max_poly_degree, + include_transcendental=self.include_transcendental, + include_ratios=self.include_ratios, + strategy=self.strategy, + information_criterion=self.information_criterion, + ) + + def fit(self, X: jnp.ndarray, y: jnp.ndarray) -> BackfittingSymbolicRegressor: + """ + Fit the backfitting additive model. + + Parameters + ---------- + X : array-like of shape (n_samples, n_features) + Training inputs. + y : array-like of shape (n_samples,) + Target values. + + Returns + ------- + self : BackfittingSymbolicRegressor + The fitted estimator. + + Raises + ------ + ValueError + If parameters are invalid, ``X``/``y`` are mismatched or non-finite. + NotImplementedError + If ``loss`` is not ``"squared_error"``. + """ + self._validate_params() + loss_fn = get_loss(self.loss) + if not isinstance(loss_fn, SquaredError): + raise NotImplementedError( + "BackfittingSymbolicRegressor currently supports only " + f"loss='squared_error', got {loss_fn.name!r}. Use " + "StagewiseSymbolicRegressor for other losses." + ) + + X = jnp.atleast_2d(jnp.asarray(X)) + y = jnp.asarray(y).ravel() + if X.shape[0] != y.shape[0]: + raise ValueError( + f"X and y must have the same number of samples. " + f"Got X: {X.shape[0]}, y: {y.shape[0]}." + ) + if not bool(jnp.all(jnp.isfinite(X))): + raise ValueError("X contains non-finite values (NaN or inf).") + if not bool(jnp.all(jnp.isfinite(y))): + raise ValueError("y contains non-finite values (NaN or inf).") + + n_features = X.shape[1] + feature_names = self.feature_names or [f"x{i}" for i in range(n_features)] + if len(feature_names) != n_features: + raise ValueError( + f"feature_names has {len(feature_names)} entries but X has " + f"{n_features} features." + ) + + # ------------------------------------------------------------------ + # Warm start: a stagewise fit with the same per-term budget provides + # the initial set of terms that backfitting then refines. + # ------------------------------------------------------------------ + warm = StagewiseSymbolicRegressor( + n_terms=self.n_terms, + max_complexity=self.max_complexity, + refit_coefficients=True, + loss="squared_error", + max_poly_degree=self.max_poly_degree, + include_transcendental=self.include_transcendental, + include_ratios=self.include_ratios, + strategy=self.strategy, + information_criterion=self.information_criterion, + feature_names=feature_names, + random_state=self.random_state, + ).fit(X, y) + + terms = list(warm.terms_) + intercept = float(warm.intercept_) + coefficients = list(warm.coefficients_) + + def full_prediction() -> jnp.ndarray: + pred = jnp.full((X.shape[0],), intercept) + for c, t in zip(coefficients, terms, strict=False): + pred = pred + float(c) * t.predict(X) + return pred + + history: list[dict] = [] + prev_loss = loss_fn.loss(y, full_prediction()) + history.append({"sweep": 0, "train_loss": float(prev_loss)}) + + # Snapshot of the best (lowest-loss) iterate. Structure re-discovery + # makes the sweep a heuristic with no monotonicity guarantee, so we + # keep the best state rather than trusting the final one. + best_loss = prev_loss + best_state = (intercept, list(coefficients), list(terms)) + + # ------------------------------------------------------------------ + # Backfitting sweeps + # ------------------------------------------------------------------ + for sweep in range(1, self.n_sweeps + 1): + for j in range(len(terms)): + # Partial residual: target minus every OTHER term's contribution. + partial = y - intercept + for i, (c, t) in enumerate(zip(coefficients, terms, strict=False)): + if i != j: + partial = partial - float(c) * t.predict(X) + + terms[j] = self._fit_term(X, partial, feature_names) + + # Re-solve intercept and all coefficients jointly (stable OLS). + Phi = jnp.stack([t.predict(X) for t in terms], axis=1) + intercept, coef_arr = refit_ols(Phi, y) + coefficients = [float(c) for c in coef_arr] + + sweep_loss = loss_fn.loss(y, full_prediction()) + history.append({"sweep": sweep, "train_loss": float(sweep_loss)}) + + if sweep_loss < best_loss: + best_loss = sweep_loss + best_state = (intercept, list(coefficients), list(terms)) + + if abs(prev_loss - sweep_loss) < self.tol: + break + prev_loss = sweep_loss + + intercept, coefficients, terms = best_state + self.model_ = AdditiveSymbolicModel( + intercept=float(intercept), + terms=terms, + coefficients=coefficients, + learning_rates=[1.0] * len(terms), + feature_names=feature_names, + training_history=history, + ) + self._is_fitted = True + return self diff --git a/src/jaxsr/additive/base.py b/src/jaxsr/additive/base.py new file mode 100644 index 0000000..76ba1d5 --- /dev/null +++ b/src/jaxsr/additive/base.py @@ -0,0 +1,251 @@ +""" +Shared base class for additive symbolic regressors. + +Both :class:`~jaxsr.additive.stagewise.StagewiseSymbolicRegressor` and +:class:`~jaxsr.additive.backfitting.BackfittingSymbolicRegressor` produce an +:class:`~jaxsr.additive.ensemble.AdditiveSymbolicModel` and share the same +fitted-attribute accessors, prediction, interpretation, and JSON +serialization. Only the fitting strategy differs, so everything except +``fit`` lives here. +""" + +from __future__ import annotations + +import json + +import jax.numpy as jnp + +from .._compat import _SklearnCompatMixin +from ..regressor import SymbolicRegressor +from .ensemble import AdditiveSymbolicModel +from .losses import get_loss, loss_from_config + + +class _BaseAdditiveRegressor(_SklearnCompatMixin): + """ + Common machinery for additive symbolic regressors. + + Subclasses must: + + * store their constructor parameters as same-named attributes (including a + ``loss`` attribute), + * implement ``fit`` to populate ``self.model_`` with an + :class:`AdditiveSymbolicModel` and set ``self._is_fitted = True``. + """ + + model_: AdditiveSymbolicModel | None + _is_fitted: bool + + # ------------------------------------------------------------------ + # Fitted-attribute accessors + # ------------------------------------------------------------------ + def _check_is_fitted(self) -> None: + """Raise if the model has not been fitted yet.""" + if not getattr(self, "_is_fitted", False) or self.model_ is None: + raise RuntimeError("Model not fitted. Call fit() first.") + + @property + def intercept_(self) -> float: + """Fitted intercept.""" + self._check_is_fitted() + return self.model_.intercept + + @property + def coefficients_(self) -> list[float]: + """Fitted per-term coefficients.""" + self._check_is_fitted() + return self.model_.coefficients + + @property + def expressions_(self) -> list[str]: + """Human-readable expression for each term.""" + self._check_is_fitted() + return self.model_.expressions + + @property + def terms_(self) -> list[SymbolicRegressor]: + """The fitted symbolic terms.""" + self._check_is_fitted() + return self.model_.terms + + @property + def learning_rates_(self) -> list[float]: + """Per-term learning rate / step scale recorded during fitting.""" + self._check_is_fitted() + return self.model_.learning_rates + + @property + def training_history_(self) -> list[dict]: + """Per-iteration diagnostics.""" + self._check_is_fitted() + return self.model_.training_history + + @property + def n_terms_(self) -> int: + """Number of terms in the fitted model.""" + self._check_is_fitted() + return self.model_.n_terms + + # ------------------------------------------------------------------ + # Prediction / interpretation + # ------------------------------------------------------------------ + def predict(self, X: jnp.ndarray) -> jnp.ndarray: + """ + Predict with the fitted additive model. + + Parameters + ---------- + X : array-like of shape (n_samples, n_features) + Input data. + + Returns + ------- + jnp.ndarray of shape (n_samples,) + Predicted values. + """ + self._check_is_fitted() + return self.model_.predict(X) + + def score(self, X: jnp.ndarray, y: jnp.ndarray) -> float: + """ + Compute the R^2 score on ``(X, y)``. + + Parameters + ---------- + X : array-like of shape (n_samples, n_features) + Input data. + y : array-like of shape (n_samples,) + True target values. + + Returns + ------- + float + R^2 score. + """ + self._check_is_fitted() + y = jnp.asarray(y).ravel() + y_pred = self.predict(X) + ss_res = jnp.sum((y - y_pred) ** 2) + ss_tot = jnp.sum((y - jnp.mean(y)) ** 2) + return float(1 - ss_res / (ss_tot + 1e-10)) + + def to_expression(self): + """ + Return the combined model as a single simplified SymPy expression. + + Returns + ------- + sympy.Expr + Combined symbolic expression for the whole ensemble. + """ + self._check_is_fitted() + return self.model_.to_expression() + + # ------------------------------------------------------------------ + # Serialization + # ------------------------------------------------------------------ + def _state_dict(self) -> dict: + """ + Return a JSON-serialisable dictionary of the fitted model state. + + Each symbolic term is serialised via the underlying + :class:`jaxsr.SymbolicRegressor` state dictionary, which avoids the + (unpicklable) basis-function closures. + + Returns + ------- + dict + Dictionary containing constructor config and fitted state. + + Raises + ------ + RuntimeError + If the model has not been fitted. + """ + self._check_is_fitted() + config = self.get_params(deep=False) + # Serialise the loss faithfully: a plain name when it has no parameters, + # otherwise a {"name", "params"} dict (e.g. quantile / huber). + loss_config = get_loss(self.loss).to_config() + config["loss"] = loss_config["name"] if not loss_config["params"] else loss_config + return { + "config": config, + "intercept": float(self.model_.intercept), + "coefficients": [float(c) for c in self.model_.coefficients], + "learning_rates": [float(lr) for lr in self.model_.learning_rates], + "feature_names": list(self.model_.feature_names), + "training_history": self.model_.training_history, + "terms": [term._state_dict() for term in self.model_.terms], + } + + @classmethod + def _from_dict(cls, data: dict) -> _BaseAdditiveRegressor: + """ + Reconstruct a fitted regressor from a state dictionary. + + Parameters + ---------- + data : dict + Dictionary produced by :meth:`_state_dict`. + + Returns + ------- + _BaseAdditiveRegressor + The reconstructed fitted model (an instance of ``cls``). + """ + config = dict(data["config"]) + config["loss"] = loss_from_config(config["loss"]) + model = cls(**config) + terms = [SymbolicRegressor._from_dict(t) for t in data["terms"]] + model.model_ = AdditiveSymbolicModel( + intercept=float(data["intercept"]), + terms=terms, + coefficients=[float(c) for c in data["coefficients"]], + learning_rates=[float(lr) for lr in data["learning_rates"]], + feature_names=list(data["feature_names"]), + training_history=data.get("training_history", []), + ) + model._is_fitted = True + return model + + def save(self, filepath: str) -> None: + """ + Save the fitted model to a JSON file. + + Parameters + ---------- + filepath : str + Destination path. + + Raises + ------ + RuntimeError + If the model has not been fitted. + """ + with open(filepath, "w") as f: + json.dump(self._state_dict(), f, indent=2) + + @classmethod + def load(cls, filepath: str) -> _BaseAdditiveRegressor: + """ + Load a fitted model from a JSON file. + + Parameters + ---------- + filepath : str + Path to a file created by :meth:`save`. + + Returns + ------- + _BaseAdditiveRegressor + The loaded fitted model (an instance of ``cls``). + """ + with open(filepath) as f: + data = json.load(f) + return cls._from_dict(data) + + def __repr__(self) -> str: + """Pretty structural repr when fitted, sklearn-style otherwise.""" + if getattr(self, "_is_fitted", False) and self.model_ is not None: + return self.model_.describe(name=type(self).__name__) + return _SklearnCompatMixin.__repr__(self) diff --git a/src/jaxsr/additive/coefficient_refit.py b/src/jaxsr/additive/coefficient_refit.py new file mode 100644 index 0000000..353fc43 --- /dev/null +++ b/src/jaxsr/additive/coefficient_refit.py @@ -0,0 +1,79 @@ +""" +Coefficient refitting for additive symbolic regression. + +After a new symbolic term is discovered, the stagewise model can optionally +re-solve the linear coefficients over *all* discovered symbolic features at +once, treating each term's prediction as a single feature column:: + + Phi[:, j] = g_j(X) + y ~= intercept + Phi @ coefficients + +This decouples term *discovery* (nonlinear, greedy) from term *weighting* +(linear, global), which typically improves accuracy relative to fixed +learning-rate-scaled stagewise weights. + +Only ordinary least squares is implemented for the first milestone. Ridge, +lasso, and sparse variants can be added later behind the same interface. +""" + +from __future__ import annotations + +import jax.numpy as jnp + + +def refit_ols(Phi: jnp.ndarray, y: jnp.ndarray) -> tuple[float, jnp.ndarray]: + """ + Refit an intercept and per-term coefficients by ordinary least squares. + + Solves ``y ~= intercept + Phi @ coefficients`` using a least-squares + solver that is robust to rank-deficient / collinear design matrices + (later boosting stages fit residuals of earlier ones, so the term columns + can be highly correlated). + + Parameters + ---------- + Phi : jnp.ndarray of shape (n_samples, n_terms) + Design matrix whose column ``j`` is ``g_j(X)`` for term ``j``. May + have zero columns (no terms discovered yet). + y : jnp.ndarray of shape (n_samples,) + Target values. + + Returns + ------- + intercept : float + Fitted intercept. + coefficients : jnp.ndarray of shape (n_terms,) + Fitted per-term coefficients. + + Raises + ------ + ValueError + If ``Phi`` is not 2-D or its number of rows does not match ``len(y)``. + """ + Phi = jnp.asarray(Phi) + y = jnp.asarray(y).ravel() + + if Phi.ndim != 2: + raise ValueError(f"Phi must be 2-D, got shape {Phi.shape}.") + if Phi.shape[0] != y.shape[0]: + raise ValueError(f"Phi has {Phi.shape[0]} rows but y has {y.shape[0]} samples.") + + n_samples, n_terms = Phi.shape + + # No terms yet: the best constant model is the mean. + if n_terms == 0: + return float(jnp.mean(y)), jnp.zeros((0,)) + + # Fit the intercept by centering rather than augmenting Phi with a + # column of ones. Augmentation mixes an O(1) ones-column with the term + # columns, whose scale is arbitrary (it tracks the scale of y); when the + # terms are tiny (e.g. y ~ 1e-6) that disparity makes the system severely + # ill-conditioned in float32 and lstsq returns garbage. Centering keeps + # the design matrix on a single scale and is the standard, stable way to + # estimate an intercept. Still SVD-based via lstsq -- never inv(). + Phi_mean = jnp.mean(Phi, axis=0) + y_mean = jnp.mean(y) + coefficients, _, _, _ = jnp.linalg.lstsq(Phi - Phi_mean, y - y_mean, rcond=None) + + intercept = float(y_mean - Phi_mean @ coefficients) + return intercept, coefficients diff --git a/src/jaxsr/additive/ensemble.py b/src/jaxsr/additive/ensemble.py new file mode 100644 index 0000000..d283b77 --- /dev/null +++ b/src/jaxsr/additive/ensemble.py @@ -0,0 +1,185 @@ +""" +Core additive symbolic model representation. + +An additive symbolic model has the form:: + + f(x) = intercept + sum_j coefficients[j] * terms[j](x) + +where each ``terms[j]`` is a small symbolic expression (a fitted +:class:`jaxsr.SymbolicRegressor`) discovered by the existing JAXSR machinery. +This is analogous to gradient boosting, except each weak learner is an +interpretable symbolic expression rather than a decision tree. + +The :class:`AdditiveSymbolicModel` is a plain data container: the fitting +strategy (stagewise, backfitting, ...) lives in the regressor classes and +produces one of these objects. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import jax.numpy as jnp + +if TYPE_CHECKING: + from ..regressor import SymbolicRegressor + + +def additive_predict( + X: jnp.ndarray, + intercept: float, + terms: list[SymbolicRegressor], + coefficients: list[float] | jnp.ndarray, +) -> jnp.ndarray: + """ + Evaluate an additive model ``intercept + sum_j coef_j * term_j(X)``. + + Parameters + ---------- + X : jnp.ndarray of shape (n_samples, n_features) + Input data. + intercept : float + Additive intercept. + terms : list of SymbolicRegressor + Fitted symbolic terms ``g_j``. + coefficients : list of float or jnp.ndarray + Per-term coefficients, aligned with ``terms``. + + Returns + ------- + jnp.ndarray of shape (n_samples,) + Predicted values. + """ + X = jnp.atleast_2d(jnp.asarray(X)) + prediction = jnp.full((X.shape[0],), float(intercept)) + for coef, term in zip(coefficients, terms, strict=False): + prediction = prediction + float(coef) * term.predict(X) + return prediction + + +@dataclass(repr=False) +class AdditiveSymbolicModel: + """ + Container for an additive symbolic model. + + Prediction is ``intercept + sum_j coefficients[j] * terms[j](X)``. + + Parameters + ---------- + intercept : float + Additive intercept ``c``. + terms : list of SymbolicRegressor + Fitted symbolic expressions ``g_j`` (the boosting "weak learners"). + coefficients : list of float + Per-term weights ``eta_j``. When coefficients are not refit these are + the learning-rate-scaled stagewise weights; when refit they are the + least-squares solution over all discovered terms. + learning_rates : list of float + Learning rate used at each stage (recorded for reproducibility). + feature_names : list of str + Feature names, shared across all terms. + training_history : list of dict + Per-stage diagnostics (train loss, validation loss, coefficients, ...). + """ + + intercept: float + terms: list[SymbolicRegressor] + coefficients: list[float] + learning_rates: list[float] + feature_names: list[str] + training_history: list[dict[str, Any]] = field(default_factory=list) + + @property + def n_terms(self) -> int: + """Number of symbolic terms in the model.""" + return len(self.terms) + + def predict(self, X: jnp.ndarray) -> jnp.ndarray: + """ + Predict with the additive model. + + Parameters + ---------- + X : jnp.ndarray of shape (n_samples, n_features) + Input data. + + Returns + ------- + jnp.ndarray of shape (n_samples,) + Predicted values. + + Raises + ------ + ValueError + If ``X`` has a different number of features than the model was + fit with. + """ + X = jnp.atleast_2d(jnp.asarray(X)) + n_expected = len(self.feature_names) + if X.shape[1] != n_expected: + raise ValueError( + f"X has {X.shape[1]} features but the model was fit with {n_expected}." + ) + return additive_predict(X, self.intercept, self.terms, self.coefficients) + + @property + def expressions(self) -> list[str]: + """Human-readable expression string for each symbolic term.""" + return [term.expression_ for term in self.terms] + + def to_expression(self): + """ + Combine all terms into a single simplified SymPy expression. + + Returns + ------- + sympy.Expr + ``intercept + sum_j coefficients[j] * g_j`` after simplification. + Falls back to the unsimplified sum if simplification fails. + + Raises + ------ + ImportError + If SymPy is not installed. + """ + import sympy + + expr = sympy.Float(float(self.intercept)) + for coef, term in zip(self.coefficients, self.terms, strict=False): + expr = expr + sympy.Float(float(coef)) * term.to_sympy() + try: + return sympy.simplify(expr) + except (TypeError, ValueError, AttributeError): + return expr + + def describe(self, name: str = "AdditiveSymbolicModel") -> str: + """ + Return a multi-line human-readable summary of the model. + + Parameters + ---------- + name : str + Class/label to show as the heading. + + Returns + ------- + str + Pretty-printed model structure. + """ + lines = [f"{name}("] + lines.append(f" intercept = {float(self.intercept):.4g}") + if self.terms: + lines.append(" terms =") + for coef, term in zip(self.coefficients, self.terms, strict=False): + coef = float(coef) + sign = "+" if coef >= 0 else "-" + lines.append(f" {sign} {abs(coef):.4g} * ({term.expression_})") + else: + lines.append(" terms = (none)") + lines.append(")") + return "\n".join(lines) + + def __repr__(self) -> str: + """Pretty multi-line representation.""" + return self.describe() diff --git a/src/jaxsr/additive/losses.py b/src/jaxsr/additive/losses.py new file mode 100644 index 0000000..e59db19 --- /dev/null +++ b/src/jaxsr/additive/losses.py @@ -0,0 +1,326 @@ +""" +Loss functions for additive symbolic regression. + +Additive (boosting-style) symbolic regression fits each new symbolic term to +the *pseudo-residual* of the current ensemble. For squared error the +pseudo-residual is simply ``y - y_pred``; for other differentiable losses it +is the negative gradient ``-dL/dy_pred``, which is what gradient boosting +fits. This lets JAXSR learn symbolic models under losses that ordinary +least-squares selection cannot target directly: + +* :class:`SquaredError` -- standard regression (mean). +* :class:`AbsoluteError` -- robust to outliers (median). +* :class:`HuberLoss` -- quadratic near zero, linear in the tails (robust). +* :class:`QuantileLoss` -- pinball loss for quantile / interval estimation. + +New losses can be added by subclassing :class:`Loss` and registering them in +``_LOSSES``. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +import jax.numpy as jnp + + +class Loss(ABC): + """ + Abstract base class for additive-regression loss functions. + + A concrete loss defines three things: + + * :meth:`initial_prediction` -- the constant that minimises the loss and + is used to initialise the ensemble intercept. + * :meth:`negative_gradient` -- the pseudo-residual each weak learner fits. + * :meth:`loss` -- the scalar training/validation loss for reporting and + early stopping. + """ + + #: Human-readable name used in the loss registry. + name: str = "loss" + + @abstractmethod + def initial_prediction(self, y: jnp.ndarray) -> float: + """ + Return the optimal constant prediction for ``y``. + + Parameters + ---------- + y : jnp.ndarray of shape (n_samples,) + Target values. + + Returns + ------- + float + Constant used to initialise the additive model intercept. + """ + + @abstractmethod + def negative_gradient(self, y: jnp.ndarray, y_pred: jnp.ndarray) -> jnp.ndarray: + """ + Return the pseudo-residual ``-dL/dy_pred`` the next term should fit. + + Parameters + ---------- + y : jnp.ndarray of shape (n_samples,) + Target values. + y_pred : jnp.ndarray of shape (n_samples,) + Current ensemble prediction. + + Returns + ------- + jnp.ndarray of shape (n_samples,) + Pseudo-residuals. + """ + + @abstractmethod + def loss(self, y: jnp.ndarray, y_pred: jnp.ndarray) -> float: + """ + Return the scalar loss between ``y`` and ``y_pred``. + + Parameters + ---------- + y : jnp.ndarray of shape (n_samples,) + Target values. + y_pred : jnp.ndarray of shape (n_samples,) + Predicted values. + + Returns + ------- + float + Scalar loss value. + """ + + def to_config(self) -> dict[str, Any]: + """ + Return a JSON-serialisable ``{"name", "params"}`` description. + + Returns + ------- + dict + ``name`` is the registry key; ``params`` are the constructor + keyword arguments needed to rebuild this loss. + """ + return {"name": self.name, "params": {}} + + def __repr__(self) -> str: + params = self.to_config()["params"] + inner = ", ".join(f"{k}={v!r}" for k, v in params.items()) + return f"{type(self).__name__}({inner})" + + +class SquaredError(Loss): + """ + Squared-error loss ``L = mean((y - y_pred)**2)``. + + The optimal constant prediction is the mean of ``y`` and the + pseudo-residual is the ordinary residual ``y - y_pred``. + """ + + name = "squared_error" + + def initial_prediction(self, y: jnp.ndarray) -> float: + """Return ``mean(y)``.""" + return float(jnp.mean(jnp.asarray(y))) + + def negative_gradient(self, y: jnp.ndarray, y_pred: jnp.ndarray) -> jnp.ndarray: + """Return the residual ``y - y_pred``.""" + return jnp.asarray(y) - jnp.asarray(y_pred) + + def loss(self, y: jnp.ndarray, y_pred: jnp.ndarray) -> float: + """Return the mean squared error.""" + y = jnp.asarray(y) + y_pred = jnp.asarray(y_pred) + return float(jnp.mean((y - y_pred) ** 2)) + + +class AbsoluteError(Loss): + """ + Absolute-error (L1) loss ``L = mean(|y - y_pred|)``. + + Robust to outliers. The optimal constant is the median and the negative + gradient is ``sign(y - y_pred)``. + """ + + name = "absolute_error" + + def initial_prediction(self, y: jnp.ndarray) -> float: + """Return ``median(y)``.""" + return float(jnp.median(jnp.asarray(y))) + + def negative_gradient(self, y: jnp.ndarray, y_pred: jnp.ndarray) -> jnp.ndarray: + """Return ``sign(y - y_pred)``.""" + return jnp.sign(jnp.asarray(y) - jnp.asarray(y_pred)) + + def loss(self, y: jnp.ndarray, y_pred: jnp.ndarray) -> float: + """Return the mean absolute error.""" + return float(jnp.mean(jnp.abs(jnp.asarray(y) - jnp.asarray(y_pred)))) + + +class HuberLoss(Loss): + """ + Huber loss: quadratic for small residuals, linear beyond ``delta``. + + Combines the efficiency of squared error near zero with the robustness of + absolute error in the tails. + + Parameters + ---------- + delta : float + Threshold at which the loss transitions from quadratic to linear. + Must be positive. + + Raises + ------ + ValueError + If ``delta`` is not positive. + """ + + name = "huber" + + def __init__(self, delta: float = 1.35): + if delta <= 0: + raise ValueError(f"delta must be positive, got {delta}.") + self.delta = float(delta) + + def initial_prediction(self, y: jnp.ndarray) -> float: + """Return ``median(y)`` (a robust location estimate).""" + return float(jnp.median(jnp.asarray(y))) + + def negative_gradient(self, y: jnp.ndarray, y_pred: jnp.ndarray) -> jnp.ndarray: + """Return ``r`` where ``|r| <= delta`` else ``delta * sign(r)``.""" + r = jnp.asarray(y) - jnp.asarray(y_pred) + return jnp.where(jnp.abs(r) <= self.delta, r, self.delta * jnp.sign(r)) + + def loss(self, y: jnp.ndarray, y_pred: jnp.ndarray) -> float: + """Return the mean Huber loss.""" + r = jnp.asarray(y) - jnp.asarray(y_pred) + abs_r = jnp.abs(r) + quad = 0.5 * r**2 + lin = self.delta * (abs_r - 0.5 * self.delta) + return float(jnp.mean(jnp.where(abs_r <= self.delta, quad, lin))) + + def to_config(self) -> dict[str, Any]: + """Return ``{"name": "huber", "params": {"delta": delta}}``.""" + return {"name": self.name, "params": {"delta": self.delta}} + + +class QuantileLoss(Loss): + """ + Quantile (pinball) loss for estimating the ``quantile``-th conditional + quantile of the target. + + Useful for asymmetric costs and for building prediction intervals (fit one + model per quantile). + + Parameters + ---------- + quantile : float + Target quantile in the open interval ``(0, 1)``. + + Raises + ------ + ValueError + If ``quantile`` is not in ``(0, 1)``. + """ + + name = "quantile" + + def __init__(self, quantile: float = 0.5): + if not 0.0 < quantile < 1.0: + raise ValueError(f"quantile must be in (0, 1), got {quantile}.") + self.quantile = float(quantile) + + def initial_prediction(self, y: jnp.ndarray) -> float: + """Return the empirical ``quantile``-th quantile of ``y``.""" + return float(jnp.quantile(jnp.asarray(y), self.quantile)) + + def negative_gradient(self, y: jnp.ndarray, y_pred: jnp.ndarray) -> jnp.ndarray: + """Return ``q`` where ``y > y_pred`` else ``q - 1``.""" + r = jnp.asarray(y) - jnp.asarray(y_pred) + return jnp.where(r > 0, self.quantile, self.quantile - 1.0) + + def loss(self, y: jnp.ndarray, y_pred: jnp.ndarray) -> float: + """Return the mean pinball loss.""" + r = jnp.asarray(y) - jnp.asarray(y_pred) + return float(jnp.mean(jnp.maximum(self.quantile * r, (self.quantile - 1.0) * r))) + + def to_config(self) -> dict[str, Any]: + """Return ``{"name": "quantile", "params": {"quantile": quantile}}``.""" + return {"name": self.name, "params": {"quantile": self.quantile}} + + +# Registry of available losses. Add future losses here. +_LOSSES: dict[str, type[Loss]] = { + "squared_error": SquaredError, + "absolute_error": AbsoluteError, + "huber": HuberLoss, + "quantile": QuantileLoss, +} + + +def get_loss(loss: str | Loss) -> Loss: + """ + Resolve a loss name (or instance) to a :class:`Loss` instance. + + Parameters + ---------- + loss : str or Loss + Either a registered loss name (``"squared_error"``, ``"absolute_error"``, + ``"huber"``, ``"quantile"``) or an already-constructed :class:`Loss` + instance. Names build losses with their default parameters; pass an + instance (e.g. ``QuantileLoss(0.9)``) to customise. + + Returns + ------- + Loss + A loss instance. + + Raises + ------ + ValueError + If ``loss`` is a string that is not a registered loss name. + TypeError + If ``loss`` is neither a string nor a :class:`Loss` instance. + """ + if isinstance(loss, Loss): + return loss + if isinstance(loss, str): + if loss not in _LOSSES: + raise ValueError(f"Unknown loss {loss!r}. Available losses: {sorted(_LOSSES)}.") + return _LOSSES[loss]() + raise TypeError(f"loss must be a str or Loss instance, got {type(loss).__name__}.") + + +def loss_from_config(config: str | dict[str, Any] | Loss) -> Loss: + """ + Rebuild a :class:`Loss` from a name, an instance, or a ``to_config`` dict. + + Parameters + ---------- + config : str or dict or Loss + A registry name, a ``{"name", "params"}`` dictionary produced by + :meth:`Loss.to_config`, or a :class:`Loss` instance. + + Returns + ------- + Loss + A loss instance. + + Raises + ------ + ValueError + If the config names an unknown loss. + TypeError + If ``config`` is not a str, dict, or :class:`Loss`. + """ + if isinstance(config, (str, Loss)): + return get_loss(config) + if isinstance(config, dict): + name = config["name"] + if name not in _LOSSES: + raise ValueError(f"Unknown loss {name!r}. Available losses: {sorted(_LOSSES)}.") + return _LOSSES[name](**config.get("params", {})) + raise TypeError(f"config must be a str, dict, or Loss, got {type(config).__name__}.") diff --git a/src/jaxsr/additive/recursive.py b/src/jaxsr/additive/recursive.py new file mode 100644 index 0000000..b6f26de --- /dev/null +++ b/src/jaxsr/additive/recursive.py @@ -0,0 +1,368 @@ +""" +Experimental: residual-guided recursive basis expansion. + +This is a deterministic, bounded cousin of genetic-programming symbolic +regression -- and a partial escape from the fixed-library ceiling. Instead of +enumerating a huge depth-``d`` composition space up front (which explodes +combinatorially), it *grows the basis library lazily along the residual*: + +1. Fit a sparse symbolic model over the current library. +2. Take the residual. +3. Build candidate basis functions by composing the currently useful + "building blocks" (selected terms + raw features) with a small operator set + (unary functions, products, ratios) -- one new layer of composition. +4. Screen hard: drop non-finite candidates, deduplicate, keep the top few by + correlation with the residual. +5. Add the survivors to the library and refit. Repeat. + +The effective composition depth is the number of expansion rounds, because a +term discovered in one round becomes an input to the operators in the next. +This is essentially Fast Function Extraction (FFX) / symbolic feature +construction: it can reach compositional targets (e.g. ``x0*sin(x1)``) that a +single flat library misses, at bounded cost -- but it re-enters search-based +territory and will not match a mature GP engine (PySR, Operon). + +The result is an ordinary fitted :class:`jaxsr.SymbolicRegressor` over the grown +library, so it inherits predict/expression/scoring (and the non-finite-basis +guard and negligible-term pruning of the base regressor). Note: the composed +bases are Python closures, so the fitted model is not serialisable via +``save``/``load``, and ``to_sympy`` may not parse deeply nested term names. +""" + +from __future__ import annotations + +from collections.abc import Callable + +import jax.numpy as jnp +import numpy as np + +from .._compat import _SklearnCompatMixin +from ..basis import BasisLibrary, _safe_exp, _safe_log, _safe_sqrt +from ..regressor import SymbolicRegressor + +# Unary operators: (name template, function, complexity cost). +_UNARY_OPS: dict[str, tuple[str, Callable, int]] = { + "sin": ("sin({})", jnp.sin, 2), + "cos": ("cos({})", jnp.cos, 2), + "exp": ("exp({})", _safe_exp, 2), + "log": ("log({})", _safe_log, 2), + "sqrt": ("sqrt({})", _safe_sqrt, 2), + "square": ("({})^2", jnp.square, 2), +} + +# A "block" is (function X->(n,), complexity, set of feature indices used). +_Block = tuple[Callable, int, frozenset] + + +def _compose_unary(func: Callable, unary: Callable) -> Callable: + """Return X -> unary(func(X)).""" + return lambda X: unary(func(X)) + + +def _compose_product(f1: Callable, f2: Callable) -> Callable: + """Return X -> f1(X) * f2(X).""" + return lambda X: f1(X) * f2(X) + + +def _compose_ratio(f1: Callable, f2: Callable) -> Callable: + """Return X -> f1(X) / f2(X).""" + return lambda X: f1(X) / f2(X) + + +def _abs_corr(a: np.ndarray, b: np.ndarray) -> float: + """Absolute Pearson correlation, robust to constant inputs.""" + a = a - a.mean() + b = b - b.mean() + na, nb = np.linalg.norm(a), np.linalg.norm(b) + if na < 1e-12 or nb < 1e-12: + return 0.0 + return abs(float((a @ b) / (na * nb))) + + +class RecursiveSymbolicRegressor(_SklearnCompatMixin): + """ + Residual-guided recursive basis expansion (experimental). + + Grows a symbolic basis library over several rounds, composing the most + useful discovered terms with unary functions and products/ratios, guided by + correlation with the current residual. Produces a fitted + :class:`jaxsr.SymbolicRegressor` over the grown library. + + Parameters + ---------- + n_expansions : int + Number of expansion rounds (roughly the maximum composition depth). + max_terms : int + Maximum number of terms in the sparse model fit each round. + beam_width : int + Number of new candidate bases (highest residual correlation) kept per + round. Controls the combinatorial cost. + base_degree : int + Degree of the initial polynomial seed terms (before any composition). + unary_ops : tuple of str + Unary operators to compose with. Subset of ``sin``, ``cos``, ``exp``, + ``log``, ``sqrt``, ``square``. + binary_ops : tuple of str + Binary operators: any of ``mul`` (products) and ``div`` (ratios). + strategy : str + Selection strategy for the per-round sparse fit. + information_criterion : str + Information criterion for the per-round sparse fit. + feature_names : list of str, optional + Names for the input features. + random_state : int, optional + Unused placeholder for API symmetry (the search is deterministic). + + Attributes + ---------- + model_ : SymbolicRegressor + The fitted sparse model over the final grown library. + library_size_ : int + Number of candidate basis functions in the final library. + history_ : list of dict + Per-round diagnostics (library size, number of terms, train R^2). + + Examples + -------- + >>> from jaxsr.additive import RecursiveSymbolicRegressor + >>> model = RecursiveSymbolicRegressor(n_expansions=3) + >>> model.fit(X, y) # doctest: +SKIP + >>> print(model.expression_) # doctest: +SKIP + """ + + def __init__( + self, + n_expansions: int = 3, + max_terms: int = 8, + beam_width: int = 25, + base_degree: int = 2, + unary_ops: tuple[str, ...] = ("sin", "cos", "exp", "log", "sqrt", "square"), + binary_ops: tuple[str, ...] = ("mul", "div"), + strategy: str = "greedy_forward", + information_criterion: str = "bic", + feature_names: list[str] | None = None, + random_state: int | None = None, + ): + self.n_expansions = n_expansions + self.max_terms = max_terms + self.beam_width = beam_width + self.base_degree = base_degree + self.unary_ops = unary_ops + self.binary_ops = binary_ops + self.strategy = strategy + self.information_criterion = information_criterion + self.feature_names = feature_names + self.random_state = random_state + + self.model_: SymbolicRegressor | None = None + self.library_size_: int | None = None + self.history_: list[dict] = [] + self._is_fitted = False + + # ------------------------------------------------------------------ + def _validate_params(self) -> None: + """Validate constructor parameters.""" + if self.n_expansions < 0: + raise ValueError(f"n_expansions must be >= 0, got {self.n_expansions}.") + if self.max_terms < 1: + raise ValueError(f"max_terms must be >= 1, got {self.max_terms}.") + if self.beam_width < 1: + raise ValueError(f"beam_width must be >= 1, got {self.beam_width}.") + bad_unary = set(self.unary_ops) - set(_UNARY_OPS) + if bad_unary: + raise ValueError(f"Unknown unary_ops {sorted(bad_unary)}; valid: {sorted(_UNARY_OPS)}.") + bad_binary = set(self.binary_ops) - {"mul", "div"} + if bad_binary: + raise ValueError(f"Unknown binary_ops {sorted(bad_binary)}; valid: ['div', 'mul'].") + + def _seed_library(self, feature_names: list[str]) -> dict[str, _Block]: + """Build the depth-0 library: constant, features, and polynomial seeds.""" + library: dict[str, _Block] = { + "1": (lambda X: jnp.ones(X.shape[0]), 0, frozenset()), + } + for i, name in enumerate(feature_names): + library[name] = (lambda X, i=i: X[:, i], 1, frozenset({i})) + for d in range(2, self.base_degree + 1): + for i, name in enumerate(feature_names): + library[f"{name}^{d}"] = (lambda X, i=i, d=d: X[:, i] ** d, d, frozenset({i})) + return library + + def _generate(self, blocks: dict[str, _Block]) -> dict[str, _Block]: + """Compose one new layer of candidates from the current blocks.""" + new: dict[str, _Block] = {} + items = list(blocks.items()) + + for name, (func, comp, feats) in items: + for op in self.unary_ops: + template, unary_fn, cost = _UNARY_OPS[op] + new[template.format(name)] = (_compose_unary(func, unary_fn), comp + cost, feats) + + if "mul" in self.binary_ops: + for i, (n1, (f1, c1, ft1)) in enumerate(items): + for n2, (f2, c2, ft2) in items[i:]: + new[f"({n1})*({n2})"] = (_compose_product(f1, f2), c1 + c2 + 1, ft1 | ft2) + + if "div" in self.binary_ops: + for n1, (f1, c1, ft1) in items: + for n2, (f2, c2, ft2) in items: + if n1 != n2: + new[f"({n1})/({n2})"] = (_compose_ratio(f1, f2), c1 + c2 + 1, ft1 | ft2) + + return new + + def _screen( + self, + candidates: dict[str, _Block], + X: jnp.ndarray, + residual: np.ndarray, + existing: dict[str, _Block], + ) -> dict[str, _Block]: + """Keep the finite, non-duplicate, most residual-correlated candidates.""" + scored: list[tuple[float, str, _Block, np.ndarray]] = [] + for name, block in candidates.items(): + if name in existing: + continue + values = np.asarray(block[0](X), dtype=float) + if values.shape != residual.shape or not np.all(np.isfinite(values)): + continue + if values.std() < 1e-12: + continue + scored.append((_abs_corr(values, residual), name, block, values)) + + scored.sort(key=lambda t: -t[0]) + + kept: dict[str, _Block] = {} + kept_values: list[np.ndarray] = [] + for corr, name, block, values in scored: + if corr < 1e-6: + break + # Deduplicate against already-kept candidates by near-perfect correlation. + if any(_abs_corr(values, kv) > 0.9999 for kv in kept_values): + continue + kept[name] = block + kept_values.append(values) + if len(kept) >= self.beam_width: + break + return kept + + def _fit_library( + self, library: dict[str, _Block], X: jnp.ndarray, y: jnp.ndarray, n_features: int + ) -> SymbolicRegressor: + """Fit a sparse SymbolicRegressor over the current library.""" + basis = BasisLibrary(n_features, self.feature_names) + for name, (func, complexity, feats) in library.items(): + basis.add_custom( + name, func, complexity=complexity, feature_indices=tuple(sorted(feats)) + ) + model = SymbolicRegressor( + basis_library=basis, + max_terms=self.max_terms, + strategy=self.strategy, + information_criterion=self.information_criterion, + ) + return model.fit(X, y) + + def fit(self, X: jnp.ndarray, y: jnp.ndarray) -> RecursiveSymbolicRegressor: + """ + Fit by residual-guided recursive basis expansion. + + Parameters + ---------- + X : array-like of shape (n_samples, n_features) + Training inputs. + y : array-like of shape (n_samples,) + Target values. + + Returns + ------- + self : RecursiveSymbolicRegressor + The fitted estimator. + + Raises + ------ + ValueError + If parameters are invalid or ``X``/``y`` are mismatched. + """ + self._validate_params() + X = jnp.atleast_2d(jnp.asarray(X)) + y = jnp.asarray(y).ravel() + if X.shape[0] != y.shape[0]: + raise ValueError( + f"X and y must have the same number of samples. " + f"Got X: {X.shape[0]}, y: {y.shape[0]}." + ) + + n_features = X.shape[1] + feature_names = self.feature_names or [f"x{i}" for i in range(n_features)] + + library = self._seed_library(feature_names) + model = self._fit_library(library, X, y, n_features) + self.history_ = [ + { + "round": 0, + "library_size": len(library), + "n_terms": len(model.selected_features_), + "train_r2": float(model.score(X, y)), + } + ] + + for rnd in range(1, self.n_expansions + 1): + residual = np.asarray(y - model.predict(X), dtype=float) + if np.linalg.norm(residual) < 1e-9 * (np.linalg.norm(np.asarray(y)) + 1e-12): + break + + # Building blocks for composition: raw features/seeds plus the + # terms the model actually selected (bounded and useful). + blocks = self._seed_library(feature_names) + for name in model.selected_features_: + if name in library: + blocks[name] = library[name] + + candidates = self._generate(blocks) + survivors = self._screen(candidates, X, residual, library) + if not survivors: + break + + library.update(survivors) + model = self._fit_library(library, X, y, n_features) + self.history_.append( + { + "round": rnd, + "library_size": len(library), + "n_terms": len(model.selected_features_), + "train_r2": float(model.score(X, y)), + } + ) + + self.model_ = model + self.library_size_ = len(library) + self._is_fitted = True + return self + + # ------------------------------------------------------------------ + def _check_is_fitted(self) -> None: + """Raise if the model has not been fitted yet.""" + if not self._is_fitted or self.model_ is None: + raise RuntimeError("Model not fitted. Call fit() first.") + + def predict(self, X: jnp.ndarray) -> jnp.ndarray: + """Predict with the fitted model.""" + self._check_is_fitted() + return self.model_.predict(X) + + def score(self, X: jnp.ndarray, y: jnp.ndarray) -> float: + """Return the R^2 score of the fitted model on ``(X, y)``.""" + self._check_is_fitted() + return self.model_.score(X, y) + + @property + def expression_(self) -> str: + """Human-readable expression of the fitted model.""" + self._check_is_fitted() + return self.model_.expression_ + + @property + def selected_features_(self) -> list[str]: + """Names of the selected (possibly composed) basis functions.""" + self._check_is_fitted() + return self.model_.selected_features_ diff --git a/src/jaxsr/additive/stagewise.py b/src/jaxsr/additive/stagewise.py new file mode 100644 index 0000000..69653a6 --- /dev/null +++ b/src/jaxsr/additive/stagewise.py @@ -0,0 +1,449 @@ +""" +Stagewise (boosting-style) additive symbolic regression. + +:class:`StagewiseSymbolicRegressor` builds a model of the form:: + + f(x) = intercept + sum_k coefficients[k] * g_k(x) + +by repeatedly fitting a small symbolic expression ``g_k`` to the current +residual (pseudo-residual for general losses). This is conceptually +"gradient boosting, but the weak learners are symbolic expressions instead of +trees". + +Once discovered, a term is *frozen* -- its internal structure never changes. +Only the linear weights may be re-estimated (see ``refit_coefficients``). The +future :class:`~jaxsr.additive.backfitting.BackfittingSymbolicRegressor` will +instead revise terms in place. +""" + +from __future__ import annotations + +import warnings + +import jax.numpy as jnp +import numpy as np +from scipy.optimize import minimize_scalar + +from ..regressor import SymbolicRegressor, fit_symbolic +from .base import _BaseAdditiveRegressor +from .coefficient_refit import refit_ols +from .ensemble import AdditiveSymbolicModel, additive_predict +from .losses import Loss, SquaredError, get_loss + + +class StagewiseSymbolicRegressor(_BaseAdditiveRegressor): + """ + Stagewise additive symbolic regression (symbolic gradient boosting). + + Fits an additive ensemble of small symbolic expressions by iteratively + fitting each new expression to the residual of the current model. Old + terms are frozen; optionally the linear coefficients over all discovered + terms are refit by least squares after each stage. + + Parameters + ---------- + n_terms : int + Maximum number of boosting stages (symbolic terms) to add. + learning_rate : float + Shrinkage applied to each stage's contribution when + ``refit_coefficients=False``. Ignored when ``refit_coefficients=True`` + (the weights are then chosen by least squares), but still recorded. + max_complexity : int + Complexity budget for each stage's symbolic expression, expressed as + the maximum number of basis terms (passed as ``max_terms`` to the + underlying :func:`jaxsr.fit_symbolic`). Keep this small to favour many + simple, interpretable terms over one large expression. + refit_coefficients : bool + If True, after each new term is added, re-solve the intercept and all + per-term coefficients by ordinary least squares over the discovered + symbolic features. If False, use learning-rate-scaled stagewise + weights. OLS refit targets squared error, so it is only applied for + ``loss="squared_error"``; with any other loss it is ignored (a warning + is issued) and gradient boosting with a per-stage line search is used. + loss : str or Loss + Loss function. One of ``"squared_error"``, ``"absolute_error"``, + ``"huber"``, ``"quantile"``, or a :class:`~jaxsr.additive.Loss` + instance for custom parameters (e.g. ``QuantileLoss(0.9)`` or + ``HuberLoss(delta=2.0)``). Non-squared losses are fit by gradient + boosting: each term fits the negative gradient and its step size is + chosen by a line search that minimises the loss. + early_stopping : bool + If True, hold out a validation split and stop adding terms once the + validation loss stops improving. + validation_fraction : float + Fraction of the training data held out for early-stopping validation. + Only used when ``early_stopping=True``. + patience : int + Number of consecutive non-improving stages tolerated before stopping. + min_delta : float + Minimum decrease in validation loss to count as an improvement. + max_poly_degree : int + Maximum polynomial degree available to each stage. + include_transcendental : bool + If True, allow ``log``/``exp``/``sqrt``/``inv`` terms in each stage. + include_ratios : bool + If True, allow ratio terms ``x_i / x_j`` in each stage. + strategy : str + Selection strategy for each stage (see :class:`jaxsr.SymbolicRegressor`). + information_criterion : str + Information criterion used to control complexity within each stage: + ``"aic"``, ``"aicc"``, or ``"bic"``. + feature_names : list of str, optional + Names for the input features. Defaults to ``["x0", "x1", ...]``. + random_state : int, optional + Seed controlling the early-stopping validation split. + + Attributes + ---------- + model_ : AdditiveSymbolicModel + The fitted additive model. + intercept_ : float + Fitted intercept. + coefficients_ : list of float + Fitted per-term coefficients. + expressions_ : list of str + Human-readable expression for each term. + terms_ : list of SymbolicRegressor + The fitted symbolic terms. + learning_rates_ : list of float + Learning rate recorded at each stage. + training_history_ : list of dict + Per-stage diagnostics. + n_terms_ : int + Number of terms in the fitted model. + + Examples + -------- + >>> import numpy as np + >>> from jaxsr.additive import StagewiseSymbolicRegressor + >>> X = np.random.randn(200, 2) + >>> y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + >>> model = StagewiseSymbolicRegressor(n_terms=5, refit_coefficients=True) + >>> model.fit(X, y) # doctest: +SKIP + >>> print(model) # doctest: +SKIP + """ + + def __init__( + self, + n_terms: int = 10, + learning_rate: float = 0.1, + max_complexity: int = 4, + refit_coefficients: bool = True, + loss: str = "squared_error", + early_stopping: bool = False, + validation_fraction: float = 0.2, + patience: int = 3, + min_delta: float = 1e-8, + max_poly_degree: int = 3, + include_transcendental: bool = False, + include_ratios: bool = False, + strategy: str = "greedy_forward", + information_criterion: str = "bic", + feature_names: list[str] | None = None, + random_state: int | None = None, + ): + self.n_terms = n_terms + self.learning_rate = learning_rate + self.max_complexity = max_complexity + self.refit_coefficients = refit_coefficients + self.loss = loss + self.early_stopping = early_stopping + self.validation_fraction = validation_fraction + self.patience = patience + self.min_delta = min_delta + self.max_poly_degree = max_poly_degree + self.include_transcendental = include_transcendental + self.include_ratios = include_ratios + self.strategy = strategy + self.information_criterion = information_criterion + self.feature_names = feature_names + self.random_state = random_state + + # Fitted attributes + self.model_: AdditiveSymbolicModel | None = None + self._is_fitted = False + + # ------------------------------------------------------------------ + # Fitting + # ------------------------------------------------------------------ + def _validate_params(self) -> None: + """Validate constructor parameters.""" + if self.n_terms < 1: + raise ValueError(f"n_terms must be >= 1, got {self.n_terms}.") + if self.learning_rate <= 0: + raise ValueError(f"learning_rate must be > 0, got {self.learning_rate}.") + if self.max_complexity < 1: + raise ValueError(f"max_complexity must be >= 1, got {self.max_complexity}.") + if self.patience < 1: + raise ValueError(f"patience must be >= 1, got {self.patience}.") + if self.early_stopping and not (0.0 < self.validation_fraction < 1.0): + raise ValueError( + f"validation_fraction must be in (0, 1), got {self.validation_fraction}." + ) + + def _fit_stage( + self, + X: jnp.ndarray, + residual: jnp.ndarray, + feature_names: list[str], + ) -> SymbolicRegressor: + """ + Fit a single symbolic term to the residual using existing machinery. + + Parameters + ---------- + X : jnp.ndarray of shape (n_samples, n_features) + Training inputs. + residual : jnp.ndarray of shape (n_samples,) + Pseudo-residual target for this stage. + feature_names : list of str + Feature names shared across stages. + + Returns + ------- + SymbolicRegressor + A fitted symbolic regressor representing ``g_k`` that produces + finite predictions on the training data. + + Notes + ----- + A transcendental or ratio basis (e.g. ``log``/``sqrt``/``1/x``) can be + invalid on part of the data domain and yield non-finite predictions + even though it was selectable during fitting. If that happens, the + stage is refit without transcendental and ratio bases so the ensemble + never contains a NaN-producing term. + """ + term = fit_symbolic( + X, + residual, + feature_names=feature_names, + max_terms=self.max_complexity, + max_poly_degree=self.max_poly_degree, + include_transcendental=self.include_transcendental, + include_ratios=self.include_ratios, + strategy=self.strategy, + information_criterion=self.information_criterion, + ) + uses_risky_basis = self.include_transcendental or self.include_ratios + if uses_risky_basis and not bool(jnp.all(jnp.isfinite(term.predict(X)))): + warnings.warn( + "A fitted term produced non-finite predictions on the training " + "data (a transcendental/ratio basis is invalid on this domain); " + "refitting this stage without transcendental and ratio bases.", + stacklevel=2, + ) + term = fit_symbolic( + X, + residual, + feature_names=feature_names, + max_terms=self.max_complexity, + max_poly_degree=self.max_poly_degree, + include_transcendental=False, + include_ratios=False, + strategy=self.strategy, + information_criterion=self.information_criterion, + ) + return term + + def _train_val_split( + self, X: jnp.ndarray, y: jnp.ndarray + ) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray | None, jnp.ndarray | None]: + """Split off a validation set for early stopping (if enabled).""" + if not self.early_stopping: + return X, y, None, None + + n = X.shape[0] + n_val = max(1, int(round(n * self.validation_fraction))) + if n_val >= n: + raise ValueError("validation_fraction too large: no training samples remain.") + + rng = np.random.default_rng(self.random_state) + perm = rng.permutation(n) + val_idx = jnp.asarray(perm[:n_val]) + train_idx = jnp.asarray(perm[n_val:]) + return X[train_idx], y[train_idx], X[val_idx], y[val_idx] + + @staticmethod + def _line_search_step( + loss_fn: Loss, + y: jnp.ndarray, + base_pred: jnp.ndarray, + direction: jnp.ndarray, + ) -> float: + """ + Find the step size that minimises the loss along a term's direction. + + Solves ``argmin_gamma loss(y, base_pred + gamma * direction)`` for the + newly added term (gradient-boosting line search). All supported losses + are convex in ``gamma``, so a 1-D scalar minimisation suffices. + + Parameters + ---------- + loss_fn : Loss + Loss to minimise. + y : jnp.ndarray + Target values. + base_pred : jnp.ndarray + Current ensemble prediction. + direction : jnp.ndarray + The new term's predictions ``g_k(X)``. + + Returns + ------- + float + Optimal step size (``0.0`` if the direction carries no signal). + """ + direction = jnp.asarray(direction) + if not bool(jnp.any(jnp.abs(direction) > 1e-12)): + return 0.0 + + def objective(gamma: float) -> float: + return loss_fn.loss(y, base_pred + gamma * direction) + + result = minimize_scalar(objective) + if not result.success or not np.isfinite(result.x): + return 1.0 + return float(result.x) + + def fit(self, X: jnp.ndarray, y: jnp.ndarray) -> StagewiseSymbolicRegressor: + """ + Fit the stagewise additive symbolic model. + + Parameters + ---------- + X : array-like of shape (n_samples, n_features) + Training inputs. + y : array-like of shape (n_samples,) + Target values. + + Returns + ------- + self : StagewiseSymbolicRegressor + The fitted estimator. + + Raises + ------ + ValueError + If parameters are invalid, ``X`` and ``y`` have mismatched sample + counts, or ``X``/``y`` contain non-finite values. + """ + self._validate_params() + + X = jnp.atleast_2d(jnp.asarray(X)) + y = jnp.asarray(y).ravel() + if X.shape[0] != y.shape[0]: + raise ValueError( + f"X and y must have the same number of samples. " + f"Got X: {X.shape[0]}, y: {y.shape[0]}." + ) + if not bool(jnp.all(jnp.isfinite(X))): + raise ValueError("X contains non-finite values (NaN or inf).") + if not bool(jnp.all(jnp.isfinite(y))): + raise ValueError("y contains non-finite values (NaN or inf).") + + n_features = X.shape[1] + feature_names = self.feature_names or [f"x{i}" for i in range(n_features)] + if len(feature_names) != n_features: + raise ValueError( + f"feature_names has {len(feature_names)} entries but X has " + f"{n_features} features." + ) + + loss_fn = get_loss(self.loss) + is_squared = isinstance(loss_fn, SquaredError) + + # OLS refit minimises squared error, so it is inconsistent with any + # other loss. Fall back to gradient boosting (with a line search) and + # tell the user. + use_refit = self.refit_coefficients and is_squared + if self.refit_coefficients and not is_squared: + warnings.warn( + "refit_coefficients=True re-solves coefficients by ordinary " + "least squares, which targets squared error and is inconsistent " + f"with loss={loss_fn.name!r}; using gradient-boosting stagewise " + "weights with a per-stage line search instead.", + stacklevel=2, + ) + + X_tr, y_tr, X_val, y_val = self._train_val_split(X, y) + + intercept = loss_fn.initial_prediction(y_tr) + terms: list[SymbolicRegressor] = [] + coefficients: list[float] = [] + learning_rates: list[float] = [] + history: list[dict] = [] + + prediction_tr = jnp.full((X_tr.shape[0],), intercept) + + best_val = np.inf + best_iter = -1 + n_no_improve = 0 + + for k in range(self.n_terms): + residual = loss_fn.negative_gradient(y_tr, prediction_tr) + term = self._fit_stage(X_tr, residual, feature_names) + terms.append(term) + learning_rates.append(self.learning_rate) + + if use_refit: + Phi_tr = jnp.stack([t.predict(X_tr) for t in terms], axis=1) + intercept, coef_arr = refit_ols(Phi_tr, y_tr) + coefficients = [float(c) for c in coef_arr] + prediction_tr = intercept + Phi_tr @ coef_arr + else: + g_tr = term.predict(X_tr) + if is_squared: + step = self.learning_rate + else: + # Gradient boosting: shrink the loss-optimal step size. + gamma = self._line_search_step(loss_fn, y_tr, prediction_tr, g_tr) + step = self.learning_rate * gamma + coefficients.append(float(step)) + prediction_tr = prediction_tr + step * g_tr + + train_loss = loss_fn.loss(y_tr, prediction_tr) + + val_loss = None + if X_val is not None: + val_pred = additive_predict(X_val, intercept, terms, coefficients) + val_loss = loss_fn.loss(y_val, val_pred) + + history.append( + { + "n_terms": k + 1, + "train_loss": float(train_loss), + "val_loss": None if val_loss is None else float(val_loss), + "intercept": float(intercept), + "coefficients": [float(c) for c in coefficients], + } + ) + + # Early stopping bookkeeping + if val_loss is not None: + if val_loss < best_val - self.min_delta: + best_val = val_loss + best_iter = k + n_no_improve = 0 + else: + n_no_improve += 1 + if n_no_improve >= self.patience: + break + + # Roll back to the best validation iteration if early stopping was used. + if self.early_stopping and best_iter >= 0 and best_iter < len(terms) - 1: + best_snapshot = history[best_iter] + terms = terms[: best_iter + 1] + learning_rates = learning_rates[: best_iter + 1] + coefficients = list(best_snapshot["coefficients"]) + intercept = float(best_snapshot["intercept"]) + + self.model_ = AdditiveSymbolicModel( + intercept=float(intercept), + terms=terms, + coefficients=coefficients, + learning_rates=learning_rates, + feature_names=feature_names, + training_history=history, + ) + self._is_fitted = True + return self diff --git a/src/jaxsr/additive/uncertainty.py b/src/jaxsr/additive/uncertainty.py new file mode 100644 index 0000000..dbd732b --- /dev/null +++ b/src/jaxsr/additive/uncertainty.py @@ -0,0 +1,182 @@ +""" +Bootstrap structural uncertainty for additive symbolic models. + +Refitting an additive symbolic regressor on bootstrap resamples of the training +data quantifies how stable the *discovered structure* is: + +* **Inclusion probabilities** -- how often each basis function is selected across + resamples (a cheap, frequentist approximation to a posterior inclusion + probability). When these are all near 0 or 1 the structure is identifiable + and a single expression is trustworthy; diffuse values (e.g. under collinear + features) signal genuine structural uncertainty, where no single expression is + well determined. +* **Predictive ensemble** -- the spread of predictions across resamples, giving + intervals that reflect structural variability, not just coefficient noise. + +This is a model-agnostic stand-in for a full Bayesian treatment: it works for +any additive regressor (stagewise or backfitting) and reuses the estimator's own +fitting machinery. It also serves as a decision gate -- if the inclusion +probabilities are already crisp, a heavier Bayesian model buys little. +""" + +from __future__ import annotations + +import warnings +from collections import Counter + +import jax.numpy as jnp +import numpy as np + + +def _clone(estimator): + """Return an unfitted copy of ``estimator`` with the same configuration.""" + return type(estimator)(**estimator.get_params(deep=False)) + + +def bootstrap_additive( + estimator, + X: jnp.ndarray, + y: jnp.ndarray, + n_bootstrap: int = 100, + random_state: int | None = None, +) -> dict: + """ + Refit an additive regressor on bootstrap resamples to assess structure. + + Parameters + ---------- + estimator : StagewiseSymbolicRegressor or BackfittingSymbolicRegressor + A configured (fitted or unfitted) additive regressor. It is cloned via + its constructor parameters and refit on each resample; the original is + not modified. + X : array-like of shape (n_samples, n_features) + Training inputs. + y : array-like of shape (n_samples,) + Target values. + n_bootstrap : int + Number of bootstrap resamples. + random_state : int, optional + Seed for the resampling for reproducibility. + + Returns + ------- + dict + Plain dictionary with keys: + + ``"inclusion_probabilities"`` : dict[str, float] + Basis-function name -> fraction of resamples selecting it (in any + term), sorted descending. + ``"n_terms"`` : numpy.ndarray + Number of terms in each successful resample fit. + ``"models"`` : list + The fitted bootstrap estimators (use with + :func:`bootstrap_predict_additive`). + ``"n_bootstrap"`` : int + Number of resamples that fit successfully. + + Raises + ------ + ValueError + If ``X``/``y`` are mismatched or ``n_bootstrap < 1``. + RuntimeError + If every bootstrap fit fails. + """ + if n_bootstrap < 1: + raise ValueError(f"n_bootstrap must be >= 1, got {n_bootstrap}.") + + X = np.asarray(jnp.atleast_2d(jnp.asarray(X))) + y = np.asarray(jnp.asarray(y).ravel()) + if X.shape[0] != y.shape[0]: + raise ValueError( + f"X and y must have the same number of samples. " + f"Got X: {X.shape[0]}, y: {y.shape[0]}." + ) + + n = X.shape[0] + rng = np.random.default_rng(random_state) + + counts: Counter = Counter() + n_terms: list[int] = [] + models: list = [] + n_failed = 0 + + for _ in range(n_bootstrap): + idx = rng.integers(0, n, n) + model = _clone(estimator) + try: + model.fit(X[idx], y[idx]) + except (ValueError, RuntimeError, np.linalg.LinAlgError): + n_failed += 1 + continue + names: set[str] = set() + for term in model.terms_: + names.update(term.selected_features_) + counts.update(names) + n_terms.append(model.n_terms_) + models.append(model) + + n_ok = len(models) + if n_ok == 0: + raise RuntimeError("All bootstrap fits failed.") + if n_failed: + warnings.warn( + f"{n_failed}/{n_bootstrap} bootstrap fits failed and were skipped.", + stacklevel=2, + ) + + inclusion = {name: counts[name] / n_ok for name in counts} + inclusion = dict(sorted(inclusion.items(), key=lambda kv: -kv[1])) + + return { + "inclusion_probabilities": inclusion, + "n_terms": np.array(n_terms), + "models": models, + "n_bootstrap": n_ok, + } + + +def bootstrap_predict_additive( + models: list, + X: jnp.ndarray, + alpha: float = 0.1, +) -> dict: + """ + Predictive ensemble from a bootstrap set of additive models. + + Parameters + ---------- + models : list + Fitted additive estimators, e.g. the ``"models"`` entry returned by + :func:`bootstrap_additive`. + X : array-like of shape (n_samples, n_features) + Inputs to predict. + alpha : float + Significance level; the interval covers ``1 - alpha`` (default 0.1 for + a 90% interval). + + Returns + ------- + dict + Plain dictionary with keys ``"mean"``, ``"std"``, ``"median"``, + ``"lower"``, ``"upper"`` (each of shape ``(n_samples,)``) and + ``"predictions"`` of shape ``(n_models, n_samples)``. + + Raises + ------ + ValueError + If ``models`` is empty or ``alpha`` is not in ``(0, 1)``. + """ + if not models: + raise ValueError("models must be a non-empty list of fitted estimators.") + if not 0.0 < alpha < 1.0: + raise ValueError(f"alpha must be in (0, 1), got {alpha}.") + + preds = np.stack([np.asarray(m.predict(X)) for m in models]) + return { + "mean": preds.mean(axis=0), + "std": preds.std(axis=0), + "median": np.median(preds, axis=0), + "lower": np.quantile(preds, alpha / 2, axis=0), + "upper": np.quantile(preds, 1 - alpha / 2, axis=0), + "predictions": preds, + } diff --git a/src/jaxsr/classifier.py b/src/jaxsr/classifier.py index aa6548f..ae2ec48 100644 --- a/src/jaxsr/classifier.py +++ b/src/jaxsr/classifier.py @@ -12,6 +12,7 @@ from __future__ import annotations +import dataclasses import json import warnings from collections.abc import Callable @@ -26,6 +27,7 @@ from .metrics import ( compute_accuracy, compute_all_classification_metrics, + compute_classification_ic, compute_log_loss, ) from .selection import ( @@ -262,16 +264,23 @@ def fit( self._X_train = X self._y_train = y - # Evaluate basis functions + # Evaluate basis functions. A basis non-finite on the training data + # must never remain selected: predict re-evaluates each basis from + # scratch, so it would return NaN. Zero the whole invalid column (it + # then contributes nothing to the fitted logits), and drop any that + # still get selected afterwards (an exact operation -- the column was + # zero during fitting). Phi = self.basis_library.evaluate(X) invalid_mask = ~jnp.all(jnp.isfinite(Phi), axis=0) - if jnp.any(invalid_mask): + has_invalid = bool(jnp.any(invalid_mask)) + if has_invalid: n_invalid = int(jnp.sum(invalid_mask)) warnings.warn( - f"Removing {n_invalid} basis functions with non-finite values", + f"Excluding {n_invalid} basis function(s) with non-finite values on the " + f"training data; they will not be selected.", stacklevel=2, ) - Phi = jnp.where(jnp.isfinite(Phi), Phi, 0) + Phi = jnp.where(invalid_mask, 0.0, Phi) if len(classes) == 2: self._is_binary = True @@ -280,9 +289,55 @@ def fit( self._is_binary = False self._fit_multiclass(Phi, y, classes) + if has_invalid: + self._drop_nonfinite_terms(invalid_mask) + self._is_fitted = True return self + def _drop_nonfinite_terms(self, invalid_mask: jnp.ndarray): + """ + Remove selected terms whose basis is non-finite on the training data. + + The removed columns were zeroed during fitting, so they contributed + nothing to the training logits; dropping them (and their coefficients) + leaves predictions on in-domain data unchanged while guaranteeing + :meth:`predict` stays finite. Information criteria are recomputed from + the unchanged negative log-likelihood. + + Parameters + ---------- + invalid_mask : jnp.ndarray of bool + Per-column mask; True marks a basis non-finite on the training data. + """ + invalid = np.asarray(invalid_mask) + + def clean(result: ClassificationResult) -> ClassificationResult: + idx = [int(i) for i in np.asarray(result.selected_indices)] + keep_pos = [p for p, i in enumerate(idx) if not invalid[i]] + if len(keep_pos) == len(idx): + return result + keep = [idx[p] for p in keep_pos] + coeffs = np.asarray(result.coefficients) + n, k = result.n_samples, len(keep) + nll = result.neg_log_likelihood + return dataclasses.replace( + result, + coefficients=jnp.asarray(coeffs[keep_pos]), + selected_indices=jnp.asarray(keep), + selected_names=[result.selected_names[p] for p in keep_pos], + complexity=int(sum(self.basis_library.complexities[i] for i in keep)), + aic=compute_classification_ic(n, k, nll, "aic"), + bic=compute_classification_ic(n, k, nll, "bic"), + aicc=compute_classification_ic(n, k, nll, "aicc"), + ) + + if self._is_binary: + self._result = clean(self._result) + else: + self._ovr_results = [clean(r) for r in self._ovr_results] + self._result = self._ovr_results[0] + def _fit_binary( self, Phi: jnp.ndarray, diff --git a/src/jaxsr/regressor.py b/src/jaxsr/regressor.py index 16faad0..5716c47 100644 --- a/src/jaxsr/regressor.py +++ b/src/jaxsr/regressor.py @@ -22,6 +22,7 @@ SelectionPath, SelectionResult, compute_pareto_front, + fit_ols, select_features, ) from .uncertainty import ( @@ -129,6 +130,14 @@ class SymbolicRegressor(_SklearnCompatMixin): information criterion during term selection, biasing selection towards models that better satisfy constraints. Default 0.0 (no constraint-aware selection). + prune_tol : float + After fitting, drop any selected term whose contribution to the fit + (``|coef| * ||basis||`` on the training data) is smaller than + ``prune_tol`` times the largest term's contribution, then refit. This + removes numerically negligible terms and prevents a spuriously selected + basis that diverges out of the training domain (e.g. ``exp(x0/x1)`` near + ``x1 = 0``) from making ``predict`` non-finite. Set to 0 to disable. + Default 1e-6. Attributes ---------- @@ -173,6 +182,7 @@ def __init__( param_optimization_budget: int = 50, constraint_enforcement: str = "penalty", constraint_selection_weight: float = 0.0, + prune_tol: float = 1e-6, ): if constraint_enforcement not in ("penalty", "constrained", "exact"): raise ValueError( @@ -191,6 +201,7 @@ def __init__( self.param_optimization_budget = param_optimization_budget self.constraint_enforcement = constraint_enforcement self.constraint_selection_weight = constraint_selection_weight + self.prune_tol = prune_tol # Fitted attributes self._result: SelectionResult | None = None @@ -306,15 +317,23 @@ def fit( # Evaluate basis functions Phi = self.basis_library.evaluate(X) - # Handle invalid values in design matrix + # Handle invalid values in the design matrix. A basis that is + # non-finite on the training data (e.g. log(x) with x <= 0) must never + # remain in the final model: predict() re-evaluates each selected basis + # from scratch, so a non-finite basis yields NaN predictions. Zero the + # whole invalid column first so the selection math stays finite, then + # drop any such column that still gets selected (below) -- greedy search + # can otherwise fill a term slot with a zeroed, useless column. invalid_mask = ~jnp.all(jnp.isfinite(Phi), axis=0) - if jnp.any(invalid_mask): + has_invalid = bool(jnp.any(invalid_mask)) + if has_invalid: n_invalid = int(jnp.sum(invalid_mask)) warnings.warn( - f"Removing {n_invalid} basis functions with non-finite values", stacklevel=2 + f"Excluding {n_invalid} basis function(s) with non-finite values on the " + f"training data; they will not be selected.", + stacklevel=2, ) - # Replace invalid columns with zeros (they won't be selected) - Phi = jnp.where(jnp.isfinite(Phi), Phi, 0) + Phi = jnp.where(invalid_mask, 0.0, Phi) # Run selection extra_kw: dict[str, Any] = {} @@ -351,6 +370,22 @@ def fit( self._result = self._selection_path.best + # Post-selection cleanups operate on ordinary (non-parametric) terms + # via OLS. Parametric bases carry their own refit/resolution machinery + # and their contribution is only meaningful after their internal + # parameters are optimised, so skip these steps for parametric libraries. + if not self.basis_library.has_parametric: + # Drop any non-finite-on-training basis that slipped into the model + # so predict() can never return NaN because of it. + if has_invalid: + self._drop_nonfinite_terms(Phi, y, invalid_mask) + + # Drop numerically negligible terms (also removes spuriously + # selected bases that would diverge out of the training domain at + # predict time). + if self.prune_tol and self.prune_tol > 0: + self._prune_negligible_terms(Phi, y) + # Apply constraints if specified if self.constraints is not None: self._apply_constraints(Phi, y, X) @@ -362,6 +397,123 @@ def fit( self._is_fitted = True return self + def _drop_nonfinite_terms( + self, + Phi: jnp.ndarray, + y: jnp.ndarray, + invalid_mask: jnp.ndarray, + ): + """ + Remove selected terms whose basis is non-finite on the training data. + + Such a term would make :meth:`predict` return NaN (the basis is + re-evaluated from scratch there). The remaining terms are refit by + ordinary least squares and the information criteria recomputed. + + Parameters + ---------- + Phi : jnp.ndarray + Design matrix with invalid columns already zeroed. + y : jnp.ndarray + Target values. + invalid_mask : jnp.ndarray of bool + Per-column mask; True marks a basis non-finite on training data. + """ + invalid = np.asarray(invalid_mask) + indices = [int(i) for i in np.asarray(self._result.selected_indices)] + keep = [i for i in indices if not invalid[i]] + if len(keep) == len(indices): + return # nothing invalid was selected + + if not keep: + # Degenerate: fall back to the constant term if the library has one. + keep = [i for i, name in enumerate(self.basis_library.names) if name == "1"][:1] + if not keep: + return + + self._result = self._refit_subset(keep, Phi, y) + + def _prune_negligible_terms(self, Phi: jnp.ndarray, y: jnp.ndarray): + """ + Drop selected terms whose contribution to the fit is negligible. + + A term with a near-zero coefficient can still make :meth:`predict` + non-finite if its basis diverges out of the training domain (e.g. + ``exp(x0/x1)`` near ``x1 = 0``). Terms contributing less than + ``prune_tol`` times the largest term's contribution are removed and the + remainder refit. + + Parameters + ---------- + Phi : jnp.ndarray + Design matrix (invalid columns already zeroed). + y : jnp.ndarray + Target values. + """ + indices = [int(i) for i in np.asarray(self._result.selected_indices)] + if len(indices) <= 1: + return + + coeffs = np.asarray(self._result.coefficients) + Phi_sub = np.asarray(Phi[:, jnp.asarray(indices)]) + contribution = np.abs(coeffs) * np.linalg.norm(Phi_sub, axis=0) + max_contribution = float(contribution.max()) + if max_contribution <= 0: + return + + keep = [ + idx + for idx, contrib in zip(indices, contribution, strict=False) + if contrib >= self.prune_tol * max_contribution + ] + if len(keep) == len(indices): + return # nothing negligible + + self._result = self._refit_subset(keep, Phi, y) + + def _refit_subset(self, keep: list[int], Phi: jnp.ndarray, y: jnp.ndarray) -> SelectionResult: + """ + Refit ordinary least squares on a subset of basis columns. + + Parameters + ---------- + keep : list of int + Library indices of the terms to retain. + Phi : jnp.ndarray + Design matrix. + y : jnp.ndarray + Target values. + + Returns + ------- + SelectionResult + Result for the retained terms with recomputed MSE and information + criteria. + """ + keep_arr = jnp.asarray(keep) + coeffs, mse = fit_ols(Phi[:, keep_arr], y) + names = [self.basis_library.names[i] for i in keep] + complexity = int(sum(self.basis_library.complexities[i] for i in keep)) + n, k = len(y), len(coeffs) + + parametric_params = self._result.parametric_params + if parametric_params: + parametric_params = {i: parametric_params[i] for i in keep if i in parametric_params} + parametric_params = parametric_params or None + + return SelectionResult( + coefficients=coeffs, + selected_indices=keep_arr, + selected_names=names, + mse=mse, + complexity=complexity, + aic=compute_information_criterion(n, k, mse, "aic"), + bic=compute_information_criterion(n, k, mse, "bic"), + aicc=compute_information_criterion(n, k, mse, "aicc"), + n_samples=n, + parametric_params=parametric_params, + ) + def _apply_constraints( self, Phi: jnp.ndarray, @@ -1135,6 +1287,7 @@ def _state_dict(self) -> dict: "param_optimizer": self.param_optimizer, "param_optimization_budget": self.param_optimization_budget, "constraint_enforcement": self.constraint_enforcement, + "prune_tol": self.prune_tol, }, "basis_library": self.basis_library.to_dict(), "result": self._result.to_dict(), @@ -1164,6 +1317,7 @@ def _from_dict(cls, data: dict) -> SymbolicRegressor: config = data["config"] config.setdefault("constraint_enforcement", "penalty") + config.setdefault("prune_tol", 1e-6) model = cls( basis_library=basis_library, diff --git a/src/jaxsr/skill/SKILL.md b/src/jaxsr/skill/SKILL.md index 4bbd7eb..5b1c908 100644 --- a/src/jaxsr/skill/SKILL.md +++ b/src/jaxsr/skill/SKILL.md @@ -227,6 +227,67 @@ next_pts = study.suggest_next(n_points=5, strategy="uncertainty") study.save("catalyst.jaxsr") ``` +### Additive (Boosting-Style) Symbolic Regression + +Use when the signal is a sum of several simple effects and you want many small +interpretable terms instead of one large expression. Fits `f(x) = c + Σ_k η_k · +g_k(x)` by stagewise residual fitting (analogous to gradient boosting with +symbolic weak learners). Reuses `fit_symbolic` for each term. + +```python +from jaxsr.additive import StagewiseSymbolicRegressor + +model = StagewiseSymbolicRegressor( + n_terms=10, # number of boosting stages (terms) + learning_rate=0.2, # shrinkage (used when refit_coefficients=False) + max_complexity=4, # max basis terms per stage — keep small + refit_coefficients=True, # re-solve all linear weights by OLS each stage + early_stopping=False, # stop on a validation split when it stops improving + validation_fraction=0.2, +) +model.fit(X, y) + +print(model) # pretty structural summary +model.expressions_ # per-term expression strings +model.intercept_, model.coefficients_ +model.predict(X_new) +model.to_expression() # single combined SymPy expression +model.save("additive.json") # JSON round-trip (models are NOT picklable) +loaded = StagewiseSymbolicRegressor.load("additive.json") +``` + +Notes: +- Prefer `refit_coefficients=True` for accuracy with squared error; keep + `max_complexity` small (2–4) to favour many simple terms. +- **Robust / quantile regression:** set `loss` to `"absolute_error"`, + `"huber"`, or `"quantile"` (or an instance like `QuantileLoss(0.9)` / + `HuberLoss(delta=2.0)`). These are fit by gradient boosting with a per-stage + line search; use `refit_coefficients=False` (OLS refit only applies to + squared error and is auto-disabled with a warning otherwise). Fit several + quantiles to build prediction intervals. +- **Structural uncertainty:** `bootstrap_additive(model, X, y, n_bootstrap=...)` + refits on resamples and returns `["inclusion_probabilities"]` (how often each + basis is selected — a posterior-inclusion-probability proxy) and `["models"]`; + pass those to `bootstrap_predict_additive(models, X_new)` for an ensemble + prediction interval. Diffuse probabilities (~0.5) flag that the data don't + determine one expression (common with collinear features). +- `include_transcendental`/`include_ratios` are off by default; if enabled, a + stage that would produce non-finite predictions falls back to a finite basis. +- **Compositional discovery (experimental):** `RecursiveSymbolicRegressor` + grows the basis library along the residual (composing unary funcs, products, + ratios of discovered terms) to reach compositions a flat library misses + (e.g. `exp(x0*x1)`). Deterministic FFX-style search; competitive with GP on + simple targets but won't match PySR/Operon on hard ones. Not serialisable + (composed bases are closures). +- `BackfittingSymbolicRegressor` (GAM-style: a fixed set of terms revised + across sweeps, warm-started from stagewise; squared error only) is available. + It is never worse than stagewise+refit on training and helps specifically + when `max_complexity` is small (single-basis terms) and features are + collinear — where greedy forward selection gets stuck and re-discovery + escapes it. With larger per-term budgets it matches stagewise+refit, so + prefer `StagewiseSymbolicRegressor` there. A Bayesian (BART/iBART) variant is + future work. + ## Quick Reference: CLI ```bash @@ -325,6 +386,11 @@ See `guides/rsm.md` for RSM designs, canonical analysis, and optimization. See `guides/active-learning.md` for acquisition functions and adaptive sampling. +### "One expression isn't enough / the signal is a sum of many effects" + +See `guides/additive.md` for boosting-style additive symbolic regression +(`StagewiseSymbolicRegressor`): fit residuals stagewise into many small terms. + ## Templates Ready-to-use scripts and notebook starters are in `templates/`: @@ -332,6 +398,7 @@ Ready-to-use scripts and notebook starters are in `templates/`: | Template | Use Case | |----------|----------| | `basic-regression.py` | Discover an equation from X, y data | +| `additive-regression.py` | Boosting-style additive SR: robust/quantile losses, bootstrap UQ, backfitting, recursive expansion | | `constrained-model.py` | Add physical constraints to model | | `doe-study.py` | Full DOE workflow from design to report | | `uncertainty-analysis.py` | Compare all UQ methods | diff --git a/src/jaxsr/skill/guides/additive.md b/src/jaxsr/skill/guides/additive.md new file mode 100644 index 0000000..838f582 --- /dev/null +++ b/src/jaxsr/skill/guides/additive.md @@ -0,0 +1,332 @@ +# Additive Symbolic Regression + +Additive symbolic regression fits a model as a **sum of small symbolic +expressions**: + +``` +f(x) = c + eta_1 * g_1(x) + eta_2 * g_2(x) + ... + eta_K * g_K(x) +``` + +where each `g_k(x)` is a small, interpretable symbolic expression discovered by +the existing JAXSR machinery. This is analogous to **gradient boosting**, except +each weak learner is a symbolic expression rather than a decision tree. + +The submodule lives in `jaxsr.additive`. + +## Three flavours of symbolic regression + +| Approach | What it does | Status | +|----------|--------------|--------| +| **Single-expression** (`jaxsr.SymbolicRegressor`) | Fits one sparse expression over a fixed basis library. | Available | +| **Stagewise additive** (`jaxsr.additive.StagewiseSymbolicRegressor`) | Repeatedly fits a small expression to the *residual* and adds it to the ensemble. Old terms are **frozen**. | Available | +| **Backfitting additive** (`jaxsr.additive.BackfittingSymbolicRegressor`) | Maintains a fixed set of terms and **revises** each one in place across sweeps (GAM-style). | Available (squared error); Bayesian variant planned | + +The key distinction between the two additive variants: + +- **Stagewise**: once a term is discovered it never changes; only its linear + weight may be re-estimated. +- **Backfitting**: terms are revised repeatedly, each conditioned on the current + fit of all the others. + +## Scope: what this is (and isn't) good for + +**In one line:** JAXSR is a *linear method over a fixed feature space* — it +selects a sparse combination of basis functions you supply. It is not a +free-composition equation discoverer. + +That distinction decides whether it is the right tool: + +- **Good fit:** the right building blocks are on the menu (or you can add them + with `BasisLibrary.add_custom`), and you want an interpretable, robust, + uncertainty-aware additive model. On targets that live in the library it is + fast and accurate, and the additive layer adds robust/quantile losses and + structural-uncertainty bootstrapping that genetic-programming tools don't + offer out of the box. +- **Wrong fit:** you want to *discover* an unknown compositional law such as + `exp(x0*x1)`, `x0 / (1 + x1**2)`, or `sin(2*x0)`. These are not single basis + functions, and the space of such compositions is infinite and continuously + parameterized, so no fixed library enumerates them in advance. For that, + reach for a genetic-programming or neural symbolic-regression tool (PySR, + Operon, AI-Feynman), which *search* the space of expressions instead of + selecting from a fixed dictionary — or try the experimental + [`RecursiveSymbolicRegressor`](#recursive-basis-expansion-experimental), which + grows compositions along the residual and partially lifts this ceiling. + +The limit is one of **discovery, not representation**: the linear-in-basis model +fits any of those targets perfectly the moment the exact term is in the library +(e.g. `library.add_custom("exp(x0*x1)", lambda X: jnp.exp(X[:, 0] * X[:, 1]))`) — +it simply cannot figure out *which* composition it needs without being told. +JAXSR's parametric bases can additionally fit a few constants *inside* a +pre-specified nonlinearity (e.g. `sin(a*x0)` with `a` optimized), but that still +requires you to name the functional form. The additive extensions in this guide +raise the *statistical* sophistication (boosting, robust/quantile losses, +backfitting, structural UQ); they do not change this expressiveness boundary. + +## Quick start + +```python +import numpy as np +from jaxsr.additive import StagewiseSymbolicRegressor + +rng = np.random.default_rng(0) +X = rng.uniform(-2, 2, size=(200, 2)) +y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + 0.1 * rng.normal(size=200) + +model = StagewiseSymbolicRegressor( + n_terms=5, + learning_rate=0.2, + max_complexity=6, + refit_coefficients=True, +) +model.fit(X, y) + +print(model) # pretty structural summary +print(model.expressions_) # per-term expression strings +print(model.coefficients_) # per-term weights +print(model.intercept_) # additive intercept +y_pred = model.predict(X) +``` + +The `print(model)` output looks like: + +``` +StagewiseSymbolicRegressor( + intercept = 1.07 + terms = + + 1 * (y = 2*x0 - 1.07 + 0.5*x1^2) + ... +) +``` + +## The stagewise algorithm + +1. Initialise the intercept to `mean(y)` and the prediction to that constant. +2. Compute the residual `y - prediction`. +3. Fit a small symbolic expression `g_k` to the residual (via + `jaxsr.fit_symbolic`). +4. Append `g_k` to the ensemble. +5. If `refit_coefficients=True`, rebuild the design matrix `Phi[:, j] = g_j(X)` + and re-solve `y ~= intercept + Phi @ coefficients` by least squares. + Otherwise, update `prediction += learning_rate * g_k(X)`. +6. Record train (and optional validation) loss. +7. Repeat until `n_terms` terms are added or early stopping triggers. + +## Key parameters + +| Parameter | Meaning | +|-----------|---------| +| `n_terms` | Maximum number of boosting stages (terms). | +| `learning_rate` | Shrinkage on each stage when `refit_coefficients=False`. | +| `max_complexity` | Complexity budget per term (max basis terms). Keep small to favour many simple terms. | +| `refit_coefficients` | Re-solve all linear weights by OLS after each stage. | +| `loss` | `"squared_error"` (default), `"absolute_error"`, `"huber"`, `"quantile"`, or a `Loss` instance. See [Losses](#losses-robust-and-quantile-regression). | +| `early_stopping` | Hold out a validation split and stop when it stops improving. | +| `validation_fraction`, `patience`, `min_delta` | Early-stopping controls. | +| `max_poly_degree`, `include_transcendental`, `include_ratios` | Which basis functions each term may use. | +| `information_criterion` | Complexity control within each term (`"aic"`, `"aicc"`, `"bic"`). | + +## Coefficient refitting + +- `refit_coefficients=False`: the weights are the learning-rate-scaled stagewise + weights (`coefficients_[k] == learning_rate`). +- `refit_coefficients=True`: after each new term, the intercept and *all* per-term + weights are re-solved by ordinary least squares over the discovered symbolic + features. This decouples term discovery (nonlinear, greedy) from term weighting + (linear, global) and typically improves accuracy. + +The refit uses `jnp.linalg.lstsq` (SVD-based, minimum-norm), so the highly +correlated columns produced by later boosting stages do not cause instability. + +## Combined expression + +`model.to_expression()` returns a single simplified SymPy expression combining +all terms: + +```python +expr = model.to_expression() # requires sympy +``` + +## Saving and loading + +Fitted models serialize to JSON (each term is stored via the underlying +`SymbolicRegressor` state), mirroring the rest of jaxsr. Note that the models +are **not picklable** — the basis-function closures cannot be pickled — so use +`save`/`load` rather than `pickle`: + +```python +model.save("additive_model.json") +loaded = StagewiseSymbolicRegressor.load("additive_model.json") +``` + +## Structural uncertainty (bootstrap) + +A single fitted expression can hide the fact that the *structure* itself is +uncertain — several different basis sets may explain the data about equally +well (this is common with collinear features). `bootstrap_additive` refits the +model on bootstrap resamples and reports, for each basis function, how often it +is selected — a cheap approximation to a posterior inclusion probability — +together with a predictive ensemble: + +```python +from jaxsr.additive import ( + StagewiseSymbolicRegressor, + bootstrap_additive, + bootstrap_predict_additive, +) + +est = StagewiseSymbolicRegressor(n_terms=3, max_complexity=2) +res = bootstrap_additive(est, X, y, n_bootstrap=100, random_state=0) + +# How stable is the discovered structure? +for name, prob in res["inclusion_probabilities"].items(): + print(f"{name:10s} selected in {prob:.0%} of resamples") + +# Prediction intervals that reflect *structural* variability, not just noise +pi = bootstrap_predict_additive(res["models"], X_new, alpha=0.1) +pi["mean"], pi["lower"], pi["upper"] +``` + +**How to read it.** Inclusion probabilities near 0 or 1 mean the structure is +identifiable and the single fitted expression is trustworthy. **Diffuse** values +(e.g. a basis selected 50–60% of the time) mean the data do not determine one +expression — no single symbolic model should be over-trusted, and the bootstrap +intervals are the honest summary. This also works as a decision gate for heavier +Bayesian modelling: if the probabilities are already crisp, there is little +structural uncertainty left to quantify. It works for both the stagewise and +backfitting regressors. + +## Early stopping + +With `early_stopping=True`, a validation split (`validation_fraction`) is held +out. After each stage the validation loss is recorded; training stops once it +fails to improve by at least `min_delta` for `patience` consecutive stages, and +the model rolls back to the best iteration. + +## Losses: robust and quantile regression + +This is where additive symbolic regression goes beyond ordinary least-squares +symbolic regression. Each weak learner fits the negative gradient `-dL/dy_pred` +(gradient boosting), so you can target losses that OLS selection cannot: + +| `loss` | Class | Use when | +|--------|-------|----------| +| `"squared_error"` (default) | `SquaredError` | Standard regression | +| `"absolute_error"` | `AbsoluteError` | Outliers present (fits the median) | +| `"huber"` | `HuberLoss(delta=1.35)` | Outliers, but keep efficiency near zero | +| `"quantile"` | `QuantileLoss(quantile=0.5)` | Quantiles / prediction intervals / asymmetric cost | + +Pass a name for defaults, or an instance to customise: + +```python +from jaxsr.additive import StagewiseSymbolicRegressor, QuantileLoss, HuberLoss + +# Robust regression: heavy outliers barely move the fit +robust = StagewiseSymbolicRegressor(loss="huber", learning_rate=0.5).fit(X, y) + +# 90th-percentile regression (build intervals by fitting several quantiles) +q90 = StagewiseSymbolicRegressor(loss=QuantileLoss(0.9), learning_rate=0.5).fit(X, y) +``` + +**How non-squared losses are fit.** Each stage fits a symbolic term to the +negative gradient, then a **line search** picks the step size that minimises the +loss (`learning_rate` shrinks that step). Because the ordinary least-squares +coefficient refit targets squared error, `refit_coefficients=True` is ignored for +non-squared losses (a warning is issued) and gradient boosting is used instead — +so set `refit_coefficients=False` explicitly for robust/quantile models. + +The optimal constant initialisation adapts to the loss: mean for squared error, +median for absolute/Huber, and the empirical quantile for quantile loss. + +Add further losses (Poisson, logistic, ...) by subclassing `Loss` and +registering them in `jaxsr.additive.losses._LOSSES`. + +## Backfitting (GAM-style) + +`BackfittingSymbolicRegressor` maintains a **fixed** number of terms and +*revises* each one across sweeps, instead of freezing them. Each sweep removes a +term, re-discovers its expression on the partial residual, and puts it back: + +```python +from jaxsr.additive import BackfittingSymbolicRegressor + +model = BackfittingSymbolicRegressor(n_terms=4, n_sweeps=6, max_complexity=3) +model.fit(X, y) # warm-started from a stagewise fit, then refined by sweeps +``` + +``` +for sweep in 1..n_sweeps: + for term j: + partial_residual = y - intercept - sum_{i != j} coef_i * g_i(X) + g_j = fit_symbolic(X, partial_residual, ...) # re-discover structure + intercept, coef = OLS refit over all terms +(stop when the training loss stops improving by `tol`) +``` + +It is warm-started from a stagewise fit and currently supports **squared error +only**. Structure re-discovery makes the sweep a heuristic (no monotonicity +guarantee), so the best-loss iterate is kept. + +**When does it actually help?** Backfitting starts from the stagewise+refit fit +and keeps the best-loss iterate, so **it is never worse than +`StagewiseSymbolicRegressor(refit_coefficients=True)` on the training data** — +the only question is whether the sweeps improve on it. That hinges entirely on +whether re-discovery changes the *set* of selected basis functions: + +- **Generous per-term budget** (`max_complexity` ≥ 2–3): greedy usually already + finds a sufficient basis set, so the joint least-squares refit makes the two + essentially identical. Backfitting adds nothing here — prefer the stagewise + regressor. +- **Small per-term budget and collinear features** (`max_complexity=1`, the + GAM-style single-basis regime): greedy forward selection can lock into a + *suboptimal* basis set that a single forward pass cannot undo. Backfitting's + coordinate-descent re-discovery escapes it, changing the basis union and + improving the fit — we have measured up to roughly **+0.04 train / +0.06 test + R²** in this regime, with no downside in the cases where it does not help. + +So reach for backfitting when you want **small, revisable single-basis terms +over correlated features** (or a fixed-size GAM-style decomposition); use the +stagewise regressor for larger per-term expressions. Its other forward-looking +value is as the foundation for a future **Bayesian backfitting** variant +(BART/iBART-style), which would sample a *posterior over symbolic structure* — +genuinely beyond point-estimate SR — using the same partial-residual sweep with +conjugate marginal likelihoods. + +## Recursive basis expansion (experimental) + +The [Scope](#scope-what-this-is-and-isnt-good-for) section notes that a fixed +library cannot *discover* compositional forms like `x0*sin(x1)` or `exp(x0*x1)`. +`RecursiveSymbolicRegressor` is an experimental step past that ceiling. Instead +of enumerating a huge composition space up front (which explodes +combinatorially), it **grows the library lazily along the residual**: + +1. Fit a sparse model over the current library; take the residual. +2. Compose the currently useful terms (selected terms + features) with a small + operator set (unary functions, products, ratios) — one new layer. +3. Screen: drop non-finite candidates, deduplicate, keep the top few by + correlation with the residual. +4. Add the survivors and refit. Repeat. The effective composition depth is the + number of rounds, because a term found in one round feeds the next. + +```python +from jaxsr.additive import RecursiveSymbolicRegressor + +model = RecursiveSymbolicRegressor(n_expansions=3, max_terms=6, beam_width=25) +model.fit(X, y) +print(model.expression_) # e.g. recovers "exp((x0)*(x1))" exactly +print(model.history_) # library size / n_terms / train R^2 per round +``` + +This is essentially Fast Function Extraction (FFX) / symbolic feature +construction — a deterministic, bounded cousin of genetic programming. On simple +compositional targets it substantially beats a flat library (e.g. `exp(x0*x1)`: +R² 1.00 vs 0.70) and is competitive with a strong GP engine (matched Operon on +`x0*sin(x1)` in our tests). + +**Caveats.** It re-enters search-based territory: cost grows with `beam_width`, +`n_expansions`, and feature count, and it will not match a mature GP (PySR, +Operon) on hard, high-dimensional, or deeply nested targets. The result is an +ordinary `SymbolicRegressor` over the grown library (so predict/`expression_`/ +scoring and the base regressor's non-finite-basis guard and negligible-term +pruning all apply), but the composed bases are Python closures, so the fitted +model is **not** serialisable via `save`/`load`, and `to_sympy` may not parse +deeply nested term names. diff --git a/src/jaxsr/skill/templates/additive-regression.py b/src/jaxsr/skill/templates/additive-regression.py new file mode 100644 index 0000000..f1c6d2b --- /dev/null +++ b/src/jaxsr/skill/templates/additive-regression.py @@ -0,0 +1,99 @@ +""" +Additive Symbolic Regression — boosting-style ensembles of small expressions. + +This template shows the ``jaxsr.additive`` workflows: +1. Stagewise additive regression (gradient boosting with symbolic weak learners) +2. Robust regression under outliers (Huber / absolute-error losses) +3. Quantile regression (prediction intervals via the pinball loss) +4. Structural uncertainty (bootstrap basis-inclusion probabilities) +5. Backfitting (GAM-style: revise terms instead of freezing them) +6. Recursive expansion (reach compositions a flat library misses) + +Pick the section you need; each block is self-contained after the imports. +""" + +import numpy as np + +from jaxsr import fit_symbolic +from jaxsr.additive import ( + BackfittingSymbolicRegressor, + QuantileLoss, + RecursiveSymbolicRegressor, + StagewiseSymbolicRegressor, + bootstrap_additive, + bootstrap_predict_additive, +) + +# Replace with your own data. X shape (n_samples, n_features), y shape (n_samples,). +rng = np.random.default_rng(0) +X = rng.uniform(-2, 2, size=(300, 2)) +y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + 0.1 * rng.normal(size=300) + +# ============================================================================= +# 1. Stagewise additive regression +# Many small interpretable terms; keep max_complexity small. +# refit_coefficients=True re-solves all weights by least squares each stage. +# ============================================================================= +model = StagewiseSymbolicRegressor( + n_terms=10, + learning_rate=0.2, + max_complexity=4, + refit_coefficients=True, + early_stopping=False, # set True + validation_fraction to guard overfitting +) +model.fit(X, y) +print(model) # pretty structural summary +print(model.expressions_) # per-term expression strings +print(model.intercept_, model.coefficients_) +print("combined:", model.to_expression()) +model.save("additive_model.json") # JSON round-trip (models are NOT picklable) + +# ============================================================================= +# 2. Robust regression (outliers) — use refit_coefficients=False for any +# non-squared loss (OLS refit only applies to squared error). +# ============================================================================= +robust = StagewiseSymbolicRegressor( + loss="huber", # or "absolute_error" + n_terms=8, + max_complexity=3, + learning_rate=0.5, + refit_coefficients=False, +).fit(X, y) + +# ============================================================================= +# 3. Quantile regression — fit several quantiles to build a prediction band. +# ============================================================================= +q_models = { + q: StagewiseSymbolicRegressor( + loss=QuantileLoss(q), + n_terms=10, + max_complexity=3, + learning_rate=0.5, + refit_coefficients=False, + ).fit(X, y) + for q in (0.1, 0.5, 0.9) +} +lower, median, upper = (q_models[q].predict(X) for q in (0.1, 0.5, 0.9)) + +# ============================================================================= +# 4. Structural uncertainty — how stable is the discovered structure? +# ============================================================================= +res = bootstrap_additive(model, X, y, n_bootstrap=100, random_state=0) +print(res["inclusion_probabilities"]) # {basis: fraction selected} +pi = bootstrap_predict_additive(res["models"], X) # mean / lower / upper / ... + +# ============================================================================= +# 5. Backfitting (GAM-style) — fixed set of terms, revised across sweeps. +# ============================================================================= +bf = BackfittingSymbolicRegressor(n_terms=4, n_sweeps=6, max_complexity=3).fit(X, y) + +# ============================================================================= +# 6. Recursive expansion (experimental) — discover compositions like exp(x0*x1) +# that a flat library cannot reach. +# ============================================================================= +Xc = rng.uniform(-2, 2, size=(400, 2)) +yc = np.exp(Xc[:, 0] * Xc[:, 1]) +rec = RecursiveSymbolicRegressor(n_expansions=3, max_terms=6, beam_width=25).fit(Xc, yc) +print(rec.expression_) # e.g. "exp((x0)*(x1))" +# Compare with a flat library (which misses it): +flat = fit_symbolic(Xc, yc, max_terms=6, include_transcendental=True) diff --git a/tests/test_additive.py b/tests/test_additive.py new file mode 100644 index 0000000..9909ff4 --- /dev/null +++ b/tests/test_additive.py @@ -0,0 +1,714 @@ +"""Tests for the additive symbolic regression submodule.""" + +import jax.numpy as jnp +import numpy as np +import pytest + + +def _additive_data(n=200, noise=0.0, seed=0): + """y = 2.0 * x0 + 0.5 * x1**2 (+ optional noise).""" + rng = np.random.default_rng(seed) + X = rng.uniform(-2, 2, size=(n, 2)) + y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + if noise: + y = y + rng.normal(0, noise, size=n) + return jnp.array(X), jnp.array(y) + + +def test_module_imports(): + """The additive module and its public API import correctly.""" + from jaxsr.additive import ( + AdditiveSymbolicModel, + BackfittingSymbolicRegressor, + Loss, + SquaredError, + StagewiseSymbolicRegressor, + get_loss, + refit_ols, + ) + + assert StagewiseSymbolicRegressor is not None + assert BackfittingSymbolicRegressor is not None + assert AdditiveSymbolicModel is not None + assert issubclass(SquaredError, Loss) + assert callable(get_loss) + assert callable(refit_ols) + + +def test_fit_runs_and_predict_shape(): + """fit runs on a simple dataset and predict returns the right shape.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data() + model = StagewiseSymbolicRegressor(n_terms=3, max_complexity=4) + model.fit(X, y) + + assert model._is_fitted + y_pred = model.predict(X) + assert y_pred.shape == y.shape + assert model.n_terms_ >= 1 + + +def test_training_loss_decreases(): + """Training loss should not increase as terms are added (refit path).""" + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data(noise=0.05) + model = StagewiseSymbolicRegressor(n_terms=4, max_complexity=3, refit_coefficients=True) + model.fit(X, y) + + losses = [h["train_loss"] for h in model.training_history_] + assert len(losses) >= 2 + # OLS refit over a growing feature set is monotone non-increasing. + for prev, curr in zip(losses[:-1], losses[1:], strict=False): + assert curr <= prev + 1e-8 + # And it genuinely improved over the intercept-only model. + assert losses[-1] < losses[0] + + +def test_recovers_simple_additive_function(): + """The model should recover y = 2*x0 + 0.5*x1**2 accurately.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data(noise=0.0) + model = StagewiseSymbolicRegressor(n_terms=5, max_complexity=4, refit_coefficients=True) + model.fit(X, y) + + r2 = model.score(X, y) + assert r2 > 0.99 + + +def test_refit_coefficients_false_predicts(): + """refit_coefficients=False produces a working model.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data(noise=0.05) + model = StagewiseSymbolicRegressor( + n_terms=6, learning_rate=0.5, max_complexity=3, refit_coefficients=False + ) + model.fit(X, y) + + y_pred = model.predict(X) + assert y_pred.shape == y.shape + # Stagewise weights are the learning rate. + assert all(abs(c - 0.5) < 1e-12 for c in model.coefficients_) + # Should still reduce error relative to the mean baseline. + assert model.score(X, y) > 0.5 + + +def test_repr_includes_terms(): + """The string representation lists the learned terms.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data() + model = StagewiseSymbolicRegressor(n_terms=3, max_complexity=4) + model.fit(X, y) + + text = str(model) + assert "StagewiseSymbolicRegressor" in text + assert "intercept" in text + assert "terms" in text + # At least one coefficient/term line rendered. + assert "*" in text + + +def test_repr_before_fit_is_sklearn_style(): + """Before fitting, repr falls back to the sklearn-style parameter form.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + model = StagewiseSymbolicRegressor(n_terms=7) + text = repr(model) + assert text.startswith("StagewiseSymbolicRegressor(") + assert "n_terms=7" in text + + +def test_early_stopping(): + """Early stopping stops before n_terms on a small validation split.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data(n=120, noise=0.02) + model = StagewiseSymbolicRegressor( + n_terms=20, + max_complexity=3, + refit_coefficients=True, + early_stopping=True, + validation_fraction=0.25, + patience=2, + random_state=0, + ) + model.fit(X, y) + + # Every stage recorded a validation loss. + assert all(h["val_loss"] is not None for h in model.training_history_) + # It should stop early rather than using all 20 terms. + assert model.n_terms_ < 20 + assert model.score(X, y) > 0.9 + + +def test_to_expression(): + """to_expression returns a combined SymPy expression.""" + pytest.importorskip("sympy") + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data() + model = StagewiseSymbolicRegressor(n_terms=3, max_complexity=4) + model.fit(X, y) + + expr = model.to_expression() + assert expr is not None + # Should reference at least one feature symbol. + assert any(sym.name in {"x0", "x1"} for sym in expr.free_symbols) + + +def test_fitted_attributes_consistent(): + """coefficients_, expressions_, and terms_ have matching lengths.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data() + model = StagewiseSymbolicRegressor(n_terms=4, max_complexity=3) + model.fit(X, y) + + n = model.n_terms_ + assert len(model.coefficients_) == n + assert len(model.expressions_) == n + assert len(model.terms_) == n + assert len(model.learning_rates_) == n + assert isinstance(model.intercept_, float) + + +def test_predict_before_fit_raises(): + """Calling predict before fit raises a clear error.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + model = StagewiseSymbolicRegressor() + with pytest.raises(RuntimeError): + model.predict(np.zeros((3, 2))) + + +def test_invalid_params_raise(): + """Invalid constructor parameters are rejected at fit time.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data(n=40) + with pytest.raises(ValueError): + StagewiseSymbolicRegressor(n_terms=0).fit(X, y) + with pytest.raises(ValueError): + StagewiseSymbolicRegressor(learning_rate=0.0).fit(X, y) + + +def test_non_finite_inputs_raise(): + """Non-finite X or y are rejected rather than silently producing NaN.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data(n=40) + + y_nan = y.at[0].set(jnp.nan) + with pytest.raises(ValueError, match="non-finite"): + StagewiseSymbolicRegressor(n_terms=2, max_complexity=2).fit(X, y_nan) + + X_inf = X.at[0, 0].set(jnp.inf) + with pytest.raises(ValueError, match="non-finite"): + StagewiseSymbolicRegressor(n_terms=2, max_complexity=2).fit(X_inf, y) + + +def test_predict_feature_count_mismatch_raises(): + """predict rejects inputs with the wrong number of features.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data(n=50, seed=1) + model = StagewiseSymbolicRegressor(n_terms=2, max_complexity=2).fit(X, y) + + with pytest.raises(ValueError, match="features"): + model.predict(np.zeros((5, 5))) + with pytest.raises(ValueError, match="features"): + model.predict(np.zeros((5, 1))) + + +def test_single_feature_and_tiny_sample(): + """The model handles a single feature and very small sample sizes.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + rng = np.random.default_rng(3) + X = jnp.array(rng.uniform(-2, 2, size=(60, 1))) + y = 2.0 * X[:, 0] + model = StagewiseSymbolicRegressor(n_terms=3, max_complexity=2).fit(X, y) + assert model.score(X, y) > 0.99 + + X_small, y_small = _additive_data(n=3) + model2 = StagewiseSymbolicRegressor(n_terms=2, max_complexity=2).fit(X_small, y_small) + assert model2.predict(X_small).shape == (3,) + + +def test_transcendental_terms_stay_finite(): + """Transcendental bases invalid on the domain must not yield NaN models. + + A ``log``/``sqrt``/``1/x`` basis can be selected yet produce NaN on data + with the wrong sign; the stage must fall back to a finite basis so the + ensemble never predicts NaN. + """ + from jaxsr.additive import StagewiseSymbolicRegressor + + rng = np.random.default_rng(6) + X = jnp.array(rng.uniform(-1.5, 1.5, size=(300, 2))) + y = jnp.exp(0.5 * X[:, 0]) + X[:, 1] + model = StagewiseSymbolicRegressor( + n_terms=5, max_complexity=3, include_transcendental=True, include_ratios=True + ).fit(X, y) + + preds = np.array(model.predict(X)) + assert np.all(np.isfinite(preds)), "model produced non-finite predictions" + assert model.score(X, y) > 0.9 + + +def test_determinism_with_random_state(): + """Two fits with the same random_state produce identical predictions.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data(n=80, noise=0.05, seed=2) + kw = { + "n_terms": 5, + "max_complexity": 3, + "early_stopping": True, + "validation_fraction": 0.25, + "random_state": 42, + } + a = StagewiseSymbolicRegressor(**kw).fit(X, y) + b = StagewiseSymbolicRegressor(**kw).fit(X, y) + assert np.allclose(np.array(a.predict(X)), np.array(b.predict(X))) + + +def test_save_load_roundtrip(tmp_path): + """save/load reconstructs a model with identical predictions and state.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data(n=120, noise=0.05, seed=4) + model = StagewiseSymbolicRegressor(n_terms=4, max_complexity=3, feature_names=["a", "b"]).fit( + X, y + ) + + path = tmp_path / "additive_model.json" + model.save(str(path)) + loaded = StagewiseSymbolicRegressor.load(str(path)) + + assert np.allclose(np.array(model.predict(X)), np.array(loaded.predict(X)), atol=1e-6) + assert loaded.intercept_ == pytest.approx(model.intercept_) + assert np.allclose(loaded.coefficients_, model.coefficients_) + assert loaded.expressions_ == model.expressions_ + assert loaded.n_terms_ == model.n_terms_ + + +def test_backfitting_fit_and_recover(): + """Backfitting fits, predicts the right shape, and recovers a simple target.""" + from jaxsr.additive import BackfittingSymbolicRegressor + + X, y = _additive_data(n=300, noise=0.05, seed=1) + model = BackfittingSymbolicRegressor(n_terms=3, n_sweeps=5, max_complexity=3).fit(X, y) + + assert model._is_fitted + assert model.predict(X).shape == y.shape + assert model.n_terms_ == 3 + assert model.score(X, y) > 0.99 + # A per-sweep training-loss history is recorded (sweep 0 = warm start). + assert len(model.training_history_) >= 2 + assert all("train_loss" in h for h in model.training_history_) + + +def test_backfitting_never_worse_than_warm_start(): + """Backfitting starts from stagewise+refit and keeps the best iterate, so + its training fit is never worse than stagewise+refit's.""" + from jaxsr.additive import BackfittingSymbolicRegressor, StagewiseSymbolicRegressor + + X, y = _additive_data(n=300, noise=0.1, seed=2) + cfg = {"n_terms": 4, "max_complexity": 3} + sw = StagewiseSymbolicRegressor(refit_coefficients=True, **cfg).fit(X, y) + bf = BackfittingSymbolicRegressor(n_sweeps=6, **cfg).fit(X, y) + assert bf.score(X, y) >= sw.score(X, y) - 1e-6 + + +def test_backfitting_helps_with_collinear_single_basis_terms(): + """With single-basis terms and collinear features, greedy selection can get + stuck; backfitting re-discovery escapes it and is at least as good.""" + from jaxsr.additive import BackfittingSymbolicRegressor, StagewiseSymbolicRegressor + + rng = np.random.default_rng(0) + x0 = rng.normal(0, 1, 500) + x1 = 0.9 * x0 + 0.1 * rng.normal(0, 1, 500) + x2 = 0.8 * x0 + 0.2 * rng.normal(0, 1, 500) + X = jnp.array(np.column_stack([x0, x1, x2])) + y = jnp.array(x0 - x1 + 0.5 * x2 + rng.normal(0, 0.1, 500)) + + cfg = {"n_terms": 3, "max_complexity": 1} + sw = StagewiseSymbolicRegressor(refit_coefficients=True, **cfg).fit(X, y) + bf = BackfittingSymbolicRegressor(n_sweeps=10, **cfg).fit(X, y) + # Backfitting is never worse, and in this stuck-greedy regime it improves. + assert bf.score(X, y) >= sw.score(X, y) - 1e-6 + assert bf.score(X, y) > 0.9 + + +def test_backfitting_nonsquared_loss_not_implemented(): + """Backfitting only supports squared error for now.""" + from jaxsr.additive import BackfittingSymbolicRegressor + + X, y = _additive_data(n=60) + with pytest.raises(NotImplementedError): + BackfittingSymbolicRegressor(loss="huber").fit(X, y) + + +def test_backfitting_invalid_params_and_inputs(): + """Backfitting validates parameters and inputs.""" + from jaxsr.additive import BackfittingSymbolicRegressor + + X, y = _additive_data(n=60) + with pytest.raises(ValueError): + BackfittingSymbolicRegressor(n_terms=0).fit(X, y) + with pytest.raises(ValueError): + BackfittingSymbolicRegressor(n_sweeps=0).fit(X, y) + with pytest.raises(ValueError, match="non-finite"): + BackfittingSymbolicRegressor().fit(X, y.at[0].set(jnp.nan)) + + +def test_bootstrap_additive_structure_and_reproducibility(): + """Bootstrap yields valid inclusion probabilities and is reproducible.""" + from jaxsr.additive import StagewiseSymbolicRegressor, bootstrap_additive + + X, y = _additive_data(n=250, noise=0.1, seed=1) + est = StagewiseSymbolicRegressor(n_terms=3, max_complexity=3) + res = bootstrap_additive(est, X, y, n_bootstrap=25, random_state=0) + + assert res["n_bootstrap"] == 25 + assert len(res["models"]) == 25 + assert res["n_terms"].shape == (25,) + probs = res["inclusion_probabilities"] + assert all(0.0 <= p <= 1.0 for p in probs.values()) + # On identifiable data the true bases are selected almost every time. + assert probs.get("x0", 0.0) > 0.9 + assert probs.get("x1^2", 0.0) > 0.9 + + # Reproducible for a fixed random_state. + res2 = bootstrap_additive(est, X, y, n_bootstrap=25, random_state=0) + assert res["inclusion_probabilities"] == res2["inclusion_probabilities"] + + +def test_bootstrap_additive_detects_structural_instability(): + """Collinear features give diffuse (non-degenerate) inclusion probs.""" + from jaxsr.additive import StagewiseSymbolicRegressor, bootstrap_additive + + rng = np.random.default_rng(0) + x0 = rng.normal(0, 1, 400) + x1 = 0.9 * x0 + 0.1 * rng.normal(0, 1, 400) + x2 = 0.8 * x0 + 0.2 * rng.normal(0, 1, 400) + X = jnp.array(np.column_stack([x0, x1, x2])) + y = jnp.array(x0 - x1 + 0.5 * x2 + rng.normal(0, 0.1, 400)) + + est = StagewiseSymbolicRegressor(n_terms=3, max_complexity=1, refit_coefficients=True) + probs = bootstrap_additive(est, X, y, n_bootstrap=40, random_state=1)["inclusion_probabilities"] + # At least one basis is genuinely uncertain (selected sometimes, not always). + assert any(0.15 < p < 0.85 for p in probs.values()) + + +def test_bootstrap_predict_additive_intervals(): + """The predictive ensemble returns well-ordered intervals.""" + from jaxsr.additive import ( + StagewiseSymbolicRegressor, + bootstrap_additive, + bootstrap_predict_additive, + ) + + X, y = _additive_data(n=200, noise=0.2, seed=2) + est = StagewiseSymbolicRegressor(n_terms=3, max_complexity=3) + res = bootstrap_additive(est, X, y, n_bootstrap=30, random_state=0) + + pi = bootstrap_predict_additive(res["models"], X, alpha=0.1) + n = X.shape[0] + assert pi["mean"].shape == (n,) + assert pi["predictions"].shape == (30, n) + assert np.all(pi["lower"] <= pi["median"] + 1e-9) + assert np.all(pi["median"] <= pi["upper"] + 1e-9) + + +def test_bootstrap_additive_works_with_backfitting(): + """Bootstrap uncertainty works for the backfitting regressor too.""" + from jaxsr.additive import BackfittingSymbolicRegressor, bootstrap_additive + + X, y = _additive_data(n=200, noise=0.1, seed=3) + est = BackfittingSymbolicRegressor(n_terms=3, n_sweeps=3, max_complexity=2) + res = bootstrap_additive(est, X, y, n_bootstrap=12, random_state=0) + assert res["n_bootstrap"] == 12 + assert all(0.0 <= p <= 1.0 for p in res["inclusion_probabilities"].values()) + + +def test_bootstrap_additive_validation(): + """Bootstrap functions validate their inputs.""" + from jaxsr.additive import ( + StagewiseSymbolicRegressor, + bootstrap_additive, + bootstrap_predict_additive, + ) + + X, y = _additive_data(n=60) + est = StagewiseSymbolicRegressor(n_terms=2, max_complexity=2) + with pytest.raises(ValueError): + bootstrap_additive(est, X, y, n_bootstrap=0) + with pytest.raises(ValueError): + bootstrap_predict_additive([], X) + res = bootstrap_additive(est, X, y, n_bootstrap=5, random_state=0) + with pytest.raises(ValueError): + bootstrap_predict_additive(res["models"], X, alpha=1.5) + + +def test_backfitting_save_load_roundtrip(tmp_path): + """Backfitting models round-trip through save/load.""" + from jaxsr.additive import BackfittingSymbolicRegressor + + X, y = _additive_data(n=150, noise=0.05, seed=3) + model = BackfittingSymbolicRegressor(n_terms=3, n_sweeps=4, max_complexity=3).fit(X, y) + + path = tmp_path / "backfit_model.json" + model.save(str(path)) + loaded = BackfittingSymbolicRegressor.load(str(path)) + + assert np.allclose(np.array(model.predict(X)), np.array(loaded.predict(X)), atol=1e-6) + assert loaded.n_terms_ == model.n_terms_ + assert loaded.expressions_ == model.expressions_ + + +def test_get_loss_and_squared_error(): + """The loss registry resolves names and squared error behaves correctly.""" + from jaxsr.additive import SquaredError, get_loss + + loss = get_loss("squared_error") + assert isinstance(loss, SquaredError) + + y = jnp.array([1.0, 2.0, 3.0]) + assert loss.initial_prediction(y) == pytest.approx(2.0) + resid = loss.negative_gradient(y, jnp.array([0.0, 0.0, 0.0])) + assert np.allclose(np.array(resid), np.array(y)) + assert loss.loss(y, y) == pytest.approx(0.0) + + with pytest.raises(ValueError): + get_loss("not_a_loss") + + +def test_loss_registry_and_parameters(): + """All registered losses resolve; parameterized losses validate inputs.""" + from jaxsr.additive import ( + AbsoluteError, + HuberLoss, + QuantileLoss, + SquaredError, + get_loss, + ) + + assert isinstance(get_loss("squared_error"), SquaredError) + assert isinstance(get_loss("absolute_error"), AbsoluteError) + assert isinstance(get_loss("huber"), HuberLoss) + assert isinstance(get_loss("quantile"), QuantileLoss) + # instances pass through unchanged + q = QuantileLoss(0.9) + assert get_loss(q) is q + + with pytest.raises(ValueError): + HuberLoss(delta=0.0) + with pytest.raises(ValueError): + QuantileLoss(quantile=0.0) + with pytest.raises(ValueError): + QuantileLoss(quantile=1.0) + + +def test_loss_gradients_and_initial_predictions(): + """Loss initial predictions and pseudo-residuals are correct.""" + from jaxsr.additive import AbsoluteError, HuberLoss, QuantileLoss + + y = jnp.array([1.0, 2.0, 3.0, 100.0]) # last is an outlier + + mae = AbsoluteError() + assert mae.initial_prediction(y) == pytest.approx(2.5) # median + grad = np.array(mae.negative_gradient(y, jnp.full(4, 2.5))) + assert np.allclose(grad, [-1.0, -1.0, 1.0, 1.0]) # sign(y - pred) + + huber = HuberLoss(delta=1.0) + r = np.array(huber.negative_gradient(jnp.array([0.0, 0.0]), jnp.array([-0.5, -5.0]))) + assert r[0] == pytest.approx(0.5) # within delta -> raw residual + assert r[1] == pytest.approx(1.0) # beyond delta -> clipped to delta + + q = QuantileLoss(0.9) + assert q.initial_prediction(jnp.arange(0.0, 101.0)) == pytest.approx(90.0, abs=1.0) + g = np.array(q.negative_gradient(jnp.array([1.0, -1.0]), jnp.array([0.0, 0.0]))) + assert g[0] == pytest.approx(0.9) # y > pred + assert g[1] == pytest.approx(-0.1) # y < pred (q - 1) + + +def test_huber_and_absolute_error_robust_to_outliers(): + """Robust losses recover the clean signal far better under contamination.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + rng = np.random.default_rng(1) + X = jnp.array(rng.uniform(-2, 2, size=(500, 2))) + clean = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + y = np.array(clean) + rng.normal(0, 0.1, 500) + idx = rng.choice(500, 40, replace=False) + y[idx] += 30.0 # heavy one-sided outliers + y = jnp.array(y) + + kw = { + "n_terms": 8, + "max_complexity": 4, + "learning_rate": 0.5, + "refit_coefficients": False, + } + + def mae_vs_clean(loss): + m = StagewiseSymbolicRegressor(loss=loss, **kw).fit(X, y) + return float(np.mean(np.abs(np.array(m.predict(X)) - np.array(clean)))) + + squared = mae_vs_clean("squared_error") + huber = mae_vs_clean("huber") + absolute = mae_vs_clean("absolute_error") + + assert huber < squared + assert absolute < squared + + +def test_quantile_coverage(): + """Quantile regression yields approximately calibrated coverage.""" + from jaxsr.additive import QuantileLoss, StagewiseSymbolicRegressor + + rng = np.random.default_rng(2) + X = jnp.array(rng.uniform(-2, 2, size=(500, 2))) + y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + jnp.array(rng.normal(0, 1.0, 500)) + + for target in (0.1, 0.9): + m = StagewiseSymbolicRegressor( + n_terms=10, + max_complexity=3, + learning_rate=0.5, + loss=QuantileLoss(target), + refit_coefficients=False, + ).fit(X, y) + coverage = float(np.mean(np.array(y) <= np.array(m.predict(X)))) + assert abs(coverage - target) < 0.06, f"q={target}: coverage {coverage:.3f}" + + +def test_refit_with_nonsquared_loss_warns_and_fits(): + """refit_coefficients=True with a non-squared loss warns and still fits.""" + from jaxsr.additive import StagewiseSymbolicRegressor + + X, y = _additive_data(n=200, noise=0.1, seed=3) + with pytest.warns(UserWarning, match="least squares"): + model = StagewiseSymbolicRegressor( + n_terms=4, max_complexity=3, loss="huber", refit_coefficients=True + ).fit(X, y) + assert np.all(np.isfinite(np.array(model.predict(X)))) + + +def test_save_load_preserves_parameterized_loss(tmp_path): + """save/load round-trips a quantile loss with its parameter.""" + from jaxsr.additive import QuantileLoss, StagewiseSymbolicRegressor + + X, y = _additive_data(n=150, noise=0.2, seed=5) + model = StagewiseSymbolicRegressor( + n_terms=4, max_complexity=3, loss=QuantileLoss(0.75), refit_coefficients=False + ).fit(X, y) + + path = tmp_path / "quantile_model.json" + model.save(str(path)) + loaded = StagewiseSymbolicRegressor.load(str(path)) + + assert isinstance(loaded.loss, QuantileLoss) + assert loaded.loss.quantile == pytest.approx(0.75) + assert np.allclose(np.array(model.predict(X)), np.array(loaded.predict(X)), atol=1e-6) + + +def test_scale_equivariance_and_tiny_targets(): + """Predictions scale with y, and tiny-magnitude targets still fit. + + Guards against float32 ill-conditioning in the coefficient refit: fitting + an intercept by augmenting with a ones-column fails when the term columns + are tiny (y ~ 1e-6); centering keeps the fit stable at any scale. + """ + from jaxsr.additive import StagewiseSymbolicRegressor + + rng = np.random.default_rng(17) + X = jnp.array(rng.uniform(-2, 2, size=(300, 2))) + y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + jnp.array(rng.normal(0, 0.1, 300)) + + base = StagewiseSymbolicRegressor(n_terms=5, max_complexity=4).fit(X, y) + pbase = np.array(base.predict(X)) + + scaled = StagewiseSymbolicRegressor(n_terms=5, max_complexity=4).fit(X, 3.7 * y) + rel = np.max(np.abs(np.array(scaled.predict(X)) - 3.7 * pbase)) / (np.max(np.abs(pbase)) + 1e-9) + assert rel < 1e-3, f"scale-equivariance broken: rel={rel:.2e}" + + for factor in (1e-6, 1e6): + m = StagewiseSymbolicRegressor(n_terms=5, max_complexity=4).fit(X, factor * y) + assert m.score(X, factor * y) > 0.98, f"failed at scale {factor:g}" + + +def test_refit_ols_zero_terms(): + """refit_ols with no columns returns the mean and empty coefficients.""" + from jaxsr.additive import refit_ols + + y = jnp.array([1.0, 3.0, 5.0]) + intercept, coefs = refit_ols(jnp.zeros((3, 0)), y) + assert intercept == pytest.approx(3.0) + assert coefs.shape == (0,) + + +def test_recursive_recovers_composition_flat_library_misses(): + """Recursive expansion reaches exp(x0*x1), which a flat library cannot.""" + from jaxsr import fit_symbolic + from jaxsr.additive import RecursiveSymbolicRegressor + + X = np.asarray(np.random.default_rng(0).uniform(-2, 2, size=(300, 2))) + y = np.exp(X[:, 0] * X[:, 1]) + Xtr, ytr, Xte, yte = X[:150], y[:150], X[150:], y[150:] + + rec = RecursiveSymbolicRegressor(n_expansions=3, max_terms=6, beam_width=20).fit(Xtr, ytr) + flat = fit_symbolic(jnp.array(Xtr), jnp.array(ytr), max_terms=6, include_transcendental=True) + + def r2(yt, p): + yt, p = np.asarray(yt, float), np.asarray(p, float) + return 1 - np.sum((yt - p) ** 2) / np.sum((yt - yt.mean()) ** 2) + + rec_r2 = r2(yte, rec.predict(Xte)) + assert rec_r2 > 0.95 + assert rec_r2 > r2(yte, flat.predict(jnp.array(Xte))) + 0.1 + assert np.all(np.isfinite(np.array(rec.predict(Xte)))) + + +def test_recursive_history_and_growth(): + """The library grows across rounds and history is recorded.""" + from jaxsr.additive import RecursiveSymbolicRegressor + + X, y = _additive_data(n=200, noise=0.05, seed=1) + rec = RecursiveSymbolicRegressor(n_expansions=2, max_terms=5, beam_width=15).fit(X, y) + assert rec.model_ is not None + assert rec.library_size_ >= 3 + assert len(rec.history_) >= 1 + assert rec.predict(X).shape == y.shape + assert isinstance(rec.expression_, str) + + +def test_recursive_zero_expansions_is_flat(): + """n_expansions=0 fits only the seed library (no composition).""" + from jaxsr.additive import RecursiveSymbolicRegressor + + X, y = _additive_data(n=150, noise=0.05, seed=2) + rec = RecursiveSymbolicRegressor(n_expansions=0, base_degree=2, max_terms=5).fit(X, y) + assert len(rec.history_) == 1 + assert rec.score(X, y) > 0.9 + + +def test_recursive_validation(): + """Invalid recursive parameters are rejected.""" + from jaxsr.additive import RecursiveSymbolicRegressor + + X, y = _additive_data(n=40) + with pytest.raises(ValueError): + RecursiveSymbolicRegressor(n_expansions=-1).fit(X, y) + with pytest.raises(ValueError): + RecursiveSymbolicRegressor(unary_ops=("tan",)).fit(X, y) + with pytest.raises(ValueError): + RecursiveSymbolicRegressor(binary_ops=("pow",)).fit(X, y) + with pytest.raises(RuntimeError): + RecursiveSymbolicRegressor().predict(np.zeros((3, 2))) diff --git a/tests/test_classifier.py b/tests/test_classifier.py index 4eb0112..c04da85 100644 --- a/tests/test_classifier.py +++ b/tests/test_classifier.py @@ -853,3 +853,52 @@ def test_vectorized(self): assert float(result[0]) < 0.01 assert abs(float(result[1]) - 0.5) < 1e-7 assert float(result[2]) > 0.99 + + +class TestClassifierNonFiniteBasis: + """Non-finite-on-training bases must not make predict_proba return NaN.""" + + def test_invalid_basis_excluded_binary(self): + """log/sqrt over sign-spanning data are dropped; probabilities finite.""" + rng = np.random.default_rng(0) + X_train = jnp.array(rng.uniform(-2, 2, size=(300, 2))) + X_test = jnp.array(rng.uniform(-2, 2, size=(300, 2))) + y = jnp.array( + (np.array(X_train[:, 0]) + 0.5 * np.array(X_train[:, 1]) ** 2 > 0).astype(float) + ) + + library = ( + BasisLibrary(n_features=2) + .add_constant() + .add_linear() + .add_polynomials(max_degree=3) + .add_transcendental(funcs=["log", "sqrt", "exp"]) + ) + with pytest.warns(UserWarning, match="non-finite"): + clf = SymbolicClassifier(basis_library=library, max_terms=5).fit(X_train, y) + + assert not any(("log" in n or "sqrt" in n) for n in clf.selected_features_) + proba = np.array(clf.predict_proba(X_test)) + assert np.all(np.isfinite(proba)) + assert clf.score(X_train, y) > 0.8 + + def test_invalid_basis_excluded_multiclass(self): + """Multiclass (OVR) predictions stay finite with invalid bases present.""" + rng = np.random.default_rng(1) + X_train = jnp.array(rng.uniform(-2, 2, size=(300, 2))) + X_test = jnp.array(rng.uniform(-2, 2, size=(300, 2))) + y = jnp.array(np.digitize(np.array(X_train[:, 0]), [-0.7, 0.7]).astype(float)) + + library = ( + BasisLibrary(n_features=2) + .add_constant() + .add_linear() + .add_polynomials(max_degree=3) + .add_transcendental(funcs=["log", "sqrt"]) + ) + with pytest.warns(UserWarning, match="non-finite"): + clf = SymbolicClassifier(basis_library=library, max_terms=5).fit(X_train, y) + + proba = np.array(clf.predict_proba(X_test)) + assert np.all(np.isfinite(proba)) + assert np.allclose(proba.sum(axis=1), 1.0, atol=1e-5) diff --git a/tests/test_regressor.py b/tests/test_regressor.py index 9f6c3f7..9bb0fa9 100644 --- a/tests/test_regressor.py +++ b/tests/test_regressor.py @@ -573,3 +573,89 @@ def test_cross_validate_compatible(self, multi_data, template): result = cross_validate(mo, X, Y, cv=3, scoring="neg_mse") assert "mean_test_score" in result assert len(result["test_scores"]) == 3 + + +class TestNonFiniteBasisExclusion: + """A basis that is non-finite on the training data must never be selected.""" + + def test_invalid_basis_excluded_and_predict_finite(self): + """log/sqrt over sign-spanning data are excluded; predict stays finite.""" + rng = np.random.default_rng(0) + X_train = jnp.array(rng.uniform(-2, 2, size=(300, 2))) + X_test = jnp.array(rng.uniform(-2, 2, size=(300, 2))) + y = 2 * X_train[:, 0] + 0.5 * X_train[:, 1] ** 2 + + library = ( + BasisLibrary(n_features=2) + .add_constant() + .add_linear() + .add_polynomials(max_degree=3) + .add_transcendental(funcs=["log", "sqrt", "exp"]) + ) + with pytest.warns(UserWarning, match="non-finite"): + model = SymbolicRegressor(basis_library=library, max_terms=5).fit(X_train, y) + + # No basis that is NaN on negative inputs should have been selected. + assert not any(("log" in name or "sqrt" in name) for name in model.selected_features_) + # Predictions on fresh (also sign-spanning) data are finite. + preds = np.array(model.predict(X_test)) + assert np.all(np.isfinite(preds)) + assert model.score(X_train, y) > 0.99 + + +class TestNegligibleTermPruning: + """Terms that contribute negligibly are pruned, keeping predict finite.""" + + def test_out_of_domain_pole_term_pruned(self): + """A spurious pole-bearing term (exp(x0/x1)) is pruned; predict finite.""" + rng = np.random.default_rng(0) + X_train = jnp.array(rng.uniform(-2, 2, size=(300, 2))) + X_test = jnp.array(rng.uniform(-2, 2, size=(300, 2))) + y = jnp.exp(X_train[:, 0] * X_train[:, 1]) + + library = ( + BasisLibrary(n_features=2) + .add_constant() + .add_linear() + .add_polynomials(max_degree=3) + .add_interactions() + .add_compositions(["exp"], ["product", "ratio"]) + ) + model = SymbolicRegressor(basis_library=library, max_terms=6).fit(X_train, y) + + preds = np.array(model.predict(X_test)) + assert np.all(np.isfinite(preds)) + assert "exp(x0/x1)" not in model.selected_features_ + assert "exp(x0*x1)" in model.selected_features_ + + def test_prune_disabled_keeps_terms(self): + """prune_tol=0 disables pruning (restores prior behavior).""" + rng = np.random.default_rng(0) + X = jnp.array(rng.uniform(-2, 2, size=(200, 2))) + y = jnp.exp(X[:, 0] * X[:, 1]) + library = ( + BasisLibrary(n_features=2) + .add_constant() + .add_linear() + .add_interactions() + .add_compositions(["exp"], ["product", "ratio"]) + ) + pruned = SymbolicRegressor(basis_library=library, max_terms=6, prune_tol=1e-6).fit(X, y) + kept = SymbolicRegressor(basis_library=library, max_terms=6, prune_tol=0.0).fit(X, y) + assert len(kept.selected_features_) >= len(pruned.selected_features_) + + def test_pruning_preserves_real_terms(self): + """A genuine multi-term model is not pruned.""" + rng = np.random.default_rng(1) + X = jnp.array(rng.uniform(-2, 2, size=(300, 2))) + y = 2.5 * X[:, 0] + 1.2 * X[:, 0] * X[:, 1] - 0.8 * X[:, 1] ** 2 + library = ( + BasisLibrary(n_features=2) + .add_constant() + .add_linear() + .add_polynomials(max_degree=3) + .add_interactions(max_order=2) + ) + model = SymbolicRegressor(basis_library=library, max_terms=5).fit(X, y) + assert len(model.selected_features_) >= 3 + assert model.score(X, y) > 0.99 diff --git a/tests/test_sklearn_compat.py b/tests/test_sklearn_compat.py index 8b291e6..777bc18 100644 --- a/tests/test_sklearn_compat.py +++ b/tests/test_sklearn_compat.py @@ -84,6 +84,7 @@ def test_symbolic_regressor_params(self, model): "param_optimization_budget", "constraint_enforcement", "constraint_selection_weight", + "prune_tol", } assert set(params.keys()) == expected_keys