| 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() |
|
|
