File size: 3,478 Bytes
31d1d63 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | import numpy as np
import torch
def _to_torch_dtype(dt):
if isinstance(dt, torch.dtype) or dt is None:
return dt or torch.get_default_dtype()
mapping = {
np.dtype("float32"): torch.float32,
np.float32: torch.float32,
np.dtype("float64"): torch.float64,
np.float64: torch.float64,
np.dtype("int32"): torch.int32,
np.int32: torch.int32,
np.dtype("int64"): torch.int64,
np.int64: torch.int64,
}
return mapping.get(dt, torch.get_default_dtype())
def mb_params(device=None, dtype=torch.float32):
dtype = _to_torch_dtype(dtype)
A = torch.tensor([-200.0, -100.0, -170.0, 15.0], device=device, dtype=dtype)
a = torch.tensor([-1.0, -1.0, -6.5, 0.7], device=device, dtype=dtype)
b = torch.tensor([0.0, 0.0, 11.0, 0.6], device=device, dtype=dtype)
c = torch.tensor([-10.0, -10.0, -6.5, 0.7], device=device, dtype=dtype)
x0 = torch.tensor([1.0, 0.0, -0.5, -1.0], device=device, dtype=dtype)
y0 = torch.tensor([0.0, 0.5, 1.5, 1.0], device=device, dtype=dtype)
return A, a, b, c, x0, y0
def mb_potential_xy(x, y):
"""Compute Müller-Brown potential U(x,y). Args: x, y: torch tensors (broadcastable to same shape) Returns: U: torch tensor, same broadcasted shape"""
device = (
x.device
if isinstance(x, torch.Tensor)
else (y.device if isinstance(y, torch.Tensor) else None)
)
dt = (
x.dtype
if isinstance(x, torch.Tensor)
else (y.dtype if isinstance(y, torch.Tensor) else None)
)
dt = _to_torch_dtype(dt)
x = torch.as_tensor(x, dtype=dt, device=device)
y = torch.as_tensor(y, dtype=dt, device=device)
A, a, b, c, x0, y0 = mb_params(device=x.device, dtype=dt)
X = x[..., None]
Y = y[..., None]
exp_arg = a * (X - x0) ** 2 + b * (X - x0) * (Y - y0) + c * (Y - y0) ** 2
U = torch.sum(A * torch.exp(exp_arg), dim=-1)
return U
def mb_grad_xy(x, y):
"""
Gradient of U: returns (dU/dx, dU/dy)
"""
device, dtype = x.device, x.dtype
A, a, b, c, x0, y0 = mb_params(device=device, dtype=dtype)
X = x[..., None]
Y = y[..., None]
dx = X - x0
dy = Y - y0
exp_arg = a * dx * dx + b * dx * dy + c * dy * dy
E = torch.exp(exp_arg) * A
dUdx = torch.sum(E * (2 * a * dx + b * dy), dim=-1)
dUdy = torch.sum(E * (b * dx + 2 * c * dy), dim=-1)
return dUdx, dUdy
def simulate_overdamped_mb(
n_steps=100000,
dt=1e-3,
beta=1.0,
mu=1.0,
x0=None,
y0=None,
noise=True,
device="cpu",
dtype=torch.float32,
seed=None,
):
"""
Generate unconditioned trajectories on the MB surface.
Overdamped Langevin: dx = -mu * grad U dt + sqrt(2 mu / beta) dW
Returns numpy arrays (xs, ys) of shape [n_steps].
A sequence of coordinates with a fixed dt
"""
if seed is not None:
torch.manual_seed(seed)
np.random.seed(seed)
x = torch.tensor([x0], device=device, dtype=dtype)
y = torch.tensor([y0], device=device, dtype=dtype)
xs = torch.empty(n_steps, dtype=dtype, device=device)
ys = torch.empty(n_steps, dtype=dtype, device=device)
sigma = (2.0 * mu / beta * dt) ** 0.5
for t in range(n_steps):
dUdx, dUdy = mb_grad_xy(x, y)
x = x - mu * dUdx * dt + sigma * torch.randn_like(x)
y = y - mu * dUdy * dt + sigma * torch.randn_like(y)
xs[t] = x
ys[t] = y
return xs.cpu().numpy(), ys.cpu().numpy()
|