1import json
  2import math
  3import random
  4from pathlib import Path
  5
  6import numpy as np
  7import torch
  8
  9SEED = 474
 10random.seed(SEED)
 11np.random.seed(SEED)
 12torch.manual_seed(SEED)
 13
 14
 15def device_or_cpu():
 16    if torch.cuda.is_available():
 17        try:
 18            torch.zeros(1, device="cuda")
 19            return "cuda"
 20        except Exception:
 21            pass
 22    return "cpu"
 23
 24
 25class MonotoneSplineGate(torch.nn.Module):
 26    """delta(x)=s(|x|)x, with increasing knot values in [0,1]."""
 27    def __init__(self, knots, init_identity=False):
 28        super().__init__()
 29        self.register_buffer("knots", torch.tensor(knots, dtype=torch.float32))
 30        k = len(knots)
 31        # Positive increments in unconstrained space guarantee monotonicity.
 32        if init_identity:
 33            self.bias = torch.nn.Parameter(torch.tensor(5.0))
 34            self.raw_inc = torch.nn.Parameter(torch.full((k - 1,), -5.0))
 35        else:
 36            self.bias = torch.nn.Parameter(torch.tensor(-1.0))
 37            self.raw_inc = torch.nn.Parameter(torch.full((k - 1,), -1.0))
 38
 39    def knot_values(self):
 40        # sigmoid(bias + cumulative positive increments), hence monotone.
 41        increments = torch.nn.functional.softplus(self.raw_inc)
 42        logits = self.bias + torch.cat([torch.zeros(1, device=increments.device),
 43                                        torch.cumsum(increments, dim=0)])
 44        return torch.sigmoid(logits)
 45
 46    def spline(self, magnitude):
 47        k = self.knots.to(magnitude.device)
 48        vals = self.knot_values()
 49        # endpoint extrapolation is constant, as required for a bounded gate.
 50        idx = torch.bucketize(magnitude.detach().reshape(-1), k[1:-1]).reshape(magnitude.shape)
 51        idx = idx.clamp(0, len(k) - 2)
 52        lo, hi = k[idx], k[idx + 1]
 53        frac = ((magnitude - lo) / (hi - lo).clamp_min(1e-6)).clamp(0, 1)
 54        out = vals[idx] * (1 - frac) + vals[idx + 1] * frac
 55        out = torch.where(magnitude >= k[-1], vals[-1], out)
 56        return out
 57
 58    def slope(self, magnitude):
 59        k = self.knots.to(magnitude.device)
 60        vals = self.knot_values()
 61        idx = torch.bucketize(magnitude.detach().reshape(-1), k[1:-1]).reshape(magnitude.shape)
 62        idx = idx.clamp(0, len(k) - 2)
 63        out = (vals[idx + 1] - vals[idx]) / (k[idx + 1] - k[idx]).clamp_min(1e-6)
 64        return torch.where(magnitude >= k[-1], torch.zeros_like(out), out)
 65
 66    def forward(self, x):
 67        return self.spline(x.abs()) * x
 68
 69    def divergence(self, x):
 70        m = x.abs()
 71        return (self.spline(m) + m * self.slope(m)).sum(dim=-1)
 72
 73    def sure(self, x):
 74        d = x.shape[-1]
 75        return ((self(x) - x).square().sum(dim=-1) + 2 * self.divergence(x) - d).mean()
 76
 77
 78def finite_difference_divergence(gate, x, eps=1e-4):
 79    # Explicit Jacobian trace for a small tensor, used only as a math check.
 80    vals = []
 81    for j in range(x.numel()):
 82        xp, xm = x.clone(), x.clone()
 83        xp.reshape(-1)[j] += eps
 84        xm.reshape(-1)[j] -= eps
 85        vals.append(((gate(xp) - gate(xm)).reshape(-1)[j] / (2 * eps)).item())
 86    return sum(vals)
 87
 88
 89def fit_gate(x, steps=250):
 90    gate = MonotoneSplineGate(np.linspace(0, 4, 9)).to(x.device)
 91    opt = torch.optim.Adam(gate.parameters(), lr=0.04)
 92    for step in range(steps):
 93        opt.zero_grad()
 94        # Mild high-magnitude anchor prevents the unconstrained solution from
 95        # attenuating every coordinate when the batch is nearly all noise.
 96        loss = gate.sure(x) + 0.01 * (1 - gate.knot_values()[-1]).square()
 97        loss.backward()
 98        torch.nn.utils.clip_grad_norm_(gate.parameters(), 5.0)
 99        opt.step()
