jaxsr 0.1.0__tar.gz

jaxsr-0.1.0/.claude/skills/jaxsr/SKILL.md

@@ -0,0 +1,311 @@
+# 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:
+
+| Scenario | Recommended Basis |
+|----------|-------------------|
+| Unknown relationship, few features | `add_constant + add_linear + add_polynomials(3) + add_interactions(2) + add_transcendental()` |
+| Known polynomial behavior | `add_constant + add_linear + add_polynomials(max_degree)` |
+| Engineering correlation (Nusselt, friction) | `add_constant + add_linear + add_polynomials(2) + add_transcendental(funcs=["log","exp","sqrt","inv"])` |
+| Chemical kinetics (rate laws) | `add_constant + add_linear + add_transcendental(funcs=["exp","inv","log"]) + add_ratios() + add_parametric(Arrhenius)` |
+| Large feature space (screening) | `add_constant + add_linear + add_interactions(2)` then use `lasso_path` strategy |
+| Response surface (DOE) | `add_constant + add_linear + add_polynomials(2) + add_interactions(2)` — or use `ResponseSurface` directly |
+| Categorical factors present | Add `add_categorical_indicators() + add_categorical_interactions()` to any of the above |
+
+**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
+
+| Data Size | Library Size | Recommended Strategy |
+|-----------|-------------|---------------------|
+| Any | < 20 basis functions | `exhaustive` (exact optimal) |
+| Any | 20-200 basis functions | `greedy_forward` (default, fast, reliable) |
+| Small n | Large library | `lasso_path` (regularized screening) |
+| Many terms expected | Any | `greedy_backward` (start full, prune) |
+
+**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
+
+| Scenario | Recommended Criterion |
+|----------|----------------------|
+| Small sample (n < 40) | `aicc` (corrected AIC, penalizes overfitting more) |
+| Medium sample (40 < n < 200) | `bic` (stronger complexity penalty, sparser models) |
+| Large sample (n > 200) | `aic` or `bic` (both work well) |
+| Want simplest model | `bic` (always penalizes complexity more) |
+| Want best prediction | `aicc` (balances fit and complexity) |
+
+**Default recommendation:** Use `"bic"` for interpretable models, `"aicc"` for predictive models.
+Only `"aic"`, `"aicc"`, and `"bic"` are supported — not `"cv"`.
+
+### Step 5: Recommend Constraints (if applicable)
+
+Ask: "Do you have any physical knowledge about the system?"
+
+| Physical Knowledge | Constraint to Add |
+|-------------------|-------------------|
+| Output must be positive | `.add_bounds("y", lower=0)` |
+| Output in known range | `.add_bounds("y", lower=lo, upper=hi)` |
+| Increasing in temperature | `.add_monotonic("T", direction="increasing")` |
+| Diminishing returns | `.add_concave(feature)` |
+| Accelerating growth | `.add_convex(feature)` |
+| Coefficient must be positive | `.add_sign_constraint(basis_name, sign="positive")` |
+| Known intercept or slope | `.add_known_coefficient(name, value)` |
+
+Use `hard=True` for strict enforcement; `hard=False` (default) for soft penalty.
+
+### Step 6: Recommend Uncertainty Quantification
+
+| Need | Method | When to Use |
+|------|--------|-------------|
+| Quick confidence intervals | `model.predict_interval()` | Default. OLS-based. Assumes normality. |
+| Coefficient significance | `model.coefficient_intervals()` | Check which terms are statistically significant |
+| Robust to model uncertainty | `model.predict_bma()` | Averages over Pareto-front models weighted by criterion |
+| No distributional assumptions | `model.predict_conformal()` | Distribution-free. Needs enough data (n > 30). |
+| Assess model stability | `bootstrap_predict()` | Resamples data. Shows sensitivity to individual points. |
+| Compare model structures | `model.predict_ensemble()` | Returns predictions from all Pareto-front models |
+| Variable importance | `anova()` | Decomposes variance by term. Shows which factors matter. |
+
+**Default recommendation:** Start with `predict_interval()` (built-in, fast). Add `predict_bma()` if you have multiple competing models. Use `predict_conformal()` for publication-quality intervals.
+
+### Step 7: Recommend Reporting Format
+
+| Goal | Action |
+|------|--------|
+| Quick look at results | `model.summary()` or `jaxsr status study.jaxsr` |
+| Share with collaborators | `jaxsr report study.jaxsr -o report.xlsx` (Excel) |
+| Formal report | `jaxsr report study.jaxsr -o report.docx` (Word) |
+| Paper/presentation | `model.to_latex()` for equation, `plot_pareto_front()` for figures |
+| Deploy model | `model.to_callable()` (pure NumPy, no JAX dependency) |
+| Archive/reproduce | `model.save("model.json")` and `study.save("study.jaxsr")` |
+
+---
+
+## Quick Reference: Installation
+
+```bash
+# Core library
+pip install jaxsr
+
+# 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)
+
+```python
+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
+
+```python
+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
+
+```python
+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
+
+```bash
+# 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
+```
+
+## Decision Trees
+
+### "Which basis functions should I use?"
+
+See `guides/basis-library.md` for the complete decision guide.
+
+### "Which selection strategy should I use?"
+
+See `guides/model-fitting.md` for strategy comparison and benchmarks.
+
+### "Which UQ method should I use?"
+
+See `guides/uncertainty.md` for method comparison and selection flowchart.
+
+### "How do I set up a DOE study?"
+
+See `guides/doe-workflow.md` for the complete lifecycle guide.
+
+### "How do I add physical constraints?"
+
+See `guides/constraints.md` for constraint types and examples.
+
+### "How do I use the CLI?"
+
+See `guides/cli.md` for full CLI reference with examples.
+
+### "I already know the model form. How do I estimate parameters?"
+
+See `guides/known-model-fitting.md` for a worked example using the Langmuir isotherm,
+including parametric basis functions, experiment design, ANOVA, and uncertainty analysis.
+Generalizes to Arrhenius, Michaelis-Menten, power laws, and other known models.
+
+### "How do I use Response Surface Methodology?"
+
+See `guides/rsm.md` for RSM designs, canonical analysis, and optimization.
+
+### "How do I set up active learning?"
+
+See `guides/active-learning.md` for acquisition functions and adaptive sampling.
+
+## Templates
+
+Ready-to-use scripts and notebook starters are in `templates/`:
+
+| Template | Use Case |
+|----------|----------|
+| `basic-regression.py` | Discover an equation from X, y data |
+| `constrained-model.py` | Add physical constraints to model |
+| `doe-study.py` | Full DOE workflow from design to report |
+| `uncertainty-analysis.py` | Compare all UQ methods |
+| `active-learning-loop.py` | Iterative experiment-model loop |
+| `langmuir-isotherm.py` | Known-model parameter estimation (Langmuir) |
+| `notebook-starter.py` | Jupyter notebook cell structure |
+
+## 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.

