Restructure meta

.github/workflows/Format.yml

@@ -1,14 +1,31 @@
-name: Format suggestions
+name: format-pr
 on:
-  pull_request:
-    types: [ opened, reopened, synchronize, labeled, unlabeled ]
+  schedule:
+    - cron: '0 0 * * *'
 jobs:
-  code-style:
+  build:
     runs-on: ubuntu-latest
     steps:
-      - name: "Check out repository"
-        uses: actions/checkout@v6
-      - uses: julia-actions/julia-format@v4
+      - uses: actions/checkout@v4
+      - uses: julia-actions/cache@v2
+      - 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:
-          version: '1'
-          suggestion-label: 'format-suggest'
+          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 }}"

src/GraphNetSim.jl

@@ -37,7 +37,7 @@ include("../convert_csv/csvToh5.jl")
 
 export SingleShooting, MultipleShooting, DerivativeTraining, BatchingStrategy
 
-export train_network, eval_network, data_minmax, data_meanstd
+export train_network, eval_network, data_minmax, data_meanstd, update_meta!
 export init_train_step, train_step, validation_step, batchTrajectory
 # export prepare_training, get_delta
 export visualize
@@ -77,6 +77,11 @@ Configuration structure for training and evaluating Graph Neural Network simulat
 - `optimizer_learning_rate_start::Float32=1.0f-4`: Initial learning rate
 - `optimizer_learning_rate_stop::Union{Nothing,Float32}=nothing`: Final learning rate (for decay schedule)
 
+### Normalization
+- `norm_type::Symbol=:online`: Normalization strategy for Float32 features.
+  `:online` (accumulate stats during training), `:minmax` (requires data_min/data_max in meta.json),
+  `:meanstd` (requires data_mean/data_std in meta.json).
+
 ### Validation
 - `show_progress_bars::Bool=true`: Show training progress bars
 - `use_valid::Bool=true`: Load validation checkpoint (best loss) instead of final checkpoint
@@ -107,6 +112,7 @@ Configuration structure for training and evaluating Graph Neural Network simulat
     reset_valid::Bool = false
     optimizer_learning_rate_start::Float32 = 1.0f-4
     optimizer_learning_rate_stop::Union{Nothing,Float32} = nothing
+    norm_type::Symbol = :online
     save_step::Bool = false
     on_grad::Union{Nothing,Function} = nothing
     on_valid::Union{Nothing,Function} = nothing
@@ -207,8 +213,27 @@ function calc_norms(dataset, device, args)
                 )
             end
         else
-            if haskey(dataset.meta["features"][feature], "data_min") &&
-                haskey(dataset.meta["features"][feature], "data_max")
+            if args.norm_type == :online
+                if feature in input_features
+                    n_norms[feature] = NormaliserOnline(
+                        feature_dim, device; max_acc=Float32(args.norm_steps)
+                    )
+                elseif feature in output_features
+                    o_norms[feature] = NormaliserOnline(
+                        feature_dim, device; max_acc=Float32(args.norm_steps)
+                    )
+                end
+            elseif args.norm_type == :minmax
+                if !haskey(dataset.meta["features"][feature], "data_min") ||
+                    !haskey(dataset.meta["features"][feature], "data_max")
+                    throw(
+                        ArgumentError(
+                            "norm_type=:minmax requires 'data_min' and 'data_max' in " *
+                            "meta.json for feature \"$feature\". Run " *
+                            "`update_meta!(path, :minmax)` to compute and write these statistics.",
+                        ),
+                    )
+                end
                 if haskey(dataset.meta["features"][feature], "target_min") &&
                     haskey(dataset.meta["features"][feature], "target_max")
                     if feature in input_features
@@ -253,8 +278,17 @@ function calc_norms(dataset, device, args)
                         end
                     end
                 end
-            elseif haskey(dataset.meta["features"][feature], "data_mean") &&
-                haskey(dataset.meta["features"][feature], "data_std")
+            elseif args.norm_type == :meanstd
+                if !haskey(dataset.meta["features"][feature], "data_mean") ||
+                    !haskey(dataset.meta["features"][feature], "data_std")
+                    throw(
+                        ArgumentError(
+                            "norm_type=:meanstd requires 'data_mean' and 'data_std' in " *
+                            "meta.json for feature \"$feature\". Run " *
+                            "`update_meta!(path, :meanstd)` to compute and write these statistics.",
+                        ),
+                    )
+                end
                 if feature in input_features
                     n_norms[feature] = NormaliserOfflineMeanStd(
                         Float32.(dataset.meta["features"][feature]["data_mean"]),
@@ -269,15 +303,12 @@ function calc_norms(dataset, device, args)
                     )
                 end
             else