100    return gate
101
102
103def fixed_soft(x, threshold=1.0):
104    return torch.sign(x) * torch.relu(x.abs() - threshold)
105
106
107def risk(est, theta):
108    return (est - theta).square().mean().item()
109
110
111def run(device):
112    # Core calculus check: analytic divergence agrees with finite differences.
113    check_gate = MonotoneSplineGate(np.linspace(0, 4, 9)).to(device)
114    check_x = torch.tensor([[0.31, 1.17, 2.63, 3.71]], device=device)
115    div_formula = check_gate.divergence(check_x).item()
116    div_fd = finite_difference_divergence(check_gate, check_x).item() if False else finite_difference_divergence(check_gate, check_x)
117
118    # SURE unbiasedness check for a fixed smooth spline under Z=theta+N(0,I).
119    theta = torch.tensor([0.0, 0.5, 1.5, 3.0], device=device).repeat(256, 1)
120    z = theta + torch.randn_like(theta)
121    sure = ((check_gate(z) - z).square().sum(1) + 2 * check_gate.divergence(z) - z.shape[1]).mean().item() / z.shape[1]
122    actual = risk(check_gate(z), theta)
123    # Repeated draws check the expected SURE identity, not just one batch.
124    sure_rep, risk_rep = [], []
125    for _ in range(20):
126        zz = theta + torch.randn_like(theta)
127        sure_rep.append((((check_gate(zz) - zz).square().sum(1) + 2 * check_gate.divergence(zz) - zz.shape[1]).mean() / zz.shape[1]).item())
128        risk_rep.append(risk(check_gate(zz), theta))
129
130    # Adaptive normal-mean experiment: sparse and dense regimes, held-out risk.
131    results = {}
132    for name, scale in [("sparse", 0.35), ("dense", 2.0)]:
133        train_theta = scale * torch.randn(64, 128, device=device)
134        train_z = train_theta + torch.randn_like(train_theta)
135        gate = fit_gate(train_z)
136        test_theta = scale * torch.randn(256, 128, device=device)
137        test_z = test_theta + torch.randn_like(test_theta)
138        with torch.no_grad():
139            r_gate = risk(gate(test_z), test_theta)
140            r_id = risk(test_z, test_theta)
141            r_soft = risk(fixed_soft(test_z), test_theta)
142            train_sure = gate.sure(train_z).item() / train_z.shape[-1]
143            vals = gate.knot_values().detach().cpu().numpy().round(3).tolist()
144        results[name] = {"identity": r_id, "fixed_soft": r_soft,
145                         "risk_fitted_gate": r_gate, "train_sure_per_dim": train_sure,
146                         "knot_values": vals}
147
148    return {"device": device, "divergence_formula": div_formula,
149            "divergence_finite_difference": div_fd,
150            "fixed_gate_sure": sure, "fixed_gate_actual_risk": actual,
151            "fixed_gate_sure_repeated_mean": float(np.mean(sure_rep)),
152            "fixed_gate_actual_risk_repeated_mean": float(np.mean(risk_rep)),
153            "fixed_gate_sure_minus_risk_repeated": float(np.mean(sure_rep) - np.mean(risk_rep)),
154            "regimes": results}
155
156
157if __name__ == "__main__":
158    out = run(device_or_cpu())
159    print(json.dumps(out, indent=2))