jaxsr.acquisition

jaxsr.acquisition#

Active Learning Acquisition Functions for JAXSR.

Provides a composable framework for design-of-experiment (DOE) and Bayesian optimization-style active learning on top of JAXSR symbolic regression models.

Because JAXSR models are linear-in-parameters (y = Phi @ beta), we have closed-form expressions for prediction variance, leverage, and coefficient covariance. This means acquisition functions that rely on posterior uncertainty are exact—no GP approximations needed.

Classes#

AcquisitionFunction: Abstract base class. Subclasses implement score().
PredictionVariance: Pure exploration: sigma^2(x) from OLS posterior.
ConfidenceBandWidth: Width of the confidence band on E[y|x].
EnsembleDisagreement: Standard deviation across Pareto-front model predictions.
BMAUncertainty: Bayesian Model Averaging uncertainty (within + between variance).
ModelDiscrimination: Maximum pairwise disagreement among top Pareto models.
UCB / LCB: Upper / Lower Confidence Bound: y_hat +/- kappa * sigma.
ExpectedImprovement: EI(x) = E[max(0, f_best - y_hat(x))] (Gaussian).
ProbabilityOfImprovement: PI(x) = P(y(x) < f_best - xi).
ThompsonSampling: Draw beta ~ N(beta_hat, Cov), score = sampled prediction.
AOptimal: Minimize tr(Cov) after adding a candidate point.
DOptimal: Maximize det(Phi^T Phi) after adding a candidate point.
Composite: Weighted combination of acquisition functions.
ActiveLearner: Orchestrator that combines candidate generation, scoring, batch selection, and the iterative update loop.

Example

>>> from jaxsr.acquisition import ActiveLearner, UCB, ExpectedImprovement
>>> learner = ActiveLearner(model, bounds=[(0, 5)], acquisition=UCB(kappa=2.0))
>>> result = learner.suggest(n_points=5)
>>> # run experiments on result.points ...
>>> learner.update(result.points, y_new)

class jaxsr.acquisition.AcquisitionResult(points: ~jax.jaxlib._jax.Array, scores: ~jax.jaxlib._jax.Array, acquisition: str, metadata: dict[str, ~typing.Any] = <factory>)#

Bases: object

Result returned by ActiveLearner.suggest().

points#

Suggested points, shape (n_points, n_features).

Type:: jnp.ndarray

scores#

Acquisition score for each suggested point.

Type:: jnp.ndarray

acquisition#

Name of the acquisition function used.

Type:: str

metadata#

Extra info (e.g. y_pred, sigma, batch diversity stats).

Type:: dict

points: Array#

scores: Array#

acquisition: str#

metadata: dict[str, Any]#

class jaxsr.acquisition.BatchStrategy(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)#

Bases: Enum

Strategies for selecting diverse batches.

GREEDY = 'greedy'#

PENALIZED = 'penalized'#

KRIGING_BELIEVER = 'kriging_believer'#

D_OPTIMAL = 'd_optimal'#

class jaxsr.acquisition.AcquisitionFunction#

Bases: ABC

Abstract base for all acquisition functions.

Subclasses must implement score(). The convention is higher score = more desirable to sample.

Supports composition via + and scalar *:

combined = 0.7 * UCB(kappa=2) + 0.3 * PredictionVariance()

abstract score(X_candidates: jnp.ndarray, model: SymbolicRegressor) → jnp.ndarray#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

property name: str#

class jaxsr.acquisition.PredictionVariance#

Bases: AcquisitionFunction

Score = prediction variance sigma^2(x).

For an OLS model, the prediction variance at a new point x is

\[\sigma^2(x) = \hat\sigma^2 \, \varphi(x)^\top (\Phi^\top\Phi)^{-1} \varphi(x)\]

where \(\hat\sigma^2\) is the unbiased noise variance estimate.

When to use: Default choice for pure exploration. Sample where the model is least certain about the mean prediction. Good for improving model accuracy uniformly across the input space.

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.ConfidenceBandWidth(alpha: float = 0.05)#

Bases: AcquisitionFunction

Score = width of the 1-alpha confidence band on E[y|x].

The width is \(2 \, t_{\alpha/2,\nu} \, \sigma(x)\).

Parameters:

