Merge pull request #65 from heal-research/feat/sample-weight

feat(sklearn): add sample_weight support to SymbolicRegressor.fit()

Author: foolnotion

README.md

@@ -25,6 +25,11 @@ reg.fit(X, y)
 ```
 Subclass `pyoperon.Callback` (`on_fit_begin`/`on_generation_end`/`on_fit_end`) for custom monitoring; return `True` from `on_generation_end` to request early termination. See `pyoperon/callback.py` for the full API.
 
+`fit()` also accepts `sample_weight=`, matching the scikit-learn convention, for training on non-uniformly sampled data (e.g. stratified/reweighted survey data, or down-weighting outliers) without distorting the fitted model:
+```python
+reg.fit(X, y, sample_weight=sample_weight)
+```
+
 The [example](https://github.com/heal-research/pyoperon/tree/main/example) folder contains sample code for using either the Python bindings directly or the **pyoperon.sklearn** module.
 
 # Installation

flake.lock

@@ -779,7 +779,7 @@ def get_model_string(
         names_map = {k: names[i] for i, k in enumerate(self.variables_)}
         return op.InfixFormatter.Format(model, names_map, precision)
 
-    def fit(self, X, y):
+    def fit(self, X, y, sample_weight=None):
         """Fit the symbolic regression model.
 
         Parameters
@@ -788,6 +788,22 @@ def fit(self, X, y):
             Training input samples.
         y : array-like of shape (n_samples,)
             Target values.
+        sample_weight : array-like of shape (n_samples,), default=None
+            Per-sample weights applied to the fitness/selection loss during
+            evolution, to coefficient optimization when `optimizer_iterations
+            > 0` and `optimizer_likelihood='gaussian'` (the default), and to
+            the post-hoc scale/intercept refit and the reported
+            `mean_squared_error` stat. `None` weights every sample equally.
+            Must be a proper 1D array of length `n_samples` - unlike some
+            scikit-learn estimators, a scalar or an (n_samples, 1) column
+            vector is not broadcast.
+
+            Poisson likelihoods (`optimizer_likelihood='poisson'` or
+            `'poisson_log'`) do not yet support sample_weight in coefficient
+            optimization: `w` there is already used for the exposure/offset
+            term, a different semantic from a precision weight, and
+            reconciling the two is unresolved. A warning is raised in that
+            case when `optimizer_iterations > 0`.
 
         Returns
         -------
@@ -814,10 +830,45 @@ def fit(self, X, y):
 
         X, y = check_X_y(X, y, accept_sparse=False)
         self.n_features_in_ = X.shape[1]
+        if sample_weight is not None:
+            # Public-API equivalent of sklearn's private _check_sample_weight:
+            # avoid depending on underscore-prefixed internals given
+            # pyproject's loose `scikit-learn>=1.6.1` pin. dtype=None just
+            # avoids needlessly forcing float32 - the nanobind SetWeights/
+            # FitLeastSquares bindings convert to whatever Operon::Scalar
+            # actually is regardless of the precision passed in here.
+            # Unlike _check_sample_weight, this does not broadcast a scalar
+            # or accept an (n_samples, 1) column vector - a real 1D array
+            # matching X's length is required, by choice, not omission.
+            sample_weight = check_array(
+                sample_weight, ensure_2d=False, dtype=None,
+                input_name='sample_weight',
+            )
+            if sample_weight.ndim != 1 or sample_weight.shape[0] != X.shape[0]:
+                raise ValueError(
+                    f'sample_weight has shape {sample_weight.shape}, '
+                    f'expected ({X.shape[0]},)'
+                )
+            if np.any(sample_weight < 0):
+                raise ValueError('sample_weight must be non-negative')
+            if not np.any(sample_weight > 0):
+                raise ValueError('sample_weight must not be all zero')
+            if optimizer_iterations > 0 and self.optimizer_likelihood in ('poisson', 'poisson_log'):
+                warnings.warn(
+                    'sample_weight is set but optimizer_iterations > 0 with '
+                    f'optimizer_likelihood={self.optimizer_likelihood!r}: '
+                    'coefficient optimization does not yet support '
+                    'sample_weight for Poisson likelihoods, so coefficients '
+                    'may be tuned against a different objective than the '
+                    'one used for selection.',
+                    stacklevel=2,
+                )
 
         # Build dataset and problem
         D                     = np.asfortranarray(np.column_stack((X, y)))
         ds                    = op.Dataset(D)
+        if sample_weight is not None:
+            ds.SetWeights(sample_weight)
         target                = max(ds.Variables, key=lambda x: x.Index)
         self.variables_       = {
             v.Hash: v.Name
@@ -970,7 +1021,10 @@ def report() -> bool:
 
         def get_solution_stats(solution):
             y_pred = op.Evaluate(dtable, solution.Genotype, ds, training_range)
-            scale, offset = op.FitLeastSquares(y_pred, y)
+            if sample_weight is None:
+                scale, offset = op.FitLeastSquares(y_pred, y)
+            else:
+                scale, offset = op.FitLeastSquares(y_pred, y, sample_weight)
             nodes = solution.Genotype.Nodes
 
             if not add_scale:
@@ -996,7 +1050,7 @@ def get_solution_stats(solution):
                 'tree': solution.Genotype,
                 'objective_values': evaluator(rng, solution),
                 'mean_squared_error': mean_squared_error(
-                    y, scale * y_pred + offset,
+                    y, scale * y_pred + offset, sample_weight=sample_weight,
                 ),
                 'minimum_description_length': mdl_eval(rng, solution)[0],
                 'bayesian_information_criterion': bic_eval(rng, solution)[0],

pyoperon/sklearn.py

@@ -104,6 +104,15 @@ void InitDataset(nb::module_ &m)
         .def("GetVariable", nb::overload_cast<const std::string&>(&Operon::Dataset::GetVariable, nb::const_))
         .def("GetVariable", nb::overload_cast<Operon::Hash>(&Operon::Dataset::GetVariable, nb::const_))
         .def_prop_ro("Variables", &Operon::Dataset::GetVariables)
+        .def("SetWeights", [](Operon::Dataset& self, std::vector<Operon::Scalar> const& w) { self.SetWeights(w); })
+        // Like GetValues/Values above, this is a view into weights_; a
+        // subsequent SetWeights() reallocates the backing vector, so
+        // re-fetch Weights after calling SetWeights rather than holding
+        // onto a stale view.
+        .def_prop_ro("Weights", [](Operon::Dataset const& self) -> nb::object {
+            auto w = self.Weights();
+            return w ? nb::cast(MakeView(*w, nb::find(self))) : nb::none();
+        })
         .def("Shuffle", &Operon::Dataset::Shuffle)
         .def("Normalize", &Operon::Dataset::Normalize)
         .def("Standardize", &Operon::Dataset::Standardize)

source/dataset.cpp

@@ -23,6 +23,15 @@ auto FitLeastSquares(nb::ndarray<T> lhs, nb::ndarray<T> rhs) -> std::pair<double
     return Operon::FitLeastSquares(s1, s2);
 }
 
+template<typename T>
+auto FitLeastSquares(nb::ndarray<T> lhs, nb::ndarray<T> rhs, nb::ndarray<T> weights) -> std::pair<double, double>
+{
+    auto s1 = MakeSpan(lhs);
+    auto s2 = MakeSpan(rhs);
+    auto s3 = MakeSpan(weights);
+    return Operon::FitLeastSquares(s1, s2, s3);
+}
+
 template<typename T>
 auto PoissonLikelihood(nb::ndarray<T> x, nb::ndarray<T> y) {
     return TPoissonLikelihood::ComputeLikelihood(MakeSpan(x), MakeSpan(y), {});
@@ -169,6 +178,14 @@ void InitEval(nb::module_ &m)
         return detail::FitLeastSquares<double>(lhs, rhs);
     });
 
+    m.def("FitLeastSquares", [](nb::ndarray<float> lhs, nb::ndarray<float> rhs, nb::ndarray<float> weights) -> std::pair<double, double> {
+        return detail::FitLeastSquares<float>(lhs, rhs, weights);
+    });
+
+    m.def("FitLeastSquares", [](nb::ndarray<double> lhs, nb::ndarray<double> rhs, nb::ndarray<double> weights) -> std::pair<double, double> {
+        return detail::FitLeastSquares<double>(lhs, rhs, weights);
+    });
+
     // dispatch table
     nb::class_<TDispatch>(m, "DispatchTable")
         .def(nb::init<>());