Skip to main content
`Ctrl`+`K`

#

# Source code for ada.short.unified

    from dataclasses import dataclass, field
    from abc import ABC, abstractmethod
    from typing import Tuple, Any
    import numpy as np
    import scipy.signal as ss
    from tqdm import tqdm
    from multiprocess.pool import Pool
    from math import ceil, isclose
    import tomllib
    import os
    from functools import cache
    # from scipy.optimize import minimize

    from ada.data_containers.scored import ScoredShort
    from ada.data_containers._base import _Epoched, _Raw
    from ada.data_containers.generic import GenericData
    from ada.io.acti_eeg import ActiPSG

    def matthews_corrcoef(psg, acti):
        tp = np.nansum(np.logical_and(psg == 0, acti == 0))
        tn = np.nansum(np.logical_and(psg == 1, acti == 1))
        fp = np.nansum(np.logical_and(psg == 1, acti == 0))
        fn = np.nansum(np.logical_and(psg == 0, acti == 1))
        if tp == 0 and fp == 0:
            if fn == 0 and tn != 0:
                return 1
            elif fn != 0 and tn == 0:
                return -1
            return 0
        elif tp == 0 and fn == 0:
            if fp == 0 and tn != 0:
                return 1
            elif fp != 0 and tn == 0:
                return -1
            return 0
        elif tn == 0 and fp == 0:
            if fn == 0 and tp != 0:
                return 1
            elif fn != 0 and tp == 0:
                return -1
            return 0
        elif tn == 0 and fn == 0:
            if fp == 0 and tp != 0:
                return 1
            elif fp != 0 and tp == 0:
                return -1
            return 0
        temp = (tp * tn - fp * fn) / ((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn)) ** .5
        if np.isnan(temp):
            return 0
        return temp

    @dataclass(slots=True, eq=False)
    class _ShortScorer(ABC):
        _threshold: float | None
        _filter: Any
        _scoring_method_metadata: dict = field(init=False)
        _default_thresholds: dict = field(init=False)

        @property
        def scoring_method_metadata(self) -> dict:
            """Metadata asssociated with the scoring algorithm and its parameters."""
            return self._scoring_method_metadata

        @abstractmethod
        def to_score(self, acti_data: Any) -> ScoredShort:
            pass

        @abstractmethod
        def transmittance(self, n_points: int = 2048, db: bool = True) -> Tuple[np.ndarray, np.ndarray]:
            pass

        def fit_threshold(self, acti_psg: list[ActiPSG], thresholds: Tuple[float, float] = (.001, .03), n_thresholds: int = 1000, full_data: bool = False) -> float | list[Tuple[float, float]]:
            """Fit a threshold that is maximazing mean correlation between PSG scorings and actigraphic scorings.

            Args:
                acti_psg (list[ActiPSG]): List of objects containing actigraphic data and PSG sleep/wake scorings.
                thresholds (Tuple[float, float], optional): Range of thresholds to fit. Defaults to (.001, .03).
                n_thresholds (int, optional): Number of thresholds to fit. Defaults to 1000.
                full_data (bool, optional): If True, list with all correlations will be returned. If False, the maximal correlation will be returned. Defaults to False.

            Returns:
                float | list[Tuple[float, float]]: List of all correlations (with corresponding thresholds) or maximal correlation.
            """
            def find_mcc(threshold: float) -> Tuple[float, float]:
                self._threshold = threshold
                mcc = 0
                for e in acti_psg:
                    score = self.to_score(e.acti_data).score
                    idx = np.argwhere(~np.isnan(e.eeg_data.psg_stages))
                    mcc += matthews_corrcoef(e.eeg_data.psg_stages[idx], score[idx])

                return (mcc / len(acti_psg), threshold)

            mcc = []
            ts = np.linspace(thresholds[0], thresholds[1], n_thresholds)
            pool = Pool()
            chunksize = ceil(n_thresholds / len(pool._pool))
            mcc = list(tqdm(pool.imap(find_mcc, ts, chunksize=chunksize), total=len(ts), desc="Fitting threshold to data"))
            mccs = [e[0] for e in mcc]
            self._threshold = mcc[np.argmax(mccs)][1]

            if full_data:
                return mcc

            return max(mccs)

            # def find_mcc(threshold: float) -> float:
            #     mcc = 0
            #     self._threshold = threshold
            #     for e in acti_psg:
            #         score = self.to_score(e.acti_data).score
            #         idx = np.argwhere(~np.isnan(e.eeg_data.psg_stages))
            #         mcc += matthews_corrcoef(score[idx], e.eeg_data.psg_stages[idx])

            #     return -mcc / len(acti_psg)

            # best_mcc = 0
            # best_threshold = 0
            # for _ in tqdm(range(10), desc="Fitting threshold to data"):
            #     start = np.random.uniform(1e-15, 1e15)
            #     optimized = minimize(find_mcc, start, method='Nelder-Mead')
            #     if -optimized.fun > best_mcc:
            #         best_mcc = -optimized.fun
            #         best_threshold = optimized.x

            # self._threshold = best_threshold

            # return best_mcc

        def _interpolate_threshold(self, data: _Raw | _Epoched | GenericData) -> float:
            device = GenericData._get_device(data).__name__
            if not isinstance(data, _Raw) and data.epoching_method_metadata is not None:
                epocher = data.epoching_method_metadata['epoching method']
            else:
                raise RuntimeError("No default thresholds for Raw data.")
            try:
                if isclose(data.fs, 1 / 60):
                    return self._default_thresholds[device][epocher][type(self).__name__]['long']

                return 1 / (self._default_thresholds[device][epocher][type(self).__name__]['slope'] / data.fs + self._default_thresholds[device][epocher][type(self).__name__]['const'])
            except KeyError:
                if device in self._default_thresholds.keys():
                    raise RuntimeError(f"Threshold interpolation for this particular epoching method ({epocher}) is not supported.")
                else:
                    raise RuntimeError(f"Threshold interpolation for this particular device ({device}) is not supported.")

        @staticmethod
        @cache
        def _load_defaults():
            with open(os.path.join(os.path.dirname(__file__), 'thresholds.toml'), 'rb') as f:
                temp = tomllib.load(f)
            return temp

    @dataclass(slots=True, eq=False)
    class _UnifiedScorer(_ShortScorer):

        _filter: list[dict]

        def __post_init__(self):
            self._scoring_method_metadata = {'algorithm': type(self).__name__,
                                             'threshold': self._threshold,
                                             'filter data': self._filter}
            self._default_thresholds = _ShortScorer._load_defaults()

        @property
        def filter(self) -> list[dict]:
            """Characteristics of the filters used by scoring algorithm."""
            return self._filter

        def to_score(self, acti_data: _Epoched | GenericData | _Raw) -> ScoredShort:
            """Apply scoring algorithm to the input data.

            Args:
                acti_data (_Epoched | GenericData | _Raw): Epoched data to be scored by the algorithm. Scoring by unified filter does not require previous epoching.

            Returns:
                ScoredShort: Object containng scored data.
            """
            if isinstance(acti_data, _Raw):
                epoching_method_metadata = None
            else:
                epoching_method_metadata = acti_data.epoching_method_metadata

            if self._threshold is None:
                threshold = self._interpolate_threshold(acti_data)
            else:
                threshold = self._threshold

            filters = []
            for filter in self._filter:
                sos = ss.iirdesign(filter['wp'], filter['ws'], filter['gpass'], filter['gstop'],
                                   ftype=filter['ftype'], output='sos', fs=acti_data.fs)
                filters.append(sos)

            out = np.empty((2, acti_data.data.shape[1]))
            out[1, :] = acti_data.timestamp
            temp = acti_data.to_score
            for sos in filters:
                temp = ss.sosfiltfilt(sos, temp)

            out[0, :] = temp >= (1 / threshold)
            del temp
            metadata = self._scoring_method_metadata.copy()
            metadata['threshold'] = threshold
            return ScoredShort(out, acti_data.metadata, acti_data.fs, epoching_method_metadata, metadata)

        def transmittance(self, n_points: int = 2048, db: bool = True, fs: float = 1 / 30) -> Tuple[np.ndarray, np.ndarray]:
            """Transmittance of filter used in the algorithm.

            Args:
                n_points (int, optional): Number of point at which filter will be evaluated. Defaults to 2048.
                db (bool, optional): Whether to return in dB scale. Defaults to True.
                fs (float, optional): Sampling frequency for which transmittance will be created.

            Returns:
                Tuple[np.ndarray, np.ndarray]: Vector of frequencies and vector of frequency responses.
            """
            f = np.linspace(0, fs / 2, n_points, endpoint=False)
            h = np.ones(n_points)
            for filt in self._filter:
                sos = ss.iirdesign(filt['wp'], filt['ws'], filt['gpass'], filt['gstop'], filt['analog'], filt['ftype'], 'sos', fs=fs)
                _, h_temp = ss.sosfreqz(sos, f, fs=fs)
                h = h * h_temp

            if db:
                return f, 20 * np.log10(np.abs(h))
            return f, np.abs(h)

    class GenericUnified(_UnifiedScorer):
        """A class for scoring actigraphic data using using a custom filter, designed via scipy.iirdesign.

            Attributes:
                threshold (float): Threshold value separating sleep and wake.
                wp, ws (float | tuple[float, float]): Passband and stopband edge frequencies. Possible values are scalars (for lowpass and highpass filters) or ranges (for bandpass and bandstop filters). For digital filters, these are in the same units as fs. By default, fs is 2 half-cycles/sample, so these are normalized from 0 to 1, where 1 is the Nyquist frequency.
                gpass (float): The maximum loss in the passband (dB).
                gstop (float): The minimum attenuation in the stopband (dB).
                analog (bool): When True, return an analog filter, otherwise a digital filter is returned. Defaults to False.
                ftype (str): The type of IIR filter to design.
        """
        def __init__(self, threshold: float, wp: float | tuple[float, float], ws: float | tuple[float, float], gpass: float, gstop: float, analog: bool = False, ftype: str = 'ellip'):
            filter = {'wp': wp,
                      'ws': ws,
                      'gpass': gpass,
                      'gstop': gstop,
                      'analog': analog,
                      'ftype': ftype}
            super().__init__(threshold, [filter])

        def add_filter(self, wp: float | tuple[float, float], ws: float | tuple[float, float], gpass: float, gstop: float, analog: bool = False, ftype: str = 'ellip'):
            """Add filter to the filter algorithm. Data will be filtered using all the filters provided.

            Args:
                wp (float | tuple[float, float]): Passband edge frequency.
                ws (float | tuple[float, float]): Stopband edge frequency.
                gpass (float): The maximum loss in the passband (dB).
                gstop (float): The minimum attenuation in the stopband (dB)/
                analog (bool, optional): When True, return an analog filter, otherwise a digital filter is returned. Defaults to False.
                ftype (str, optional): The type of IIR filter to design. Defaults to 'ellip'.
            """
            filter = {'wp': wp,
                      'ws': ws,
                      'gpass': gpass,
                      'gstop': gstop,
                      'analog': analog,
                      'ftype': ftype}
            self._filter.append(filter)

    class Dummy(_UnifiedScorer):
        """A dummy class for scoring actigraphic data using without any filtration.

        Attributes:
            threshold (float): Threshold value separating sleep and wake.

        """
        def __init__(self, threshold: float):
            filter = [{None: None}]
            super().__init__(threshold, filter)

        def to_score(self, acti_data: _Epoched | GenericData | _Raw) -> ScoredShort:
            """Apply scoring algorithm to the input data.

            Args:
                acti_data (_Epoched | GenericData | _Raw): Epoched data to be scored by the algorithm. Scoring by unified filter does not require previous epoching.

            Returns:
                ScoredShort: Object containng scored data.
            """
            if isinstance(acti_data, _Raw):
                epoching_method_metadata = None
            else:
                epoching_method_metadata = acti_data.epoching_method_metadata

            if self._threshold is None:
                threshold = self._interpolate_threshold(acti_data)
            else:
                threshold = self._threshold

            out = np.empty((2, acti_data.data.shape[1]))
            out[1, :] = acti_data.timestamp
            out[0, :] = acti_data.to_score >= (1 / threshold)
            metadata = self._scoring_method_metadata.copy()
            metadata['threshold'] = threshold
            return ScoredShort(out, acti_data.metadata, acti_data.fs, epoching_method_metadata, metadata)

    class UFA(_UnifiedScorer):
        """A class for scoring actigraphic data using Unified Filter Approach (for details see Biegański et al. 2026).

        Attributes:
            threshold (float | None = None): Threshold value separating sleep and wake. Autodetected based on input data if None. Defaults to None.

        """
        def __init__(self, threshold: float | None = None):
            new_filter = {'wp': 0.0006410425509217069,
                          'ws': 0.006805922349578073,
                          'gpass': 3.5661765768415044,
                          'gstop': 14.594650542815454,
                          'analog': False,
                          'ftype': 'ellip'}
            super().__init__(threshold, [new_filter])