File size: 3,063 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 | from __future__ import annotations
from pathlib import Path
import numpy as np
import torch
from ..mb_potential import simulate_overdamped_mb
BASIN_A = (-0.56, 1.44)
BASIN_B = (0.62, 0.03)
DEFAULT_DATA = Path(__file__).resolve().parents[2] / "data" / "mb_reactive"
K = 100
START_XY = BASIN_A
def generate(n_reps=12, steps=250_000, dt=2e-4, beta=0.06, mu=1.0, stride=5, seed=0):
"""Return a list of (T, 2) reactive MB trajectories (after striding).
``stride`` thins the saved frames (the integrator step dt is tiny for
stability at high T; the saved frame spacing is dt*stride).
"""
trajs = []
for r in range(n_reps):
x0, y0 = BASIN_A if r % 2 == 0 else BASIN_B
xs, ys = simulate_overdamped_mb(
n_steps=steps,
dt=dt,
beta=beta,
mu=mu,
x0=x0,
y0=y0,
noise=True,
seed=seed + r,
dtype=torch.float64,
)
xy = np.stack([xs, ys], axis=1)[::stride]
if np.all(np.isfinite(xy)):
trajs.append(xy.astype(np.float64))
return trajs
def save(trajs, out_dir):
out = Path(out_dir)
out.mkdir(parents=True, exist_ok=True)
for f in out.glob("traj_*.npy"):
f.unlink()
for i, t in enumerate(trajs):
np.save(out / f"traj_{i:04d}.npy", t)
with open(out / "index.txt", "w") as fh:
fh.write("\n".join(f"traj_{i:04d}.npy" for i in range(len(trajs))) + "\n")
return out
def load_trajs(data_dir):
files = sorted(Path(data_dir).glob("traj_*.npy"))
if not files:
print(f" no reactive data in {data_dir} -- generating trajectories...")
trajs = generate()
save(trajs, data_dir)
files = sorted(Path(data_dir).glob("traj_*.npy"))
trajs = [np.load(f).astype(np.float64) for f in files]
return trajs
def orient_endpoints(model, A, B):
"""Return (start_set, end_set) so start is the basin nearest START_XY."""
xyA = model.active_to_xy(A).mean(axis=0)
xyB = model.active_to_xy(B).mean(axis=0)
dA = np.linalg.norm(xyA - np.array(START_XY))
dB = np.linalg.norm(xyB - np.array(START_XY))
return (A, B) if dA <= dB else (B, A)
if __name__ == "__main__":
import argparse
ap = argparse.ArgumentParser(description="Generate reactive MB trajectories")
ap.add_argument("--out", default="data/mb_reactive")
ap.add_argument("--n_reps", type=int, default=12)
ap.add_argument("--steps", type=int, default=250_000)
ap.add_argument("--dt", type=float, default=2e-4)
ap.add_argument("--beta", type=float, default=0.06)
ap.add_argument("--stride", type=int, default=5)
ap.add_argument("--seed", type=int, default=0)
args = ap.parse_args()
trajs = generate(
n_reps=args.n_reps,
steps=args.steps,
dt=args.dt,
beta=args.beta,
stride=args.stride,
seed=args.seed,
)
out = save(trajs, args.out)
nf = sum(len(t) for t in trajs)
print(f"wrote {len(trajs)} trajectories ({nf} frames) -> {out}")
|