Skip to content

Commit 3b244f0

Browse files
JoKircherJoKircher
JoKircher
authored and
JoKircher
committed
WIP: original implementation
1 parent ceeedcc commit 3b244f0

9 files changed

Lines changed: 684 additions & 132 deletions

File tree

‎src/GraphNetSim.jl‎

Lines changed: 134 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ using HDF5
1717
using Plots
1818

1919
import SciMLBase: ODEProblem, SecondOrderODEProblem
20-
import OrdinaryDiffEq: OrdinaryDiffEqAlgorithm, Tsit5
20+
import OrdinaryDiffEq: OrdinaryDiffEqAlgorithm, Tsit5, Euler
2121
import ProgressMeter: Progress
2222

2323
import Base: @kwdef
@@ -30,6 +30,7 @@ import Printf: @sprintf
3030
include("utils.jl")
3131
include("graph.jl")
3232
include("solve.jl")
33+
include("rollout_history.jl")
3334
include("dataset.jl")
3435
include("visualize.jl")
3536
include("config.jl")
@@ -114,11 +115,53 @@ Configuration structure for training and evaluating Graph Neural Network simulat
114115
optimizer_learning_rate_start::Float32 = 1.0f-4
115116
optimizer_learning_rate_stop::Union{Nothing,Float32} = nothing
116117
norm_type::Symbol = :online
118+
history_size::Int = 1
117119
save_step::Bool = false
118120
on_grad::Union{Nothing,Function} = nothing
119121
on_valid::Union{Nothing,Function} = nothing
120122
end
121123

