1import json, math, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6SEED = 2359
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8try:
  9    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 10    if device.type == 'cuda': torch.cuda.get_device_properties(0)
 11except Exception:
 12    device = torch.device('cpu')
 13
 14
 15def exact_gaussian_ot(x, m1, s1, m2, s2):
 16    # In Euclidean space, for diagonal covariances, T(x)=m2+D2 D1^-1 (x-m1).
 17    return (x - m1) * (s2 / s1) + m2
 18
 19
 20def gaussian_sanity():
 21    g = torch.Generator().manual_seed(SEED)
 22    m1 = torch.tensor([-0.7, 0.4]); s1 = torch.tensor([1.3, 0.7])
 23    m2 = torch.tensor([1.1, -0.8]); s2 = torch.tensor([0.8, 1.4])
 24    x = torch.randn(20000, 2, generator=g) * s1 + m1
 25    y = exact_gaussian_ot(x, m1, s1, m2, s2)
 26    mean_err = float((y.mean(0) - m2).norm())
 27    cov = torch.cov(y.T)
 28    cov_err = float((cov - torch.diag(s2*s2)).norm())
 29    delta = torch.tensor([2.0, -1.5]); d2 = float(delta @ delta)
 30    scaling = []
 31    for L in [1, 2, 4, 8, 16]:
 32        observed = L * 0.5 * float((delta / L @ (delta / L)))
 33        predicted = d2 / (2 * L)
 34        scaling.append({'L': L, 'predicted': predicted, 'observed': observed,
 35                        'relative_error': abs(predicted-observed)/predicted})
 36    return {'gaussian_mean_error': mean_err, 'gaussian_covariance_error': cov_err,
 37            'translation_cost_scaling': scaling}
 38
 39
 40def moons(n, noise=0.08, seed=0):
 41    rng = np.random.RandomState(seed)
 42    t = rng.rand(n) * math.pi
 43    a = np.stack([np.cos(t), np.sin(t)], 1)
 44    b = np.stack([1-np.cos(t), 1-np.sin(t)-0.45], 1)
 45    z = np.concatenate([a, b], 0)[:n]
 46    z += noise * rng.randn(*z.shape)
 47    return torch.tensor(z, dtype=torch.float32)
 48
 49
 50def mmd_rbf(x, y):
 51    dxx = torch.cdist(x, x).square(); dyy = torch.cdist(y, y).square(); dxy = torch.cdist(x, y).square()
 52    ans = 0.
 53    for h in (0.3, 0.6, 1.0, 1.8, 3.0):
 54        ans = ans + torch.exp(-dxx/(2*h*h)).mean() + torch.exp(-dyy/(2*h*h)).mean() - 2*torch.exp(-dxy/(2*h*h)).mean()
 55    return ans / 5
 56
 57
 58def sliced_wasserstein(x, y, projections=32):
 59    th = torch.arange(projections, device=x.device, dtype=x.dtype) * (math.pi / projections)
 60    dirs = torch.stack([torch.cos(th), torch.sin(th)], 1)
 61    a = torch.sort(x @ dirs.T, 0).values
 62    b = torch.sort(y @ dirs.T, 0).values
 63    return float((a-b).abs().mean().item())
 64
 65
 66class PotentialLayer(nn.Module):
 67    # T(x)=x+grad psi(x), the Euclidean Exp_x(grad psi) adaptation.
 68    def __init__(self, hidden=24):
 69        super().__init__()
 70        self.net = nn.Sequential(nn.Linear(2, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 1))
 71        # Small initial displacement makes the composition near identity.
 72        nn.init.zeros_(self.net[-1].weight); nn.init.zeros_(self.net[-1].bias)
 73
 74    def forward(self, x, create_graph=True):
 75        x = x.requires_grad_(True)
 76        psi = self.net(x).sum()
 77        v = torch.autograd.grad(psi, x, create_graph=create_graph, retain_graph=create_graph)[0]
 78        return x + v
 79
 80
 81class ResidualLayer(nn.Module):
 82    def __init__(self, hidden=24):
 83        super().__init__()
 84        self.net = nn.Sequential(nn.Linear(2, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 2))
 85        nn.init.zeros_(self.net[-1].weight); nn.init.zeros_(self.net[-1].bias)
 86
 87    def forward(self, x): return x + self.net(x)
 88
 89
 90class OTFlow(nn.Module):
 91    def __init__(self, potential=True, layers=2):
 92        super().__init__()
 93        cls = PotentialLayer if potential else ResidualLayer
 94        self.layers = nn.ModuleList([cls() for _ in range(layers)])
 95        self.potential = potential
 96
 97    def forward(self, x, create_graph=True):
 98        for layer in self.layers:
 99            x = layer(x, create_graph) if self.potential else layer(x)
100        return x
101
102
103def train_and_measure(potential, source, target, steps=260):
104    model = OTFlow(potential=potential).to(device)
105    opt = torch.optim.Adam(model.parameters(), lr=0.008)
106    x, ytar = source.to(device), target.to(device)
107    for step in range(steps):
108        y = model(x, True)
109        # OT displacement plus the implementation's distribution discrepancy.
110        loss = 0.5*(y-x).square().sum(1).mean() + 12*mmd_rbf(y, ytar)
111        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
112    # Potential maps require input derivatives at evaluation time; disable only parameter grads.
113    with torch.enable_grad():
114        y = model(x.detach(), False)
115        mm = float(mmd_rbf(y, ytar).item())
116        sw = sliced_wasserstein(y, ytar)
117        disp = float(0.5*(y-x.detach()).square().sum(1).mean().item())
118    # Jacobian determinant on a small held-out batch tests local invertibility.
119    q = source[:48].to(device).requires_grad_(True)
120    dets = []
121    for i in range(2):
122        row = torch.autograd.grad(model(q, True)[:, i].sum(), q, retain_graph=True)[0]
123        if i == 0: jac0 = row
124        else: jac1 = row
125    J = torch.stack([jac0, jac1], 2)
126    dets = torch.linalg.det(J).detach()
127    return {'mmd': mm, 'sliced_wasserstein': sw, 'displacement_cost': disp,
128            'jacobian_positive_fraction': float((dets > 0).float().mean().item()),
129            'jacobian_min': float(dets.min().item()), 'jacobian_max': float(dets.max().item())}
130
131
132def main():
133    global device
134    sanity = gaussian_sanity()
135    source = moons(192, seed=SEED+1)
136    gen = torch.Generator().manual_seed(SEED+2)
137    target = torch.randn(192, 2, generator=gen) * 1.05
138    try:
139        baseline = train_and_measure(False, source, target)
140        idea = train_and_measure(True, source, target)
141    except Exception:
142        device = torch.device('cpu')
143        baseline = train_and_measure(False, source, target)
144        idea = train_and_measure(True, source, target)
145    result = {'device': str(device), 'sanity_check': sanity,
146              'mini_experiment': {'baseline_unconstrained_residual': baseline, 'ot_gradient_potential': idea},
147              'settings': {'seed': SEED, 'samples': 192, 'layers': 2, 'steps': 260, 'mmd_weight': 12}}
148    with open('results.json', 'w') as f: json.dump(result, f, indent=2)
149    print(json.dumps(result, indent=2))
150
151if __name__ == '__main__': main()