Analysis API Reference

def column_cycle_array(posterior, amt=None):
    """
    Cycle each column of the posterior matrix by a random or specified amount.

    Parameters
    ----------
    posterior : np.ndarray
        Posterior probability matrix (position x time).
    amt : array-like or None, optional
        Amount to cycle each column. If None, random cycling is used.

    Returns
    -------
    out : np.ndarray
        Cycled posterior matrix.
    """
    out = copy.deepcopy(posterior)
    rows, cols = posterior.shape

    if amt is None:
        for col in range(cols):
            if np.isnan(np.sum(posterior[:, col])):
                continue
            else:
                out[:, col] = np.roll(posterior[:, col], np.random.randint(1, rows))
    else:
        if len(amt) == cols:
            for col in range(cols):
                if np.isnan(np.sum(posterior[:, col])):
                    continue
                else:
                    out[:, col] = np.roll(posterior[:, col], int(amt[col]))
        else:
            raise TypeError("amt does not seem to be the correct shape!")
    return out

def fmpt(P):
    """
    Calculates the matrix of first mean passage times for an
    ergodic transition probability matrix.
    Parameters
    ----------
    P    : array (kxk)
           an ergodic Markov transition probability matrix
    Returns
    -------
    M    : array (kxk)
           elements are the expected value for the number of intervals
           required for  a chain starting in state i to first enter state j
           If i=j then this is the recurrence time.

    Examples
    --------
    >>> import numpy as np
    >>> p = np.array([[0.5, 0.25, 0.25], [0.5, 0, 0.5], [0.25, 0.25, 0.5]])
    >>> fm = fmpt(p)
    >>> fm
    array([[ 2.5       ,  4.        ,  3.33333333],
           [ 2.66666667,  5.        ,  2.66666667],
           [ 3.33333333,  4.        ,  2.5       ]])
    Thus, if it is raining today in Oz we can expect a nice day to come
    along in another 4 days, on average, and snow to hit in 3.33 days. We can
    expect another rainy day in 2.5 days. If it is nice today in Oz, we would
    experience a change in the weather (either rain or snow) in 2.67 days from
    today. (That wicked witch can only die once so I reckon that is the
    ultimate absorbing state).

    Notes -----
    Uses formulation (and examples on p. 218) in Kemeny and Snell (1976) [1]_
    References
    ----------

    .. [1] Kemeny, John, G. and J. Laurie Snell (1976) Finite Markov
       Chains. Springer-Verlag. Berlin
    """
    P = np.asarray(P)
    A = np.zeros_like(P)
    ss = steady_state(P)
    k = ss.shape[0]
    for i in range(k):
        A[:, i] = ss.flatten()
    A = A.T
    identity_matrix = np.identity(k)
    Z = la.inv(identity_matrix - P + A)
    E = np.ones_like(Z)
    D = np.diag(1.0 / np.diag(A))
    Zdg = np.diag(np.diag(Z))
    M = (identity_matrix - Z + np.multiply(E, Zdg)) @ D
    return M

def get_significant_events(scores, shuffled_scores, q=95):
    """
    Return the significant events based on percentiles.

    Parameters
    ----------
    scores : np.ndarray
        Scores for each event.
    shuffled_scores : np.ndarray
        Shuffled scores for each event and shuffle.
    q : float, optional
        Percentile to compute (default is 95).

    Returns
    -------
    sig_event_idx : np.ndarray
        Indices of significant events.
    pvalues : np.ndarray
        Monte Carlo p-values for each event.
    """

    n, _ = shuffled_scores.shape
    r = np.sum(shuffled_scores >= scores, axis=0)
    pvalues = (r + 1) / (n + 1)

    sig_event_idx = np.argwhere(
        scores > np.percentile(shuffled_scores, axis=0, q=q)
    ).squeeze()

    return np.atleast_1d(sig_event_idx), np.atleast_1d(pvalues)

def linregress_array(posterior):
    """
    Perform linear regression on the posterior matrix, and return the slope, intercept, and R^2 value.

    Parameters
    ----------
    posterior : np.ndarray
        Posterior probability matrix (position x time).

    Returns
    -------
    slope : float
        Slope of the best-fit line.
    intercept : float
        Intercept of the best-fit line.
    r2 : float
        R^2 value of the fit.
    """

    mode_pth = get_mode_pth_from_array(posterior)

    y = mode_pth
    x = np.arange(len(y))
    x = x[~np.isnan(y)]
    y = y[~np.isnan(y)]

    if len(y) > 0:
        slope, intercept, rvalue, pvalue, stderr = stats.linregress(x, y)
        return slope, intercept, rvalue**2
    else:
        return np.nan, np.nan, np.nan

