File size: 3,036 Bytes
cc86714 | 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 | 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()
|