-                if feature in input_features
-                    n_norms[feature] = NormaliserOnline(
-                        feature_dim, device; max_acc=Float32(args.norm_steps)
-                    )
-                elseif feature in output_features
-                    o_norms[feature] = NormaliserOnline(
-                        feature_dim, device; max_acc=Float32(args.norm_steps)
-                    )
-                end
+                throw(
+                    ArgumentError(
+                        "Invalid norm_type=:$(args.norm_type). " *
+                        "Must be one of :online, :minmax, :meanstd.",
+                    ),
+                )
             end
         end
     end
@@ -353,6 +384,7 @@ function train_network(opt, ds_path, cp_path; kws...)
                 mps=existing_cfg.mps,
                 layer_size=existing_cfg.layer_size,
                 hidden_layers=existing_cfg.hidden_layers,
+                norm_type=existing_cfg.norm_type,
             ),
             NamedTuple(kws),
         )
@@ -369,6 +401,7 @@ function train_network(opt, ds_path, cp_path; kws...)
             types_updated=args.types_updated,
             types_noisy=args.types_noisy,
             noise_stddevs=args.noise_stddevs,
+            norm_type=args.norm_type,
         ),
         cp_path,
     )
@@ -707,7 +740,7 @@ function train_gns!(
                 end
 
                 if valid_error / ds_valid.meta["n_trajectories"] < min_validation_loss
-                    push!(df_valid, [step, valid_error / ds_valid.meta["n_trajectories"]])
+                    # push!(df_valid, [step, valid_error / ds_valid.meta["n_trajectories"]])
                     save!(
                         gns,
                         opt_state,

src/config.jl

@@ -26,6 +26,7 @@ as documentation and may legitimately differ between training phases.
 - `types_updated`: Node types whose outputs are predicted.
 - `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`).
 """
 @kwdef struct ModelConfig
     format_version::Int = 1
@@ -36,6 +37,7 @@ as documentation and may legitimately differ between training phases.
     types_updated::Vector{Int}
     types_noisy::Vector{Int}
     noise_stddevs::Vector{Float32}
+    norm_type::Symbol = :online
 end
 
 """
@@ -89,6 +91,7 @@ function save_model_config(cfg::ModelConfig, cp_path::String)
                     "types_updated" => cfg.types_updated,
                     "types_noisy" => cfg.types_noisy,
                     "noise_stddevs" => cfg.noise_stddevs,
+                    "norm_type" => String(cfg.norm_type),
                 ),
             ),
             2,
@@ -121,6 +124,7 @@ function load_model_config(cp_path::String)::Union{ModelConfig,Nothing}
             types_updated=Int.(train["types_updated"]),
             types_noisy=Int.(train["types_noisy"]),
             noise_stddevs=Float32.(train["noise_stddevs"]),
+            norm_type=Symbol(get(train, "norm_type", "online")),
         )
     catch e
         @warn "Could not parse model config at \"$path\": $e. Falling back to supplied arguments."

src/dataset.jl

@@ -209,7 +209,8 @@ function MLUtils.getobs!(buffer, ds::Dataset, idx)
 
     prepare_trajectory!(buffer, ds.meta, ds.meta["device"])
     n_particles = size(buffer["node_type"], 2)
-    if n_node_types(ds.meta) == 1
+    single_type = ds.meta["features"]["node_type"]["data_min"]
+    if n_node_types(ds.meta) == 1 && single_type in ds.meta["types_updated"]
         buffer["mask"] = ds.meta["device"](Int32.(1:n_particles))
         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))

src/utils.jl

@@ -4,6 +4,23 @@
 #
 
 import Printf: @sprintf
