Skip to main content
`Ctrl`+`K`

#

# Source code for ada.long.ar

    import numpy as np
    import matplotlib.pyplot as plt
    import tqdm
    import scipy.signal as ss
    import pprint
    import csv
    from typing import Tuple

    from ada.data_containers._base import _ActiData
    from ada.data_containers._long import _Long, _LongContainer

    class ARResults(_LongContainer):
        """A class for storing results of AR model fitting. Do not construct manually! The fit_params gives access to the following fields:
           - coefficients: of fitted model
           - noise std: standard deviations of model noise
           - order: order of the model
           - peak: position of the main peak near to the 24 hours
           - preprocessing: whether it was used
           - period ci: confidence interval of the peak position estimated using bootstrap or None.
        """
        def __init__(self, a: np.ndarray, sigma: float, f: np.ndarray, s: np.ndarray, filter: bool, ci: Tuple[float, float] | None, id: str):
            # find peak near 24 h
            s_around_24 = s[(f >= 1 / (32 * 3600)) & (f <= 1 / (16 * 3600))]
            f_around_24 = f[(f >= 1 / (32 * 3600)) & (f <= 1 / (16 * 3600))]
            peak_pos = 1 / f_around_24[np.argmax(s_around_24)] / 3600

            self._fit_params = {'coefficients': a,
                                'noise std': sigma,
                                'order': len(a),
                                'peak': peak_pos,
                                'preprocessing': filter,
                                'period ci': ci}
            self._data = np.empty((2, len(f)))
            self._data[0] = f
            self._data[1] = s
            self._fit_params['id'] = id

        def __repr__(self):
            to_print = {key: self._fit_params[key] for key in self._fit_params.keys() - {'coefficients', 'noise std'}}
            return pprint.pformat(to_print, indent=4)

        def plot(self, out_path: str | None = None):
            """Plot the power spectrum in the low frequency range together with the position of the main peak.

            Args:
                out_path (str | None): Path to save the plot. If None, the plot will be opened in interactive window. Defaults to None.
            """
            plt.figure("Autoregressive model spectrum")
            plt.semilogx(self._data[0], self._data[1])
            plt.vlines([1 / 24 / 3600], [0], [np.max(self._data[1])], colors=['red'])
            plt.title("order {}, peak in {} h".format(self._fit_params['order'], round(self._fit_params['peak'], 3)))
            plt.xlabel('Frequency [Hz]')
            plt.ylabel(r'Power density [V$^2$/Hz]')
            if out_path is None:
                plt.show()
            else:
                plt.savefig(out_path)
                plt.close()

        @staticmethod
        def _construct_from_load(data: np.ndarray, fit: dict) -> "ARResults":
            out = ARResults(np.array([]), 0., np.array([1e-5]), np.array([1]), True, None, '')
            out._data = data
            out._fit_params = fit
            return out

        @staticmethod
        def load_file(path: str) -> "ARResults":
            """Load model fitting results from .ada.long file (created via export method).

            Args:
                path (str): Path to file.

            Returns:
                ARResults: Object containing fitted model.
            """
            data, fit_params = ARResults._generic_load(path)
            return ARResults._construct_from_load(data, fit_params)

        def export(self, out_path: str):
            """Save data (both parameters of model and spectrum data) to compressed generic file.

            Args:
                out_path (str): Path to the out file.
            """
            to_export = [self._data, self._fit_params]
            ARResults._generic_export(out_path, to_export)

        def save_csv(self, out_path: str):
            """Save parameters of fitted model to human-readable csv.

            Args:
                out_path (str): Path to the out file.
            """
            with open(out_path, 'w', newline='') as f:
                writer = csv.writer(f)
                for key in ['order', 'peak', 'period ci', 'preprocessing']:
                    writer.writerow([key, self._fit_params[key]])

        @property
        def freq(self) -> np.ndarray:
            """Frequency vector (x-axis of spectrum)."""
            return self._data[0]

        @property
        def spectrum(self) -> np.ndarray:
            """Power density calculated for given frequencies."""
            return self._data[1]

        @property
        def id(self) -> str:
            """ID of a recording to which long estimate was fit."""
            return self._fit_params['id']

    class Spectrum(_Long):
        """A class for fitting an autoregressive (AR) model to the actigraphic data, and estimating circadian rhythm on this basis.

        Args:
            order (int | None, optional): Order of the model. If None, it will be selected automatically based on a heuristic (order = number of samples spanning 30 hours in the data). Defaults to None.
            preprocessing (bool, optional): Whether to filter data before model fitting. Defaults to False.
            use_bootstrap (bool, optional): Whether to estimate period confidence interval using bootrstrap. Might take some time, especially for shorter epoch lenghts. Defaults to False.
        """
        def __init__(self, order: int | None = None, preprocessing: bool = False, use_bootstrap: bool = False):
            self._filter = preprocessing
            self._order = order
            self._use_bootstrap = use_bootstrap

        def fit(self, data: _ActiData, ch_name: str | None = None):
            """Fit an AR model to the provided actiraphic data. It is highly recomended to epoch them using Downsampler, as other methods might not work at all with AR.

            Args:
                data (_ActiData): Actigraphic data.
                ch_name (str | None, optional): Channel to which model will be fitted. to_score channel, when None. Defaults to None.

            Returns:
                ARResults: Object containing results of the model fitting.
            """
            y = self._get_data(data, ch_name)
            y = y - np.mean(y)
            if self._filter:
                sos = ss.iirdesign(1 / (3600 * 9), 1 / (3600 * 3), 2, 30, ftype='butter', output='sos', fs=data.fs)
                y = ss.sosfiltfilt(sos, y)

            if self._order is None:
                order = int(30 * 3600 * data.fs)
            else:
                order = self._order

            n = 75000
            f = np.linspace(1 / (8 * 3600), 1 / (40 * 3600), n)
            a, sigma = Spectrum._ar_params(y, order)
            _, s = Spectrum._ar_spectrum(a, sigma, n, data.fs, f)

            if self._use_bootstrap:
                p_value = self._bootstrap(y, n, f, order, data.fs)
            else:
                p_value = None

            return ARResults(a, sigma, f, s, self._filter, p_value, data.id)

        def _bootstrap(self, y: np.ndarray, n: int, f: np.ndarray, order: int, fs: float):
            n_iterations = int(order / 10)
            orders = np.arange(-n_iterations // 2, n_iterations // 2) + order
            out = np.empty(n_iterations)
            for i, ord in tqdm.tqdm(enumerate(orders), desc="Calculating confidence intervals using bootstrap", total=n_iterations):
                a, sigma = Spectrum._ar_params(y, ord)
                _, s = Spectrum._ar_spectrum(a, sigma, n, fs, f)
                s_around_24 = s[(f >= 1 / (32 * 3600)) & (f <= 1 / (16 * 3600))]
                f_around_24 = f[(f >= 1 / (32 * 3600)) & (f <= 1 / (16 * 3600))]
                peak_pos = 1 / f_around_24[np.argmax(s_around_24)] / 3600
                out[i] = peak_pos

            high = np.nanpercentile(out, 97.5)
            low = np.nanpercentile(out, 2.5)
            return (float(low), float(high))

        @staticmethod
        def _ar_params(x: np.ndarray, p: int):
            n = len(x)
            ak = np.correlate(x, x, mode='full')
            norm_ak = np.hstack((np.arange(1, n + 1, 1), np.arange(n - 1, 0, -1)))
            ak = ak / norm_ak
            r = ak[n - 1:]
            rl = r[1:1 + p]
            rp = np.zeros((p, p))
            for i in range(p):
                aa = ak[n - 1 - i:n - 1 - i + p]
                rp[i, :] = aa
            a = np.linalg.solve(rp, rl)
            sigma = (ak[n - 1] - np.sum(a * ak[n:n + p])) ** 0.5
            return a, sigma

        @staticmethod
        def _aic_criterium(x: np.ndarray, p_max: int):
            n = len(x)
            aic = np.zeros(p_max)
            for p in range(1, p_max + 1):
                _, sigma = Spectrum._ar_params(x, p)
                aic[p - 1] = 2 * p / n + np.log(sigma)
            return aic

        @staticmethod
        def _generate_ar(a: np.ndarray, sigma: float, n: int):
            x = np.zeros(n)
            p = len(a)
            for i in range(p, n):
                for j in range(len(a)):
                    x[i] += a[j] * x[i - (j + 1)]
                x[i] += sigma * np.random.randn()
            return x

        @staticmethod
        def _ar_spectrum(a: np.ndarray, sigma: float, n: int, fs: float, f: np.ndarray | None = None):
            if f is None:
                f = np.linspace(0, 0.001, n)
            z = np.exp(1j * 2 * np.pi * f / fs)
            A = 1 * np.ones(n) + 1j * np.zeros(n)
            for i in range(len(a)):
                A -= a[i] * z ** (-(i + 1))
            H = 1. / A
            sp = H * sigma ** 2 * H.conj()
            sp = sp / fs
            sp = sp.real
            return f, sp