jaxsr.regressor

Contents

jaxsr.regressor#

Main SymbolicRegressor Class for JAXSR.

Provides a scikit-learn compatible interface for symbolic regression.

class jaxsr.regressor.SymbolicRegressor(basis_library: BasisLibrary | None = None, max_terms: int = 5, strategy: str = 'greedy_forward', information_criterion: str = 'bic', cv_folds: int = 5, regularization: float | None = None, constraints: Constraints | None = None, random_state: int | None = None, param_optimizer: str = 'scipy', param_optimization_budget: int = 50, constraint_enforcement: str = 'penalty', constraint_selection_weight: float = 0.0, prune_tol: float = 1e-06)#

Bases: _SklearnCompatMixin

JAX-accelerated symbolic regression using sparse selection.

Discovers interpretable algebraic expressions from data by selecting sparse subsets of basis functions using information criteria.

Parameters:

  • basis_library (BasisLibrary, optional) – Library of candidate basis functions. If None, must be specified when calling fit() via the library_config parameter.

  • max_terms (int) – Maximum number of terms in the expression.

  • strategy (str) – Selection strategy: “greedy_forward”, “greedy_backward”, “exhaustive”, or “lasso_path”.

  • information_criterion (str) – Criterion for model selection: “aic”, “aicc”, “bic”.

  • cv_folds (int) – Number of cross-validation folds (if using CV for selection).

  • regularization (float, optional) – L2 regularization parameter (ridge penalty).

  • constraints (Constraints, optional) – Physical constraints to enforce.

  • random_state (int, optional) – Random seed for reproducibility.

  • constraint_selection_weight (float) – Weight for constraint penalty during model selection. When > 0, a penalty proportional to constraint violations is added to the 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.

expression_#

Human-readable expression after fitting.

Type:

str

coefficients_#

Fitted coefficients.

Type:

jnp.ndarray

selected_features_#

Names of selected basis functions.

Type:

list of str

complexity_#

Total complexity score of the expression.

Type:

int

metrics_#

Dictionary of evaluation metrics.

Type:

dict

pareto_front_#

Pareto-optimal models.

Type:

list of SelectionResult

Examples

>>> from jaxsr import BasisLibrary, SymbolicRegressor
>>> library = (BasisLibrary(n_features=2)
...     .add_constant()
...     .add_linear()
...     .add_polynomials(max_degree=3)
... )
>>> model = SymbolicRegressor(basis_library=library, max_terms=5)
>>> model.fit(X, y)
>>> print(model.expression_)
>>> y_pred = model.predict(X_new)
property expression_: str#

Human-readable expression.

property coefficients_: Array#

Fitted coefficients.

property selected_features_: list[str]#

Names of selected basis functions.

property selected_indices_: Array#

Indices of selected basis functions.

property complexity_: int#

Total complexity score.

property metrics_: dict[str, float]#

Evaluation metrics (weighted when the model was fitted with weights).

property sample_weight_: Array | None#

Normalised training weights, or None if the fit was unweighted.

Returns:

Weights of shape (n_samples,) scaled to sum to n_samples. These are the weights actually used, not the array passed to fit() – only the ratios between weights carry meaning.

Return type:

jnp.ndarray or None

Raises:

RuntimeError – If the model has not been fitted.

property effective_sample_size_: float#

Kish effective sample size of the training weights.

Returns:

(sum w)^2 / sum w^2, equal to n for an unweighted fit and smaller when the weights are uneven.

Return type:

float

Raises:

RuntimeError – If the model has not been fitted.

Notes

This is a diagnostic, not an input: AIC/BIC/AICc use the nominal n (see jaxsr.metrics.compute_information_criterion()). When this number is far below n, the criteria are comparing models on much less information than n suggests, and the usual asymptotic intuitions about BIC’s penalty are correspondingly weaker.

property pareto_front_: list[SelectionResult]#

Pareto-optimal models.

property selection_path_: SelectionPath#

Every candidate model the search evaluated, with the winner marked.

Returns:

Holds results (a list of SelectionResult, each carrying its own coefficients, MSE, complexity and AIC/BIC/AICc), the strategy name, and best_index – the entry chosen by the configured information_criterion.

Return type:

SelectionPath

Raises:

RuntimeError – If the model has not been fitted.

Notes

What results contains depends on the strategy. The greedy searches record one model per step, so the path reads as “best model of each size”. exhaustive records every subset it evaluated, so ranking the path by an information criterion is meaningful there. Use pareto_front_ for the accuracy/complexity trade-off curve.

