| 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}") |
|
|