def linregress_bst(bst, tuningcurve):
    """
    Perform linear regression on all the events in bst, and return the slopes, intercepts, and R^2 values.

    Parameters
    ----------
    bst : BinnedSpikeTrainArray
        Binned spike train array containing all candidate events.
    tuningcurve : TuningCurve1D
        Tuning curve for decoding.

    Returns
    -------
    slopes : np.ndarray
        Slopes for each event.
    intercepts : np.ndarray
        Intercepts for each event.
    r2values : np.ndarray
        R^2 values for each event.
    """

    posterior, bdries, mode_pth, mean_pth = decode(bst=bst, ratemap=tuningcurve)

    slopes = np.zeros(bst.n_epochs)
    intercepts = np.zeros(bst.n_epochs)
    r2values = np.zeros(bst.n_epochs)
    for idx in range(bst.n_epochs):
        y = mode_pth[bdries[idx] : bdries[idx + 1]]
        x = np.arange(bdries[idx], bdries[idx + 1], step=1)
        x = x[~np.isnan(y)]
        y = y[~np.isnan(y)]

        if len(y) > 0:
            slope, intercept, rvalue, pvalue, stderr = stats.linregress(x, y)
            slopes[idx] = slope
            intercepts[idx] = intercept
            r2values[idx] = rvalue**2
        else:
            slopes[idx] = np.nan
            intercepts[idx] = np.nan
            r2values[idx] = np.nan  #
    #     if bst.n_epochs == 1:
    #         return np.asscalar(slopes), np.asscalar(intercepts), np.asscalar(r2values)
    return slopes, intercepts, r2values

def linregress_ting(bst, tuningcurve, n_shuffles=250):
    """
    Perform linear regression on all the events in bst, and return the R^2 values.

    Parameters
    ----------
    bst : BinnedSpikeTrainArray
        Binned spike train array containing all candidate events.
    tuningcurve : TuningCurve1D
        Tuning curve for decoding.
    n_shuffles : int, optional
        Number of shuffles for the null distribution (default is 250).

    Returns
    -------
    r2values : np.ndarray
        R^2 values for each event.
    r2values_shuffled : np.ndarray
        Shuffled R^2 values for each event and shuffle.
    """

    if float(n_shuffles).is_integer:
        n_shuffles = int(n_shuffles)
    else:
        raise ValueError("n_shuffles must be an integer!")

    posterior, bdries, mode_pth, mean_pth = decode(bst=bst, ratemap=tuningcurve)

    #     bdries = np.insert(np.cumsum(bst.lengths), 0, 0)
    r2values = np.zeros(bst.n_epochs)
    r2values_shuffled = np.zeros((n_shuffles, bst.n_epochs))
    for idx in range(bst.n_epochs):
        y = mode_pth[bdries[idx] : bdries[idx + 1]]
        x = np.arange(bdries[idx], bdries[idx + 1], step=1)
        x = x[~np.isnan(y)]
        y = y[~np.isnan(y)]

        if len(y) > 0:
            slope, intercept, rvalue, pvalue, stderr = stats.linregress(x, y)
            r2values[idx] = rvalue**2
        else:
            r2values[idx] = np.nan  #
        for ss in range(n_shuffles):
            if len(y) > 0:
                slope, intercept, rvalue, pvalue, stderr = stats.linregress(
                    np.random.permutation(x), y
                )
                r2values_shuffled[ss, idx] = rvalue**2
            else:
                r2values_shuffled[ss, idx] = (
                    np.nan
                )  # event contained NO decoded activity... unlikely or even impossible with current code

    #     sig_idx = np.argwhere(r2values[0,:] > np.percentile(r2values, q=q, axis=0))
    #     np.argwhere(((R2[1:,:] >= R2[0,:]).sum(axis=0))/(R2.shape[0]-1)<0.05) # equivalent to above
    if n_shuffles > 0:
        return r2values, r2values_shuffled
    return r2values

def pooled_time_swap_bst(bst):
    """
    Time swap on BinnedSpikeTrainArray, swapping within entire bst.

    Parameters
    ----------
    bst : BinnedSpikeTrainArray
        Binned spike train array to swap.

    Returns
    -------
    out : BinnedSpikeTrainArray
        Time-swapped spike train array.
    """
    out = copy.deepcopy(bst)  # should this be deep? YES! Oh my goodness, yes!
    shuffled = np.random.permutation(bst.n_bins)
    out._data = out._data[:, shuffled]
    return out