124+
function _validate_history_args(args::Args)
125+
args.history_size ≥ 1 || throw(
126+
ArgumentError("history_size must be ≥ 1, got $(args.history_size)")
127+
)
128+
args.history_size == 1 && return
129+
args.training_strategy isa DerivativeTraining || throw(
130+
ArgumentError(
131+
"history_size > 1 is only supported with DerivativeTraining; got " *
132+
"$(typeof(args.training_strategy)). ODE-based strategies will be " *
133+
"extended in a follow-up plan.",
134+
),
135+
)
136+
args.solver_valid isa Euler || throw(
137+
ArgumentError(
138+
"history_size > 1 requires solver_valid = Euler() (sliding-buffer " *
139+
"rollout is fixed-step only); got $(typeof(args.solver_valid)).",
140+
),
141+
)
142+
isnothing(args.solver_valid_dt) && throw(
143+
ArgumentError(
144+
"history_size > 1 requires solver_valid_dt to be set explicitly " *
145+
"(Euler is fixed-step).",
146+
),
147+
)
148+
return
149+
end
150+
151+
function _validate_history_meta(meta::Dict, args::Args)
152+
args.history_size == 1 && return
153+
allowed = ("velocity", "wall_distance")
154+
bad = [f for f in meta["input_features"] if !(f in allowed)]
155+
isempty(bad) || throw(
156+
ArgumentError(
157+
"history_size > 1 requires input_features ⊆ $(allowed); got " *
158+
"extras $(bad). Drop position from input_features or set " *
159+
"history_size = 1.",
160+
),
161+
)
162+
return
163+
end
164+
122165
"""
123166
calc_norms(dataset::Dataset, device::Function, args::Args)
124167
@@ -163,7 +206,11 @@ function calc_norms(dataset, device, args)
163206
for feature in dataset.meta["feature_names"]
164207
feature_dim = dataset.meta["features"][feature]["dim"]
165208
if feature in input_features
166-
quantities += feature_dim
209+
if feature == "velocity"
210+
quantities += feature_dim * args.history_size
211+
else
212+
quantities += feature_dim
213+
end
167214
end
168215

169216
if getfield(
@@ -313,7 +360,7 @@ function calc_norms(dataset, device, args)
313360
end
314361
end
315362
end
316-
if n_node_types(dataset.meta) > 1
363+
if n_node_types(dataset.meta) > 1 || "wall_distance" in input_features
317364
quantities += length(dataset.meta["bounds"]) * 2
318365
end
319366

@@ -386,12 +433,14 @@ function train_network(opt, ds_path, cp_path; kws...)
386433
layer_size=existing_cfg.layer_size,
387434
hidden_layers=existing_cfg.hidden_layers,
388435
norm_type=existing_cfg.norm_type,
436+
history_size=existing_cfg.history_size,
389437
),
390438
NamedTuple(kws),
391439
)
392440
end
393441

394442
args = Args(; kws...)
443+
_validate_history_args(args)
395444

396445
save_model_config(
397446
ModelConfig(;
@@ -403,6 +452,7 @@ function train_network(opt, ds_path, cp_path; kws...)
403452
types_noisy=args.types_noisy,
404453
noise_stddevs=args.noise_stddevs,
405454
norm_type=args.norm_type,
455+
history_size=args.history_size,
406456
),
407457
cp_path,
408458
)
@@ -425,12 +475,15 @@ function train_network(opt, ds_path, cp_path; kws...)
425475
ds_train.meta["types_noisy"] = args.types_noisy
426476
ds_train.meta["noise_stddevs"] = args.noise_stddevs
427477
ds_train.meta["device"] = device
478+
ds_train.meta["history_size"] = args.history_size
428479
ds_valid = Dataset(:valid, ds_path, args)
429480
ds_valid.meta["types_updated"] = args.types_updated
430481
ds_valid.meta["types_noisy"] = args.types_noisy
431482
ds_valid.meta["noise_stddevs"] = args.noise_stddevs
432483
ds_valid.meta["device"] = device
484+
ds_valid.meta["history_size"] = args.history_size
433485
ds_valid.meta["training_strategy"] = nothing
486+
_validate_history_meta(ds_train.meta, args)
434487

435488
@info "Training data loaded!"
436489
Threads.nthreads() < 2 &&
@@ -884,12 +937,14 @@ function eval_network(
884937
mps=existing_cfg.mps,
885938
layer_size=existing_cfg.layer_size,
886939
hidden_layers=existing_cfg.hidden_layers,
940+
history_size=existing_cfg.history_size,
887941
),
888942
NamedTuple(kws),
889943
)
890944
end
891945

892946
args = Args(; kws...)
947+
_validate_history_args(args)
893948

894949
if CUDA.functional() && args.use_cuda
895950
@info "Evaluating on CUDA GPU..."
@@ -906,7 +961,9 @@ function eval_network(
906961
println("Loading evaluation data...")
907962
ds_test = Dataset(:test, ds_path, args)
908963
ds_test.meta["device"] = device
964+
ds_test.meta["history_size"] = args.history_size
909965
ds_test.meta["training_strategy"] = nothing
966+
_validate_history_meta(ds_test.meta, args)
910967

911968
# clear_log(1, false)
912969
@info "Evaluation data loaded!"
@@ -1022,8 +1079,11 @@ function eval_network!(
10221079
println("Rollout trajectory $ti...")
10231080

10241081
if length(test_loader) > 1
1025-
start = 0.0f0
10261082
dt = data["dt"] # TODO dt can be an array?
1083+
C = get(ds_test.meta, "history_size", 1)
1084+
# With paper-faithful warmup, the first prediction frame is C; the C-1
1085+
# frames before it seed the velocity buffer.
1086+
start = Float32((C - 1) * dt)
10271087
stop = round((data["trajectory_length"] - 1) * dt; digits=6)
10281088
saves = start:dt:stop
10291089
mse_steps = saves
@@ -1039,23 +1099,39 @@ function eval_network!(
10391099
enabled=args.show_progress_bars,
10401100
)
10411101

1042-
sol = rollout(
1043-
solver,
1044-
gns,
1045-
initial_state,
1046-
output_features,
1047-
ds_test.meta,
1048-
target_features,
1049-
node_type,
1050-
data["mask"],
1051-
data["val_mask"],
1052-
start,
1053-
stop,
1054-
dt,
1055-
saves,
1056-
device,
1057-
pr,
1058-
)
1102+
sol = if get(ds_test.meta, "history_size", 1) > 1
1103+
rollout_history(
1104+
gns,
1105+
initial_state,
1106+
output_features,
1107+
ds_test.meta,
1108+
target_features,
1109+
node_type,
1110+
data["mask"],
1111+
data["val_mask"],
1112+
saves,
1113+
device,
1114+
pr,
1115+
)
1116+
else
1117+
rollout(
1118+
solver,
1119+
gns,
1120+
initial_state,
1121+
output_features,
1122+
ds_test.meta,
1123+
target_features,
1124+
node_type,
1125+
data["mask"],
1126+
data["val_mask"],
1127+
start,
1128+
stop,
1129+
dt,
1130+
saves,
1131+
device,
1132+
pr,
1133+
)
1134+
end
10591135

10601136
sol_t, prediction = _extract_trajectory_arrays(sol)
10611137
timesteps[(ti, "timesteps")] = sol_t
@@ -1198,12 +1274,14 @@ function extrapolate_network(
11981274
mps=existing_cfg.mps,
11991275
layer_size=existing_cfg.layer_size,
12001276
hidden_layers=existing_cfg.hidden_layers,
1277+
history_size=existing_cfg.history_size,
12011278
),
12021279
NamedTuple(kws),
12031280
)
12041281
end
12051282

12061283
args = Args(; kws...)
1284+
_validate_history_args(args)
12071285

12081286
if CUDA.functional() && args.use_cuda
12091287
@info "Extrapolating on CUDA GPU..."
@@ -1220,7 +1298,9 @@ function extrapolate_network(
12201298
println("Loading evaluation data...")
12211299
ds_test = Dataset(:test, ds_path, args)
12221300
ds_test.meta["device"] = device
1301+
ds_test.meta["history_size"] = args.history_size
12231302
ds_test.meta["training_strategy"] = nothing
1303+
_validate_history_meta(ds_test.meta, args)
12241304

12251305
@info "Evaluation data loaded!"
12261306
Threads.nthreads() < 2 &&
@@ -1346,23 +1426,39 @@ function extrapolate_network!(
13461426
enabled=args.show_progress_bars,
13471427
)
13481428

1349-
sol = rollout(
1350-
solver,
1351-
gns,
1352-
initial_state,
1353-
output_features,
1354-
ds_test.meta,
1355-
target_features,
1356-
node_type,
1357-
data["mask"],
1358-
data["val_mask"],
1359-
start,
1360-
stop,
1361-
dt,
1362-
saves,
1363-
device,
1364-
pr,
1365-
)
1429+
sol = if get(ds_test.meta, "history_size", 1) > 1
1430+
rollout_history(
1431+
gns,
1432+
initial_state,
1433+
output_features,
1434+
ds_test.meta,
1435+
target_features,
1436+
node_type,
1437+
data["mask"],
1438+
data["val_mask"],
1439+
saves,
1440+
device,
1441+
pr,
1442+
)
1443+
else
1444+
rollout(
1445+
solver,
1446+
gns,
1447+
initial_state,
1448+
output_features,
1449+
ds_test.meta,
1450+
target_features,
1451+
node_type,
1452+
data["mask"],
1453+
data["val_mask"],
1454+
start,
1455+
stop,
1456+
dt,
1457+
saves,
1458+
device,
1459+
pr,
1460+
)
1461+
end
13661462

13671463
sol_t, prediction = _extract_trajectory_arrays(sol)
13681464
npred = size(prediction.pos, 3)

‎src/config.jl‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ as documentation and may legitimately differ between training phases.
2929
- `norm_type`: Normalization strategy for Float32 features (`:online`, `:minmax`, `:meanstd`).
3030
"""
3131
@kwdef struct ModelConfig
32-
format_version::Int = 1
32+
format_version::Int = 2
3333
mps::Int
3434
layer_size::Int
3535
hidden_layers::Int
@@ -38,6 +38,7 @@ as documentation and may legitimately differ between training phases.
3838
types_noisy::Vector{Int}
3939
noise_stddevs::Vector{Float32}
4040
norm_type::Symbol = :online
41+
history_size::Int = 1
4142
end
4243

