JAXSR Skill — Symbolic Regression Assistant

JAXSR is a JAX-based symbolic regression library that discovers interpretable algebraic expressions from data using sparse optimization. It follows ALAMO-style methodology: build a rich candidate basis, then select the simplest model that explains the data.

Skill Activation

Activate this skill when the user wants to:

• Discover algebraic expressions or equations from data
• Set up a Design of Experiments (DOE) study
• Fit, interpret, or export symbolic regression models
• Choose between basis functions, strategies, UQ methods, or design methods
• Generate reports from experimental data
• Use the

  jaxsr

  CLI tool
• Build notebooks or scripts for symbolic regression workflows

Assistant Mode

When the user asks for help deciding how to set up, analyze, or report on their problem, enter assistant mode. In this mode, ask diagnostic questions to guide them to the right configuration. Do not dump all options at once — walk through decisions sequentially.

Step 1: Characterize the Problem

Ask the user:

1. What are you modeling? (physical system, chemical process, ML feature engineering, etc.)
2. How many input features? (1-3 is small, 4-8 is medium, 9+ is large)
3. How many data points? (< 20 is very small, 20-100 is typical, 100+ is large)
4. Do you have domain knowledge? (known physics, monotonicity, bounds, symmetry)
5. What is the goal? (interpretable equation, prediction, optimization, screening)

Step 2: Recommend Basis Library

Based on the answers, recommend a basis library configuration:

add_constant + add_linear + add_polynomials(3) + add_interactions(2) + add_transcendental()

add_constant + add_linear + add_polynomials(max_degree)

add_constant + add_linear + add_polynomials(2) + add_transcendental(funcs=["log","exp","sqrt","inv"])

add_constant + add_linear + add_transcendental(funcs=["exp","inv","log"]) + add_ratios() + add_parametric(Arrhenius)

add_constant + add_linear + add_interactions(2)

lasso_path

add_constant + add_linear + add_polynomials(2) + add_interactions(2)

ResponseSurface

add_categorical_indicators() + add_categorical_interactions()

Key guidance:

• Start simple. You can always add complexity.

• add_transcendental(safe=True)

  guards against log(0), 1/0, sqrt(<0). Always use

  safe=True

  .

• add_ratios(safe=True)

  adds x_i/x_j terms. Doubles the library size — only use when ratios are physically meaningful.

• add_parametric()

  enables nonlinear parameters (e.g.,

  exp(-a*x)

  ). Powerful but slower to fit.
• If n_features > 5, avoid

  add_polynomials(degree>2)

  — the library becomes enormous.

Step 3: Recommend Selection Strategy

exhaustive

greedy_forward

lasso_path

greedy_backward

When to change from defaults:

• greedy_forward

  is the right choice 80% of the time. It's the default.
• Use

  exhaustive

  only when the basis library is small enough (< 20 terms). It guarantees the global optimum but scales as O(2^n).
• Use

  lasso_path

  when you have a very large library and want fast screening. It may miss interaction effects.
• Use

  greedy_backward

  when you suspect many terms matter and want to start from the full model.

Step 4: Recommend Information Criterion

aicc

bic

aic

bic

bic

aicc

Default recommendation: Use

"bic"

"aicc"

"aic"

"aicc"

"bic"

"cv"

Step 5: Recommend Constraints (if applicable)

Ask: "Do you have any physical knowledge about the system?"

.add_bounds("y", lower=0)

.add_bounds("y", lower=lo, upper=hi)

.add_monotonic("T", direction="increasing")

.add_concave(feature)

.add_convex(feature)

.add_sign_constraint(basis_name, sign="positive")

.add_known_coefficient(name, value)

Use

hard=True

hard=False

Step 6: Recommend Uncertainty Quantification

model.predict_interval()

model.coefficient_intervals()

model.predict_bma()

model.predict_conformal()

bootstrap_predict()

model.predict_ensemble()

anova()

Default recommendation: Start with

predict_interval()

predict_bma()

predict_conformal()

Step 7: Recommend Reporting Format

model.summary()

jaxsr status study.jaxsr

jaxsr report study.jaxsr -o report.xlsx

jaxsr report study.jaxsr -o report.docx

model.to_latex()

plot_pareto_front()

model.to_callable()

model.save("model.json")

study.save("study.jaxsr")

"How do I use JAXSR with scikit-learn?"

See

guides/sklearn-integration.md

Quick Reference: Installation

# Core library
pip install jaxsr

# With sklearn integration (cross_val_score, GridSearchCV, Pipeline)
pip install "jaxsr[sklearn]"

# With CLI support
pip install "jaxsr[cli]"

# With Excel reporting
pip install "jaxsr[excel]"

# With Word reports
pip install "jaxsr[reports]"

# Everything for development
pip install -e ".[dev,cli,excel,reports]"

Quick Reference: Python API

Minimal Example (5 lines)

from jaxsr import fit_symbolic
import numpy as np

X = np.column_stack([x1, x2])  # shape (n_samples, n_features)
model = fit_symbolic(X, y, feature_names=["x1", "x2"], max_terms=5)
print(model.expression_)

Full Control Example

from jaxsr import BasisLibrary, SymbolicRegressor, Constraints

# 1. Build basis library
library = (BasisLibrary(n_features=2, feature_names=["T", "P"])
    .add_constant()
    .add_linear()
    .add_polynomials(max_degree=3)
    .add_interactions(max_order=2)
    .add_transcendental(funcs=["log", "exp", "sqrt"])
)

