Add additive symbolic regression (stagewise & backfitting)

[jkitchin]

Summary

Introduces a new jaxsr.additive submodule implementing additive (boosting-style) symbolic regression, where models are fit as sums of small symbolic expressions: f(x) = c + Σ_k η_k · g_k(x). This complements the existing single-expression SymbolicRegressor by enabling discovery of many interpretable terms instead of one large expression.

Key Changes

Core Implementation

• StagewiseSymbolicRegressor: Boosting-style regressor that iteratively fits small symbolic expressions to residuals and freezes them. Supports:

  • Optional OLS coefficient refitting after each stage (decouples discovery from weighting)
  • Multiple loss functions (squared error, absolute error, Huber, quantile) via gradient boosting
  • Early stopping with validation split and patience
  • Per-stage learning rates and training history tracking

• BackfittingSymbolicRegressor: GAM-style regressor that maintains a fixed set of terms and revises each in place across sweeps, conditioned on partial residuals. Warm-started from stagewise fit.

• AdditiveSymbolicModel: Plain data container holding intercept, terms, coefficients, and metadata. Shared by both regressors.

Supporting Infrastructure

• Loss functions (losses.py): Abstract Loss base class with concrete implementations for squared error, absolute error, Huber, and quantile losses. Supports both OLS refit (squared error only) and gradient boosting with line search (all losses).
• Coefficient refitting (coefficient_refit.py): OLS solver for re-estimating all linear weights over discovered symbolic features.
• Uncertainty quantification (uncertainty.py): Bootstrap resampling to assess structural stability and compute inclusion probabilities.
• Base class (base.py): Shared fitted-attribute accessors, prediction, serialization, and sklearn compatibility.

Public API

Exports: StagewiseSymbolicRegressor, BackfittingSymbolicRegressor, AdditiveSymbolicModel, Loss, SquaredError, get_loss, refit_ols, bootstrap_additive.

Testing & Documentation

• 654 tests in tests/test_additive.py covering:

  • Module imports and public API
  • Fit/predict mechanics and shape validation
  • Loss monotonicity and early stopping
  • Recovery of synthetic additive functions
  • Coefficient refitting vs. stagewise weights
  • Feature count validation and edge cases
  • Serialization and expression export

• Jupyter notebook (docs/examples/additive_symbolic_regression.ipynb) demonstrating stagewise fitting on a synthetic additive target with visualization and interpretation.

• Comprehensive guides in docs/guides/additive-symbolic-regression.md and skill documentation explaining the algorithm, key parameters, and use cases.

• API documentation (docs/api/additive.rst) with Sphinx autodoc.

Notable Implementation Details

• Residual fitting: Each stage fits to the negative gradient (pseudo-residual) of the current loss, enabling gradient boosting for non-squared losses.
• OLS refit path: When refit_coefficients=True and loss is squared error, all coefficients are re-solved globally after each term, ensuring monotone non-increasing training loss.
• Line search: For non-squared losses, each term's step size is chosen by scalar line search minimizing the loss.
• Early stopping: Validation split is stratified by random state; patience counter resets on improvement.
• Expression export: to_expression() collapses the ensemble into a single simplified SymPy expression.
• Sklearn compatibility: Full get_params/set_params protocol; works with cross_val_score, GridSearchCV, Pipeline.

Integration

• Updated README.md with additive SR feature highlight
• Updated SKILL.md and CLAUDE.md with additive SR guidance
• Updated docs/_toc.yml and docs/api/index.rst to include new guide and API docs

https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz