1import itertools
  2import json
  3import math
  4import random
  5from pathlib import Path
  6
  7import numpy as np
  8import torch
  9
 10SEED = 466
 11np.random.seed(SEED)
 12random.seed(SEED)
 13torch.manual_seed(SEED)
 14
 15
 16def subset_scores(A, m):
 17    """Return every m-row subset and its least singular value."""
 18    subsets = list(itertools.combinations(range(A.shape[0]), m))
 19    scores = np.array([
 20        np.linalg.svd(A[list(T), :], compute_uv=False)[-1] for T in subsets
 21    ], dtype=np.float64)
 22    return subsets, scores
 23
 24
 25def gaussian_scaling_check():
 26    """Exhaustively check the theorem's log M / m behavior at gamma=2."""
 27    gamma = 2.0
 28    h = gamma * math.log(gamma) - (gamma - 1) * math.log(gamma - 1)
 29    theory = -h
 30    rows = []
 31    for m in range(2, 9):
 32        vals = []
 33        for rep in range(10):
 34            rng = np.random.default_rng(SEED + 1000 * m + rep)
 35            A = rng.standard_normal((2 * m, m))
 36            _, scores = subset_scores(A, m)
 37            vals.append(math.log(max(float(scores.min()), 1e-300)) / m)
 38        rows.append({"m": m, "mean": float(np.mean(vals)),
 39                     "std": float(np.std(vals)), "theory": theory})
 40    rng = np.random.default_rng(SEED + 999)
 41    A = rng.standard_normal((12, 6))
 42    _, scores = subset_scores(A, 6)
 43    return {"gamma": gamma, "theory_limit": theory, "scaling": rows,
 44            "directional": {"min": float(scores.min()),
 45                            "median": float(np.median(scores)),
 46                            "min_over_median": float(scores.min() / np.median(scores))}}
 47
 48
 49def train(mask_mode, W, X, y, masks, steps=500, lr=0.04):
 50    """Train a small downstream network on fixed channel subsets."""
 51    device = "cuda" if torch.cuda.is_available() else "cpu"
 52    try:
 53        torch.manual_seed(SEED)
 54        model = torch.nn.Sequential(torch.nn.Linear(int(masks[0].numel()), 32),
 55                                    torch.nn.Tanh(), torch.nn.Linear(32, 1)).to(device)
 56        opt = torch.optim.Adam(model.parameters(), lr=lr)
 57        Xt, yt = X.to(device), y.to(device)
 58        Wt = W.to(device)
 59        losses = []
 60        for step in range(steps):
 61            if mask_mode == "random":
 62                # deterministic pseudo-random schedule, with the same pool size
 63                j = (step * 7919 + 104729) % len(masks)
 64            else:
 65                # Scores are computed from the fixed readout matrix. The
 66                # adversary chooses the lowest-score candidate each step.
 67                j = 0
 68                # masks are preordered by increasing score; cycle through the
 69                # hardest few to avoid training on one pathological subset only.
 70                j = (step // 25) % max(1, min(5, len(masks)))
 71            idx = masks[j].to(device)
 72            # Select surviving measurements and project to a common m-vector.
 73            xb = Xt[:, idx]
 74            pred = model(xb)
 75            loss = torch.nn.functional.mse_loss(pred, yt)
 76            opt.zero_grad(); loss.backward(); opt.step()
 77            losses.append(float(loss.detach().cpu()))
 78        return model, losses
 79    except Exception:
 80        # CPU fallback is required for shared/fragile CUDA environments.
 81        model = torch.nn.Sequential(torch.nn.Linear(int(masks[0].numel()), 32),
 82                                    torch.nn.Tanh(), torch.nn.Linear(32, 1))
 83        opt = torch.optim.Adam(model.parameters(), lr=lr)
 84        losses = []
 85        for step in range(steps):
 86            j = step % len(masks) if mask_mode == "random" else (step // 25) % min(5, len(masks))
 87            pred = model(X[:, masks[j]])
 88            loss = torch.nn.functional.mse_loss(pred, y)
 89            opt.zero_grad(); loss.backward(); opt.step(); losses.append(float(loss))
 90        return model, losses
 91
 92
 93def regression_check():
 94    """Compare random masks with low-condition-score masks at equal steps."""
 95    rng = np.random.default_rng(SEED + 77)
 96    n, N, m, latent = 1200, 12, 6, 6
 97    Z = rng.standard_normal((n, latent)).astype(np.float32)
 98    W_np = rng.standard_normal((N, latent)).astype(np.float32)
 99    X_np = Z @ W_np.T + 0.05 * rng.standard_normal((n, N)).astype(np.float32)
100    y_np = (Z[:, :1] - 0.6 * Z[:, 1:2] + 0.2 * Z[:, 2:3]).astype(np.float32)
101    _, scores = subset_scores(W_np, m)
102    all_subsets = list(itertools.combinations(range(N), m))
103    order = np.argsort(scores)
104    ordered = [all_subsets[i] for i in order]
105    # Both methods have the same 20-mask training budget. Random uses a
106    # deterministic spread of the full pool; adversarial uses the 20 lowest scores.
107    random_subsets = [all_subsets[i] for i in np.linspace(0, len(all_subsets)-1, 20, dtype=int)]
108    masks = [torch.tensor(t, dtype=torch.long) for t in random_subsets]
109    hard_masks = [torch.tensor(t, dtype=torch.long) for t in ordered[:20]]
110    X, y, W = torch.tensor(X_np), torch.tensor(y_np), torch.tensor(W_np)
111    rmodel, rloss = train("random", W, X, y, masks)
112    hmodel, hloss = train("hard", W, X, y, hard_masks)
113    eval_device = next(rmodel.parameters()).device
114    Xe, ye = X.to(eval_device), y.to(eval_device)
115    # Identical exhaustive evaluation: this is the relevant robustness test.
116    with torch.no_grad():
117        random_eval = [float(torch.nn.functional.mse_loss(rmodel(Xe[:, torch.tensor(t, device=eval_device)]), ye).cpu()) for t in all_subsets]
118        hard_eval = [float(torch.nn.functional.mse_loss(hmodel(Xe[:, torch.tensor(t, device=eval_device)]), ye).cpu()) for t in all_subsets]
119    def summary(a):
120        a = np.asarray(a)
121        return {"mean_all": float(a.mean()), "worst_1pct": float(np.quantile(a, .99)),
122                "worst": float(a.max()), "best": float(a.min())}
123    return {"random_final_train_loss": rloss[-1], "hard_final_train_loss": hloss[-1],
124            "random_eval": summary(random_eval), "hard_eval": summary(hard_eval),
125            "random_low_score20_mean": float(np.mean([random_eval[i] for i in order[:20]])),
126            "hard_low_score20_mean": float(np.mean([hard_eval[i] for i in order[:20]])),
127            "hardest_score": float(scores[order[0]]),
128            "median_score": float(np.median(scores)),
129            "score_ratio": float(scores[order[0]] / np.median(scores)),
130            "random_curve": rloss, "hard_curve": hloss}
131
132
133def main():
134    out = {"seed": SEED, "gaussian": gaussian_scaling_check(), "regression": regression_check()}
135    Path("results.json").write_text(json.dumps(out, indent=2))
136    print(json.dumps({"gaussian": out["gaussian"], "regression_summary":
137                      {k: v for k, v in out["regression"].items() if "curve" not in k}}, indent=2))
138
139
140if __name__ == "__main__":
141    main()