jaxsr.selection

Contents

jaxsr.selection#

Model Selection Algorithms for JAXSR.

Implements multiple strategies for selecting sparse subsets of basis functions: - Greedy forward selection - Greedy backward elimination - Exhaustive search - LASSO path screening - Coordinate descent LASSO - Logistic (classification) selection via IRLS and FISTA

class jaxsr.selection.SelectionResult(coefficients: Array, selected_indices: Array, selected_names: list[str], mse: float, complexity: int, aic: float, bic: float, aicc: float, n_samples: int, parametric_params: dict[int, dict[str, float]] | None = None)#

Bases: object

Result of model selection for a single model configuration.

Parameters:

  • coefficients (jnp.ndarray) – Fitted coefficients for selected features.

  • selected_indices (jnp.ndarray) – Indices of selected basis functions.

  • selected_names (list of str) – Names of selected basis functions.

  • mse (float) – Mean squared error.

  • complexity (int) – Total complexity score.

  • aic (float) – Akaike Information Criterion.

  • bic (float) – Bayesian Information Criterion.

  • aicc (float) – Corrected AIC.

  • n_samples (int) – Number of training samples.

coefficients: Array#
selected_indices: Array#
selected_names: list[str]#
mse: float#
complexity: int#
aic: float#
bic: float#
aicc: float#
n_samples: int#
parametric_params: dict[int, dict[str, float]] | None = None#
expression() → str#

Return human-readable expression.

property n_terms: int#

Number of terms in the model.

to_dict() → dict[str, Any]#

Serialize to dictionary.

class jaxsr.selection.SelectionPath(results: list[SelectionResult], strategy: str, best_index: int)#

Bases: object

Full path of model selection (for visualization).

Parameters:

  • results (list of SelectionResult) – Results at each step.

  • strategy (str) – Selection strategy used.

  • best_index (int) – Index of best model according to information criterion.

results: list[SelectionResult]#
strategy: str#
best_index: int#
property best: SelectionResult#

Return the best model.

jaxsr.selection.fit_ols(Phi: Array, y: Array, sample_weight: Array | None = None) → tuple[Array, float]#

Fit ordinary (or weighted) least squares.

Parameters:

  • Phi (jnp.ndarray) – Design matrix of shape (n_samples, n_features).

  • y (jnp.ndarray) – Target vector.

  • sample_weight (jnp.ndarray, optional) – Per-sample weights of shape (n_samples,). When given, minimises sum_i w_i (y_i - phi_i @ c)^2 and returns the weighted MSE. Weights are normalised to sum to n_samples first, so the returned MSE stays on the same scale as the unweighted one.

Returns:

  • coefficients (jnp.ndarray) – Fitted coefficients.

  • mse (float) – (Weighted) mean squared error.

Raises:

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

jaxsr.selection.fit_ridge(Phi: Array, y: Array, alpha: float, sample_weight: Array | None = None) → tuple[Array, float]#

Fit ridge regression (L2 regularized OLS).

Solves: minimize ||y - Phi @ w||^2 + alpha * ||w||^2

Parameters:

  • Phi (jnp.ndarray) – Design matrix of shape (n_samples, n_features).

  • y (jnp.ndarray) – Target vector.

  • alpha (float) – L2 regularization strength.

  • sample_weight (jnp.ndarray, optional) – Per-sample weights of shape (n_samples,). When given, the data term becomes sum_i w_i (y_i - phi_i @ c)^2.

Returns:

  • coefficients (jnp.ndarray) – Fitted coefficients.

  • mse (float) – (Weighted) mean squared error (without regularization term).

Raises:

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

jaxsr.selection.fit_subset(Phi: Array, y: Array, indices: list[int] | Array, basis_names: list[str], complexities: Array, regularization: float | None = None, X: Array | None = None, basis_library=None, param_optimizer: str = 'scipy', param_optimization_budget: int = 50, _param_cache: dict | None = None, sqrt_sample_weight: Array | None = None) → SelectionResult#

Fit model using only selected basis functions.

