jaxsr.additive

Contents

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).

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.

class jaxsr.additive.AbsoluteError#

Bases: 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: str = 'absolute_error'#

Human-readable name used in the loss registry.

initial_prediction(y: Array) → float#

Return median(y).

negative_gradient(y: Array, y_pred: Array) → Array#

Return sign(y - y_pred).

loss(y: Array, y_pred: Array) → float#

Return the mean absolute error.

class jaxsr.additive.AdditiveSymbolicModel(intercept: float, terms: list[SymbolicRegressor], coefficients: list[float], learning_rates: list[float], feature_names: list[str], training_history: list[dict[str, Any]] = <factory>)#

Bases: object

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]]#
property n_terms: int#

Number of symbolic terms in the model.

predict(X: Array) → Array#

Predict with the additive model.

Parameters:

X (jnp.ndarray of shape (n_samples, n_features)) – Input data.

Returns:

Predicted values.

Return type:

jnp.ndarray of shape (n_samples,)

Raises:

ValueError – If X has a different number of features than the model was fit with.

property expressions: list[str]#

Human-readable expression string for each symbolic term.

to_expression()#

Combine all terms into a single simplified SymPy expression.

Returns:

intercept + sum_j coefficients[j] * g_j after simplification. Falls back to the unsimplified sum if simplification fails.

Return type:

sympy.Expr

Raises:

ImportError – If SymPy is not installed.

describe(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:

Pretty-printed model structure.

Return type:

str

class jaxsr.additive.BackfittingSymbolicRegressor(n_terms: int = 5, n_sweeps: int = 10, max_complexity: int = 4, loss: str = 'squared_error', tol: float = 1e-06, 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)#

Bases: _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 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.

model_#

The fitted additive model.

Type:

AdditiveSymbolicModel

intercept_#

Fitted intercept.

Type:

float

coefficients_#

Fitted per-term coefficients.

Type:

list of float

expressions_#

Human-readable expression for each term.

Type:

list of str

terms_#

The fitted symbolic terms.

Type:

list of SymbolicRegressor

training_history_#

Per-sweep diagnostics (sweep index and train_loss).

Type:

list of dict

n_terms_#

Number of terms in the fitted model.

Type:

int

Examples

>>> from jaxsr.additive import BackfittingSymbolicRegressor
>>> model = BackfittingSymbolicRegressor(n_terms=3, n_sweeps=5)
>>> model.fit(X, y)  
>>> print(model)
fit(X: Array, y: Array) → 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 – The fitted estimator.

Return type:

BackfittingSymbolicRegressor

Raises:

  • ValueError – If parameters are invalid, X/y are mismatched or non-finite.

  • NotImplementedError – If loss is not "squared_error".

class jaxsr.additive.HuberLoss(delta: float = 1.35)#

Bases: 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: str = 'huber'#

Human-readable name used in the loss registry.

initial_prediction(y: Array) → float#

Return median(y) (a robust location estimate).

negative_gradient(y: Array, y_pred: Array) → Array#

Return r where |r| <= delta else delta * sign(r).

loss(y: Array, y_pred: Array) → float#

Return the mean Huber loss.

to_config() → dict[str, Any]#

Return {"name": "huber", "params": {"delta": delta}}.

class jaxsr.additive.Loss#

Bases: ABC

Abstract base class for additive-regression loss functions.

A concrete loss defines three things:

  • initial_prediction() – the constant that minimises the loss and is used to initialise the ensemble intercept.

  • negative_gradient() – the pseudo-residual each weak learner fits.

  • loss() – the scalar training/validation loss for reporting and early stopping.

name: str = 'loss'#

Human-readable name used in the loss registry.

abstract initial_prediction(y: Array) → float#

Return the optimal constant prediction for y.

Parameters:

y (jnp.ndarray of shape (n_samples,)) – Target values.

Returns:

Constant used to initialise the additive model intercept.

Return type:

float

abstract negative_gradient(y: Array, y_pred: Array) → Array#

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:

Pseudo-residuals.

Return type:

jnp.ndarray of shape (n_samples,)

abstract loss(y: Array, y_pred: Array) → 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:

Scalar loss value.

Return type:

float

to_config() → dict[str, Any]#

Return a JSON-serialisable {"name", "params"} description.

Returns:

name is the registry key; params are the constructor keyword arguments needed to rebuild this loss.

Return type:

dict

class jaxsr.additive.QuantileLoss(quantile: float = 0.5)#

Bases: 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: str = 'quantile'#

Human-readable name used in the loss registry.

initial_prediction(y: Array) → float#

Return the empirical quantile-th quantile of y.

negative_gradient(y: Array, y_pred: Array) → Array#

Return q where y > y_pred else q - 1.

loss(y: Array, y_pred: Array) → float#

Return the mean pinball loss.

to_config() → dict[str, Any]#

Return {"name": "quantile", "params": {"quantile": quantile}}.

class jaxsr.additive.RecursiveSymbolicRegressor(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)#

Bases: _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 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).