The path records the search as it happened, before the post-selection steps in fit() – dropping non-finite terms, pruning negligible ones under prune_tol, and constraint refitting. So path.best can carry more terms than selected_features_; the fitted model is always a subset of it. Read the final model from coefficients_ and selected_features_, not from the path.

Examples

>>> model = SymbolicRegressor(basis_library=library).fit(X, y)  
>>> path = model.selection_path_  
>>> for i, r in enumerate(path.results):  
...     mark = "*" if i == path.best_index else " "
...     print(f"{mark} {r.n_terms} terms  BIC={r.bic:.1f}  {r.expression()}")
fit(X: Array, y: Array, sample_weight: Array | None = None) → SymbolicRegressor#

Fit symbolic regression model.

Parameters:

  • X (array-like of shape (n_samples, n_features)) – Training data.

  • y (array-like of shape (n_samples,)) – Target values.

  • sample_weight (array-like of shape (n_samples,), optional) – Per-sample weights. Observation i is treated as having variance sigma^2 / w_i, so the fit minimises sum_i w_i (y_i - f(x_i))^2. Typical sources are measurement variances (w = 1 / var), replicate counts, or a smooth taper that de-emphasises rows carrying little information.

Returns:

self – Fitted model.

Return type:

SymbolicRegressor

Raises:

ValueError – If X and y disagree on the number of samples, if basis_library is not set, if X has the wrong number of features, or if sample_weight has the wrong length or contains negative or non-finite values.

Notes

Weights are used consistently, not just for the final coefficients: term selection, the reported MSE, AIC/BIC/AICc, constraint refitting, R², the classical intervals and the bootstrap all run on the weighted problem.

Only the ratios between weights matter. They are normalised to average 1 internally, so w, 2 * w and w / 1000 all give the same model and the same information criteria, and the effective sample size in AIC/BIC stays the nominal n. Weighting therefore never manufactures observations – which is also why duplicating rows is not an equivalent trick: it inflates n and shifts every criterion. See effective_sample_size_ for a diagnostic of how much information the weights actually leave.

A weight of exactly 0 excludes a row from the fit while keeping it in n. If you mean to drop the row entirely, drop it from X and y instead.

Examples

>>> variances = np.array([0.1, 0.1, 4.0, 4.0])  
>>> model.fit(X, y, sample_weight=1.0 / variances)
predict(X: Array) → Array#

Predict using fitted model.

Parameters:

X (array-like of shape (n_samples, n_features)) – Samples to predict.

Returns:

y_pred – Predicted values.

Return type:

jnp.ndarray of shape (n_samples,)

score(X: Array, y: Array, sample_weight: Array | None = None) → float#

Compute R² score.

Parameters:

  • X (array-like) – Test samples.

  • y (array-like) – True values.

  • sample_weight (array-like of shape (n_samples,), optional) – Weights for these samples. Not inherited from fit() – X and y here are usually different data, so their weights have to come with them.

Returns:

score – (Weighted) R² score.

Return type:

float

Raises:

  • RuntimeError – If the model has not been fitted.

  • ValueError – If sample_weight is invalid.

property sigma_: float#

sqrt(SSR / (n - p)).

Returns:

sigma – Noise standard deviation estimate. Under sample weights this is the noise level of a unit-weight observation: row i is modelled with variance sigma^2 / w_i.

Return type:

float

Type:

Estimated noise standard deviation

property covariance_matrix_: Array#

s^2 * (Phi^T W Phi)^{-1}.

Returns:

cov – Covariance matrix of shape (p, p). W is the diagonal matrix of training weights (the identity for an unweighted fit).

Return type:

jnp.ndarray

Type:

Coefficient covariance matrix

coefficient_intervals(alpha: float = 0.05) → dict[str, tuple]#

Confidence intervals for all coefficients.

Parameters:

alpha (float) – Significance level (default 0.05 for 95% CIs).

Returns:

intervals – {name: (estimate, lower, upper, se)} for each coefficient. Standard errors come from the weighted covariance matrix when the model was fitted with sample_weight.

Return type:

dict

predict_interval(X: Array, alpha: float = 0.05) → tuple[Array, Array, Array]#

Prediction intervals for new observations.

Parameters:

  • X (jnp.ndarray) – New input data of shape (n_samples, n_features).

  • alpha (float) – Significance level (default 0.05 for 95% intervals).

Returns:

  • y_pred (jnp.ndarray) – Predicted values.

  • lower (jnp.ndarray) – Lower prediction interval bound.

  • upper (jnp.ndarray) – Upper prediction interval bound.

