diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml new file mode 100644 index 0000000..f9e6ff7 --- /dev/null +++ b/.github/workflows/CI.yml @@ -0,0 +1,47 @@ +name: CI + +on: + push: + branches: + - main + tags: ['*'] + pull_request: + branches: + - main + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +jobs: + test: + name: Test - Julia ${{ matrix.version }} + runs-on: ubuntu-latest + timeout-minutes: 120 + strategy: + fail-fast: false + matrix: + version: + - '1.12' # minimum supported (Octopus requires Julia ≥ 1.12; see Project.toml [compat]) + - '1' # latest stable 1.x + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: ${{ matrix.version }} + arch: x64 + - uses: julia-actions/cache@v3 + - uses: julia-actions/julia-buildpkg@v1 + # The suite auto-detects the absence of a GPU (CUDA.functional() == false + # on GitHub-hosted runners) and runs the CPU code path, so no GPU is needed. + - uses: julia-actions/julia-runtest@v1 + - uses: julia-actions/julia-processcoverage@v1 + if: always() + - uses: codecov/codecov-action@v4 + if: always() + with: + files: lcov.info + fail_ci_if_error: false + env: + CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} diff --git a/.github/workflows/Documenter.yml b/.github/workflows/Documenter.yml index 7988cab..fdb8b24 100644 --- a/.github/workflows/Documenter.yml +++ b/.github/workflows/Documenter.yml @@ -28,7 +28,7 @@ jobs: - name: "Set up Julia" uses: julia-actions/setup-julia@v3 with: - version: '1.11' + version: '1.12' # Octopus 0.2 requires Julia ≥ 1.12 (see Project.toml [compat]) arch: x64 - name: "Copy readme to doc" diff --git a/Project.toml b/Project.toml index 7ab247c..1e61b3f 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "GraphNetSim" uuid = "5ff66f56-808c-48e7-ac84-dd29877231f8" -version = "0.1.2" +version = "0.2.0" authors = ["Josef Jouaux ", "JT "] [deps] @@ -79,7 +79,7 @@ Statistics = "1" Test = "1" Zygote = "0.6, 0.7" cuDNN = "1.4 - 1" -julia = "1.11" +julia = "1.12" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" diff --git a/README.md b/README.md index 21bf3eb..bd416df 100644 --- a/README.md +++ b/README.md @@ -23,7 +23,7 @@ The package is build upon [**GraphNetCore.jl**](https://github.com/una-auxme/Gra ## Requirements -- **Julia ≥ 1.11** +- **Julia ≥ 1.12** - Built on [GraphNetCore.jl](https://github.com/una-auxme/GraphNetCore.jl) **v0.4**, which uses a [Lux.jl](https://github.com/LuxDL/Lux.jl) `TrainState` and pulls in a CUDA-capable stack. - A CUDA-capable GPU is recommended for training (falls back to CPU when CUDA is unavailable). diff --git a/docs/Project.toml b/docs/Project.toml index b283db2..efc8e8f 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,6 +1,7 @@ [deps] Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +GraphNetSim = "5ff66f56-808c-48e7-ac84-dd29877231f8" Optuna = "a5d0552b-b2dc-4f08-ac5c-85ca7d701b92" [compat] -julia = "1.11" +julia = "1.12" diff --git a/docs/make.jl b/docs/make.jl index 6548261..8ac9767 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -21,6 +21,8 @@ makedocs(; pages=[ "Home" => "index.md", "Loading Data" => "loading_data.md", + "Examples" => "examples.md", + "Hyperparameter Optimization" => "hyperparameter_optimization.md", "API Reference" => "api.md", ], ) diff --git a/docs/src/examples.md b/docs/src/examples.md new file mode 100644 index 0000000..ad9ab10 --- /dev/null +++ b/docs/src/examples.md @@ -0,0 +1,87 @@ +# Examples + +Runnable examples live in the [`example/`](https://github.com/una-auxme/GraphNetSim.jl/tree/main/example) +directory of the repository. The two smallest — [BallisticSmall](#BallisticSmall) and +[DamBreakSmall](#DamBreakSmall) — are self-contained: on first run they generate their dataset into +`data/` via the generators in `test/generators.jl`, so no external download is required. Run them from +the repository root with the package's own project environment: + +```bash +julia --project example/BallisticSmall/BallisticSmall.jl +``` + +## BallisticSmall + +A tiny ballistic dataset (10 particles, no boundary nodes, linear drag physics) — the simplest +end-to-end example, and a good first run to confirm your setup works. + +[`example/BallisticSmall/BallisticSmall.jl`](https://github.com/una-auxme/GraphNetSim.jl/tree/main/example/BallisticSmall/BallisticSmall.jl) +walks through the recommended multi-phase workflow: + +1. **DerivativeTraining** — fast initial training against precomputed derivatives (no ODE solve per step). +2. **[`BatchingStrategy`](@ref GraphNetSim.BatchingStrategy)** fine-tuning — ODE-based loss over the trajectory. +3. **[`MultipleShooting`](@ref GraphNetSim.MultipleShooting)** fine-tuning — trajectory split into intervals with a continuity penalty. +4. **[`eval_network`](@ref GraphNetSim.eval_network)** — long-horizon rollout on the test split, then `visualize_eval` to export VTK HDF5 for ParaView. + +Because there are no boundary particles, `types_updated = [1]` predicts every particle. + +## DamBreakSmall + +A tiny 2D weakly-compressible SPH dam break (9 fluid + 9 boundary particles). Like +[BallisticSmall](#BallisticSmall), but with boundary nodes — so it is the reference environment for +both a complete training run and a hyperparameter search. + +### Full training pipeline + +[`example/DamBreakSmall/DamBreakSmall.jl`](https://github.com/una-auxme/GraphNetSim.jl/tree/main/example/DamBreakSmall/DamBreakSmall.jl) +runs the same four-step pipeline as BallisticSmall, updating only the fluid particles +(`types_updated = [2]`). Offline normalization statistics are precomputed once with +[`data_minmax`](@ref GraphNetSim.data_minmax) and [`data_meanstd`](@ref GraphNetSim.data_meanstd) so +training can run with `norm_steps=0`. + +### Hyperparameter optimization with Optuna + +[`example/DamBreakSmall/DamBreakSmall_optuna.jl`](https://github.com/una-auxme/GraphNetSim.jl/tree/main/example/DamBreakSmall/DamBreakSmall_optuna.jl) +runs an automated hyperparameter search over the same dataset using +[Optuna.jl](https://github.com/una-auxme/Optuna.jl). It uses the ask/tell interface: each trial trains +a GNN with [`DerivativeTraining`](@ref GraphNetSim.DerivativeTraining) for a fixed number of steps and +reports the best validation loss returned by [`train_network`](@ref GraphNetSim.train_network). + +Optuna is an extra dependency, provided by this example's own `Project.toml`, so run it with that +environment activated: + +```bash +julia --project=example/DamBreakSmall example/DamBreakSmall/DamBreakSmall_optuna.jl +``` + +Searched hyperparameters: + +| Group | Parameters | +| --- | --- | +| Architecture | `mps`, `layer_size`, `hidden_layers` | +| Optimiser | `optimizer` (Adam / AdamW / RAdam), `lr`, `lr_decay_ratio`, `weight_decay` (AdamW only) | +| Regularisation | `noise_std` | +| Normalisation | `norm_type` (`:minmax` / `:meanstd`) | +| Training | `random_sampling`, `window_size` | + +Key properties: + +- **Sampler / pruner** — a TPE sampler with a median pruner drops unpromising trials early. +- **Both normalization statistics are precomputed** with [`update_meta!`](@ref GraphNetSim.update_meta!) + (once for `:minmax`, once for `:meanstd`), so a trial only selects between them via `norm_type`. +- **Resumable** — the study is persisted in a SQLite database and trial artifacts on disk, so re-running + the script continues from where it left off until the target trial count is reached. + +When the run finishes, the best trial, its parameters, and its validation loss are printed. Adjust +`n_trials` and the per-trial `n_steps` at the top of the script to trade search breadth against +wall-clock time. + +## Further scripts + +The remaining subfolders of [`example/`](https://github.com/una-auxme/GraphNetSim.jl/tree/main/example) +— `Ballistics`, `DamBreak`, `Duese`, `GradientDiagnostics`, `RuntimeBenchmark`, and `WaterRamps` — +hold research and benchmarking material: training variants, ablations, evaluation/visualization +utilities, SLURM (`.sbatch`) cluster job scripts, and comparison harnesses. They target larger +datasets that are **not** bundled with the repository and often assume specific hardware, so treat +them as references rather than turnkey tutorials. Notably, `WaterRamps/WaterRamps_optuna.jl` mirrors +the DamBreakSmall Optuna search for the (external) WaterRamps dataset. diff --git a/docs/src/hyperparameter_optimization.md b/docs/src/hyperparameter_optimization.md new file mode 100644 index 0000000..7ad2045 --- /dev/null +++ b/docs/src/hyperparameter_optimization.md @@ -0,0 +1,162 @@ +# Hyperparameter Optimization + +GraphNetSim integrates with [Optuna.jl](https://github.com/una-auxme/Optuna.jl) to automate the search +for good GNN-simulator hyperparameters. This page documents the workflow implemented by the runnable +example +[`example/DamBreakSmall/DamBreakSmall_optuna.jl`](https://github.com/una-auxme/GraphNetSim.jl/tree/main/example/DamBreakSmall/DamBreakSmall_optuna.jl); +the same pattern applies to any dataset. + +## Requirements + +Optuna is an extra dependency, so activate an environment that provides it. The DamBreakSmall example +ships its own `Project.toml` with Optuna, OrdinaryDiffEq, and Optimisers, so run the script with that +environment from the repository root: + +```bash +julia --project=example/DamBreakSmall example/DamBreakSmall/DamBreakSmall_optuna.jl +``` + +## Workflow overview + +The workflow has four parts: a **persistent study**, an **objective** that trains one model per trial, +an **ask/tell loop** that samples the search space, and **result inspection**. A single trial trains a +GNN with sampled hyperparameters and reports the best validation loss returned by +[`train_network`](@ref GraphNetSim.train_network) — that is, the loss of the periodic ODE rollout on +the validation split. Minimizing that value across trials is the optimization objective. + +### 1. A persistent, resumable study + +The study is backed by a SQLite database and an on-disk artifact store, so re-running the script +continues an existing study rather than starting over: + +```julia +storage_url = create_sqlite_url(database_url, database_name) +storage = RDBStorage(storage_url) +artifact_store = FileSystemArtifactStore(artifact_path) + +study = Study( + study_name, + artifact_store, + storage; + sampler=TPESampler(), # Tree-structured Parzen Estimator + pruner=MedianPruner(5, 1), # stop trials worse than the running median + direction="minimize", + load_if_exists=true, # resume an existing study of the same name +) +``` + +- **Sampler** — `TPESampler` models the relationship between hyperparameters and loss and proposes + promising configurations; swap in another sampler to change the search strategy. +- **Pruner** — `MedianPruner` terminates unpromising trials early (after a startup grace period) by + comparing a trial's reported loss against previous trials. +- **`load_if_exists=true`** — combined with the SQLite storage, this is what makes the run resumable. + +### 2. The objective — one training run per trial + +The objective converts sampled parameters into a training configuration, runs +[`train_network`](@ref GraphNetSim.train_network), and returns the best validation loss. Each trial +trains into a fresh temporary checkpoint directory so trials do not interfere: + +```julia +function objective(trial::Trial; params) + cp_path = mktempdir() + + opt = if params[:optimizer] == "Adam" + Adam(params[:lr]) + elseif params[:optimizer] == "AdamW" + AdamW(; eta=params[:lr], lambda=params[:weight_decay]) + else + RAdam(params[:lr]) + end + + min_val_loss = train_network( + opt, ds_path, cp_path; + training_strategy=DerivativeTraining(; + random=params[:random_sampling], window_size=params[:window_size] + ), + steps=n_steps, checkpoint=cp_interval, + mps=params[:mps], layer_size=params[:layer_size], hidden_layers=params[:hidden_layers], + noise_stddevs=[params[:noise_std]], + norm_steps=0, norm_type=params[:norm_type], + optimizer_learning_rate_start=params[:lr], + optimizer_learning_rate_stop=params[:lr] * params[:lr_decay_ratio], + # ... fixed args: types_updated, types_noisy, solver_valid, use_cuda, ... + ) + + report(trial, Float64(min_val_loss), 1) # feed the pruner + should_prune(trial) && return nothing + + upload_artifact(study, trial, Dict(String(k) => v for (k, v) in pairs(params))) + return Float64(min_val_loss) +end +``` + +`report` hands the trial's loss to the pruner; `should_prune` then decides whether to abandon it; +`upload_artifact` records the trial's hyperparameters for later inspection. + +### 3. The ask/tell loop — sampling the search space + +Each iteration `ask`s the study for a trial, draws hyperparameters with the `suggest_*` family, runs +the objective, and `tell`s the study the result (or that it was pruned): + +```julia +trial = ask(study) + +mps = suggest_int(trial, "mps", 3, 10) +layer_size = suggest_categorical(trial, "layer_size", [32, 64, 128]) +lr = suggest_float(trial, "lr", 1.0e-5, 1.0e-3; log=true) # log-scale +norm_type = Symbol(suggest_categorical(trial, "norm_type", ["minmax", "meanstd"])) +# ... remaining suggestions ... + +params = (; mps, layer_size, lr=Float32(lr), norm_type, #= ... =#) +score = objective(trial; params) + +if isnothing(score) + tell(study, trial; prune=true) +else + tell(study, trial, score) +end +``` + +Use `suggest_int` / `suggest_categorical` / `suggest_float` (with `log=true` for scale-free +quantities like learning rates) to declare each dimension. The example searches architecture +(`mps`, `layer_size`, `hidden_layers`), optimiser (`optimizer`, `lr`, `lr_decay_ratio`, +`weight_decay`), regularisation (`noise_std`), normalisation (`norm_type`), and training strategy +(`random_sampling`, `window_size`). + +### 4. Resume-awareness and results + +Because the study is persistent, the loop counts already-completed trials so restarts converge on a +fixed total instead of adding a fresh batch each time: + +```julia +n_completed = length(study.study.trials) +n_remaining = max(0, n_trials - n_completed) +``` + +When the run finishes, inspect the outcome with `best_trial(study)`, `best_params(study)`, and +`best_value(study)`. + +## Adapting it to your own dataset + +1. **Point at your data** — set `ds_path`, and generate or provide `train.h5` / `valid.h5` / + `test.h5` + `meta.json`. +2. **Precompute normalization** — if you search `norm_type`, run [`update_meta!`](@ref GraphNetSim.update_meta!) + once per statistic you want available (`:minmax` and/or `:meanstd`); a trial then only selects + between them. +3. **Set the fixed budget** — `n_steps` per trial, `cp_interval` (validation cadence), `n_trials`, + and the simulation interval (`dt`, `tstop`). These trade search breadth against wall-clock time. +4. **Edit the search space** — add or remove `suggest_*` calls, mirror them in the `params` + NamedTuple, and forward them to `train_network`. + +!!! warning "Only tune parameters the API actually exposes" + Every keyword forwarded to `train_network` must be an [`Args`](@ref GraphNetSim.Args) field, and + every keyword to a strategy constructor must exist on that strategy (for example, + [`DerivativeTraining`](@ref GraphNetSim.DerivativeTraining) accepts only `window_size` and + `random`). Passing an unknown keyword errors when the trial builds its configuration. + +## Another instance + +[`example/WaterRamps/WaterRamps_optuna.jl`](https://github.com/una-auxme/GraphNetSim.jl/tree/main/example/WaterRamps/WaterRamps_optuna.jl) +applies this same workflow to the (external) WaterRamps dataset, which is useful as a larger-scale +reference — see [Examples](@ref) for why those research scripts are not turnkey. diff --git a/example/DamBreakSmall/DamBreakSmall_optuna.jl b/example/DamBreakSmall/DamBreakSmall_optuna.jl new file mode 100644 index 0000000..ae086d4 --- /dev/null +++ b/example/DamBreakSmall/DamBreakSmall_optuna.jl @@ -0,0 +1,226 @@ +# +# Copyright (c) 2026 Josef Kircher +# Licensed under the MIT license. See LICENSE file in the project root for details. +# +# Hyperparameter optimization for the DamBreakSmall dataset using Optuna.jl. +# +# Uses the ask/tell interface for single-step trials: each trial trains a +# GNN with sampled hyperparameters and reports the best validation loss. +# +# Searched hyperparameters: +# Architecture: mps, layer_size, hidden_layers +# Optimiser: optimizer, lr, lr_decay_ratio, weight_decay (AdamW only) +# Regularisation: noise_std +# Normalisation: norm_type +# Training: random_sampling, window_size +# +# Usage: +# julia --project=example/DamBreakSmall example/DamBreakSmall/DamBreakSmall_optuna.jl +# +# The study is persisted in a SQLite database, so re-running the script +# continues from where it left off. +# + +using GraphNetSim +using Optuna +import OrdinaryDiffEq: Euler +import Optimisers: Adam, AdamW, RAdam + +# ── Dataset setup ──────────────────────────────────────────────────────── + +include(joinpath(dirname(dirname(@__DIR__)), "test", "generators.jl")) +let _gen_dir = joinpath(dirname(dirname(@__DIR__)), "data", "dam_break_small") + if _needs_generation(_gen_dir) + @info "Generating dataset: dam_break_small" + _GenDamBreak.generate(_gen_dir) + @info " Done." + end +end + +ds_path = "data/dam_break_small" + +# Precompute both normalization statistics (reusable across trials). +# `update_meta!` writes one statistic type per call, so run both — the trial +# then selects between them via `norm_type` (:minmax / :meanstd). +update_meta!(ds_path, :minmax) +update_meta!(ds_path, :meanstd) + +# ── Fixed training parameters ──────────────────────────────────────────── + +types_updated = [2] +types_noisy = [2] +cuda = true +tstart = 0.0f0 +dt = 0.001f0 +tstop = 0.079f0 +n_steps = 1500 # derivative training steps per trial +cp_interval = 500 # checkpoint (and validate) every N steps + +# ── Optuna study setup ────────────────────────────────────────────────── + +database_url = "data/dam_break_small/optuna" +database_name = "hpo_db" +study_name = "dam-break-small-hpo" +artifact_path = "data/dam_break_small/optuna/artifacts" + +storage_url = create_sqlite_url(database_url, database_name) +storage = RDBStorage(storage_url) +artifact_store = FileSystemArtifactStore(artifact_path) + +study = Study( + study_name, + artifact_store, + storage; + sampler=TPESampler(), + pruner=MedianPruner(5, 1), + direction="minimize", + load_if_exists=true, +) + +# ── Objective function ────────────────────────────────────────────────── + +function objective(trial::Trial; params) + cp_path = mktempdir() + + # Build optimizer + opt = if params[:optimizer] == "Adam" + Adam(params[:lr]) + elseif params[:optimizer] == "AdamW" + AdamW(; eta=params[:lr], lambda=params[:weight_decay]) + else # RAdam + RAdam(params[:lr]) + end + + lr_stop = params[:lr] * params[:lr_decay_ratio] + + min_val_loss = train_network( + opt, + ds_path, + cp_path; + training_strategy=DerivativeTraining(; + random=params[:random_sampling], window_size=params[:window_size] + ), + steps=n_steps, + checkpoint=cp_interval, + types_updated=types_updated, + types_noisy=types_noisy, + noise_stddevs=[params[:noise_std]], + mps=params[:mps], + layer_size=params[:layer_size], + hidden_layers=params[:hidden_layers], + norm_steps=0, + norm_type=params[:norm_type], + use_cuda=cuda, + solver_valid=Euler(), + solver_valid_dt=dt, + optimizer_learning_rate_start=params[:lr], + optimizer_learning_rate_stop=lr_stop, + show_progress_bars=false, + ) + + # Report validation loss for pruning + report(trial, Float64(min_val_loss), 1) + + if should_prune(trial) + return nothing + end + + # Upload trial hyperparameters and result as artifact + upload_artifact(study, trial, Dict(String(k) => v for (k, v) in pairs(params))) + + return Float64(min_val_loss) +end + +# ── Optimization loop ─────────────────────────────────────────────────── + +n_trials = 50 + +# Resume-aware: count already-completed trials so restarts reach the +# target total rather than adding n_trials on top of previous runs. +n_completed = length(study.study.trials) +n_remaining = max(0, n_trials - n_completed) +if n_completed > 0 + println( + "Resuming study: $n_completed trials already completed, $n_remaining remaining." + ) +end + +for i in 1:n_remaining + trial_num = n_completed + i + println("\n" * "="^60) + println(" Trial $trial_num / $n_trials") + println("="^60) + + trial = ask(study) + + # --- Architecture --- + mps = suggest_int(trial, "mps", 3, 10) + layer_size = suggest_categorical(trial, "layer_size", [32, 64, 128]) + hidden_layers = suggest_int(trial, "hidden_layers", 1, 3) + + # --- Optimizer --- + optimizer = suggest_categorical(trial, "optimizer", ["Adam", "AdamW", "RAdam"]) + lr = suggest_float(trial, "lr", 1.0e-5, 1.0e-3; log=true) + lr_decay_ratio = suggest_float(trial, "lr_decay_ratio", 0.001, 0.1; log=true) + weight_decay = if optimizer == "AdamW" + suggest_float(trial, "weight_decay", 1.0e-6, 1.0e-2; log=true) + else + 0.0 + end + + # --- Regularisation --- + noise_std = suggest_float(trial, "noise_std", 1.0e-6, 1.0e-3; log=true) + + # --- Normalisation --- + norm_type = Symbol(suggest_categorical(trial, "norm_type", ["minmax", "meanstd"])) + + # --- Training strategy --- + random_sampling = suggest_categorical(trial, "random_sampling", [true, false]) + window_size = suggest_int(trial, "window_size", 0, 5) + + params = (; + mps, + layer_size, + hidden_layers, + optimizer, + lr=Float32(lr), + lr_decay_ratio=Float32(lr_decay_ratio), + weight_decay=Float32(weight_decay), + noise_std=Float32(noise_std), + norm_type, + random_sampling, + window_size, + ) + + println( + " Architecture: mps=$mps, layer_size=$layer_size, hidden_layers=$hidden_layers" + ) + println( + " Optimizer: $optimizer, lr=$(round(lr; sigdigits=3)), decay_ratio=$(round(lr_decay_ratio; sigdigits=2))", + ) + if optimizer == "AdamW" + println(" Weight decay: $(round(weight_decay; sigdigits=3))") + end + println(" Noise std: $(round(Float64(noise_std); sigdigits=3))") + println(" Norm type: $norm_type") + println(" Training: random=$random_sampling, window=$window_size") + + score = objective(trial; params) + + if isnothing(score) + tell(study, trial; prune=true) + println(" -> PRUNED") + else + tell(study, trial, score) + println(" -> val_loss = $score") + end +end + +# ── Results ───────────────────────────────────────────────────────────── + +println("\n" * "="^60) +println(" OPTIMIZATION COMPLETE ($n_trials trials)") +println("="^60) +println("Best trial: ", best_trial(study)) +println("Best params: ", best_params(study)) +println("Best value: ", best_value(study)) diff --git a/example/DamBreakSmall/Project.toml b/example/DamBreakSmall/Project.toml new file mode 100644 index 0000000..ee2be45 --- /dev/null +++ b/example/DamBreakSmall/Project.toml @@ -0,0 +1,13 @@ +[deps] +GraphNetSim = "5ff66f56-808c-48e7-ac84-dd29877231f8" +Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" +Optuna = "a5d0552b-b2dc-4f08-ac5c-85ca7d701b92" +OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" + +[sources.GraphNetSim] +path = "../.." + +[compat] +Optimisers = "0.4, 1" +Optuna = "0.2.1" +OrdinaryDiffEq = "6.85 - 6" diff --git a/src/GraphNetSim.jl b/src/GraphNetSim.jl index b453489..82941fe 100644 --- a/src/GraphNetSim.jl +++ b/src/GraphNetSim.jl @@ -56,7 +56,6 @@ import Setfield: @set! include("utils.jl") include("graph.jl") include("solve.jl") -include("rollout_history.jl") include("dataset.jl") include("visualize.jl") include("config.jl") @@ -143,52 +142,38 @@ Configuration structure for training and evaluating Graph Neural Network simulat optimizer_learning_rate_start::Float32 = 1.0f-4 optimizer_learning_rate_stop::Union{Nothing,Float32} = nothing norm_type::Symbol = :online - history_size::Int = 1 neighbor_backend::Symbol = :pointneighbors save_step::Bool = false on_grad::Union{Nothing,Function} = nothing on_valid::Union{Nothing,Function} = nothing end -function _validate_history_args(args::Args) - args.history_size ≥ 1 || throw( - ArgumentError("history_size must be ≥ 1, got $(args.history_size)") - ) - args.history_size == 1 && return - args.training_strategy isa DerivativeTraining || throw( - ArgumentError( - "history_size > 1 is only supported with DerivativeTraining; got " * - "$(typeof(args.training_strategy)). ODE-based strategies will be " * - "extended in a follow-up plan.", - ), - ) - args.solver_valid isa Euler || throw( - ArgumentError( - "history_size > 1 requires solver_valid = Euler() (sliding-buffer " * - "rollout is fixed-step only); got $(typeof(args.solver_valid)).", - ), - ) - isnothing(args.solver_valid_dt) && throw( - ArgumentError( - "history_size > 1 requires solver_valid_dt to be set explicitly " * - "(Euler is fixed-step).", - ), - ) - return -end +""" + _check_bounds_consistency(model_bounded::Bool, dataset_bounded::Bool) -function _validate_history_meta(meta::Dict, args::Args) - args.history_size == 1 && return - allowed = ("velocity", "wall_distance") - bad = [f for f in meta["input_features"] if !(f in allowed)] - isempty(bad) || throw( - ArgumentError( - "history_size > 1 requires input_features ⊆ $(allowed); got " * - "extras $(bad). Drop position from input_features or set " * - "history_size = 1.", - ), - ) - return +Guard against loading a checkpoint whose boundary (wall-distance) feature does not match +the dataset. A bounded model has a wider encoder input than an unbounded one, so the two +must agree. Throws an `ArgumentError` describing the mismatch; a no-op when they match. +""" +function _check_bounds_consistency(model_bounded::Bool, dataset_bounded::Bool) + if model_bounded && !dataset_bounded + throw( + ArgumentError( + "The loaded model was trained with boundary (wall-distance) " * + "features, but the dataset has no \"bounds\" in its meta.json. " * + "Use a bounded dataset, or train a new model at a different cp_path.", + ), + ) + elseif !model_bounded && dataset_bounded + throw( + ArgumentError( + "The dataset defines \"bounds\" in its meta.json, but the loaded " * + "model was trained without boundary (wall-distance) features. " * + "Remove \"bounds\", or train a new model at a different cp_path.", + ), + ) + end + return nothing end """ @@ -235,11 +220,7 @@ function calc_norms(dataset, device, args) for feature in dataset.meta["feature_names"] feature_dim = dataset.meta["features"][feature]["dim"] if feature in input_features - if feature == "velocity" - quantities += feature_dim * args.history_size - else - quantities += feature_dim - end + quantities += feature_dim end if getfield( @@ -401,7 +382,10 @@ function calc_norms(dataset, device, args) end end end - if n_node_types(dataset.meta) > 1 || "wall_distance" in input_features + # Wall distance is driven purely by the presence of a domain box (`meta["bounds"]`): + # reserve its `2 * dims` rows iff the dataset is bounded. An unbounded dataset simply + # produces a shorter input vector (no wall block). + if haskey(dataset.meta, "bounds") quantities += length(dataset.meta["bounds"]) * 2 end @@ -474,29 +458,12 @@ function train_network(opt, ds_path, cp_path; kws...) layer_size=existing_cfg.layer_size, hidden_layers=existing_cfg.hidden_layers, norm_type=existing_cfg.norm_type, - history_size=existing_cfg.history_size, ), NamedTuple(kws), ) end args = Args(; kws...) - _validate_history_args(args) - - save_model_config( - ModelConfig(; - mps=args.mps, - layer_size=args.layer_size, - hidden_layers=args.hidden_layers, - norm_steps=args.norm_steps, - types_updated=args.types_updated, - types_noisy=args.types_noisy, - noise_stddevs=args.noise_stddevs, - norm_type=args.norm_type, - history_size=args.history_size, - ), - cp_path, - ) if CUDA.functional() && args.use_cuda @info "Training on CUDA GPU..." @@ -516,17 +483,35 @@ function train_network(opt, ds_path, cp_path; kws...) ds_train.meta["types_noisy"] = args.types_noisy ds_train.meta["noise_stddevs"] = args.noise_stddevs ds_train.meta["device"] = device - ds_train.meta["history_size"] = args.history_size ds_train.meta["neighbor_backend"] = args.neighbor_backend ds_valid = Dataset(:valid, ds_path, args) ds_valid.meta["types_updated"] = args.types_updated ds_valid.meta["types_noisy"] = args.types_noisy ds_valid.meta["noise_stddevs"] = args.noise_stddevs ds_valid.meta["device"] = device - ds_valid.meta["history_size"] = args.history_size ds_valid.meta["neighbor_backend"] = args.neighbor_backend ds_valid.meta["training_strategy"] = nothing - _validate_history_meta(ds_train.meta, args) + + # Boundary (wall-distance) features are driven by `meta["bounds"]`. Record whether the + # dataset is bounded and, when resuming, reject a mismatch against the saved model. + dataset_bounded = haskey(ds_train.meta, "bounds") + if !isnothing(existing_cfg) + _check_bounds_consistency(existing_cfg.bounded, dataset_bounded) + end + save_model_config( + ModelConfig(; + mps=args.mps, + layer_size=args.layer_size, + hidden_layers=args.hidden_layers, + norm_steps=args.norm_steps, + types_updated=args.types_updated, + types_noisy=args.types_noisy, + noise_stddevs=args.noise_stddevs, + norm_type=args.norm_type, + bounded=dataset_bounded, + ), + cp_path, + ) @info "Training data loaded!" Threads.nthreads() < 2 && @@ -969,14 +954,12 @@ function eval_network( mps=existing_cfg.mps, layer_size=existing_cfg.layer_size, hidden_layers=existing_cfg.hidden_layers, - history_size=existing_cfg.history_size, ), NamedTuple(kws), ) end args = Args(; kws...) - _validate_history_args(args) if CUDA.functional() && args.use_cuda @info "Evaluating on CUDA GPU..." @@ -993,10 +976,12 @@ function eval_network( println("Loading evaluation data...") ds_test = Dataset(:test, ds_path, args) ds_test.meta["device"] = device - ds_test.meta["history_size"] = args.history_size ds_test.meta["neighbor_backend"] = args.neighbor_backend ds_test.meta["training_strategy"] = nothing - _validate_history_meta(ds_test.meta, args) + + if !isnothing(existing_cfg) + _check_bounds_consistency(existing_cfg.bounded, haskey(ds_test.meta, "bounds")) + end # clear_log(1, false) @info "Evaluation data loaded!" @@ -1121,10 +1106,7 @@ function eval_network!( if length(test_loader) > 1 dt = data["dt"] # TODO dt can be an array? - C = get(ds_test.meta, "history_size", 1) - # With paper-faithful warmup, the first prediction frame is C; the C-1 - # frames before it seed the velocity buffer. - start = Float32((C - 1) * dt) + start = 0.0f0 stop = round((data["trajectory_length"] - 1) * dt; digits=6) saves = start:dt:stop mse_steps = saves @@ -1140,39 +1122,23 @@ function eval_network!( enabled=args.show_progress_bars, ) - sol = if get(ds_test.meta, "history_size", 1) > 1 - rollout_history( - gns, - initial_state, - output_features, - ds_test.meta, - target_features, - node_type, - data["mask"], - data["val_mask"], - saves, - device, - pr, - ) - else - rollout( - solver, - gns, - initial_state, - output_features, - ds_test.meta, - target_features, - node_type, - data["mask"], - data["val_mask"], - start, - stop, - dt, - saves, - device, - pr, - ) - end + sol = rollout( + solver, + gns, + initial_state, + output_features, + ds_test.meta, + target_features, + node_type, + data["mask"], + data["val_mask"], + start, + stop, + dt, + saves, + device, + pr, + ) sol_t, prediction = _extract_trajectory_arrays(sol) timesteps[(ti, "timesteps")] = sol_t @@ -1315,14 +1281,12 @@ function extrapolate_network( mps=existing_cfg.mps, layer_size=existing_cfg.layer_size, hidden_layers=existing_cfg.hidden_layers, - history_size=existing_cfg.history_size, ), NamedTuple(kws), ) end args = Args(; kws...) - _validate_history_args(args) if CUDA.functional() && args.use_cuda @info "Extrapolating on CUDA GPU..." @@ -1339,10 +1303,12 @@ function extrapolate_network( println("Loading evaluation data...") ds_test = Dataset(:test, ds_path, args) ds_test.meta["device"] = device - ds_test.meta["history_size"] = args.history_size ds_test.meta["neighbor_backend"] = args.neighbor_backend ds_test.meta["training_strategy"] = nothing - _validate_history_meta(ds_test.meta, args) + + if !isnothing(existing_cfg) + _check_bounds_consistency(existing_cfg.bounded, haskey(ds_test.meta, "bounds")) + end @info "Evaluation data loaded!" Threads.nthreads() < 2 && @@ -1470,39 +1436,23 @@ function extrapolate_network!( enabled=args.show_progress_bars, ) - sol = if get(ds_test.meta, "history_size", 1) > 1 - rollout_history( - gns, - initial_state, - output_features, - ds_test.meta, - target_features, - node_type, - data["mask"], - data["val_mask"], - saves, - device, - pr, - ) - else - rollout( - solver, - gns, - initial_state, - output_features, - ds_test.meta, - target_features, - node_type, - data["mask"], - data["val_mask"], - start, - stop, - dt, - saves, - device, - pr, - ) - end + sol = rollout( + solver, + gns, + initial_state, + output_features, + ds_test.meta, + target_features, + node_type, + data["mask"], + data["val_mask"], + start, + stop, + dt, + saves, + device, + pr, + ) sol_t, prediction = _extract_trajectory_arrays(sol) npred = size(prediction.pos, 3) diff --git a/src/config.jl b/src/config.jl index 7ef82cb..c0e3c22 100644 --- a/src/config.jl +++ b/src/config.jl @@ -13,10 +13,11 @@ const MODEL_CONFIG_FILENAME = "model_config.json" Persists the minimal set of parameters required to reconstruct the GNN model from a checkpoint without re-specifying them at the call site. -Only the three architecture fields (`mps`, `layer_size`, `hidden_layers`) are -strictly required for model reconstruction. All other fields are derived from -`meta.json` or the JLD2 checkpoint at load time. The training fields are saved -as documentation and may legitimately differ between training phases. +The architecture fields (`mps`, `layer_size`, `hidden_layers`, `bounded`) affect +weight shapes and must match on resume; `bounded` additionally must agree with the +dataset's `meta["bounds"]` presence at load time (a bounded model has a wider encoder +input). The remaining fields are derived from `meta.json` or the JLD2 checkpoint and +are saved as documentation; they may legitimately differ between training phases. ## Fields - `mps`: Number of message passing steps. @@ -27,6 +28,10 @@ as documentation and may legitimately differ between training phases. - `types_noisy`: Node types receiving noise injection during training. - `noise_stddevs`: Per-type noise standard deviations. - `norm_type`: Normalization strategy for Float32 features (`:online`, `:minmax`, `:meanstd`). +- `bounded`: Whether the model was trained with the boundary (wall-distance) node feature, + i.e. whether the training dataset defined `meta["bounds"]`. Part of the architecture: a + bounded model has a wider encoder input than an unbounded one, so this must match the + dataset at load time. """ @kwdef struct ModelConfig format_version::Int = 2 @@ -38,7 +43,7 @@ as documentation and may legitimately differ between training phases. types_noisy::Vector{Int} noise_stddevs::Vector{Float32} norm_type::Symbol = :online - history_size::Int = 1 + bounded::Bool = false end """ @@ -60,15 +65,14 @@ function save_model_config(cfg::ModelConfig, cp_path::String) if !isnothing(existing) if existing.mps != cfg.mps || existing.layer_size != cfg.layer_size || - existing.hidden_layers != cfg.hidden_layers || - existing.history_size != cfg.history_size + existing.hidden_layers != cfg.hidden_layers error( "Architecture mismatch between supplied arguments and saved " * "model config at \"$path\".\n" * " Saved: mps=$(existing.mps), layer_size=$(existing.layer_size), " * - "hidden_layers=$(existing.hidden_layers), history_size=$(existing.history_size)\n" * + "hidden_layers=$(existing.hidden_layers)\n" * " Supplied: mps=$(cfg.mps), layer_size=$(cfg.layer_size), " * - "hidden_layers=$(cfg.hidden_layers), history_size=$(cfg.history_size)\n" * + "hidden_layers=$(cfg.hidden_layers)\n" * "These parameters must match the existing checkpoint. " * "Use a different cp_path to start a new training run.", ) @@ -87,7 +91,7 @@ function save_model_config(cfg::ModelConfig, cp_path::String) "mps" => cfg.mps, "layer_size" => cfg.layer_size, "hidden_layers" => cfg.hidden_layers, - "history_size" => cfg.history_size, + "bounded" => cfg.bounded, ), "training" => Dict( "norm_steps" => cfg.norm_steps, @@ -123,7 +127,7 @@ function load_model_config(cp_path::String)::Union{ModelConfig,Nothing} mps=arch["mps"], layer_size=arch["layer_size"], hidden_layers=arch["hidden_layers"], - history_size=Int(get(arch, "history_size", 1)), + bounded=Bool(get(arch, "bounded", false)), norm_steps=train["norm_steps"], types_updated=Int.(train["types_updated"]), types_noisy=Int.(train["types_noisy"]), diff --git a/src/dataset.jl b/src/dataset.jl index 8009431..41328c5 100644 --- a/src/dataset.jl +++ b/src/dataset.jl @@ -90,7 +90,7 @@ function Dataset(datafile::String, metafile::String, args) ) end - meta = parse(Base.read(metafile), String) + meta = parse(Base.read(metafile, String)) keys_traj = keystraj(datafile) meta["n_trajectories"] = length(keys_traj) meta["keys_trajectories"] = keys_traj @@ -237,6 +237,30 @@ end MLUtils.numobs(ds::Dataset) = ds.meta["n_trajectories"] +""" + _updated_particle_indices(node_types, types_updated, idx, key) + +Indices of the particles whose node type is in `types_updated` — the per-trajectory +training/prediction mask. Throws an `ArgumentError` if the trajectory has none, because an +empty mask silently yields `NaN` losses and zero gradient signal. Non-updated particles +(e.g. boundaries) are expected and fine, as long as at least one updated particle exists. +""" +function _updated_particle_indices(node_types, types_updated, idx, key) + updated = findall(x -> x in types_updated, node_types) + if isempty(updated) + present = sort(unique(node_types)) + throw( + ArgumentError( + "Trajectory $idx (key \"$key\") has no particles of any updated type " * + "(present node types: $present; types_updated=$types_updated). An empty " * + "mask yields NaN losses and no gradient signal. Adjust types_updated, or " * + "exclude this trajectory.", + ), + ) + end + return updated +end + """ MLUtils.getobs!(buffer::Dict{String,Any}, ds::Dataset, idx::Int) @@ -272,11 +296,10 @@ function MLUtils.getobs!(buffer, ds::Dataset, idx) n_out = sum(size(buffer[field], 1) for field in ds.meta["output_features"]) buffer["val_mask"] = ds.meta["device"](ones(Float32, n_out, n_particles)) else - buffer["mask"] = ds.meta["device"]( - Int32.( - findall(x -> x in ds.meta["types_updated"], buffer["node_type"][1, :, 1]) - ), + updated = _updated_particle_indices( + buffer["node_type"][1, :, 1], ds.meta["types_updated"], idx, key ) + buffer["mask"] = ds.meta["device"](Int32.(updated)) val_mask = Float32.( map(x -> x in ds.meta["types_updated"], buffer["node_type"][:, :, 1]) ) @@ -821,7 +844,6 @@ function prepare_trajectory!( if !isnothing(meta["training_strategy"]) && (typeof(meta["training_strategy"]) <: DerivativeStrategy) add_targets!(data, meta["derivative_target_features"], device) - _stack_velocity_history!(data, meta, device) preprocess!( data, meta["input_features"], @@ -846,34 +868,3 @@ function prepare_trajectory!( end return data, meta end - -function _stack_velocity_history!( - data::Dict{String,Any}, meta::Dict{String,Any}, device::Function -) - C = get(meta, "history_size", 1) - C == 1 && return - haskey(data, "velocity") || return - vel = data["velocity"] - T = size(vel, 3) - T >= C || throw( - ArgumentError("trajectory_length=$T is shorter than history_size=$C"), - ) - M = T - C + 1 - - dim, np = size(vel, 1), size(vel, 2) - slices = [reshape(vel[:, :, c:(c + M - 1)], dim, np, 1, M) for c in 1:C] - data["velocity_history"] = device(cat(slices...; dims=3)) - - for key in collect(keys(data)) - key == "velocity_history" && continue - v = data[key] - (v isa AbstractArray) || continue - ndims(v) >= 3 || continue - size(v, ndims(v)) == T || continue - idx = ntuple(i -> i == ndims(v) ? (C:T) : Colon(), ndims(v)) - data[key] = v[idx...] - end - - data["trajectory_length"] = M - return -end diff --git a/src/graph.jl b/src/graph.jl index 84b11c8..aaad493 100644 --- a/src/graph.jl +++ b/src/graph.jl @@ -60,15 +60,25 @@ function build_graph( # fluid_position = data["position"][:, data["mask"], datapoint] current_position = data["position"][:, :, datapoint] - velocity = if get(meta, "history_size", 1) > 1 && haskey(data, "velocity_history") - data["velocity_history"][:, :, :, datapoint] - else - data["velocity"][:, :, datapoint] - end + velocity = data["velocity"][:, :, datapoint] build_graph(gns, current_position, velocity, meta, node_type, data["mask"], device) end +# Materialise `position` into a concrete, device-native matrix for neighbor search. +# +# In `build_graph`, `position === x.x` is a view into the ODE-state ComponentArray +# (on CPU a `ReshapedArray` over a 1-D `SubArray`). Neighbor search dispatches on the +# concrete array type (`point_neighbor_ns(::Array)` vs `(::CuArray)`), and the +# SingleShooting/MultipleShooting backward pass differentiates through this call, so the +# result must be (a) a concrete `Array`/`CuArray` and (b) produced by a +# Zygote-differentiable op. On GPU, `device(x)` yields a `CuArray` (unchanged from the +# original code). On CPU, `collect(x)` yields a dense `Array` via a differentiable copy — +# unlike `adapt_structure(::CPUDevice, ::SubArray)`, whose `SubArray` constructor has no +# adjoint (`Need an adjoint for constructor SubArray`). +_neighbor_positions(device, x) = device(x) +_neighbor_positions(::CPUDevice, x) = collect(x) + """ build_graph(gns::GraphNetCore.GraphNetwork, position, velocity, meta, node_type, mask, device) @@ -94,16 +104,21 @@ function build_graph( gns::GraphNetCore.GraphNetwork, position, velocity, meta, node_type, mask, device ) # TODO check ODE solve and if this is really repeatedly done senders, receivers, rel_displacement, rel_dist_norm = neighbor_search( - device(position), + _neighbor_positions(device, position), Float32(meta["default_connectivity_radius"]), get(meta, "neighbor_backend", :pointneighbors), ) multi_type = n_node_types(meta) > 1 - use_wall = multi_type || "wall_distance" in meta["input_features"] + # Wall distance is driven purely by the presence of a domain box in the dataset meta + # (`meta["bounds"]`) — independent of the node-type count and of `input_features`. A + # bounded dataset gives every particle a distance-to-walls block; an unbounded one + # (no `bounds`) yields a correspondingly shorter input vector. `multi_type` only + # controls the `node_type` one-hot block below. + use_wall = haskey(meta, "bounds") use_position = "position" in meta["input_features"] - vel_norm = _normalize_velocity(gns, velocity) + vel_norm = gns.n_norm["velocity"](velocity) edge_features = device(vcat(rel_displacement, rel_dist_norm) .+ 1.0f-8) @@ -111,16 +126,15 @@ function build_graph( # [position?, velocity, wall_distance?, node_type?] # Build this as a single `vcat` over the present blocks rather than chaining one # `vcat` per block. Chaining re-allocates and re-copies the growing feature matrix - # at every link — it allocated 1.7-2.1x the final array in transient GPU garbage - # (worse with velocity history, since the tall velocity block is recopied each - # link), whereas a single `vcat` allocates exactly the output once (~4.8x faster + # at every link — it allocated 1.7-2.1x the final array in transient GPU garbage, + # whereas a single `vcat` allocates exactly the output once (~4.8x faster # forward, ~1/3 less total fwd+bwd allocation; values/gradients bit-identical). # The block tuple uses immutable splats (no `push!`) so the expression stays # Zygote-differentiable on the training RHS; a lone present block returns untouched. nf_blocks = ( (use_position ? (gns.n_norm["position"](position),) : ())..., vel_norm, - (use_wall ? (_wall_distance(position, mask, meta, device),) : ())..., + (use_wall ? (_wall_distance(position, meta, device),) : ())..., (multi_type ? (node_type,) : ())..., ) nf = length(nf_blocks) == 1 ? nf_blocks[1] : vcat(nf_blocks...) @@ -131,33 +145,17 @@ function build_graph( ) end -function _normalize_velocity(gns::GraphNetCore.GraphNetwork, velocity) - nv = gns.n_norm["velocity"] - if ndims(velocity) == 2 - return nv(velocity) - elseif ndims(velocity) == 3 - C = size(velocity, 3) - slices = map(1:C) do c - slc = velocity[:, :, c] - return nv isa NormaliserOnline ? nv(slc, c == C) : nv(slc) - end - stacked = cat(slices...; dims=3) - permuted = permutedims(stacked, (1, 3, 2)) - return reshape(permuted, :, size(velocity, 2)) - else - throw( - ArgumentError( - "velocity must be 2D (dim, particles) or 3D (dim, particles, C); " * - "got ndims=$(ndims(velocity))", - ), - ) - end -end - -function _wall_distance(position, mask, meta, device) - if length(mask) == size(position, 2) - return device(ones(Float32, size(position)...)) - end +# Clipped per-particle distance to the domain box `meta["bounds"]`, following +# DeepMind's "distance to walls" node feature. This depends only on `meta["bounds"]` +# and positions — NOT on node types or on the presence of boundary particles — so it +# is computed for every particle whenever the feature is enabled. Returns `2 * dims` +# rows (low + high bound per spatial dimension), matching the width reserved by +# `calc_norms`. An implicit domain box with no boundary particles is fully supported; +# the only hard requirement is that `meta["bounds"]` is defined. +function _wall_distance(position, meta, device) + haskey(meta, "bounds") || throw( + ArgumentError("wall_distance requires meta[\"bounds\"] to be defined."), + ) boundaries = device(Float32.(vcat(permutedims.(meta["bounds"])...))) dist_low_bound = position .- boundaries[:, 1] dist_up_bound = boundaries[:, 2] .- position diff --git a/src/rollout_history.jl b/src/rollout_history.jl deleted file mode 100644 index 515a2db..0000000 --- a/src/rollout_history.jl +++ /dev/null @@ -1,109 +0,0 @@ -# -# Copyright (c) 2026 Josef Jouaux, Julian Trommer -# Licensed under the MIT license. See LICENSE file in the project root for details. -# - -""" - rollout_history(gns, initial_state, output_fields, meta, target_fields, - node_type, mask, val_mask, sim_interval, device, pr=nothing) - -Hand-written fixed-step Euler rollout that maintains a `(dim, particles, C)` velocity -history buffer for `meta["history_size"] > 1`. - -Each step: -1. Build graph from current `position` and the C-step velocity buffer. -2. Run the GNN, denormalize outputs, apply `val_mask`. -3. Integrate: `velocity += dt * accel`, `position += dt * velocity`. -4. Slide the buffer one step (drop index 1, append new velocity at index C). - -# Buffer warmup - -If `initial_state` contains a 3D `"velocity_window"` of shape `(dim, particles, C)`, -the buffer is seeded with those C ground-truth velocities (paper-faithful: the model -sees an in-distribution history at step 1). The caller is responsible for pairing -`"position"` with the last frame of the window so the newest buffer slot and the -position correspond to the same time step (matching the training-time pairing). - -If `"velocity_window"` is absent the buffer falls back to `repeat(initial_state["velocity"], C)`, -which is OOD vs. the training distribution — only intended as a debug/legacy path. - -Returns `(t = times, u = states, acc = accelerations)` mirroring the subset of the -`sol.u` interface that `_validation_step` and `_extract_trajectory_arrays` consume. -""" -function rollout_history( - gns::GraphNetCore.GraphNetwork, - initial_state, - output_fields, - meta, - target_fields, - node_type, - mask, - val_mask, - sim_interval, - device, - pr=nothing, -) - C = get(meta, "history_size", 1) - pos = initial_state["position"] - vh = if haskey(initial_state, "velocity_window") - vw = initial_state["velocity_window"] - size(vw, 3) == C || throw( - ArgumentError( - "velocity_window must have last dim == history_size = $C; got " * - "size $(size(vw)).", - ), - ) - device(vw) - else - v0 = initial_state["velocity"] - d, n = size(v0, 1), size(v0, 2) - device(repeat(reshape(v0, d, n, 1); outer=(1, 1, C))) - end - vel = vh[:, :, end] - dim, np = size(vel, 1), size(vel, 2) - dt = Float32(sim_interval[2] - sim_interval[1]) - - indices = [meta["features"][tf]["dim"] for tf in target_fields] - saved = [(; x=copy(pos), dx=copy(vel))] - accs = [device(zeros(Float32, dim, np))] - times = Float32[Float32(sim_interval[1])] - - for k in 1:(length(sim_interval) - 1) - graph = build_graph(gns, pos, vh, meta, node_type, mask, device) - output, st = gns.model(graph, gns.ps, gns.st) - gns.st = st - - denorm = similar(output) - for i in eachindex(output_fields) - r = (sum(indices[1:(i - 1)]) + 1):sum(indices[1:i]) - denorm[r, :] = inverse_data(gns.o_norm[output_fields[i]], output[r, :]) - end - accel = denorm .* val_mask - - vel = vel .+ dt .* accel - pos = pos .+ dt .* vel - vh = cat(vh[:, :, 2:end], reshape(vel, dim, np, 1); dims=3) - - push!(saved, (; x=copy(pos), dx=copy(vel))) - push!(accs, copy(accel)) - push!(times, Float32(sim_interval[k + 1])) - if !isnothing(pr) - next!(pr; showvalues=[(:t, "$(length(saved))")]) - end - end - - if !isnothing(pr) - finish!(pr) - end - - return (t=times, u=saved, acc=accs) -end - -function _extract_trajectory_arrays( - sol::NamedTuple{(:t, :u, :acc)} -) - sol_pos = cpu_device()(cat([u.x for u in sol.u]...; dims=3)) - sol_vel = cat([u.dx for u in sol.u]...; dims=3) - sol_acc = cat(sol.acc...; dims=3) - return sol.t, (pos=sol_pos, vel=sol_vel, acc=sol_acc) -end diff --git a/src/solve.jl b/src/solve.jl index 6960c6c..18f74b7 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -388,31 +388,11 @@ using the bounds stored in `ds_test.meta["features"]["node_type"]`. """ function _prepare_rollout_inputs(data, ds_test, start, dt, device) stepstart = round(Int, ((start / dt) + 1)) - C = get(ds_test.meta, "history_size", 1) - - if C > 1 - # Paper-faithful warmup: `stepstart` is the first prediction frame, and - # the velocity buffer is seeded from the C ground-truth velocities ending - # at `stepstart` (so the buffer's newest slot and the integrator's current - # position align — the same pairing used at training time). - stepstart >= C || throw( - ArgumentError( - "history_size=$C requires start to map to a frame >= $C " * - "(got stepstart=$stepstart). Pass start=(C-1)*dt or later.", - ), - ) - window = (stepstart - C + 1):stepstart - initial_state = Dict( - "position" => data["position"][:, :, stepstart], - "velocity" => data["velocity"][:, :, stepstart], - "velocity_window" => data["velocity"][:, :, window], - ) - else - initial_state = Dict( - "position" => data["position"][:, :, stepstart], - "velocity" => data["velocity"][:, :, stepstart], - ) - end + + initial_state = Dict( + "position" => data["position"][:, :, stepstart], + "velocity" => data["velocity"][:, :, stepstart], + ) node_type = device( Float32.( diff --git a/src/strategies.jl b/src/strategies.jl index b67e43b..1142974 100644 --- a/src/strategies.jl +++ b/src/strategies.jl @@ -149,29 +149,10 @@ Inner function for validation of a single trajectory. function _validation_step(t::Tuple, sim_interval, data_interval) gns, data, meta, _, solver, solver_dt, node_type, pr = t - C = get(meta, "history_size", 1) - initial_state = if C > 1 - # First prediction frame is `data_interval[1] = C`; the C velocities ending - # there form the warmup buffer (paper-faithful: the model sees an - # in-distribution velocity history at step 1 of the rollout). - first_frame = first(data_interval) - first_frame >= C || throw( - ArgumentError( - "history_size=$C requires data_interval[1] >= $C; got $first_frame.", - ), - ) - window = (first_frame - C + 1):first_frame - Dict( - "position" => data["position"][:, :, first_frame], - "velocity" => data["velocity"][:, :, first_frame], - "velocity_window" => data["velocity"][:, :, window], - ) - else - Dict( - "position" => data["position"][:, :, 1], - "velocity" => data["velocity"][:, :, 1], - ) - end + initial_state = Dict( + "position" => data["position"][:, :, 1], + "velocity" => data["velocity"][:, :, 1], + ) target_dict = Dict{String,Int32}() for tf in meta["solver_target_features"] @@ -181,46 +162,29 @@ function _validation_step(t::Tuple, sim_interval, data_interval) gt = vcat([data[tf] for tf in meta["solver_target_features"]]...)[ :, data["mask"], data_interval ] - sol = if C > 1 - rollout_history( - gns, - initial_state, - meta["output_features"], - meta, - meta["solver_target_features"], - node_type, - data["mask"], - data["val_mask"], - sim_interval, - meta["device"], - pr, - ) - else - rollout( - solver, - gns, - initial_state, - meta["output_features"], - meta, - meta["solver_target_features"], - node_type, - data["mask"], - data["val_mask"], - Float32(sim_interval[1]), - Float32(sim_interval[end]), - solver_dt, - sim_interval, - meta["device"], - pr, - ) - end + sol = rollout( + solver, + gns, + initial_state, + meta["output_features"], + meta, + meta["solver_target_features"], + node_type, + data["mask"], + data["val_mask"], + Float32(sim_interval[1]), + Float32(sim_interval[end]), + solver_dt, + sim_interval, + meta["device"], + pr, + ) GC.gc() # Run Julia's garbage collector first if CUDA.functional() CUDA.reclaim() # Force garbage collection and free unused memory end sol_pos = [u.x for u in sol.u] - # `sol.u[k]` is the predicted state at frame `data_interval[k]` (paper-faithful) - # or frame `k` (legacy). In both cases, the comparison takes the first + # `sol.u[k]` is the predicted state at frame `k`; the comparison takes the first # `length(data_interval)` saved entries. prediction = cat(sol_pos...; dims=3)[:, data["mask"], 1:length(data_interval)] @@ -1143,11 +1107,6 @@ derivatives. Computes gradients via backpropagation. """ function train_step(strategy::DerivativeStrategy, t::Tuple) gns, data, meta, target_quantities_change, node_type, mask, device, datapoint = t # TODO here own function - velocity_arg = if haskey(data, "velocity_history") - data["velocity_history"][:, :, :, datapoint] - else - data["velocity"][:, :, datapoint] - end loss, gs = Zygote.withgradient( ps -> train_loss( strategy, @@ -1155,7 +1114,7 @@ function train_step(strategy::DerivativeStrategy, t::Tuple) ps, gns, data["position"][:, :, datapoint], - velocity_arg, + data["velocity"][:, :, datapoint], meta, target_quantities_change, node_type, @@ -1245,33 +1204,16 @@ and comparing derivatives with ground truth. - `t::Tuple`: Validation data tuple. """ function validation_step(::DerivativeStrategy, t::Tuple) - _, data, meta, delta, _, _, _, _ = t + _, data, _, delta, _, _, _, _ = t dt = data["dt"] - C = get(meta, "history_size", 1) - if C > 1 - # Paper-faithful: warmup uses frames 1..C of ground truth. The first - # prediction frame is C; the integration covers frames C..T (length - # T - C + 1). We size `sim_interval` with one trailing extra step (matching - # the C=1 convention of integrating one step past the last compared frame). - T = data["trajectory_length"] - L = T - C + 1 - if typeof(dt) <: AbstractArray - base_dt = dt[2] - dt[1] - sim_interval = 0.0:base_dt:(base_dt * L) - else - sim_interval = 0.0:dt:(dt * L) - end - data_interval = C:T # length L + if typeof(dt) <: AbstractArray + sim_interval = dt[1]:(dt[2] - dt[1]):dt[delta] else - if typeof(dt) <: AbstractArray - sim_interval = dt[1]:(dt[2] - dt[1]):dt[delta] - else - # sim_interval = 0.0: dt: dt * t[2]["trajectory_length"] - sim_interval = 0.0:dt:(dt * delta) - end - # data_interval = 1:t[2]["trajectory_length"] - data_interval = 1:(length(sim_interval) - 1) + # sim_interval = 0.0: dt: dt * t[2]["trajectory_length"] + sim_interval = 0.0:dt:(dt * delta) end + # data_interval = 1:t[2]["trajectory_length"] + data_interval = 1:(length(sim_interval) - 1) return _validation_step(t, sim_interval, data_interval) end diff --git a/test/Project.toml b/test/Project.toml index 99bb091..c8afe2e 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,6 +1,7 @@ [deps] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" GraphNetCore = "7809f980-de1b-4f9a-8451-85f041491431" HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" diff --git a/test/runtests.jl b/test/runtests.jl index 6c61a91..44b9cd5 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -22,5 +22,7 @@ include("generate_fixtures.jl") include("test_converters.jl") include("test_normalizer.jl") include("test_datasets.jl") - include("test_history_stack.jl") + include("test_wall_distance.jl") + include("test_trajectory_mask.jl") + include("test_visualize.jl") end diff --git a/test/test_datasets.jl b/test/test_datasets.jl index b7c7f03..8b716fa 100644 --- a/test/test_datasets.jl +++ b/test/test_datasets.jl @@ -422,14 +422,18 @@ for cfg in CONFIGS end # ───────────────────────────────────────────────────────────────── - # Group I: Regression anchor — two-arg Dataset constructor bug + # Group I: three-arg Dataset(datafile, metafile, args) constructor + # Regression anchor for the dataset.jl:93 meta-parse bug (read then + # parse) — with valid files this constructor must build a Dataset. # ───────────────────────────────────────────────────────────────── - @testset "I: Two-arg Dataset constructor bug (dataset.jl:57)" begin - println("Running: I — Two-arg Dataset constructor bug ($(cfg.name))") + @testset "I: three-arg Dataset(datafile, metafile, args) constructor" begin + println("Running: I — three-arg Dataset constructor ($(cfg.name))") args = make_args(cfg) - @test_throws ArgumentError GraphNetSim.Dataset( + ds = GraphNetSim.Dataset( joinpath(cfg.path, "train.h5"), joinpath(cfg.path, "meta.json"), args ) + @test ds isa GraphNetSim.Dataset + @test ds.meta["n_trajectories"] > 0 end # ───────────────────────────────────────────────────────────────── diff --git a/test/test_history_stack.jl b/test/test_history_stack.jl deleted file mode 100644 index 13cc257..0000000 --- a/test/test_history_stack.jl +++ /dev/null @@ -1,223 +0,0 @@ -# -# Regression tests for the Sanchez-Gonzalez 2020 history-stack input -# (Args.history_size + per-particle C-velocity stacking for DerivativeTraining). -# -# Run in isolation: -# julia --project test/test_history_stack.jl -# - -using Test -using GraphNetSim -using Lux -using CUDA -using JLD2 -using JSON -using MLUtils -import OrdinaryDiffEq: Euler, Tsit5 -import Optimisers: Adam - -include(joinpath(@__DIR__, "generate_fixtures.jl")) - -const HAS_CUDA = CUDA.functional() -const DEVICE = HAS_CUDA ? gpu_device() : cpu_device() - -const BALLISTIC = joinpath(@__DIR__, "fixtures", "ballistic_small") -const DAM_BREAK = joinpath(@__DIR__, "fixtures", "dam_break_small") - -function _args_for(path; history_size=1) - is_dam_break = path == DAM_BREAK - return GraphNetSim.Args(; - use_cuda=HAS_CUDA, - show_progress_bars=false, - mps=4, - layer_size=32, - hidden_layers=2, - training_strategy=DerivativeTraining(), - solver_valid=!is_dam_break && history_size == 1 ? Tsit5() : Euler(), - solver_valid_dt=is_dam_break ? 0.001f0 : 0.002f0, - types_updated=is_dam_break ? [2] : [1], - types_noisy=is_dam_break ? [2] : [1], - noise_stddevs=[0.0f0], - norm_steps=0, - history_size=history_size, - ) -end - -function _dam_break_train_kwargs(; history_size=1, kws...) - base = ( - use_cuda=HAS_CUDA, - show_progress_bars=false, - mps=4, - layer_size=32, - hidden_layers=2, - training_strategy=DerivativeTraining(), - solver_valid=Euler(), - solver_valid_dt=0.001f0, - types_updated=[2], - types_noisy=[2], - noise_stddevs=[0.0f0], - norm_steps=0, - history_size=history_size, - ) - return merge(base, NamedTuple(kws)) -end - -@testset "Sanchez-Gonzalez history-stack input" begin - @testset "H1: history_size=1 keeps quantities unchanged (back-compat)" begin - # ballistic_small: single-type, input_features=[velocity], no wall_distance. - # Single-type one-hot is omitted, so quantities = velocity dim = 3. - args = _args_for(BALLISTIC; history_size=1) - ds = GraphNetSim.Dataset(:train, BALLISTIC, args) - ds.meta["device"] = DEVICE - q, _, _, _ = GraphNetSim.calc_norms(ds, DEVICE, args) - @test q == 3 - - # dam_break_small: 2 types, input_features=[velocity], multi-type wall_distance. - # quantities = velocity (2) + onehot (2) + 2*length(bounds) (4) = 8. - args2 = _args_for(DAM_BREAK; history_size=1) - ds2 = GraphNetSim.Dataset(:train, DAM_BREAK, args2) - ds2.meta["device"] = DEVICE - q2, _, _, _ = GraphNetSim.calc_norms(ds2, DEVICE, args2) - @test q2 == 8 - end - - @testset "H2: history_size=5 quantities formula" begin - args = _args_for(BALLISTIC; history_size=5) - ds = GraphNetSim.Dataset(:train, BALLISTIC, args) - ds.meta["device"] = DEVICE - q, _, _, _ = GraphNetSim.calc_norms(ds, DEVICE, args) - # single-type, history=5: 5 * vel_dim(3) = 15 - @test q == 15 - - args2 = _args_for(DAM_BREAK; history_size=5) - ds2 = GraphNetSim.Dataset(:train, DAM_BREAK, args2) - ds2.meta["device"] = DEVICE - q2, _, _, _ = GraphNetSim.calc_norms(ds2, DEVICE, args2) - # multi-type, history=5: 5*2 + 4(walls) + 2(onehot) = 16 - @test q2 == 16 - end - - @testset "H3: training smoke (history_size=5) runs without OOB" begin - mktempdir() do cp_path - # dam_break_small: 4 train trajectories, traj_length=150. - # After C=5 stacking, effective M = 146. Run ~1 outer iteration. - min_val_loss = train_network( - Adam(1.0f-3), - DAM_BREAK, - cp_path; - _dam_break_train_kwargs(; history_size=5)..., - steps=146 * 2, - checkpoint=146, - ) - @test isfinite(min_val_loss) - - # Confirm the persisted ModelConfig recorded history_size. - cfg = GraphNetSim.load_model_config(cp_path) - @test !isnothing(cfg) - @test cfg.history_size == 5 - end - end - - @testset "H4: wall_distance flag on a single-type dataset" begin - mktempdir() do tmpdir - for f in ("meta.json", "train.h5", "valid.h5", "test.h5") - cp(joinpath(BALLISTIC, f), joinpath(tmpdir, f)) - end - meta = JSON.parsefile(joinpath(tmpdir, "meta.json")) - push!(meta["input_features"], "wall_distance") - open(joinpath(tmpdir, "meta.json"), "w") do io - JSON.print(io, meta, 2) - end - - args = _args_for(tmpdir; history_size=5) - ds = GraphNetSim.Dataset(:train, tmpdir, args) - ds.meta["device"] = DEVICE - q, _, _, _ = GraphNetSim.calc_norms(ds, DEVICE, args) - # ballistic dims=3, single-type, history=5, wall_distance enabled: - # 5*3 + 2*length(bounds)(6) = 21. - @test q == 21 - end - end - - @testset "H6: _prepare_rollout_inputs paper-faithful warmup" begin - args = _args_for(DAM_BREAK; history_size=5) - ds = GraphNetSim.Dataset(:test, DAM_BREAK, args) - ds.meta["device"] = DEVICE - ds.meta["history_size"] = 5 - ds.meta["training_strategy"] = nothing - - # Materialize a single trajectory the way DataLoader would - traj = MLUtils.getobs(ds, 1) - traj["dt"] = ds.meta["features"]["acceleration"]["dim"] == 2 ? 0.001f0 : 0.002f0 - - # Test the C>1 path: start = (C-1)*dt should produce stepstart = C and - # velocity_window of shape (dim, np, C) covering frames 1..C. - dt = 0.001f0 - start = Float32((5 - 1) * dt) - initial_state, _, stepstart = GraphNetSim._prepare_rollout_inputs( - traj, ds, start, dt, DEVICE - ) - @test stepstart == 5 - @test haskey(initial_state, "velocity_window") - vw = initial_state["velocity_window"] - @test size(vw, 3) == 5 - @test size(initial_state["position"]) == size(traj["position"])[1:2] - # vw[:, :, end] must equal data["velocity"][:, :, stepstart] - last_slice = Array(vw[:, :, end]) - expected = Array(traj["velocity"][:, :, stepstart]) - @test last_slice ≈ expected - - # start < (C-1)*dt must throw (insufficient warmup history) - @test_throws ArgumentError GraphNetSim._prepare_rollout_inputs( - traj, ds, 0.0f0, dt, DEVICE - ) - - # C = 1 path: stepstart = 1, no velocity_window - args1 = _args_for(DAM_BREAK; history_size=1) - ds1 = GraphNetSim.Dataset(:test, DAM_BREAK, args1) - ds1.meta["device"] = DEVICE - ds1.meta["training_strategy"] = nothing - traj1 = MLUtils.getobs(ds1, 1) - is1, _, ss1 = GraphNetSim._prepare_rollout_inputs(traj1, ds1, 0.0f0, dt, DEVICE) - @test ss1 == 1 - @test !haskey(is1, "velocity_window") - end - - @testset "H5: checkpoint round-trip with history_size=5" begin - mktempdir() do cp_path - min1 = train_network( - Adam(1.0f-3), - DAM_BREAK, - cp_path; - _dam_break_train_kwargs(; history_size=5)..., - steps=146, - checkpoint=146, - ) - cfg = GraphNetSim.load_model_config(cp_path) - @test !isnothing(cfg) - @test cfg.history_size == 5 - - # Architecture-mismatch guard: same cp_path with a different history_size - # must error (proves history_size is part of the persisted architecture). - @test_throws ErrorException train_network( - Adam(1.0f-3), - DAM_BREAK, - cp_path; - _dam_break_train_kwargs(; history_size=3)..., - steps=10, - checkpoint=10, - ) - - # Resume with the original history_size. - min2 = train_network( - Adam(1.0f-3), - DAM_BREAK, - cp_path; - _dam_break_train_kwargs(; history_size=5)..., - steps=146 * 2, - checkpoint=146, - ) - @test isfinite(min2) - end - end -end diff --git a/test/test_trajectory_mask.jl b/test/test_trajectory_mask.jl new file mode 100644 index 0000000..a131f76 --- /dev/null +++ b/test/test_trajectory_mask.jl @@ -0,0 +1,34 @@ +# +# Copyright (c) 2026 Josef Jouaux, Julian Trommer +# Licensed under the MIT license. See LICENSE file in the project root for details. +# +# Regression tests for the per-trajectory updated-particle mask +# (`_updated_particle_indices`, src/dataset.jl). +# +# A multi-type dataset (meta declares >1 node type) may contain individual trajectories +# that hold none of the `types_updated` — e.g. an all-boundary scene when +# `types_updated=[1]` (the default). That produces an EMPTY mask, which used to flow +# silently into training/validation as `NaN` losses (mean over zero elements) and zero +# gradient signal. The mask builder now turns that into a clear, actionable error. +# Non-updated particles (boundaries) are expected and fine as long as ≥1 updated +# particle is present. + +using Test +using GraphNetSim + +@testset "updated-particle mask" begin + @testset "boundaries are fine when an updated particle exists" begin + # types_updated=[1]; particles 2,3 are non-updated boundaries — mask = [1, 4]. + @test GraphNetSim._updated_particle_indices([1, 2, 2, 1], [1], 3, "traj_3") == + [1, 4] + # All particles updated (types_updated covers every present type). + @test GraphNetSim._updated_particle_indices([2, 2], [1, 2], 1, "t") == [1, 2] + end + + @testset "a trajectory with no updated particles is a hard error" begin + # Only type-2 particles present, but types_updated=[1] → empty mask. + @test_throws ArgumentError GraphNetSim._updated_particle_indices( + [2, 2, 2], [1], 7, "traj_7" + ) + end +end diff --git a/test/test_visualize.jl b/test/test_visualize.jl new file mode 100644 index 0000000..a22fe78 --- /dev/null +++ b/test/test_visualize.jl @@ -0,0 +1,116 @@ +# +# Round-trip tests for the VTK HDF5 export (src/visualize.jl). +# +# Build a synthetic trajectory .h5 with the structure `visualize` expects, run +# `visualize` / `visualize_eval`, then re-read the emitted `.vtkhdf` files and +# check the Points/PointData match the input. Only needs HDF5 (no WriteVTK), so +# it runs both standalone under `--project` and inside `Pkg.test`. +# +using GraphNetSim +using Test +using HDF5 + +# Dataset key exactly as visualize.jl reads it: `name * "[$t]"`. +_dsname(name, t) = name * "[" * string(t) * "]" + +# Write an input file laid out as: +# //timesteps (scalar upper bound) +# ///[t] (dim × npts, for t in 1:T) +# Returns a dict (traj, subgroup, name, t) => array of the exact data written. +function _write_synthetic_traj(path; ntraj=1, T=3, npts=5, dim=2, Position="pos", subgroups) + data = Dict{Tuple{Int,String,String,Int},Matrix{Float64}}() + HDF5.h5open(path, "w") do f + for traj in 1:ntraj + g = HDF5.create_group(f, string(traj)) + g["timesteps"] = T + for (sg, params) in subgroups + sgg = HDF5.create_group(g, sg) + for t in 1:T + pos = rand(dim, npts) + sgg[_dsname(Position, t)] = pos + data[(traj, sg, Position, t)] = pos + for p in params + v = rand(dim, npts) + sgg[_dsname(p, t)] = v + data[(traj, sg, p, t)] = v + end + end + end + end + end + return data +end + +@testset "visualize.jl VTK HDF5 export" begin + println("Running: visualize.jl VTK export round-trip") + T, npts, dim = 3, 5, 2 + subgroups = Dict("prediction" => ["vel", "acc", "err"], "gt" => ["vel", "acc"]) + + mktempdir() do dir + inPath = joinpath(dir, "trajectories.h5") + data = _write_synthetic_traj(inPath; T=T, npts=npts, dim=dim, subgroups=subgroups) + + @testset "V1: visualize returns the read datasets" begin + readDict = visualize( + inPath, joinpath(dir, "v1"), "pos", "prediction", ["vel", "acc", "err"] + ) + # pos (1) + 3 params, per timestep + @test length(readDict) == T * 4 + # params are not padded → exact round-trip of what we wrote + @test readDict[(1, "vel", 1)] == data[(1, "prediction", "vel", 1)] + @test readDict[(1, "acc", T)] == data[(1, "prediction", "acc", T)] + end + + @testset "V2: emitted .vtkhdf re-reads to padded Points + PointData" begin + outFolder = joinpath(dir, "v2") + visualize(inPath, outFolder, "pos", "prediction", ["vel", "acc", "err"]) + vtk = joinpath( + outFolder, "1Trajectory", "prediction", "prediction_1Trajectory_1.vtkhdf" + ) + @test isfile(vtk) + HDF5.h5open(vtk, "r") do f + pts = HDF5.read(f, "VTKHDF/Points") + @test size(pts) == (3, npts) # 2D padded to 3D + @test pts[1:2, :] == data[(1, "prediction", "pos", 1)] + @test all(pts[3, :] .== 0) # z padding + @test HDF5.read(f, "VTKHDF/PointData/vel") == + data[(1, "prediction", "vel", 1)] + @test HDF5.read(f, "VTKHDF/NumberOfPoints") == [npts] + @test HDF5.read(f, "VTKHDF/NumberOfCells") == [npts] + end + end + + @testset "V3: visualize_eval writes both prediction and gt trees" begin + outFolder = joinpath(dir, "v3") + visualize_eval(inPath, outFolder) # defaults: pos / prediction / gt + @test isfile( + joinpath( + outFolder, + "1Trajectory", + "prediction", + "prediction_1Trajectory_1.vtkhdf", + ), + ) + @test isfile( + joinpath(outFolder, "1Trajectory", "gt", "gt_1Trajectory_$(T).vtkhdf") + ) + end + + @testset "V4: small point cloud (<= 3 points) Points branch" begin + small = joinpath(dir, "small.h5") + _write_synthetic_traj( + small; T=1, npts=3, dim=2, subgroups=Dict("prediction" => ["vel"]) + ) + out = joinpath(dir, "v4") + visualize(small, out, "pos", "prediction", ["vel"]) + vtk = joinpath( + out, "1Trajectory", "prediction", "prediction_1Trajectory_1.vtkhdf" + ) + @test isfile(vtk) + HDF5.h5open(vtk, "r") do f + @test HDF5.read(f, "VTKHDF/NumberOfPoints") == [3] + @test size(HDF5.read(f, "VTKHDF/Points")) == (3, 3) + end + end + end +end diff --git a/test/test_wall_distance.jl b/test/test_wall_distance.jl new file mode 100644 index 0000000..89be0a9 --- /dev/null +++ b/test/test_wall_distance.jl @@ -0,0 +1,89 @@ +# +# Copyright (c) 2026 Josef Jouaux, Julian Trommer +# Licensed under the MIT license. See LICENSE file in the project root for details. +# +# Regression tests for the wall-distance node feature (`_wall_distance`, src/graph.jl). +# +# Guards the boundary-particle-free / implicit-domain-box case. `_wall_distance` +# depends only on `meta["bounds"]` (+ positions) — never on node types or on the +# presence of boundary particles — so it must be computed for EVERY particle and must +# return `2 * dims` rows (low + high bound per spatial dimension), matching the width +# `calc_norms` reserves (`length(meta["bounds"]) * 2`). +# +# The pre-fix code short-circuited to `ones(dims, n)` whenever `length(mask) == +# size(position, 2)` (a single-type dataset with no boundary particles). That produced +# a constant feature of the WRONG width (`dims` instead of `2*dims`) that disagreed +# with `calc_norms` and threw away the geometric box information the feature exists to +# provide. These tests fail against that old behaviour. + +using Test +using GraphNetSim +using JSON + +@testset "wall_distance" begin + device = identity # keep everything on CPU; `device` only places `meta["bounds"]` + radius = 1.0f0 + + # Three particles, ALL fluid (no boundary particles → mirrors the single-type + # `mask == 1:n` case). p2 and p3 lie outside the unit box, so the feature must + # contain signed distances — something the old constant-`ones` fallback can never + # produce. + # p1 p2 p3 + position = Float32[0.3 -0.5 0.5 # x + 0.7 0.5 1.5] # y + meta = Dict{String,Any}( + "bounds" => [[0.0, 1.0], [0.0, 1.0]], + "default_connectivity_radius" => radius, + ) + dims = length(meta["bounds"]) + + wd = Array(GraphNetSim._wall_distance(position, meta, device)) + + @testset "shape is 2*dims, not dims" begin + @test size(wd, 1) == 2 * dims # regression: old code returned `dims` + @test size(wd, 2) == size(position, 2) + end + + @testset "values are the clamped box distances (not constant ones)" begin + # rows: [low-x; low-y; high-x; high-y]; distances / radius, clamped to [-1, 1]. + expected = Float32[ + 0.3 -0.5 0.5 # low-x = pos_x - lo_x + 0.7 0.5 1.0 # low-y = pos_y - lo_y (1.5 clamped to 1.0) + 0.7 1.0 0.5 # high-x = hi_x - pos_x (1.5 clamped to 1.0) + 0.3 0.5 -0.5 # high-y = hi_y - pos_y + ] + @test wd ≈ expected + @test any(<(0), wd) # signed distances, never constant ones + end + + @testset "missing meta[\"bounds\"] is a hard error" begin + bad_meta = Dict{String,Any}("default_connectivity_radius" => radius) + @test_throws ArgumentError GraphNetSim._wall_distance(position, bad_meta, device) + end +end + +# The wall-distance feature is driven purely by `meta["bounds"]`: a bounded dataset +# reserves `2 * dims` extra input rows; an unbounded one yields a shorter vector with no +# error. `calc_norms` only reads `dataset.meta`, so a lightweight `(; meta=...)` suffices. +@testset "wall_distance gating (calc_norms, bounds-driven)" begin + dev = GraphNetSim.cpu_device() + args = GraphNetSim.Args() + base = JSON.parsefile(joinpath(@__DIR__, "fixtures", "ballistic_small", "meta.json")) + + q_bounded, = GraphNetSim.calc_norms((; meta=deepcopy(base)), dev, args) + + unbounded = deepcopy(base) + delete!(unbounded, "bounds") + q_unbounded, = GraphNetSim.calc_norms((; meta=unbounded), dev, args) + + @test q_bounded == q_unbounded + 2 * length(base["bounds"]) +end + +# The only bounds-related error path: loading a model whose boundary feature disagrees +# with the dataset (bounded model + unbounded dataset, or vice versa). +@testset "bounds/model consistency" begin + @test GraphNetSim._check_bounds_consistency(true, true) === nothing + @test GraphNetSim._check_bounds_consistency(false, false) === nothing + @test_throws ArgumentError GraphNetSim._check_bounds_consistency(true, false) + @test_throws ArgumentError GraphNetSim._check_bounds_consistency(false, true) +end