Parameters:

  • Phi (jnp.ndarray) – Full design matrix.

  • y (jnp.ndarray) – Target vector.

  • indices (array-like) – Indices of selected basis functions.

  • basis_names (list of str) – Names of all basis functions.

  • complexities (jnp.ndarray) – Complexity scores for all basis functions.

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

  • X (jnp.ndarray, optional) – Raw input data – required when parametric basis functions are present.

  • basis_library (BasisLibrary, optional) – The basis library – required when parametric basis functions are present.

  • param_optimizer (str) – "scipy" (default) or "optuna".

  • param_optimization_budget (int) – Number of optuna trials (ignored for scipy).

  • _param_cache (dict, optional) – Cache for optimised parametric parameter values.

  • sqrt_sample_weight (jnp.ndarray, optional) – sqrt(sample_weight) already applied to Phi and y by the caller. Only needed for parametric libraries, whose columns are rebuilt from the raw X and must be whitened to match.

Returns:

result – Fitting result.

Return type:

SelectionResult

jaxsr.selection.greedy_forward_selection(Phi: Array, y: Array, basis_names: list[str], complexities: Array, max_terms: int = 5, information_criterion: str = 'bic', early_stop: bool = True, candidate_indices: list[int] | None = None, regularization: float | None = None, X: Array | None = None, basis_library=None, param_optimizer: str = 'scipy', param_optimization_budget: int = 50, constraint_scorer=None, sample_weight: Array | None = None) → SelectionPath#

Greedy forward stepwise selection.

Starting from empty model, iteratively add the basis function that most improves the information criterion until no improvement or max_terms reached.

Parameters:

  • Phi (jnp.ndarray) – Design matrix.

  • y (jnp.ndarray) – Target vector.

  • basis_names (list of str) – Names of basis functions.

  • complexities (jnp.ndarray) – Complexity scores.

  • max_terms (int) – Maximum number of terms.

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

  • early_stop (bool) – If True, stop when IC stops improving.

  • candidate_indices (list of int, optional) – Indices of candidate basis functions to consider.

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

  • sample_weight (jnp.ndarray, optional) – Per-sample weights. Every candidate model is fitted by weighted least squares and scored by the weighted MSE, so the weights steer term selection, not just the final coefficients.

Returns:

path – Selection path with all intermediate results.

Return type:

SelectionPath

Raises:

ValueError – If sample_weight is invalid.

jaxsr.selection.greedy_backward_elimination(Phi: Array, y: Array, basis_names: list[str], complexities: Array, min_terms: int = 1, information_criterion: str = 'bic', start_indices: list[int] | None = None, max_terms: int = 5, regularization: float | None = None, X: Array | None = None, basis_library=None, param_optimizer: str = 'scipy', param_optimization_budget: int = 50, constraint_scorer=None, sample_weight: Array | None = None) → SelectionPath#

Greedy backward elimination.

Starting from full model (or specified subset), iteratively remove the term whose removal most improves the information criterion.

Parameters:

  • Phi (jnp.ndarray) – Design matrix.

  • y (jnp.ndarray) – Target vector.

  • basis_names (list of str) – Names of basis functions.

  • complexities (jnp.ndarray) – Complexity scores.

  • min_terms (int) – Minimum number of terms to keep.

  • information_criterion (str) – Criterion for selection.

  • start_indices (list of int, optional) – Starting indices. If None, uses all basis functions.

  • sample_weight (jnp.ndarray, optional) – Per-sample weights; every candidate model is fitted and scored by weighted least squares.

Returns:

path – Selection path.

Return type:

SelectionPath

Raises:

ValueError – If sample_weight is invalid.

Exhaustive search over all combinations.

Enumerate all combinations up to max_terms and select the best.

Parameters:

  • Phi (jnp.ndarray) – Design matrix.

  • y (jnp.ndarray) – Target vector.

  • basis_names (list of str) – Names of basis functions.

  • complexities (jnp.ndarray) – Complexity scores.

  • max_terms (int) – Maximum number of terms.

  • information_criterion (str) – Criterion for selection.

  • candidate_indices (list of int, optional) – Indices to consider. If None, uses all.

  • max_combinations (int) – Maximum combinations to evaluate (safety limit).

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

  • sample_weight (jnp.ndarray, optional) – Per-sample weights; every subset is fitted and scored by weighted least squares.

