jaxsr.metrics

jaxsr.metrics#

Metrics and Model Comparison for JAXSR.

Provides information criteria, cross-validation scores, and model comparison utilities.

jaxsr.metrics.compute_aic(n_samples: int, n_params: int, mse: float, variance: float | None = None) → float#

Compute Akaike Information Criterion.

AIC = n * log(MSE) + 2 * k

Parameters:

n_samples (int) – Number of samples.
n_params (int) – Number of model parameters.
mse (float) – Mean squared error.
variance (float, optional) – Known error variance (if None, estimated from MSE).

Returns:

aic – AIC value (lower is better).

Return type:

float

jaxsr.metrics.compute_bic(n_samples: int, n_params: int, mse: float, variance: float | None = None) → float#

Compute Bayesian Information Criterion.

BIC = n * log(MSE) + k * log(n)

Parameters:

n_samples (int) – Number of samples.
n_params (int) – Number of model parameters.
mse (float) – Mean squared error.
variance (float, optional) – Known error variance (if None, estimated from MSE).

Returns:

bic – BIC value (lower is better).

Return type:

float

jaxsr.metrics.compute_aicc(n_samples: int, n_params: int, mse: float, variance: float | None = None) → float#

Compute corrected Akaike Information Criterion (AICc).

AICc = AIC + 2*k*(k+1) / (n-k-1)

Parameters:

n_samples (int) – Number of samples.
n_params (int) – Number of model parameters.
mse (float) – Mean squared error.
variance (float, optional) – Known error variance (if None, estimated from MSE).

Returns:

aicc – AICc value (lower is better).

Return type:

float

Notes

AICc includes a correction for small sample sizes. It should be preferred when n/k < 40.

jaxsr.metrics.compute_hqc(n_samples: int, n_params: int, mse: float) → float#

Compute Hannan-Quinn Criterion.

HQC = n * log(MSE) + 2 * k * log(log(n))

Parameters:

n_samples (int) – Number of samples.
n_params (int) – Number of model parameters.
mse (float) – Mean squared error.

Returns:

hqc – HQC value (lower is better).

Return type:

float

Notes

HQC is an alternative to BIC that penalizes complexity less severely.

jaxsr.metrics.compute_mdl(n_samples: int, n_params: int, mse: float) → float#

Compute Minimum Description Length criterion.

MDL = n/2 * log(MSE) + k/2 * log(n)

Parameters:

n_samples (int) – Number of samples.
n_params (int) – Number of model parameters.
mse (float) – Mean squared error.

Returns:

mdl – MDL value (lower is better).

Return type:

float

jaxsr.metrics.compute_information_criterion(n_samples: int, n_params: int, mse: float, criterion: str = 'bic') → float#

Compute the specified information criterion.

Parameters:

n_samples (int) – Number of samples.
n_params (int) – Number of model parameters.
mse (float) – Mean squared error.
criterion (str) – One of “aic”, “aicc”, “bic”, “hqc”, “mdl”.

Returns:

ic – Information criterion value (lower is better).

Return type:

float

Notes

Under sample weights, pass the nominal number of observations as n_samples and the weighted MSE sum_i w_i r_i^2 / n as mse. See the effective-sample-size note at the top of this module.

jaxsr.metrics.group_indices(groups: Any) → tuple[ndarray, list[ndarray]]#

Split row indices by group label.

Parameters:

groups (array-like of shape (n_samples,)) – Group label for each row. Labels may be integers, floats, or strings.

Returns:

unique (np.ndarray) – Unique group labels, sorted.
row_indices (list of np.ndarray) – row_indices[i] holds the row indices belonging to unique[i].

Raises:

ValueError – If groups is empty or not one-dimensional after raveling.

jaxsr.metrics.cross_validate(model: SymbolicRegressor, X: jnp.ndarray, y: jnp.ndarray, cv: int = 5, scoring: str = 'neg_mse', random_state: int | None = None, groups: Any | None = None, strategy: str = 'kfold', sample_weight: jnp.ndarray | None = None) → dict[str, Any]#

Perform cross-validation, optionally holding out whole groups.

