# Pocock boundaries for K equally spaced looks, two-sided overall alpha 0.05 (normal approximation):
# find the constant c so that P(|Z_k| < c at every look) = 0.95 when there is no effect, where Z_k is the
# cumulative z-statistic at look k (correlation sqrt(j/k) between looks j < k). Python 3.13, NumPy, SciPy.
import numpy as np
from scipy import stats, optimize

def pocock_p(K, alpha=0.05):
    cov = np.array([[np.sqrt(min(i, j) / max(i, j)) for j in range(1, K + 1)] for i in range(1, K + 1)])
    mvn = stats.multivariate_normal(np.zeros(K), cov)
    def none_cross(c):                                  # P(all |Z_k| < c) by inclusion-exclusion over signs
        lo, hi = -c * np.ones(K), c * np.ones(K)
        total = 0.0
        for mask in range(2 ** K):
            pt = np.where([(mask >> k) & 1 for k in range(K)], lo, hi)
            total += (-1) ** bin(mask).count("1") * mvn.cdf(pt)
        return total
    c = optimize.brentq(lambda c: none_cross(c) - (1 - alpha), 1.5, 3.5, xtol=1e-6)
    return 2 * stats.norm.sf(c)

for K in (2, 3):
    print(f"{K} looks: reject at any look if p < {pocock_p(K):.4f}")
