Parametric Generalized Hamiltonian Neural Networks

[benedict-96]

Implements parametric generalized Hamiltonian neural networks (PGHNNs): architectures whose forward pass takes the parameters of the system alongside the state, so one network covers a parameter range rather than a single problem instance.

Rebased onto main (0.6.0). The branch's 53 commits are collapsed into one; see the commit message for the detail. In short:

• cc1a786f was 123 MB of build output — docs/build.zip (77 MB), docs/build 2.zip (47 MB), eleven LaTeX artefacts, two go_migration_inspection_*.txt — and about four lines of source. All of it except the inspection logs is already in .gitignore.
• 8d243592 was main's own HDF5 migration re-applied by hand (byte-identical blobs), redundant after a rebase.
• 43d89206 merged main only as far as d07b4c26, so the branch never saw the GeometricOptimizers separation, and its resolution left src/GeometricMachineLearning.jl half from each side — go_bridges.jl, src/arrays/ and src/manifolds/ included again next to main's import GeometricOptimizers. That is the Method overwriting is not permitted during Module precompilation failure all 13 CI jobs hit.
• 4c0e2684 is kept: a real seven-line fix to the symbolic Jacobian broadcast and two build_nn_function calls.

The repeated bad merges have one cause: main reformatted src/GeometricMachineLearning.jl from a 4-space-indented module body to column 0 and the branch never did, so every merge conflicted on the whole file. Here main's version is taken as-is and only the branch's includes and exports re-applied. Nothing that moved to GeometricOptimizers comes back — git diff --name-only main...HEAD -- src/arrays src/manifolds src/optimizers/go_bridges.jl test/manifolds test/optimizers is empty.

What it adds

GeneralizedHamiltonianArchitecture is implemented — it used to be a stub whose constructor threw error("GHNN still has to be implemented!"). It composes n_integrators symplectic Euler steps, each differentiating a learned kinetic or potential energy. Around it: ForcedGeneralizedHamiltonianArchitecture and ForcedSympNet with ForcingLayers, ParametricDataLoader, ParametricLoss, a SymbolicPullback for the parametric case, ParametricResNet as the non-structure-preserving baseline, and QPT2/QPTOAT2.

Flattening the system parameters

The system parameters are flattened into the network input. The branch used ParameterHandling and pirated three flatten methods on it. That could not work: GeometricOptimizers defines ParameterHandling.flatten(x) with an unbound type parameter, and that method wins — UndefVarError: T not defined in static parameter matching, which is D6 in NeuralNetworkParameters' PLAN.md, hit in practice.

flatten/unflatten from NeuralNetworkParameters replace it, and because AbstractNeuralNetworks exports params this needs no method on a foreign type; the layout a layer stores is a value rather than a closure. NeuralNetworkParameters is registered as of 0.1.1 and main depends on it directly, so this needs no dependency main does not already have — Project.toml on this branch is byte-identical to main's.

Five bugs, all on paths nothing ran

test/training_phnn.jl is not in runtests.jl, so the loop the PR is for had never executed:

1. concatenate_array_with_parameters(::AbstractMatrix, ::AbstractVector) used vcat where it needs hcat, collapsing a batch into one long vector.
2. optimize_for_one_epoch! called _unpack_tuple, which has never existed in this package.
3. Zygote differentiates through the NeuralNetworkParameters struct, so every nesting level of the gradient comes back wrapped in (params = …,). These architectures nest — a SymplecticEuler layer holds a whole sub-network — and _get_params only unwraps the top, so _unwrap_gradient recurses.

Codecov then showed six of the new files at 0% — nothing constructed a forcing layer, a wide or parametric ResNet layer, a ParametricResNet, a ForcedSympNet or a ForcedGeneralizedHamiltonianArchitecture. Writing a construct-and-evaluate test for each turned up two more:

4. ForcedGeneralizedHamiltonianArchitecture could not be evaluated at all. The parameter-dependent NeuralNetwork functor and the Optimizer entry point were written for GeneralizedHamiltonianArchitecture, and the forced architecture is a sibling of it under HamiltonianArchitecture, not a subtype — so nn(x, μ) fell through to AbstractNeuralNetworks' generic functor, which read the system parameters as the network parameters and reached the first layer as a Float64.
5. ParametricResNet(::DataLoader, n_blocks, width; parameters = …) accepted parameters and did not forward it, always building a network with no parameter dependence.

