JAXSR Documentation

JAX-based Symbolic Regression

JAXSR is a Python library for discovering interpretable algebraic expressions from data using sparse optimization techniques.

Overview

JAXSR provides tools for symbolic regression - the task of finding mathematical expressions that describe relationships in data. Unlike black-box machine learning methods, symbolic regression produces human-readable equations that can provide scientific insight.

Key features:

• Flexible Basis Functions: Build custom libraries of candidate functions

• Multiple Selection Strategies: Choose from greedy, exhaustive, or LASSO-based methods

• Uncertainty Quantification: Prediction intervals, Bayesian Model Averaging, conformal prediction, and bootstrap methods

• Physical Constraints: Incorporate domain knowledge through constraints

• Additive Symbolic Regression: Boosting-style ensembles of small symbolic expressions (jaxsr.additive)

• JAX-Powered: GPU acceleration, JIT compilation, automatic differentiation

• Scikit-learn Compatible: Familiar fit/predict interface

Installation

pip install jaxsr

For development:

git clone https://github.com/jkitchin/jaxsr.git
cd jaxsr
pip install -e ".[dev]"

Quick Start

from jaxsr import BasisLibrary, SymbolicRegressor
import jax.numpy as jnp

# Create basis library
library = (BasisLibrary(n_features=2, feature_names=["x", "y"])
    .add_constant()
    .add_linear()
    .add_polynomials(max_degree=3)
    .add_interactions()
)

# Fit model
model = SymbolicRegressor(basis_library=library, max_terms=5)
model.fit(X, y)

# Results
print(model.expression_)
print(f"R² = {model.metrics_['r2']:.4f}")

Documentation Contents

• Quickstart Guide - Get started quickly

• Design of Experiments Guide - Adaptive DOE and active learning

• Acquisition Functions - Detailed acquisition function reference

• CLI Guide - Command-line interface reference

• Claude Code Skills - AI-assisted workflows with Claude Code

• API Reference - Complete API documentation

• Literature Review - Background on symbolic regression

• Examples - Worked examples for various applications

How It Works

JAXSR follows the ALAMO (Automated Learning of Algebraic Models for Optimization) methodology:

1. Basis Library Construction: Define a library of candidate basis functions (polynomials, transcendentals, interactions, etc.)

2. Design Matrix Evaluation: Evaluate all basis functions on training data to create a design matrix Φ

3. Sparse Selection: Use information criteria (BIC, AIC) to select a sparse subset of basis functions

4. Coefficient Fitting: Fit coefficients via least squares, optionally with constraints

5. Model Analysis: Examine Pareto front, export to various formats

When to Use JAXSR

JAXSR is ideal when you:

• Want interpretable models rather than black boxes

• Have domain knowledge to constrain the solution space

• Need to discover physical laws or empirical correlations

• Require reproducible results (deterministic algorithms)

• Want to explore the accuracy-complexity trade-off

Comparison with Other Tools

Feature	JAXSR	ALAMO	PySR	GP
Open Source	✓	✗	✓	✓
Deterministic	✓	✓	✗	✗
UQ / Intervals	✓	Limited	✗	✗
Constraints	✓	✓	Limited	Limited
GPU Support	✓	✗	✓	Varies

License

JAXSR is released under the MIT License.