jaxsr.classifier

jaxsr.classifier#

Symbolic Classification for JAXSR.

Provides a SymbolicClassifier that discovers interpretable logistic models from data. Binary classification uses the sigmoid link; multiclass uses one-vs-rest (OVR) with per-class expressions.

The classifier reuses BasisLibrary for design-matrix construction and IRLS for coefficient fitting, keeping the same sparse-selection workflow as SymbolicRegressor.

class jaxsr.classifier.SymbolicClassifier(basis_library: BasisLibrary | None = None, max_terms: int = 5, strategy: str = 'greedy_forward', information_criterion: str = 'bic', regularization: float | None = None, constraints: Constraints | None = None, random_state: int | None = None, max_iter: int = 100, tol: float = 1e-06)#

Bases: _SklearnCompatMixin

JAX-accelerated symbolic classification using sparse selection.

Discovers interpretable logistic models by selecting sparse subsets of basis functions and fitting via IRLS. Binary problems use a single sigmoid link; multiclass problems are handled via one-vs-rest (OVR), giving each class its own interpretable expression.

Parameters:

  • basis_library (BasisLibrary, optional) – Library of candidate basis functions.

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

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

  • information_criterion (str) – IC for model selection: "aic", "aicc", "bic".

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

  • constraints (Constraints, optional) – Physical constraints (applied to linear predictor).

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

  • max_iter (int) – Maximum IRLS iterations.

  • tol (float) – IRLS convergence tolerance.

expression_#

Human-readable linear-predictor expression (binary) or dict of expressions (multiclass).

Type:

str

coefficients_#

Fitted coefficients.

Type:

jnp.ndarray

selected_features_#

Selected basis function names.

Type:

list of str

classes_#

Unique class labels found during fit().

Type:

jnp.ndarray

metrics_#

Classification metrics computed on training data.

Type:

dict

Examples

>>> from jaxsr import BasisLibrary
>>> from jaxsr.classifier import SymbolicClassifier
>>> lib = BasisLibrary(2).add_constant().add_linear().add_polynomials()
>>> clf = SymbolicClassifier(basis_library=lib, max_terms=4)
>>> clf.fit(X_train, y_train)
>>> print(clf.expression_)
>>> proba = clf.predict_proba(X_test)
property expression_: str | dict[int, str]#

Human-readable expression for the linear predictor.

For binary classification, returns a string. For multiclass, returns a dict mapping class label to expression string.

property coefficients_: Array#

1-D; multiclass: list).

Type:

Fitted coefficients (binary

property selected_features_: list[str]#

Names of selected basis functions (binary model).

property classes_: Array#

Unique class labels found during fit.

property metrics_: dict[str, float]#

Training classification metrics.

property pareto_front_: list[ClassificationResult]#

Pareto-optimal models from the selection path.

fit(X: Array, y: Array) → SymbolicClassifier#

Fit the symbolic classifier.

Automatically detects binary vs multiclass from unique(y).

Parameters:

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

  • y (array-like of shape (n_samples,)) – Class labels (integers or floats castable to int).

Returns:

self – Fitted model.

Return type:

SymbolicClassifier

Raises:

ValueError – If basis_library is None, shapes are inconsistent, or labels contain fewer than 2 classes.

decision_function(X: Array) → Array#

Compute raw logits (linear predictor values).

Parameters:

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

Returns:

logits – Shape (n,) for binary, (n, K) for multiclass.

Return type:

jnp.ndarray

predict_proba(X: Array) → Array#

Predict class probabilities.

Parameters:

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

Returns:

proba – Shape (n, K) probability matrix. Columns correspond to self.classes_.

Return type:

jnp.ndarray

predict_log_proba(X: Array) → Array#

Predict log class probabilities.

Parameters:

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

Returns:

log_proba – Shape (n, K) log-probability matrix.

Return type:

jnp.ndarray

predict(X: Array) → Array#

Predict class labels.

Parameters:

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

Returns:

y_pred – Predicted class labels of shape (n,).

Return type:

jnp.ndarray

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

Compute classification accuracy (sklearn convention).

Parameters:

  • X (array-like) – Test data.

  • y (array-like) – True labels.

Returns:

accuracy – Fraction of correct predictions.

Return type:

float

to_sympy()#

Convert expression to SymPy.

For binary: returns 1 / (1 + exp(-linear_predictor)). For multiclass: returns a dict mapping class → SymPy expression.

Returns:

expr – SymPy representation of the probability model.

Return type:

sympy.Expr or dict

to_latex() → str#

Convert expression to LaTeX string.

Returns:

latex – LaTeX representation.

Return type:

str

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

Convert to a pure NumPy callable that returns probabilities.

Returns:

func – func(X) -> proba where proba has shape (n, K).

Return type:

callable

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

Wald confidence intervals for coefficients (binary only).

Parameters:

alpha (float) – Significance level.

Returns:

intervals – {name: (estimate, lower, upper, se)}.

Return type:

dict

predict_conformal(X: Array, alpha: float = 0.05, X_cal: Array | None = None, y_cal: Array | None = None) → dict[str, Any]#

Conformal prediction sets for classification.

Parameters:

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

  • alpha (float) – Significance level.

  • X_cal (jnp.ndarray, optional) – Calibration features. If None, uses training data.

  • y_cal (jnp.ndarray, optional) – Calibration labels. If None, uses training data.

Returns:

result – Keys: "prediction_sets", "quantile", "y_pred".

Return type:

dict

save(filepath: str) → None#

Save model to JSON file.

Parameters:

filepath (str) – Path to save the model.

classmethod load(filepath: str) → SymbolicClassifier#

Load model from JSON file.

Parameters:

filepath (str) – Path to the model file.

Returns:

model – Loaded model.

Return type:

SymbolicClassifier

summary() → str#

Return a human-readable summary of the fitted model.

Returns:

summary – Model summary string.

Return type:

str

jaxsr.classifier.fit_symbolic_classification(X: Array, y: Array, feature_names: list[str] | None = None, max_terms: int = 5, max_poly_degree: int = 3, include_transcendental: bool = False, strategy: str = 'greedy_forward', information_criterion: str = 'bic') → SymbolicClassifier#

Convenience function for quick symbolic classification.

Builds a default BasisLibrary, constructs a SymbolicClassifier, and fits it.

Parameters:

  • X (array-like) – Input features.

  • y (array-like) – Class labels.

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

  • strategy (str) – Selection strategy.

  • information_criterion (str) – Information criterion.

Returns:

model – Fitted classifier.

Return type:

SymbolicClassifier