Returns:

path – Selection path (all Pareto-optimal models).

Return type:

SelectionPath

Raises:

ValueError – If too many combinations would be evaluated, or if sample_weight is invalid.

jaxsr.selection.coordinate_descent_lasso(Phi: Array, y: Array, alpha: float, max_iter: int = 1000, tol: float = 1e-06, warm_start: Array | None = None) → Array#

Solve LASSO using coordinate descent.

minimize (1/2n) ||y - Phi @ w||^2 + alpha * ||w||_1

Parameters:

  • Phi (jnp.ndarray) – Design matrix (should be standardized).

  • y (jnp.ndarray) – Target vector (should be centered).

  • alpha (float) – Regularization parameter.

  • max_iter (int) – Maximum iterations.

  • tol (float) – Convergence tolerance.

  • warm_start (jnp.ndarray, optional) – Initial coefficients.

Returns:

coefficients – LASSO coefficients.

Return type:

jnp.ndarray

jaxsr.selection.lasso_path_selection(Phi: Array, y: Array, basis_names: list[str], complexities: Array, max_terms: int = 5, information_criterion: str = 'bic', n_alphas: int = 100, alpha_min_ratio: float = 0.0001, regularization: float | None = None, X: Array | None = None, basis_library=None, param_optimizer: str = 'scipy', param_optimization_budget: int = 50, constraint_scorer=None, sample_weight: Array | None = None) → SelectionPath#

LASSO path screening for variable selection.

Use LASSO regularization path to identify promising subsets, then refit with OLS for final selection.

Parameters:

  • Phi (jnp.ndarray) – Design matrix.

  • y (jnp.ndarray) – Target vector.

  • basis_names (list of str) – Names of basis functions.

  • complexities (jnp.ndarray) – Complexity scores.

  • max_terms (int) – Maximum number of terms.

  • information_criterion (str) – Criterion for selection.

  • n_alphas (int) – Number of regularization values.

  • alpha_min_ratio (float) – Ratio of min to max alpha.

  • sample_weight (jnp.ndarray, optional) – Per-sample weights. The screening LASSO and the OLS refit of each active set both run on the weighted problem.

Returns:

path – Selection path.

Return type:

SelectionPath

Raises:

ValueError – If sample_weight is invalid.

jaxsr.selection.coordinate_descent_elastic_net(Phi: Array, y: Array, alpha: float, l1_ratio: float = 0.5, max_iter: int = 1000, tol: float = 1e-06, warm_start: Array | None = None) → Array#

Solve Elastic Net using coordinate descent.

minimize (1/2n) ||y - Phi @ w||^2 + alpha * (l1_ratio * ||w||_1 + (1-l1_ratio)/2 * ||w||_2^2)

Parameters:

  • Phi (jnp.ndarray) – Design matrix.

  • y (jnp.ndarray) – Target vector.

  • alpha (float) – Overall regularization parameter.

  • l1_ratio (float) – Mix between L1 (1) and L2 (0).

  • max_iter (int) – Maximum iterations.

  • tol (float) – Convergence tolerance.

  • warm_start (jnp.ndarray, optional) – Initial coefficients.

Returns:

coefficients – Elastic net coefficients.

Return type:

jnp.ndarray

jaxsr.selection.compute_pareto_front(results: list[SelectionResult]) → list[SelectionResult]#

Extract Pareto-optimal models (complexity vs MSE).

Parameters:

results (list of SelectionResult) – All candidate models.

Returns:

pareto – Pareto-optimal models (sorted by complexity).

Return type:

list of SelectionResult

jaxsr.selection.compute_pareto_front_multi(results: list[SelectionResult], objectives: list[str] | None = None) → list[SelectionResult]#

Compute Pareto front for multiple objectives.

Parameters:

  • results (list of SelectionResult) – All candidate models.

  • objectives (list of str) – Objective names (attributes of SelectionResult to minimize).