Notes

For a weighted fit the interval is for a new observation of unit weight, i.e. one measured as precisely as an average training point. Rescale by 1 / sqrt(w_new) if the new observation has a known precision of its own.

confidence_band(X: Array, alpha: float = 0.05) → tuple[Array, Array, Array]#

Confidence band on mean response E[y|x].

Parameters:

  • X (jnp.ndarray) – New input data of shape (n_samples, n_features).

  • alpha (float) – Significance level (default 0.05 for 95% band).

Returns:

  • y_pred (jnp.ndarray) – Predicted mean values.

  • lower (jnp.ndarray) – Lower confidence band bound.

  • upper (jnp.ndarray) – Upper confidence band bound.

predict_ensemble(X: Array) → dict[str, Any]#

Predictions from Pareto-front models.

Parameters:

X (jnp.ndarray) – New input data.

Returns:

result – Keys: y_mean, y_std, y_min, y_max, y_all, models.

Return type:

dict

predict_bma(X: Array, criterion: str = 'bic', alpha: float = 0.05) → tuple[Array, Array, Array]#

Bayesian Model Averaging prediction with intervals.

Parameters:

  • X (jnp.ndarray) – New input data.

  • criterion (str) – IC for weighting (“bic” or “aic”).

  • alpha (float) – Significance level.

Returns:

  • y_pred (jnp.ndarray) – BMA mean prediction.

  • lower (jnp.ndarray) – Lower bound.

  • upper (jnp.ndarray) – Upper bound.

predict_conformal(X: Array, alpha: float = 0.05, method: str = 'jackknife+', X_cal: Array | None = None, y_cal: Array | None = None) → tuple[Array, Array, Array]#

Conformal prediction intervals.

Parameters:

  • X (jnp.ndarray) – New input data.

  • alpha (float) – Significance level.

  • method (str) – “jackknife+” or “split”.

  • X_cal (jnp.ndarray, optional) – Calibration features (required for “split” method).

  • y_cal (jnp.ndarray, optional) – Calibration targets (required for “split” method).

Returns:

  • y_pred (jnp.ndarray) – Point predictions.

  • lower (jnp.ndarray) – Lower interval bound.

  • upper (jnp.ndarray) – Upper interval bound.

update(X_new: Array, y_new: Array, refit: bool = True, sample_weight_new: Array | None = None) → SymbolicRegressor#

Update model with new data points.

Parameters:

  • X_new (array-like) – New input samples.

  • y_new (array-like) – New target values.

  • refit (bool) – If True, rerun full selection. If False, only refit coefficients.

  • sample_weight_new (array-like of shape (n_new,), optional) – Weights for the new samples, on the same scale as the weights passed to fit(). Defaults to 1 per new sample. Omitting it for a weighted model therefore adds the new points at unit weight, which is what “as precise as an average existing point” means.

Returns:

self – Updated model.

Return type:

SymbolicRegressor

Raises:

  • RuntimeError – If the model has not been fitted.

  • ValueError – If sample_weight_new has the wrong length or contains negative or non-finite values.

to_sympy()#

Convert expression to SymPy for symbolic manipulation.

Returns:

expr – SymPy expression.

Return type:

sympy.Expr

to_latex() → str#

Convert expression to LaTeX string.

Returns:

latex – LaTeX representation.

Return type:

str

to_callable() → Callable[[ndarray], ndarray]#

Convert to pure Python/NumPy callable (no JAX dependency).

Returns:

func – Function that takes NumPy array and returns predictions.

Return type:

callable

save(filepath: str) → None#

Save model to JSON file.

Parameters:

filepath (str) – Path to save the model.

classmethod load(filepath: str) → SymbolicRegressor#

Load model from JSON file.

Parameters:

filepath (str) – Path to the model file.

Returns:

model – Loaded model.

Return type:

SymbolicRegressor

summary() → str#

Return a summary of the fitted model.

Returns:

summary – Model summary.

Return type:

str

class jaxsr.regressor.MultiOutputSymbolicRegressor(estimator: SymbolicRegressor, target_names: list[str] | None = None)#

Bases: _SklearnCompatMixin

Multi-output symbolic regression via per-column fitting.

Wraps a SymbolicRegressor template and clones it once per output column, discovering a separate algebraic expression for each target.

Parameters:

  • estimator (SymbolicRegressor) – Template estimator whose configuration is cloned for every output.

  • target_names (list of str, optional) – Human-readable names for each output column. If None, defaults to ["y0", "y1", ...].

estimators_#

