Add predict_proba output option to SklearnModelProcessor

src/ezmsg/learn/process/sklearn.py

@@ -1,6 +1,7 @@
 import importlib
 import pickle
 import typing
+from enum import Enum
 
 import ezmsg.core as ez
 import numpy as np
@@ -14,6 +15,11 @@
 from ezmsg.util.messages.util import replace
 
 
+class PredictMethod(str, Enum):
+    PREDICT = "predict"
+    PREDICT_PROBA = "predict_proba"
+
+
 class SklearnModelSettings(ez.Settings):
     model_class: str
     """
@@ -36,6 +42,18 @@ class SklearnModelSettings(ez.Settings):
     The full list of classes to use for partial_fit calls.
     This must be provided to use `partial_fit` with classifiers.
     """
+    predict_method: PredictMethod = PredictMethod.PREDICT
+    """
+    Inference method used to produce the output.
+
+    - ``"predict"`` (default): hard predictions, one value per sample. Output is
+      unchanged from earlier versions.
+    - ``"predict_proba"``: per-class posterior probabilities. The output ``ch``
+      axis has one column per class and is relabeled to the model's ``classes_``.
+      Requires a classifier exposing ``predict_proba`` (sklearn) or
+      ``predict_proba_many`` (River). Use this when a downstream consumer needs
+      the full posterior vector rather than a single hard label.
+    """
 
 
 @processor_state
@@ -87,6 +105,8 @@ def load_checkpoint(self, path: str) -> None:
             raise RuntimeError(f"Failed to load model from {path}: {str(e)}") from e
 
     def _reset_state(self, message: AxisArray) -> None:
+        # Validate: raises ValueError on an unknown method; plain strings are accepted.
+        PredictMethod(self.settings.predict_method)
         # Try loading from checkpoint first
         if self.settings.checkpoint_path:
             self.load_checkpoint(self.settings.checkpoint_path)
@@ -171,7 +191,18 @@ def _process(self, message: AxisArray) -> AxisArray | None:
             if expected != n_input:
                 raise ValueError(f"Model expects {expected} features, but got {n_input}")
 
-        if hasattr(self._state.model, "predict"):
+        if self.settings.predict_method == PredictMethod.PREDICT_PROBA:
+            if hasattr(self._state.model, "predict_proba"):
+                y_pred = np.asarray(self._state.model.predict_proba(X))
+            elif hasattr(self._state.model, "predict_proba_many"):  # River classifiers
+                df_X = pd.DataFrame(X, columns=[f"f{i}" for i in range(X.shape[1])])
+                y_pred = np.asarray(self._state.model.predict_proba_many(df_X))
+            else:
+                raise NotImplementedError(
+                    f"Model {type(self._state.model).__name__} does not support "
+                    "predict_proba (or River predict_proba_many)."
+                )
+        elif hasattr(self._state.model, "predict"):
             y_pred = self._state.model.predict(X)
         elif hasattr(self._state.model, "predict_many"):
             df_X = pd.DataFrame(X, columns=[f"f{i}" for i in range(X.shape[1])])
@@ -191,7 +222,14 @@ def _process(self, message: AxisArray) -> AxisArray | None:
         y_pred = y_pred.reshape(output_shape)
 
         if self._state.chan_ax is None:
-            self._state.chan_ax = AxisArray.CoordinateAxis(data=np.arange(output_shape[1]), dims=["ch"])
+            if self.settings.predict_method == PredictMethod.PREDICT_PROBA and (
+                getattr(self._state.model, "classes_", None) is not None
+                and len(self._state.model.classes_) == output_shape[1]
+            ):
+                chan_labels = np.asarray(self._state.model.classes_)
+            else:
+                chan_labels = np.arange(output_shape[1])
+            self._state.chan_ax = AxisArray.CoordinateAxis(data=chan_labels, dims=["ch"])
 
         return replace(
             message,

tests/unit/test_sklearn.py

@@ -225,3 +225,75 @@ def test_hmmlearn_gaussianhmm_predict(input_axisarray):
     assert output.data.shape[0] == input_axisarray.data.shape[0]
     assert output.data.ndim == 2
     assert output.data.shape[1] == 1
+
+
+@pytest.fixture
+def labels_two_class():
+    return np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1])
+
+
+@pytest.mark.parametrize(
+    "model_class",
+    [
+        "sklearn.discriminant_analysis.LinearDiscriminantAnalysis",
+        "sklearn.linear_model.LogisticRegression",
+    ],
+)
+def test_predict_proba_output(model_class, input_axisarray, labels_two_class):
+    proc = SklearnModelProcessor(model_class=model_class, predict_method="predict_proba")
+    proc._reset_state(input_axisarray)
+    proc.fit(input_axisarray.data, labels_two_class)
+
+    output = proc._process(input_axisarray)
+
+    n_classes = len(np.unique(labels_two_class))
+    assert output.data.shape == (input_axisarray.data.shape[0], n_classes)
+    assert np.all(output.data >= 0.0)
+    np.testing.assert_allclose(output.data.sum(axis=1), 1.0, rtol=1e-6, atol=1e-6)
+    # Output ch axis is relabeled to the model's classes_.
+    np.testing.assert_array_equal(output.axes["ch"].data, proc._state.model.classes_)
+
+
+def test_predict_proba_partial_fit_classifier(input_axisarray, labels_two_class):
+    proc = SklearnModelProcessor(
+        model_class="sklearn.linear_model.SGDClassifier",
+        model_kwargs={"loss": "log_loss"},
+        partial_fit_classes=np.array([0, 1]),
+        predict_method="predict_proba",
+    )
+    proc._reset_state(input_axisarray)
+    sample_msg = replace(
+        input_axisarray,
+        attrs={**input_axisarray.attrs, "trigger": SampleTriggerMessage(timestamp=0.0, value=labels_two_class)},
+    )
+    proc.partial_fit(sample_msg)
+
+    output = proc._process(input_axisarray)
+    assert output.data.shape == (input_axisarray.data.shape[0], 2)
+    np.testing.assert_allclose(output.data.sum(axis=1), 1.0, rtol=1e-6, atol=1e-6)
+
+
+def test_predict_method_default_unchanged(input_axisarray, labels_two_class):
+    # Default predict_method must match an explicit predict_method="predict".
+    proc_default = SklearnModelProcessor(model_class="sklearn.linear_model.LogisticRegression")
+    proc_default._reset_state(input_axisarray)
+    proc_default.fit(input_axisarray.data, labels_two_class)
+
+    proc_explicit = SklearnModelProcessor(
+        model_class="sklearn.linear_model.LogisticRegression", predict_method="predict"
+    )
+    proc_explicit._reset_state(input_axisarray)
+    proc_explicit.fit(input_axisarray.data, labels_two_class)
+
+    out_default = proc_default._process(input_axisarray)
+    out_explicit = proc_explicit._process(input_axisarray)
+    np.testing.assert_array_equal(out_default.data, out_explicit.data)
+    assert out_default.data.shape == (input_axisarray.data.shape[0], 1)
+
+
+def test_predict_proba_unsupported_model_raises(input_axisarray, labels_regression):
+    proc = SklearnModelProcessor(model_class="sklearn.linear_model.Ridge", predict_method="predict_proba")
+    proc._reset_state(input_axisarray)
+    proc.fit(input_axisarray.data, labels_regression)
+    with pytest.raises(NotImplementedError, match="predict_proba"):
+        proc._process(input_axisarray)