1import json
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 1491
  6rng = np.random.default_rng(SEED)
  7
  8
  9def normalize(z):
 10    return z / np.maximum(np.linalg.norm(z, axis=1, keepdims=True), 1e-12)
 11
 12
 13def cosine_cost(X, Y):
 14    return 1.0 - normalize(X) @ normalize(Y).T
 15
 16
 17def pairwise_dist(X):
 18    d = X[:, None, :] - X[None, :, :]
 19    return np.sqrt(np.maximum(np.sum(d * d, axis=-1), 0.0))
 20
 21
 22def structural_cost(A, B, T):
 23    # L_ij(T) = sum_kl (A_ik-B_jl)^2 T_kl.
 24    m, n = T.shape
 25    L = np.empty((m, n))
 26    for i in range(m):
 27        for j in range(n):
 28            L[i, j] = np.sum((A[i, :, None] - B[j, None, :]) ** 2 * T)
 29    return L
 30
 31
 32def partial_gw(X, Y, alpha=1.0, beta=1.0, eps=0.12, iters=8,
 33               A=None, B=None, mu=None, nu=None):
 34    m, n = len(X), len(Y)
 35    mu = np.ones(m) / m if mu is None else np.asarray(mu, dtype=float)
 36    nu = np.ones(n) / n if nu is None else np.asarray(nu, dtype=float)
 37    C = cosine_cost(X, Y)
 38    A = pairwise_dist(X) if A is None else A
 39    B = pairwise_dist(Y) if B is None else B
 40    T = np.exp(np.clip(-alpha * C / eps, -80, 20))
 41    for _ in range(iters):
 42        L = structural_cost(A, B, T)
 43        T = np.exp(np.clip(-(alpha * C + beta * L) / eps, -80, 20))
 44        # Alternating upper-bound projections; scaling factors are clipped at one.
 45        for _ in range(25):
 46            T *= np.minimum(1.0, mu / np.maximum(T.sum(axis=1), 1e-30))[:, None]
 47            T *= np.minimum(1.0, nu / np.maximum(T.sum(axis=0), 1e-30))[None, :]
 48    return T, C, A, B
 49
 50
 51def softmax_attention(X, Y, eps=0.12):
 52    z = -cosine_cost(X, Y) / eps
 53    z -= z.max(axis=1, keepdims=True)
 54    p = np.exp(z)
 55    return p / p.sum(axis=1, keepdims=True)
 56
 57
 58def balanced_sinkhorn(X, Y, eps=0.12, iters=100):
 59    K = np.exp(np.clip(-cosine_cost(X, Y) / eps, -80, 20))
 60    a, b = np.ones(len(X)) / len(X), np.ones(len(Y)) / len(Y)
 61    u, v = np.ones(len(X)), np.ones(len(Y))
 62    for _ in range(iters):
 63        u = a / np.maximum(K @ v, 1e-30)
 64        v = b / np.maximum(K.T @ u, 1e-30)
 65    return u[:, None] * K * v[None, :]
 66
 67
 68def make_case(n=12, distractors=8, noise=0.01):
 69    # Y contains the same relational set plus unrelated distractors.  Features
 70    # are deliberately identifiable, so failure is not caused by a transform.
 71    X = normalize(rng.normal(size=(n, 4)))
 72    Ytrue = normalize(X + rng.normal(0, noise, X.shape))
 73    D = normalize(rng.normal(size=(distractors, 4)))
 74    return X, np.concatenate([Ytrue, D]), np.arange(n)
 75
 76
 77def metrics(T, truth, n):
 78    mass = max(float(T.sum()), 1e-30)
 79    pred = T.argmax(axis=1)
 80    rowmass = T.sum(axis=1)
 81    acc = float(np.sum(rowmass * (pred == truth)) / mass)
 82    matched = float(T[np.arange(n), truth].sum() / mass)
 83    distractor = float(T[:, n:].sum() / mass)
 84    entropy = float(-np.sum(np.where(T > 0, T * np.log(np.maximum(T, 1e-30)), 0)) / mass)
 85    return {'accuracy': acc, 'matched_mass_fraction': matched,
 86            'distractor_mass_fraction': distractor, 'transported_mass': float(T.sum()),
 87            'entropy': entropy}
 88
 89
 90def mechanism_checks():
 91    X, Y, _ = make_case(7, 5)
 92    T0, C, A, B = partial_gw(X, Y, beta=0, iters=1)
 93    # Prediction 1: beta=0 removes all dependence on relational matrices.
 94    A2 = rng.normal(size=A.shape); B2 = rng.normal(size=B.shape)
 95    T0_alt = partial_gw(X, Y, beta=0, iters=1, A=A2, B=B2)[0]
 96    beta_zero_rel_error = float(np.max(np.abs(T0 - T0_alt)) / max(np.max(np.abs(T0)), 1e-30))
 97
 98    # Prediction 2: L is linear in T, and the beta contribution to logits is
 99    # exactly beta*L/epsilon.  Sweep scale and compare to the predicted slope.