model_#

The fitted sparse model over the final grown library.

Type:

SymbolicRegressor

library_size_#

Number of candidate basis functions in the final library.

Type:

int

history_#

Per-round diagnostics (library size, number of terms, train R^2).

Type:

list of dict

Examples

>>> from jaxsr.additive import RecursiveSymbolicRegressor
>>> model = RecursiveSymbolicRegressor(n_expansions=3)
>>> model.fit(X, y)  
>>> print(model.expression_)
fit(X: Array, y: Array) → 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 – The fitted estimator.

Return type:

RecursiveSymbolicRegressor

Raises:

ValueError – If parameters are invalid or X/y are mismatched.

predict(X: Array) → Array#

Predict with the fitted model.

score(X: Array, y: Array) → float#

Return the R^2 score of the fitted model on (X, y).

property expression_: str#

Human-readable expression of the fitted model.

property selected_features_: list[str]#

Names of the selected (possibly composed) basis functions.

class jaxsr.additive.SquaredError#

Bases: 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: str = 'squared_error'#

Human-readable name used in the loss registry.

initial_prediction(y: Array) → float#

Return mean(y).

negative_gradient(y: Array, y_pred: Array) → Array#

Return the residual y - y_pred.

loss(y: Array, y_pred: Array) → float#

Return the mean squared error.

class jaxsr.additive.StagewiseSymbolicRegressor(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-08, 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)#

Bases: _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 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 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 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.

model_#

The fitted additive model.

Type:

AdditiveSymbolicModel

intercept_#

Fitted intercept.

Type:

float

coefficients_#

Fitted per-term coefficients.

Type:

list of float

expressions_#

Human-readable expression for each term.

Type:

list of str

terms_#

The fitted symbolic terms.

Type:

list of SymbolicRegressor

learning_rates_#

Learning rate recorded at each stage.

Type:

list of float

training_history_#

Per-stage diagnostics.

Type:

list of dict

n_terms_#

Number of terms in the fitted model.

Type:

int

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)  
>>> print(model)
fit(X: Array, y: Array) → 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 – The fitted estimator.

Return type:

StagewiseSymbolicRegressor

Raises:

ValueError – If parameters are invalid, X and y have mismatched sample counts, or X/y contain non-finite values.

jaxsr.additive.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:

Predicted values.

Return type:

jnp.ndarray of shape (n_samples,)

jaxsr.additive.bootstrap_additive(estimator, X: Array, y: Array, 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:

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 bootstrap_predict_additive()).

"n_bootstrap"int

Number of resamples that fit successfully.

Return type:

dict

Raises:

  • ValueError – If X/y are mismatched or n_bootstrap < 1.

  • RuntimeError – If every bootstrap fit fails.

jaxsr.additive.bootstrap_predict_additive(models: list, X: Array, 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 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:

Plain dictionary with keys "mean", "std", "median", "lower", "upper" (each of shape (n_samples,)) and "predictions" of shape (n_models, n_samples).

Return type:

dict

Raises:

ValueError – If models is empty or alpha is not in (0, 1).

jaxsr.additive.get_loss(loss: str | Loss) → Loss#

Resolve a loss name (or instance) to a Loss instance.

Parameters:

loss (str or Loss) – Either a registered loss name ("squared_error", "absolute_error", "huber", "quantile") or an already-constructed Loss instance. Names build losses with their default parameters; pass an instance (e.g. QuantileLoss(0.9)) to customise.

Returns:

A loss instance.

Return type:

Loss

Raises:

  • ValueError – If loss is a string that is not a registered loss name.

  • TypeError – If loss is neither a string nor a Loss instance.

jaxsr.additive.loss_from_config(config: str | dict[str, Any] | Loss) → Loss#

Rebuild a 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 Loss.to_config(), or a Loss instance.

Returns:

A loss instance.

Return type:

Loss

Raises:

  • ValueError – If the config names an unknown loss.

  • TypeError – If config is not a str, dict, or Loss.

jaxsr.additive.refit_ols(Phi: Array, y: Array) → tuple[float, Array]#

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).

Stagewise (boosting-style) additive symbolic regression.

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 BackfittingSymbolicRegressor will instead revise terms in place.

class jaxsr.additive.stagewise.StagewiseSymbolicRegressor(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-08, 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)#

Bases: _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 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 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 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.

model_#

The fitted additive model.

Type:

AdditiveSymbolicModel

intercept_#

Fitted intercept.

Type:

float

coefficients_#

Fitted per-term coefficients.

