Get breaking index of state persistence for sequences of length more than tau

persistent_sequence.py

import numpy as np
import numpy.typing as npt


def simulate(P:npt.NDArray, T:int, s0:int = 0) -> np.ndarray:
    'Simulates with transition matrix P for T steps, starting at s0 = 0.'

    u = np.random.uniform(0., 1., T)
    S = np.empty(T, dtype = int)
    S[0] = s0

    for i in range(1, T):
        s = S[i - 1]
        S[i] = s if u[i] < P[s, s] else 1 - s
    
    return S


def persistent_trials(S:npt.ArrayLike, tau:int) -> np.ndarray:
    'Get breaking index of state persistence for sequences of length more than tau'
    
    breakpoints = np.where(np.diff(S) != 0)[0] + 1
    relevant_chunks = [ len(chunk) >= tau for chunk in np.split(S, breakpoints)[:-1] ]

    return breakpoints[relevant_chunks]
    

if __name__ == "__main__":
    S = np.array([0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 1])
    persistent_trials(S, 3)

    S = simulate(np.ones((2, 2)) / 2, 10_000_000)
    persistent_trials(S, 3)