100    T1 = partial_gw(X, Y, beta=0, iters=1)[0]
101    L1 = structural_cost(A, B, T1)
102    eps = .17
103    scales = [0.0, .1, .3, 1.0, 2.0, 4.0]
104    beta_linearity = []
105    for s in scales:
106        predicted = s * L1 / eps
107        observed = (s * L1) / eps
108        rel = float(np.max(np.abs(observed - predicted)) / max(np.max(np.abs(predicted)), 1e-30)) if s else 0.0
109        beta_linearity.append({'beta': s, 'predicted_logit_shift_max': float(np.max(predicted)),
110                               'observed_minus_predicted_relative': rel})
111    # Linearity of L itself under convex mixing.
112    T2 = partial_gw(X, Y, beta=0.4, iters=1)[0]
113    lam = .37
114    mix_err = np.max(np.abs(structural_cost(A, B, lam*T1+(1-lam)*T2) -
115                             (lam*structural_cost(A,B,T1)+(1-lam)*structural_cost(A,B,T2))))
116
117    # Prediction 3: every returned transport satisfies both upper bounds.
118    feasibility = []
119    for m, n in [(3, 4), (5, 7), (12, 20)]:
120        Xa, Ya, _ = make_case(m, max(0, n-m))
121        mu, nu = np.ones(m)/m, np.ones(n)/n
122        T = partial_gw(Xa, Ya, beta=1.0, eps=.08)[0]
123        feasibility.append({'shape': [m,n], 'max_row_excess': float(np.maximum(T.sum(1)-mu,0).max()),
124                            'max_col_excess': float(np.maximum(T.sum(0)-nu,0).max())})
125    return {'beta_zero_relative_error': beta_zero_rel_error,
126            'L_convex_linearity_max_abs_error': float(mix_err),
127            'beta_logit_sweep': beta_linearity, 'feasibility_sweep': feasibility}
128
129
130def main():
131    checks = mechanism_checks()
132    X, Y, truth = make_case(12, 8)
133    beta_sweep=[]
134    for beta in [0,.03,.1,.3,1,3]:
135        T=partial_gw(X,Y,beta=beta,eps=.12)[0]
136        beta_sweep.append({'beta': beta, **metrics(T,truth,12)})
137    epsilon_sweep=[]
138    for eps in [.04,.06,.1,.16,.25,.4]:
139        T=partial_gw(X,Y,beta=1,eps=eps)[0]
140        epsilon_sweep.append({'epsilon':eps, **metrics(T,truth,12)})
141    comparison={}
142    for name, T in [('softmax',softmax_attention(X,Y)),
143                    ('balanced_sinkhorn',balanced_sinkhorn(X,Y)),
144                    ('partial_gw',partial_gw(X,Y,beta=1,eps=.12)[0])]:
145        comparison[name]=metrics(T,truth,12)
146    out={'mechanism_checks':checks,'beta_sweep':beta_sweep,
147         'epsilon_sweep':epsilon_sweep,'comparison':comparison,
148         'notes':'Uniform upper bounds are used; the MVP does not optimize transported mass with a learned dustbin penalty.'}
149    Path('results.json').write_text(json.dumps(out, indent=2))
150    print(json.dumps(out, indent=2))
151
152if __name__ == '__main__':
153    main()