1import json, math, time
  2from itertools import combinations
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED = 7
  8np.random.seed(SEED)
  9torch.manual_seed(SEED)
 10torch.set_num_threads(4)
 11DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 12try:
 13    if DEVICE == "cuda":
 14        torch.cuda.set_device(0)
 15        torch.zeros(1, device=DEVICE)
 16except Exception:
 17    DEVICE = "cpu"
 18
 19
 20def barycentric_matrix(n):
 21    d = n - 1
 22    H = np.zeros((n, n), dtype=np.float64)
 23    for i in range(n):
 24        for j in range(n):
 25            for r in range(i + 1):
 26                H[i, j] += r**j * (r + 1)**(d-j) * (-1)**(i-r) * math.comb(d+1, i-r)
 27    return H
 28
 29
 30def all_minors(A, max_order=None):
 31    n, m = A.shape
 32    q = min(n, m) if max_order is None else min(max_order, n, m)
 33    out = []
 34    for k in range(1, q + 1):
 35        for rows in combinations(range(n), k):
 36            for cols in combinations(range(m), k):
 37                out.append((k, rows, cols, float(np.linalg.det(A[np.ix_(rows, cols)]))))
 38    return out
 39
 40
 41class PositiveBidiagonalMixer(nn.Module):
 42    def __init__(self, n, depth=2):
 43        super().__init__()
 44        self.n, self.depth = n, depth
 45        self.ell = nn.Parameter(torch.full((depth, n-1), -4.0))
 46        self.u = nn.Parameter(torch.full((depth, n-1), -4.0))
 47        self.s = nn.Parameter(torch.zeros(n))
 48
 49    def forward(self, x):
 50        y = x * torch.exp(self.s)
 51        for k in range(self.depth):
 52            z = torch.zeros_like(y)
 53            z[..., 1:] = y[..., :-1] * torch.nn.functional.softplus(self.ell[k])
 54            y = y + z
 55        for k in range(self.depth - 1, -1, -1):
 56            z = torch.zeros_like(y)
 57            z[..., :-1] = y[..., 1:] * torch.nn.functional.softplus(self.u[k])
 58            y = y + z
 59        return y
 60
 61    def matrix(self):
 62        return self(torch.eye(self.n, device=self.s.device)).T
 63
 64
 65class DenseMap(nn.Module):
 66    def __init__(self, n):
 67        super().__init__()
 68        self.weight = nn.Parameter(torch.randn(n, n) / math.sqrt(n))
 69    def forward(self, x):
 70        return x @ self.weight.T
 71
 72
 73class LowRankMap(nn.Module):
 74    def __init__(self, n, rank):
 75        super().__init__()
 76        self.a = nn.Parameter(torch.randn(n, rank) / math.sqrt(n))
 77        self.b = nn.Parameter(torch.randn(rank, n) / math.sqrt(rank))
 78    def forward(self, x):
 79        return x @ (self.a @ self.b).T
 80
 81
 82def fit(model, xtr, ytr, xva, yva, steps=500, lr=0.03):
 83    model.to(DEVICE)
 84    opt = torch.optim.Adam(model.parameters(), lr=lr)
 85    t0 = time.perf_counter()
 86    for _ in range(steps):
 87        opt.zero_grad(set_to_none=True)
 88        loss = ((model(xtr) - ytr)**2).mean()
 89        loss.backward()
 90        opt.step()
 91    with torch.no_grad():
 92        train = ((model(xtr) - ytr)**2).mean().item()
 93        val = ((model(xva) - yva)**2).mean().item()
 94    return train, val, time.perf_counter() - t0, sum(p.numel() for p in model.parameters())
 95
 96
 97def main():
 98    # Core formula and total-positivity sanity check.
 99    H = barycentric_matrix(5)
100    hm = all_minors(H)
101    negative = [v for v in hm if v[3] < -1e-8]
102    # The practical parameterization has strictly positive entries, but not every
103    # product of arbitrary positive bidiagonals is TP; test the claimed network map.
104    torch.manual_seed(SEED)
105    pos = PositiveBidiagonalMixer(5, depth=4).to(DEVICE)
106    P = pos.matrix().detach().cpu().numpy()
107    pm = all_minors(P)
108    pnegative = [v for v in pm if v[3] < -1e-8]
109
110    # Approximation/learning test: target is a dense linear map.
111    torch.manual_seed(SEED)
112    n, ns = 12, 1024
113    x = torch.randn(ns, n, device=DEVICE)
114    target = torch.randn(n, n, device=DEVICE) / math.sqrt(n)
115    y = x @ target.T
116    xtr, ytr, xva, yva = x[:768], y[:768], x[768:], y[768:]
117    results = {}
118    for name, model in [
119        ("dense", DenseMap(n)),
120        ("positive_K2", PositiveBidiagonalMixer(n, 2)),
121        ("positive_K6", PositiveBidiagonalMixer(n, 6)),
122        ("lowrank_r4", LowRankMap(n, 4)),
123    ]:
124        torch.manual_seed(SEED)
125        tr, va, secs, params = fit(model, xtr, ytr, xva, yva)
126        results[name] = {"train_mse": tr, "val_mse": va, "seconds": secs, "parameters": params}
127
128    report = {
129        "device": DEVICE,
130        "barycentric_H_5": H.tolist(),
131        "barycentric_minors": {"count": len(hm), "negative": len(negative), "minimum": min(v[3] for v in hm)},
132        "positive_product_minors": {"count": len(pm), "negative": len(pnegative), "minimum": min(v[3] for v in pm)},
133        "fit": results,
134    }
135    with open("results.json", "w") as f:
136        json.dump(report, f, indent=2)
137    print(json.dumps(report, indent=2))
138
139
140if __name__ == "__main__":
141    main()