+import JSON: print as json_print, parsefile as json_parsefile
+
+"""
+    n_node_types(meta)
+
+Returns the number of distinct node types in the dataset.
+
+## Arguments
+- `meta`: Feature metadata dictionary containing `node_type` feature spec.
+
+## Returns
+- `Int`: `data_max - data_min + 1` for the `node_type` feature.
+"""
+function n_node_types(meta)
+    meta["features"]["node_type"]["data_max"] - meta["features"]["node_type"]["data_min"] +
+    1
+end
 
 """
     n_node_types(meta)
@@ -133,21 +150,9 @@ function data_minmax(path)
                 end
             end
         end
-        for tf in target_features #TODO as target features are part of the data no FD has to be computed
+        for tf in target_features
             if !haskey(ds_train.meta["features"][tf], "onehot") &&
                 isnumber(ds_train.meta, tf)
-                # ddiff = data[tf][:, :, 2:end] - data[tf][:, :, 1:(end - 1)]
-                # if size(data["dt"]) == 1
-                #     dts = data["dt"]
-                # else
-                #     dts = Float32.(data["dt"][2:end] - data["dt"][1:(end - 1)])
-                # end
-                # for i in eachindex(dts)
-                #     ddiff[:, :, i] ./= dts[i]
-                # end
-                # ddiff_min = minimum(ddiff)
-                # ddiff_max = maximum(ddiff)
-
                 ddiff_min = minimum(data[tf])
                 ddiff_max = maximum(data[tf])
                 if ddiff_min < result["target|$tf"][1]
@@ -381,3 +386,99 @@ function clear_log(lines::Integer, move_up=true)
         clear_line()
     end
 end
+
+"""
+    update_meta!(path, norm_type)
+
+Compute normalization statistics and write them into the dataset's `meta.json`.
+
+Calls [`data_minmax`](@ref) or [`data_meanstd`](@ref) depending on `norm_type`,
+then updates the per-feature entries in `meta.json` with the computed statistics.
+Conflicting statistics from a different normalization type are removed.
+
+## Arguments
+- `path::String`: Dataset directory containing `meta.json`, `train.h5`, `valid.h5`, `test.h5`.
+- `norm_type::Symbol`: One of `:online`, `:minmax`, or `:meanstd`.
+  `:online` is a no-op (no precomputed statistics needed).
+
+## Returns
+- `String`: Path to the written `meta.json` file.
+"""
+function update_meta!(path::String, norm_type::Symbol)
+    if norm_type ∉ (:online, :minmax, :meanstd)
+        throw(
+            ArgumentError(
+                "Invalid norm_type=:$norm_type. Must be one of :online, :minmax, :meanstd."
+            ),
+        )
+    end
+
+    meta_path = joinpath(path, "meta.json")
+    if !isfile(meta_path)
+        throw(ArgumentError("meta.json not found at \"$meta_path\"."))
+    end
+
+    if norm_type == :online
+        @info "norm_type=:online requires no precomputed statistics. meta.json unchanged."
+        return meta_path
+    end
+
+    meta = json_parsefile(meta_path)
+
+    conflicting_minmax = ("data_min", "data_max", "output_min", "output_max")
+    conflicting_meanstd = ("data_mean", "data_std")
+
+    if norm_type == :minmax
+        stats = data_minmax(path)
+
+        for (key, val) in stats
+            if startswith(key, "target|")
+                feat = key[8:end]
+                if !haskey(meta["features"], feat)
+                    continue
+                end
+                # Write output-specific keys only if they differ from input stats
+                if haskey(stats, feat) && val != stats[feat]
+                    meta["features"][feat]["output_min"] = Float64(val[1])
+                    meta["features"][feat]["output_max"] = Float64(val[2])
+                end
+            else
+                if !haskey(meta["features"], key)
+                    continue
+                end
+                meta["features"][key]["data_min"] = Float64(val[1])
+                meta["features"][key]["data_max"] = Float64(val[2])
+                # Remove conflicting meanstd keys
+                for ck in conflicting_meanstd
+                    delete!(meta["features"][key], ck)
+                end
+            end
+        end
+
+    elseif norm_type == :meanstd
+        stats = data_meanstd(path)
+
+        for (key, val) in stats
+            if startswith(key, "target|")
+                # meanstd has no separate output keys; skip target entries
+                continue
+            end
+            if !haskey(meta["features"], key)
+                continue
+            end
+            meta["features"][key]["data_mean"] = Float64.(val[1])
+            meta["features"][key]["data_std"] = Float64.(val[2])
+            # Remove conflicting minmax keys
+            for ck in conflicting_minmax
+                delete!(meta["features"][key], ck)
+            end
+        end
+    end
+
+    open(meta_path, "w") do f
+        json_print(f, meta, 2)
+    end
+
+    @info "Updated meta.json with $norm_type statistics at \"$meta_path\"."
+    return meta_path
+end