Returns:

pareto – Pareto-optimal models.

Return type:

list of SelectionResult

jaxsr.selection.select_features(Phi: Array, y: Array, basis_names: list[str], complexities: Array, strategy: str = 'greedy_forward', max_terms: int = 5, information_criterion: str = 'bic', sample_weight: Array | None = None, **kwargs) → SelectionPath#

Run feature selection with the specified strategy.

Parameters:

  • Phi (jnp.ndarray) – Design matrix.

  • y (jnp.ndarray) – Target vector.

  • basis_names (list of str) – Names of basis functions.

  • complexities (jnp.ndarray) – Complexity scores.

  • strategy (str) – One of “greedy_forward”, “greedy_backward”, “exhaustive”, “lasso_path”.

  • max_terms (int) – Maximum number of terms.

  • information_criterion (str) – Information criterion for model selection.

  • sample_weight (jnp.ndarray, optional) – Per-sample weights of shape (n_samples,). Pass Phi and y unweighted – the strategy whitens them internally. Every strategy supports weights, so selection, coefficients and the reported MSE are all consistent with the weighted objective.

  • **kwargs – Additional arguments for specific strategies.

Returns:

path – Selection results.

Return type:

SelectionPath

Raises:

ValueError – If strategy is unknown or sample_weight is invalid.

class jaxsr.selection.ClassificationResult(coefficients: Array, selected_indices: Array, selected_names: list[str], neg_log_likelihood: float, complexity: int, aic: float, bic: float, aicc: float, n_samples: int, n_iter: int = 0, converged: bool = True)#

Bases: object

Result of model selection for a single logistic model.

Parameters:

  • coefficients (jnp.ndarray) – Fitted coefficients for selected features.

  • selected_indices (jnp.ndarray) – Indices of selected basis functions.

  • selected_names (list of str) – Names of selected basis functions.

  • neg_log_likelihood (float) – Negative log-likelihood (sum, not mean).

  • complexity (int) – Total complexity score.

  • aic (float) – Akaike Information Criterion (from Bernoulli likelihood).

  • bic (float) – Bayesian Information Criterion.

  • aicc (float) – Corrected AIC.

  • n_samples (int) – Number of training samples.

  • n_iter (int) – Number of IRLS iterations used.

  • converged (bool) – Whether IRLS converged within tolerance.

coefficients: Array#
selected_indices: Array#
selected_names: list[str]#
neg_log_likelihood: float#
complexity: int#
aic: float#
bic: float#
aicc: float#
n_samples: int#
n_iter: int = 0#
converged: bool = True#
expression() → str#

Return human-readable expression for the linear predictor.

property n_terms: int#

Number of terms in the model.

to_dict() → dict[str, Any]#

Serialize to dictionary.

class jaxsr.selection.ClassificationPath(results: list[ClassificationResult], strategy: str, best_index: int)#

Bases: object

Full path of classification model selection.

Parameters:

  • results (list of ClassificationResult) – Results at each step.

  • strategy (str) – Selection strategy used.

  • best_index (int) – Index of best model according to information criterion.

results: list[ClassificationResult]#
strategy: str#
best_index: int#
property best: ClassificationResult#

Return the best model.

jaxsr.selection.fit_irls(Phi: Array, y: Array, max_iter: int = 100, tol: float = 1e-06, regularization: float | None = None) → tuple[Array, float, int, bool]#

Fit logistic regression via Iteratively Reweighted Least Squares.

Solves binary logistic regression by iteratively solving weighted least-squares problems. Each iteration has quadratic convergence near the optimum.

Parameters:

  • Phi (jnp.ndarray) – Design matrix of shape (n, p).

  • y (jnp.ndarray) – Binary labels in {0, 1} of shape (n,).

  • max_iter (int) – Maximum number of IRLS iterations.

  • tol (float) – Convergence tolerance on max(|w_new - w_old|).

  • regularization (float, optional) – L2 ridge penalty to add to the diagonal of the weighted normal equations. Helps when data are (nearly) separable.

