Merge pull request #8 from ezmsg-org/partial_fit

Directly use `partial_fit` instead of deprecated routing through `__call__`

Author: cboulay

docs/source/guides/classification.rst

@@ -125,7 +125,8 @@ For models that support ``partial_fit``, you can update them during streaming:
 .. code-block:: python
 
    from ezmsg.learn.process.sklearn import SklearnModelProcessor, SklearnModelSettings
-   from ezmsg.sigproc.sampler import SampleMessage
+   from ezmsg.baseproc import SampleTriggerMessage
+   from ezmsg.util.messages.util import replace
 
    # Create processor with online learning support
    processor = SklearnModelProcessor(
@@ -137,9 +138,9 @@ For models that support ``partial_fit``, you can update them during streaming:
    )
 
    # Training with labeled samples
-   sample_msg = SampleMessage(
-       sample=feature_array,  # AxisArray with features
-       trigger=label_value,   # The class label
+   sample_msg = replace(
+       feature_array,  # AxisArray with features
+       attrs={"trigger": SampleTriggerMessage(value=label_value)}
    )
    processor.partial_fit(sample_msg)
 

pyproject.toml

@@ -9,7 +9,7 @@ license = "MIT"
 requires-python = ">=3.10.15"
 dynamic = ["version"]
 dependencies = [
-    "ezmsg-baseproc>=1.3.0",
+    "ezmsg-baseproc>=1.4.0",
     "ezmsg-sigproc>=2.15.0",
     "river>=0.22.0",
     "scikit-learn>=1.6.0",

src/ezmsg/learn/process/ssr.py

@@ -265,11 +265,6 @@ def fit(self, X: np.ndarray) -> None:
         self._state.weights = self._solve_weights(self._state.cxx)
         self._on_weights_updated()
 
-    def fit_transform(self, message: AxisArray) -> AxisArray:
-        """Convenience: ``partial_fit`` then ``_process``."""
-        self.partial_fit(message)
-        return self._process(message)
-
     # -- abstract hooks for subclasses ---------------------------------------
 
     @abstractmethod

tests/unit/test_adaptive_linear_regressor.py

@@ -45,7 +45,7 @@ def test_adaptive_linear_regressor(model_type: str):
     samp = replace(sig_axarr, attrs={"trigger": samp_trig})
 
     proc = AdaptiveLinearRegressorTransformer(model_type=model_type)
-    _ = proc.send(samp)
-    preds = proc.send(replace(sig_axarr, data=X + np.random.randn(*X.shape)))
+    proc.partial_fit(samp)
+    preds = proc(replace(sig_axarr, data=X + np.random.randn(*X.shape)))
     assert isinstance(preds, AxisArray)
     assert preds.data.shape == (n_times, 1)

tests/unit/test_linear_regressor.py

@@ -42,9 +42,9 @@ def test_linear_regressor(model_type: str):
     )
     samp = replace(sig_axarr, attrs={"trigger": samp_trig})
 
-    gen = LinearRegressorTransformer(model_type=model_type)
-    _ = gen.send(samp)
-    preds = gen.send(replace(sig_axarr, data=X + 0.1 * np.random.randn(*X.shape)))
+    proc = LinearRegressorTransformer(model_type=model_type)
+    proc.partial_fit(samp)
+    preds = proc(replace(sig_axarr, data=X + 0.1 * np.random.randn(*X.shape)))
     rss = ((samp_trig.value.data - preds.data) ** 2).sum()
     tss = ((samp_trig.value.data - samp_trig.value.data.mean()) ** 2).sum()
     rsq = 1 - rss / tss

tests/unit/test_mlp_old.py

@@ -186,8 +186,7 @@ def xy_gen(set: int = 0):
         # Train: This is unrealistic in that we would normally do inference on many axisarray messages throughout
         #  the trial, and only do training infrequently at the end of a trial if we can infer the labels.
         #  But I'm too lazy to split the data into many small axarrs and one large SampleMessage per trial.
-        # Note: We don't have to call `partial_fit` because `__call__` inspects the message type and calls it for us.
-        proc(sample_msg)
+        proc.partial_fit(sample_msg)
 
     def eval_test(processor, set: int = 1):
         # Run the test inference

tests/unit/test_rnn.py

@@ -143,7 +143,7 @@ def test_rnn_partial_fit(simple_message):
         attrs={**simple_message.attrs, "trigger": SampleTriggerMessage(timestamp=0.0, value=target_value)},
     )
 
-    proc(sample_message)
+    proc.partial_fit(sample_message)
 
     assert not proc.state.model.training
     updated_weights = [p.detach() for p in proc.state.model.parameters()]

tests/unit/test_sgd.py

@@ -54,7 +54,7 @@ def test_sgd():
 
     # Now let's try training on all samples
     for sample in samples:
-        decoder(sample)  # Should invoke partial_fit
+        decoder.partial_fit(sample)
     # Then doing inference on all multi-wins
     probas = [decoder(win) for win in windows]
 
@@ -72,7 +72,7 @@ def test_sgd():
     )
     probas = []
     for samp_ix, samp in enumerate(samples):
-        decoder(samp)  # Should invoke partial_fit
+        decoder.partial_fit(samp)
         probas.append(decoder(windows[samp_ix * 2]))
         probas.append(decoder(windows[samp_ix * 2 + 1]))
     class_ids = []

tests/unit/test_ssr.py

@@ -230,15 +230,15 @@ def test_non_last_axis(self):
         assert out.data.shape == X.shape
 
 
-class TestFitTransform:
-    def test_fit_transform(self):
-        """fit_transform matches separate partial_fit + process."""
+class TestPartialFitTransform:
+    def test_partial_fit_transform(self):
+        """partial_fit_transform matches separate partial_fit + process."""
         rng = np.random.default_rng(10)
         X = _random_data(rng=rng)
         msg = _make_axisarray(X)
 
         proc1 = LRRTransformer(LRRSettings())
-        out1 = proc1.fit_transform(msg)
+        out1 = proc1.partial_fit_transform(msg)
 
         proc2 = LRRTransformer(LRRSettings())
         proc2.partial_fit(msg)