1import json, time
  2import numpy as np
  3
  4
  5def make_dictionary(m, n, d, seed=0):
  6    rng = np.random.default_rng(seed)
  7    rows = np.empty(d*n, dtype=np.int64)
  8    cols = np.repeat(np.arange(n), d)
  9    weights = np.empty(d*n, dtype=np.float64)
 10    for j in range(n):
 11        # Sampling without replacement gives exactly d active coordinates per column.
 12        rr = rng.choice(m, size=d, replace=False)
 13        sl = slice(j*d, (j+1)*d)
 14        rows[sl] = rr
 15        weights[sl] = rng.normal(size=d)
 16        weights[sl] /= np.linalg.norm(weights[sl])
 17    W = np.zeros((m, n), dtype=np.float64)
 18    W[rows, cols] = weights
 19    return W, rows, cols, weights
 20
 21
 22def edge_correlations(R, rows, cols, weights):
 23    # R is [batch,m]. This is the direct gather/multiply/reduce realization.
 24    C = np.zeros((R.shape[0], int(cols.max()) + 1), dtype=R.dtype)
 25    for e in range(len(rows)):
 26        C[:, cols[e]] += R[:, rows[e]] * weights[e]
 27    return C
 28
 29
 30def dense_omp(H, W, k, edge=None, lam=1e-9):
 31    B, m = H.shape
 32    n = W.shape[1]
 33    R = H.copy()
 34    X = np.zeros((B, n))
 35    supports = [[] for _ in range(B)]
 36    for _ in range(k):
 37        C = W.T @ R.T
 38        C = C.T
 39        for b in range(B):
 40            if supports[b]: C[b, supports[b]] = 0.0
 41            j = int(np.argmax(np.abs(C[b])))
 42            supports[b].append(j)
 43            S = supports[b]
 44            WS = W[:, S]
 45            # Ridge only protects the intentionally tiny toy cases from singularity.
 46            G = WS.T @ WS + lam*np.eye(len(S))
 47            coef = np.linalg.solve(G, WS.T @ H[b])
 48            X[b, S] = coef
 49            R[b] = H[b] - WS @ coef
 50    return X, R, supports
 51
 52
 53def edge_omp(H, W, rows, cols, weights, k, lam=1e-9):
 54    B = H.shape[0]
 55    n = W.shape[1]
 56    R = H.copy()
 57    X = np.zeros((B, n))
 58    supports = [[] for _ in range(B)]
 59    for _ in range(k):
 60        C = edge_correlations(R, rows, cols, weights)
 61        for b in range(B):
 62            if supports[b]: C[b, supports[b]] = 0.0
 63            j = int(np.argmax(np.abs(C[b])))
 64            supports[b].append(j)
 65            S = supports[b]
 66            WS = W[:, S]
 67            coef = np.linalg.solve(WS.T @ WS + lam*np.eye(len(S)), WS.T @ H[b])
 68            X[b, S] = coef
 69            R[b] = H[b] - WS @ coef
 70    return X, R, supports
 71
 72
 73def timed(fn, repeats=3):
 74    vals = []
 75    for _ in range(repeats):
 76        t = time.perf_counter(); fn(); vals.append(time.perf_counter()-t)
 77    return float(np.median(vals))
 78
 79
 80def main():
 81    rng = np.random.default_rng(123)
 82    # Core algebra check, including batched residuals.
 83    W, rows, cols, weights = make_dictionary(48, 96, 5, 4)
 84    R = rng.normal(size=(7, 48))
 85    cd = R @ W
 86    ce = edge_correlations(R, rows, cols, weights)
 87    corr_err = float(np.max(np.abs(cd-ce)))
 88
 89    # Same observations and exact same refit rule for both implementations.
 90    m, n, d, B, k = 96, 384, 6, 24, 6
 91    W, rows, cols, weights = make_dictionary(m, n, d, 9)
 92    true_x = np.zeros((B, n))
 93    for b in range(B):
 94        true_x[b, rng.choice(n, k, replace=False)] = rng.normal(size=k)
 95    H = true_x @ W.T + 0.01*rng.normal(size=(B,m))
 96    xd, rd, sd = dense_omp(H, W, k)
 97    xe, re, se = edge_omp(H, W, rows, cols, weights, k)
 98    dense_mse = float(np.mean(rd**2)); edge_mse = float(np.mean(re**2))
 99    support_equal = float(np.mean([set(a)==set(b) for a,b in zip(sd,se)]))
100    xdiff = float(np.max(np.abs(xd-xe)))
101
102    # Correlation-only scaling/cost comparison. Dense BLAS is the control;
103    # edge implementation is intentionally transparent Python gather-reduce.
104    scaling = []
105    for nn in (128, 256, 512, 1024):
106        mm, dd, bb = 128, 8, 32
107        WW, rr, cc, ww = make_dictionary(mm, nn, dd, nn+1)
108        RR = rng.normal(size=(bb, mm))
109        td = timed(lambda: RR @ WW)
110        te = timed(lambda: edge_correlations(RR, rr, cc, ww))
111        scaling.append({"n":nn, "dense_sec":td, "edge_sec":te,
112                        "edge_over_dense":te/td,
113                        "dense_flops_proxy":mm*nn*bb,
114                        "edge_flops_proxy":dd*nn*bb})
115    out = {"max_correlation_abs_error":corr_err,
116           "omp_max_coefficient_abs_diff":xdiff,
117           "support_set_agreement":support_equal,
118           "dense_reconstruction_mse":dense_mse,
119           "edge_reconstruction_mse":edge_mse,
120           "theoretical_correlation_work_ratio_d_over_m":d/m,
121           "scaling":scaling}
122    print(json.dumps(out, indent=2))
123
124if __name__ == '__main__': main()