Repository: neuroinformatics-unit/spikewrap
Pull request #205: [Fix-200] Expose whiten function
Author: login: chandrikarj
Created: 2025-03-11T14:47:14Z | Merged: 2025-03-21T00:31:41Z


## Description

**What is this PR**

- [ ] Bug fix
- [x] Addition of a new feature
- [ ] Other

**Why is this PR needed?**
The whiten function from SpikeInterface preprocessing was not exposed in Spikewrap. This PR adds it to the exposed functions, ensuring it can be used in preprocessing pipelines.

**What does this PR do?**
Exposes the whiten function in Spikewrap
Adds tests to verify whiten
Updates the documentation

## References

#200 

## Is this a breaking change?

No

## Does this PR require an update to the documentation?
Yes, documentation has been updated to:
Include whiten as an available preprocessing step.

## Checklist:

- [x] The code has been tested locally
- [x] Tests have been added to cover all new functionality
- [x] The documentation has been updated to reflect any changes
- [x] The code has been formatted with [pre-commit](https://pre-commit.com/)




Diff:
diff --git a/spikewrap/process/_preprocessing.py b/spikewrap/process/_preprocessing.py
index 5567b58b..7adc05d8 100644
--- a/spikewrap/process/_preprocessing.py
+++ b/spikewrap/process/_preprocessing.py
@@ -121,6 +121,7 @@ def _get_pp_funcs() -> dict[str, Callable]:
         "phase_shift": si.phase_shift,
         "bandpass_filter": si.bandpass_filter,
         "common_reference": si.common_reference,
+        "whiten": si.whiten,
         "remove_bad_channels": remove_bad_channels,
         "interpolate_bad_channels": interpolate_bad_channels,
         "remove_channels": remove_channels,
diff --git a/tests/test_integration/test_whiten.py b/tests/test_integration/test_whiten.py
new file mode 100644
index 00000000..52e0d925
--- /dev/null
+++ b/tests/test_integration/test_whiten.py
@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+from base import BaseTest
+from spikeinterface.preprocessing import whiten
+
+import spikewrap as sw
+
+
+class TestWhiten(BaseTest):
+
+    @pytest.fixture(scope="function")
+    def session(self, tmp_path):
+        """
+        Sets up a session fixture to load raw spike data.
+        This fixture ensures that each test runs in isolation with temporary data.
+        """
+        session = sw.Session(
+            subject_path=sw.get_example_data_path("spikeglx") / "rawdata" / "sub-001",
+            session_name="ses-001",
+            file_format="spikeglx",
+            run_names=["run-001_g0_imec0", "run-002_g0_imec0"],
+            output_path=tmp_path,
+        )
+        session.load_raw_data()
+        return session
+
+    def test_whitening_consistency(self, session):
+        """
+        Compare the output of applying SpikeInterface's `whiten()` directly on the raw recording
+        with the output obtained after running Spikewrap’s whitening preprocessing.
+
+        This ensures that Spikewrap's implementation of whitening is consistent with the
+        standard whitening approach from SpikeInterface.
+        """
+
+        # Whitening configuration used for preprocessing
+        whitening_config = {
+            "1": [
+                "whiten",
+                {
+                    "apply_mean": True,  # Subtract mean before whitening
+                    "mode": "global",  # Apply whitening across all channels
+                    "dtype": "float32",  # Ensure floating-point precision
+                    "int_scale": 1.0,  # Internal scaling factor
+                    "chunk_size": 1000,  # Process data in chunks of 500 samples
+                    "num_chunks_per_segment": 150,  # Number of chunks per segment
+                    "seed": 42,  # Seed
+                    "regularize": True,  # Apply regularisation
+                },
+            ]
+        }
+
+        # Apply whitening using Spikewrap preprocessing
+        session.preprocess(whitening_config)
+        spikewrap_whitened = session._pp_runs[0]
+
+        # Extract whitened data and transformation matrix from Spikewrap
+        underlying_recording = spikewrap_whitened._preprocessed["grouped"][
+            "1-raw-whiten"
+        ]
+        sw_data = underlying_recording.get_traces()  # Get the whitened traces
+        W_sw = np.array(
+            spikewrap_whitened._preprocessed["grouped"]["1-raw-whiten"]._kwargs["W"]
+        )  # Whitening matrix
+
+        # Apply whitening using SpikeInterface directly on raw recording
+        raw_recording = session._raw_runs[0]._raw["grouped"]
+        whitening_kwargs = whitening_config["1"][1]
+        si_whitened = whiten(
+            raw_recording, **whitening_kwargs
+        )  # Apply SpikeInterface whitening
+        si_data = si_whitened.get_traces()  # Get whitened traces
+        W_si = np.array(si_whitened._kwargs["W"])  # Extract whitening matrix
+
+        # Verify consistency between Spikewrap and SpikeInterface results
+        assert (
+            si_data.shape == sw_data.shape
+        ), "Shape mismatch between SpikeInterface and Spikewrap outputs"
+        assert np.allclose(
+            si_data, sw_data, atol=1e-5
+        ), "Data mismatch: whitened traces differ"
+        assert np.allclose(
+            W_si, W_sw, atol=1e-5
+        ), f"Whitening matrix mismatch: difference norm {np.linalg.norm(W_si - W_sw):.4f}"


Reviews:
user: login: JoeZiminski
body: Hey @chandrikarj thanks a lot for this! Please see comments on the code. For the test, that is a nice implementation to check the whitening itself, however in this case we only really need to check the preprocessing through spikewrap matches the preprocessing through spikeinterface (the validity of the spikeinterface preprocessing is checked on spikeinterface tests). Please let me know if you have any questions!
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-03-11T20:57:22Z
user: login: JoeZiminski
body: Brilliant, thanks @chandrikarj !
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-21T00:31:22Z

Review comments:
diff_hunk: @@ -46,9 +46,10 @@
 if __name__ == "__main__":
 
     pp_steps = {
-        "1": ["phase_shift", {}],  
-        "2": ["bandpass_filter", {"freq_min": 300, "freq_max": 6000}],
-        "3": ["common_reference", {"operator": "median", "reference": "global"}],
+    "1": ["phase_shift", {}],
+    "2": ["bandpass_filter", {"freq_min": 300, "freq_max": 6000}],
+    "3": ["common_reference", {"operator": "median"}],
+    "4": ["whiten", {"apply_mean": True, "mode": "global", "dtype": "float32"}]  # ✅ Added whiten properly
path: docs/source/galleries/how_to/01_preprocess_a_session.py
user: login: JoeZiminski
body: In this case, we will not highlight the whitening in the example scripts. The reason for this is quite obscure, basically many sorters run the whitening step internally. We could however run it on the spikeinterface side and turn it off in the sorter, although it is not clear (to me at least) how much this might effect the results. In theory spikeinterface whiten should be the same as the kilosort whiten but it is not easy to test (against the matlab versions). For ease of example it makes sense to show only some basic preprocessing steps here.
created_at: 2025-03-11T20:40:46Z
author_association: MEMBER
diff_hunk: @@ -180,10 +180,11 @@
 
 
 config_dict = {
-    "preprocessing": {
+       "preprocessing": {
         "1": ["phase_shift", {}],
         "2": ["bandpass_filter", {"freq_min": 300, "freq_max": 6000}],
-        "3": ["common_reference", {"operator": "median"}]
+        "3": ["common_reference", {"operator": "median"}],
+        "4": ["whiten", {"apply_mean": True, "mode": "global", "dtype": "float32"}]  # ✅ Added whiten properly
path: docs/source/galleries/tutorials/04_spike_sorting.py
user: login: JoeZiminski
body: same as above for the whiten step (and the other examples here)
created_at: 2025-03-11T20:41:03Z
author_association: MEMBER
diff_hunk: @@ -197,9 +198,11 @@
 sorting_dict = {"sorting": {"mountainsort5": {"filter": False}}}
 
 session.sort(
-    configs=sorting_dict, run_sorter_method="local"
+    configs={"mountainsort5": {"filter": False}},  
path: docs/source/galleries/tutorials/04_spike_sorting.py
user: login: JoeZiminski
body: I think it's okay as it was as `sorting_dict` is defined on line 198. It could go either way, but is slightly more readable to define it first as a variable.
created_at: 2025-03-11T20:42:12Z
author_association: MEMBER
diff_hunk: @@ -32,21 +33,19 @@ def _preprocess_recording(
     checked_pp_steps, pp_step_names = _check_and_sort_pp_steps(pp_steps, pp_funcs)
 
     for step_num, pp_info in checked_pp_steps.items():
-
         pp_name, pp_options = pp_info
 
-        last_pp_step_output, _ = _utils._get_dict_value_from_step_num(
+        last_pp_step_output, __ = _utils._get_dict_value_from_step_num(
             preprocess_data, step_num=str(int(step_num) - 1)
         )
 
         preprocessed_recording = pp_funcs[pp_name](last_pp_step_output, **pp_options)
 
         new_name = f"{step_num}-" + "-".join(["raw"] + pp_step_names[: int(step_num)])
+        # print(f"Step {step_num}: {pp_name} output type ->", type(preprocessed_recording))
 
         preprocess_data[new_name] = preprocessed_recording
 
-    return preprocess_data
path: spikewrap/process/_preprocessing.py
user: login: JoeZiminski
body: I think these changes might be because a branch was merged while this PR was open (let me know if this is not the case).

If so, you can do an interactive rebase (see [this guide](https://hackernoon.com/beginners-guide-to-interactive-rebasing-346a3f9c3a6d) for detailed information) on running and handling conflicts. Let me know if you have any questions!

```
git checkout main
git pull
git checkout expose-whiten-func
git rebase -i main
```

Nonetheless the way this function is used on [`main`](https://github.com/neuroinformatics-unit/spikewrap/blob/d1db4277b8114c0dea5932ea46f48286b08f0b00/spikewrap/process/_preprocessing.py#L11) is all wrong, because it is editing `preprocess_dict` in place and then returning it. This can be fixed in another issue / PR! A mistake from a refactoring that wasn't properly tidied.

created_at: 2025-03-11T20:50:28Z
author_association: MEMBER
diff_hunk: @@ -120,6 +119,12 @@ def _get_pp_funcs() -> dict[str, Callable]:
         "phase_shift": si.phase_shift,
         "bandpass_filter": si.bandpass_filter,
         "common_reference": si.common_reference,
+        "whiten": si.whiten,  # Add whiten function
     }
 
     return pp_funcs
+
+
+def _preprocess_recording(recording, params=None):
path: spikewrap/process/_preprocessing.py
user: login: JoeZiminski
body: I think this was needed due to the change of function name above. This function should be preprocessing the passed dictionary but has disappeared with the function rename.
created_at: 2025-03-11T20:51:16Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,77 @@
+import numpy as np
+from spikeinterface.core import BaseRecording, NumpyRecording
+
+from spikewrap.process._preprocessing import _get_pp_funcs
+
+
+def test_whiten():
path: tests/test_integration/test_preprocessing.py
user: login: JoeZiminski
body: This is a good start at the test, I think it will be easier to directly test whitening in spikewrap against the spikeinterface (the spikeinterface implementation is tested on the spikeinterface side). The idea would be to use the example kilosort data (like in the spikewrap example script), and run whiten only on the runs. Then the spikeinterface `whiten` function 

```
from spikeinterface.preprocessing import whiten
```

can be used on the recording objects in `session._raw`. and compared with the recording objects in `session.pp_runs`. Let me know if you have any questions about this! It would be similar to [this test](https://github.com/neuroinformatics-unit/spikewrap/blob/d1db4277b8114c0dea5932ea46f48286b08f0b00/tests/test_integration/test_sync.py#L26) but would also run `session.preprocess({<configs to run just whitening>})` and then compare the spikeinterface `whiten(session._raw_runs[0]["grouped"])` recording to the recording is `session._pp_runs[0]`. 

 


created_at: 2025-03-11T20:56:03Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+from base import BaseTest
+from spikeinterface.preprocessing import whiten
+
+import spikewrap as sw
+
+
+class TestWhiten(BaseTest):
+
+    @pytest.fixture(scope="function")
+    def session(self, tmp_path):
+        session = sw.Session(
+            subject_path=sw.get_example_data_path("spikeglx") / "rawdata" / "sub-001",
+            session_name="ses-001",
+            file_format="spikeglx",
+            run_names=["run-001_g0_imec0", "run-002_g0_imec0"],
+            output_path=tmp_path,
+        )
+        session.load_raw_data()
+        return session
+
+    def test_whitening_consistency(self, session):
+        """
+        Compare the output of applying SpikeInterface's whiten() directly on the raw recording
+        with the output obtained after running Spikewrap’s whitening preprocessing.
+        """
+        # Apply whitening using preprocessing
path: tests/test_integration/test_whiten.py
user: login: JoeZiminski
body: I would move this comment under the configs and change it to include a reference to spikewrap e.g:

```
# Apply whitening in spikewrap
```

And then above the configs I would make a comment on their selection (similar to as you have below with `# Set to True if setup is computationally efficient, use regularize_kwargs with caution` which I think can be deleted). Maybe something minimal like `Choose a set of kwargs (that are not all spikeinterface defaults) to pass to whiten functions`

created_at: 2025-03-19T12:47:29Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+from base import BaseTest
+from spikeinterface.preprocessing import whiten
+
+import spikewrap as sw
+
+
+class TestWhiten(BaseTest):
+
+    @pytest.fixture(scope="function")
+    def session(self, tmp_path):
+        session = sw.Session(
+            subject_path=sw.get_example_data_path("spikeglx") / "rawdata" / "sub-001",
+            session_name="ses-001",
+            file_format="spikeglx",
+            run_names=["run-001_g0_imec0", "run-002_g0_imec0"],
+            output_path=tmp_path,
+        )
+        session.load_raw_data()
+        return session
+
+    def test_whitening_consistency(self, session):
+        """
+        Compare the output of applying SpikeInterface's whiten() directly on the raw recording
+        with the output obtained after running Spikewrap’s whitening preprocessing.
+        """
+        # Apply whitening using preprocessing
+        whitening_config = {
+            "1": [
+                "whiten",
+                {
+                    "apply_mean": False,
+                    "mode": "global",
+                    "dtype": "float32",
+                    "int_scale": 1.0,
+                    "chunk_size": 1000,
+                    "num_chunks_per_segment": 150,
+                    "seed": 42,
+                    # Set to True if setup is computationally efficient, use regularize_kwargs with caution
+                    "regularize": False,
+                },
+            ]
+        }
+        session.preprocess(whitening_config)
+        spikewrap_whitened = session._pp_runs[0]
+
+        underlying_recording = spikewrap_whitened._preprocessed["grouped"][
+            "1-raw-whiten"
+        ]
+        sw_data = underlying_recording.get_traces()
+        W_sw = np.array(
+            spikewrap_whitened._preprocessed["grouped"]["1-raw-whiten"]._kwargs["W"]
+        )
+
+        # Apply whitening using SpikeInterface
+        raw_recording = session._raw_runs[0]._raw["grouped"]
+        whitening_kwargs = whitening_config["1"][1]
+        si_whitened = whiten(raw_recording, **whitening_kwargs)
+        si_data = si_whitened.get_traces()
+        W_si = np.array(si_whitened._kwargs["W"])
+
+        assert si_data.shape == sw_data.shape, "Shape mismatch"
path: tests/test_integration/test_whiten.py
user: login: JoeZiminski
body: Maybe final comment here like `#Compare spikewrap vs. spikeinterface outputs`
created_at: 2025-03-19T12:49:12Z
author_association: MEMBER

PR comments:
user: login: chandrikarj
created_at: 2025-03-11T14:48:16Z
body: Hey @JoeZiminski,
I've implemented the whiten function and updated the docs/configs accordingly. However, the docs are getting stuck without showing errors. I am trying some workarounds, but can't confirm the docs building. Could you check this once? Thanks.
author_association: CONTRIBUTOR
user: login: chandrikarj
created_at: 2025-03-12T11:33:59Z
body: @JoeZiminski, I will get to making the changes as soon as possible, and will ask if I can't resolve conflicts or otherwise. Thanks.
author_association: CONTRIBUTOR
user: login: chandrikarj
created_at: 2025-03-12T20:46:51Z
body: Hi @JoeZiminski , this might be a slightly long message, apologies.

I have made the changes to directly test whitening in Spikewrap against the SpikeInterface's whiten(). 
I’m using the example data (as in the spikewrap example script), applying whitening only on the runs, and then comparing using the same parameters (e.g., mode="global", apply_mean=False, dtype="float32", int_scale=1.0) on both Spikewrap’s session.preprocess({<configs to run just whitening>}) and SpikeInterface’s whiten(session._raw_runs[0]["grouped"]). I’m using chunk size = 1000 and num_chunks_per_segment = 150 to ensure enough data is sampled for the covariance estimation.

The good news is that everything runs, and the mean and std are close. The test currently prints all of this as I wanted to see and compare the outputs.
But, the overall covariance difference still diverges more than I think it should. (can be wrong here)🤔 
I think it might be because : whitening relies on estimating the covariance matrix from randomly sampled data chunks and  because minor differences in floating-point calculations are accumulating when looking at the full covariance matrix. (across all 384 channels, the overall difference comes out to be larger)

I tried to minimise it by adding regularisation, and setting apply_mean to True.
SpikeInterface’s whitening function currently supports regularisation of the covariance matrix via scikit‑learn’s GraphicalLasso methods. By default, if one enables regularisation), it uses GraphicalLassoCV (cross‑validated) via the regularize_kwargs parameter. One can also choose to use GraphicalLasso (with fixed parameters) by setting, which I tried, but GraphicalLasso sometimes fails to converge or ends up with a high difference.

So, I ended up reading more about it. 
One alternative approach (which might help as GraphicalLasso keeps struggling) is to compute the covariance matrix externally using a different estimator such as LedoitWolf or Oracle Approximating Shrinkage, then invert it (take its inverse square root) to get a whitening matrix \(W\), and compute the mean vector \(M\). 

After that, to pass W2d=W and M1d=M directly into the whiten() function, effectively bypassing the internal covariance estimation in SpikeInterface. This is said to often yield more stable results, especially for large channel counts or more challenging data as in our case (15000,384).

Now, this is all in theory and reading for now. 

Let me know if you’d like me to try this approach in the pipeline, or if there is any other way you think can be easier, or any feedback. I would like to hear your thoughts on this, as you'd be the best judge on what can be added. In any case, I’m happy to keep iterating on the parameters more if we want to stick with the built-in method, just wanted to flag the persistent covariance difference and a potential workaround. 

Thanks. 
author_association: CONTRIBUTOR
user: login: JoeZiminski
created_at: 2025-03-12T22:46:10Z
body: Hi @chandrikarj thanks a lot for this, the approach in the test is looking good, and thanks for the investigation into what might be causing the issue. That's a good strategy to try and stabilize as much as possible the covariance estimate to make it robust to differences that may be caused by the chunking. 

I think this might be due to the `random_chunk_kwargs` argument on `spikeinterface`, I see you have correctly set the seed but I am not 100% sure that is actually doing anything due to how spikeinterface handles seed-setting under the hood. To compute the covariance matrix, as you say spikeinterface is taking some random chunks from the data [here](https://github.com/SpikeInterface/spikeinterface/blob/17c3d328f1c039e9cb70b576086eb19beb2c4ffb/src/spikeinterface/preprocessing/whiten.py#L242). Within the chunking function, (and after another function call) the main work is done [here](https://github.com/SpikeInterface/spikeinterface/blob/17c3d328f1c039e9cb70b576086eb19beb2c4ffb/src/spikeinterface/core/recording_tools.py#L517). You can see in that function spikeinterface is re-setting the seed. I think if you pass 

```
random_chunk_kwargs = {"seed": 42}
```
to the whiten function calls, the seed should be set equally for both runs and I believe the results should be identical (even if there is instability, it is the same on both calls).

Manually calculating the mean / covariance is a great approach, but for convenience the whitening matrix is stored on the spikeinterface recording object `_kwargs` [here](https://github.com/SpikeInterface/spikeinterface/blob/17c3d328f1c039e9cb70b576086eb19beb2c4ffb/src/spikeinterface/preprocessing/whiten.py#L118). So to test, I would suggest comparing the traces as you do + the two whitening matrices on the recording objects.

This should be enough to convince us that the underlying spikeinterface function is called correctly, you can then go ahead and remove the other calculations and the print statements. I hope the suggested random seed fix works to make the outputs identical, otherwise please let me know! 

BTW another trick you can do is, define the kwargs once:

```
        whitening_config = {
            "1": [
                "whiten",
                {
                    "apply_mean": False,
                    "mode": "global",
                    "dtype": "float32",
                    "int_scale": 1.0,
                    "chunk_size": 1000,
                    "num_chunks_per_segment": 150,
                    "regularize": False,
                    "regularize_kwargs": {
                        "method": "GraphicalLasso",
                        "alpha": 1.0,
                        "max_iter": 500,
                    },
                    "eps": 1e-3,
                },
            ]
        }
 ```
 
 and then use these to call the spikeinterface function with `si.whiten(recording, **whitening_config["1"][1])` which will unpack the dictionary (to save having to define the kwargs twice). Let me know if you have any questions!
author_association: MEMBER
user: login: chandrikarj
created_at: 2025-03-15T10:51:39Z
body: Hi @JoeZiminski, 
I think I have incorporated all the feedback into the changes. As you said, the seed wasn't getting passed correctly to the seed-setting on spikeinterface and adding it as a separate config solved the whole issue nicely. It was because of get_random_recording_slices still setting the seed under the hood and setting it separately not doing much (as you said).

There's just one more thing I'd like to add here : since spikeinterface was using GraphicalLassoCV as default [here](https://github.com/SpikeInterface/spikeinterface/blob/17c3d328f1c039e9cb70b576086eb19beb2c4ffb/src/spikeinterface/preprocessing/whiten.py#L245)  so I removed using regularize_kwargs to pass GraphicalLasso and having to set alpha manually. It's all working well but takes a bit longerto complete (this is only in the case if regularise is set to true). 
But I think it's still better than getting into trying to find a well-tuned alpha value to avoid ill-fitted covariance matrices. 
And one can always add regularize_kwargs if they are sure of their regularisation parameters.

Please let me know what you think, and if any other changes are needed, thanks.

author_association: CONTRIBUTOR
user: login: JoeZiminski
created_at: 2025-03-19T12:56:02Z
body: Hey @chandrikarj thanks for this! This looks great, it is basically ready to merge. I've requested some extremely minor changes around comments, just to make it super obvious what is going on without reading the code. In general test code can be a little hard to read / figure out what's going on, so it can be particularly useful to have comments just indicating what each step does. 

In this case you can quickly skim the test function without actually reading any code to know what it does. Non-test code should in general be more readable and require less comments, but because test code is usually doing strange things for the purposes of testing edge cases, or accessing deep into the internals of public objects, it is often harder to skim and understand.

That's a good point about the `reguarlise_kwargs`. In this case, I think it makes sense to set it to `True` (and just use spikeinerface default). The reason is that it is `False` as default in spikeinterface, and so there is some chance that the kwarg is not passed from spikewrap to spikeinterface, but we wouldn't know because we set the test case to the spikeinterface default. For this test, we don't mind if the matrix is ill conditioned (unless I am missing something) as long as the output is identical between spikewrap and spikeinterface. For this reason it might also be worth setting `apply_mean=True`.

Thanks again, the test is looking brilliant, these are just super-pernickety things 😅 and this will be the last round of review. Cheers!


author_association: MEMBER
user: login: chandrikarj
created_at: 2025-03-19T16:56:44Z
body: Thanks for the review, @JoeZiminski, and I agree with all you mentioned here. As you said, if the ill-fitted covariance matrix isn't a problem here, then I don't see any other concerns in particular. I have added the changes and they were minor fixes, no trouble at all, always happy to iterate towards improvement.
author_association: CONTRIBUTOR