When rows are not independent observations – replicates of one experimental condition, points along one measured curve, samples from one subject – splitting at the row level leaks information from the training set into the test set and reports an optimistic score. Passing groups keeps every row of a group on the same side of the split.

Parameters:

model (SymbolicRegressor) – Model to evaluate. Cloned (unfitted) for every fold.
X (jnp.ndarray) – Feature matrix of shape (n_samples, n_features).
y (jnp.ndarray) – Target values of shape (n_samples,) or (n_samples, n_outputs).
cv (int) – Number of folds. Ignored when strategy="leave-one-group-out".
scoring (str) – Scoring metric: “neg_mse”, “neg_mae”, “r2”.
random_state (int, optional) – Random seed for fold splitting (used by strategy="kfold" only; the group strategies are deterministic).
groups (array-like of shape (n_samples,), optional) – Group label for each row. Required by the group strategies. If given with the default strategy="kfold", the strategy is promoted to "group-kfold".
strategy (str) – One of "kfold" (random rows), "group-kfold" (whole groups distributed over cv folds), or "leave-one-group-out" (one fold per group).
sample_weight (jnp.ndarray, optional) – Per-sample weights. Weights follow their rows into the folds under every strategy: each fold is fitted on its training weights and scored with its test weights, so down-weighted observations neither drive the fit nor dominate the score. Per-group scores are weighted the same way.

Returns:

results – Dictionary with keys:

”test_scores”: array of test scores for each fold
”train_scores”: array of train scores for each fold
”mean_test_score”, “std_test_score”: summary of the test scores
”mean_train_score”, “std_train_score”: summary of the train scores
”strategy”: the strategy actually used
”n_splits”: number of folds evaluated
”groups_out”: list of the held-out group labels per fold (empty lists for "kfold")
”per_group_scores”: dict mapping group label to the score on that group’s rows while it was held out (empty for "kfold")
”edge_groups”: the lowest and highest group labels when labels are numeric, else an empty list. These are the extrapolation cases – read their entries in “per_group_scores” separately from the mean, since interpolation and extrapolation carry different risk.
”edge_group_scores”: “per_group_scores” restricted to “edge_groups”

Return type:

dict

Raises:

ValueError – If scoring or strategy is unknown, if cv < 2 for a k-fold strategy, if groups is missing for a group strategy or has the wrong length, if there are fewer groups than folds, if sample_weight is invalid, or if a fold ends up with zero total weight on either side of its split.

Examples

>>> result = cross_validate(model, X, y, cv=5)  
>>> result = cross_validate(  
...     model, X, y, groups=temperature_id, strategy="leave-one-group-out"
... )

jaxsr.metrics.compute_cv_score(Phi: Array, y: Array, cv: int = 5, random_state: int | None = None, sample_weight: Array | None = None) → float#

Compute cross-validation MSE for a design matrix.

This is a lower-level function that works directly with the design matrix.

Parameters:

Phi (jnp.ndarray) – Design matrix of shape (n_samples, n_features).
y (jnp.ndarray) – Target values.
cv (int) – Number of folds.
random_state (int, optional) – Random seed.
sample_weight (jnp.ndarray, optional) – Per-sample weights. Each fold is fitted by weighted least squares on its training weights and scored by weighted MSE on its test weights.

Returns:

cv_mse – Mean cross-validation MSE.

Return type:

float

Raises:

ValueError – If sample_weight is invalid.

jaxsr.metrics.compute_loo_mse(Phi: Array, y: Array, coefficients: Array, sample_weight: Array | None = None) → float#

Compute leave-one-out MSE efficiently using Sherman-Morrison formula.

This avoids refitting the model n times by using the hat matrix.

Parameters:

Phi (jnp.ndarray) – Design matrix of shape (n_samples, n_features).
y (jnp.ndarray) – Target values.
coefficients (jnp.ndarray) – Fitted coefficients.
sample_weight (jnp.ndarray, optional) – Per-sample weights. The leverages become the weighted-least-squares leverages w_i * phi_i^T (Phi^T W Phi)^-1 phi_i and the LOO residuals are averaged with the same weights.

Returns:

