Speedup rollout

Project.toml

@@ -1,6 +1,6 @@
 name = "GraphNetSim"
 uuid = "5ff66f56-808c-48e7-ac84-dd29877231f8"
-version = "0.1.1"
+version = "0.1.2"
 authors = ["Josef Jouaux <Josef.Kircher@uni-a.de>", "JT <julian.trommer@uni-a.de>"]
 
 [deps]
@@ -17,12 +17,12 @@ GraphNetCore = "7809f980-de1b-4f9a-8451-85f041491431"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
-JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LightXML = "9c8b4983-aa76-5018-a973-4c85ecc9e179"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+Octopus = "49bb3e92-56cb-4d34-aed6-8ee8ef58b458"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
@@ -61,6 +61,7 @@ LightXML = "0.9"
 Lux = "1.13 - 1"
 LuxCUDA = "0.3"
 MLUtils = "0.4.4 - 0.4"
+Octopus = "0.2"
 Optimisers = "0.4, 1"
 OrdinaryDiffEq = "6.85 - 6"
 Plots = "1.40.18"
@@ -82,7 +83,8 @@ julia = "1.11"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Test"]
+test = ["Aqua", "JuliaFormatter", "Test"]

src/GraphNetSim.jl

@@ -60,8 +60,8 @@ include("rollout_history.jl")
 include("dataset.jl")
 include("visualize.jl")
 include("config.jl")
-include("../convert_csv/csvToh5.jl")
-include("../convert_csv/vtkToh5.jl")
+include("import_data/csvToh5.jl")
+include("import_data/vtkToh5.jl")
 
 export SingleShooting, MultipleShooting, DerivativeTraining, BatchingStrategy
 
@@ -144,6 +144,7 @@ Configuration structure for training and evaluating Graph Neural Network simulat
     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
@@ -516,12 +517,14 @@ function train_network(opt, ds_path, cp_path; kws...)
     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)
 
@@ -991,6 +994,7 @@ function eval_network(
     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)
 