Type:

list of float

expressions_#

Human-readable expression for each term.

Type:

list of str

terms_#

The fitted symbolic terms.

Type:

list of SymbolicRegressor

learning_rates_#

Learning rate recorded at each stage.

Type:

list of float

training_history_#

Per-stage diagnostics.

Type:

list of dict

n_terms_#

Number of terms in the fitted model.

Type:

int

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)  
>>> print(model)
fit(X: Array, y: Array) → 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 – The fitted estimator.

Return type:

StagewiseSymbolicRegressor

Raises:

ValueError – If parameters are invalid, X and y have mismatched sample counts, or X/y contain non-finite values.

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 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 AdditiveSymbolicModel is a plain data container: the fitting strategy (stagewise, backfitting, …) lives in the regressor classes and produces one of these objects.

jaxsr.additive.ensemble.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:

Predicted values.

Return type:

jnp.ndarray of shape (n_samples,)

class jaxsr.additive.ensemble.AdditiveSymbolicModel(intercept: float, terms: list[SymbolicRegressor], coefficients: list[float], learning_rates: list[float], feature_names: list[str], training_history: list[dict[str, Any]] = <factory>)#

Bases: object

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]]#
property n_terms: int#

Number of symbolic terms in the model.

predict(X: Array) → Array#

Predict with the additive model.

Parameters:

X (jnp.ndarray of shape (n_samples, n_features)) – Input data.

Returns:

Predicted values.

Return type:

jnp.ndarray of shape (n_samples,)

Raises:

ValueError – If X has a different number of features than the model was fit with.

property expressions: list[str]#

Human-readable expression string for each symbolic term.

to_expression()#

Combine all terms into a single simplified SymPy expression.

Returns:

intercept + sum_j coefficients[j] * g_j after simplification. Falls back to the unsimplified sum if simplification fails.

Return type:

sympy.Expr

Raises:

ImportError – If SymPy is not installed.

describe(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:

Pretty-printed model structure.

Return type:

str

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:

New losses can be added by subclassing Loss and registering them in _LOSSES.

class jaxsr.additive.losses.Loss#

Bases: ABC

Abstract base class for additive-regression loss functions.

A concrete loss defines three things:

  • initial_prediction() – the constant that minimises the loss and is used to initialise the ensemble intercept.

  • negative_gradient() – the pseudo-residual each weak learner fits.

  • loss() – the scalar training/validation loss for reporting and early stopping.

name: str = 'loss'#

Human-readable name used in the loss registry.

abstract initial_prediction(y: Array) → float#

Return the optimal constant prediction for y.

Parameters:

y (jnp.ndarray of shape (n_samples,)) – Target values.

Returns:

Constant used to initialise the additive model intercept.

Return type:

float

abstract negative_gradient(y: Array, y_pred: Array) → Array#

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:

Pseudo-residuals.

Return type:

jnp.ndarray of shape (n_samples,)

abstract loss(y: Array, y_pred: Array) → 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:

Scalar loss value.

Return type:

float

to_config() → dict[str, Any]#

Return a JSON-serialisable {"name", "params"} description.

Returns:

name is the registry key; params are the constructor keyword arguments needed to rebuild this loss.

Return type:

dict

class jaxsr.additive.losses.SquaredError#

Bases: 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: str = 'squared_error'#

Human-readable name used in the loss registry.

initial_prediction(y: Array) → float#

Return mean(y).

negative_gradient(y: Array, y_pred: Array) → Array#

Return the residual y - y_pred.

loss(y: Array, y_pred: Array) → float#

Return the mean squared error.

class jaxsr.additive.losses.AbsoluteError#

Bases: 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: str = 'absolute_error'#

Human-readable name used in the loss registry.

initial_prediction(y: Array) → float#

Return median(y).

negative_gradient(y: Array, y_pred: Array) → Array#

Return sign(y - y_pred).

loss(y: Array, y_pred: Array) → float#

Return the mean absolute error.

class jaxsr.additive.losses.HuberLoss(delta: float = 1.35)#

Bases: 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: str = 'huber'#

Human-readable name used in the loss registry.

initial_prediction(y: Array) → float#

Return median(y) (a robust location estimate).

negative_gradient(y: Array, y_pred: Array) → Array#

Return r where |r| <= delta else delta * sign(r).

loss(y: Array, y_pred: Array) → float#

Return the mean Huber loss.

to_config() → dict[str, Any]#

Return {"name": "huber", "params": {"delta": delta}}.

class jaxsr.additive.losses.QuantileLoss(quantile: float = 0.5)#

Bases: 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: str = 'quantile'#

Human-readable name used in the loss registry.

initial_prediction(y: Array) → float#

Return the empirical quantile-th quantile of y.

negative_gradient(y: Array, y_pred: Array) → Array#

Return q where y > y_pred else q - 1.

loss(y: Array, y_pred: Array) → float#

Return the mean pinball loss.

to_config() → dict[str, Any]#

Return {"name": "quantile", "params": {"quantile": quantile}}.

jaxsr.additive.losses.get_loss(loss: str | Loss) → Loss#

Resolve a loss name (or instance) to a Loss instance.

Parameters:

loss (str or Loss) – Either a registered loss name ("squared_error", "absolute_error", "huber", "quantile") or an already-constructed Loss instance. Names build losses with their default parameters; pass an instance (e.g. QuantileLoss(0.9)) to customise.

Returns:

A loss instance.

Return type:

Loss

Raises:

  • ValueError – If loss is a string that is not a registered loss name.

  • TypeError – If loss is neither a string nor a Loss instance.

jaxsr.additive.losses.loss_from_config(config: str | dict[str, Any] | Loss) → Loss#

Rebuild a 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 Loss.to_config(), or a Loss instance.

Returns:

A loss instance.

Return type:

Loss

Raises:

  • ValueError – If the config names an unknown loss.

  • TypeError – If config is not a str, dict, or Loss.

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.

jaxsr.additive.coefficient_refit.refit_ols(Phi: Array, y: Array) → tuple[float, Array]#

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).