loo_mse – Leave-one-out mean squared error.

Return type:

float

Raises:

ValueError – If sample_weight is invalid.

jaxsr.metrics.compute_press(Phi: Array, y: Array, coefficients: Array, sample_weight: Array | None = None) → float#

Compute PRESS (Predicted Residual Error Sum of Squares).

PRESS = sum_i (e_i / (1 - h_ii))^2

Parameters:

Phi (jnp.ndarray) – Design matrix.
y (jnp.ndarray) – Target values.
coefficients (jnp.ndarray) – Fitted coefficients.
sample_weight (jnp.ndarray, optional) – Per-sample weights; see compute_loo_mse().

Returns:

press – PRESS statistic.

Return type:

float

Raises:

ValueError – If sample_weight is invalid.

jaxsr.metrics.compute_mse(y_true: Array, y_pred: Array, sample_weight: Array | None = None) → float#

Compute mean squared error, optionally weighted.

Parameters:

y_true (jnp.ndarray) – True values of shape (n_samples,).
y_pred (jnp.ndarray) – Predicted values of shape (n_samples,).
sample_weight (jnp.ndarray, optional) – Per-sample weights. Weights are normalised to sum to n_samples, so the result is sum_i w_i r_i^2 / n.

Returns:

(Weighted) mean squared error.

Return type:

float

Raises:

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

jaxsr.metrics.compute_rmse(y_true: Array, y_pred: Array, sample_weight: Array | None = None) → float#

Compute root mean squared error, optionally weighted.

Parameters:

y_true (jnp.ndarray) – True values.
y_pred (jnp.ndarray) – Predicted values.
sample_weight (jnp.ndarray, optional) – Per-sample weights.

Returns:

(Weighted) root mean squared error.

Return type:

float

Raises:

ValueError – If sample_weight is invalid.

jaxsr.metrics.compute_mae(y_true: Array, y_pred: Array, sample_weight: Array | None = None) → float#

Compute mean absolute error, optionally weighted.

Parameters:

y_true (jnp.ndarray) – True values.
y_pred (jnp.ndarray) – Predicted values.
sample_weight (jnp.ndarray, optional) – Per-sample weights.

Returns:

(Weighted) mean absolute error.

Return type:

float

Raises:

ValueError – If sample_weight is invalid.

jaxsr.metrics.compute_r2(y_true: Array, y_pred: Array, sample_weight: Array | None = None) → float#

Compute R-squared (coefficient of determination), optionally weighted.

Parameters:

y_true (jnp.ndarray) – True values.
y_pred (jnp.ndarray) – Predicted values.
sample_weight (jnp.ndarray, optional) – Per-sample weights. Both the residual and the total sum of squares are weighted, and the total is taken about the weighted mean of y_true.

Returns:

(Weighted) R-squared.

Return type:

float

Raises:

ValueError – If sample_weight is invalid.

jaxsr.metrics.compute_adjusted_r2(y_true: Array, y_pred: Array, n_params: int, sample_weight: Array | None = None) → float#

Compute adjusted R-squared.

Adjusted R² = 1 - (1 - R²) * (n - 1) / (n - k - 1)

Parameters:

y_true (jnp.ndarray) – True values.
y_pred (jnp.ndarray) – Predicted values.
n_params (int) – Number of model parameters.
sample_weight (jnp.ndarray, optional) – Per-sample weights. n remains the nominal sample count; only R² itself is weighted.

Returns:

adj_r2 – Adjusted R-squared.

Return type:

float

Raises:

ValueError – If sample_weight is invalid.

jaxsr.metrics.compute_max_error(y_true: Array, y_pred: Array) → float#: Compute maximum absolute error (unweighted – a max is not an average).

jaxsr.metrics.compute_mape(y_true: Array, y_pred: Array, sample_weight: Array | None = None) → float#

Compute mean absolute percentage error, optionally weighted.

MAPE = mean(|y_true - y_pred| / |y_true|) * 100

Parameters:

y_true (jnp.ndarray) – True values. Entries with |y_true| <= 1e-10 are skipped.
y_pred (jnp.ndarray) – Predicted values.
sample_weight (jnp.ndarray, optional) – Per-sample weights.