alpha (float) – Significance level (default 0.05 for 95% band).
use (**When to)
target). (confidence band specifically (e.g. you have a coverage)

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.EnsembleDisagreement#

Bases: AcquisitionFunction

Score = standard deviation across Pareto-front model predictions.

This captures structural (model-form) uncertainty: where do models of different complexity disagree?

When to use: When you suspect model misspecification and want to identify regions where the choice of model complexity matters most. Helps resolve “which model is right?” rather than “how noisy is the data?”.

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.BMAUncertainty(criterion: str = 'bic', top_k: int | None = None)#

Bases: AcquisitionFunction

Score = BMA posterior standard deviation.

Combines within-model variance (noise) and between-model variance (structural uncertainty) via IC-weighted Bayesian Model Averaging.

Parameters:

criterion (str) – Information criterion for weights ("bic" or "aic").
top_k (int or None) – Number of top models to include in the average.
use (**When to)
and (Use when you want a single number that accounts for both noise)
than (model selection uncertainty. More expensive)

:param PredictionVariance because it evaluates multiple models.:

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.ModelDiscrimination(top_k: int | None = None)#

Bases: AcquisitionFunction

Score = maximum pairwise disagreement among Pareto-front models.

For each candidate x, computes the max absolute difference between any pair of Pareto-front model predictions. Points with high scores are the most informative for deciding which model structure is correct.

Parameters:

top_k (int or None) – Only compare the top_k Pareto models. None uses all.
use (**When to)
is (correct. Particularly useful in early stages when model structure)
is
uncertain. (still)

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.ModelMin#

Bases: AcquisitionFunction

Score = -y_hat(x). Finds the minimum of the surrogate.

When to use: Pure exploitation. You fully trust the model and want to sample at its predicted optimum. Combine with an exploration term (e.g. 0.8 * ModelMin() + 0.2 * PredictionVariance()) to avoid getting stuck.

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.ModelMax#

Bases: AcquisitionFunction

Score = y_hat(x). Finds the maximum of the surrogate.

When to use: Pure exploitation when maximising the response.

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.UCB(kappa: float = 2.0)#

Bases: AcquisitionFunction

Upper Confidence Bound.

\[\text{UCB}(x) = \hat y(x) + \kappa \, \sigma(x)\]

Parameters:

kappa (float) – Exploration weight. Larger values favour exploration. A common starting point is kappa=2.0.
use (**When to)
exploration-exploitation (The kappa parameter directly controls the)
balance.
0 (- kappa =)
2 (- kappa ~)
3 (- kappa >)

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.LCB(kappa: float = 2.0)#

Bases: AcquisitionFunction

Lower Confidence Bound.

\[\text{LCB}(x) = -\hat y(x) + \kappa \, \sigma(x)\]

Score convention: higher is more desirable, so the sign of the predicted mean is flipped. Minimising y corresponds to maximising -y_hat + kappa * sigma.

Parameters:

kappa (float) – Exploration weight (same semantics as UCB).
use (**When to)
UCB. (Mirror image of)

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.ExpectedImprovement(y_best: float | None = None, xi: float = 0.01, minimize: bool = True)#

Bases: AcquisitionFunction

Expected Improvement over the current best.

\[\text{EI}(x) = (\hat y_\text{best} - \hat y(x))\, \Phi(z) + \sigma(x)\,\phi(z), \quad z = \frac{\hat y_\text{best} - \hat y(x) - \xi}{\sigma(x)}\]

for minimisation (flip signs for maximisation).

Parameters:

y_best (float or None) – Current best observed value. If None, the best value from the training set is used automatically.
xi (float) – Jitter for encouraging exploration (default 0.01).
minimize (bool) – If True (default), seek points that reduce y. If False, seek points that increase y.
use (**When to)
(predicted (balances exploration (high sigma) and exploitation)
UCB. (improvement over best). Less sensitive to hyperparameters than)

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.ProbabilityOfImprovement(y_best: float | None = None, xi: float = 0.01, minimize: bool = True)#

Bases: AcquisitionFunction

Probability of improving over the current best.

\[\text{PI}(x) = \Phi\!\left( \frac{\hat y_\text{best} - \hat y(x) - \xi}{\sigma(x)} \right)\]

Parameters:

y_best (float or None) – Current best observed value. None = auto from training data.
xi (improvement. Tends to be more exploitative than EI for the same) – Improvement threshold (default 0.01).
minimize (bool) – If True, seek reduction in y.
use (**When to)
of (*probability* of beating a threshold rather than the magnitude)
xi
magnitude. (because it ignores improvement)

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.ThompsonSampling(minimize: bool = True, seed: int | None = None)#

Bases: AcquisitionFunction

Thompson Sampling via posterior coefficient draw.

Draws \(\beta \sim \mathcal{N}(\hat\beta, \text{Cov}(\hat\beta))\) and scores each candidate with the sampled model.

Parameters:

minimize (bool) – If True, score = -y_sampled (prefer smaller predicted y).
seed (int or None) – Random seed for the posterior draw.
use (**When to)
posterior (Each call to score() draws a different model from the)

:param : :param so repeated calls produce diverse batches without explicit: :param diversification. Theoretically elegant; matches Bayes-optimal: :param one-step-ahead strategy.:

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.AOptimal#

Bases: AcquisitionFunction

A-Optimal design: reduce average parameter uncertainty.

Score for adding candidate x is the reduction in \(\text{tr}(\text{Cov}(\hat\beta))\).

\[\Delta\text{tr} = \frac{ \varphi(x)^\top (\Phi^\top\Phi)^{-1} \text{Cov}(\hat\beta) (\Phi^\top\Phi)^{-1} \varphi(x) }{1 + \varphi(x)^\top (\Phi^\top\Phi)^{-1} \varphi(x)}\]

When to use: When your goal is to tighten the confidence intervals on all coefficients. Good for model building where every parameter matters (e.g. reporting a physics equation).

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.DOptimal#

Bases: AcquisitionFunction

D-Optimal design: maximise information gain.

Score = leverage \(h(x) = \varphi(x)^\top (\Phi^\top\Phi)^{-1} \varphi(x)\), which is proportional to the increase in \(\det(\Phi^\top\Phi)\) from adding the point.

When to use: When you want maximum information per experiment. Classic choice for building precise models with minimal data. Tends to place points at the “edges” of the design space.

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.Composite(functions: list[tuple[float, AcquisitionFunction]])#

Bases: AcquisitionFunction

Weighted combination of acquisition functions.

Typically created via operator overloading rather than directly:

acq = 0.7 * UCB(kappa=2) + 0.3 * PredictionVariance()

Each component is scored independently and the weighted sum is returned after min-max normalisation of each component.

functions#

Type:: list of (weight, AcquisitionFunction)

property name: str#

score(X_candidates, model)#

Score each candidate (higher = more desirable).

Parameters:

X_candidates (jnp.ndarray, shape (n_candidates, n_features)) – Candidate points to evaluate.
model (SymbolicRegressor) – Fitted symbolic regression model.

Returns:

scores

Return type:

jnp.ndarray, shape (n_candidates,)

class jaxsr.acquisition.ActiveLearner(model: SymbolicRegressor, bounds: list[tuple[float, float]], acquisition: AcquisitionFunction, n_candidates: int = 1000, candidate_method: str = 'lhs', random_state: int | None = None)#

Bases: object

Orchestrator for iterative active learning / DOE.

Combines candidate generation, acquisition scoring, batch selection, and model updating into a single workflow.

Parameters:

model (SymbolicRegressor) – A fitted model to improve.
bounds (list of (float, float)) – Input bounds [(lo, hi), ...] for each feature.
acquisition (AcquisitionFunction) – The acquisition function (or composite) to use for scoring.
n_candidates (int) – Number of space-filling candidates to generate internally.
candidate_method (str) – Candidate generation method: "lhs" (default), "sobol", "halton", or "random".
random_state (int or None) – Seed for reproducibility.

history_X#

Points suggested at each iteration.

Type:: list of jnp.ndarray

history_y#

Observations received at each iteration.

Type:: list of jnp.ndarray

iteration#

Number of suggest-update cycles completed.

Type:: int

Examples

>>> learner = ActiveLearner(model, bounds=[(0, 5)],
...     acquisition=UCB(kappa=2))
>>> for _ in range(10):
...     result = learner.suggest(n_points=3)
...     y_new = oracle(result.points)
...     learner.update(result.points, y_new)

suggest(n_points: int = 5, batch_strategy: str = 'greedy', min_distance: float = 0.01, candidates: Array | None = None) → AcquisitionResult#

Suggest the next batch of points to evaluate.

Parameters:

n_points (int) – Number of points to return.
batch_strategy (str) – How to select a batch: "greedy" (top-k by score), "penalized" (top-1, penalise nearby, repeat), "kriging_believer" (top-1, fantasise observation, repeat), "d_optimal" (select batch maximising det(Phi^T Phi)).
min_distance (float) – Minimum normalised distance from training data.
candidates (jnp.ndarray, optional) – Pre-specified candidate pool. If None, candidates are generated internally.

Returns:

With .points, .scores, and .metadata.

Return type:

AcquisitionResult

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

Add new observations and refit the model.

Parameters:

X_new (jnp.ndarray) – New input points.
y_new (jnp.ndarray) – Observed responses.
refit (bool) – If True, rerun full feature selection. If False, keep the same basis terms and refit coefficients.

converged(tol: float = 0.001, window: int = 3, metric: str = 'mse') → bool#

Check whether model improvement has stalled.

Parameters:

tol (float) – Minimum relative improvement to consider “not converged”.
window (int) – Number of recent iterations to compare.
metric (str) – Metric to track: "mse" or "r2".

Returns:

True if improvement over the last window iterations is below tol.

Return type:

bool

property best_y: float#: Best observed y in the training set.

property best_X: Array#: Input corresponding to the best observed y.

property n_observations: int#: Total number of observations in the training set.

jaxsr.acquisition.suggest_points(model: SymbolicRegressor, bounds: list[tuple[float, float]], acquisition: AcquisitionFunction, n_points: int = 5, batch_strategy: str = 'greedy', n_candidates: int = 1000, random_state: int | None = None) → AcquisitionResult#

One-shot convenience: suggest points without creating an ActiveLearner.

Parameters:

model (SymbolicRegressor) – Fitted model.
bounds (list of (float, float)) – Feature bounds.
acquisition (AcquisitionFunction) – Acquisition function to use.
n_points (int) – Number of points.
batch_strategy (str) – Batch strategy.
n_candidates (int) – Candidate pool size.
random_state (int or None) – Seed.

Return type:

AcquisitionResult

jaxsr.acquisition

Contents

jaxsr.acquisition#

Classes#