Rework single multiple shooting

.github/workflows/Format.yml

@@ -2,6 +2,8 @@ name: format-pr
 on:
   schedule:
     - cron: '0 0 * * *'
+  pull_request:    
+    types: [ opened, reopened, synchronize, labeled, unlabeled ]
 jobs:
   build:
     runs-on: ubuntu-latest
@@ -11,21 +13,4 @@ jobs:
       - name: Install JuliaFormatter and format
         run: |
           julia  -e 'import Pkg; Pkg.add("JuliaFormatter")'
-          julia  -e 'import JuliaFormatter; JuliaFormatter.format(".", BlueStyle())'
-
-      # https://github.com/marketplace/actions/create-pull-request
-      # https://github.com/peter-evans/create-pull-request#reference-example
-      - name: Create Pull Request
-        id: cpr
-        uses: peter-evans/create-pull-request@v3
-        with:
-          token: ${{ secrets.GITHUB_TOKEN }}
-          commit-message: Format .jl files
-          title: 'Automatic JuliaFormatter.jl run'
-          branch: auto-juliaformatter-pr
-          delete-branch: true
-          labels: formatting, automated pr, no changelog
-      - name: Check outputs
-        run: |
-          echo "Pull Request Number - ${{ steps.cpr.outputs.pull-request-number }}"
-          echo "Pull Request URL - ${{ steps.cpr.outputs.pull-request-url }}"
+          julia  -e 'import JuliaFormatter; JuliaFormatter.format(".", JuliaFormatter.BlueStyle())'

.gitignore

@@ -31,4 +31,7 @@ plots/
 
 .vscode
 
-storage/
+storage/
+
+example/Ballistics
+example/DamBreak

Project.toml

@@ -21,7 +21,6 @@ KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
-NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
@@ -59,7 +58,6 @@ KernelAbstractions = "0.9.39"
 Lux = "1.4 - 1"
 LuxCUDA = "0.3"
 MLUtils = "0.4.4 - 0.4"
-NearestNeighbors = "0.4.27"
 Optimisers = "0.4, 1"
 OrdinaryDiffEq = "6.85 - 6"
 Plots = "1.40.18"
@@ -79,7 +77,8 @@ julia = "1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Test"]
+test = ["Aqua", "Statistics", "Test"]

example/BallisticSmall/BallisticSmall.jl

