SSTT-visuals 0.1.0__py3-none-any.whl

SSTT_visuals/__init__.py

@@ -0,0 +1,5 @@
+"""SSTT visualisation helpers for scan-statistic and occupancy summaries."""
+
+from .plot import statisticsVisual
+
+__all__ = ["statisticsVisual"]

SSTT_visuals/clusterStats.py

@@ -0,0 +1,298 @@
+"""
+Two goodness-of-fit tests against the null hypothesis that a 5x5 array of
+nonnegative-integer counts arises from N i.i.d. uniform draws over the 25
+cells (equivalently: the count vector is Multinomial(N, p=1/25 each)).
+
+1. occupancy_test:  T = number of empty cells. Reports P(T >= t_obs), exact.
+2. scan_statistic_test: Q = min over "2-by-n" rectangles of the rectangle's
+   own upper-tail binomial p-value. Reports P(Q <= q_obs) via simulation,
+   since the exact null distribution of a minimum over overlapping,
+   dependent windows has no tractable closed form.
+"""
+
+from fractions import Fraction
+from math import comb
+
+import scipy.stats as stats
+
+from .mos import normal_isf_asymptotic
+
+import numpy as np
+from scipy.stats import binom
+
+GRID_SIZE = 5
+CENTER = (GRID_SIZE // 2, GRID_SIZE // 2)  # (2, 2), zero-indexed
+
+
+# ---------------------------------------------------------------------------
+# Test 1: Occupancy test
+# ---------------------------------------------------------------------------
+
+def _surjections(n: int, j: int) -> int:
+    """Number of surjective functions from an n-set onto a j-set, via
+    inclusion-exclusion. Surj(n, 0) = 1 iff n == 0 (handled naturally by the
+    formula since 0**0 = 1 in Python)."""
+    total = 0
+    for i in range(j + 1):
+        total += (-1) ** i * comb(j, i) * (j - i) ** n
+    return total
+
+
+def occupancyStatisticTest(data: np.ndarray):
+    """
+    Exact p-value P(T >= t_obs) for the occupancy test, where T is the
+    number of empty cells among the m = 25 cells of `data`, under the null
+    that N = sum(data) observations are i.i.d. uniform over the 25 cells.
+
+    Returns the p-value (float).
+    """
+    data = np.asarray(data)
+    if data.shape != (GRID_SIZE, GRID_SIZE):
+        raise ValueError(f"expected a {GRID_SIZE}x{GRID_SIZE} array")
+
+    m = data.size            # 25 boxes
+    n = int(data.sum())      # total balls
+    t_obs = int(np.sum(data == 0))
+
+    p_tail = Fraction(0)
+    for k in range(t_obs, m + 1):
+        j = m - k  # number of "occupied" boxes required among the chosen k... 
+        # careful: k here is the number of EMPTY boxes, so m-k boxes must
+        # all be nonempty -> surjection of n balls onto (m-k) boxes
+        occupied = m - k
+        surj = _surjections(n, occupied)
+        p_tail += Fraction(comb(m, k) * surj, m ** n)
+
+    p_value = float(p_tail)
+    if p_value <= 0: 
+        z_score = normal_isf_asymptotic(p_tail)
+        lp = None
+    else:
+        z_score = stats.norm.isf(p_value)
+        lp = -np.log(p_value)
+    return (p_value, z_score, lp)
+
+
+# ---------------------------------------------------------------------------
+# Test 2: Scan statistic test
+# ---------------------------------------------------------------------------
+
+def _enumerate_rectangles(size: int = GRID_SIZE, center=CENTER):
+    """
+    All axis-aligned rectangles with one side exactly length 2 and the other
+    side length 2..size (both orientations: 2-tall / n-wide, and n-tall /
+    2-wide, with n >= 2 so both dimensions are at least 2), excluding any
+    rectangle that contains `center`.
+ 
+    Returns a list of (r0, r1, c0, c1, flat_indices, p) tuples where
+    (r0..r1, c0..c1) are inclusive cell ranges, flat_indices is the list of
+    flattened (row-major) cell indices in the rectangle, and p = ncells/25
+    is the null probability mass covered by the rectangle.
+    """
+    rects = set()
+ 
+    for w in range(2, size + 1):
+        h = 2
+        for r0 in range(size - h + 1):
+            for c0 in range(size - w + 1):
+                r1, c1 = r0 + h - 1, c0 + w - 1
+                if r0 <= center[0] <= r1 and c0 <= center[1] <= c1:
+                    continue
+                rects.add((r0, r1, c0, c1))
+ 
+    for h in range(2, size + 1):
+        w = 2
+        for r0 in range(size - h + 1):
+            for c0 in range(size - w + 1):
+                r1, c1 = r0 + h - 1, c0 + w - 1
+                if r0 <= center[0] <= r1 and c0 <= center[1] <= c1:
+                    continue
+                rects.add((r0, r1, c0, c1))
+ 
+    rect_info = []
+    for (r0, r1, c0, c1) in sorted(rects):
+        idx = [r * size + c for r in range(r0, r1 + 1) for c in range(c0, c1 + 1)]
+        p = len(idx) / (size * size)
+        rect_info.append((r0, r1, c0, c1, idx, p))
+    return rect_info
+ 
+ 
+# These are the two composite regions in grid coordinates. Each tuple keeps
+# the normal rectangle fields, followed by the cells in the union and the
+# component rectangles used by the plotter for highlighting.
+_SPECIAL_COMPOSITE_RECTS = [
+    (
+        0, 4, 0, 2,
+        [0, 1, 5, 6, 10, 11, 15, 16, 17, 20, 21, 22],
+        12 / 25,
+        ((3, 4, 0, 2), (0, 2, 0, 1)),
+    ),
+    (
+        0, 4, 2, 4,
+        [2, 3, 4, 7, 8, 9, 13, 14, 18, 19, 23, 24],
+        12 / 25,
+        ((0, 1, 2, 4), (2, 4, 3, 4)),
+    ),
+]
+
+
+_RECTS = _enumerate_rectangles() + _SPECIAL_COMPOSITE_RECTS
+ 
+
+def _compute_Q(flat_counts: np.ndarray, n_total: int) -> float:
+    """flat_counts: 1D array of length 25 (row-major). Returns Q = min_r q_r."""
+    q_min = 1.0
+    min_rect = None
+    for rect in _RECTS:
+        idx = rect[4]
+        p = rect[5]
+        s = flat_counts[idx].sum()
+        q_r = binom.sf(s - 1, n_total, p)  # P(X >= s)
+        if q_r < q_min:
+            q_min = q_r
+            min_rect = rect
+    return q_min, min_rect
+
+
+def _compute_Q_batch(flat_counts_batch: np.ndarray, n_total: int) -> np.ndarray:
+    """Vectorised version: flat_counts_batch has shape (n_sims, 25)."""
+    n_sims = flat_counts_batch.shape[0]
+    q_min = np.ones(n_sims)
+    for rect in _RECTS:
+        idx = rect[4]
+        p = rect[5]
+        s = flat_counts_batch[:, idx].sum(axis=1)
+        q_r = binom.sf(s - 1, n_total, p)
+        np.minimum(q_min, q_r, out=q_min)
+    return q_min
+
+
+def scanStatisticTest(data, n_sims: int = 2000000, random_state=None, return_rect: bool = False):
+    """
+    Monte Carlo p-value P(Q <= q_obs) for the scan statistic test.
+
+    Q = min over all valid "2-by-n" rectangles r of q_r, where q_r is the
+    upper-tail binomial probability P(count in r >= observed count in r)
+    under the null that N = sum(data) observations are i.i.d. uniform over
+    the 25 cells.
+
+    Returns the p-value, z-score, log-p-value, q_obs, and optionally the
+    rectangle with the minimum q-value.
+    """
+    data = np.asarray(data)
+    if data.shape != (GRID_SIZE, GRID_SIZE):
+        raise ValueError(f"expected a {GRID_SIZE}x{GRID_SIZE} array")
+
+    n_total = int(data.sum())
+    flat_obs = data.reshape(-1)
+    q_obs, min_rect = _compute_Q(flat_obs, n_total)
+
+    if n_total == 0:
+        # Every rectangle sum is trivially 0 >= 0, so q_r = 1 for all r,
+        # Q = 1 always, under the null too -> degenerate p-value of 1.
+        result = (q_obs, 1.0)
+        if return_rect:
+            return (1.0, 0.0, 0.0, q_obs, min_rect)
+        return result
+
+    rng = np.random.default_rng(random_state)
+    flat_p = np.full(GRID_SIZE * GRID_SIZE, 1.0 / (GRID_SIZE * GRID_SIZE))
+    sims = rng.multinomial(n_total, flat_p, size=n_sims)  # shape (n_sims, 25)
+
+    q_sims = _compute_Q_batch(sims, n_total)
+    p_value = float(np.mean(q_sims <= q_obs))
+
+    z_score = stats.norm.isf(p_value)
+    lp = np.inf if p_value == 0 else -np.log(p_value)
+
+    if return_rect:
+        return (p_value, z_score, lp, q_obs, min_rect)
+    return (p_value, z_score, lp, q_obs)
+
+def _min_tail_prob_at_least_as_extreme(n_total: int, p: float, q_threshold: float) -> float:
+    """
+    For S ~ Binomial(n_total, p), let q(s) = P(S >= s) be the upper-tail
+    p-value function (non-increasing, right-continuous step function of s).
+    Returns P(q(S) <= q_threshold), exactly, with no simulation: this is
+    just sf(s* - 1) where s* is the smallest integer count whose own tail
+    probability already drops to or below q_threshold.
+    """
+    s_values = np.arange(0, n_total + 1)
+    tail = binom.sf(s_values - 1, n_total, p)  # monotone non-increasing in s
+    mask = tail <= q_threshold
+    if not mask.any():
+        return 0.0
+    s_star = int(s_values[mask][0])
+    return float(binom.sf(s_star - 1, n_total, p))
+
+
+def scanStatisticTestBonferroni(data: np.ndarray, return_rect: bool = False):
+    """
+    Exact (analytic, no simulation) UPPER BOUND on P(Q <= q_obs), via the
+    union bound over the fixed set of 38 rectangles:
+
+        P(Q <= q) = P(exists r: q_r <= q) <= sum_r P(q_r <= q)
+
+    Unlike the Monte Carlo version, this has no floor on how small a
+    p-value it can report -- useful when the true p-value is far smaller
+    than 1/n_sims. It is conservative: the true p-value is <= this bound.
+
+    Returns (bound, z_score, lp, q_obs) by default, or includes the
+    rectangle with the minimum q-value when return_rect=True.
+    """
+    data = np.asarray(data)
+    if data.shape != (GRID_SIZE, GRID_SIZE):
+        raise ValueError(f"expected a {GRID_SIZE}x{GRID_SIZE} array")
+
+    n_total = int(data.sum())
+    flat_obs = data.reshape(-1)
+    q_obs, min_rect = _compute_Q(flat_obs, n_total)
+
+    if n_total == 0:
+        raise ValueError("cannot compute scan statistic for empty data (n_total=0)")
+
+    bound = 0.0
+    for rect in _RECTS:
+        p = rect[5]
+        bound += _min_tail_prob_at_least_as_extreme(n_total, p, q_obs)
+    bound = min(bound, 1.0)  # a union bound can nominally exceed 1
+
+    z_score = stats.norm.isf(bound)
+    lp = np.inf if bound == 0 else -np.log(bound)
+    if return_rect:
+        return (bound, z_score, lp, q_obs, min_rect)
+    return (bound, z_score, lp, q_obs)
+
+# ---------------------------------------------------------------------------
+# Demo
+# ---------------------------------------------------------------------------
+
+# if __name__ == "__main__":
+    # rng = np.random.default_rng(0)
+
+    # # Example: 40 observations thrown uniformly at random (should look "null-ish")
+    # example = rng.multinomial(40, np.full(25, 1 / 25)).reshape(5, 5)
+    # print("Example data:\n", example)
+
+    # t_obs = int(np.sum(example == 0))
+    # p_occ = occupancyStatisticTest(example)
+    # print(f"\nOccupancy test: t_obs = {t_obs}, P(T >= t_obs) = {p_occ:.4f}")
+
+    # q_obs, p_scan = scanStatisticTest(example, n_sims=2000000, random_state=1)
+    # print(f"Scan statistic test: q_obs = {q_obs:.4f}, P(Q <= q_obs) = {p_scan:.4f}")
+
+    # # Example designed to trip the tests: a clustered/skewed pattern
+    # skewed = np.array([
+    #     [10, 10, 0, 0, 0],
+    #     [10, 10, 0, 0, 0],
+    #     [0,  0,  0, 0, 0],
+    #     [0,  0,  0, 0, 0],
+    #     [0,  0,  0, 0, 0],
+    # ])
+    # print("\nSkewed data:\n", skewed)
+    # t_obs2 = int(np.sum(skewed == 0))
+    # p_occ2 = occupancy_test(skewed)
+    # print(f"\nOccupancy test: t_obs = {t_obs2}, P(T >= t_obs) = {p_occ2:.6f}")
+
+    # q_obs2, p_scan2 = scan_statistic_test(skewed, n_sims=20000, random_state=1)
+    # print(f"Scan statistic test: q_obs = {q_obs2:.6f}, P(Q <= q_obs) = {p_scan2:.6f}")

SSTT_visuals/mos.py

@@ -0,0 +1,293 @@
+import scipy.stats as stats
+import gmpy2
+from math import log, factorial, perm, sqrt, pi
+from typing import Self, overload
+from functools import cache
+from sys import getsizeof
+from time import perf_counter
+try:
+    from .primeslist import PRIMES
+except ImportError:  # pragma: no cover - compatibility fallback for direct script execution
+    from primeslist import PRIMES
+
+class PDFrac():
+    __slots__ = ('numerator', 'denominator', 'dmax', 'q')
+
+    @overload
+    def __init__(self, numerator : int): pass
+    @overload
+    def __init__(self, numerator : int, denominator : int): pass
+
+
+    def __init__(self, numerator, denominator = 1, dmax = None, q = 1):
+        # if not (type(numerator) is int\
+        # or type(numerator) is gmpy2.mpz)\
+        # or not (type(denominator) is int\
+        # or type(denominator) is gmpy2.mpz):
+        #     raise TypeError(f"PDFrac recieved {type(numerator)}, {type(denominator)}")
+        
+        # dmax represents the largest witnessed multiple of the denominator. This is
+        # a simple upper bound for the largest prime that could divide the denominator.
+
+        self.q = q
+        self.dmax = dmax if dmax is not None else denominator
+        self.numerator = gmpy2.mpz(numerator)
+        self.denominator = gmpy2.mpz(denominator)
+
+    def __add__(self, frac : Self|int):
+        if type(self) != PDFrac \
+        or type(frac) not in (PDFrac, int):
+            err_str = f"Addition defined only between PDFrac and PDFrac|int, "\
+                    + f"but received {type(self)}, {type(frac)}"
+            raise TypeError(err_str)
+        
+        if type(frac) == int:
+            frac = PDFrac(frac)
+
+        if self.numerator == 0:
+            return frac.copy()
+        if frac.numerator == 0:
+            return self.copy()
+        
+        res = self.copy()
+
+        res.numerator = res.numerator * frac.denominator + frac.numerator * res.denominator
+        res.denominator = res.denominator * frac.denominator
+        res.dmax = max(self.denominator, frac.denominator)
+        res.q = max(self.q, frac.q)
+        return res
+    
+    def __mul__(self, frac : Self|int):
+        if type(self) != PDFrac \
+        or type(frac) not in (PDFrac, int):
+            err_str = f"Multiplication defined only between PDFrac and PDFrac|int, "\
+                    + f"but received {type(self)}, {type(frac)}"
+            raise TypeError(err_str)
+        
+        if type(frac) == int:
+            frac = PDFrac(frac)
+
+        if self.numerator == 0 or frac.numerator == 0:
+            return PDFrac(0)
+        
+        if self.numerator == 1 and self.denominator == 1:
+            return frac.copy()
+        if frac.numerator == 1 and frac.denominator == 1:
+            return self.copy()
+        
+        res = self.copy()
+        
+        res.numerator *= frac.numerator
+        res.denominator *= frac.denominator
+        res.dmax = max(self.denominator, frac.denominator)
+        res.q = self.q + frac.q
+        return res
+        
+    def simplify(self):
+        largest_undivided_prime = 1
+        for i in range(len(PRIMES)):
+            p = PRIMES[i]
+            if p > self.dmax or p > self.numerator:
+                break
+
+            # if (p == 2 or self.q > 2) and p <= MAXP:
+            if self.q > 4 and p <= MAXP:
+                self.pOrderBSearch(p = p, i = i)
+                if self.denominator % p == 0:
+                    largest_undivided_prime = p
+                continue
+
+            while True:
+                if self.denominator % p != 0:
+                    break
+                if self.numerator % p != 0:
+                    largest_undivided_prime = p
+                    break
+                self.numerator //= p
+                self.denominator //= p
+        self.dmax = largest_undivided_prime
+    
+    def pOrderBSearch(self, p, i):
+        if self.numerator % p != 0 or self.denominator % p != 0:
+            return
+        
+        pk = PPOW[i][-1]
+        while self.numerator % pk == 0 and self.denominator % pk == 0:
+            self.denominator //= pk
+            self.numerator //= pk
+
+        self.pobs(i)
+        return
+
+    def pobs(self, i, d_low = None):
+        # d_low == True => (d | pk => n | pk).  Hence only check d divis.
+        # d_low == False => (n | pk => d | pk). Hence only check n divis.
+        # d_low is None implies nothing.        Hence check both
+
+        for pk in reversed(PPOW[i][:-1]):
+            if d_low is False:
+                d_divided = True
+                n_divided = (self.numerator % pk == 0)
+            elif d_low is True:
+                d_divided = (self.denominator % pk == 0)
+                n_divided = True
+            else:
+                d_divided = (self.denominator % pk == 0)
+                n_divided = (self.numerator % pk == 0)
+
+            if d_divided and n_divided:
+                self.denominator //= pk
+                self.numerator //= pk
+            elif d_low is None:
+                if d_divided:
+                    d_low = False
+                elif n_divided:
+                    d_low = True
+
+        return 
+
+    def to_float(self) -> float:
+        return float(self.numerator / self.denominator)
+
+    def copy(self) -> Self:
+        new = PDFrac.__new__(PDFrac)
+        new.numerator = self.numerator
+        new.denominator = self.denominator
+        new.dmax = self.dmax
+        new.q = self.q
+        return new
+
+def A(q: int, r: int, n: int, m: int) -> PDFrac:
+    q = gmpy2.mpz(q); r = gmpy2.mpz(r); n = gmpy2.mpz(n); m = gmpy2.mpz(m)
+    if q == 0:
+        return PDFrac(m**n)
+    rq = r * q
+    numerator = perm(n, rq)  # n! / (n-rq)!
+    denominator = factorial(r)**q * factorial(q)  # (r!)^q * q!
+    if m == q:
+        numerator *= factorial(m) # m!
+    else:
+        numerator *= perm(m, q) * (m - q)**(n - rq) # (m! / (m-q)!) * (m - q)^(n-rq)
+    return PDFrac(numerator, denominator, dmax=max(q, r), q=q)
+
+@cache # Cache output for dynamic programming
+def maxOrderStatistic(r : int, n : int, m : int) -> PDFrac:
+    # P(max urn <= r | marbles = n, urns = m)
+    #r : cdf
+    #n : number of balls
+    #m : number of urns
+    if r==0 and n!=0:
+        return PDFrac(0)
+    if r==0 and n==0:
+        return PDFrac(1)
+    if r >= n:
+        return PDFrac(1)
+    if n==0:
+        return PDFrac(1)
+
+    if r == 1:
+        if m >= n:
+            return PDFrac(perm(m, n), m**n)
+        else:
+            return PDFrac(0)
+    else:
+        probability = PDFrac(0)
+        qMax = n//r + 1
+        qMin = max(0, n - r*m + m)
+        for q in range(qMin, qMax):
+            probability += A(q, r, n, m) * maxOrderStatistic(r-1, n-r*q, m-q)
+            # probability.simplify()
+
+        probability *= PDFrac(1, m**n)
+        probability.simplify()
+
+        return probability
+    
+# Very small p_values <~ 1e-16 can't be captured by conjugating a float due to the scale mismatch
+# between the p_value and 1. Hence, conjugation must be performed on the fraction before converting
+# it to a float.
+def conjMaxOrderStatistic(s, marbles, urns):
+    q_frac = maxOrderStatistic(s - 1, marbles, urns) # q = P(max urn < s | marbles, urns)
+    p_frac = q_frac.copy() # Cached mutable must be copied pre-edit
+    p_frac.numerator = p_frac.denominator - p_frac.numerator
+    return p_frac, q_frac
+
+def PandZ(urnMax : int, marbles : int, urns : int, initialise = True) -> tuple[float, float, float]:
+    # urnMax = the number of marbles in the urn with the maximum number
+    # marbles = how many marbles did you place in the urns
+    # urns = the number of urns the marbles were distributed between
+
+    urnMax = int(urnMax); marbles = int(marbles); urns = int(urns)
+
+    if urnMax > 3000 or marbles > 6000:
+        raise NotImplementedError("Only supports urnMax <= P and marbles <= 2P, where P " +\
+                                  "is the largest prime in primeslist.PRIMES\n" +\
+                                  "Currently, P = 3000.")
+
+    if initialise:
+        initPPOW(prime_max = max(marbles, urns))
+
+    p_value_frac, _ = conjMaxOrderStatistic(urnMax, marbles, urns)
+    p_value = p_value_frac.to_float()
+    if p_value <= 0: 
+        # print(f"Warning: nonpositive p-value {p_value:.2e}")
+        # p_value = 1e-307 
+        z_score = normal_isf_asymptotic(p_value_frac)
+        lp = None
+    else:
+        z_score = stats.norm.isf(p_value)
+        lp = -log(p_value)
+    return (p_value, z_score, lp)
+
+def initPPOW(prime_max, pow_max = 1e100, metrics = False):
+    global PPOW
+    global MAXP
+    PPOW = []
+    MAXP = prime_max
+
+    if metrics:
+        size = 0
+        pow_start = perf_counter()
+    for prime in PRIMES:
+        if prime > MAXP:
+            break
+        pow_list = []
+        pk = prime
+        while pk < pow_max:
+            if metrics:
+                size += getsizeof(pk)
+            pow_list.append(pk)
+            pk = pk**2
+        PPOW.append(pow_list)
+
+    if metrics:
+        pow_end = perf_counter()
+        print(f"Prime time: {pow_end-pow_start:.3f}")
+        print(f"Prime memory: {size}")
+
+    return
+
+def normal_isf_asymptotic(frac, iters=4):
+    L = gmpy2.log(frac.denominator) - gmpy2.log(frac.numerator)
+    x2 = 2 * L
+    for _ in range(iters):
+        x = sqrt(x2)
+        x2 = 2*L - 2*log(x) - log(2*pi) + 2/x2
+    return sqrt(x2)
+
+def clear_caches():
+    maxOrderStatistic.cache_clear()
+    return
+
+def cache_size_MB():
+    return maxOrderStatistic.cache_info()[3]/ (1024 * 1024)
+
+def main():
+    initPPOW(100, metrics = True)
+    start = perf_counter()
+    PandZ(50, 100, 100) # Takes 1.6-1.7 seconds on a 3.8GHz Intel Core i5
+    end = perf_counter()
+    print(f"Total time: {end-start:.3f}")
+
+if __name__ == "__main__":
+    main()