py-neuromodulation 0.0.3__py3-none-any.whl → 0.0.5__py3-none-any.whl

py_neuromodulation/nm_coherence.py

@@ -1,205 +1,251 @@
-from scipy import signal
-import numpy as np
-from typing import Iterable
-import warnings
-
-from py_neuromodulation import nm_features_abc
-
-
-class CoherenceObject:
-    def __init__(
-        self,
-        sfreq,
-        window,
-        fbands,
-        fband_names,
-        ch_1_name,
-        ch_2_name,
-        ch_1_idx,
-        ch_2_idx,
-        coh,
-        icoh,
-        features_coh,
-    ) -> None:
-        self.sfreq = sfreq
-        self.window = window
-        self.Pxx = None
-        self.Pyy = None
-        self.Pxy = None
-        self.f = None
-        self.coh = coh
-        self.icoh = icoh
-        self.coh_val = None
-        self.icoh_val = None
-        self.ch_1 = ch_1_name
-        self.ch_2 = ch_2_name
-        self.ch_1_idx = ch_1_idx
-        self.ch_2_idx = ch_2_idx
-        self.fbands = fbands  # list of lists, e.g. [[10, 15], [15, 20]]
-        self.fband_names = fband_names
-        self.features_coh = features_coh
-
-    def get_coh(self, features_compute, x, y):
-        self.f, self.Pxx = signal.welch(x, self.sfreq, self.window, nperseg=128)
-        self.Pyy = signal.welch(y, self.sfreq, self.window, nperseg=128)[1]
-        self.Pxy = signal.csd(x, y, self.sfreq, self.window, nperseg=128)[1]
-
-        if self.coh is True:
-            self.coh_val = np.abs(self.Pxy**2) / (self.Pxx * self.Pyy)
-        if self.icoh is True:
-            self.icoh_val = np.array(self.Pxy / (self.Pxx * self.Pyy)).imag
-
-        for coh_idx, coh_type in enumerate([self.coh, self.icoh]):
-            if coh_type is True:
-                if coh_idx == 0:
-                    coh_val = self.coh_val
-                    coh_name = "coh"
-                else:
-                    coh_val = self.icoh_val
-                    coh_name = "icoh"
-
-            for idx, fband in enumerate(self.fbands):
-                if self.features_coh["mean_fband"] is True:
-                    feature_calc = np.mean(
-                        coh_val[
-                            np.bitwise_and(self.f > fband[0], self.f < fband[1])
-                        ]
-                    )
-                    feature_name = "_".join(
-                        [
-                            coh_name,
-                            self.ch_1,
-                            "to",
-                            self.ch_2,
-                            "mean_fband",
-                            self.fband_names[idx],
-                        ]
-                    )
-                    features_compute[feature_name] = feature_calc
-                if self.features_coh["max_fband"] is True:
-                    feature_calc = np.max(
-                        coh_val[
-                            np.bitwise_and(self.f > fband[0], self.f < fband[1])
-                        ]
-                    )
-                    feature_name = "_".join(
-                        [
-                            coh_name,
-                            self.ch_1,
-                            "to",
-                            self.ch_2,
-                            "max_fband",
-                            self.fband_names[idx],
-                        ]
-                    )
-                    features_compute[feature_name] = feature_calc
-            if self.features_coh["max_allfbands"] is True:
-                feature_calc = self.f[np.argmax(coh_val)]
-                feature_name = "_".join(
-                    [
-                        coh_name,
-                        self.ch_1,
-                        "to",
-                        self.ch_2,
-                        "max_allfbands",
-                        self.fband_names[idx],
-                    ]
-                )
-                features_compute[feature_name] = feature_calc
-        return features_compute
-
-
-class NM_Coherence(nm_features_abc.Feature):
-
-    
-
-    def __init__(
-        self, settings: dict, ch_names: Iterable[str], sfreq: float
-    ) -> None:
-        self.s = settings
-        self.sfreq = sfreq
-        self.ch_names = ch_names
-        self.coherence_objects: Iterable[CoherenceObject] = []
-
-        for idx_coh in range(len(self.s["coherence"]["channels"])):
-            fband_names = self.s["coherence"]["frequency_bands"]
-            fband_specs = []
-            for band_name in fband_names:
-                fband_specs.append(self.s["frequency_ranges_hz"][band_name])
-
-            ch_1_name = self.s["coherence"]["channels"][idx_coh][0]
-            ch_1_name_reref = [
-                ch for ch in self.ch_names if ch.startswith(ch_1_name)
-            ][0]
-            ch_1_idx = self.ch_names.index(ch_1_name_reref)
-
-            ch_2_name = self.s["coherence"]["channels"][idx_coh][1]
-            ch_2_name_reref = [
-                ch for ch in self.ch_names if ch.startswith(ch_2_name)
-            ][0]
-            ch_2_idx = self.ch_names.index(ch_2_name_reref)
-
-            self.coherence_objects.append(
-                CoherenceObject(
-                    sfreq,
-                    "hann",
-                    fband_specs,
-                    fband_names,
-                    ch_1_name,
-                    ch_2_name,
-                    ch_1_idx,
-                    ch_2_idx,
-                    self.s["coherence"]["method"]["coh"],
-                    self.s["coherence"]["method"]["icoh"],
-                    self.s["coherence"]["features"],
-                )
-            )
-
-    @staticmethod
-    def test_settings(
-        s: dict,
-        ch_names: Iterable[str],
-        sfreq: int | float,
-    ):
-
-        assert (
-            len(s["coherence"]["frequency_bands"]) > 0
-        ), "coherence frequency_bands list needs to specify at least one frequency band"
-        assert (ch_coh in ch_names for ch_coh in s["coherence"]["channels"]), (
-            f"coherence selected channels don't match the ones in nm_channels"
-            f"ch_names: {ch_names} settings['coherence']['channels']: {s['coherence']['channels']}"
-        )
-
-        assert (
-            f_band_coh in s["frequency_ranges_hz"]
-            for f_band_coh in s["coherence"]["frequency_bands"]
-        ), (
-            "coherence selected frequency bands don't match the ones"
-            "specified in s['frequency_ranges_hz']"
-            f"coherence frequency bands: {s['coherence']['frequency_bands']}"
-            f"specified frequency_ranges_hz: {s['frequency_ranges_hz']}"
-        )
-
-        assert (
-            s["frequency_ranges_hz"][fb][0] < sfreq / 2
-            and s["frequency_ranges_hz"][fb][1] < sfreq / 2
-            for fb in s["coherence"]["frequency_bands"]
-        ), (
-            "the coherence frequency band ranges need to be smaller than the nyquist frequency"
-            f"got sfreq = {sfreq} and fband ranges {s['coherence']['frequency_bands']}"
-        )
-
-        if sum(list(s["coherence"]["method"].values())) == 0:
-            warnings.warn(
-                "feature coherence enabled, but no coherence['method'] selected"
-            )
-
-    def calc_feature(self, data: np.array, features_compute: dict) -> dict:
-        for coh_obj in self.coherence_objects:
-            features_compute = coh_obj.get_coh(
-                features_compute,
-                data[coh_obj.ch_1_idx, :],
-                data[coh_obj.ch_2_idx, :],
-            )
-
-        return features_compute
+import numpy as np
+from collections.abc import Iterable
+
+
+from typing import TYPE_CHECKING, Annotated
+from pydantic import Field, field_validator
+
+from py_neuromodulation.nm_features import NMFeature
+from py_neuromodulation.nm_types import BoolSelector, FrequencyRange, NMBaseModel
+from py_neuromodulation import logger
+
+if TYPE_CHECKING:
+    from py_neuromodulation.nm_settings import NMSettings
+
+
+class CoherenceMethods(BoolSelector):
+    coh: bool = True
+    icoh: bool = True
+
+
+class CoherenceFeatures(BoolSelector):
+    mean_fband: bool = True
+    max_fband: bool = True
+    max_allfbands: bool = True
+
+
+ListOfTwoStr = Annotated[list[str], Field(min_length=2, max_length=2)]
+
+
+class CoherenceSettings(NMBaseModel):
+    features: CoherenceFeatures = CoherenceFeatures()
+    method: CoherenceMethods = CoherenceMethods()
+    channels: list[ListOfTwoStr] = []
+    frequency_bands: list[str] = Field(default=["high_beta"], min_length=1)
+
+    @field_validator("frequency_bands")
+    def fbands_spaces_to_underscores(cls, frequency_bands):
+        return [f.replace(" ", "_") for f in frequency_bands]
+
+
+class CoherenceObject:
+    def __init__(
+        self,
+        sfreq: float,
+        window: str,
+        fbands: list[FrequencyRange],
+        fband_names: list[str],
+        ch_1_name: str,
+        ch_2_name: str,
+        ch_1_idx: int,
+        ch_2_idx: int,
+        coh: bool,
+        icoh: bool,
+        features_coh: CoherenceFeatures,
+    ) -> None:
+        self.sfreq = sfreq
+        self.window = window
+        self.fbands = fbands
+        self.fband_names = fband_names
+        self.ch_1 = ch_1_name
+        self.ch_2 = ch_2_name
+        self.ch_1_idx = ch_1_idx
+        self.ch_2_idx = ch_2_idx
+        self.coh = coh
+        self.icoh = icoh
+        self.features_coh = features_coh
+
+        self.Pxx = None
+        self.Pyy = None
+        self.Pxy = None
+        self.f = None
+        self.coh_val = None
+        self.icoh_val = None
+
+    def get_coh(self, feature_results, x, y):
+        from scipy.signal import welch, csd
+
+        self.f, self.Pxx = welch(x, self.sfreq, self.window, nperseg=128)
+        self.Pyy = welch(y, self.sfreq, self.window, nperseg=128)[1]
+        self.Pxy = csd(x, y, self.sfreq, self.window, nperseg=128)[1]
+
+        if self.coh:
+            self.coh_val = np.abs(self.Pxy**2) / (self.Pxx * self.Pyy)
+        if self.icoh:
+            self.icoh_val = np.array(self.Pxy / (self.Pxx * self.Pyy)).imag
+
+        for coh_idx, coh_type in enumerate([self.coh, self.icoh]):
+            if coh_type:
+                if coh_idx == 0:
+                    coh_val = self.coh_val
+                    coh_name = "coh"
+                else:
+                    coh_val = self.icoh_val
+                    coh_name = "icoh"
+
+            for idx, fband in enumerate(self.fbands):
+                if self.features_coh.mean_fband:
+                    feature_calc = np.mean(
+                        coh_val[np.bitwise_and(self.f > fband[0], self.f < fband[1])]
+                    )
+                    feature_name = "_".join(
+                        [
+                            coh_name,
+                            self.ch_1,
+                            "to",
+                            self.ch_2,
+                            "mean_fband",
+                            self.fband_names[idx],
+                        ]
+                    )
+                    feature_results[feature_name] = feature_calc
+                if self.features_coh.max_fband:
+                    feature_calc = np.max(
+                        coh_val[np.bitwise_and(self.f > fband[0], self.f < fband[1])]
+                    )
+                    feature_name = "_".join(
+                        [
+                            coh_name,
+                            self.ch_1,
+                            "to",
+                            self.ch_2,
+                            "max_fband",
+                            self.fband_names[idx],
+                        ]
+                    )
+                    feature_results[feature_name] = feature_calc
+            if self.features_coh.max_allfbands:
+                feature_calc = self.f[np.argmax(coh_val)]
+                feature_name = "_".join(
+                    [
+                        coh_name,
+                        self.ch_1,
+                        "to",
+                        self.ch_2,
+                        "max_allfbands",
+                        self.fband_names[idx],
+                    ]
+                )
+                feature_results[feature_name] = feature_calc
+        return feature_results
+
+
+class NMCoherence(NMFeature):
+    def __init__(
+        self, settings: "NMSettings", ch_names: list[str], sfreq: float
+    ) -> None:
+        self.settings = settings.coherence
+        self.frequency_ranges_hz = settings.frequency_ranges_hz
+        self.sfreq = sfreq
+        self.ch_names = ch_names
+        self.coherence_objects: Iterable[CoherenceObject] = []
+
+        self.test_settings(settings, ch_names, sfreq)
+
+        for idx_coh in range(len(self.settings.channels)):
+            fband_names = self.settings.frequency_bands
+            fband_specs = []
+            for band_name in fband_names:
+                fband_specs.append(self.frequency_ranges_hz[band_name])
+
+            ch_1_name = self.settings.channels[idx_coh][0]
+            ch_1_name_reref = [ch for ch in self.ch_names if ch.startswith(ch_1_name)][
+                0
+            ]
+            ch_1_idx = self.ch_names.index(ch_1_name_reref)
+
+            ch_2_name = self.settings.channels[idx_coh][1]
+            ch_2_name_reref = [ch for ch in self.ch_names if ch.startswith(ch_2_name)][
+                0
+            ]
+            ch_2_idx = self.ch_names.index(ch_2_name_reref)
+
+            self.coherence_objects.append(
+                CoherenceObject(
+                    sfreq,
+                    "hann",
+                    fband_specs,
+                    fband_names,
+                    ch_1_name,
+                    ch_2_name,
+                    ch_1_idx,
+                    ch_2_idx,
+                    self.settings.method.coh,
+                    self.settings.method.icoh,
+                    self.settings.features,
+                )
+            )
+
+    @staticmethod
+    def test_settings(
+        settings: "NMSettings",
+        ch_names: Iterable[str],
+        sfreq: float,
+    ):
+        flat_channels = [
+            ch for ch_pair in settings.coherence.channels for ch in ch_pair
+        ]
+
+        valid_coh_channel = [
+            sum(ch.startswith(ch_coh) for ch in ch_names) for ch_coh in flat_channels
+        ]
+        for ch_idx, ch_coh in enumerate(flat_channels):
+            if valid_coh_channel[ch_idx] == 0:
+                raise RuntimeError(
+                    f"Coherence selected channel {ch_coh} does not match any channel name: \n"
+                    f"  - settings.coherence.channels: {settings.coherence.channels}\n"
+                    f"  - ch_names: {ch_names} \n"
+                )
+
+            if valid_coh_channel[ch_idx] > 1:
+                raise RuntimeError(
+                    f"Coherence selected channel {ch_coh} is ambigous and matches more than one channel name: \n"
+                    f"  - settings.coherence.channels: {settings.coherence.channels}\n"
+                    f"  - ch_names: {ch_names} \n"
+                )
+
+        assert all(
+            f_band_coh in settings.frequency_ranges_hz
+            for f_band_coh in settings.coherence.frequency_bands
+        ), (
+            "coherence selected frequency bands don't match the ones"
+            "specified in s['frequency_ranges_hz']"
+            f"coherence frequency bands: {settings.coherence.frequency_bands}"
+            f"specified frequency_ranges_hz: {settings.frequency_ranges_hz}"
+        )
+
+        assert all(
+            settings.frequency_ranges_hz[fb][0] < sfreq / 2
+            and settings.frequency_ranges_hz[fb][1] < sfreq / 2
+            for fb in settings.coherence.frequency_bands
+        ), (
+            "the coherence frequency band ranges need to be smaller than the Nyquist frequency"
+            f"got sfreq = {sfreq} and fband ranges {settings.coherence.frequency_bands}"
+        )
+
+        if not settings.coherence.method.get_enabled():
+            logger.warn(
+                "feature coherence enabled, but no coherence['method'] selected"
+            )
+
+    def calc_feature(self, data: np.ndarray) -> dict:
+        feature_results = {}
+
+        for coh_obj in self.coherence_objects:
+            feature_results = coh_obj.get_coh(
+                feature_results,
+                data[coh_obj.ch_1_idx, :],
+                data[coh_obj.ch_2_idx, :],
+            )
+
+        return feature_results