4344
"""
@@ -59,14 +60,15 @@ function save_model_config(cfg::ModelConfig, cp_path::String)
5960
if !isnothing(existing)
6061
if existing.mps != cfg.mps ||
6162
existing.layer_size != cfg.layer_size ||
62-
existing.hidden_layers != cfg.hidden_layers
63+
existing.hidden_layers != cfg.hidden_layers ||
64+
existing.history_size != cfg.history_size
6365
error(
6466
"Architecture mismatch between supplied arguments and saved " *
6567
"model config at \"$path\".\n" *
6668
" Saved: mps=$(existing.mps), layer_size=$(existing.layer_size), " *
67-
"hidden_layers=$(existing.hidden_layers)\n" *
69+
"hidden_layers=$(existing.hidden_layers), history_size=$(existing.history_size)\n" *
6870
" Supplied: mps=$(cfg.mps), layer_size=$(cfg.layer_size), " *
69-
"hidden_layers=$(cfg.hidden_layers)\n" *
71+
"hidden_layers=$(cfg.hidden_layers), history_size=$(cfg.history_size)\n" *
7072
"These parameters must match the existing checkpoint. " *
7173
"Use a different cp_path to start a new training run.",
7274
)
@@ -85,6 +87,7 @@ function save_model_config(cfg::ModelConfig, cp_path::String)
8587
"mps" => cfg.mps,
8688
"layer_size" => cfg.layer_size,
8789
"hidden_layers" => cfg.hidden_layers,
90+
"history_size" => cfg.history_size,
8891
),
8992
"training" => Dict(
9093
"norm_steps" => cfg.norm_steps,
@@ -120,6 +123,7 @@ function load_model_config(cp_path::String)::Union{ModelConfig,Nothing}
120123
mps=arch["mps"],
121124
layer_size=arch["layer_size"],
122125
hidden_layers=arch["hidden_layers"],
126+
history_size=Int(get(arch, "history_size", 1)),
123127
norm_steps=train["norm_steps"],
124128
types_updated=Int.(train["types_updated"]),
125129
types_noisy=Int.(train["types_noisy"]),

0 commit comments

Comments
 (0)