def score_hmm_logprob_cumulative(bst, hmm, normalize=False):
    """
    Score events in a BinnedSpikeTrainArray by computing the cumulative log probability under the model.

    Parameters
    ----------
    bst : BinnedSpikeTrainArray
        Binned spike train array containing all candidate events.
    hmm : PoissonHMM
        Trained hidden Markov model.
    normalize : bool, optional
        If True, log probabilities will be normalized by their sequence lengths.

    Returns
    -------
    logprob : np.ndarray
        Cumulative log probabilities for each event in bst.
    """

    logprob = np.atleast_1d(hmm._cum_score_per_bin(bst))
    if normalize:
        cumlengths = []
        for evt in bst.lengths:
            cumlengths.extend(np.arange(1, evt + 1).tolist())
        cumlengths = np.array(cumlengths)
        logprob = np.atleast_1d(logprob) / cumlengths

    return logprob

def score_hmm_time_resolved(bst, hmm, n_shuffles=250, normalize=False):
    """
    Score sequences using a hidden Markov model and a model where the transition probability matrix has been shuffled (time-resolved).

    Parameters
    ----------
    bst : BinnedSpikeTrainArray
        Binned spike train array containing all candidate events.
    hmm : PoissonHMM
        Trained hidden Markov model.
    n_shuffles : int, optional
        Number of shuffles for the null distribution. Default is 250.
    normalize : bool, optional
        If True, normalize the scores by event lengths.

    Returns
    -------
    scores : np.ndarray
        Log probabilities for each event.
    shuffled : np.ndarray
        Shuffled log probabilities for each event and shuffle.
    """

    if float(n_shuffles).is_integer:
        n_shuffles = int(n_shuffles)
    else:
        raise ValueError("n_shuffles must be an integer!")

    hmm_shuffled = copy.deepcopy(hmm)
    Lbraw = score_hmm_logprob_cumulative(bst=bst, hmm=hmm, normalize=normalize)

    # per event, compute L(:b|raw) - L(:b-1|raw)
    Lb = copy.deepcopy(Lbraw)

    cumLengths = np.cumsum(bst.lengths)
    cumLengths = np.insert(cumLengths, 0, 0)

    for ii in range(bst.n_epochs):
        LE = cumLengths[ii]
        RE = cumLengths[ii + 1]
        Lb[LE + 1 : RE] -= Lbraw[LE : RE - 1]

    n_bins = bst.n_bins
    shuffled = np.zeros((n_shuffles, n_bins))
    for ii in range(n_shuffles):
        hmm_shuffled.transmat_ = shuffle_transmat(hmm_shuffled.transmat_)
        Lbtmat = score_hmm_logprob_cumulative(
            bst=bst, hmm=hmm_shuffled, normalize=normalize
        )

        # per event, compute L(:b|tmat) - L(:b-1|raw)
        NL = copy.deepcopy(Lbtmat)
        for jj in range(bst.n_epochs):
            LE = cumLengths[jj]
            RE = cumLengths[jj + 1]
            NL[LE + 1 : RE] -= Lbraw[LE : RE - 1]

        shuffled[ii, :] = NL

    scores = Lb

    return scores, shuffled

def steady_state(P):
    """
    Calculates the steady state probability vector for a regular Markov
    transition matrix P
    Parameters
    ----------
    P        : array (kxk)
               an ergodic Markov transition probability matrix
    Returns
    -------
    implicit : array (kx1)
               steady state distribution
    Examples
    --------
    Taken from Kemeny and Snell. [1]_ Land of Oz example where the states are
    Rain, Nice and Snow - so there is 25 percent chance that if it
    rained in Oz today, it will snow tomorrow, while if it snowed today in
    Oz there is a 50 percent chance of snow again tomorrow and a 25
    percent chance of a nice day (nice, like when the witch with the monkeys
    is melting).
    >>> import numpy as np
    >>> p = np.array([[0.5, 0.25, 0.25], [0.5, 0, 0.5], [0.25, 0.25, 0.5]])
    >>> steady_state(p)
    array([[ 0.4],
           [ 0.2],
           [ 0.4]])
    Thus, the long run distribution for Oz is to have 40 percent of the
    days classified as Rain, 20 percent as Nice, and 40 percent as Snow
    (states are mutually exclusive).
    """

    v, d = la.eig(np.transpose(P))

    # for a regular P maximum eigenvalue will be 1
    mv = max(v.real)  # Use real part for comparison
    # find its position
    i = v.real.tolist().index(mv)

    # normalize eigenvector corresponding to the eigenvalue 1
    # Take real part to avoid complex warning
    eigenvector = d[:, i].real
    return eigenvector / np.sum(eigenvector)

