| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import pandas as pd |
| import torch |
| import trackio |
| from model import CalibratedDiscreteTransition, NeuralVectorField, parameter_count |
| from safetensors.torch import save_file |
| from torch.nn import functional as F |
|
|
| PROJECT_DIR = Path(__file__).resolve().parent |
| ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "neural-ode-pocket" |
| DATA_DIR = PROJECT_DIR / "data" |
| SEED = 2213 |
|
|
|
|
| def true_field(state: torch.Tensor) -> torch.Tensor: |
| x, velocity = state[:, :1], state[:, 1:] |
| return torch.cat( |
| [velocity, 1.5 * (1 - x.square()) * velocity - x], |
| dim=1, |
| ) |
|
|
|
|
| def true_step(state: torch.Tensor, delta_time: torch.Tensor) -> torch.Tensor: |
| k1 = true_field(state) |
| k2 = true_field(state + 0.5 * delta_time * k1) |
| k3 = true_field(state + 0.5 * delta_time * k2) |
| k4 = true_field(state + delta_time * k3) |
| return state + delta_time * (k1 + 2 * k2 + 2 * k3 + k4) / 6 |
|
|
|
|
| @torch.inference_mode() |
| def rollout( |
| model: torch.nn.Module, |
| initial: torch.Tensor, |
| delta_time: float, |
| steps: int, |
| *, |
| continuous: bool, |
| ) -> torch.Tensor: |
| state = initial |
| output = [] |
| dt = torch.full((len(initial), 1), delta_time) |
| for _ in range(steps): |
| state = model.step(state, dt) if continuous else model(state, dt) |
| output.append(state) |
| return torch.stack(output, dim=1) |
|
|
|
|
| @torch.inference_mode() |
| def true_rollout( |
| initial: torch.Tensor, |
| delta_time: float, |
| steps: int, |
| ) -> torch.Tensor: |
| state = initial |
| output = [] |
| dt = torch.full((len(initial), 1), delta_time) |
| for _ in range(steps): |
| state = true_step(state, dt) |
| output.append(state) |
| return torch.stack(output, dim=1) |
|
|
|
|
| def first_failure(error: torch.Tensor, threshold: float = 0.5) -> float: |
| horizons = [] |
| for row in error: |
| failures = torch.nonzero(row > threshold) |
| horizons.append(int(failures[0]) if len(failures) else len(row)) |
| return float(sum(horizons) / len(horizons)) |
|
|
|
|
| @torch.inference_mode() |
| def evaluate( |
| model: torch.nn.Module, |
| delta_time: float, |
| steps: int, |
| *, |
| continuous: bool, |
| ) -> dict: |
| generator = torch.Generator().manual_seed(SEED + int(delta_time * 1000)) |
| initial = -2.5 + 5 * torch.rand(500, 2, generator=generator) |
| truth = true_rollout(initial, delta_time, steps) |
| prediction = rollout( |
| model, initial, delta_time, steps, continuous=continuous |
| ) |
| error = (prediction - truth).square().sum(2).sqrt() |
| return { |
| "trajectory_rmse": float((prediction - truth).square().mean().sqrt()), |
| "final_state_rmse": float( |
| (prediction[:, -1] - truth[:, -1]).square().mean().sqrt() |
| ), |
| "mean_horizon_before_error_0.5": first_failure(error), |
| "delta_time": delta_time, |
| "steps": steps, |
| "initial_conditions": len(initial), |
| } |
|
|
|
|
| def train_variant( |
| name: str, |
| model: torch.nn.Module, |
| *, |
| continuous: bool, |
| ) -> torch.nn.Module: |
| optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-6) |
| best = float("inf") |
| best_state = None |
| validation_generator = torch.Generator().manual_seed(SEED + 10_000) |
| validation_state = -3 + 6 * torch.rand(5_000, 2, generator=validation_generator) |
| validation_dt = torch.full((len(validation_state), 1), 0.05) |
| validation_target = true_step(validation_state, validation_dt) |
| for step in range(1, 5_001): |
| generator = torch.Generator().manual_seed(SEED + step) |
| state = -3 + 6 * torch.rand(512, 2, generator=generator) |
| dt = torch.full((len(state), 1), 0.05) |
| target = true_step(state, dt) |
| prediction = model.step(state, dt) if continuous else model(state, dt) |
| loss = F.mse_loss(prediction, target) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| optimizer.step() |
| if step % 250 == 0: |
| with torch.inference_mode(): |
| validation_prediction = ( |
| model.step(validation_state, validation_dt) |
| if continuous |
| else model(validation_state, validation_dt) |
| ) |
| validation_rmse = float( |
| (validation_prediction - validation_target) |
| .square() |
| .mean() |
| .sqrt() |
| ) |
| trackio.log( |
| { |
| "variant": name, |
| "training_step": step, |
| "training_mse": float(loss.detach()), |
| "validation_one_step_rmse": validation_rmse, |
| } |
| ) |
| if validation_rmse < best: |
| best = validation_rmse |
| best_state = { |
| key: value.detach().cpu().clone() |
| for key, value in model.state_dict().items() |
| } |
| assert best_state is not None |
| model.load_state_dict(best_state) |
| return model |
|
|
|
|
| def main() -> None: |
| torch.manual_seed(SEED) |
| torch.set_num_threads(1) |
| models = { |
| "neural_ode_rk4": (NeuralVectorField(), True), |
| "discrete_transition": (CalibratedDiscreteTransition(), False), |
| } |
| assert {parameter_count(item[0]) for item in models.values()} == {1_218} |
| trackio.init( |
| project="neural-ode-pocket", |
| name="van-der-pol-rk4-v1", |
| config={ |
| "parameters_per_model": 1_218, |
| "training_delta_time": 0.05, |
| "training_steps": 5_000, |
| }, |
| ) |
| results = {} |
| ARTIFACT_DIR.mkdir(parents=True, exist_ok=True) |
| DATA_DIR.mkdir(parents=True, exist_ok=True) |
| for name, (model, continuous) in models.items(): |
| model = train_variant(name, model, continuous=continuous) |
| models[name] = (model, continuous) |
| results[name] = { |
| "parameters": parameter_count(model), |
| "trained_timestep": evaluate( |
| model, 0.05, 400, continuous=continuous |
| ), |
| "unseen_four_x_timestep": evaluate( |
| model, 0.20, 100, continuous=continuous |
| ), |
| } |
| save_file(model.state_dict(), ARTIFACT_DIR / f"{name}.safetensors") |
| initial = torch.tensor([[2.0, 0.0]]) |
| truth = true_rollout(initial, 0.05, 400)[0] |
| frame = { |
| "step": list(range(400)), |
| "true_x": truth[:, 0].numpy(), |
| "true_velocity": truth[:, 1].numpy(), |
| } |
| for name, (model, continuous) in models.items(): |
| prediction = rollout(model, initial, 0.05, 400, continuous=continuous)[0] |
| frame[f"{name}_x"] = prediction[:, 0].numpy() |
| frame[f"{name}_velocity"] = prediction[:, 1].numpy() |
| report = { |
| "experiment": "Neural ODE versus discrete transition model", |
| "system": "Van der Pol oscillator, mu=1.5", |
| "results": results, |
| } |
| (ARTIFACT_DIR / "evaluation.json").write_text( |
| json.dumps(report, indent=2), encoding="utf-8" |
| ) |
| pd.DataFrame(frame).to_parquet(DATA_DIR / "phase_rollout.parquet", index=False) |
| trackio.log( |
| { |
| "ode_unseen_rmse": results["neural_ode_rk4"][ |
| "unseen_four_x_timestep" |
| ]["trajectory_rmse"], |
| "discrete_unseen_rmse": results["discrete_transition"][ |
| "unseen_four_x_timestep" |
| ]["trajectory_rmse"], |
| } |
| ) |
| trackio.finish() |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
