# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/utils_metrics.ipynb (unless otherwise specified).

__all__ = ['TAMSD', 'get_D', 'get_exponent']

# Cell
import numpy as np

# Cell

def TAMSD(traj, t_lags):
    '''
    Calculates the time average mean squared displacement of a set of trajectories
    Inputs: - traj: trajectory to calculate TA-MSD
            - t_lags: time lags used for the TA-MSD
    Outputs: - TA-MSD
    '''
    tamsd = np.zeros_like(t_lags, dtype= float)
    for idx, t in enumerate(t_lags):
        for p in range(len(traj)-t):
            tamsd[idx] += (traj[p]-traj[p+t])**2
        tamsd[idx] /= len(traj)-t
    return tamsd

def get_D(traj, t_lags = None):
    if not t_lags:
        N_t_lags = max(4, int(len(traj)*0.1))
        t_lags = np.arange(1,N_t_lags)

    tasmd = TAMSD(traj, t_lags)
    return np.polyfit(t_lags, tasmd, deg = 1)[0]

def get_exponent(traj, t_lags = None):
    if not t_lags:
        N_t_lags = max(4, int(len(traj)*0.1))
        t_lags = np.arange(1,N_t_lags)

    tasmd = TAMSD(traj, t_lags)
    return np.polyfit(np.log(t_lags), np.log(tasmd), deg = 1)[0]