py_neuromodulation/nm_database.py

@@ -0,0 +1,149 @@
+import sqlite3
+from pathlib import Path
+import pandas as pd
+from py_neuromodulation.nm_types import _PathLike
+from py_neuromodulation.nm_IO import generate_unique_filename
+
+
+class NMDatabase:
+    """
+    Class to create a database and insert data into it.
+    Parameters
+    ----------
+    out_dir : _PathLike
+        The directory to save the database.
+    csv_path : str, optional
+        The path to save the csv file. If not provided, it will be saved in the same folder as the database.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        out_dir: _PathLike,
+        csv_path: _PathLike | None = None,
+    ):
+        # Make sure out_dir exists
+        Path(out_dir).mkdir(parents=True, exist_ok=True)
+
+        self.db_path = Path(out_dir, f"{name}.db")
+
+        self.table_name = f"{name}_data"  # change to param?
+        self.table_created = False
+
+        if self.db_path.exists():
+            self.db_path = generate_unique_filename(self.db_path)
+            name = self.db_path.stem
+
+        if csv_path is None:
+            self.csv_path = Path(out_dir, f"{name}.csv")
+        else:
+            self.csv_path = Path(csv_path)
+
+        self.csv_path.parent.mkdir(parents=True, exist_ok=True)
+
+        self.conn = sqlite3.connect(self.db_path)
+        self.cursor = self.conn.cursor()
+
+        # Database config and optimization, prioritize data integrity
+        self.cursor.execute("PRAGMA journal_mode=WAL")  # Write-Ahead Logging mode
+        self.cursor.execute("PRAGMA synchronous=FULL")  # Sync on every commit
+        self.cursor.execute("PRAGMA temp_store=MEMORY")  # Store temp tables in memory
+        self.cursor.execute(
+            "PRAGMA wal_autocheckpoint = 1000"
+        )  # WAL checkpoint every 1000 pages (default, 4MB, might change)
+        self.cursor.execute(
+            f"PRAGMA mmap_size = {2 * 1024 * 1024 * 1024}"
+        )  # 2GB of memory mapped
+
+    def infer_type(self, value):
+        """Infer the type of the value to create the table schema.
+        Parameters
+        ----------
+        value : int, float, str
+            The value to infer the type."""
+
+        if isinstance(value, (int, float)):
+            return "REAL"
+        elif isinstance(value, str):
+            return "TEXT"
+        else:
+            return "BLOB"
+
+    def create_table(self, feature_dict: dict):
+        """
+        Create a table in the database.
+        Parameters
+        ----------
+        feature_dict : dict
+            The dictionary with the feature names and values.
+        """
+        columns_schema = ", ".join(
+            [
+                f'"{column}" {self.infer_type(value)}'
+                for column, value in feature_dict.items()
+            ]
+        )
+
+        self.cursor.execute(
+            f'CREATE TABLE IF NOT EXISTS "{self.table_name}" ({columns_schema})'
+        )
+
+        # Create column names and placeholders for insert statement
+        self.columns: str = ", ".join([f'"{column}"' for column in feature_dict.keys()])
+        # Use named placeholders for more resiliency against unexpected change in column order
+        self.placeholders = ", ".join([f":{key}" for key in feature_dict.keys()])
+
+    def insert_data(self, feature_dict: dict):
+        """
+        Insert data into the database.
+        Parameters
+        ----------
+        feature_dict : dict
+            The dictionary with the feature names and values.
+        """
+
+        if not self.table_created:
+            self.create_table(feature_dict)
+            self.table_created = True
+
+        insert_sql = f'INSERT INTO "{self.table_name}" ({self.columns}) VALUES ({self.placeholders})'
+
+        self.cursor.execute(insert_sql, feature_dict)
+
+    def commit(self):
+        self.conn.commit()
+
+    def fetch_all(self):
+        """ "
+        Fetch all the data from the database.
+        Returns
+        -------
+        pd.DataFrame
+            The data in a pandas DataFrame.
+        """
+        return pd.read_sql_query(f'SELECT * FROM "{self.table_name}"', self.conn)
+
+    def head(self, n: int = 5):
+        """ "
+        Returns the first N rows of the database.
+        Parameters
+        ----------
+        n : int, optional
+            The number of rows to fetch, by default 1
+        -------
+        pd.DataFrame
+            The data in a pandas DataFrame.
+        """
+        return pd.read_sql_query(
+            f'SELECT * FROM "{self.table_name}" LIMIT {n}', self.conn
+        )
+
+    def save_as_csv(self):
+        df = self.fetch_all()
+        df.to_csv(self.csv_path, index=False)
+
+    def close(self):
+        # Optimize before closing is recommended:
+        # https://www.sqlite.org/pragma.html#pragma_optimize
+        self.cursor.execute("PRAGMA optimize")
+        self.conn.close()