Returns:

  • coefficients (jnp.ndarray) – Fitted coefficient vector of shape (p,).

  • nll (float) – Negative log-likelihood (sum, not mean).

  • n_iter (int) – Number of iterations performed.

  • converged (bool) – Whether the algorithm converged within tol.

jaxsr.selection.fit_logistic_lasso(Phi: Array, y: Array, alpha: float, l1_ratio: float = 1.0, max_iter: int = 500, tol: float = 1e-06, warm_start: Array | None = None) → Array#

L1/elastic-net penalised logistic regression via FISTA.

Minimises:

(1/n) * NLL(w) + alpha * (l1_ratio * ||w||_1
                           + (1-l1_ratio)/2 * ||w||_2^2)
Parameters:

  • Phi (jnp.ndarray) – Design matrix of shape (n, p).

  • y (jnp.ndarray) – Binary labels of shape (n,).

  • alpha (float) – Overall regularisation strength.

  • l1_ratio (float) – Mix between L1 (1.0) and L2 (0.0).

  • max_iter (int) – Maximum FISTA iterations.

  • tol (float) – Convergence tolerance.

  • warm_start (jnp.ndarray, optional) – Initial coefficient vector.

Returns:

coefficients – Fitted (sparse) coefficient vector.

Return type:

jnp.ndarray

jaxsr.selection.fit_classification_subset(Phi: Array, y: Array, indices: list[int] | Array, basis_names: list[str], complexities: Array, information_criterion: str = 'bic', regularization: float | None = None, max_iter: int = 100, tol: float = 1e-06) → ClassificationResult#

Fit a logistic regression model on a subset of basis functions.

Parameters:

  • Phi (jnp.ndarray) – Full design matrix.

  • y (jnp.ndarray) – Binary labels.

  • indices (array-like) – Indices of selected basis functions.

  • basis_names (list of str) – Names of all basis functions.

  • complexities (jnp.ndarray) – Complexity scores for all basis functions.

  • information_criterion (str) – IC for model evaluation.

  • regularization (float, optional) – L2 regularisation strength.

  • max_iter (int) – Maximum IRLS iterations.

  • tol (float) – IRLS convergence tolerance.

Returns:

result – Fitting result.

Return type:

ClassificationResult

jaxsr.selection.greedy_forward_classification(Phi: Array, y: Array, basis_names: list[str], complexities: Array, max_terms: int = 5, information_criterion: str = 'bic', early_stop: bool = True, candidate_indices: list[int] | None = None, regularization: float | None = None, max_iter: int = 100, tol: float = 1e-06) → ClassificationPath#

Greedy forward stepwise selection for logistic regression.

Parameters:

  • Phi (jnp.ndarray) – Design matrix.

  • y (jnp.ndarray) – Binary labels.

  • basis_names (list of str) – Names of basis functions.

  • complexities (jnp.ndarray) – Complexity scores.

  • max_terms (int) – Maximum number of terms.

  • information_criterion (str) – Criterion for selection ("aic", "bic", "aicc").

  • early_stop (bool) – If True, stop when IC stops improving.

  • candidate_indices (list of int, optional) – Indices of candidate basis functions to consider.

  • regularization (float, optional) – L2 regularisation strength.

  • max_iter (int) – Maximum IRLS iterations.

  • tol (float) – IRLS convergence tolerance.

Returns:

path – Selection path with all intermediate results.

Return type:

ClassificationPath

jaxsr.selection.greedy_backward_classification(Phi: Array, y: Array, basis_names: list[str], complexities: Array, min_terms: int = 1, information_criterion: str = 'bic', start_indices: list[int] | None = None, max_terms: int = 5, regularization: float | None = None, max_iter: int = 100, tol: float = 1e-06) → ClassificationPath#

Greedy backward elimination for logistic regression.

