1import json, math, random
  2import numpy as np
  3import torch
  4
  5SEED = 540
  6np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  7torch.set_default_dtype(torch.float64)
  8try:
  9    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 10except Exception:
 11    device = torch.device('cpu')
 12
 13# 2D nonequilibrium Langevin system: double well plus rotational current.
 14omega, T, dt, K = 1.25, 0.08, 0.05, 50
 15A = torch.tensor([-1.0, 0.0], device=device)
 16B = torch.tensor([1.0, 0.0], device=device)
 17
 18def mobility(z):
 19    # Positive, state-dependent diagonal mobility (known ground truth).
 20    x, y = z[..., 0], z[..., 1]
 21    return torch.stack((0.35 + 0.85*torch.sigmoid(2.0*x),
 22                        0.55 + 0.45*torch.sigmoid(-1.5*y)), dim=-1)
 23
 24def potential_grad(z):
 25    x, y = z[..., 0], z[..., 1]
 26    return torch.stack((x*(x*x-1.0), y), dim=-1)
 27
 28def drift(z):
 29    g = potential_grad(z)
 30    # nonequilibrium rotational component, not a gradient or time reversal
 31    rot = torch.stack((-z[..., 1], z[..., 0]), dim=-1)
 32    return -g + omega*rot
 33
 34def action(path):
 35    z0, z1 = path[:-1], path[1:]
 36    m = mobility(z0)
 37    v = (z1-z0)/dt
 38    residual = v-drift(z0)
 39    # a = 2 T M; FW action = .5 dt residual^T a^-1 residual
 40    return (0.25*dt/T * (residual*residual/m).sum(dim=-1)).sum()
 41
 42def optimize_path(init, steps=1200, lr=0.025):
 43    inner = torch.nn.Parameter(init[1:-1].clone())
 44    opt = torch.optim.Adam([inner], lr=lr)
 45    history=[]
 46    for i in range(steps):
 47        path = torch.cat((A[None], inner, B[None]), dim=0)
 48        loss = action(path)
 49        opt.zero_grad(); loss.backward(); opt.step()
 50        if i % 100 == 0: history.append(float(loss.detach().cpu()))
 51    path = torch.cat((A[None], inner.detach(), B[None]), dim=0)
 52    return path, history
 53
 54def check_sde(n=150000):
 55    # Verify both claims in the stochastic equation: conditional mean drift and variance 2TMdt.
 56    z = torch.tensor([0.35, -0.7], device=device)
 57    m = mobility(z); b = drift(z)
 58    eps = torch.randn(n, 2, device=device)
 59    inc = b*dt + torch.sqrt(2*T*m*dt)*eps
 60    mean_err = float(torch.max(torch.abs(inc.mean(0)/dt-b)).cpu())
 61    var_ratio = (inc.var(0)/(2*T*m*dt)).detach().cpu().numpy()
 62    return {'mean_abs_error_per_dt': mean_err, 'variance_ratio': var_ratio.tolist(),
 63            'expected_variance_ratio': [1.0,1.0]}
 64
 65def stochastic_success(path, mode, n=1200, gain=7.0, seed=1234):
 66    gen = torch.Generator(device=device); gen.manual_seed(seed)
 67    z = A[None,:].repeat(n,1)
 68    min_dist = torch.full((n,), 99.0, device=device)
 69    for k in range(K):
 70        if mode == 'optimized':
 71            # local tube controller following the optimized nonreversible path
 72            target = path[k]
 73            r = (path[k+1]-path[k])/dt + gain*(target-z)
 74        elif mode == 'reverse_relaxation':
 75            # standard time-reversed-relaxation control, with no path optimization
 76            r = -drift(z)
 77        else:
 78            raise ValueError(mode)
 79        m = mobility(z)
 80        noise = torch.randn(z.shape, generator=gen, device=device)
 81        z = z + r*dt + torch.sqrt(2*T*m*dt)*noise
 82        min_dist = torch.minimum(min_dist, torch.linalg.vector_norm(z-B, dim=1))
 83    finite = torch.isfinite(z).all(dim=1)
 84    dist = torch.linalg.vector_norm(torch.nan_to_num(z, nan=1e6, posinf=1e6, neginf=-1e6)-B, dim=1)
 85    reach = finite & (dist < 0.35)
 86    finite_dist = dist[finite]
 87    mean_dist = float(finite_dist.mean().cpu()) if finite_dist.numel() else float('inf')
 88    finite_min = min_dist[torch.isfinite(min_dist)]
 89    mean_min = float(finite_min.mean().cpu()) if finite_min.numel() else float('inf')
 90    return float(reach.float().mean().cpu()), mean_dist, mean_min, int(finite.sum().cpu())
 91
 92def main():
 93    # Core action sanity check: constant path velocity has the stated quadratic residual.
 94    straight = torch.linspace(0,1,K+1,device=device)[:,None]*B + torch.linspace(1,0,K+1,device=device)[:,None]*A
 95    straight_s = float(action(straight).cpu())
 96    # Add small deterministic perturbations and check action rises around this fixed path.
 97    perturb = straight.clone(); perturb[1:-1,1] += 0.08*torch.sin(torch.arange(1,K,device=device))
 98    pert_s = float(action(perturb).cpu())
 99    path, hist = optimize_path(straight + torch.cat((torch.zeros(1,2,device=device), 0.03*torch.randn(K-1,2,device=device), torch.zeros(1,2,device=device))))
100    opt_s = float(action(path).cpu())
101    check = check_sde()
102    base = stochastic_success(path, 'reverse_relaxation')
103    idea = stochastic_success(path, 'optimized')
104    result = {'device':str(device), 'K':K, 'dt':dt, 'T':T, 'omega':omega,
105      'math_check':check, 'action_check':{'straight':straight_s,'perturbed':pert_s,'optimized':opt_s,
106                                           'optimization_history':hist},
107      'baseline_reverse_relaxation':{'success_rate':base[0],'mean_final_distance':base[1],'mean_min_distance':base[2], 'finite_trajectories':base[3]},
108      'idea_action_optimized':{'success_rate':idea[0],'mean_final_distance':idea[1],'mean_min_distance':idea[2], 'finite_trajectories':idea[3]}}
109    with open('results.json','w') as f: json.dump(result,f,indent=2)
110    np.savetxt('optimized_path.csv', path.cpu().numpy(), delimiter=',', header='x,y', comments='')
111    print(json.dumps(result, indent=2))
112
113if __name__ == '__main__': main()