Backfitting additive symbolic regression (GAM-style).

Unlike 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.

class jaxsr.additive.backfitting.BackfittingSymbolicRegressor(n_terms: int = 5, n_sweeps: int = 10, max_complexity: int = 4, loss: str = 'squared_error', tol: float = 1e-06, 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)#

Bases: _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 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.

model_#

The fitted additive model.

Type:

AdditiveSymbolicModel

intercept_#

Fitted intercept.

Type:

float

coefficients_#

Fitted per-term coefficients.

Type:

list of float

expressions_#

Human-readable expression for each term.

Type:

list of str

terms_#

The fitted symbolic terms.

Type:

list of SymbolicRegressor

training_history_#

Per-sweep diagnostics (sweep index and train_loss).

Type:

list of dict

n_terms_#

Number of terms in the fitted model.

Type:

int

Examples

>>> from jaxsr.additive import BackfittingSymbolicRegressor
>>> model = BackfittingSymbolicRegressor(n_terms=3, n_sweeps=5)
>>> model.fit(X, y)  
>>> print(model)
fit(X: Array, y: Array) → 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 – The fitted estimator.

Return type:

BackfittingSymbolicRegressor

Raises:

  • ValueError – If parameters are invalid, X/y are mismatched or non-finite.

  • NotImplementedError – If loss is not "squared_error".

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 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.

class jaxsr.additive.recursive.RecursiveSymbolicRegressor(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)#

Bases: _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 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).

model_#

The fitted sparse model over the final grown library.

Type:

SymbolicRegressor

library_size_#

Number of candidate basis functions in the final library.

Type:

int

history_#

Per-round diagnostics (library size, number of terms, train R^2).

Type:

list of dict

Examples

>>> from jaxsr.additive import RecursiveSymbolicRegressor
>>> model = RecursiveSymbolicRegressor(n_expansions=3)
>>> model.fit(X, y)  
>>> print(model.expression_)
fit(X: Array, y: Array) → 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 – The fitted estimator.

Return type:

RecursiveSymbolicRegressor

Raises:

ValueError – If parameters are invalid or X/y are mismatched.

predict(X: Array) → Array#

Predict with the fitted model.

score(X: Array, y: Array) → float#

Return the R^2 score of the fitted model on (X, y).

property expression_: str#

Human-readable expression of the fitted model.

property selected_features_: list[str]#

Names of the selected (possibly composed) basis functions.

Shared base class for additive symbolic regressors.

Both StagewiseSymbolicRegressor and BackfittingSymbolicRegressor produce an AdditiveSymbolicModel and share the same fitted-attribute accessors, prediction, interpretation, and JSON serialization. Only the fitting strategy differs, so everything except fit lives here.

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.

jaxsr.additive.uncertainty.bootstrap_additive(estimator, X: Array, y: Array, 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:

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 bootstrap_predict_additive()).

"n_bootstrap"int

Number of resamples that fit successfully.

Return type:

dict

Raises:

  • ValueError – If X/y are mismatched or n_bootstrap < 1.

  • RuntimeError – If every bootstrap fit fails.

jaxsr.additive.uncertainty.bootstrap_predict_additive(models: list, X: Array, 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 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:

Plain dictionary with keys "mean", "std", "median", "lower", "upper" (each of shape (n_samples,)) and "predictions" of shape (n_models, n_samples).

Return type:

dict

Raises:

ValueError – If models is empty or alpha is not in (0, 1).

Contents