def time_swap_array(posterior):
    """
    Time swap.

    Note: it is often possible to simply shuffle the time bins, and not the actual data, for computational
    efficiency. Still, this function works as expected.

    Parameters
    ----------
    posterior : np.ndarray
        Posterior probability matrix (position x time).

    Returns
    -------
    out : np.ndarray
        Time-swapped posterior matrix.
    """
    out = copy.deepcopy(posterior)
    rows, cols = posterior.shape

    colidx = np.arange(cols)
    shuffle_cols = np.random.permutation(colidx)
    out = out[:, shuffle_cols]

    return out

def trajectory_score_array(
    posterior, slope=None, intercept=None, w=None, weights=None, normalize=False
):
    """
    Compute the trajectory score for a given posterior matrix and line parameters.

    Parameters
    ----------
    posterior : np.ndarray
        Posterior probability matrix (position x time).
    slope : float, optional
        Slope of the line. If None, estimated from the data.
    intercept : float, optional
        Intercept of the line. If None, estimated from the data.
    w : int, optional
        Half band width for calculating the trajectory score. Default is 0.
    weights : array-like, optional
        Weights for the band around the line (not yet implemented).
    normalize : bool, optional
        If True, normalize the score by the number of non-NaN bins.

    Returns
    -------
    score : float
        Trajectory score for the event.
    """

    rows, cols = posterior.shape

    if w is None:
        w = 0
    if not float(w).is_integer:
        raise ValueError("w has to be an integer!")
    if slope is None or intercept is None:
        slope, intercept, _ = linregress_array(posterior=posterior)

    x = np.arange(cols)
    line_y = np.round((slope * x + intercept))  # in position bin #s

    # idea: cycle each column so that the top w rows are the band surrounding the regression line

    if np.isnan(slope):  # this will happen if we have 0 or only 1 decoded bins
        return np.nan
    else:
        temp = column_cycle_array(posterior, -line_y + w)

    if normalize:
        num_non_nan_bins = round(np.nansum(posterior))
    else:
        num_non_nan_bins = 1

    return np.nansum(temp[: 2 * w + 1, :]) / num_non_nan_bins

def trajectory_score_bst(
    bst, tuningcurve, w=None, n_shuffles=250, weights=None, normalize=False
):
    """
    Compute the trajectory scores from Davidson et al. for each event in the BinnedSpikeTrainArray.

    Parameters
    ----------
    bst : BinnedSpikeTrainArray
        Binned spike train array containing all candidate events.
    tuningcurve : TuningCurve1D
        Tuning curve to decode events in bst.
    w : int, optional
        Half band width for calculating the trajectory score. Default is 0.
    n_shuffles : int, optional
        Number of shuffles for the null distribution. Default is 250.
    weights : array-like, optional
        Weights for the band around the line (not yet implemented).
    normalize : bool, optional
        If True, normalize the score by the number of non-NaN bins.

    Returns
    -------
    scores : np.ndarray
        Trajectory scores for each event.
    scores_time_swap : np.ndarray
        Shuffled scores using time swap.
    scores_col_cycle : np.ndarray
        Shuffled scores using column cycle.
    """

    if w is None:
        w = 0
    if not float(w).is_integer:
        raise ValueError("w has to be an integer!")

    if float(n_shuffles).is_integer:
        n_shuffles = int(n_shuffles)
    else:
        raise ValueError("n_shuffles must be an integer!")

    posterior, bdries, mode_pth, mean_pth = decode(bst=bst, ratemap=tuningcurve)

    # idea: cycle each column so that the top w rows are the band
    # surrounding the regression line

    scores = np.zeros(bst.n_epochs)
    if n_shuffles > 0:
        scores_time_swap = np.zeros((n_shuffles, bst.n_epochs))
        scores_col_cycle = np.zeros((n_shuffles, bst.n_epochs))

    for idx in range(bst.n_epochs):
        posterior_array = posterior[:, bdries[idx] : bdries[idx + 1]]
        scores[idx] = trajectory_score_array(
            posterior=posterior_array, w=w, normalize=normalize
        )
        for shflidx in range(n_shuffles):
            # time swap:

            posterior_ts = time_swap_array(posterior_array)
            posterior_cs = column_cycle_array(posterior_array)
            scores_time_swap[shflidx, idx] = trajectory_score_array(
                posterior=posterior_ts, w=w, normalize=normalize
            )
            scores_col_cycle[shflidx, idx] = trajectory_score_array(
                posterior=posterior_cs, w=w, normalize=normalize
            )

    if n_shuffles > 0:
        return scores, scores_time_swap, scores_col_cycle
    return scores