# 2. Define constraints (optional)
constraints = (Constraints()
    .add_monotonic("T", direction="increasing")
    .add_bounds("y", lower=0)
)

# 3. Fit model
model = SymbolicRegressor(
    basis_library=library,
    max_terms=5,
    strategy="greedy_forward",
    information_criterion="bic",
    constraints=constraints,
)
model.fit(X_train, y_train)

# 4. Inspect results
print(model.expression_)
print(model.metrics_)
print(model.summary())

# 5. Predict with uncertainty
y_pred, lower, upper = model.predict_interval(X_test, alpha=0.05)

# 6. Export
model.save("model.json")
latex_eq = model.to_latex()
predict_fn = model.to_callable()  # pure NumPy function

DOE Workflow

from jaxsr import DOEStudy

# Create study
study = DOEStudy("catalyst", ["T", "P", "flow"],
                 bounds=[(300, 500), (1, 10), (0.1, 2.0)])
X_design = study.create_design(method="latin_hypercube", n_points=20)
study.save("catalyst.jaxsr")

# After collecting data
study = DOEStudy.load("catalyst.jaxsr")
study.add_observations(X_measured, y_measured)
model = study.fit(max_terms=5)

# Get next experiments
next_pts = study.suggest_next(n_points=5, strategy="uncertainty")
study.save("catalyst.jaxsr")

Quick Reference: CLI

# Create study with factors
jaxsr init my_study -f "temp:300:500" -f "pressure:1:10" -f "catalyst:A,B,C"

# Generate experimental design → Excel template
jaxsr design my_study.jaxsr -m latin_hypercube -n 20 --format xlsx -o template.xlsx

# Import completed experiments
jaxsr add my_study.jaxsr completed.xlsx

# Fit model
jaxsr fit my_study.jaxsr --max-terms 5 --strategy greedy_forward --criterion bic

# Suggest next experiments
jaxsr suggest my_study.jaxsr -n 5 --strategy uncertainty

# Generate reports
jaxsr report my_study.jaxsr -o report.xlsx
jaxsr report my_study.jaxsr -o report.docx

# Check study status
jaxsr status my_study.jaxsr

Scikit-learn Integration

# JAXSR estimators work with sklearn's clone, cross_val_score, GridSearchCV, Pipeline
from sklearn.model_selection import cross_val_score, GridSearchCV
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

# Cross-validation
scores = cross_val_score(model, X, y, cv=5, scoring="r2")

# Hyperparameter tuning
grid = GridSearchCV(model, {"max_terms": [3, 5, 7]}, cv=3, scoring="r2", n_jobs=1)
grid.fit(X, y)

# Pipeline
pipe = Pipeline([("scaler", StandardScaler()), ("sr", model)])
pipe.fit(X, y)

# get_params / set_params protocol
model.get_params()                                # all constructor params
model.set_params(max_terms=7, strategy="exhaustive")  # modify params

# Nested params for MultiOutputSymbolicRegressor
mo.set_params(estimator__max_terms=3)

Key points:

• scikit-learn

  is an optional dependency:

  pip install jaxsr[sklearn]

• Always use

  n_jobs=1

  with sklearn tools — JAX handles parallelism internally

• BasisLibrary

  is not a sklearn transformer; configure it before creating the model

Decision Trees

"Which basis functions should I use?"

See

guides/basis-library.md

"Which selection strategy should I use?"

See

guides/model-fitting.md

"Which UQ method should I use?"

See

guides/uncertainty.md

"How do I set up a DOE study?"

See

guides/doe-workflow.md

"How do I add physical constraints?"

See

guides/constraints.md

"How do I use the CLI?"

See

guides/cli.md

"I already know the model form. How do I estimate parameters?"

See

guides/known-model-fitting.md

"How do I use Response Surface Methodology?"

See

guides/rsm.md

"How do I set up active learning?"

See

guides/active-learning.md

Templates

Ready-to-use scripts and notebook starters are in

templates/

basic-regression.py

constrained-model.py

doe-study.py

uncertainty-analysis.py

active-learning-loop.py

langmuir-isotherm.py

notebook-starter.py

Common Pitfalls

1. Library too large for exhaustive search. If you have > 20 basis functions, use

   greedy_forward

   instead of

   exhaustive

   .
2. Using

   information_criterion="cv"

   . Only

   "aic"

   ,

   "aicc"

   ,

   "bic"

   are supported.
3. Forgetting

   safe=True

   for transcendental. Without it,

   log(0)

   and

   1/0

   produce NaN.
4. Over-specifying the basis. A library with 500+ terms is slow and prone to overfitting. Start simple.
5. Not checking collinearity. Use

   from jaxsr.utils import check_collinearity

   before fitting if terms are nearly redundant.
6. Stale metrics after refit. After applying constraints, metrics are automatically recalculated. Do not copy metrics from a previous result.
7. Python control flow in JIT. If writing custom basis functions with

   @jit

   , use

   jnp.where

   instead of

   if/else

   .
8. Calling

   float()

   on JAX arrays inside JIT. Use

   .item()

   outside JIT or keep values as JAX arrays.
9. Using

   n_jobs=-1

   with sklearn tools. JAX handles parallelism internally. Always use

   n_jobs=1

   with

   GridSearchCV

   ,

   cross_val_score

   , etc.
10. Treating

    BasisLibrary

    as a sklearn transformer. It is not a transformer; configure it before creating the estimator. It cannot be included as a

    Pipeline

    step.