Dambreak optuna example

.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 }}

.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"

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 <Josef.Kircher@uni-a.de>", "JT <julian.trommer@uni-a.de>"]
 
 [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"

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).

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"

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",
     ],
 )

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.

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.