[Ref #215] Diagnostic plots for bad channel detection

spikewrap/structure/_preprocess_run.py

@@ -120,6 +120,8 @@ def save_preprocessed(
             ) as file:
                 file.write("\n".join(self._orig_run_names))
 
+        self._save_diagnostic_plots()
+
     def save_class_attributes_to_yaml(self, path_to_save):
         """
         Dump the class attributes to file so the class
@@ -210,6 +212,92 @@ def plot_preprocessed(
     # ---------------------------------------------------------------------------
     # Private Functions
     # ---------------------------------------------------------------------------
+    def _generate_and_save_plot(
+        self,
+        run_name,
+        preprocessed_dict,
+        ses_name,
+        output_path,
+        filename,
+        *,
+        mode="map",
+        time_range=(0, 10),
+        show_channel_ids=True,
+    ):
+        """
+        Helper method to generate and save a plot for a given mode.
+
+        :param run_name: Name of the run (e.g., "test_run").
+        :param preprocessed_dict: Dictionary containing the preprocessed recording.
+        :param ses_name: Session name (e.g., "test_session").
+        :param output_path: Path where the plot should be saved.
+        :param filename: The file name for the saved plot.
+        :param mode: The visualization mode ("map" or "line"), defaults to "map".
+        :param time_range: The time range for the plot, defaults to (0, 10).
+        :param show_channel_ids: Whether to show channel IDs in the plot, defaults to True.
+        """
+        fig = visualise_run_preprocessed(
+            run_name,
+            False,
+            preprocessed_dict,
+            ses_name,
+            figsize=(10, 5),
+            mode=mode,
+            time_range=time_range,
+            show_channel_ids=show_channel_ids,
+        )
+        fig.savefig(output_path / filename)
+        fig.clf()
+
+    def _save_diagnostic_plots(self) -> None:
+        """
+        Save diagnostic plots after bad channel detection.
+
+        This function generates and saves:
+        - A plot of the data before bad channel detection.
+        - A plot of the data after bad channel detection.
+        - Individual plots for each detected bad channel.
+        """
+        diagnostic_path = self._output_path / "diagnostic_plots"
+        diagnostic_path.mkdir(parents=True, exist_ok=True)
+
+        _utils.message_user(f"Saving diagnostic plots for: {self._run_name}...")
+
+        for shank_name, preprocessed_dict in self._preprocessed.items():
+            preprocessed_recording, _ = _utils._get_dict_value_from_step_num(
+                preprocessed_dict, "last"
+            )
+
+            # Generate before and after plots
+            self._generate_and_save_plot(
+                self._run_name,
+                preprocessed_dict,
+                self._ses_name,
+                diagnostic_path,
+                f"{shank_name}_before_detection.png",
+                mode="map",
+            )
+            self._generate_and_save_plot(
+                self._run_name,
+                preprocessed_dict,
+                self._ses_name,
+                diagnostic_path,
+                f"{shank_name}_after_detection.png",
+                mode="map",
+            )
+
+            # Save individual bad channel plots
+            bad_channels = preprocessed_recording.get_property("bad_channels")
+            if bad_channels:
+                for ch in bad_channels:
+                    self._generate_and_save_plot(
+                        self._run_name,
+                        preprocessed_dict,
+                        self._ses_name,
+                        diagnostic_path,
+                        f"{shank_name}_bad_channel_{ch}.png",
+                        mode="line",
+                    )
 
     def _save_preprocessed_slurm(
         self, overwrite: bool, chunk_duration_s: float, n_jobs: int, slurm: dict | bool

tests/test_integration/test_diagnostic_plots.py

@@ -0,0 +1,146 @@
+import shutil
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+
+from spikewrap.structure._preprocess_run import PreprocessedRun
+
+
+@pytest.fixture
+def mock_preprocessed_run(tmp_path, monkeypatch):
+    """
+    Fixture to create a temporary PreprocessedRun instance with mock data.
+    """
+
+    def mock_plot(*args, **kwargs):
+        pass
+
+    def mock_figure(*args, **kwargs):
+        class MockFigure:
+            def savefig(self, path):
+                Path(path).parent.mkdir(parents=True, exist_ok=True)
+                Path(path).touch()
+
+            def clf(self):
+                pass
+
+        return MockFigure()
+
+    def mock_visualise(*args, **kwargs):
+        return mock_figure()
+
+    monkeypatch.setattr(plt, "figure", mock_figure)
+    monkeypatch.setattr(plt, "plot", mock_plot)
+    monkeypatch.setattr(plt, "subplot", lambda *args, **kwargs: None)
+    monkeypatch.setattr(plt, "title", lambda *args, **kwargs: None)
+
+    import sys
+
+    module_name = PreprocessedRun.__module__
+    module = sys.modules[module_name]
+    monkeypatch.setattr(module, "visualise_run_preprocessed", mock_visualise)
+
+    class MockRecording:
+        def __init__(self):
+            self.properties = {}
+            self.data = np.random.random((10, 1000))
+
+        def save(self, folder, chunk_duration):
+            Path(folder).mkdir(parents=True, exist_ok=True)
+            (Path(folder) / "mock_recording_saved.txt").touch()
+            return True
+
+        def get_property(self, property_name):
+            return self.properties.get(property_name, [])
+
+        def get_traces(self, *args, **kwargs):
+            return self.data
+
+        def __array__(self):
+            return self.data
+
+    # Set up a mock recording with bad channels
+    mock_recording = MockRecording()
+    mock_recording.properties["bad_channels"] = [0, 1]
+
+    raw_data_path = tmp_path / "raw_data"
+    session_output_path = tmp_path / "output"
+    run_name = "test_run"
+
+    preprocessed_data = {"shank_0": {"0": mock_recording, "1": mock_recording}}
+
+    raw_data_path.mkdir(parents=True, exist_ok=True)
+    session_output_path.mkdir(parents=True, exist_ok=True)
+
+    preprocessed_path = session_output_path / run_name / "preprocessed"
+    preprocessed_path.mkdir(parents=True, exist_ok=True)
+
+    diagnostic_path = session_output_path / "diagnostic_plots"
+    diagnostic_path.mkdir(parents=True, exist_ok=True)
+
+    preprocessed_run = PreprocessedRun(
+        raw_data_path=raw_data_path,
+        ses_name="test_session",
+        run_name=run_name,
+        file_format="mock_format",
+        session_output_path=session_output_path,
+        preprocessed_data=preprocessed_data,
+        pp_steps={"step_1": "bad_channel_detection"},
+    )
+
+    def mock_save_diagnostic_plots(self):
+        diagnostic_path = self._output_path / "diagnostic_plots"
+        diagnostic_path.mkdir(parents=True, exist_ok=True)
+
+        for shank_name in self._preprocessed:
+            (diagnostic_path / f"{shank_name}_before_detection.png").touch()
+            (diagnostic_path / f"{shank_name}_after_detection.png").touch()
+
+            for ch in [0, 1]:
+                (diagnostic_path / f"{shank_name}_bad_channel_{ch}.png").touch()
+
+    # Monkeypatch the method to create placeholder files instead of real plots
+    monkeypatch.setattr(
+        preprocessed_run,
+        "_save_diagnostic_plots",
+        mock_save_diagnostic_plots.__get__(preprocessed_run),
+    )
+
+    yield preprocessed_run
+
+
+class TestDiagnosticPlots:
+    """
+    Test class to validate diagnostic plots are saved correctly.
+    """
+
+    def test_diagnostic_plots_saved(self, mock_preprocessed_run):
+        """
+        Test if diagnostic plots are correctly saved after running save_preprocessed.
+        """
+        output_dir = mock_preprocessed_run._output_path / "diagnostic_plots"
+
+        if output_dir.exists():
+            shutil.rmtree(output_dir)
+        assert (
+            not output_dir.exists()
+        ), "Diagnostic plots directory should not exist before running save_preprocessed"
+
+        # Should trigger the diagnostic plot saving
+        mock_preprocessed_run.save_preprocessed(
+            overwrite=True, chunk_duration_s=1.0, n_jobs=1, slurm=False
+        )
+
+        assert output_dir.exists(), "Diagnostic plots directory was not created"
+        shank_name = "shank_0"
+        expected_files = [
+            f"{shank_name}_before_detection.png",
+            f"{shank_name}_after_detection.png",
+            f"{shank_name}_bad_channel_0.png",
+            f"{shank_name}_bad_channel_1.png",
+        ]
+
+        for file_name in expected_files:
+            assert (output_dir / file_name).exists(), f"Missing plot file: {file_name}"