From 32e3a04126ef5ec60139c60217d24485e28f556e Mon Sep 17 00:00:00 2001 From: Chandrika Singh Jadon Date: Tue, 18 Mar 2025 01:05:43 +0530 Subject: [PATCH 1/2] plots for bad channel detection --- spikewrap/structure/_preprocess_run.py | 88 +++++++++++ .../test_integration/test_diagnostic_plots.py | 138 ++++++++++++++++++ 2 files changed, 226 insertions(+) create mode 100644 tests/test_integration/test_diagnostic_plots.py diff --git a/spikewrap/structure/_preprocess_run.py b/spikewrap/structure/_preprocess_run.py index 5032adf..c95c6c2 100644 --- a/spikewrap/structure/_preprocess_run.py +++ b/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 diff --git a/tests/test_integration/test_diagnostic_plots.py b/tests/test_integration/test_diagnostic_plots.py new file mode 100644 index 0000000..ec8f76e --- /dev/null +++ b/tests/test_integration/test_diagnostic_plots.py @@ -0,0 +1,138 @@ +import pytest +from pathlib import Path + +from spikewrap.structure._preprocess_run import PreprocessedRun +import spikewrap.visualise +import numpy as np +import matplotlib.pyplot as plt +import shutil + +@pytest.fixture +def mock_preprocessed_run(tmp_path, monkeypatch): + """ + Fixture to create a temporary PreprocessedRun instance with mock data. + """ + + from spikewrap.structure._preprocess_run import PreprocessedRun + + 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}" From 16cd13353e3224623f6581de59897dfbef60d97a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 17 Mar 2025 20:11:25 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../test_integration/test_diagnostic_plots.py | 76 ++++++++++--------- 1 file changed, 42 insertions(+), 34 deletions(-) diff --git a/tests/test_integration/test_diagnostic_plots.py b/tests/test_integration/test_diagnostic_plots.py index ec8f76e..d76ebf8 100644 --- a/tests/test_integration/test_diagnostic_plots.py +++ b/tests/test_integration/test_diagnostic_plots.py @@ -1,11 +1,12 @@ -import pytest +import shutil from pathlib import Path -from spikewrap.structure._preprocess_run import PreprocessedRun -import spikewrap.visualise -import numpy as np import matplotlib.pyplot as plt -import shutil +import numpy as np +import pytest + +from spikewrap.structure._preprocess_run import PreprocessedRun + @pytest.fixture def mock_preprocessed_run(tmp_path, monkeypatch): @@ -13,73 +14,72 @@ def mock_preprocessed_run(tmp_path, monkeypatch): Fixture to create a temporary PreprocessedRun instance with mock data. """ - from spikewrap.structure._preprocess_run import PreprocessedRun - 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)) - + 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] - + 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) + 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", @@ -89,20 +89,24 @@ def __array__(self): 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)) + monkeypatch.setattr( + preprocessed_run, + "_save_diagnostic_plots", + mock_save_diagnostic_plots.__get__(preprocessed_run), + ) yield preprocessed_run @@ -120,10 +124,14 @@ def test_diagnostic_plots_saved(self, mock_preprocessed_run): if output_dir.exists(): shutil.rmtree(output_dir) - assert not output_dir.exists(), "Diagnostic plots directory should not exist before running save_preprocessed" + 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) + 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"