Fitted estimators (one per output column).

Type:

list of SymbolicRegressor

expressions_#

Discovered expressions (one per output column).

Type:

list of str

n_outputs_#

Number of output columns.

Type:

int

target_names_#

Resolved target names.

Type:

list of str

Examples

>>> from jaxsr import BasisLibrary, SymbolicRegressor, MultiOutputSymbolicRegressor
>>> lib = BasisLibrary(n_features=2).add_constant().add_linear().add_polynomials(2)
>>> template = SymbolicRegressor(basis_library=lib, max_terms=4)
>>> mo = MultiOutputSymbolicRegressor(template)
>>> mo.fit(X, Y)  # Y has shape (n, m)
>>> Y_pred = mo.predict(X)  # shape (n, m)
property estimators_: list[SymbolicRegressor]#

Fitted per-output estimators.

property expressions_: list[str]#

Discovered expressions (one per output).

property n_outputs_: int#

Number of output columns.

property target_names_: list[str]#

Resolved target names.

fit(X: Array, y: Array, sample_weight: Array | None = None) → MultiOutputSymbolicRegressor#

Fit one symbolic regressor per output column.

Parameters:

  • X (array-like of shape (n_samples, n_features)) – Training data.

  • y (array-like of shape (n_samples, n_outputs)) – Target matrix. Must be 2-D; for single-output problems use SymbolicRegressor directly.

  • sample_weight (array-like of shape (n_samples,), optional) – Per-sample weights, shared by every output. Weights that differ per output are not supported here – fit separate SymbolicRegressor instances for that.

Returns:

self – Fitted model.

Return type:

MultiOutputSymbolicRegressor

Raises:

ValueError – If y is not 2-D, if shapes are inconsistent, if target_names length does not match the number of columns, or if sample_weight is invalid.

predict(X: Array) → Array#

Predict all outputs.

Parameters:

X (array-like of shape (n_samples, n_features)) – Samples to predict.

Returns:

y_pred – Predicted values.

Return type:

jnp.ndarray of shape (n_samples, n_outputs)

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

Mean R² across all outputs (sklearn convention).

Parameters:

  • X (array-like) – Test samples.

  • y (array-like of shape (n_samples, n_outputs)) – True values.

Returns:

score – Mean R² across outputs.

Return type:

float

to_sympy() → list#

Convert each output’s expression to SymPy.

Returns:

exprs – One SymPy expression per output.

Return type:

list of sympy.Expr

to_latex() → list[str]#

Convert each output’s expression to LaTeX.

Returns:

latex_strs – One LaTeX string per output.

Return type:

list of str

to_callable() → Callable[[ndarray], ndarray]#

Return a pure NumPy callable that predicts all outputs.

Returns:

func – Function (n, d) -> (n, m) using NumPy only.

Return type:

callable

summary() → str#

Return a combined summary of all per-output models.

Returns:

summary – Concatenated summaries.

Return type:

str

save(filepath: str) → None#

Save multi-output model to JSON file.

Parameters:

filepath (str) – Path to save the model.

classmethod load(filepath: str) → MultiOutputSymbolicRegressor#

Load multi-output model from JSON file.

Parameters:

filepath (str) – Path to the model file.

Returns:

model – Loaded model.

Return type:

MultiOutputSymbolicRegressor

jaxsr.regressor.fit_symbolic(X: Array, y: Array, feature_names: list[str] | None = None, max_terms: int = 5, max_poly_degree: int = 3, include_transcendental: bool = True, include_ratios: bool = False, strategy: str = 'greedy_forward', information_criterion: str = 'bic', sample_weight: Array | None = None) → SymbolicRegressor#

Convenience function for quick symbolic regression.

Parameters:

  • X (array-like) – Input features.

  • y (array-like) – Target values.

  • feature_names (list of str, optional) – Names for features.

  • max_terms (int) – Maximum terms in expression.

  • max_poly_degree (int) – Maximum polynomial degree.

  • include_transcendental (bool) – Include log, exp, sqrt, inv.

  • include_ratios (bool) – Include x_i / x_j terms.

  • strategy (str) – Selection strategy.

  • information_criterion (str) – Information criterion.

  • sample_weight (array-like of shape (n_samples,), optional) – Per-sample weights; see SymbolicRegressor.fit().

Returns:

model – Fitted model.

Return type:

SymbolicRegressor

Raises:

ValueError – If sample_weight is invalid.

Examples

>>> model = fit_symbolic(X, y, feature_names=["T", "P"])
>>> print(model.expression_)
Contents