Parameters:

  • Phi (jnp.ndarray) – Design matrix.

  • y (jnp.ndarray) – Binary labels.

  • basis_names (list of str) – Names of basis functions.

  • complexities (jnp.ndarray) – Complexity scores.

  • min_terms (int) – Minimum number of terms to keep.

  • information_criterion (str) – Criterion for selection.

  • start_indices (list of int, optional) – Starting indices. If None, uses all basis functions.

  • max_terms (int) – For API compatibility.

  • regularization (float, optional) – L2 regularisation strength.

  • max_iter (int) – Maximum IRLS iterations.

  • tol (float) – IRLS convergence tolerance.

Returns:

path – Selection path.

Return type:

ClassificationPath

jaxsr.selection.exhaustive_classification(Phi: Array, y: Array, basis_names: list[str], complexities: Array, max_terms: int = 5, information_criterion: str = 'bic', candidate_indices: list[int] | None = None, max_combinations: int = 100000, regularization: float | None = None, max_iter: int = 100, tol: float = 1e-06) → ClassificationPath#

Exhaustive search over all combinations for logistic regression.

Parameters:

  • Phi (jnp.ndarray) – Design matrix.

  • y (jnp.ndarray) – Binary labels.

  • basis_names (list of str) – Names of basis functions.

  • complexities (jnp.ndarray) – Complexity scores.

  • max_terms (int) – Maximum number of terms.

  • information_criterion (str) – Criterion for selection.

  • candidate_indices (list of int, optional) – Indices to consider.

  • max_combinations (int) – Safety limit on combinations.

  • regularization (float, optional) – L2 regularisation strength.

  • max_iter (int) – Maximum IRLS iterations.

  • tol (float) – IRLS convergence tolerance.

Returns:

path – Selection path.

Return type:

ClassificationPath

Raises:

ValueError – If too many combinations would be evaluated.

jaxsr.selection.lasso_path_classification(Phi: Array, y: Array, basis_names: list[str], complexities: Array, max_terms: int = 5, information_criterion: str = 'bic', n_alphas: int = 50, alpha_min_ratio: float = 0.001, regularization: float | None = None, max_iter: int = 100, tol: float = 1e-06) → ClassificationPath#

LASSO path screening for logistic classification.

Traces an L1-penalised logistic regression path, identifies unique active sets, and refits each with IRLS for unbiased IC comparison.

Parameters:

  • Phi (jnp.ndarray) – Design matrix.

  • y (jnp.ndarray) – Binary labels.

  • basis_names (list of str) – Names of basis functions.

  • complexities (jnp.ndarray) – Complexity scores.

  • max_terms (int) – Maximum number of terms.

  • information_criterion (str) – Criterion for selection.

  • n_alphas (int) – Number of regularisation values.

  • alpha_min_ratio (float) – Ratio of minimum to maximum alpha.

  • regularization (float, optional) – L2 regularisation for IRLS refitting.

  • max_iter (int) – Maximum IRLS iterations.

  • tol (float) – IRLS convergence tolerance.

Returns:

path – Selection path.

Return type:

ClassificationPath

jaxsr.selection.compute_pareto_front_classification(results: list[ClassificationResult]) → list[ClassificationResult]#

Extract Pareto-optimal classification models (complexity vs NLL).

Parameters:

results (list of ClassificationResult) – All candidate models.

Returns:

pareto – Pareto-optimal models sorted by complexity.

Return type:

list of ClassificationResult

jaxsr.selection.select_features_classification(Phi: Array, y: Array, basis_names: list[str], complexities: Array, strategy: str = 'greedy_forward', max_terms: int = 5, information_criterion: str = 'bic', **kwargs) → ClassificationPath#

Run classification feature selection with the specified strategy.

Parameters:

  • Phi (jnp.ndarray) – Design matrix.

  • y (jnp.ndarray) – Binary labels.

  • basis_names (list of str) – Names of basis functions.

  • complexities (jnp.ndarray) – Complexity scores.

  • strategy (str) – One of "greedy_forward", "greedy_backward", "exhaustive", "lasso_path".

  • max_terms (int) – Maximum number of terms.

  • information_criterion (str) – Information criterion for model selection.

  • **kwargs – Additional arguments for specific strategies.

Returns:

path – Selection results.

Return type:

ClassificationPath

Raises:

ValueError – If strategy is not recognised.

Contents