New pghnn_training_test.jl and parametric_layers_and_architectures_test.jl cover those paths. The latter also pins the ForcingLayer convention, which is the opposite of what the names suggest: Q/P/QP say what the forcing depends on, not what it changes — all three add to p and leave q alone, which is what a force does to the ṗ equation. The two existing tests are rewritten: the data-loader one pinned which shuffled batch holds which parameters, which depends on the Julia version's RNG stream, and now asserts the correspondence itself; the pullback one had no @test at all and re-defined in the test file what src/ now provides, and now checks the symbolic gradient against the summed per-sample Zygote gradients (they agree to 4e-16).

Adapting to the current dependencies

The branch predates SymbolicNeuralNetworks 0.5. symbolize! is gone (symbolic_variables replaces it) and symbolic variables are scalar Nums rather than Symbolics.Arrs (SNN#14), so the parametric SymbolicPullback is rewritten in the shape of SymbolicNeuralNetworks.SymbolicPullback, using symbolic_derivative and ParameterGradient instead of a hand-rolled closure and semi_flatten_network_parameters. Five Symbolics.Arr special cases — three of them type piracy — are unreachable now and deleted. GeometricProblems.default_parameters is a function since 0.8, so the tests call it.

src/optimizers/optimizer.jl keeps main's version: the branch's rgrad methods for NeuralNetworkParameters and nothing gradients are all covered by the 0.5.0 rewrite, whose _tree_optim_step! already skips a nothing block.

The three items from the original description

• Type piracy. Six of the nine methods are gone. Three were fixed upstream and are simply deleted: applychain (AbstractNeuralNetworks 0.7 leaves its data argument untyped for this reason), networkbackend(::LazyArrays.ApplyArray) (LazyArrays is not a GML dependency any more) and h5save(::HDF5.Group, ::NetworkParameters, …) (NeuralNetworkParameters' extension has the ::H5DataStore method — and GML's, being more specific, was shadowing it and losing the kind attribute). The four three-argument Chain functors became apply_parametric, a function GML owns, which the applychain change made possible. What is left — SymbolicNeuralNetworks.Jacobian(f, nn, dim2) and the four three-tuple SymbolicPullback operators — is genuinely upstream, marked with a TODO, and collected in Type piracy introduced by the parametric Hamiltonian architectures (#207) #243.
• QPT2. Introduced and documented, and QPT widened. Replacing QPT with it is still open.
• Parameter-dependent DataLoader. Done, as ParametricDataLoader.

Also

• GeometricProblems and Printf dropped from [deps], where the branch had duplicated them out of [extras]; nothing in src/ uses either.
• Breaking: SymplecticEuler, SymplecticEulerA and SymplecticEulerB are no longer exported. Those names are the layer types of the generalized architectures now; the training methods that used to carry them are SymplecticEulerIntegrator, SymplecticEulerIntegratorA and SymplecticEulerIntegratorB. SEuler, SEulerA and SEulerB are unchanged.
• The training loop takes Union{DataLoader, ParametricDataLoader} rather than an untyped argument.

Dependencies

Follows main onto AbstractNeuralNetworks 0.7, SymbolicNeuralNetworks 0.6 and Zygote 0.7. The parameter container is NeuralNetworkParameters.NetworkParameters now — ANN 0.7 removed the old name outright rather than aliasing it — so every type position is renamed.

On #245. SymbolicNeuralNetworks 0.6 shipped the layerwise pullback (SNN #49), which its changelog says fixes this case. Measured, it does not reach it yet: SNN's seam is a plain Vector{Num}, so a SymplecticEuler built with return_parameters = true — which passes the system parameters on and returns a Tuple — cannot be seeded, and layerwise = :auto raises rather than declining, which is SNN #54. The n_integrators > 1 guard therefore stays; the findings are in #245.

Verification

Pkg.test() passes (exit 0, no failures) and docs/make.jl builds clean, both against a deved NeuralNetworkParameters. The four new testsets:

testset	assertions
Generalized Hamiltonian Neural Network	2
Symbolic pullback for a single-layer PGHNN	8
PGHNN training on a ParametricDataLoader	4
Parametric and forced layers and architectures	61
Parametric DataLoader	2

The previous tip of this branch was 43d89206, still reachable by SHA if anything needs recovering.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 88.01843% with 52 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.46%. Comparing base (6f735e9) to head (6db0b98).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...tectures/generalized_hamiltonian_neural_network.jl	86.84%	15 Missing ⚠️
src/pullbacks/symbolic_hnn_pullback.jl	74.19%	8 Missing ⚠️
src/layers/parametric_resnet_layer.jl	80.00%	7 Missing ⚠️
src/training_method/symplectic_euler.jl	16.66%	5 Missing ⚠️
...s/forced_generalized_hamiltonian_neural_network.jl	85.00%	3 Missing ⚠️
src/layers/forcing_dissipation_layers.jl	94.44%	3 Missing ⚠️
src/architectures/forced_sympnet.jl	84.61%	2 Missing ⚠️
src/architectures/parametric_resnet.jl	87.50%	2 Missing ⚠️
src/architectures/resnet.jl	60.00%	2 Missing ⚠️
src/data_loader/parametric_data_loader.jl	96.66%	2 Missing ⚠️
... and 2 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #207      +/-   ##
==========================================
+ Coverage   66.51%   69.46%   +2.95%     
==========================================
  Files          99      108       +9     
  Lines        2888     3285     +397     
==========================================
+ Hits         1921     2282     +361     
- Misses        967     1003      +36     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[benedict-96]

Commented code with the thought of handing this review over to claude in the first iteration.
Something that I haven't mentioned yet is missing documentation on the parametric forced generalized Hamiltonian neural network, the forced SympNet and the WideResNet. For making this documentation more visually appealing I could provide the tikz pictures that are part of the paper pertaining to this pr.

[benedict-96]

This functionality should most likely be in GeometricProblems or GMLDataSets.

[benedict-96]

The data set should be generated with GeometricProblems (I would think). The problem may be that there is no good way at the moment to obtain analytic solutions (which is what happens here).

[benedict-96]

This should be in GMLDataSets (I think).

[benedict-96]

Again, probably better in GMLDataSets.

[benedict-96]

This seems to be a copy of the code that I commented already above. We should try to avoid duplicates like this (and probably outsource this functionality to GMLDataSets).

[benedict-96]

Is this still needed (potentially). If it is obsolete we can just delete it.

[benedict-96]

Confer comment above.

[benedict-96]

Is this obsolete? Confer comment above.

[benedict-96]

With the SymbolicPullback there seem to be massive performance issues that have to be investigated. Also: should parts of this be in SymplecticNeuralNetworks?

[benedict-96]

This is used for symbolic expressions. It is probably debatable if this is the cleanest way of implementing this functionality. May deserve a new issue but may not be solved in this pr.

[michakraus]

Force-pushed a rebase onto main (0.5.0)

The previous tip was 43d89206; nothing is lost, that commit is still reachable by SHA on GitHub and I have it as a local branch. The branch is now one commit, 633878f, and the PR is MERGEABLE again.

What was dropped, and why

commit	verdict
cc1a786f "Fix GML compilation and add generated artifacts"	dropped. 123 MB of build output — docs/build.zip (77 MB), docs/build 2.zip (47 MB), eleven LaTeX .aux/.log/.fls/.fdb_latexmk files, two go_migration_inspection_*.txt — and about four lines of source. .gitignore already covers all of it bar the inspection logs. The four lines are kept. It did not, in fact, fix compilation.
8d243592 "Resolve merge conflicts"	dropped. Not this branch's work: it is main's own HDF5 migration re-applied by hand. ext/HDF5Ext.jl, test/hdf5_support.jl, docs/utils/adjust_image_size.jl, scripts/convert_jld2_to_h5.jl, src/architectures/symplectic_autoencoder.jl and src/layers/bias_layer.jl were byte-identical blobs to main's. Redundant once rebased.
43d89206 "Merge main into…"	dropped. It merged main only as far as d07b4c26, so the branch never saw the GeometricOptimizers separation, and its resolution left src/GeometricMachineLearning.jl half from each side: go_bridges.jl, src/arrays/ and src/manifolds/ included again next to main's import GeometricOptimizers. That is the Method overwriting is not permitted during Module precompilation that all 13 jobs were failing on.
4c0e2684 "Fix Hamiltonian symbolic code generation"	kept. A real seven-line fix — the Jacobian broadcast (dx.(scalarize(f)) rather than scalarize(dx(f))) and inplace = false on two build_nn_function calls.

The other 53 commits are collapsed into one, still authored by @benedict-96.

Why the merges kept going wrong: main reformatted src/GeometricMachineLearning.jl from a 4-space-indented module body to column 0, and this branch never did — so every merge of main conflicts on the whole file, and each has been resolved a bit differently by hand. Here I took main's version as-is and re-applied only this branch's includes and exports. Worth keeping in mind for any other long-lived branch.

Nothing that moved to GeometricOptimizers came back: git diff --name-only main...HEAD -- src/arrays src/manifolds src/optimizers/go_bridges.jl test/manifolds test/optimizers is empty, and no file main deleted is present.

⚠️ This cannot merge until NeuralNetworkParameters.jl is registered

Every job now fails with one error, at dependency resolution — nothing compiles or tests wrong:

ERROR: LoadError: expected package `NeuralNetworkParameters [67f4d93a]` to be registered

The branch used ParameterHandling to flatten the system parameters into the network input, and pirated three flatten methods on it. That does not work at all: GeometricOptimizers defines ParameterHandling.flatten(x) with an unbound type parameter, and that method is the one that wins —

ERROR: UndefVarError: `T` not defined in static parameter matching
  [1] flatten(x::@NamedTuple{})
    @ GeometricOptimizers …/src/optimizers/named_tuple_wrapper.jl:14

which is D6 in NeuralNetworkParameters' PLAN.md, hit in practice. flatten/unflatten from NeuralNetworkParameters replace it — and because AbstractNeuralNetworks already exports params, that needs no method on a foreign type, so three pirated methods disappear rather than move. GML supports julia = "1.10", where [sources] does not exist, so registering NNP 0.1.0 is a hard prerequisite.

Three bugs, all on the training path

test/training_phnn.jl is not in runtests.jl, so the loop this PR exists for had never run:

1. concatenate_array_with_parameters(::AbstractMatrix, ::AbstractVector) used vcat where it needs hcat, collapsing a batch into one long vector.
2. optimize_for_one_epoch! called _unpack_tuple, which has never existed in this package.
3. Zygote differentiates through the NeuralNetworkParameters struct, so every nesting level of the gradient comes back wrapped in (params = …,). These architectures nest — a SymplecticEuler layer holds a whole sub-network — and _get_params only unwraps the top, so _unwrap_gradient recurses.

New pghnn_training_test.jl covers that path. The two existing tests are rewritten: parametric_data_loader_test.jl pinned which shuffled batch holds which parameters, which depends on the Julia version's RNG stream, and now asserts the correspondence itself; pghnn_symbolic_pullback_single_layer_test.jl had no @test at all and re-defined in the test file what src/ now provides, and now checks the symbolic gradient against the summed per-sample Zygote gradients — they agree to 4e-16.

Dependency drift

The branch predates SymbolicNeuralNetworks 0.5: symbolize! is gone (symbolic_variables replaces it) and symbolic variables are scalar Nums rather than Symbolics.Arrs (SNN#14). The parametric SymbolicPullback is rewritten in the shape of SymbolicNeuralNetworks.SymbolicPullback, and five Symbolics.Arr special cases — three of them piracy — became unreachable and are deleted. GeometricProblems.default_parameters is a function since 0.8. src/optimizers/optimizer.jl keeps main's version, since the 0.5.0 rewrite covers what the branch added there.

The three items from the original description

• Type piracy — removed where it was free (Base.NamedTuple(::NeuralNetworkParameters) is just params; the ParameterHandling and Symbolics.Arr methods went with their callers). The rest each carry a TODO naming their proper home and are collected in Type piracy introduced by the parametric Hamiltonian architectures (#207) #243; they need releases of AbstractNeuralNetworks and SymbolicNeuralNetworks, so they are deliberately deferred. NNP's PLAN.md did not cover any of this — the two entries that are its to own are now D9 and D10 there.
• QPT2 — introduced and documented, QPT widened. Replacing QPT with it is still open.
• Parameter-dependent DataLoader — done, as ParametricDataLoader.

Verified locally

Pkg.test() passes (exit 0; the four new testsets are 2 + 8 + 4 + 2 assertions, no failures) and docs/make.jl builds clean, both against a deved NeuralNetworkParameters.

🤖 Generated with Claude Code

[michakraus]

Thanks — worked through all of these. Summary, then the two where the answer turned out to be different from what the comment assumed.

Done

Dead code (parametric_data_loader.jl:52, :61, parametric_resnet_layer.jl:38) — all three deleted. They were not just unused but obsolete: rearrange_parameters was never finished (zeros(_eltype(parameters), ) has no dimensions) and never called; the commented batch_over_two_axes is superseded by the n_batches::Integer method in batch.jl that the live code calls; and the commented three-dimensional ParametricResNetLayer methods were wrong — both compute input and then use x — so they were not a starting point for a real implementation either.

The D:\RESEARCH - UTWENTE\… path — gone, now joinpath(get(ENV, "GML_OUTPUT_DIR", @__DIR__), "damped_oscillator_network.jld2"). It was the only absolute path in the PR.

Script duplication — you were right, and it was worse than the comments suggest: turn_q_p_data_into_correct_format was byte-identical in all three scripts, and the two TimeDependentHarmonicOscillator* scripts shared a ~110-line prelude that differed only in the constants F and T. All of it now lives in scripts/parametric_data_helpers.jl, including the shared analytic solution as forced_harmonic_oscillator_solution(t, IC; omega, Omega, F) — about 200 lines less duplication. The file's header says it is a staging post, and GMLDatasets #5 tracks moving the helpers where they belong; I wrote up there the two open questions, including your point that the data has to come from closed-form solutions because there is no good way to get analytic solutions out of GeometricProblems at the moment.

Missing documentation — partly. ForcedGeneralizedHamiltonianArchitecture and ForcedSympNet already had docstrings and manual entries. But WideResNetLayer, ParametricResNetLayer and ParametricResNet had no docstring at all and were absent from the manual; they slipped past missing_docs because Documenter only checks bindings that already carry one. All three are written and in the manual now. The narrative documentation — and your tikz figures from the paper, which I would very much like — I have deliberately left; say the word and I will write the section around them.

_flatten_system_parameters (utils.jl:204)

Agreed it is debatable. One thing that may change the picture: it is no longer ParameterHandling. That could not work at all — GeometricOptimizers defines ParameterHandling.flatten(x) with an unbound type parameter and that method wins, so any call errors with UndefVarError: T not defined in static parameter matching. It is NeuralNetworkParameters.flatten/unflatten now, which also means the layer stores a ParameterLayout — a value, so inferable and storable — rather than a closure. Happy to open an issue if you still want the design revisited.

SymbolicPullback performance — measured, and it is the other way round

This is the one worth reading. Per call the symbolic pullback is 100–1000× faster than Zygote, and flat:

network (dim, width, n_int)	params	ZygotePullback	SymbolicPullback
2, 2, 1	32	10.4 ms	0.1 ms
2, 4, 1	80	23.6 ms	0.1 ms
4, 4, 1	88	33.6 ms	0.1 ms
4, 4, 2	176	140.1 ms	— never builds

What explodes is the build, and multiplicatively in n_integrators. Phase by phase at dim = 4, width = 4, nhidden = 1:

phase	n_int = 1	n_int = 2
build the symbolic loss	0.87 s, 3.4 ⋅ 10⁵ chars	0.03 s, 1.4 ⋅ 10⁹ chars
symbolic_derivative	0.66 s, 1.5 ⋅ 10⁸ chars	never finishes
build_nn_function	1.34 s	—

The loss expression grows ~4,100× going from one integrator to two, before any differentiation, and differentiating multiplies by a further ~440×. At two integrators the process passes 8 GB and does not return.

Cause: each SymplecticEuler layer's forward pass calls the executable gradient build_gradient produced for its energy network. Traced on symbolic Nums that inlines the entire gradient expression, so stacking integrators inlines it inside itself. cse = true cannot help — it runs at code generation, long after the expression has been materialised.

Fix: stop tracing the inlined chain and compose the pullback layer by layer from the gradients each SymplecticEuler has already built, which keeps the expression linear in depth. That answers your other question too — yes, parts of this belong in SymbolicNeuralNetworks: it needs either a layerwise pullback protocol there, or a build_nn_function that can take a composition without inlining it. Written up with all the numbers in #245.

I did not fix it here — it is a redesign of the symbolic pipeline, not a tweak. What this PR does is make the failure honest: SymbolicPullback(nn, ::ParametricLoss, μ) now throws an ArgumentError for n_integrators > 1 explaining why and pointing at ZygotePullback, instead of appearing to hang. One integrator builds in ≈1.4 s and is fully supported; the default training path uses ZygotePullback and is unaffected.

Everything is on 2463e03b; Pkg.test() and docs/check_references.jl both pass locally.

🤖 Generated with Claude Code