Returns:

(Weighted) MAPE in percent, or inf if every target is ~0.

Return type:

float

Raises:

ValueError – If sample_weight is invalid.

jaxsr.metrics.compute_all_metrics(y_true: Array, y_pred: Array, n_params: int, sample_weight: Array | None = None) → dict[str, float]#

Compute all standard regression metrics.

Parameters:

y_true (jnp.ndarray) – True values.
y_pred (jnp.ndarray) – Predicted values.
n_params (int) – Number of model parameters.
sample_weight (jnp.ndarray, optional) – Per-sample weights. Every averaged metric is weighted; "max_error" is taken over the samples with non-zero weight, since a maximum has no weighted analogue.

Returns:

metrics – Dictionary containing all metrics.

Return type:

dict

Raises:

ValueError – If sample_weight is invalid.

class jaxsr.metrics.ModelComparison(models: list[SymbolicRegressor], names: list[str], train_metrics: list[dict[str, float]], test_metrics: list[dict[str, float]] | None, rankings: dict[str, list[int]])#

Bases: object

Container for model comparison results.

models: list[SymbolicRegressor]#

names: list[str]#

train_metrics: list[dict[str, float]]#

test_metrics: list[dict[str, float]] | None#

rankings: dict[str, list[int]]#

jaxsr.metrics.compare_models(models: list[SymbolicRegressor], X_train: jnp.ndarray, y_train: jnp.ndarray, X_test: jnp.ndarray | None = None, y_test: jnp.ndarray | None = None, names: list[str] | None = None) → ModelComparison#

Compare multiple fitted models.

Parameters:

models (list of SymbolicRegressor) – Fitted models to compare.
X_train (jnp.ndarray) – Training features.
y_train (jnp.ndarray) – Training targets.
X_test (jnp.ndarray, optional) – Test features.
y_test (jnp.ndarray, optional) – Test targets.
names (list of str, optional) – Names for each model.

Returns:

comparison – Comparison results including metrics and rankings.

Return type:

ModelComparison

jaxsr.metrics.format_comparison_table(comparison: ModelComparison) → str#

Format model comparison as a text table.

Parameters:: comparison (ModelComparison) – Comparison results.
Returns:: table – Formatted table string.
Return type:: str

jaxsr.metrics.compute_classification_ic(n_samples: int, n_params: int, neg_log_likelihood: float, criterion: str = 'bic') → float#

Compute information criterion from Bernoulli negative log-likelihood.

Unlike regression IC (which starts from MSE/Gaussian likelihood), this uses the Bernoulli log-likelihood directly:

AIC = 2*NLL + 2*k BIC = 2*NLL + k*log(n) AICc = AIC + 2*k*(k+1)/(n-k-1)

Parameters:

n_samples (int) – Number of samples.
n_params (int) – Number of model parameters.
neg_log_likelihood (float) – Negative log-likelihood (sum, not mean).
criterion (str) – One of "aic", "aicc", "bic".

Returns:

ic – Information criterion value (lower is better).

Return type:

float

Raises:

ValueError – If criterion is not one of "aic", "aicc", "bic".

jaxsr.metrics.compute_accuracy(y_true: Array, y_pred: Array) → float#

Compute classification accuracy.

Parameters:

y_true (jnp.ndarray) – True class labels.
y_pred (jnp.ndarray) – Predicted class labels.

Returns:

accuracy – Fraction of correct predictions.

Return type:

float

jaxsr.metrics.compute_log_loss(y_true: Array, y_pred_proba: Array, eps: float = 1e-15) → float#

Compute binary or multiclass log-loss (cross-entropy).

Parameters:

y_true (jnp.ndarray) – True class labels of shape (n,).
y_pred_proba (jnp.ndarray) – Predicted probabilities. Shape (n,) for binary or (n, K) for multiclass.
eps (float) – Clipping bound for numerical safety.

Returns:

loss – Mean negative log-likelihood per sample.

Return type:

float

jaxsr.metrics.compute_precision(y_true: Array, y_pred: Array, pos_label: int = 1) → float#