@@ -1104,7 +1108,13 @@ function eval_network!(
 
     test_loader = DataLoader(ds_test; batchsize=-1, buffer=false, parallel=true)
 
+    # Optional cap on the number of evaluated trajectories (default: all). Lets bounded eval and
+    # A/B timing runs stay tractable — a full rollout over every test trajectory is very slow.
+    # NB: bare `parse` resolves to `JSON.parse` in this module (see dataset.jl), so qualify Base.parse.
+    n_eval_traj = Base.parse(Int, get(ENV, "GNS_EVAL_NTRAJ", string(typemax(Int))))
+
     for (ti, data) in enumerate(test_loader)
+        ti > n_eval_traj && break
         target_features = ds_test.meta["solver_target_features"]
         output_features = ds_test.meta["output_features"]
         println("Rollout trajectory $ti...")
@@ -1330,6 +1340,7 @@ function extrapolate_network(
     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)
 

src/graph.jl

@@ -26,6 +26,7 @@ using CUDA
 import Statistics: norm
 using JLD2
 using PointNeighbors
+using Octopus: Octopus
 using ChainRulesCore
 
 """
@@ -65,15 +66,7 @@ function build_graph(
         data["velocity"][:, :, datapoint]
     end
 
-    build_graph(
-        gns,
-        current_position,
-        velocity,
-        meta,
-        node_type,
-        data["mask"],
-        device,
-    )
+    build_graph(gns, current_position, velocity, meta, node_type, data["mask"], device)
 end
 
 """
@@ -100,8 +93,10 @@ All features are normalized using the normalizers stored in the model.
 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 = point_neighbor_ns(
-        position, Float32(meta["default_connectivity_radius"])
+    senders, receivers, rel_displacement, rel_dist_norm = neighbor_search(
+        device(position),
+        Float32(meta["default_connectivity_radius"]),
+        get(meta, "neighbor_backend", :pointneighbors),
     )
 
     multi_type = n_node_types(meta) > 1
@@ -112,16 +107,23 @@ function build_graph(
 
     edge_features = device(vcat(rel_displacement, rel_dist_norm) .+ 1.0f-8)
 
-    nf = vel_norm
-    if use_position
-        nf = vcat(gns.n_norm["position"](position), nf)
-    end
-    if use_wall
-        nf = vcat(nf, _wall_distance(position, mask, meta, device))
-    end
-    if multi_type
-        nf = vcat(nf, node_type)
-    end
+    # Node features are the vertical concatenation of, in order:
+    #   [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
+    # 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),) : ())...,
+        (multi_type ? (node_type,) : ())...,
+    )
+    nf = length(nf_blocks) == 1 ? nf_blocks[1] : vcat(nf_blocks...)
     node_features = device(nf)
 
     return GraphNetCore.FeatureGraph(
@@ -160,8 +162,7 @@ function _wall_distance(position, mask, meta, device)
     dist_low_bound = position .- boundaries[:, 1]
     dist_up_bound = boundaries[:, 2] .- position
     return clamp.(
-        vcat(dist_low_bound, dist_up_bound) ./
-        Float32(meta["default_connectivity_radius"]),
+        vcat(dist_low_bound, dist_up_bound) ./ Float32(meta["default_connectivity_radius"]),
         -1.0f0,
         1.0f0,
     )
@@ -340,20 +341,26 @@ Returns normalized relative displacements and distances.
 - Supports arbitrary dimension (2D, 3D, etc.).
 """
 function point_neighbor_ns(pos::CuArray, radius::Float32)
-    system = pos#[:,mask]
+    system = pos
     min_corner = minimum(pos; dims=2)
     max_corner = maximum(pos; dims=2)
     nhs = GridNeighborhoodSearch{size(pos, 1)}(;
         search_radius=radius,
         n_points=size(pos, 2),
         cell_list=FullGridCellList(; min_corner, max_corner, search_radius=radius),
+        update_strategy=ParallelUpdate(),
     )
-    initialize!(nhs, Array(system), Array(pos))
-    backend = CUDABackend()
-    # Simple example: just count the neighbors of each particle
-    n_neighbors_gpu = CuArray(zeros(Int, size(pos, 2)))
-    nhs_gpu = adapt(backend, nhs)
+    # Build the cell list on-device: adapt the (empty) nhs to the GPU, then `initialize!` with the
+    # CuArray positions so PointNeighbors dispatches to the parallel atomic init
+    # (`default_backend(::CuArray)` => GPU; `ParallelUpdate`'s `initialize_grid!` is the parallel one).
+    # Avoids the GPU->CPU->GPU round trip of the old `initialize!(nhs, Array(...), Array(...))` +
+    # adapt-back, whose serial CPU cell-list build was ~16 ms at 33k particles vs ~2 ms here (~8x).
+    # Edge set is identical; edge order may differ (atomic push), perturbing the downstream scatter
+    # by ~eps only. Benchmarked in example/RuntimeBenchmark/.
+    nhs_gpu = adapt(CUDABackend(), nhs)
+    initialize!(nhs_gpu, pos, pos)
 
+    n_neighbors_gpu = CuArray(zeros(Int, size(pos, 2)))
     foreach_point_neighbor(system, pos, nhs_gpu) do i, _, _, _
         n_neighbors_gpu[i] += 1
     end
@@ -436,6 +443,80 @@ function point_neighbor_ns(pos::Array, radius::Float32)
     return senders, receivers, rel_displacement, rel_dist_norm
 end
 
+"""
+    octopus_ns(pos, radius::Float32)
+
+Neighbor search backend using [Octopus.jl](https://github.com/una-auxme/Octopus.jl)'s fast octree
+(`TNS`). Dispatches on the array type internally (CPU `Array` or `CuArray`), so the same code serves
+both devices. Returns `(senders, receivers, rel_displacement, rel_dist_norm)` in exactly the format
+of [`point_neighbor_ns`](@ref) — same receiver/sender convention, displacement sign, radius
+normalization, and appended self-edges — so the two backends are interchangeable in [`build_graph`](@ref).
+Gradients flow through `pos` via Octopus's differentiable `build_edges_diff` rrule; the tree build is
+non-differentiable (`@ignore_derivatives`). The octree is O(N) and ~18x leaner than PointNeighbors'
+dense grid on the GPU, so it's the backend for large point clouds where the grid OOMs.
+"""
+function octopus_ns(pos, radius::Float32)
+    D = size(pos, 1)
+    n = size(pos, 2)
+    tns = ChainRulesCore.@ignore_derivatives begin
+        t = Octopus.TNS(eltype(pos); ndims=D)
+        Octopus.set_search_radius!(t, radius)
+        pid = Octopus.add_point_set!(t, pos)
+        Octopus.set_active_search!(t, pid, pid)
+        Octopus.run!(t)
+        t
+    end
+    e = Octopus.build_edges_diff(pos, tns, 1, radius)
+
+    self_s, self_r, self_disp, self_dist = ChainRulesCore.@ignore_derivatives begin
+        s = similar(e.senders, n)
+        copyto!(s, Int32.(1:n))
+        (
+            s,
+            copy(s),
+            fill!(similar(e.rel_displacement, D, n), zero(eltype(e.rel_displacement))),
+            fill!(similar(e.rel_dist_norm, 1, n), zero(eltype(e.rel_dist_norm))),
+        )
+    end
+
+    senders = vcat(e.senders, self_s)
+    receivers = vcat(e.receivers, self_r)
+    rel_displacement = hcat(e.rel_displacement, self_disp)
+    rel_dist_norm = hcat(e.rel_dist_norm, self_dist)
+    return senders, receivers, rel_displacement, rel_dist_norm
+end
+
+"""
+    neighbor_search(pos, radius::Float32, backend::Symbol)
+
+Select the neighborhood-search implementation for graph construction. `backend` comes from
+`Args.neighbor_backend` (threaded via `meta["neighbor_backend"]`):
+
+- `:pointneighbors` — [`point_neighbor_ns`](@ref), PointNeighbors.jl grid search (default).
+- `:octopus`        — [`octopus_ns`](@ref), Octopus.jl octree search (memory-lean on GPU; needed for
+                      large clouds where the dense grid OOMs).
+- `:auto`           — PointNeighbors on CPU (`Array`), Octopus on GPU (`CuArray`).
+
+All branches return the identical `(senders, receivers, rel_displacement, rel_dist_norm)` format and
+are differentiable in `pos`.
+"""
+function neighbor_search(pos, radius::Float32, backend::Symbol)
+    if backend === :pointneighbors
+        return point_neighbor_ns(pos, radius)
+    elseif backend === :octopus
+        return octopus_ns(pos, radius)
+    elseif backend === :auto
+        return pos isa CuArray ? octopus_ns(pos, radius) : point_neighbor_ns(pos, radius)
+    else
+        throw(
+            ArgumentError(
+                "unknown neighbor_backend $(repr(backend)); " *
+                "use :pointneighbors, :octopus, or :auto",
+            ),
+        )
+    end
+end
+
 """
     ChainRulesCore.rrule(::typeof(point_neighbor_ns), pos::Array, radius::Float32)
 

src/rollout_history.jl

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2026 Josef Kircher, Julian Trommer
+# Copyright (c) 2026 Josef Jouaux, Julian Trommer
 # Licensed under the MIT license. See LICENSE file in the project root for details.
 #
 

src/solve.jl

@@ -348,21 +348,30 @@ function ode_step_eval(
     # state equals the input state — no write-back needed (previously `gns.st = st`).
     output, _ = gns.train_state.model(graph, ps, gns.train_state.states)
     indices = [meta["features"][tf]["dim"] for tf in target_fields]
-    buf = Zygote.Buffer(output)
-    for i in 1:length(output_fields)
-        buf[(sum(indices[1:(i - 1)]) + 1):sum(indices[1:i]), :] = inverse_data(
-            gns.o_norm[output_fields[i]],
-            output[(sum(indices[1:(i - 1)]) + 1):sum(indices[1:i]), :],
-        )
-    end
+
+    # Inference-only path: `rollout` solves an `ODEProblem{false}` with no sensealg, so this RHS is
+    # never differentiated. Denormalize the output slices directly instead of routing them through a
+    # `Zygote.Buffer` + `copy` (as the differentiated training `ode_step` must) — that AD bookkeeping
+    # is dead weight in a hot loop run ~`trajectory_length` times per trajectory. Values are identical
+    # to the buffer version; for a single output field (the common case) `reduce(vcat, [slice])`
+    # returns the slice untouched, so no concat/copy happens at all.
+    dx = reduce(
+        vcat,
+        [
+            inverse_data(
+                gns.o_norm[output_fields[i]],
+                output[(sum(indices[1:(i - 1)]) + 1):sum(indices[1:i]), :],
+            ) for i in 1:length(output_fields)
+        ],
+    )
 
     @ignore_derivatives begin
         if !isnothing(pr)
             next!(pr, showvalues=[(:t, "$(t)")])
         end
     end
 
-    return device(ComponentArray(; x=x.dx, dx=copy(buf) .* val_mask)) # TODO check why output is used here directly
+    return device(ComponentArray(; x=x.dx, dx=dx .* val_mask))
 end
 
 """

src/utils.jl

@@ -2,6 +2,23 @@
 # Copyright (c) 2026 Josef Jouaux, Julian Trommer
 # Licensed under the MIT license. See LICENSE file in the project root for details.
 #
+# This file contains work derived from DeepMind's "learning_to_simulate"
+# (https://github.com/google-deepmind/deepmind-research), modified from the original:
+#
+#   Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+#
 
 import Printf: @sprintf
 import JSON: print as json_print, parsefile as json_parsefile

test/runtests.jl

@@ -16,12 +16,7 @@ include("generate_fixtures.jl")
         # stale-deps check:
         #  - GPUCompiler: deps-only version pin (see Project.toml [compat]), loaded
         #    transitively via CUDA/Reactant.
-        #  - JuliaFormatter: dev/CI tool invoked as `using JuliaFormatter; format(".")`.
-        Aqua.test_all(
-            GraphNetSim;
-            ambiguities=false,
-            stale_deps=(ignore=[:GPUCompiler, :JuliaFormatter],),
-        )
+        Aqua.test_all(GraphNetSim; ambiguities=false, stale_deps=(ignore=[:GPUCompiler],))
     end
 
     include("test_converters.jl")