jaxsr-0.1.0/.claude/skills/jaxsr/guides/active-learning.md

@@ -0,0 +1,267 @@
+# Active Learning & Adaptive Sampling Guide
+
+## What is Active Learning?
+
+Instead of running all experiments upfront, active learning iteratively:
+1. Fit a model to current data
+2. Use the model to suggest the most informative next experiments
+3. Run those experiments
+4. Repeat until the model is good enough
+
+This can dramatically reduce the number of experiments needed.
+
+## AdaptiveSampler: Simple Interface
+
+```python
+from jaxsr import AdaptiveSampler
+
+sampler = AdaptiveSampler(
+    model=model,  # a fitted SymbolicRegressor
+    bounds=[(300, 500), (1, 10)],  # Feature bounds
+    strategy="uncertainty",
+    batch_size=5,
+    n_candidates=1000,
+    random_state=42
+)
+
+# Get suggestions
+result = sampler.suggest(n_points=5)
+X_next = result.points    # shape (5, n_features)
+scores = result.scores    # acquisition function values
+```
+
+## Suggestion Strategies
+
+### `uncertainty`
+
+Suggests points where the model is most uncertain (largest prediction intervals).
+
+**Best for:** General exploration, reducing model error everywhere.
+
+### `error`
+
+Suggests points where the model prediction error is expected to be highest.
+Uses ensemble disagreement or residual extrapolation.
+
+**Best for:** Improving accuracy in poorly-modeled regions.
+
+### `leverage`
+
+Suggests points with high statistical leverage — positions that would have the
+most influence on the fitted coefficients.
+
+**Best for:** Stabilizing coefficient estimates, optimal experimental design.
+
+### `gradient`
+
+Suggests points where the model gradient is steepest — regions of rapid change.
+
+**Best for:** Resolving sharp transitions, nonlinear behavior.
+
+### `space_filling`
+
+Suggests points that fill gaps in the existing design.
+Uses maximin distance criterion.
+
+**Best for:** Initial exploration, uniform coverage, no model required.
+
+### `random`
+
+Uniform random sampling within bounds.
+
+**Best for:** Baseline comparison, sanity check.
+
+## Strategy Selection Guide
+
+```
+Do you have a fitted model?
+├── NO → Use "space_filling" or "random"
+└── YES
+    ├── Want to reduce overall uncertainty?
+    │   └── Use "uncertainty"
+    ├── Want to improve accuracy in worst regions?
+    │   └── Use "error"
+    ├── Want to stabilize coefficients?
+    │   └── Use "leverage"
+    └── Want to resolve sharp features?
+        └── Use "gradient"
+```
+
+## Acquisition Functions: Advanced Interface
+
+An **acquisition function** scores candidate points by how useful they would be as
+the next experiment. Higher scores mean more informative. For fine-grained control,
+use acquisition functions directly:
+
+```python
+from jaxsr.acquisition import (
+    ActiveLearner,
+    PredictionVariance,
+    UCB, LCB,
+    ExpectedImprovement,
+    EnsembleDisagreement,
+    BMAUncertainty,
+    AOptimal, DOptimal,
+    Composite,
+    suggest_points
+)
+```
+
+### Prediction Variance
+
+```python
+acq = PredictionVariance()
+learner = ActiveLearner(model, bounds=[(0, 1), (0, 1)], acquisition=acq)
+result = learner.suggest(n_points=5)
+```
+
+### Upper/Lower Confidence Bound
+
+```python
+# Explore high-response regions (maximize)
+acq = UCB(kappa=2.0)     # mean + kappa * std
+
+# Explore low-response regions (minimize)
+acq = LCB(kappa=2.0)     # mean - kappa * std
+```
+
+`kappa` controls exploration vs. exploitation:
+- High kappa (> 2) → more exploration (wider search)
+- Low kappa (< 1) → more exploitation (stay near known good areas)
+
+### Expected Improvement
+
+```python
+acq = ExpectedImprovement()  # For finding the maximum of the response
+```
+
+Classic Bayesian optimization acquisition function. Balances exploitation
+(predicted value) and exploration (uncertainty).
+
+### Ensemble Disagreement
+
+```python
+acq = EnsembleDisagreement()  # Where Pareto-front models disagree most
+```
+
+### Optimal Design Criteria
+
+```python
+acq = AOptimal()    # Minimize average variance of coefficients
+acq = DOptimal()    # Minimize volume of coefficient confidence ellipsoid
+```
+
+### Composite Acquisition Functions
+
+Combine multiple acquisition functions with weights:
+
+```python
+# 70% uncertainty + 30% space-filling
+combined = 0.7 * PredictionVariance() + 0.3 * AOptimal()
+
+learner = ActiveLearner(model, bounds=bounds, acquisition=combined)
+result = learner.suggest(n_points=5)
+```
+
+You can also use the `Composite` class:
+```python
+combined = Composite(functions=[(0.7, PredictionVariance()), (0.3, AOptimal())])
+```
+
+## Complete Active Learning Loop
+
+```python
+import numpy as np
+from jaxsr import BasisLibrary, SymbolicRegressor, AdaptiveSampler
+
+# Define your experiment (replace with real measurements)
+def true_function(X):
+    """Placeholder — replace with your actual experiment or simulation."""
+    return 3.0 * X[:, 0] + 0.5 * X[:, 1] - 0.01 * X[:, 0]**2
+
+# Initial data (small)
+X_init = np.random.uniform([300, 1], [500, 10], size=(10, 2))
+y_init = true_function(X_init)
+
+# Build library
+library = (BasisLibrary(n_features=2, feature_names=["T", "P"])
+    .add_constant().add_linear()
+    .add_polynomials(max_degree=2)
+    .add_interactions(max_order=2)
+)
+
+# Fit initial model
+model = SymbolicRegressor(basis_library=library, max_terms=5)
+model.fit(X_init, y_init)
+
+X_all, y_all = X_init.copy(), y_init.copy()
+bounds = [(300, 500), (1, 10)]
+
+for iteration in range(5):
+    print(f"\n--- Iteration {iteration + 1} ---")
+    print(f"Model: {model.expression_}")
+    print(f"R²: {model.metrics_['r2']:.4f}")
+
+    # Suggest next experiments
+    sampler = AdaptiveSampler(model, bounds, strategy="uncertainty", batch_size=5)
+    result = sampler.suggest(n_points=5)
+    X_next = result.points
+
+    # Run experiments (replace with actual experiments)
+    y_next = true_function(X_next)
+
+    # Update model
+    X_all = np.vstack([X_all, X_next])
+    y_all = np.concatenate([y_all, y_next])
+    model.fit(X_all, y_all)
+
+print(f"\nFinal model ({len(y_all)} points): {model.expression_}")
+```
+
+## DOE Integration
+
+The `DOEStudy` class wraps active learning into the study workflow:
+
+```python
+from jaxsr import DOEStudy
+
+study = DOEStudy.load("my_study.jaxsr")
+
+# Suggest next experiments using the fitted model
+next_pts = study.suggest_next(n_points=5, strategy="uncertainty")
+
+# Or via CLI
+# jaxsr suggest my_study.jaxsr -n 5 --strategy uncertainty
+```
+
+## Discrete/Integer Variables
+
+If some features take only discrete values:
+
+```python
+sampler = AdaptiveSampler(
+    model, bounds,
+    strategy="uncertainty",
+    discrete_dims={1: [1, 2, 3, 4, 5]}  # Feature index 1: valid discrete values
+)
+```
+
+## Excluding Regions
+
+Avoid suggesting points too close to existing data:
+
+```python
+result = sampler.suggest(
+    n_points=5,
+    exclude_points=X_existing,    # Don't suggest near these
+    min_distance=0.01             # Minimum distance threshold
+)
+```
+
+## Batch vs. Sequential
+
+- **Sequential** (n_points=1): Optimal but requires one experiment at a time.
+- **Batch** (n_points=5+): Less optimal but allows parallel experiments.
+
+For batch suggestions, points are selected greedily to avoid clustering:
+each new point considers previously selected points in the batch.