Compute precision for a binary classification problem.

Parameters:

y_true (jnp.ndarray) – True class labels.
y_pred (jnp.ndarray) – Predicted class labels.
pos_label (int) – Label considered as positive.

Returns:

precision – TP / (TP + FP). Returns 0.0 when there are no positive predictions.

Return type:

float

jaxsr.metrics.compute_recall(y_true: Array, y_pred: Array, pos_label: int = 1) → float#

Compute recall (sensitivity) for a binary classification problem.

Parameters:

y_true (jnp.ndarray) – True class labels.
y_pred (jnp.ndarray) – Predicted class labels.
pos_label (int) – Label considered as positive.

Returns:

recall – TP / (TP + FN). Returns 0.0 when there are no positive samples.

Return type:

float

jaxsr.metrics.compute_f1_score(y_true: Array, y_pred: Array, pos_label: int = 1) → float#

Compute F1 score (harmonic mean of precision and recall).

Parameters:

y_true (jnp.ndarray) – True class labels.
y_pred (jnp.ndarray) – Predicted class labels.
pos_label (int) – Label considered as positive.

Returns:

f1 – F1 score. Returns 0.0 when precision + recall = 0.

Return type:

float

jaxsr.metrics.compute_auc_roc(y_true: Array, y_score: Array) → float#

Compute Area Under the ROC Curve via the trapezoidal rule.

Parameters:

y_true (jnp.ndarray) – True binary labels (0 or 1).
y_score (jnp.ndarray) – Predicted scores or probabilities for the positive class.

Returns:

auc – AUC-ROC value in [0, 1].

Return type:

float

Raises:

ValueError – If y_true contains fewer than two distinct classes.

jaxsr.metrics.compute_confusion_matrix(y_true: Array, y_pred: Array, n_classes: int | None = None) → ndarray#

Compute the confusion matrix.

Parameters:

y_true (jnp.ndarray) – True class labels.
y_pred (jnp.ndarray) – Predicted class labels.
n_classes (int, optional) – Number of classes. Inferred from data if None.

Returns:

cm – Confusion matrix of shape (n_classes, n_classes) where cm[i, j] is the count of samples with true label i and predicted label j.

Return type:

np.ndarray

jaxsr.metrics.compute_matthews_corrcoef(y_true: Array, y_pred: Array) → float#

Compute Matthews Correlation Coefficient for binary classification.

Parameters:

y_true (jnp.ndarray) – True binary labels.
y_pred (jnp.ndarray) – Predicted binary labels.

Returns:

mcc – MCC in [-1, 1]. Returns 0.0 when the denominator is zero.

Return type:

float

jaxsr.metrics.compute_all_classification_metrics(y_true: Array, y_pred: Array, y_pred_proba: Array | None = None, n_params: int = 0) → dict[str, float]#

Compute a comprehensive suite of classification metrics.

Parameters:

y_true (jnp.ndarray) – True class labels.
y_pred (jnp.ndarray) – Predicted class labels.
y_pred_proba (jnp.ndarray, optional) – Predicted probabilities (enables log-loss and AUC-ROC).
n_params (int) – Number of model parameters (for IC calculation).

Returns:

metrics – Dictionary of metric name to value.

Return type:

dict[str, float]

jaxsr.metrics.cross_validate_classification(model, X: Array, y: Array, cv: int = 5, scoring: str = 'accuracy', random_state: int | None = None) → dict[str, Any]#

Perform k-fold cross-validation for a classification model.

Parameters:

model (SymbolicClassifier) – Model to evaluate (must implement fit and predict).
X (jnp.ndarray) – Feature matrix.
y (jnp.ndarray) – Target labels.
cv (int) – Number of folds.
scoring (str) – Scoring metric: "accuracy", "neg_log_loss", "f1".
random_state (int, optional) – Random seed for fold splitting.

Returns:

results – Dictionary with keys "test_scores", "train_scores", "mean_test_score", "std_test_score", "mean_train_score", "std_train_score".

Return type:

dict

Raises:

ValueError – If scoring is not a recognised metric name.

jaxsr.metrics

Contents

jaxsr.metrics#