Multiple minor fixes:

Save bad channels and save annotations moved to separate functions.
_save_psds moved to public domain save_psds
In dashboard set reference after channel interpolation
Add return in ICA plots.

Author: Genuster

sleepeeg/base.py

@@ -141,37 +141,56 @@ def plot(
         """
         kwargs.setdefault("theme", "dark")
         kwargs.setdefault("bad_color", "r")
-        kwargs.setdefault("scalings", "auto")
         kwargs["block"] = True
 
         self.mne_raw.plot(**kwargs)
 
         if save_annotations:
-            self.mne_raw.annotations.save(
-                self.output_dir / self.__class__.__name__ / "annotations.txt",
-                overwrite=overwrite,
-            )
+            self.save_annotations(overwrite=overwrite)
 
         if save_bad_channels:
-            new_bads = self.mne_raw.info["bads"]
-
-            if new_bads:
-                from natsort import natsorted
-
-                fpath = self.output_dir / self.__class__.__name__ / "bad_channels.txt"
-                old_bads = []
-                if fpath.exists():
-                    with open(fpath, "r") as f:
-                        old_bads = f.read().split()
-
-                with open(fpath, "w") as f:
-                    bads = (
-                        natsorted(new_bads)
-                        if overwrite
-                        else natsorted(set(old_bads + new_bads))
-                    )
-                    for bad in bads:
-                        f.write(f"{bad}\n")
+            self.save_bad_channels(overwrite=overwrite)
+
+    def save_bad_channels(self, overwrite=False):
+        """Adds bad channels from info["bads"] to the "bad_channels.txt" file.
+
+        Args:
+            overwrite: Whether to overwrite the file if exists.
+                If False will add unique new channels to the file.
+                Defaults to False.
+        """
+        new_bads = self.mne_raw.info["bads"]
+
+        if new_bads:
+            from natsort import natsorted
+
+            fpath = self.output_dir / self.__class__.__name__ / "bad_channels.txt"
+            old_bads = []
+            if fpath.exists():
+                with open(fpath, "r") as f:
+                    old_bads = f.read().split()
+
+            with open(fpath, "w") as f:
+                bads = (
+                    natsorted(new_bads)
+                    if overwrite
+                    else natsorted(set(old_bads + new_bads))
+                )
+                for bad in bads:
+                    f.write(f"{bad}\n")
+
+    def save_annotations(self, overwrite=False):
+        """Writes annotations to "annotations.txt" file.
+
+        Args:
+            overwrite: Whether to overwrite the file if exists.
+                If False and the file exists will throw an exception.
+                Defaults to False.
+        """
+        self.mne_raw.annotations.save(
+            self.output_dir / self.__class__.__name__ / "annotations.txt",
+            overwrite=overwrite,
+        )
 
     def plot_sensors(
         self, legend: Iterable[str] = None, legend_args: dict = None, **kwargs
@@ -1143,7 +1162,7 @@ def plot_topomap_collage(
             )
 
     @logger_wraps()
-    def _save_psds(self, overwrite):
+    def save_psds(self, overwrite):
         import re
 
         for stage, spectrum in self.psds.items():

sleepeeg/dashboard.py

@@ -286,8 +286,6 @@ def create_dashboard(
         bandpass_filter_freqs[1],
     )
 
-    pipe.set_eeg_reference(ref_channels=ref_channels)
-
     if path_to_bad_channels is not None:
         pipe.read_bad_channels(path=path_to_bad_channels)
         bads = pipe.mne_raw.info["bads"]
@@ -297,6 +295,7 @@ def create_dashboard(
 
         bads = literal_eval(pipe.mne_raw.info["description"])
 
+    pipe.set_eeg_reference(ref_channels=ref_channels)
     s_pipe, sleep_stages = _init_s_pipe(pipe, path_to_hypnogram, hypno_freq)
     min_psd, max_psd = _hypno_psd(
         s_pipe,

sleepeeg/pipeline.py

@@ -222,15 +222,15 @@ def fit(self, filter_kwargs: dict = None, **fit_kwargs):
 
     def plot_sources(self, **kwargs):
         """A wrapper for :py:meth:`mne:mne.preprocessing.ICA.plot_sources`."""
-        self.mne_ica.plot_sources(inst=self.mne_raw, **kwargs)
+        return self.mne_ica.plot_sources(inst=self.mne_raw, **kwargs)
 
     def plot_components(self, **kwargs):
         """A wrapper for :py:meth:`mne:mne.preprocessing.ICA.plot_components`."""
-        self.mne_ica.plot_components(inst=self.mne_raw, **kwargs)
+        return self.mne_ica.plot_components(inst=self.mne_raw, **kwargs)
 
     def plot_properties(self, picks=None, **kwargs):
         """A wrapper for :py:meth:`mne:mne.preprocessing.ICA.plot_properties`."""
-        self.mne_ica.plot_properties(self.mne_raw, picks=picks, **kwargs)
+        return self.mne_ica.plot_properties(self.mne_raw, picks=picks, **kwargs)
 
     @logger_wraps()
     def apply(self, exclude=None, **kwargs):
@@ -321,7 +321,7 @@ def compute_psds_per_stage(
                 **psd_kwargs,
             )
         if save:
-            self._save_psds(overwrite)
+            self.save_psds(overwrite)
 
     @logger_wraps()
     def read_spectra(self, dirpath: str | None = None):
@@ -740,4 +740,4 @@ def compute_psds_per_stage(
             self.psds[stage] = avg_func(spectra, axis=0)
 
         if save:
-            self._save_psds(overwrite)
+            self.save_psds(overwrite)