@@ -0,0 +1,235 @@
+#
+# Copyright (c) 2026 Josef Kircher
+# Licensed under the MIT license. See LICENSE file in the project root for details.
+#
+# Small ballistic dataset: 10 particles, no boundary nodes, linear drag physics.
+# Two-phase training: DerivativeTraining first, then BatchingStrategy fine-tuning.
+#
+
+using GraphNetSim
+
+import OrdinaryDiffEq: Euler, Tsit5
+import Optimisers: Adam
+
+# Generate dataset into data/ballistic_small if not already present
+include(joinpath(dirname(dirname(@__DIR__)), "test", "generators.jl"))
+let _gen_dir = joinpath(dirname(dirname(@__DIR__)), "data", "ballistic_small")
+    if _needs_generation(_gen_dir)
+        @info "Generating dataset: ballistic_small"
+        _GenBallistic.generate(_gen_dir)
+        @info "  Done."
+    end
+end
+
+######################
+# Network parameters #
+######################
+
+message_steps = 5
+layer_size = 64
+hidden_layers = 2
+batch = 1
+epo = 1
+nder = 40000   # derivative training steps
+ns = 500 + nder  # total steps incl. batching fine-tune
+nms = 500 + ns   # total steps incl. multiple-shooting fine-tune
+norm_steps = 20      # small: only 5 train trajectories × ~67 steps each
+cuda = true
+cp_derivative = 2000
+cp_solver = 100
+cp_ms = 100
+
+########################
+# Node type parameters #
+########################
+
+noise_stddevs = [3.0f-7]
+types_updated = [1]   # only fluid particles (no boundary in this dataset)
+types_noisy = [1]
+
+########################
+# Optimiser parameters #
+########################
+
+learning_rate_start = 1.0f-4
+learning_rate_finish = 1.0f-6
+opt = Adam(learning_rate_start)
+
+#######################
+# Simulation interval #
+#######################
+
+tstart = 0.0f0
+dt = 0.002f0
+tstop = 0.132f0   # (T_LENGTH - 1) * dt = 66 * 0.002
+
+stepsPerTrajectory = 66   # covers the full trajectory in one batch
+
+# timesteps at which MSE is calculated during evaluation
+mse_steps = tstart:dt:tstop
+
+#########################
+# Paths to data folders #
+#########################
+
+ds_path = "data/ballistic_small"
+chk_path = "data/ballistic_small/checkpoints"
+eval_path = "data/ballistic_small/eval"
+
+data_minmax(
+    ds_path;
+    types_updated=types_updated,
+    types_noisy=types_noisy,
+    noise_stddevs=noise_stddevs,
+)
+data_meanstd(
+    ds_path;
+    types_updated=types_updated,
+    types_noisy=types_noisy,
+    noise_stddevs=noise_stddevs,
+)
+
+###########
+# Solvers #
+###########
+
+solver_train = Tsit5()
+solver_eval = Tsit5()
+
+#################
+# Train network #
+#################
+
+# Phase 1: Derivative-based training (fast, no ODE solve per step)
+
+train_network(
+    opt,
+    ds_path,
+    chk_path;
+    mps=message_steps,
+    layer_size=layer_size,
+    hidden_layers=hidden_layers,
+    batchsize=batch,
+    epochs=epo,
+    steps=Int(nder),
+    use_cuda=cuda,
+    checkpoint=cp_derivative,
+    norm_steps=norm_steps,
+    types_updated=types_updated,
+    types_noisy=types_noisy,
+    noise_stddevs=noise_stddevs,
+    training_strategy=DerivativeTraining(),
+    solver_valid=solver_eval,
+    solver_valid_dt=dt,
+    optimizer_learning_rate_start=learning_rate_start,
+    optimizer_learning_rate_stop=learning_rate_finish,
+    show_progress_bars=true,
+    save_step=false,
+)
+
+# Phase 2: BatchingStrategy fine-tuning (ODE-based loss, lower learning rate)
+
+learning_rate_start = 1.0f-6
+learning_rate_finish = nothing
+opt = Adam(learning_rate_start)
+
+train_network(
+    opt,
+    ds_path,
+    chk_path;
+    mps=message_steps,
+    layer_size=layer_size,
+    hidden_layers=hidden_layers,
+    batchsize=batch,
+    epochs=epo,
+    steps=Int(ns),
+    use_cuda=cuda,
+    checkpoint=cp_solver,
+    norm_steps=norm_steps,
+    types_updated=types_updated,
+    types_noisy=types_noisy,
+    noise_stddevs=noise_stddevs,
+    training_strategy=BatchingStrategy(
+        tstart,
+        tstop,          # intervall = full trajectory (66 steps × dt)
+        solver_train,
+        stepsPerTrajectory;
+        loss_function=:mae,
+        adaptive=false,
+    ),
+    solver_valid=solver_eval,
+    optimizer_learning_rate_start=learning_rate_start,
+    optimizer_learning_rate_stop=learning_rate_finish,
+    show_progress_bars=true,
+)
+
+# Phase 3: MultipleShooting fine-tuning (splits trajectory into intervals
+# with continuity penalty — helps escape local minima from SingleShooting)
+
+learning_rate_start = 1.0f-6
+learning_rate_finish = nothing
+opt = Adam(learning_rate_start)
+
+ms_interval_size = 22   # 3 intervals across the 66-step trajectory
+
+train_network(
+    opt,
+    ds_path,
+    chk_path;
+    mps=message_steps,
+    layer_size=layer_size,
+    hidden_layers=hidden_layers,
+    batchsize=batch,
+    epochs=epo,
+    steps=Int(nms),
+    use_cuda=cuda,
+    checkpoint=cp_ms,
+    norm_steps=norm_steps,
+    types_updated=types_updated,
+    types_noisy=types_noisy,
+    noise_stddevs=noise_stddevs,
+    training_strategy=MultipleShooting(
+        tstart, dt, tstop, solver_train, ms_interval_size, 100
+    ),
+    solver_valid=solver_eval,
+    optimizer_learning_rate_start=learning_rate_start,
+    optimizer_learning_rate_stop=learning_rate_finish,
+    show_progress_bars=true,
+)
+
+####################
+# Evaluate network #
+####################
+
+eval_network(
+    ds_path,
+    chk_path,
+    eval_path,
+    solver_eval;
+    start=tstart,
+    stop=tstop,
+    saves=tstart:dt:tstop,
+    dt=dt,
+    mse_steps=collect(mse_steps),
+    mps=message_steps,
+    types_updated=types_updated,
+    layer_size=layer_size,
+    hidden_layers=hidden_layers,
+    use_cuda=cuda,
+)
+
+visualize(
+    eval_path * "/tsit5/trajectories.h5",
+    eval_path * "/vtkhdf/",
+    "pos",
+    "gt",
+    ["vel", "acc"],
+)
+
+visualize(
+    eval_path * "/tsit5/trajectories.h5",
+    eval_path * "/vtkhdf/",
+    "pos",
+    "prediction",
+    ["vel", "acc", "err"],
+)