| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import gradio as gr |
| import plotly.graph_objects as go |
| import torch |
| from model import CalibratedDiscreteTransition, NeuralVectorField |
| from safetensors.torch import load_file |
| from train import rollout, true_rollout |
|
|
| ARTIFACT_DIR = Path(__file__).resolve().parent / "artifacts" / "neural-ode-pocket" |
| ODE = NeuralVectorField() |
| ODE.load_state_dict(load_file(ARTIFACT_DIR / "neural_ode_rk4.safetensors")) |
| ODE.eval() |
| DISCRETE = CalibratedDiscreteTransition() |
| DISCRETE.load_state_dict( |
| load_file(ARTIFACT_DIR / "discrete_transition.safetensors") |
| ) |
| DISCRETE.eval() |
| REPORT = json.loads((ARTIFACT_DIR / "evaluation.json").read_text(encoding="utf-8")) |
|
|
|
|
| @torch.inference_mode() |
| def phase_portrait( |
| initial_x: float, |
| initial_velocity: float, |
| delta_time: float, |
| physical_time: float, |
| ) -> tuple[go.Figure, dict]: |
| initial = torch.tensor([[initial_x, initial_velocity]], dtype=torch.float32) |
| steps = max(1, int(physical_time / delta_time)) |
| truth = true_rollout(initial, delta_time, steps)[0] |
| ode = rollout(ODE, initial, delta_time, steps, continuous=True)[0] |
| discrete = rollout(DISCRETE, initial, delta_time, steps, continuous=False)[0] |
| figure = go.Figure() |
| figure.add_trace(go.Scatter(x=truth[:, 0], y=truth[:, 1], name="True")) |
| figure.add_trace(go.Scatter(x=ode[:, 0], y=ode[:, 1], name="Neural ODE")) |
| figure.add_trace( |
| go.Scatter(x=discrete[:, 0], y=discrete[:, 1], name="Discrete model") |
| ) |
| figure.update_layout( |
| template="plotly_dark", |
| title="Van der Pol phase portrait", |
| xaxis_title="Position", |
| yaxis_title="Velocity", |
| ) |
| metrics = { |
| "steps": steps, |
| "neural_ode_live_rmse": float((ode - truth).square().mean().sqrt()), |
| "discrete_live_rmse": float((discrete - truth).square().mean().sqrt()), |
| "verified_unseen_dt_ode_rmse": REPORT["results"]["neural_ode_rk4"][ |
| "unseen_four_x_timestep" |
| ]["trajectory_rmse"], |
| } |
| return figure, metrics |
|
|
|
|
| with gr.Blocks(title="Neural ODE Pocket") as demo: |
| gr.Markdown( |
| "# Neural ODE Pocket\n" |
| "Compare a learned continuous vector field integrated with RK4 against an " |
| "exactly parameter-matched discrete transition network." |
| ) |
| with gr.Row(): |
| initial_x = gr.Slider(-3, 3, value=2.0, step=0.1, label="Initial position") |
| initial_v = gr.Slider(-3, 3, value=0.0, step=0.1, label="Initial velocity") |
| dt = gr.Slider(0.05, 0.20, value=0.20, step=0.05, label="Timestep") |
| duration = gr.Slider(5, 30, value=20, step=1, label="Physical time") |
| initial = phase_portrait(2.0, 0.0, 0.20, 20) |
| chart = gr.Plot(value=initial[0]) |
| metrics = gr.JSON(value=initial[1]) |
| button = gr.Button("Integrate dynamics", variant="primary") |
| button.click( |
| phase_portrait, |
| inputs=[initial_x, initial_v, dt, duration], |
| outputs=[chart, metrics], |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
|
|
