Implement sample_weight instead of silently ignoring it

[jkitchin]

Closes #19.

SymbolicRegressor.fit() accepted sample_weight, documented it as "not yet implemented", and dropped it on the floor. A user passing measurement variances got a plausible answer computed from the assumption they believed they had overridden.

The reproducer from the issue now gives:

no weights  : y = 25.28 + 2.074*x^3
with weights: y = 2.989*x        # identical to fitting the clean half alone

Mechanism

WLS on (Phi, y) is OLS on (sqrt(w)·Phi, sqrt(w)·y), so the weights are applied by whitening at each entry point rather than by growing a parallel weighted code path. The Gram-matrix fast paths, the LASSO screening, the SVD interval machinery and the closed-form MSE all stay exact and untouched.

Two places deliberately keep the raw matrix:

• constraint evaluation — a monotonicity or convexity constraint is a property of the fitted function over the design space, not of how much a given row was trusted, so a down-weighted row still has to obey it;
• Phi_new in prediction intervals — so the interval is for a new observation of unit weight (one as precise as an average training point).

Where weights now apply

Everything the issue listed, plus a few things that would otherwise have been inconsistent in a way that is hard to see:

• all four selection strategies, so weights steer which terms are chosen, not only their coefficients
• the reported MSE, R², and the AIC/BIC/AICc computed from them
• constraint refitting (fit_constrained_ols), negligible-term pruning, and the profile-likelihood optimisation of parametric basis parameters
• sigma_, covariance_matrix_, coefficient_intervals(), predict_interval(), confidence_band(), anova(), the bootstrap functions, cross_validate(), jackknife+ conformal prediction
• the active-learning posterior in acquisition.py, which was scoring candidates against the unweighted (PhiᵀPhi)⁻¹ — it would have proposed the next experiment as if the down-weighted rows had been fully informative. The kriging-believer batch path also grew _X_train without the matching weights.

The parametric case is a good illustration of why partial coverage is not enough: unweighted, the test fixture collapses to a constant; weighted, it recovers exp(-2*x).

Effective-sample-size policy

Weights are normalised to average 1, so only their ratios matter and the fit is invariant to their overall scale — w, 2*w and w/1000 give the same model, MSE and criteria. The n used by the information criteria stays the nominal sample count: weighting describes how much each measurement is trusted, it does not create or destroy measurements. That is precisely why duplicating rows is not an equivalent trick — it inflates n and shifts every criterion.

New SymbolicRegressor.effective_sample_size_ reports the Kish ESS as a diagnostic for when the nominal n overstates the evidence (100.0 for the 200-point reproducer). sample_weight_ exposes the normalised weights.

Invalid weights (wrong length, negative, non-finite, all-zero) now raise ValueError rather than being ignored.

API surface

sample_weight was also added to fit_symbolic(), MultiOutputSymbolicRegressor.fit(), SymbolicRegressor.score(), fit_ols(), fit_ridge(), select_features() and the four strategy functions, cross_validate(), compute_mse/rmse/mae/r2/adjusted_r2/mape/all_metrics(), compute_cv_score(), compute_loo_mse(), compute_press(), and bootstrap_model_selection(). SymbolicRegressor.update() gained sample_weight_new.

Deliberately not weighted

• conformal_predict_split() — its coverage comes from exchangeability of the user-supplied calibration residuals, not from the fit. Documented, with a pointer to method="jackknife+" for a weight-aware interval.
• SymbolicClassifier — it never accepted the argument, so there is no silent-ignore bug there.
• max_error in compute_all_metrics() — a maximum has no weighted analogue; it is now taken over the rows with non-zero weight instead of over rows that carry none.

Docs

New docs/guides/sample-weights.md and the matching skill guide (src/jaxsr/skill/ re-synced from .claude/skills/jaxsr/), covering weight semantics, the effective-sample-size policy, recipes for variance-derived and replicate weights, and what is left unweighted. Includes the smooth-taper recipe w = y_x² / (y_x² + eps²) for the superposition case in #14, in place of a hard |y_x| threshold.

Verification

• 674 tests pass on JAX and on the NumPy/Pyodide backend (scripts/test_under_numpy.py)
• black --check and ruff check clean
• coverage 63.8% (threshold 60%)
• 53 new tests in tests/test_sample_weight.py, built around the invariants that make weights trustworthy: uniform weights change nothing, scale invariance, a zero weight equals dropping the row, and integer weights match row duplication in the coefficients while n correctly does not move.

🤖 Generated with Claude Code

https://claude.ai/code/session_01JigoMqmadVXVs7ALgnLnvG

---

Generated by Claude Code