| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from physics import gray_scott_step, initial_state |
|
|
| PROJECT_DIR = Path(__file__).resolve().parent |
| DATA_DIR = PROJECT_DIR / "data" |
| SIZE = 32 |
| TRAJECTORIES = 36 |
| SAMPLES_PER_TRAJECTORY = 40 |
|
|
|
|
| def main() -> None: |
| DATA_DIR.mkdir(parents=True, exist_ok=True) |
| rng = np.random.default_rng(2039) |
| states = [] |
| next_states = [] |
| feeds = [] |
| kills = [] |
| trajectory_ids = [] |
| for trajectory in range(TRAJECTORIES): |
| feed = float(rng.uniform(0.025, 0.060)) |
| kill = float(rng.uniform(0.050, 0.072)) |
| feed_tensor = torch.tensor([feed]) |
| kill_tensor = torch.tensor([kill]) |
| state = initial_state(1, SIZE, seed=2039 + trajectory) |
| for step in range(SAMPLES_PER_TRAJECTORY * 2): |
| next_state = gray_scott_step(state, feed_tensor, kill_tensor) |
| if step % 2 == 0: |
| states.append(state[0].numpy()) |
| next_states.append(next_state[0].numpy()) |
| feeds.append(feed) |
| kills.append(kill) |
| trajectory_ids.append(trajectory) |
| state = next_state |
| arrays = { |
| "states": np.stack(states).astype(np.float32), |
| "next_states": np.stack(next_states).astype(np.float32), |
| "feeds": np.asarray(feeds, dtype=np.float32), |
| "kills": np.asarray(kills, dtype=np.float32), |
| "trajectory_ids": np.asarray(trajectory_ids, dtype=np.int64), |
| } |
| np.savez_compressed(DATA_DIR / "gray_scott_trajectories.npz", **arrays) |
| manifest = { |
| "grid_size": SIZE, |
| "trajectories": TRAJECTORIES, |
| "samples": len(states), |
| "samples_per_trajectory": SAMPLES_PER_TRAJECTORY, |
| "train_trajectories": list(range(0, 28)), |
| "validation_trajectories": list(range(28, 32)), |
| "test_trajectories": list(range(32, 36)), |
| "feed_range": [0.025, 0.060], |
| "kill_range": [0.050, 0.072], |
| "path": "gray_scott_trajectories.npz", |
| } |
| (DATA_DIR / "manifest.json").write_text( |
| json.dumps(manifest, indent=2), |
| encoding="utf-8", |
| ) |
| print(json.dumps(manifest, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
