From 633878fcfe9c13bc83ec42cc2da9320a4edc0409 Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Wed, 19 Aug 2026 18:14:48 +0900 Subject: [PATCH 1/3] Rebase the parametric generalized Hamiltonian neural networks onto main Reconstructs PR #207 on top of 0.5.0. The branch's own 53 commits are collapsed into this one; the three that followed them are not carried over. `cc1a786f` "Fix GML compilation and add generated artifacts" was 123 MB of build output -- `docs/build.zip` (77 MB), `docs/build 2.zip` (47 MB), eleven LaTeX `.aux`/`.log`/`.fls`/`.fdb_latexmk` files and two `go_migration_inspection_*.txt` -- around four lines of source. `.gitignore` already covers all of it bar the inspection logs. `8d243592` "Resolve merge conflicts" 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 conflict resolution left `src/GeometricMachineLearning.jl` half from each side: `go_bridges.jl`, `src/arrays/` and `src/manifolds/` were `include`d again next to main's `import GeometricOptimizers`, which is `Method overwriting is not permitted during Module precompilation` on Julia 1.10. `4c0e2684` is kept -- it is a real seven-line fix to the symbolic `Jacobian` broadcast and two `build_nn_function` calls. Root cause of the repeated bad merges: 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 `include`s and `export`s are re-applied. What the feature adds `GeneralizedHamiltonianArchitecture` is implemented; it used to be a stub whose constructor threw. It composes `n_integrators` symplectic Euler steps, each differentiating a learned kinetic or potential energy. Around it: `ForcedGeneralizedHamiltonianArchitecture` and `ForcedSympNet` with `ForcingLayer`s, `ParametricDataLoader`, `ParametricLoss`, a `SymbolicPullback` for it, `ParametricResNet`, and `QPT2`/`QPTOAT2`. Adaptation to the current dependencies The branch predates SymbolicNeuralNetworks 0.5. `symbolize!` is gone -- `symbolic_variables` replaces it -- and symbolic variables are scalar `Num`s rather than `Symbolics.Arr`s (SNN#14), so the parametric `SymbolicPullback` is rewritten in the shape of `SymbolicNeuralNetworks.SymbolicPullback`, with `symbolic_derivative` and `ParameterGradient` in place of the 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 three tests call it. `src/optimizers/optimizer.jl` keeps main's version: the branch widened `_optimization_step!` and added `rgrad` methods for `NeuralNetworkParameters` and `nothing` gradients, and the 0.5.0 rewrite covers all of it -- `_tree_optim_step!` already skips a `nothing` block. NeuralNetworkParameters instead of ParameterHandling The system parameters are flattened into the network input. The branch used `ParameterHandling` and pirated three `flatten` methods on it; that does not work, because GeometricOptimizers defines `ParameterHandling.flatten(x)` with an unbound type parameter and that method wins -- D6 in NeuralNetworkParameters' PLAN.md, hit in practice. `flatten`/`unflatten` from NeuralNetworkParameters replace it. Since AbstractNeuralNetworks already exports `params`, no method on a foreign type is needed, and the layout stored in a layer is now a value rather than a closure. This makes registering NeuralNetworkParameters a prerequisite for merging: GML supports Julia 1.10, where `[sources]` does not exist. Bugs fixed on the way -- all on the training path, which nothing ran `concatenate_array_with_parameters(::AbstractMatrix, ::AbstractVector)` used `vcat` where it needs `hcat`, collapsing a batch into one long vector. `optimize_for_one_epoch!` called `_unpack_tuple`, which has never existed in this package. And `Zygote` differentiates *through* the `NeuralNetworkParameters` struct, so every nesting level of the gradient comes back wrapped in `(params = ...,)`; these architectures nest, and `_get_params` only unwraps the top, so `_unwrap_gradient` recurses. `test/generalized_hamiltonian_neural_networks/pghnn_training_test.jl` now covers that path. The other two tests are rewritten: the data-loader one asserted which shuffled batch holds which parameters, which depends on the RNG stream of the Julia version, 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, which agree to 4e-16. Type piracy Deleted where it was free: `Base.NamedTuple(::NeuralNetworkParameters)` is `params`, and the `ParameterHandling` and `Symbolics.Arr` methods are gone with their callers. What remains -- `applychain`, four `Chain` functors, `SymbolicNeuralNetworks.Jacobian`, `networkbackend(::LazyArrays.ApplyArray)`, `h5save(::HDF5.Group, ::NeuralNetworkParameters, ...)` and the `SymbolicPullback` call operators -- carries a `TODO` naming its proper home. Also: `GeometricProblems` and `Printf` are dropped from `[deps]`, where the branch had duplicated them out of `[extras]` and nothing in `src/` uses them; `SymplecticEuler`, `SymplecticEulerA` and `SymplecticEulerB` are no longer exported, the training methods that used to carry those names being `SymplecticEulerIntegrator*` now; and the training loop takes `Union{DataLoader, ParametricDataLoader}` rather than an untyped argument. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 51 +++ Project.toml | 5 + docs/src/GeometricMachineLearning.bib | 15 +- .../hamiltonian_neural_network.md | 17 +- docs/src/architectures/sympnet.md | 1 + docs/src/reduced_order_modeling/losses.md | 1 + .../structure_preservation/symplecticity.md | 2 + ext/HDF5Ext.jl | 10 + ...rcedGeneralizedHamiltonianNeuralNetwork.jl | 161 +++++++++ ...ndentHarmonicOscillatorParametricResnet.jl | 150 ++++++++ scripts/Train_DampedOscillator_QP.jl | 114 ++++++ scripts/forcing_layers_parameter_number.jl | 11 + src/GeometricMachineLearning.jl | 27 +- ..._generalized_hamiltonian_neural_network.jl | 35 ++ src/architectures/forced_sympnet.jl | 65 ++++ .../generalized_hamiltonian_neural_network.jl | 335 ++++++++++++++++++ .../hamiltonian_neural_network.jl | 114 ------ src/architectures/parametric_resnet.jl | 37 ++ src/architectures/resnet.jl | 13 +- .../standard_hamiltonian_neural_network.jl | 85 +++++ src/data_loader/batch.jl | 7 +- src/data_loader/optimize.jl | 4 +- src/data_loader/parametric_data_loader.jl | 143 ++++++++ src/layers/forcing_dissipation_layers.jl | 163 +++++++++ src/layers/parametric_resnet_layer.jl | 64 ++++ src/layers/sympnets.jl | 12 +- src/layers/wide_resnet.jl | 32 ++ src/loss/losses.jl | 22 ++ src/pullbacks/symbolic_hnn_pullback.jl | 69 ++++ src/pullbacks/zygote_pullback.jl | 6 + src/training_method/symplectic_euler.jl | 20 +- src/utils.jl | 84 ++++- .../parametric_data_loader_test.jl | 45 +++ ...hnn_symbolic_pullback_single_layer_test.jl | 63 ++++ .../pghnn_training_test.jl | 47 +++ ...alized_hamiltonian_neural_networks_test.jl | 25 ++ test/runtests.jl | 12 + test/train!/test_method.jl | 4 +- test/training_phnn.jl | 8 +- 39 files changed, 1923 insertions(+), 156 deletions(-) create mode 100644 scripts/TimeDependentHarmonicOscillatorForcedGeneralizedHamiltonianNeuralNetwork.jl create mode 100644 scripts/TimeDependentHarmonicOscillatorParametricResnet.jl create mode 100644 scripts/Train_DampedOscillator_QP.jl create mode 100644 scripts/forcing_layers_parameter_number.jl create mode 100644 src/architectures/forced_generalized_hamiltonian_neural_network.jl create mode 100644 src/architectures/forced_sympnet.jl create mode 100644 src/architectures/generalized_hamiltonian_neural_network.jl create mode 100644 src/architectures/parametric_resnet.jl create mode 100644 src/architectures/standard_hamiltonian_neural_network.jl create mode 100644 src/data_loader/parametric_data_loader.jl create mode 100644 src/layers/forcing_dissipation_layers.jl create mode 100644 src/layers/parametric_resnet_layer.jl create mode 100644 src/layers/wide_resnet.jl create mode 100644 test/data_loader/parametric_data_loader_test.jl create mode 100644 test/generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl create mode 100644 test/generalized_hamiltonian_neural_networks/pghnn_training_test.jl create mode 100644 test/generalized_hamiltonian_neural_networks_test.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index aaf7f1130..f0759e07d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,57 @@ breaking release). > alongside the work. Where a release removed exported names the list is given; where it is a > reconstruction of intent, it says so. +## [Unreleased] + +### Added + +**Parametric generalized Hamiltonian neural networks (PGHNNs)** +([#207](https://github.com/JuliaGNI/GeometricMachineLearning.jl/pull/207)). A family of +architectures whose forward pass takes the parameters of the *system* alongside the state, so one +network covers a whole parameter range rather than a single problem instance. + +- **`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 of which differentiates a learned kinetic or potential energy — + `SymbolicKineticEnergy` and `SymbolicPotentialEnergy`, built into an executable gradient by + `build_gradient`. +- **`ForcedGeneralizedHamiltonianArchitecture`** and **`ForcedSympNet`**, which add `ForcingLayer`s + for forcing and dissipation in the `q`, `p` or both coordinates, following the + Lagrange–d'Alembert integrator of [marsden2001discrete](@cite). +- **`ParametricDataLoader`**, which carries one set of system parameters per trajectory and hands + the matching parameters to each sample of a batch. Built from an `EnsembleSolution` whose members + were integrated at different parameters. +- **`ParametricLoss`**, `FeedForwardLoss` with the system parameters threaded through, and a + `SymbolicPullback(nn, ::ParametricLoss, system_params)` that differentiates it symbolically. +- **`ParametricResNet`** and a widened **`ResNet`**, which now takes a `width` separate from the + system dimension and uses `WideResNetLayer` when the two differ. This is the non-structure-preserving + baseline the PGHNNs are compared against. +- `QPT2` and `QPTOAT2`: `QPT`/`QPTOAT` with the array rank fixed but the two array *types* allowed to + differ, which is what splitting an input array into `q` and `p` produces. + +**New dependency: [NeuralNetworkParameters][nnp].** The system parameters are flattened into the +network input, and `flatten`/`unflatten` do that. `ParameterHandling` cannot: `GeometricOptimizers` +defines `ParameterHandling.flatten(x)` with an unbound type parameter, and that method wins. + +### Changed + +- **`SymplecticEuler`, `SymplecticEulerA` and `SymplecticEulerB` are no longer exported.** The names + now belong to the layer type of the generalized architectures; the *training methods* they used to + name are `SymplecticEulerIntegrator`, `SymplecticEulerIntegratorA` and + `SymplecticEulerIntegratorB`. `SEuler`, `SEulerA` and `SEulerB`, which is how they are constructed, + are unchanged. +- `src/architectures/hamiltonian_neural_network.jl` is split: it keeps the abstract + `HamiltonianArchitecture`, and `StandardHamiltonianArchitecture` moves to + `standard_hamiltonian_neural_network.jl`. `hamiltonian_vector_field` is narrowed from + `::HamiltonianArchitecture` to `::StandardHamiltonianArchitecture` accordingly. + +### Fixed + +- `concatenate_array_with_parameters(::AbstractMatrix, ::AbstractVector)` concatenated a batch with + `vcat` rather than `hcat`, collapsing it into a single long vector. + +[nnp]: https://github.com/JuliaGNI/NeuralNetworkParameters.jl + ## [0.5.0] — 2026-08-19 **The optimizer machinery moves to [GeometricOptimizers][go].** GML no longer implements its own diff --git a/Project.toml b/Project.toml index fbc33712e..ce82c774a 100644 --- a/Project.toml +++ b/Project.toml @@ -18,6 +18,7 @@ KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" +NeuralNetworkParameters = "67f4d93a-60e9-472b-8cdd-1ccf6005724a" ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SymbolicNeuralNetworks = "aed23131-dcd0-47ca-8090-d21e605652e3" @@ -28,6 +29,9 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [weakdeps] HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" +[sources] +NeuralNetworkParameters = {path = "/Users/mkraus/Datashare/Julia/NeuralNetworkParameters"} + [extensions] HDF5Ext = "HDF5" @@ -47,6 +51,7 @@ HDF5 = "0.16, 0.17" KernelAbstractions = "0.9" LazyArrays = "=2.3.2" NNlib = "0.8, 0.9" +NeuralNetworkParameters = "0.1" ProgressMeter = "1" SafeTestsets = "0.1" SymbolicNeuralNetworks = "0.5" diff --git a/docs/src/GeometricMachineLearning.bib b/docs/src/GeometricMachineLearning.bib index e069aee6c..a4b523263 100644 --- a/docs/src/GeometricMachineLearning.bib +++ b/docs/src/GeometricMachineLearning.bib @@ -746,9 +746,10 @@ @article{bon2024optimal @article{kingma2014adam, title={Adam: a method for stochastic optimization}, - author={Kingma, DP}, + author={Kingma, Diederik P. and Ba, Jimmy Lei}, journal={arXiv preprint arXiv:1412.6980}, - year={2014} + year={2014}, + note={Published as a conference paper at ICLR 2015} } @article{toda1967vibration, @@ -923,6 +924,16 @@ @article{ge1988lie publisher={Elsevier} } +@article{marsden2001discrete, + title={Discrete mechanics and variational integrators}, + author={Marsden, Jerrold E and West, Matthew}, + journal={Acta numerica}, + volume={10}, + pages={357--514}, + year={2001}, + publisher={Cambridge University Press} +} + @article{otto2023learning, title={Learning nonlinear projections for reduced-order modeling of dynamical systems using constrained autoencoders}, author={Otto, Samuel E and Macchio, Gregory R and Rowley, Clarence W}, diff --git a/docs/src/architectures/hamiltonian_neural_network.md b/docs/src/architectures/hamiltonian_neural_network.md index be33ca27f..8c6234419 100644 --- a/docs/src/architectures/hamiltonian_neural_network.md +++ b/docs/src/architectures/hamiltonian_neural_network.md @@ -42,13 +42,28 @@ Here the derivatives (i.e. vector field data) ``\dot{q}_i^{(t)}`` and ``\dot{p}_ ## Library Functions ```@docs -GeometricMachineLearning.hamiltonian_vector_field(::HamiltonianArchitecture) +GeometricMachineLearning.hamiltonian_vector_field(::StandardHamiltonianArchitecture) GeometricMachineLearning.HamiltonianArchitecture GeometricMachineLearning.StandardHamiltonianArchitecture GeometricMachineLearning.HNNLoss GeometricMachineLearning.symbolic_hamiltonian_vector_field(::GeometricMachineLearning.SymbolicNeuralNetwork) GeometricMachineLearning.SymbolicPullback(::HamiltonianArchitecture) +GeometricMachineLearning.SymbolicEnergy +GeometricMachineLearning.SymbolicPotentialEnergy +GeometricMachineLearning.SymbolicKineticEnergy +GeometricMachineLearning.build_gradient +GeometricMachineLearning.SymplecticEulerA +GeometricMachineLearning.SymplecticEulerB GeometricMachineLearning.GeneralizedHamiltonianArchitecture +GeometricMachineLearning.ForcedGeneralizedHamiltonianArchitecture +GeometricMachineLearning.ForcingLayer +GeometricMachineLearning.ForcingLayerQ +GeometricMachineLearning.ForcingLayerP +GeometricMachineLearning.ForcingLayerQP +GeometricMachineLearning.ParametricDataLoader +GeometricMachineLearning.SymbolicPullback(::GeometricMachineLearning.NeuralNetwork, ::GeometricMachineLearning.ParametricLoss, ::GeometricMachineLearning.GeometricBase.OptionalParameters) +GeometricMachineLearning._flatten_system_parameters +GeometricMachineLearning._unwrap_gradient GeometricMachineLearning._processing GeometricMachineLearning._get_contents GeometricMachineLearning._get_params diff --git a/docs/src/architectures/sympnet.md b/docs/src/architectures/sympnet.md index 72cea0d3c..67978d107 100644 --- a/docs/src/architectures/sympnet.md +++ b/docs/src/architectures/sympnet.md @@ -215,6 +215,7 @@ is the predicted state. In the [example section](@ref "SympNets with `GeometricM SympNet LASympNet GSympNet +ForcedSympNet ``` ```@raw latex diff --git a/docs/src/reduced_order_modeling/losses.md b/docs/src/reduced_order_modeling/losses.md index ba933a840..f7398133b 100644 --- a/docs/src/reduced_order_modeling/losses.md +++ b/docs/src/reduced_order_modeling/losses.md @@ -46,6 +46,7 @@ where ``\mathbf{x}^{(t)}`` is the solution of the FOM at point ``t`` and ``\math TransformerLoss AutoEncoderLoss ReducedLoss +ParametricLoss projection_error reduction_error ``` diff --git a/docs/src/structure_preservation/symplecticity.md b/docs/src/structure_preservation/symplecticity.md index 1e46a0eaa..1fe7a3974 100644 --- a/docs/src/structure_preservation/symplecticity.md +++ b/docs/src/structure_preservation/symplecticity.md @@ -115,7 +115,9 @@ It is important to note that symplecticity is a very strong property[^2] that ma ```@docs PoissonTensor GeometricMachineLearning.QPT +GeometricMachineLearning.QPT2 GeometricMachineLearning.QPTOAT +GeometricMachineLearning.QPTOAT2 ``` ```@raw latex diff --git a/ext/HDF5Ext.jl b/ext/HDF5Ext.jl index 038263cb0..6f83df200 100644 --- a/ext/HDF5Ext.jl +++ b/ext/HDF5Ext.jl @@ -48,6 +48,16 @@ function h5save(h5::HDF5.H5DataStore, A::UpperTriangular, path::AbstractString) group["n"] = A.n end +# A `NeuralNetworkParameters` nested inside a parameter tree -- the parameter-dependent +# architectures put one per sub-network. AbstractNeuralNetworks has `save(::H5DataStore, +# ::NeuralNetworkParameters)` for the top level only. +# +# TODO: type piracy -- `h5save` and `NeuralNetworkParameters` are both AbstractNeuralNetworks'. +# This belongs in ANN's own `ext/HDF5Ext.jl`, next to `h5save(::H5DataStore, ::NamedTuple, …)`. +function h5save(h5::HDF5.Group, p::NeuralNetworkParameters, path::AbstractString) + h5save(h5, params(p), path) +end + # --------------------------------------------------------------------------- # changebackend — new methods for GML special array types # diff --git a/scripts/TimeDependentHarmonicOscillatorForcedGeneralizedHamiltonianNeuralNetwork.jl b/scripts/TimeDependentHarmonicOscillatorForcedGeneralizedHamiltonianNeuralNetwork.jl new file mode 100644 index 000000000..e9ca22ce1 --- /dev/null +++ b/scripts/TimeDependentHarmonicOscillatorForcedGeneralizedHamiltonianNeuralNetwork.jl @@ -0,0 +1,161 @@ +using HDF5 +using GeometricMachineLearning +using GeometricMachineLearning: QPT, QPT2, Activation, ParametricLoss, SymbolicNeuralNetwork, SymbolicPullback +using CairoMakie +using NNlib: relu + +# PARAMETERS +omega = 1.0 # natural frequency of the harmonic Oscillator +Omega = 3.5 # frequency of the external sinusoidal forcing +F = .9 # amplitude of the external sinusoidal forcing +ni_dim = 10 # number of initial conditions per dimension (so ni_dim^2 total) +T = 2π * 20 +nt = 1000 # number of time steps +dt = T/nt # time step + +# Generating the initial condition array +IC = vec( [(q=q0, p=p0) for q0 in range(-1, 1, ni_dim), p0 in range(-1, 1, ni_dim)] ) + +# Generating the solution array +ni = ni_dim^2 +q = zeros(Float64, ni, nt+1) +p = zeros(Float64, ni, nt+1) +t = collect(dt * range(0, nt, step=1)) + +""" +Turn a vector of numbers into a vector of `NamedTuple`s to be used by `ParametricDataLoader`. +""" +function turn_parameters_into_correct_format(t::AbstractVector, IC::AbstractVector{<:NamedTuple}) + vec_of_params = NamedTuple[] + for time_step ∈ t + time_step == t[end] || push!(vec_of_params, (t = time_step, )) + end + vcat((vec_of_params for _ in axes(IC, 1))...) +end + +for i in 1:nt+1 + for j=1:ni + q[j,i] = ( IC[j].p - Omega*F/(omega^2-Omega^2) )/ omega *sin(omega*t[i]) + IC[j].q*cos(omega*t[i]) + F/(omega^2-Omega^2)*sin(Omega*t[i]) + p[j,i] = -omega^2*IC[j].q*sin(omega*t[i]) + ( IC[j].p - Omega*F/(omega^2-Omega^2) )*cos(omega*t[i]) + Omega*F/(omega^2-Omega^2)*cos(Omega*t[i]) + # q[j,i] = ( IC[j].p - Omega*F/(omega^2-Omega^2) )/ omega *exp(-omega*t[i]) - IC[j].q*exp(-omega*t[i]) + F/(omega^2-Omega^2)*exp(-Omega*t[i]) + # p[j,i] = -omega^2*IC[j].q*exp(-omega*t[i]) + ( IC[j].p + Omega*F/(omega^2-Omega^2) )*exp(-omega*t[i]) - Omega*F/(omega^2-Omega^2)*exp(-Omega*t[i]) + end + +end + +@doc raw""" +Turn a `NamedTuple` of ``(q,p)`` data into two tensors of the correct format. + +This is the tricky part as the structure of the input array(s) needs to conform with the structure of the parameters. + +Here the data are rearranged in an array of size ``(n, 2, t_f - 1)`` where ``[t_0, t_1, \ldots, t_f]`` is the vector storing the time steps. + +If we deal with different initial conditions as well, we still put everything into the third (parameter) axis. + +# Example + +```jldoctest +using GeometricMachineLearning + +q = [1. 2. 3.; 4. 5. 6.] +p = [1.5 2.5 3.5; 4.5 5.5 6.5] +qp = (q = q, p = p) +turn_q_p_data_into_correct_format(qp) + +# output + +(q = [1.0 2.0; 4.0 5.0;;; 2.0 3.0; 5.0 6.0], p = [1.5 2.5; 4.5 5.5;;; 2.5 3.5; 5.5 6.5]) +``` +""" +function turn_q_p_data_into_correct_format(qp::QPT2{T, 2}) where {T} + number_of_time_steps = size(qp.q, 2) - 1 # not counting t₀ + number_of_initial_conditions = size(qp.q, 1) + q_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + p_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + for initial_condition_index ∈ 0:(number_of_initial_conditions - 1) + for time_index ∈ 1:number_of_time_steps + q_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index] + q_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index + 1] + p_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index] + p_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index + 1] + end + end + (q = q_array, p = p_array) +end + +# SAVING TO FILE + +# h5 = h5open(path, "w") +# write(h5, "q", q) +# write(h5, "p", p) +# write(h5, "t", t) +# +# attrs(h5)["ni"] = ni +# attrs(h5)["nt"] = nt +# attrs(h5)["dt"] = dt +# +# close(h5) + +""" +This takes time as a single additional parameter (third axis). +""" +function load_time_dependent_harmonic_oscillator_with_parametric_data_loader(qp::QPT{T}, t::AbstractVector{T}, IC::AbstractVector) where {T} + qp_reformatted = turn_q_p_data_into_correct_format(qp) + t_reformatted = turn_parameters_into_correct_format(t, IC) + ParametricDataLoader(qp_reformatted, t_reformatted) +end + +# This sets up the data loader +dl = load_time_dependent_harmonic_oscillator_with_parametric_data_loader((q = q, p = p), t, IC) + +# This sets up the neural network +width::Int = 1 +nhidden::Int = 1 +n_integrators::Int = 2 +# sigmoid_linear_unit(x::T) where {T<:Number} = x / (T(1) + exp(-x)) +arch1 = ForcedGeneralizedHamiltonianArchitecture(2; activation = tanh, width = width, nhidden = nhidden, n_integrators = n_integrators, parameters = turn_parameters_into_correct_format(t, IC)[1], forcing_type = :P) +arch2 = ForcedGeneralizedHamiltonianArchitecture(2; activation = tanh, width = width, nhidden = nhidden, n_integrators = n_integrators, parameters = turn_parameters_into_correct_format(t, IC)[1], forcing_type = :Q) +arch3 = ForcedGeneralizedHamiltonianArchitecture(2; activation = tanh, width = 2width, nhidden = nhidden, n_integrators = n_integrators, parameters = turn_parameters_into_correct_format(t, IC)[1], forcing_type = :QP) +nn1 = NeuralNetwork(arch1) +nn2 = NeuralNetwork(arch2) +nn3 = NeuralNetwork(arch3) + +# This is where training starts +batch_size = 128 +n_epochs = 200 +batch = Batch(batch_size) +o1 = Optimizer(AdamOptimizer(), nn1) +o2 = Optimizer(AdamOptimizer(), nn2) +o3 = Optimizer(AdamOptimizer(), nn3) +loss = ParametricLoss() +_pb = SymbolicPullback(nn1, loss, turn_parameters_into_correct_format(t, IC)[1]); +_pb = SymbolicPullback(nn2, loss, turn_parameters_into_correct_format(t, IC)[1]); +_pb = SymbolicPullback(nn3, loss, turn_parameters_into_correct_format(t, IC)[1]); + +function train_network() + o1(nn1, dl, batch, n_epochs, loss, _pb) + o2(nn2, dl, batch, n_epochs, loss, _pb) + o3(nn3, dl, batch, n_epochs, loss, _pb) +end + +loss_array = train_network() + +trajectory_number = 20 + +# Testing the network +initial_conditions = (q = q[trajectory_number, 1], p = p[trajectory_number, 1]) +n_steps = nt +trajectory = (q = zeros(1, n_steps), p = zeros(1, n_steps)) +trajectory.q[:, 1] .= initial_conditions.q +trajectory.p[:, 1] .= initial_conditions.p +# note that we have to supply the parameters as a named tuple as well here: +for t_step ∈ 0:(n_steps-2) + qp_temporary = nn3.model((q = [trajectory.q[1, t_step+1]], p = [trajectory.p[1, t_step+1]]), (t = t[t_step+1],), nn3.params) + trajectory.q[:, t_step+2] .= qp_temporary.q + trajectory.p[:, t_step+2] .= qp_temporary.p +end + +fig = Figure() +ax = Axis(fig[1,1]) +lines!(ax, trajectory.q[1,:]; label="nn") +lines!(ax, q[trajectory_number,:]; label="analytic") \ No newline at end of file diff --git a/scripts/TimeDependentHarmonicOscillatorParametricResnet.jl b/scripts/TimeDependentHarmonicOscillatorParametricResnet.jl new file mode 100644 index 000000000..5b6e0f31f --- /dev/null +++ b/scripts/TimeDependentHarmonicOscillatorParametricResnet.jl @@ -0,0 +1,150 @@ +using HDF5 +using GeometricMachineLearning +using GeometricMachineLearning: QPT, QPT2, Activation, ParametricLoss, SymbolicNeuralNetwork, SymbolicPullback +using CairoMakie +using NNlib: relu + +# PARAMETERS +omega = 1.0 # natural frequency of the harmonic Oscillator +Omega = 3.5 # frequency of the external sinusoidal forcing +F = .0 # .9 # amplitude of the external sinusoidal forcing +ni_dim = 10 # number of initial conditions per dimension (so ni_dim^2 total) +T = 2π * 5 +nt = 1000 # number of time steps +dt = T/nt # time step + +# Generating the initial condition array +IC = vec( [(q=q0, p=p0) for q0 in range(-1, 1, ni_dim), p0 in range(-1, 1, ni_dim)] ) + +# Generating the solution array +ni = ni_dim^2 +q = zeros(Float64, ni, nt+1) +p = zeros(Float64, ni, nt+1) +t = collect(dt * range(0, nt, step=1)) + +""" +Turn a vector of numbers into a vector of `NamedTuple`s to be used by `ParametricDataLoader`. +""" +function turn_parameters_into_correct_format(t::AbstractVector, IC::AbstractVector{<:NamedTuple}) + vec_of_params = NamedTuple[] + for time_step ∈ t + time_step == t[end] || push!(vec_of_params, (t = time_step, )) + end + vcat((vec_of_params for _ in axes(IC, 1))...) +end + +for i in 1:nt+1 + for j=1:ni + q[j,i] = ( IC[j].p - Omega*F/(omega^2-Omega^2) )/ omega *sin(omega*t[i]) + IC[j].q*cos(omega*t[i]) + F/(omega^2-Omega^2)*sin(Omega*t[i]) + p[j,i] = -omega^2*IC[j].q*sin(omega*t[i]) + ( IC[j].p - Omega*F/(omega^2-Omega^2) )*cos(omega*t[i]) + Omega*F/(omega^2-Omega^2)*cos(Omega*t[i]) + # q[j,i] = ( IC[j].p - Omega*F/(omega^2-Omega^2) )/ omega *exp(-omega*t[i]) - IC[j].q*exp(-omega*t[i]) + F/(omega^2-Omega^2)*exp(-Omega*t[i]) + # p[j,i] = -omega^2*IC[j].q*exp(-omega*t[i]) + ( IC[j].p + Omega*F/(omega^2-Omega^2) )*exp(-omega*t[i]) - Omega*F/(omega^2-Omega^2)*exp(-Omega*t[i]) + end + +end + +@doc raw""" +Turn a `NamedTuple` of ``(q,p)`` data into two tensors of the correct format. + +This is the tricky part as the structure of the input array(s) needs to conform with the structure of the parameters. + +Here the data are rearranged in an array of size ``(n, 2, t_f - 1)`` where ``[t_0, t_1, \ldots, t_f]`` is the vector storing the time steps. + +If we deal with different initial conditions as well, we still put everything into the third (parameter) axis. + +# Example + +```jldoctest +using GeometricMachineLearning + +q = [1. 2. 3.; 4. 5. 6.] +p = [1.5 2.5 3.5; 4.5 5.5 6.5] +qp = (q = q, p = p) +turn_q_p_data_into_correct_format(qp) + +# output + +(q = [1.0 2.0; 4.0 5.0;;; 2.0 3.0; 5.0 6.0], p = [1.5 2.5; 4.5 5.5;;; 2.5 3.5; 5.5 6.5]) +``` +""" +function turn_q_p_data_into_correct_format(qp::QPT2{T, 2}) where {T} + number_of_time_steps = size(qp.q, 2) - 1 # not counting t₀ + number_of_initial_conditions = size(qp.q, 1) + q_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + p_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + for initial_condition_index ∈ 0:(number_of_initial_conditions - 1) + for time_index ∈ 1:number_of_time_steps + q_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index] + q_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index + 1] + p_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index] + p_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index + 1] + end + end + (q = q_array, p = p_array) +end + +# SAVING TO FILE + +# h5 = h5open(path, "w") +# write(h5, "q", q) +# write(h5, "p", p) +# write(h5, "t", t) +# +# attrs(h5)["ni"] = ni +# attrs(h5)["nt"] = nt +# attrs(h5)["dt"] = dt +# +# close(h5) + +""" +This takes time as a single additional parameter (third axis). +""" +function load_time_dependent_harmonic_oscillator_with_parametric_data_loader(qp::QPT{T}, t::AbstractVector{T}, IC::AbstractVector) where {T} + qp_reformatted = turn_q_p_data_into_correct_format(qp) + t_reformatted = turn_parameters_into_correct_format(t, IC) + ParametricDataLoader(qp_reformatted, t_reformatted) +end + +# This sets up the data loader +dl = load_time_dependent_harmonic_oscillator_with_parametric_data_loader((q = q, p = p), t, IC) + +# This sets up the neural network +width::Int = 2 +n_blocks::Int = 1 +n_integrators::Int = 1 +# sigmoid_linear_unit(x::T) where {T<:Number} = x / (T(1) + exp(-x)) +arch = ResNet(2, n_blocks=n_blocks, width=width; activation=tanh, parameters=turn_parameters_into_correct_format(t, IC)[1]) +nn = NeuralNetwork(arch) + +# This is where training starts +batch_size = 128 +n_epochs = 200 +batch = Batch(batch_size) +o = Optimizer(AdamOptimizer(), nn) +loss = ParametricLoss() + +function train_network() + o(nn, dl, batch, n_epochs, loss) +end + +loss_array = train_network() + +trajectory_number = 20 + +# Testing the network +initial_conditions = (q = q[trajectory_number, 1], p = p[trajectory_number, 1]) +n_steps = nt +trajectory = (q = zeros(1, n_steps), p = zeros(1, n_steps)) +trajectory.q[:, 1] .= initial_conditions.q +trajectory.p[:, 1] .= initial_conditions.p +# note that we have to supply the parameters as a named tuple as well here: +for t_step ∈ 0:(n_steps-2) + qp_temporary = nn.model((q = [trajectory.q[1, t_step+1]], p = [trajectory.p[1, t_step+1]]), (t = t[t_step+1],), nn.params) + trajectory.q[:, t_step+2] .= qp_temporary.q + trajectory.p[:, t_step+2] .= qp_temporary.p +end + +fig = Figure() +ax = Axis(fig[1,1]) +lines!(ax, trajectory.q[1,:]; label="nn") +lines!(ax, q[trajectory_number,:]; label="analytic") \ No newline at end of file diff --git a/scripts/Train_DampedOscillator_QP.jl b/scripts/Train_DampedOscillator_QP.jl new file mode 100644 index 000000000..1058ef941 --- /dev/null +++ b/scripts/Train_DampedOscillator_QP.jl @@ -0,0 +1,114 @@ +using HDF5 +using GeometricMachineLearning +using GeometricMachineLearning: QPT, QPT2 +using CairoMakie +using JLD2 +using NNlib: relu + +# PARAMETERS +nu = 0.001 # friction force coefficient +ni_dim = 2 # number of initial conditions per dimension (so ni_dim^2 total) +T = 13 +nt = 100 # number of time steps +dt = T/nt # time step +n_epochs = 100000 +n_epochs = 3 +width = 4 # width of the neural network +nhidden = 3 # number of hidden layers in the neural network +batch_size = 5000 # the size of the batch + +path_out = "D:\\RESEARCH - UTWENTE\\GFHNNs\\Damped Oscillator\\network_TEST.jld2" +#path_out = "/home/tyranowskitm/GFHNNs/DampedOscillator/OUTPUTS/network.jld2" + + +# Generating the initial condition array +IC = vec( [(q=q0, p=p0) for q0 in range(-1, 1, ni_dim), p0 in range(-1, 1, ni_dim)] ) + + +# Generating the solution array +ni = ni_dim^2 +omega = sqrt(4-nu^2) / 2 + +q = zeros(Float64, ni, nt+1) +p = zeros(Float64, ni, nt+1) +t = collect(dt*range(0,nt,step=1)) + +for i in 1:nt+1 + + for j=1:ni + q[j,i] = (1/omega)*( IC[j].p + nu/2 *IC[j].q )*exp(-nu*t[i]/2)*sin(omega*t[i]) + IC[j].q*exp(-nu*t[i]/2)*cos(omega*t[i]) + p[j,i] = -(1/omega)*( IC[j].q + nu/2 *IC[j].p )*exp(-nu*t[i]/2)*sin(omega*t[i]) + IC[j].p*exp(-nu*t[i]/2)*cos(omega*t[i]) + end + +end + + + +@doc raw""" +Turn a `NamedTuple` of ``(q,p)`` data into two tensors of the correct format. + +This is the tricky part as the structure of the input array(s) needs to conform with the structure of the parameters. + +Here the data are rearranged in an array of size ``(n, 2, t_f - 1)`` where ``[t_0, t_1, \ldots, t_f]`` is the vector storing the time steps. + +If we deal with different initial conditions as well, we still put everything into the third (parameter) axis. + +# Example + +```jldoctest +using GeometricMachineLearning + +q = [1. 2. 3.; 4. 5. 6.] +p = [1.5 2.5 3.5; 4.5 5.5 6.5] +qp = (q = q, p = p) +turn_q_p_data_into_correct_format(qp) + +# output + +(q = [1.0 2.0; 4.0 5.0;;; 2.0 3.0; 5.0 6.0], p = [1.5 2.5; 4.5 5.5;;; 2.5 3.5; 5.5 6.5]) +``` +""" +function turn_q_p_data_into_correct_format(qp::QPT2{T, 2}) where {T} + number_of_time_steps = size(qp.q, 2) - 1 # not counting t₀ + number_of_initial_conditions = size(qp.q, 1) + q_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + p_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + for initial_condition_index ∈ 0:(number_of_initial_conditions - 1) + for time_index ∈ 1:number_of_time_steps + q_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index] + q_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index + 1] + p_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index] + p_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index + 1] + end + end + (q = q_array, p = p_array) +end + + +# This sets up the data loader +dl = DataLoader(turn_q_p_data_into_correct_format((q = q, p = p))) + +# This sets up the neural network +arch = ForcedSympNet(2; upscaling_dimension = width, n_layers = nhidden, forcing_type = :P) +#arch = ForcedSympNet(2; upscaling_dimension = width, n_layers = nhidden, activation=(x-> max(0,x)^2/2)) +nn = NeuralNetwork(arch) + +# This is where training starts +batch = Batch(batch_size) +o = Optimizer(AdamOptimizer(), nn) + +loss_array = o(nn, dl, batch, n_epochs) + + +# Saving the parameters of the network +println("Saving the parameters of the neural network...") +flush(stdout) + +params = GeometricMachineLearning.map_to_cpu(nn.params) + +save(path_out,"parameters", params, "training loss", loss_array, "ni_dim", ni_dim, "T", T, "nt", nt, "n_epochs", n_epochs, "width", width, "nhidden", nhidden, "batch_size", batch_size, "nu", nu) + +println(" ...Done!") +flush(stdout) + + diff --git a/scripts/forcing_layers_parameter_number.jl b/scripts/forcing_layers_parameter_number.jl new file mode 100644 index 000000000..e27376d04 --- /dev/null +++ b/scripts/forcing_layers_parameter_number.jl @@ -0,0 +1,11 @@ +using GeometricMachineLearning +using GeometricMachineLearning: ForcingLayerP, ForcingLayerQP + +forcing_layer_p = ForcingLayerP(2) +forcing_layer_qp = ForcingLayerQP(2) + +nn_p = NeuralNetwork(forcing_layer_p) +nn_qp = NeuralNetwork(forcing_layer_qp) + +println(parameterlength(nn_p)) +println(parameterlength(nn_qp)) \ No newline at end of file diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 22a5ac5aa..cec211b22 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -21,8 +21,12 @@ using TimerOutputs import LazyArrays import SymbolicNeuralNetworks import SymbolicNeuralNetworks: SymbolicPullback -using SymbolicNeuralNetworks: derivative, SymbolicNeuralNetwork +using SymbolicNeuralNetworks: derivative, SymbolicNeuralNetwork, AbstractSymbolicNeuralNetwork import Symbolics +# The system parameters of a parameter-dependent architecture are flattened into the network's +# input. Only the two conversions are brought in: the module name would clash with +# `AbstractNeuralNetworks.NeuralNetworkParameters`, the *type* GML re-exports. +using NeuralNetworkParameters: flatten, unflatten # The manifolds, the structured matrix types, the global sections and the retractions are # `GeometricOptimizers`' — GML used to carry near-verbatim copies of all eleven types, which Julia @@ -179,8 +183,11 @@ export StiefelManifold, GrassmannManifold, Manifold export rgrad, metric, check include("layers/sympnets.jl") +include("layers/forcing_dissipation_layers.jl") include("layers/bias_layer.jl") include("layers/resnet.jl") +include("layers/wide_resnet.jl") +include("layers/parametric_resnet_layer.jl") include("layers/manifold_layer.jl") include("layers/stiefel_layer.jl") include("layers/grassmann_layer.jl") @@ -289,11 +296,13 @@ export arch include("backends/backends.jl") include("backends/lux.jl") -export NetworkLoss, TransformerLoss, FeedForwardLoss, AutoEncoderLoss, ReducedLoss, HNNLoss +export NetworkLoss, TransformerLoss, FeedForwardLoss, AutoEncoderLoss, ReducedLoss, HNNLoss, + ParametricLoss #INCLUDE ARCHITECTURES include("architectures/neural_network_integrator.jl") include("architectures/resnet.jl") +include("architectures/parametric_resnet.jl") include("architectures/transformer_integrator.jl") include("architectures/standard_transformer_integrator.jl") include("architectures/sympnet.jl") @@ -302,6 +311,8 @@ include("architectures/symplectic_autoencoder.jl") include("architectures/psd.jl") include("architectures/fixed_width_network.jl") include("architectures/hamiltonian_neural_network.jl") +include("architectures/standard_hamiltonian_neural_network.jl") +include("architectures/generalized_hamiltonian_neural_network.jl") include("architectures/lagrangian_neural_network.jl") include("architectures/variable_width_network.jl") include("architectures/recurrent_neural_network.jl") @@ -319,7 +330,8 @@ export ClassificationTransformer, ClassificationLayer export VolumePreservingFeedForward export SymplecticAutoencoder, PSDArch export HamiltonianArchitecture, StandardHamiltonianArchitecture, - GeneralizedHamiltonianArchitecture + GeneralizedHamiltonianArchitecture, ForcedGeneralizedHamiltonianArchitecture +export ForcedSympNet export solve!, encoder, decoder @@ -337,13 +349,18 @@ export AbstractPullback, ZygotePullback, SymbolicPullback include("pullbacks/zygote_pullback.jl") include("pullbacks/symbolic_hnn_pullback.jl") -export DataLoader +export DataLoader, ParametricDataLoader export Batch, optimize_for_one_epoch! include("data_loader/tensor_assign.jl") include("data_loader/matrix_assign.jl") include("data_loader/batch.jl") +# before `optimize.jl`, whose training loop takes either data loader +include("data_loader/parametric_data_loader.jl") include("data_loader/optimize.jl") +include("architectures/forced_sympnet.jl") +include("architectures/forced_generalized_hamiltonian_neural_network.jl") + # INCLUDE TRAINING parameters export TrainingParameters @@ -391,8 +408,6 @@ export train! include("training/train.jl") -export SymplecticEuler -export SymplecticEulerA, SymplecticEulerB export SEuler, SEulerA, SEulerB include("training_method/symplectic_euler.jl") diff --git a/src/architectures/forced_generalized_hamiltonian_neural_network.jl b/src/architectures/forced_generalized_hamiltonian_neural_network.jl new file mode 100644 index 000000000..2bb97632e --- /dev/null +++ b/src/architectures/forced_generalized_hamiltonian_neural_network.jl @@ -0,0 +1,35 @@ +const N_FORCING_LAYERS_DEFAULT = 2 + +""" + ForcedGeneralizedHamiltonianArchitecture <: HamiltonianArchitecture + +A version of [`GeneralizedHamiltonianArchitecture`](@ref) that includes forcing/dissipation terms. Also compare this to [`ForcedSympNet`](@ref). +""" +struct ForcedGeneralizedHamiltonianArchitecture{FT, AT, PT <: OptionalParameters} <: HamiltonianArchitecture{AT} + dim::Int + width::Int + nhidden::Int + n_forcing_layers::Int + n_integrators::Int + parameters::PT + activation::AT + + function ForcedGeneralizedHamiltonianArchitecture(dim; width=dim, nhidden=HNN_nhidden_default, n_forcing_layers=N_FORCING_LAYERS_DEFAULT, n_integrators::Integer=1, activation=HNN_activation_default, parameters=NullParameters(), forcing_type::Symbol=:P) + forcing_type == :P || forcing_type == :Q || forcing_type == :QP || error("Forcing has to be either :Q or :P. It is $(forcing_type).") + activation = (typeof(activation) <: Activation) ? activation : Activation(activation) + new{forcing_type, typeof(activation), typeof(parameters)}(dim, width, nhidden, n_forcing_layers, n_integrators, parameters, activation) + end +end + +function Chain(arch::ForcedGeneralizedHamiltonianArchitecture{FT}) where {FT} + layers = () + kinetic_energy = SymbolicKineticEnergy(arch.dim, arch.width, arch.nhidden, arch.activation; parameters=arch.parameters) + potential_energy = SymbolicPotentialEnergy(arch.dim, arch.width, arch.nhidden, arch.activation; parameters=arch.parameters) + for i ∈ 1:arch.n_integrators + layers = (layers..., SymplecticEulerA(kinetic_energy; return_parameters = true)) + layers = (layers..., ForcingLayer(arch.dim, arch.width, arch.n_forcing_layers, arch.activation; parameters=arch.parameters, return_parameters=true, type=FT)) + _return_parameters = !(i == arch.n_integrators) + layers = (layers..., SymplecticEulerB(potential_energy; return_parameters = _return_parameters)) + end + Chain(layers...) +end \ No newline at end of file diff --git a/src/architectures/forced_sympnet.jl b/src/architectures/forced_sympnet.jl new file mode 100644 index 000000000..582acca44 --- /dev/null +++ b/src/architectures/forced_sympnet.jl @@ -0,0 +1,65 @@ +@doc raw""" + ForcedSympNet <: NeuralNetworkIntegrator + +`ForcedSympNet`s are based on [`SympNet`](@ref)s [jin2020sympnets](@cite) and include [`ForcingLayer`](@ref)s. They are based on [`GSympNet`](@ref)s. + +# Constructor + +```julia +ForcedSympNet(d) +``` + +Make a forced SympNet with dimension ``d.`` + +# Arguments + +Keyword arguments are: +- `upscaling_dimension::Int = 2d`: The *upscaling dimension* of the gradient layer. See the documentation for [`GradientLayerQ`](@ref) and [`GradientLayerP`](@ref) for further explanation. +- `n_layers::Int""" * "$(g_n_layers_default)`" * raw""": The number of layers (i.e. the total number of [`GradientLayerQ`](@ref) and [`GradientLayerP`](@ref)). +- `activation""" * "$(g_activation_default)`" * raw""": The activation function that is applied. +- `init_upper::Bool""" * "$(g_init_upper_default)`" * raw""": Initialize the gradient layer so that it first modifies the $q$-component. +""" +struct ForcedSympNet{FT, AT} <: NeuralNetworkIntegrator + dim::Int + upscaling_dimension::Int + n_layers::Int + n_forcing_layers::Int + act::AT + init_upper::Bool + + function ForcedSympNet(dim::Integer; + upscaling_dimension = 2 * dim, + n_layers = g_n_layers_default, + n_forcing_layers = N_FORCING_LAYERS_DEFAULT, + activation = g_activation_default, + init_upper = g_init_upper_default, + forcing_type::Symbol = :P) + new{forcing_type, typeof(activation)}(dim, upscaling_dimension, n_layers, n_forcing_layers, activation, init_upper) + end + + function ForcedSympNet(dl::DataLoader; + upscaling_dimension = 2 * dl.input_dim, + n_layers = g_n_layers_default, + n_forcing_layers = N_FORCING_LAYERS_DEFAULT, + activation = g_activation_default, + init_upper = g_init_upper_default, + forcing_type::Symbol = :P) + new{forcing_type, typeof(activation)}(dl.input_dim, upscaling_dimension, n_layers, n_forcing_layers, activation, init_upper) + end +end + +function Chain(arch::ForcedSympNet{FT}) where {FT} + layers = () + is_upper_criterion = arch.init_upper ? isodd : iseven + for i in 1:arch.n_layers + layers = + if is_upper_criterion(i) + (layers..., GradientLayerQ(arch.dim, arch.upscaling_dimension, arch.act)) + else + (layers..., + ForcingLayer(arch.dim, arch.upscaling_dimension, arch.n_forcing_layers, arch.act; return_parameters=false, type=FT), + GradientLayerP(arch.dim, arch.upscaling_dimension, arch.act)) + end + end + Chain(layers...) +end \ No newline at end of file diff --git a/src/architectures/generalized_hamiltonian_neural_network.jl b/src/architectures/generalized_hamiltonian_neural_network.jl new file mode 100644 index 000000000..447a9cd3a --- /dev/null +++ b/src/architectures/generalized_hamiltonian_neural_network.jl @@ -0,0 +1,335 @@ +""" + SymbolicEnergy + +See [`SymbolicPotentialEnergy`](@ref) and [`SymbolicKineticEnergy`](@ref). +""" +struct SymbolicEnergy{AT <: Activation, PT, Kinetic} + dim::Int + width::Int + nhidden::Int + parameter_length::Int + parameter_layout::PT + activation::AT + + function SymbolicEnergy(dim, width, nhidden, activation; parameters::OptionalParameters=NullParameters(), type) + @assert iseven(dim) "The input dimension must be an even integer!" + flat_parameters, layout = _flatten_system_parameters(parameters) + _activation = Activation(activation) + new{typeof(_activation), typeof(layout), type}(dim, width, nhidden, length(flat_parameters), layout, _activation) + end +end + +""" + SymbolicPotentialEnergy + +A `const` derived from [`SymbolicEnergy`](@ref). + +# Constructors + +```jldoctest; setup=:(using GeometricMachineLearning; using GeometricMachineLearning: Activation) +julia> params, dim, width, nhidden, activation = (m = 1., ω = π / 2), 2, 2, 1, tanh +((m = 1.0, ω = 1.5707963267948966), 2, 2, 1, tanh) + +julia> se = GeometricMachineLearning.SymbolicPotentialEnergy(dim, width, nhidden, activation; parameters = params); + +``` + +In practice we use `SymbolicPotentialEnergy` (and [`SymbolicKineticEnergy`](@ref)) together with [`build_gradient(::SymbolicEnergy)`](@ref). + +# Parameter Dependence +""" +const SymbolicPotentialEnergy{AT, PT} = SymbolicEnergy{AT, PT, :potential} + +""" + SymbolicKineticEnergy + +A `const` derived from [`SymbolicEnergy`](@ref). + +# Constructors + +See [`SymbolicPotentialEnergy`](@ref). +""" +const SymbolicKineticEnergy{AT, PT} = SymbolicEnergy{AT, PT, :kinetic} + +SymbolicPotentialEnergy(args...; kwargs...) = SymbolicEnergy(args...; type = :potential, kwargs...) +SymbolicKineticEnergy(args...; kwargs...) = SymbolicEnergy(args...; type = :kinetic, kwargs...) + +function Chain(se::SymbolicEnergy) + inner_layers = Tuple( + [Dense(se.width, se.width, se.activation) for _ in 1:se.nhidden] + ) + + Chain( + Dense(se.dim÷2 + se.parameter_length, se.width, se.activation), + inner_layers..., + Linear(se.width, 1; use_bias = false) + ) +end + +# Jacobian with respect to the *first* `dim2` input variables only: for a parameter-dependent +# network the remaining input components are the system parameters, which are not differentiated. +# +# TODO: type piracy -- `Jacobian` and `AbstractSymbolicNeuralNetwork` are both +# SymbolicNeuralNetworks'. The restricted-Jacobian variant belongs there. +function SymbolicNeuralNetworks.Jacobian(f, nn::AbstractSymbolicNeuralNetwork, dim2::Integer) + # make differential of input variables (not of parameters) + Dx = SymbolicNeuralNetworks.symbolic_differentials(nn.input)[1:dim2] + + # Evaluation of gradient + s∇f = hcat([SymbolicNeuralNetworks.expand_derivatives.(dx.(SymbolicNeuralNetworks.Symbolics.scalarize(f))) for dx in Dx]...) + + SymbolicNeuralNetworks.Jacobian(f, s∇f, nn) +end + +function SymbolicNeuralNetworks.Jacobian(nn::AbstractSymbolicNeuralNetwork, dim2::Integer) + + # Evaluation of the symbolic output + soutput = nn.model(nn.input, params(nn)) + + SymbolicNeuralNetworks.Jacobian(soutput, nn, dim2) +end + +""" + build_gradient(se) + +Build a gradient function from a [`SymbolicEnergy`](@ref) `se`. + +# Examples + +```jldoctest; setup=:(using GeometricMachineLearning; using GeometricMachineLearning: SymbolicPotentialEnergy, build_gradient, concatenate_array_with_parameters, OneInitializer; using GeometricMachineLearning.GeometricBase: OptionalParameters) +params, dim, width, nhidden, activation = (m = 1., ω = π / 2), 4, 2, 1, tanh + +se = SymbolicPotentialEnergy(dim, width, nhidden, activation; parameters = params) + +# `OneInitializer` rather than the default random one, so that the output below does not depend on +# the random number stream of the Julia version +network_params = NeuralNetwork(Chain(se); initializer = OneInitializer()).params + +built_grad = build_gradient(se) +grad(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) = built_grad(concatenate_array_with_parameters(qp, problem_params), params) + +grad([0.5, 0.25], params, network_params) + +# output + +2×1 Matrix{Float64}: + 2.7907683385233434e-5 + 2.7907683385233434e-5 +``` +""" +function build_gradient(se::SymbolicEnergy) + model = Chain(se) + nn = SymbolicNeuralNetwork(model) + □ = SymbolicNeuralNetworks.Jacobian(nn, se.dim÷2) + SymbolicNeuralNetworks.build_nn_function(SymbolicNeuralNetworks.derivative(□)', nn.params, nn.input; + inplace = false) +end + +struct SymplecticEuler{M, N, FT<:Base.Callable, MT<:Chain, type, ReturnParameters} <: AbstractExplicitLayer{M, N} + gradient_function::FT + energy_model::MT +end + +function parameterlength(integrator::SymplecticEuler) + parameterlength(integrator.energy_model) +end + +function initialparameters(rng::Random.AbstractRNG, init_weight::AbstractNeuralNetworks.Initializer, integrator::SymplecticEuler, backend::KernelAbstractions.Backend, ::Type{T}) where {T} + initialparameters(rng, init_weight, integrator.energy_model, backend, T) +end + +const SymplecticEulerA{M, N, FT, AT, ReturnParameters} = SymplecticEuler{M, N, FT, AT, :A, ReturnParameters} +const SymplecticEulerB{M, N, FT, AT, ReturnParameters} = SymplecticEuler{M, N, FT, AT, :B, ReturnParameters} + +""" +Changes ``q`` (based on the kinetic energy). +""" +function SymplecticEulerA(se::SymbolicKineticEnergy; return_parameters::Bool) + gradient_function = build_gradient(se) + c = Chain(se) + SymplecticEuler{se.dim, se.dim, typeof(gradient_function), typeof(c), :A, return_parameters}(gradient_function, c) +end + +""" +Changes ``p`` (based on the potential energy). +""" +function SymplecticEulerB(se::SymbolicPotentialEnergy; return_parameters::Bool) + gradient_function = build_gradient(se) + c = Chain(se) + SymplecticEuler{se.dim, se.dim, typeof(gradient_function), typeof(c), :B, return_parameters}(gradient_function, c) +end + +# A network with no system parameters gets its input unchanged; without this the empty flat vector +# would have to be `vcat`ed on, which loses the element type. +concatenate_array_with_parameters(qp::AbstractArray, ::NullParameters) = qp + +function concatenate_array_with_parameters(qp::AbstractVector, params::NamedTuple) + vcat(qp, _flatten_system_parameters(params)[1]) +end + +function concatenate_array_with_parameters(qp::AbstractMatrix, params::NamedTuple) + @assert size(qp, 2) == 1 + vcat(qp, repeat(_flatten_system_parameters(params)[1], 1, size(qp, 2))) +end + +function concatenate_array_with_parameters(qp::AbstractArray{T, 3}, params::AbstractVector) where {T} + @assert size(qp, 3) == length(params) + matrices = Tuple(concatenate_array_with_parameters(qp[:, :, i], params[i]) for i in axes(qp, 3)) + cat(matrices...; dims = 3) +end + +# function concatenate_array_with_parameters(qp::AbstractMatrix, params::OptionalParameters) +# hcat((concatenate_array_with_parameters(qp[:, i], params) for i in axes(qp, 2))...) +# end + +# One parameter set per column, so the columns are concatenated *side by side*: the result is a +# matrix with `size(qp, 1) + parameter_length` rows, one column per sample. +function concatenate_array_with_parameters(qp::AbstractMatrix, params::AbstractVector) + @assert _size(qp, 2) == length(params) + hcat((concatenate_array_with_parameters(@view(qp[:, i]), params[i]) for i in axes(params, 1))...) +end + +function (integrator::SymplecticEulerA{M, N, FT, AT, false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT} + input = concatenate_array_with_parameters(qp.p, problem_params) + (q = @view((qp.q + integrator.gradient_function(input, params))[:, 1]), p = qp.p) +end + +function (integrator::SymplecticEulerB{M, N, FT, AT, false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT} + input = concatenate_array_with_parameters(qp.q, problem_params) + (q = qp.q, p = @view((qp.p - integrator.gradient_function(input, params))[:, 1])) +end + +function (integrator::SymplecticEulerA{M, N, FT, AT, true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT} + input = concatenate_array_with_parameters(qp.p, problem_params) + ((q = @view((qp.q + integrator.gradient_function(input, params))[:, 1]), p = qp.p), problem_params) +end + +function (integrator::SymplecticEulerB{M, N, FT, AT, true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT} + input = concatenate_array_with_parameters(qp.q, problem_params) + ((q = qp.q, p = @view((qp.p - integrator.gradient_function(input, params))[:, 1])), problem_params) +end + +function (integrator::SymplecticEuler)(qp_params::Tuple{<:QPTOAT2, <:OptionalParameters}, params::NeuralNetworkParameters) + integrator(qp_params..., params) +end + +function (integrator::SymplecticEuler)(::TT, ::NeuralNetworkParameters) where {TT <: Tuple} + error("The input is of type $(TT). This shouldn't be the case!") +end + +function (integrator::SymplecticEuler{M, N, FT, AT, Type, true})(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT, Type} + @assert iseven(size(qp, 1)) + n = size(qp, 1)÷2 + qp_split = assign_q_and_p(qp, n) + evaluated = integrator(qp_split, problem_params, params)[1] + (vcat(evaluated.q, evaluated.p), problem_params) +end + +function (integrator::SymplecticEuler{M, N, FT, AT, Type, false})(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT, Type} + @assert iseven(size(qp, 1)) + n = size(qp, 1)÷2 + qp_split = assign_q_and_p(qp, n) + evaluated = integrator(qp_split, problem_params, params) + vcat(evaluated.q, evaluated.p) +end + +(integrator::SymplecticEuler)(qp::QPTOAT2, params::NeuralNetworkParameters) = integrator(qp, NullParameters(), params) + +""" + GeneralizedHamiltonianArchitecture <: HamiltonianArchitecture + +A realization of generalized Hamiltonian neural networks (GHNNs) as introduced in [horn2025generalized](@cite). + +Also see [`StandardHamiltonianArchitecture`](@ref). + +# Constructor + +The constructor takes the following input arguments: +1. `dim`: system dimension, +2. `width = dim`: width of the hidden layer. By default this is equal to `dim`, +3. `nhidden = $(HNN_nhidden_default)`: the number of hidden layers, +4. `n_integrators`: the number of integrators used in the GHNN. +5. `activation = $(HNN_activation_default)`: the activation function used in the GHNN, +""" +struct GeneralizedHamiltonianArchitecture{AT, PT <: OptionalParameters} <: HamiltonianArchitecture{AT} + dim::Int + width::Int + nhidden::Int + n_integrators::Int + parameters::PT + activation::AT + + function GeneralizedHamiltonianArchitecture(dim; width=dim, nhidden=HNN_nhidden_default, n_integrators::Integer=1, activation=HNN_activation_default, parameters=NullParameters()) + activation = (typeof(activation) <: Activation) ? activation : Activation(activation) + new{typeof(activation), typeof(parameters)}(dim, width, nhidden, n_integrators, parameters, activation) + end +end + +# The parameter-dependent layers pass `(state, system parameters)` down the chain, so `applychain` +# has to accept that tuple as its data argument. +# +# TODO: type piracy -- `applychain` is AbstractNeuralNetworks' and every argument type here is +# `Base`'s. ANN's own `applychain(layers, x, ps::Union{NamedTuple, NeuralNetworkParameters})` is +# already generic in `x`; widening the `@generated` method the same way would remove the need. +@generated function AbstractNeuralNetworks.applychain(layers::Tuple, x::Tuple{<:QPTOAT2, <:OptionalParameters}, ps::Tuple) + N = length(fieldtypes((layers))) + x_symbols = vcat([:x], [gensym() for _ in 1:N]) + calls = [:(($(x_symbols[i + 1])) = layers[$i]($(x_symbols[i]), ps[$i])) for i in 1:N] + push!(calls, :(return $(x_symbols[N + 1]))) + return Expr(:block, calls...) +end + +index_qpt(qp::QPT2{T, 2}, i, j) where {T} = (q = qp.q[i, j], p = qp.p[i, j]) +index_gpt(qp::QPT2{T, 3}, i, j, k) where {T} = (q = qp.q[i, j, k], p = qp.p[i, j, k]) + +function Chain(ghnn_arch::GeneralizedHamiltonianArchitecture) + c = () + kinetic_energy = SymbolicKineticEnergy(ghnn_arch.dim, ghnn_arch.width, ghnn_arch.nhidden, ghnn_arch.activation; parameters=ghnn_arch.parameters) + potential_energy = SymbolicPotentialEnergy(ghnn_arch.dim, ghnn_arch.width, ghnn_arch.nhidden, ghnn_arch.activation; parameters=ghnn_arch.parameters) + + for n in 1:ghnn_arch.n_integrators + c = (c..., SymplecticEulerA(kinetic_energy; return_parameters = true)) + c = n == ghnn_arch.n_integrators ? (c..., SymplecticEulerB(potential_energy; return_parameters=false)) : (c..., SymplecticEulerB(potential_energy; return_parameters=true)) + end + + Chain(c...) +end + +function (nn::NeuralNetwork{GT})(qp::QPTOAT2, problem_params::OptionalParameters) where {GT <: GeneralizedHamiltonianArchitecture} + nn.model(qp, problem_params, params(nn)) +end + +# TODO: type piracy -- `Chain` is AbstractNeuralNetworks' and so is every argument type of the four +# functors below. A `ParametricChain` wrapper owned by GML, or these methods upstream, would fix it. +function (model::Chain)(qp::QPTOAT2, problem_params::OptionalParameters, params::Union{NeuralNetworkParameters, NamedTuple}) + model((qp, problem_params), params) +end + +function (c::Chain)(qp::QPT2{T, 3}, system_params::AbstractVector, ps::Union{NamedTuple, NeuralNetworkParameters})::QPT2{T} where {T} + @assert size(qp.q, 3) == length(system_params) + @assert size(qp.q, 2) == 1 + output_vectorwise = [c(index_gpt(qp, :, 1, i), system_params[i], ps) for i in axes(system_params, 1)] + q_output = hcat([single_output_vectorwise.q for single_output_vectorwise ∈ output_vectorwise]...) + p_output = hcat([single_output_vectorwise.p for single_output_vectorwise ∈ output_vectorwise]...) + (q = reshape(q_output, size(q_output, 1), 1, size(q_output, 2)), p = reshape(p_output, size(p_output, 1), 1, size(p_output, 2))) +end + +function (c::Chain)(qp::AbstractArray{T, 2}, system_params::AbstractVector, ps::Union{NamedTuple, NeuralNetworkParameters}) where {T} + @assert _size(qp, 2) == length(system_params) + qp_reshaped = reshape(qp, size(qp, 1), 1, length(system_params)) + c(qp_reshaped, system_params, ps) +end + +function (c::Chain)(qp::AbstractArray{T, 3}, system_params::AbstractVector, ps::Union{NamedTuple, NeuralNetworkParameters}) where {T} + @assert size(qp, 3) == length(system_params) + @assert size(qp, 2) == 1 + @assert iseven(size(qp, 1)) + n = size(qp, 1)÷2 + qp_split = assign_q_and_p(qp, n) + c_output = c(qp_split, system_params, ps)::QPT + reshape(vcat(c_output.q, c_output.p), 2n, length(system_params)) +end + +# TODO: type piracy -- `networkbackend` is AbstractNeuralNetworks' and `ApplyArray` is LazyArrays'. +# Belongs in ANN, which already dispatches `networkbackend` on array types it does not own either. +AbstractNeuralNetworks.networkbackend(::LazyArrays.ApplyArray) = AbstractNeuralNetworks.CPU() diff --git a/src/architectures/hamiltonian_neural_network.jl b/src/architectures/hamiltonian_neural_network.jl index 2b2ab690d..ea4aaf765 100644 --- a/src/architectures/hamiltonian_neural_network.jl +++ b/src/architectures/hamiltonian_neural_network.jl @@ -13,118 +13,4 @@ function HamiltonianArchitecture(dim::Integer, width::Integer, nhidden::Integer, StandardHamiltonianArchitecture(dim, width, nhidden, activation) end -""" - StandardHamiltonianArchitecture <: HamiltonianArchitecture - -A realization of the standard Hamiltonian neural network (HNN) [greydanus2019hamiltonian](@cite). - -Also see [`GeneralizedHamiltonianArchitecture`](@ref). - -# Constructor - -The constructor takes the following input arguments: -1. `dim`: system dimension, -2. `width = dim`: width of the hidden layer. By default this is equal to `dim`, -3. `nhidden = $(HNN_nhidden_default)`: the number of hidden layers, -4. `activation = $(HNN_activation_default)`: the activation function used in the HNN. -""" -struct StandardHamiltonianArchitecture{AT} <: HamiltonianArchitecture{AT} - dim::Int - width::Int - nhidden::Int - activation::AT - - function StandardHamiltonianArchitecture(dim, width=dim, nhidden=HNN_nhidden_default, activation=HNN_activation_default) - new{typeof(activation)}(dim, width, nhidden, activation) - end -end - -GHNN_integrator_default = nothing - -""" - GeneralizedHamiltonianArchitecture <: HamiltonianArchitecture - -A realization of generalized Hamiltonian neural networks (GHNNs) as introduced in [horn2025generalized](@cite). - -Also see [`StandardHamiltonianArchitecture`](@ref). - -# Constructor - -The constructor takes the following input arguments: -1. `dim`: system dimension, -2. `width = dim`: width of the hidden layer. By default this is equal to `dim`, -3. `nhidden = $(HNN_nhidden_default)`: the number of hidden layers, -4. `activation = $(HNN_activation_default)`: the activation function used in the GHNN, -5. `integrator = $(GHNN_integrator_default)`: the integrator that is used to design the GHNN. -""" -struct GeneralizedHamiltonianArchitecture{AT, IT} <: HamiltonianArchitecture{AT} - dim::Int - width::Int - nhidden::Int - activation::AT - integrator::IT - - function GeneralizedHamiltonianArchitecture(dim, width=dim, nhidden=HNN_nhidden_default, activation=HNN_activation_default, integrator=GHNN_integrator_default) - error("GHNN still has to be implemented!") - new{typeof(activation), typeof(integrator)}(dim, width, nhidden, activation, integrator) - end -end - @inline AbstractNeuralNetworks.dim(arch::HamiltonianArchitecture) = arch.dim - -""" - symbolic_hamiltonian_vector_field(nn::SymbolicNeuralNetwork) - -Get the symbolic expression for the vector field belonging to the HNN `nn`. - -# Implementation - -This is calling `SymbolicNeuralNetworks.Jacobian` and then multiplies the result with a Poisson tensor. -""" -function symbolic_hamiltonian_vector_field(nn::SymbolicNeuralNetwork) - □ = SymbolicNeuralNetworks.Jacobian(nn) - n = input_dimension(nn.model) ÷ 2 - # The Poisson tensor is built from *integers* on purpose: a `Float64` literal in the symbolic - # expression would widen the result of a `Float32` network. `PoissonTensor`, which has the same - # convention, is an `AbstractMatrix{Float64}` and would do exactly that. - 𝕆 = zeros(Int, n, n) - 𝕀 = Matrix(1I, n, n) - 𝕁 = [𝕆 𝕀; -𝕀 𝕆] - # `Jacobian` uses the convention `□[i, j] = ∂fᵢ/∂xⱼ` and the HNN output is scalar, so the one - # row of `derivative(□)` is the gradient of the Hamiltonian. The vector field is built as a - # *vector*, so that the generated function returns what `HNNLoss` compares against: a vector - # for a single sample, and one column per sample for a batch. - ∇H = vec(derivative(□)) - 𝕁 * ∇H -end - -""" - hamiltonian_vector_field(arch::HamiltonianArchitecture) - -Compute an executable expression of the Hamiltonian vector field of a [`HamiltonianArchitecture`](@ref). - -# Implementation - -This first computes a symbolic expression of the vector field using [`symbolic_hamiltonian_vector_field`](@ref). - -The function is built with `inplace = false`: [`HNNLoss`](@ref) wraps it and is differentiated with -`Zygote`, and the in-place kernel `SymbolicNeuralNetworks.build_nn_function` builds by default -*mutates* its result, which `Zygote` does not support. -""" -function hamiltonian_vector_field(arch::HamiltonianArchitecture) - nn = SymbolicNeuralNetwork(arch) - hvf = symbolic_hamiltonian_vector_field(nn) - SymbolicNeuralNetworks.build_nn_function(hvf, nn.params, nn.input; inplace = false) -end - -function Chain(arch::HamiltonianArchitecture) - inner_layers = Tuple( - [Dense(arch.width, arch.width, arch.activation) for _ in 1:arch.nhidden] - ) - - Chain( - Dense(arch.dim, arch.width, arch.activation), - inner_layers..., - Linear(arch.width, 1; use_bias = false) - ) -end \ No newline at end of file diff --git a/src/architectures/parametric_resnet.jl b/src/architectures/parametric_resnet.jl new file mode 100644 index 000000000..5652ef4cc --- /dev/null +++ b/src/architectures/parametric_resnet.jl @@ -0,0 +1,37 @@ +struct ParametricResNet{AT <: Activation, PT <: OptionalParameters} <: NeuralNetworkIntegrator + sys_dim::Int + n_blocks::Int + width::Int + parameters::PT + activation::AT + + function ParametricResNet(dim; width=dim, n_blocks = HNN_nhidden_default, activation=HNN_activation_default, parameters=NullParameters()) + activation = (typeof(activation) <: Activation) ? activation : Activation(activation) + new{typeof(activation), typeof(parameters)}(dim, n_blocks, width, parameters, activation) + end +end + +function ParametricResNet(dl::DataLoader, n_blocks::Integer, width::Integer=dl.input_dim; activation=HNN_activation_default, parameters=NullParameters()) + ParametricResNet(dl.input_dim; width=width, n_blocks=n_blocks, activation) +end + +function ResNet(input_dim::Integer, n_blocks::Integer, width::Integer=input_dim; activation=HNN_activation_default, parameters=NullParameters()) + typeof(parameters) <: NullParameters ? ResNet(input_dim, n_blocks, width, activation) : ParametricResNet(input_dim; n_blocks=n_blocks, width=width, parameters=parameters, activation=activation) +end + +function ResNet(input_dim::Integer; n_blocks::Integer, width::Integer=input_dim, activation=HNN_activation_default, parameters=NullParameters()) + ResNet(input_dim, n_blocks, width; activation=activation, parameters=parameters) +end + +function Chain(arch::ParametricResNet{AT}) where AT + layers = () + for _ in 1:arch.n_blocks + # nonlinear layers + layers = (layers..., ParametricResNetLayer(arch.sys_dim, arch.width, arch.activation; parameters=arch.parameters, return_parameters=true)) + end + + # linear layers for the output + layers = (layers..., ParametricResNetLayer(arch.sys_dim, arch.width, identity; parameters=arch.parameters, return_parameters=false)) + + Chain(layers...) +end \ No newline at end of file diff --git a/src/architectures/resnet.jl b/src/architectures/resnet.jl index cc8e98519..bc8829b60 100644 --- a/src/architectures/resnet.jl +++ b/src/architectures/resnet.jl @@ -23,20 +23,23 @@ where `dl` is an instance of `DataLoader`. See [`iterate`](@ref) for an example of this. """ struct ResNet{AT} <: NeuralNetworkIntegrator - sys_dim::Int - n_blocks::Int + sys_dim::Int + n_blocks::Int + width::Int activation::AT end -function ResNet(dl::DataLoader, n_blocks::Integer; activation = tanh) - ResNet(dl.input_dim, n_blocks, activation) +ResNet(sys_dim::Integer, n_blocks::Integer, activation) = ResNet(sys_dim, n_blocks, sys_dim, activation) + +function ResNet(dl::DataLoader, n_blocks::Integer, width::Integer=dl.input_dim; activation = tanh) + ResNet(dl.input_dim, n_blocks, width, activation) end function Chain(arch::ResNet{AT}) where AT layers = () for _ in 1:arch.n_blocks # nonlinear layers - layers = (layers..., ResNetLayer(arch.sys_dim, arch.activation; use_bias=true)) + layers = (layers..., arch.sys_dim == arch.width ? ResNetLayer(arch.sys_dim, arch.activation; use_bias=true) : WideResNetLayer(arch.sys_dim, arch.width, arch.activation)) end # linear layers for the output diff --git a/src/architectures/standard_hamiltonian_neural_network.jl b/src/architectures/standard_hamiltonian_neural_network.jl new file mode 100644 index 000000000..a18b01062 --- /dev/null +++ b/src/architectures/standard_hamiltonian_neural_network.jl @@ -0,0 +1,85 @@ +""" + StandardHamiltonianArchitecture <: HamiltonianArchitecture + +A realization of the standard Hamiltonian neural network (HNN) [greydanus2019hamiltonian](@cite). + +Also see [`GeneralizedHamiltonianArchitecture`](@ref). + +# Constructor + +The constructor takes the following input arguments: +1. `dim`: system dimension, +2. `width = dim`: width of the hidden layer. By default this is equal to `dim`, +3. `nhidden = $(HNN_nhidden_default)`: the number of hidden layers, +4. `activation = $(HNN_activation_default)`: the activation function used in the HNN. +""" +struct StandardHamiltonianArchitecture{AT} <: HamiltonianArchitecture{AT} + dim::Int + width::Int + nhidden::Int + activation::AT + + function StandardHamiltonianArchitecture(dim::Integer, width=dim, + nhidden=HNN_nhidden_default, activation=HNN_activation_default) + @assert iseven(dim) "The input dimension must be an even integer." + new{typeof(activation)}(dim, width, nhidden, activation) + end +end + +""" + symbolic_hamiltonian_vector_field(nn::SymbolicNeuralNetwork) + +Get the symbolic expression for the vector field belonging to the HNN `nn`. + +# Implementation + +This is calling `SymbolicNeuralNetworks.Jacobian` and then multiplies the result with a Poisson tensor. +""" +function symbolic_hamiltonian_vector_field(nn::SymbolicNeuralNetwork) + □ = SymbolicNeuralNetworks.Jacobian(nn) + n = input_dimension(nn.model) ÷ 2 + # The Poisson tensor is built from *integers* on purpose: a `Float64` literal in the symbolic + # expression would widen the result of a `Float32` network. `PoissonTensor`, which has the same + # convention, is an `AbstractMatrix{Float64}` and would do exactly that. + 𝕆 = zeros(Int, n, n) + 𝕀 = Matrix(1I, n, n) + 𝕁 = [𝕆 𝕀; -𝕀 𝕆] + # `Jacobian` uses the convention `□[i, j] = ∂fᵢ/∂xⱼ` and the HNN output is scalar, so the one + # row of `derivative(□)` is the gradient of the Hamiltonian. The vector field is built as a + # *vector*, so that the generated function returns what `HNNLoss` compares against: a vector + # for a single sample, and one column per sample for a batch. + ∇H = vec(derivative(□)) + 𝕁 * ∇H +end + +""" + hamiltonian_vector_field(arch::StandardHamiltonianArchitecture) + +Compute an executable expression of the Hamiltonian vector field of a +[`StandardHamiltonianArchitecture`](@ref). + +# Implementation + +This first computes a symbolic expression of the vector field using [`symbolic_hamiltonian_vector_field`](@ref). + +The function is built with `inplace = false`: [`HNNLoss`](@ref) wraps it and is differentiated with +`Zygote`, and the in-place kernel `SymbolicNeuralNetworks.build_nn_function` builds by default +*mutates* its result, which `Zygote` does not support. +""" +function hamiltonian_vector_field(arch::StandardHamiltonianArchitecture) + nn = SymbolicNeuralNetwork(arch) + hvf = symbolic_hamiltonian_vector_field(nn) + SymbolicNeuralNetworks.build_nn_function(hvf, nn.params, nn.input; inplace = false) +end + +function Chain(arch::StandardHamiltonianArchitecture) + inner_layers = Tuple( + [Dense(arch.width, arch.width, arch.activation) for _ in 1:arch.nhidden] + ) + + Chain( + Dense(arch.dim, arch.width, arch.activation), + inner_layers..., + Linear(arch.width, 1; use_bias = false) + ) +end diff --git a/src/data_loader/batch.jl b/src/data_loader/batch.jl index c2afecbba..2f3e759fa 100644 --- a/src/data_loader/batch.jl +++ b/src/data_loader/batch.jl @@ -149,18 +149,19 @@ function number_of_batches(dl::DataLoader{T, AT, OT, :RegularData}, batch::Batch Int(ceil(dl.input_time_steps * dl.n_params / batch.batch_size)) end -function batch_over_two_axes(batch::Batch, number_columns::Int, third_dim::Int, dl::DataLoader) +function batch_over_two_axes(batch::Batch, number_columns::Integer, third_dim::Integer, n_batches::Integer) time_indices = shuffle(1:number_columns) parameter_indices = shuffle(1:third_dim) complete_indices = Iterators.product(time_indices, parameter_indices) |> collect |> vec batches = () - n_batches = number_of_batches(dl, batch) for batch_number in 1:(n_batches - 1) batches = (batches..., complete_indices[(batch_number - 1) * batch.batch_size + 1 : batch_number * batch.batch_size]) end (batches..., complete_indices[(n_batches - 1) * batch.batch_size + 1:end]) end +batch_over_two_axes(batch::Batch, number_of_columns::Integer, third_dim::Integer, dl::DataLoader) = batch_over_two_axes(batch, number_of_columns, third_dim, number_of_batches(dl, batch)) + function (batch::Batch)(dl::DataLoader{T, BT, OT, :RegularData}) where {T, AT<:AbstractArray{T, 3}, BT<:Union{AT, NamedTuple{(:q, :p), Tuple{AT, AT}}}, OT} batch_over_two_axes(batch, dl.input_time_steps, dl.n_params, dl) end @@ -195,7 +196,7 @@ end output[i, j, k] = data[i, indices[1, k] + seq_length + j - 1, indices[2, k]] end -# this is neeced if we want to use the vector of tuples in a kernel +# this is needed if we want to use the vector of tuples in a kernel function convert_vector_of_tuples_to_matrix(backend::Backend, batch_indices_tuple::Vector{Tuple{Int, Int}}) _batch_size = length(batch_indices_tuple) diff --git a/src/data_loader/optimize.jl b/src/data_loader/optimize.jl index 72af33360..a3a245026 100644 --- a/src/data_loader/optimize.jl +++ b/src/data_loader/optimize.jl @@ -83,8 +83,8 @@ _copy(qp::QPT) = (q = copy(qp.q), p = copy(qp.p)) _copy(t::Tuple{<:QPTOAT, <:QPTOAT}) = _copy.(t) function (o::Optimizer)(nn::NeuralNetwork, - dl::DataLoader, - batch::Batch, + dl::Union{DataLoader, ParametricDataLoader}, + batch::Batch, n_epochs::Integer, loss::NetworkLoss, _pullback::AbstractPullback = ZygotePullback(loss); show_progress = true) diff --git a/src/data_loader/parametric_data_loader.jl b/src/data_loader/parametric_data_loader.jl new file mode 100644 index 000000000..6a520a7c7 --- /dev/null +++ b/src/data_loader/parametric_data_loader.jl @@ -0,0 +1,143 @@ +""" + ParametricDataLoader + +Very similar to [`DataLoader`](@ref), but can deal with parametric problems. +""" +struct ParametricDataLoader{T, AT<:QPTOAT2, VT<:AbstractVector} + input::AT + input_dim::Int + input_time_steps::Int + parameters::VT + n_params::Int + + function ParametricDataLoader(data::QPTOAT2{T, 3}, parameters::AbstractVector) where {T} + input_dim, input_time_steps, n_params = _size(data) + @assert T == _eltype(parameters) "Provided data and parameters must have the same eltype!" + @assert length(parameters) == _size(data, 3) "The number of provided parameters and the parameter axis of the supplied data do not have the same length!" + + new{T, typeof(data), typeof(parameters)}(data, input_dim, input_time_steps, parameters, n_params) + end +end + +function ParametricDataLoader(input::AbstractMatrix{T}, parameters::AbstractVector) where {T} + ParametricDataLoader(reshape(input, size(input)..., 1), parameters) +end + +# The same solution shape `DataLoader(::EnsembleSolution)` takes: since GeometricSolutions 0.6 a +# solution carries the time series and the vector field alongside `q` and `p`. +function ParametricDataLoader(ensemble_solution::EnsembleSolution{T, T1, Vector{ST}}) where {T, + T1, + TuT, + TT <: TimeSeries{T1}, + ST <: GeometricSolution{T, T1, TT, NamedTuple{(:t, :q, :p, :q̇, :ṗ), TuT}} + } + + sys_dim = length(ensemble_solution.s[1].q[0]) + input_time_steps = length(ensemble_solution.t) + n_params = length(ensemble_solution.s) + params = ensemble_solution.problem.parameters + + data = (q = zeros(T, sys_dim, input_time_steps, n_params), p = zeros(T, sys_dim, input_time_steps, n_params)) + + for (solution, i) in zip(ensemble_solution.s, axes(ensemble_solution.s, 1)) + for dim in 1:sys_dim + data.q[dim, :, i] = solution.q[:, dim] + data.p[dim, :, i] = solution.p[:, dim] + end + end + + ParametricDataLoader(data, params) +end + +# """ +# rearrange_parameters(parameters) +# +# Rearrange `parameters` such that they can be used by [`ParametricDataLoader`](@ref). +# """ +# function rearrange_parameters(parameters::Vector{<:NamedTuple}) +# parameters_rearranged = zeros(_eltype(parameters), ) +# end + +# function batch_over_two_axes(batch::Batch, number_columns::Int, third_dim::Int, dl::ParametricDataLoader) +# time_indices = shuffle(1:number_columns) +# parameter_indices = shuffle(1:third_dim) +# complete_indices = Iterators.product(time_indices, parameter_indices) |> collect |> vec +# batches = () +# n_batches = number_of_batches(dl, batch) +# for batch_number in 1:(n_batches - 1) +# batches = (batches..., complete_indices[(batch_number - 1) * batch.batch_size + 1 : batch_number * batch.batch_size]) +# end +# (batches..., complete_indices[(n_batches - 1) * batch.batch_size + 1:end]) +# end + +function optimize_for_one_epoch!( opt::Optimizer, + model, + ps::Union{NeuralNetworkParameters, NamedTuple}, + dl::ParametricDataLoader{T}, + batch::Batch, + _pullback::AbstractPullback, + λY) where T + count = 0 + total_error = T(0) + batches = batch(dl) + for batch_indices in batches + count += 1 + # these `copy`s should not be necessary! coming from a Zygote problem! + _input_nt_output_nt_parameter_indices = convert_input_and_batch_indices_to_array(dl, batch, batch_indices) + # input_nt_output_nt = _input_nt_output_nt_parameter_indices[1:2] + loss_value, pullback = _pullback(ps, model, _input_nt_output_nt_parameter_indices) + total_error += loss_value + dp = _unwrap_gradient(_get_contents(pullback(one(loss_value)))) + optimization_step!(opt, λY, ps, dp) + end + total_error / count +end + +function parameter_indices(parameters::AbstractVector, parameter_indices::AbstractVector{Int}) + [parameters[parameter_index] for parameter_index in parameter_indices] +end + +function parameter_indices(parameters::AbstractVector, batch_indices::AbstractMatrix{Int}) + parameter_indices(parameters, batch_indices[2, :]) +end + +function parameter_indices(dl::ParametricDataLoader, indices::AbstractArray{Int}) + parameter_indices(dl.parameters, indices) +end + +function convert_input_and_batch_indices_to_array(dl::ParametricDataLoader{T, BT}, batch::Batch, batch_indices_tuple::Vector{Tuple{Int, Int}}) where {T, AT<:AbstractArray{T, 3}, BT<:NamedTuple{(:q, :p), Tuple{AT, AT}}} + backend = networkbackend(dl.input.q) + + # the batch size is smaller for the last batch + _batch_size = length(batch_indices_tuple) + + batch_indices = convert_vector_of_tuples_to_matrix(backend, batch_indices_tuple) + + q_input = KernelAbstractions.allocate(backend, T, dl.input_dim ÷ 2, batch.seq_length, _batch_size) + p_input = similar(q_input) + + assign_input_from_vector_of_tuples! = assign_input_from_vector_of_tuples_kernel!(backend) + assign_input_from_vector_of_tuples!(q_input, p_input, dl.input, batch_indices, ndrange=(dl.input_dim ÷ 2, batch.seq_length, _batch_size)) + + q_output = KernelAbstractions.allocate(backend, T, dl.input_dim ÷ 2, batch.prediction_window, _batch_size) + p_output = similar(q_output) + + assign_output_from_vector_of_tuples! = assign_output_from_vector_of_tuples_kernel!(backend) + assign_output_from_vector_of_tuples!(q_output, p_output, dl.input, batch_indices, batch.seq_length, ndrange=(dl.input_dim ÷ 2, batch.prediction_window, _batch_size)) + + (q = q_input, p = p_input), (q = q_output, p = p_output), parameter_indices(dl, batch_indices) +end + +function number_of_batches(dl::ParametricDataLoader, batch::Batch) + @assert dl.input_time_steps ≥ (batch.seq_length + batch.prediction_window) "The number of time steps has to be greater than sequence length + prediction window." + Int(ceil((dl.input_time_steps - (batch.seq_length - 1) - batch.prediction_window) * dl.n_params / batch.batch_size)) +end + +function (batch::Batch)(dl::ParametricDataLoader) + batch_over_two_axes(batch, dl.input_time_steps - (batch.seq_length - 1) - batch.prediction_window, dl.n_params, number_of_batches(dl, batch)) +end + +function (o::Optimizer)(nn::NeuralNetwork{<:GeneralizedHamiltonianArchitecture}, dl::ParametricDataLoader, batch::Batch{:FeedForward}, n_epochs::Integer=1, loss::NetworkLoss=ParametricLoss(); kwargs...) + _pullback::AbstractPullback = ZygotePullback(loss) + o(nn, dl, batch, n_epochs, loss, _pullback; kwargs...) +end \ No newline at end of file diff --git a/src/layers/forcing_dissipation_layers.jl b/src/layers/forcing_dissipation_layers.jl new file mode 100644 index 000000000..cc7ea16cd --- /dev/null +++ b/src/layers/forcing_dissipation_layers.jl @@ -0,0 +1,163 @@ +@doc raw""" + ForcingLayer <: AbstractExplicitLayer + +Layers that can learn dissipative or forcing terms, but not conservative ones. + +Use the constructors [`ForcingLayerQ`](@ref) and [`ForcingLayerP`](@ref) for this. + +!!! warn + The forcing is dependent on either ``q`` or ``p``, but always applied to the ``p`` component. + +The forcing layers are inspired by the Lagrange-d'Alembert integrator from [marsden2001discrete; Example 3.2.2](@cite): + +```math +\begin{aligned} + q^{(t+1)} = & q^{(t)} + & hM^{-1}p^{(t)}, \\ + p^{(t+1)} = & p^{(t)} + & -h\nabla{}U(q^{(t+1)}) + hf_H(q^{(t+1)}, p^{(t)}), +\end{aligned} +``` +for a separable Hamiltonian ``H(q, p) = T(p) + U(q) = p^TM^{-1}p + U(q)`` and external forcing ``f_H.`` +""" +struct ForcingLayer{M,N,PT,CT,type,ReturnParameters} <: AbstractExplicitLayer{M,N} + dim::Int + width::Int + nhidden::Int + parameter_length::Int + parameter_layout::PT + model::CT +end + +parameterlength(l::ForcingLayer) = parameterlength(l.model) + +function initialparameters(rng::Random.AbstractRNG, init_weight::AbstractNeuralNetworks.Initializer, integrator::ForcingLayer, backend::KernelAbstractions.Backend, ::Type{T}) where {T} + initialparameters(rng, init_weight, integrator.model, backend, T) +end + +""" + ForcingLayerQ + +A layer that is derived from the more general [`ForcingLayer`](@ref) and the resulting forcing only depends on the ``q`` component. +""" +const ForcingLayerQ{M,N,FT,AT,ReturnParameters} = ForcingLayer{M,N,FT,AT,:Q,ReturnParameters} + +""" + ForcingLayerP + +A layer that is derived from the more general [`ForcingLayer`](@ref) and the resulting forcing only depends on the ``p`` component. +""" +const ForcingLayerP{M,N,FT,AT,ReturnParameters} = ForcingLayer{M,N,FT,AT,:P,ReturnParameters} + +""" + ForcingLayerQP + +A layer that is derived from the more general [`ForcingLayer`](@ref) and the resulting forcing only depends on the ``q`` and the ``p`` component. +""" +const ForcingLayerQP{M,N,FT,AT,ReturnParameters} = ForcingLayer{M,N,FT,AT,:QP,ReturnParameters} + +function build_chain(dim::Integer, width::Integer, nhidden::Integer, parameter_length::Integer, activation, type::Symbol) + inner_layers = Tuple( + [Dense(width, width, activation) for _ in 1:nhidden] + ) + + Chain( + type == :QP ? Dense(dim + parameter_length, width, activation) : Dense(dim ÷ 2 + parameter_length, width, activation), + inner_layers..., + Linear(width, dim ÷ 2; use_bias=false) + ) +end + +function ForcingLayer(dim::Integer, width::Integer, nhidden::Integer, activation; parameters::OptionalParameters=NullParameters(), return_parameters::Bool, type::Symbol) + flat_parameters, layout = _flatten_system_parameters(parameters) + parameter_length = length(flat_parameters) + c = build_chain(dim, width, nhidden, parameter_length, activation, type) + ForcingLayer{dim,dim,typeof(layout),typeof(c),type,return_parameters}(dim, width, nhidden, parameter_length, layout, c) +end + +""" + ForcingLayerQ(dim) + +# Examples + +```julia +ForcingLayerQ(dim, width, nhidden, activation; parameters, return_parameters) +``` +""" +function ForcingLayerQ(dim::Integer, width::Integer=dim, nhidden::Integer=HNN_nhidden_default, activation=HNN_activation_default; parameters::OptionalParameters=NullParameters(), return_parameters::Bool=false) + ForcingLayer(dim, width, nhidden, activation; parameters=parameters, return_parameters=return_parameters, type=:Q) +end + +""" + ForcingLayerP(dim) + +See [`ForcingLayerQ`](@ref). + +# Examples + +```julia +ForcingLayerP(dim, width, nhidden, activation; parameters, return_parameters) +``` +""" +function ForcingLayerP(dim::Integer, width::Integer=dim, nhidden::Integer=HNN_nhidden_default, activation=HNN_activation_default; parameters::OptionalParameters=NullParameters(), return_parameters::Bool=false) + ForcingLayer(dim, width, nhidden, activation; parameters=parameters, return_parameters=return_parameters, type=:P) +end + +function ForcingLayerQP(dim::Integer, width::Integer=dim, nhidden::Integer=HNN_nhidden_default, activation=HNN_activation_default; parameters::OptionalParameters=NullParameters(), return_parameters::Bool=false) + ForcingLayer(dim, width, nhidden, activation; parameters=parameters, return_parameters=return_parameters, type=:QP) +end + +function (integrator::ForcingLayerQ{M,N,FT,AT,false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} + input = concatenate_array_with_parameters(qp.q, problem_params) + (q=qp.q, p=qp.p + integrator.model(input, params)) +end + +function (integrator::ForcingLayerP{M,N,FT,AT,false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} + input = concatenate_array_with_parameters(qp.p, problem_params) + (q=qp.q, p=qp.p + integrator.model(input, params)) +end + +function (integrator::ForcingLayerQP{M,N,FT,AT,false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} + input = concatenate_array_with_parameters(vcat(qp.q, qp.p), problem_params) + (q=qp.q, p=qp.p + integrator.model(input, params)) +end + +function (integrator::ForcingLayerQ{M,N,FT,AT,true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} + input = concatenate_array_with_parameters(qp.q, problem_params) + ((q=qp.q, p=qp.p + integrator.model(input, params)), problem_params) +end + +function (integrator::ForcingLayerP{M,N,FT,AT,true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} + input = concatenate_array_with_parameters(qp.p, problem_params) + ((q=qp.q, p=qp.p + integrator.model(input, params)), problem_params) +end + +function (integrator::ForcingLayerQP{M,N,FT,AT,true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} + input = concatenate_array_with_parameters(vcat(qp.q, qp.p), problem_params) + ((q=qp.q, p=qp.p + integrator.model(input, params)), problem_params) +end + +function (integrator::ForcingLayer)(qp_params::Tuple{<:QPTOAT2,<:OptionalParameters}, params::NeuralNetworkParameters) + integrator(qp_params..., params) +end + +function (integrator::ForcingLayer)(::TT, ::NeuralNetworkParameters) where {TT<:Tuple} + error("The input is of type $(TT). This shouldn't be the case!") +end + +function (integrator::ForcingLayer{M,N,FT,AT,Type,true})(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT,Type} + @assert iseven(size(qp, 1)) + n = size(qp, 1) ÷ 2 + qp_split = assign_q_and_p(qp, n) + evaluated = integrator(qp_split, problem_params, params)[1] + (vcat(evaluated.q, evaluated.p), problem_params) +end + +function (integrator::ForcingLayer{M,N,FT,AT,Type,false})(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT,Type} + @assert iseven(size(qp, 1)) + n = size(qp, 1) ÷ 2 + qp_split = assign_q_and_p(qp, n) + evaluated = integrator(qp_split, problem_params, params) + vcat(evaluated.q, evaluated.p) +end + +(integrator::ForcingLayer)(qp::QPTOAT2, params::NeuralNetworkParameters) = integrator(qp, NullParameters(), params) +(integrator::ForcingLayer)(qp::QPTOAT2, params::NamedTuple) = integrator(qp, NeuralNetworkParameters(params)) diff --git a/src/layers/parametric_resnet_layer.jl b/src/layers/parametric_resnet_layer.jl new file mode 100644 index 000000000..149d6c4e0 --- /dev/null +++ b/src/layers/parametric_resnet_layer.jl @@ -0,0 +1,64 @@ +struct ParametricResNetLayer{M, N, F1 <: Activation, PT, ReturnParameters} <: AbstractExplicitLayer{M, N} + width::Int + activation::F1 + parameter_length::Int + parameter_layout::PT +end + +function ParametricResNetLayer(dim::Integer, width::Integer, activation=identity; parameters::OptionalParameters=NullParameters(), return_parameters::Bool) + flat_parameters, layout = _flatten_system_parameters(parameters) + _activation = Activation(activation) + ParametricResNetLayer{dim, dim, typeof(_activation), typeof(layout), return_parameters}(width, _activation, length(flat_parameters), layout) +end + +function initialparameters(rng::Random.AbstractRNG, init_weight::AbstractNeuralNetworks.Initializer, l::ParametricResNetLayer{M, M}, backend::KernelAbstractions.Backend, ::Type{T}; init_bias = ZeroInitializer()) where {M, T} + upscale_weight = KernelAbstractions.allocate(backend, T, l.width, M + l.parameter_length) + upscale_bias = KernelAbstractions.allocate(backend, T, l.width) + downscale_weight = KernelAbstractions.allocate(backend, T, M, l.width) + bias = KernelAbstractions.allocate(backend, T, M) + init_weight(rng, upscale_weight) + init_weight(rng, downscale_weight) + init_bias(rng, upscale_bias) + init_bias(rng, bias) + (upscale_weight=upscale_weight, downscale_weight=downscale_weight, upscale_bias=upscale_bias, bias=bias) +end + +parameterlength(l::ParametricResNetLayer{M, M}) where {M} = (l.width + l.parameter_length) * (M + 1) + M * (l.width + 1) + +function (d::ParametricResNetLayer{M, M, F, PT, false})(x::AbstractVecOrMat, problem_params::OptionalParameters, ps::NamedTuple) where {M, F, PT} + input = concatenate_array_with_parameters(x, problem_params) + x + d.activation.(ps.downscale_weight * d.activation.(ps.upscale_weight * input .+ ps.upscale_bias) .+ ps.bias) +end + +function (d::ParametricResNetLayer{M, M, F, PT, true})(x::AbstractVecOrMat, problem_params::OptionalParameters, ps::NamedTuple) where {M, F, PT} + input = concatenate_array_with_parameters(x, problem_params) + (x + d.activation.(ps.downscale_weight * d.activation.(ps.upscale_weight * input .+ ps.upscale_bias) .+ ps.bias), problem_params) +end + +# function (d::ParametricResNetLayer{M, M, F, PT, false})(x::AbstractArray{T, 3}, problem_params::AbstractVector, ps::NamedTuple) where {M, F, PT, T} +# input = concatenate_array_with_parameters(x, problem_params) +# x + d.activation.(mat_tensor_mul(ps.downscale_weight, d.activation.(mat_tensor_mul(ps.upscale_weight, x) .+ ps.upscale_bias)) .+ ps.bias) +# end +# +# function (d::ParametricResNetLayer{M, M, F, PT, true})(x::AbstractArray{T, 3}, problem_params::AbstractVector, ps::NamedTuple) where {M, F, PT, T} +# input = concatenate_array_with_parameters(x, problem_params) +# (x + d.activation.(mat_tensor_mul(ps.downscale_weight, d.activation.(mat_tensor_mul(ps.upscale_weight, x) .+ ps.upscale_bias)) .+ ps.bias), problem_params) +# end + +(d::ParametricResNetLayer)(input::Tuple, ps::NamedTuple) = length(input) == 2 ? d(input..., ps) : error("The tuple must contain the input array/nt as well as the system parameters.") + +function (d::ParametricResNetLayer{M, M, F, PT, false})(z::QPT, problem_params::OptionalParameters, ps::NamedTuple) where {M, F, PT} + @assert iseven(M) + @assert size(z.q, 1) * 2 == M + N2 = M ÷ 2 + output = d(vcat(z.q, z.p), problem_params, ps) + assign_q_and_p(output, N2) +end + +function (d::ParametricResNetLayer{M, M, F, PT, true})(z::QPT, problem_params::OptionalParameters, ps::NamedTuple) where {M, F, PT} + @assert iseven(M) + @assert size(z.q, 1) * 2 == M + N2 = M ÷ 2 + output = d(vcat(z.q, z.p), problem_params, ps) + (assign_q_and_p(output[1], N2), problem_params) +end \ No newline at end of file diff --git a/src/layers/sympnets.jl b/src/layers/sympnets.jl index 5122467a3..cecbf49cd 100644 --- a/src/layers/sympnets.jl +++ b/src/layers/sympnets.jl @@ -236,36 +236,36 @@ function custom_vec_mul(scale::AbstractVector{T}, x::AbstractArray{T, 3}) where vec_tensor_mul(scale, x) end -@inline function (d::ActivationLayerQ{M, M})(x::NamedTuple, ps) where {M} +@inline function (d::ActivationLayerQ{M, M})(x::QPT2, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") return (q = x.q + custom_vec_mul(ps.scale, d.activation.(x.p)), p = x.p) end -@inline function (d::ActivationLayerP{M, M})(x::NamedTuple, ps) where {M} +@inline function (d::ActivationLayerP{M, M})(x::QPT2, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") return (q = x.q, p = x.p + custom_vec_mul(ps.scale, d.activation.(x.q))) end -@inline function (d::GradientLayerQ{M, M})(x::NamedTuple, ps) where {M} +@inline function (d::GradientLayerQ{M, M})(x::QPT2, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") (q = x.q + custom_mat_mul(ps.weight', (custom_vec_mul(ps.scale, d.activation.(custom_mat_mul(ps.weight, x.p) .+ ps.bias)))), p = x.p) end -@inline function(d::GradientLayerP{M, M})(x::NamedTuple, ps) where {M} +@inline function(d::GradientLayerP{M, M})(x::QPT2, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") (q = x.q, p = x.p + custom_mat_mul(ps.weight', (custom_vec_mul(ps.scale, d.activation.(custom_mat_mul(ps.weight, x.q) .+ ps.bias))))) end -@inline function(d::LinearLayerQ{M, M})(x::NamedTuple, ps) where {M} +@inline function(d::LinearLayerQ{M, M})(x::QPT2, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") (q = x.q + custom_mat_mul(ps.weight, x.p), p = x.p) end -@inline function(d::LinearLayerP{M, M})(x::NamedTuple, ps) where {M} +@inline function(d::LinearLayerP{M, M})(x::QPT2, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") (q = x.q, p = x.p + custom_mat_mul(ps.weight, x.q)) end diff --git a/src/layers/wide_resnet.jl b/src/layers/wide_resnet.jl new file mode 100644 index 000000000..f1a4c4682 --- /dev/null +++ b/src/layers/wide_resnet.jl @@ -0,0 +1,32 @@ +struct WideResNetLayer{M, N, F1} <: AbstractExplicitLayer{M, N} + width::Int + activation::F1 +end + +WideResNetLayer(dim::Integer, width::Integer, activation=identity) = WideResNetLayer{dim, dim, typeof(activation)}(width, activation) + +function initialparameters(rng::Random.AbstractRNG, init_weight::AbstractNeuralNetworks.Initializer, l::WideResNetLayer{M, M}, backend::KernelAbstractions.Backend, ::Type{T}; init_bias = ZeroInitializer()) where {M, T} + upscale_weight = KernelAbstractions.allocate(backend, T, l.width, M) + upscale_bias = KernelAbstractions.allocate(backend, T, l.width) + downscale_weight = KernelAbstractions.allocate(backend, T, M, l.width) + bias = KernelAbstractions.allocate(backend, T, M) + init_weight(rng, upscale_weight) + init_weight(rng, downscale_weight) + init_bias(rng, upscale_bias) + init_bias(rng, bias) + (upscale_weight=upscale_weight, downscale_weight=downscale_weight, upscale_bias=upscale_bias, bias=bias) +end + +parameterlength(l::WideResNetLayer{M, M}) where {M} = l.width * (M + 1) + M * (l.width + 1) + +(d::WideResNetLayer{M, M})(x::AbstractVecOrMat, ps::NamedTuple) where {M} = x + d.activation.(ps.downscale_weight * d.activation.(ps.upscale_weight * x .+ ps.upscale_bias) .+ ps.bias) + +(d::WideResNetLayer{M, M})(x::AbstractArray{T, 3}, ps::NamedTuple) where {M, T} = x + d.activation.(mat_tensor_mul(ps.downscale_weight, d.activation.(mat_tensor_mul(ps.upscale_weight, x) .+ ps.upscale_bias)) .+ ps.bias) + +function (d::WideResNetLayer{M, M})(z::QPT, ps::NamedTuple) where {M} + @assert iseven(M) + @assert size(z.q, 1) * 2 == M + N2 = M ÷ 2 + output = d(vcat(z.q, z.p), ps) + assign_q_and_p(output, N2) +end \ No newline at end of file diff --git a/src/loss/losses.jl b/src/loss/losses.jl index 6c97ab8ff..e6d10ccba 100644 --- a/src/loss/losses.jl +++ b/src/loss/losses.jl @@ -257,3 +257,25 @@ function (loss::ReducedLoss)(model::Chain, params::NeuralNetworkParameters, input::CT, output::CT) where {CT <: QPTOAT} _compute_loss(loss.decoder(model(loss.encoder(input), params)), output) end + +@doc raw""" + ParametricLoss() + +The loss for a network whose forward pass takes the parameters of the system alongside the input, +i.e. the parameter-dependent architectures built on [`GeneralizedHamiltonianArchitecture`](@ref). + +It is `FeedForwardLoss` with the system parameters threaded through: + +```math +L(\mathtt{input}, \mathtt{output}, \mu) = ||\mathcal{NN}(\mathtt{input}, \mu) - \mathtt{output}||. +``` + +This loss does not have any parameters. +""" +struct ParametricLoss <: NetworkLoss end + +function (loss::ParametricLoss)(model::Chain, + params::Union{NamedTuple, NeuralNetworkParameters}, input::CT, output::CT, + system_parameters::Union{NamedTuple, AbstractVector}) where {CT <: QPTOAT} + _compute_loss(model(input, system_parameters, params), output) +end diff --git a/src/pullbacks/symbolic_hnn_pullback.jl b/src/pullbacks/symbolic_hnn_pullback.jl index c5168be93..a812d4be9 100644 --- a/src/pullbacks/symbolic_hnn_pullback.jl +++ b/src/pullbacks/symbolic_hnn_pullback.jl @@ -25,3 +25,72 @@ function SymbolicPullback(arch::HamiltonianArchitecture) soutput; reduce = +) SymbolicPullback(loss, SymbolicNeuralNetworks.ParameterGradient(gradient_function)) end + +@doc raw""" + SymbolicPullback(nn, loss, system_params) + +The `SymbolicPullback` for a network whose forward pass also takes the parameters of the *system*, +i.e. one built on [`GeneralizedHamiltonianArchitecture`](@ref), with a [`ParametricLoss`](@ref). + +# Implementation + +This is `SymbolicNeuralNetworks.SymbolicPullback(nn, loss)` with the system parameters threaded +through. `build_nn_function` generates a function of *one* input array, so the flattened system +parameters are appended to the network input, and the symbolic expression splits them off again with +[`_flatten_system_parameters`](@ref) and `unflatten`. The numeric side does the same concatenation, +in the call operators below. + +`reduce = +`: the loss of a batch is the sum of the losses of its samples, so its gradient is the +sum of the per-sample gradients. +""" +function SymbolicPullback(nn::NeuralNetwork, loss::ParametricLoss, + system_params::OptionalParameters; cse::Bool = true, inplace::Bool = true) + symbolic_system_parameters = SymbolicNeuralNetworks.symbolic_variables(system_params, :S) + symbolic_network_parameters = SymbolicNeuralNetworks.symbolic_variables(params(nn), :W) + + input_dim = input_dimension(nn.model) + _, parameter_layout = _flatten_system_parameters(SymbolicNeuralNetworks.Symbolics.Num, + symbolic_system_parameters) + sinput = Symbolics.variables(:x, 1:(input_dim + length(system_params))) + soutput = Symbolics.variables(:y, 1:output_dimension(nn.model)) + symbolic_system_input = unflatten(parameter_layout, sinput[(input_dim + 1):end]) + + symbolic_loss = loss(nn.model, symbolic_network_parameters, sinput[1:input_dim], soutput, + symbolic_system_input) + differentials = SymbolicNeuralNetworks.symbolic_differentials(symbolic_network_parameters) + gradient = SymbolicNeuralNetworks.symbolic_derivative(symbolic_loss, differentials) + gradient_function = SymbolicNeuralNetworks.build_nn_function( + gradient, symbolic_network_parameters, sinput, soutput; + reduce = +, cse = cse, inplace = inplace) + SymbolicPullback(loss, SymbolicNeuralNetworks.ParameterGradient(gradient_function)) +end + +# TODO: type piracy -- `SymbolicPullback` is `SymbolicNeuralNetworks`', and so is every argument +# type here. These belong upstream, together with a `build_nn_function` that takes more than one +# data argument, which is what would make the concatenation below unnecessary. +# +# The generated pullback takes *one* input array, so the system parameters are appended to the +# network input before it is called; the loss, which knows about them, gets them separately. +function (_pullback::SymbolicPullback)(ps, model, + input_output_params::Tuple{<:AbstractMatrix, <:AbstractMatrix, + <:Union{NamedTuple, AbstractVector}})::Tuple + input, output, system_params = input_output_params + _pullback.loss(model, ps, input, output, system_params), + _pullback.fun(concatenate_array_with_parameters(input, system_params), output, ps) +end + +# A batch with a time axis: the network is applied sample-wise, so the time and parameter axes are +# folded into one before the pullback sees them. +function (_pullback::SymbolicPullback)(ps, model, + input_output_params::Tuple{AT, AT, <:Union{NamedTuple, AbstractVector}})::Tuple where {T, AT <: AbstractArray{T, 3}} + input, output, system_params = input_output_params + _input = reshape(input, size(input, 1), size(input, 2) * size(input, 3)) + _output = reshape(output, size(output, 1), size(output, 2) * size(output, 3)) + _pullback(ps, model, (_input, _output, system_params)) +end + +function (_pullback::SymbolicPullback)(ps, model, + input_output_params::Tuple{<:QPT, <:QPT, <:Union{NamedTuple, AbstractVector}})::Tuple + input, output, system_params = input_output_params + _pullback(ps, model, (vcat(input.q, input.p), vcat(output.q, output.p), system_params)) +end diff --git a/src/pullbacks/zygote_pullback.jl b/src/pullbacks/zygote_pullback.jl index dad6054c5..6aa2164b2 100644 --- a/src/pullbacks/zygote_pullback.jl +++ b/src/pullbacks/zygote_pullback.jl @@ -31,6 +31,12 @@ end ps -> _pullback.loss(model, ps, input_nt), ps) (_pullback::ZygotePullback)(ps, model, input_nt_output_nt::Tuple{<:QPTOAT, <:QPTOAT})::Tuple = Zygote.pullback( ps -> _pullback.loss(model, ps, input_nt_output_nt...), ps) +# The parameter-dependent architectures take the system parameters as a third element of the +# input tuple, either as a `NamedTuple` of parameters or as one vector entry per sample. +(_pullback::ZygotePullback)(ps, model, input_output_params::Tuple{<:QPTOAT, <:QPTOAT, <:NamedTuple})::Tuple = Zygote.pullback( + ps -> _pullback.loss(model, ps, input_output_params...), ps) +(_pullback::ZygotePullback)(ps, model, input_output_params::Tuple{<:QPTOAT, <:QPTOAT, <:AbstractVector})::Tuple = Zygote.pullback( + ps -> _pullback.loss(model, ps, input_output_params...), ps) """ _get_contents(returned_pullback) diff --git a/src/training_method/symplectic_euler.jl b/src/training_method/symplectic_euler.jl index fab5de86b..ea06c5b38 100644 --- a/src/training_method/symplectic_euler.jl +++ b/src/training_method/symplectic_euler.jl @@ -1,27 +1,27 @@ -abstract type SymplecticEuler <: HnnTrainingMethod end +abstract type SymplecticEulerIntegrator <: HnnTrainingMethod end -struct SymplecticEulerA <: SymplecticEuler end -struct SymplecticEulerB <: SymplecticEuler end +struct SymplecticEulerIntegratorA <: SymplecticEulerIntegrator end +struct SymplecticEulerIntegratorB <: SymplecticEulerIntegrator end SEuler(;sqdist = sqeuclidean) = SEulerA(sqdist = sqdist) -SEulerA(;sqdist = sqeuclidean) = TrainingMethod{SymplecticEulerA, PhaseSpaceSymbol, TrajectoryData, typeof(sqdist)}(sqdist) -SEulerB(;sqdist = sqeuclidean) = TrainingMethod{SymplecticEulerB, PhaseSpaceSymbol, TrajectoryData, typeof(sqdist)}(sqdist) +SEulerA(;sqdist = sqeuclidean) = TrainingMethod{SymplecticEulerIntegratorA, PhaseSpaceSymbol, TrajectoryData, typeof(sqdist)}(sqdist) +SEulerB(;sqdist = sqeuclidean) = TrainingMethod{SymplecticEulerIntegratorB, PhaseSpaceSymbol, TrajectoryData, typeof(sqdist)}(sqdist) -function loss_single(::TrainingMethod{SymplecticEulerA}, nn::AbstractNeuralNetwork{<:HamiltonianArchitecture}, qₙ, qₙ₊₁, pₙ, pₙ₊₁, Δt, params = params(nn)) +function loss_single(::TrainingMethod{SymplecticEulerIntegratorA}, nn::AbstractNeuralNetwork{<:HamiltonianArchitecture}, qₙ, qₙ₊₁, pₙ, pₙ₊₁, Δt, params = params(nn)) dH = vectorfield(nn, [qₙ₊₁...,pₙ...], params) sqeuclidean(dH[1],(qₙ₊₁-qₙ)/Δt) + sqeuclidean(dH[2],(pₙ₊₁-pₙ)/Δt) end -function loss_single(::TrainingMethod{SymplecticEulerB}, nn::AbstractNeuralNetwork{<:HamiltonianArchitecture}, qₙ, qₙ₊₁, pₙ, pₙ₊₁, Δt, params = params(nn)) +function loss_single(::TrainingMethod{SymplecticEulerIntegratorB}, nn::AbstractNeuralNetwork{<:HamiltonianArchitecture}, qₙ, qₙ₊₁, pₙ, pₙ₊₁, Δt, params = params(nn)) dH = vectorfield(nn, [qₙ...,pₙ₊₁...], params) sqeuclidean(dH[1],(qₙ₊₁-qₙ)/Δt) + sqeuclidean(dH[2],(pₙ₊₁-pₙ)/Δt) end -get_loss(::TrainingMethod{<:SymplecticEuler}, ::AbstractNeuralNetwork{<:HamiltonianArchitecture}, data::TrainingData{<:DataSymbol{<:PhaseSpaceSymbol}}, args) = (get_data(data,:q, args...), get_data(data,:q, next(args...)...), get_data(data,:p, args...), get_data(data,:p,next(args...)...), get_Δt(data)) +get_loss(::TrainingMethod{<:SymplecticEulerIntegrator}, ::AbstractNeuralNetwork{<:HamiltonianArchitecture}, data::TrainingData{<:DataSymbol{<:PhaseSpaceSymbol}}, args) = (get_data(data,:q, args...), get_data(data,:q, next(args...)...), get_data(data,:p, args...), get_data(data,:p,next(args...)...), get_Δt(data)) -loss(ti::TrainingMethod{<:SymplecticEuler}, nn::AbstractNeuralNetwork{<:HamiltonianArchitecture}, data::TrainingData{<:DataSymbol{<:PhaseSpaceSymbol}}, index_batch = eachindex(ti, data), params = params(nn)) = +loss(ti::TrainingMethod{<:SymplecticEulerIntegrator}, nn::AbstractNeuralNetwork{<:HamiltonianArchitecture}, data::TrainingData{<:DataSymbol{<:PhaseSpaceSymbol}}, index_batch = eachindex(ti, data), params = params(nn)) = mapreduce(args->loss_single(Zygote.ignore_derivatives(ti), nn, get_loss(ti, nn, data, args)..., params),+, index_batch) -min_length_batch(::SymplecticEuler) = 2 \ No newline at end of file +min_length_batch(::SymplecticEulerIntegrator) = 2 \ No newline at end of file diff --git a/src/utils.jl b/src/utils.jl index 9fa719360..e2e8782fb 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -151,7 +151,16 @@ qp = (q = [1, 2], p = [3, 4]) ``` """ -const QPT{T} = NamedTuple{(:q, :p), Tuple{AT, AT}} where {T, AT <: AbstractArray{T}} +const QPT{T} = NamedTuple{(:q, :p), Tuple{AT, AT}} where {T, N, AT <: AbstractArray{T, N}} + +@doc raw""" + QPT2 + +[`QPT`](@ref) with the number of dimensions of the two arrays fixed, but their types allowed to +differ. A `Chain` that splits an input array into `q` and `p` produces views of different types, so +the layers of a parameter-dependent network dispatch on this rather than on `QPT`. +""" +const QPT2{T, N} = NamedTuple{(:q, :p), Tuple{AT₁, AT₂}} where {T, N, AT₁ <: AbstractArray{T, N}, AT₂ <: AbstractArray{T, N}} @doc raw""" QPTOAT @@ -165,9 +174,82 @@ This could be data in ``(q, p)\in\mathbb{R}^{2d}`` form or come from an arbitrar """ const QPTOAT{T} = Union{QPT{T}, AbstractArray{T}} where {T} +@doc raw""" + QPTOAT2 + +[`QPTOAT`](@ref) with the number of dimensions of the arrays fixed: + +```julia +const QPTOAT2 = Union{QPT2, AbstractArray} +``` +""" +const QPTOAT2{T, N} = Union{QPT2{T, N}, AbstractArray{T, N}} where {T, N} + Base.:≈(qp₁::QPT, qp₂::QPT) = (qp₁.q ≈ qp₂.q) & (qp₁.p ≈ qp₂.p) +@doc raw""" + _flatten_system_parameters(parameters) + _flatten_system_parameters(T, parameters) + +Flatten the parameters of the *system* (not of the network) into a vector, together with the +`NeuralNetworkParameters.ParameterLayout` that puts such a vector back into the original shape. + +The parameter-dependent architectures — [`GeneralizedHamiltonianArchitecture`](@ref) and the layers +it is built from — feed the system parameters to the network as extra input components, so they have +to be a vector. `NullParameters` flattens to an empty one, which makes the parameter-free case fall +out of the same code path. + +The layout is a *value*, not a closure, so a layer can store it in a field and stay inferable. +""" +_flatten_system_parameters(parameters::NamedTuple) = flatten(parameters) +_flatten_system_parameters(::NullParameters) = flatten(NamedTuple()) +_flatten_system_parameters(::Type{T}, parameters::NamedTuple) where {T} = flatten(T, parameters) +_flatten_system_parameters(::Type{T}, ::NullParameters) where {T} = flatten(T, NamedTuple()) + +""" + _unwrap_gradient(dp) + +Strip the `NeuralNetworkParameters` wrappers and the `(params = …,)` layers out of a gradient, so +that it has the same shape as the parameters it belongs to. + +`Zygote` differentiates *through* the `NeuralNetworkParameters` struct, so the gradient of a +parameter set comes back as a `NamedTuple` with a single `params` field. [`_get_params`](@ref) undoes +that at the top level. The parameter-dependent architectures nest — a `SymplecticEuler` layer +holds the parameters of a whole sub-network — so the unwrapping has to recurse. +""" +_unwrap_gradient(dp) = dp +_unwrap_gradient(dp::NeuralNetworkParameters) = _unwrap_gradient(params(dp)) +_unwrap_gradient(dp::NamedTuple{(:params,)}) = _unwrap_gradient(dp.params) +_unwrap_gradient(dp::NamedTuple) = map(_unwrap_gradient, dp) + _eltype(x) = eltype(x) _eltype(ps::NamedTuple) = _eltype(ps[1]) _eltype(ps::Tuple) = _eltype(ps[1]) _eltype(ps::NeuralNetworkParameters) = _eltype(params(ps)[1]) + +# `ParametricDataLoader` stores one `NamedTuple` of system parameters per trajectory, and they all +# have to agree with the element type of the data. +function _eltype(parameters::AbstractVector{<:NamedTuple}) + T = _eltype(first(parameters)) + for p in parameters + _eltype(p) == T || error("The parameters do not all have the same element type.") + end + T +end + +# `size` that also works on `(q, p)` data, where the first axis is the concatenation of the two. +_size(x) = size(x) +function _size(qp::QPT) + q_size = _size(qp.q) + p_size = _size(qp.p) + @assert q_size == p_size + (2q_size[1], q_size[2:end]...) +end + +_size(x, a::Integer) = size(x, a) +function _size(qp::QPT, a::Integer) + q_size = _size(qp.q, a) + p_size = _size(qp.p, a) + @assert q_size == p_size + a == 1 ? 2q_size : q_size +end diff --git a/test/data_loader/parametric_data_loader_test.jl b/test/data_loader/parametric_data_loader_test.jl new file mode 100644 index 000000000..12c9caa3d --- /dev/null +++ b/test/data_loader/parametric_data_loader_test.jl @@ -0,0 +1,45 @@ +using GeometricMachineLearning +using GeometricMachineLearning: convert_input_and_batch_indices_to_array +using Test +using GeometricProblems.CoupledHarmonicOscillator: hodeensemble, default_parameters +using GeometricIntegrators: ImplicitMidpoint, integrate +using Random: seed! +seed!(123) + +function make_alternative_parameters_by_adding_constant(params::NamedTuple = default_parameters(), + a::Number = 1.) + NamedTuple{keys(params)}(Tuple(value .+ a for value in values(params))) +end + +all_parameters = [default_parameters(), make_alternative_parameters_by_adding_constant()] + +h_ensemble = hodeensemble(; parameters = all_parameters) +sol = integrate(h_ensemble, ImplicitMidpoint()) +dl = ParametricDataLoader(sol) +batch = Batch(2) +batch_indices = batch(dl) + +# Each entry of a batch is a `(time index, parameter index)` pair, and the third element of what +# `convert_input_and_batch_indices_to_array` returns has to be the parameters of *that* trajectory. +# The batches are shuffled, so this asserts the correspondence rather than which batch holds which +# parameters -- pinning the latter makes the test depend on the RNG stream of the Julia version. +function batch_is_consistent(n::Integer) + input, output, parameters = convert_input_and_batch_indices_to_array(dl, batch, batch_indices[n]) + all(enumerate(batch_indices[n])) do (k, (time_index, parameter_index)) + parameters[k] == all_parameters[parameter_index] && + input.q[:, 1, k] == dl.input.q[:, time_index, parameter_index] && + input.p[:, 1, k] == dl.input.p[:, time_index, parameter_index] && + output.q[:, 1, k] == dl.input.q[:, time_index + 1, parameter_index] && + output.p[:, 1, k] == dl.input.p[:, time_index + 1, parameter_index] + end +end + +@test all(batch_is_consistent, eachindex(batch_indices)) + +# Both parameter sets have to turn up somewhere, or the assertion above would also pass on a data +# loader that always returned the first one. +returned_parameters = Set(parameters + for n in eachindex(batch_indices) + for parameters in last(convert_input_and_batch_indices_to_array( + dl, batch, batch_indices[n]))) +@test returned_parameters == Set(all_parameters) diff --git a/test/generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl b/test/generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl new file mode 100644 index 000000000..a4c52604e --- /dev/null +++ b/test/generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl @@ -0,0 +1,63 @@ +# The symbolic pullback of a parameter-dependent network, on the smallest such network there is: a +# single `SymplecticEulerB` layer built from a `SymbolicPotentialEnergy`. +# +# `SymbolicPullback(nn, ::ParametricLoss, system_params)` builds its gradient with `reduce = +`, so +# what it returns is the *sum* of the per-sample gradients -- that is the convention +# `SymbolicNeuralNetworks.SymbolicPullback` uses too. The test compares against exactly that, +# computed with `Zygote` one sample at a time. + +using GeometricMachineLearning +using GeometricMachineLearning: ParametricLoss, SymbolicPotentialEnergy, SymplecticEulerB +using AbstractNeuralNetworks: params +using Random: seed! +using Test +import Zygote + +seed!(1234) + +system_parameters = (m = 1.0, ω = π / 2) +dim, width, nhidden, activation = 2, 2, 1, tanh +n_samples = 10 + +se = SymbolicPotentialEnergy(dim, width, nhidden, activation; parameters = system_parameters) +nn = NeuralNetwork(Chain(SymplecticEulerB(se; return_parameters = false))) + +loss = ParametricLoss() +pullback = SymbolicPullback(nn, loss, system_parameters) + +input = rand(dim, n_samples) +output = rand(dim, n_samples) +# one parameter set per sample, which is the shape `ParametricDataLoader` hands to the optimizer +batch_parameters = fill(system_parameters, n_samples) + +loss_value, gradient = pullback(params(nn), nn.model, (input, output, batch_parameters)) + +@test loss_value ≈ loss(nn.model, params(nn), input, output, batch_parameters) + +symbolic_gradient = gradient(1.0) + +function summed_per_sample_gradient() + total = nothing + for i in 1:n_samples + single = Zygote.gradient( + ps -> loss(nn.model, ps, input[:, i:i], output[:, i:i], [system_parameters]), + params(nn))[1] + block = single.L1.params + total = isnothing(total) ? block : map((a, b) -> map(+, a, b), total, block) + end + total +end + +reference_gradient = summed_per_sample_gradient() + +@test keys(symbolic_gradient) == (:L1,) +for layer in keys(reference_gradient) + for parameter in keys(reference_gradient[layer]) + @test symbolic_gradient.L1[layer][parameter] ≈ reference_gradient[layer][parameter] + end +end + +# A gradient of all zeros would pass the loop above if the reference were zero too, so check that +# the network actually depends on its parameters here. +@test any(any(abs.(block) .> 1e-8) for layer in keys(reference_gradient) + for block in values(reference_gradient[layer])) diff --git a/test/generalized_hamiltonian_neural_networks/pghnn_training_test.jl b/test/generalized_hamiltonian_neural_networks/pghnn_training_test.jl new file mode 100644 index 000000000..e1b4c4200 --- /dev/null +++ b/test/generalized_hamiltonian_neural_networks/pghnn_training_test.jl @@ -0,0 +1,47 @@ +# One epoch of training a `GeneralizedHamiltonianArchitecture` on a `ParametricDataLoader`, which is +# the path that ties the pieces together: the batch splitter, the parametric loss, the `Zygote` +# pullback through the symbolic gradient of the energies, and the optimizer step over a *nested* +# parameter set. + +using GeometricMachineLearning +using AbstractNeuralNetworks: params +using GeometricProblems.CoupledHarmonicOscillator: hodeensemble, default_parameters +using GeometricIntegrators: ImplicitMidpoint, integrate +using Random: seed! +using Test + +seed!(1234) + +function shift_parameters(params::NamedTuple, a::Number) + NamedTuple{keys(params)}(Tuple(value .+ a for value in values(params))) +end + +all_parameters = [default_parameters(), shift_parameters(default_parameters(), 0.5)] + +sol = integrate(hodeensemble(; parameters = all_parameters), ImplicitMidpoint()) +dl = ParametricDataLoader(sol) + +arch = GeneralizedHamiltonianArchitecture(dl.input_dim; parameters = default_parameters()) +nn = NeuralNetwork(arch) +parameters_before = deepcopy(params(nn)) + +n_epochs = 2 +loss_array = Optimizer(AdamOptimizer(), nn)(nn, dl, Batch(200), n_epochs; show_progress = false) + +@test length(loss_array) == n_epochs +@test all(isfinite, loss_array) +@test all(>(0), loss_array) + +# The optimizer has to reach every block of the *nested* parameter set: the architecture is a chain +# of `SymplecticEuler` layers, each of which holds the parameters of a whole sub-network. +function every_block_moved(before, after) + all(keys(before)) do layer + all(keys(before[layer])) do sublayer + all(keys(before[layer][sublayer])) do parameter + before[layer][sublayer][parameter] != after[layer][sublayer][parameter] + end + end + end +end + +@test every_block_moved(parameters_before, params(nn)) diff --git a/test/generalized_hamiltonian_neural_networks_test.jl b/test/generalized_hamiltonian_neural_networks_test.jl new file mode 100644 index 000000000..824ded805 --- /dev/null +++ b/test/generalized_hamiltonian_neural_networks_test.jl @@ -0,0 +1,25 @@ +using GeometricMachineLearning +using GeometricMachineLearning: OptionalParameters, OneInitializer +using GeometricProblems.HarmonicOscillator: odeproblem, default_parameters +using GeometricIntegrators +using Test + +sol = integrate(odeproblem(), ImplicitMidpoint()) +dim = length(sol.problem.ics.q) + +dl = DataLoader(sol) + +function test_ghnn_without_parameters(dim::Integer = dim) + arch = GeneralizedHamiltonianArchitecture(dim) + nn = NeuralNetwork(arch; initializer=OneInitializer()) + @test nn([1., 1.]) ≈ [1.003217200759985, 0.9968055760434815] +end + +function test_ghnn_with_parameters(dim::Integer = dim, parameters::OptionalParameters = default_parameters()) + arch = GeneralizedHamiltonianArchitecture(dim, parameters = parameters) + nn = NeuralNetwork(arch; initializer=OneInitializer()) + @test nn([1., 1.], parameters) ≈ [1.0000350420844089, 0.9999649603746436] +end + +test_ghnn_without_parameters() +test_ghnn_with_parameters() \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index d8758cd8d..b697aea21 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -37,6 +37,15 @@ end @safetestset "Hamiltonian Neural Network " begin include("hamiltonian_neural_network_tests.jl") end +@safetestset "Generalized Hamiltonian Neural Network " begin + include("generalized_hamiltonian_neural_networks_test.jl") +end +@safetestset "Symbolic pullback for a single-layer PGHNN " begin + include("generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl") +end +@safetestset "PGHNN training on a ParametricDataLoader " begin + include("generalized_hamiltonian_neural_networks/pghnn_training_test.jl") +end @safetestset "Manifold Neural Network Layers " begin include("layers/manifold_layers.jl") end @@ -128,6 +137,9 @@ end @safetestset "Test data loader for a tensor (q and p data) " begin include("data_loader/draw_batch_for_tensor_test.jl") end +@safetestset "Parametric DataLoader " begin + include("data_loader/parametric_data_loader_test.jl") +end @info "Starting network-loss and kernel tests" @safetestset "Test NetworkLoss + Optimizer " begin diff --git a/test/train!/test_method.jl b/test/train!/test_method.jl index fbce31e80..4df65d03e 100644 --- a/test/train!/test_method.jl +++ b/test/train!/test_method.jl @@ -24,7 +24,7 @@ exacthnn = ExactHnn() sympeuler = SEuler() -@test GeometricMachineLearning.type(sympeuler) == SymplecticEulerA +@test GeometricMachineLearning.type(sympeuler) == SymplecticEulerIntegratorA @test symbols(sympeuler) == PhaseSpaceSymbol @test shape(sympeuler) == TrajectoryData @test min_length_batch(sympeuler) == 2 @@ -64,7 +64,7 @@ midpointlnn = VariaMidPoint() ######################################### @testerror GeometricMachineLearning.type(default_Method(sympnet, tra_pos_data)) -@test GeometricMachineLearning.type(default_method(hnn, tra_ps_data)) == SymplecticEulerA +@test GeometricMachineLearning.type(default_method(hnn, tra_ps_data)) == SymplecticEulerIntegratorA @test GeometricMachineLearning.type(default_method(hnn, sam_dps_data)) == HnnExactMethod @test GeometricMachineLearning.type(default_method(sympnet, tra_ps_data)) == BasicSympNetMethod @test GeometricMachineLearning.type(default_method(lnn, tra_pos_data)) == VariationalMidPointMethod diff --git a/test/training_phnn.jl b/test/training_phnn.jl index 4f30fda1c..0a65aa7b0 100644 --- a/test/training_phnn.jl +++ b/test/training_phnn.jl @@ -5,7 +5,7 @@ using GeometricIntegrators: ImplicitMidpoint, integrate using Random: seed! seed!(123) -function make_alternative_parameters_by_adding_constant(params::NamedTuple=default_parameters, n::Integer=1, a::Number=1.) +function make_alternative_parameters_by_adding_constant(params::NamedTuple=default_parameters(), n::Integer=1, a::Number=1.) _keys = keys(params) values = () for (key, i) in zip(_keys, 1:length(_keys)) @@ -18,13 +18,13 @@ function make_alternative_parameters_by_adding_constant(params::NamedTuple, n::I [make_alternative_parameters_by_adding_constant(params,n, a) for a ∈ a_vals] end -alternative_parameters = make_alternative_parameters_by_adding_constant(default_parameters, 1, Vector(.1:.1:1.)) +alternative_parameters = make_alternative_parameters_by_adding_constant(default_parameters(), 1, Vector(.1:.1:1.)) h_ensemble = hodeensemble(; parameters = alternative_parameters) sol = integrate(h_ensemble, ImplicitMidpoint()) dl = ParametricDataLoader(sol) batch = Batch(100) -arch = GeneralizedHamiltonianArchitecture(4; parameters = default_parameters) +arch = GeneralizedHamiltonianArchitecture(4; parameters = default_parameters()) nn = NeuralNetwork(arch) o = Optimizer(AdamOptimizer(), nn) -o(nn, dl, batch) \ No newline at end of file +o(nn, dl, batch) From 16f6614c592ff2a2e4023103abbb7e4212d6ae82 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Wed, 19 Aug 2026 22:25:24 +0900 Subject: [PATCH 2/3] Pin NeuralNetworkParameters to its main branch until it is registered The previous commit left an absolute local path in `Project.toml`'s `[sources]` -- `Pkg.develop(path = ...)` wrote it and only `docs/Project.toml` was reverted. It is a git URL now. `[sources]` alone is not enough: it is Pkg 1.11+, and the test matrix includes Julia 1.10, where the table is ignored and resolution fails with `expected package NeuralNetworkParameters [67f4d93a] to be registered`. So CI adds it explicitly before `julia-buildpkg`, which works on every version. Verified both paths from a clean manifest: 1.13 resolves through `[sources]`, 1.10 through the added step, and neither rewrites `Project.toml`. `docs/Project.toml` gets its own entry -- `[sources]` is read from the active project only, so `dev`ing GML into the docs environment does not carry GML's across. That follows the `BrenierTwoFluid` precedent already there. The Documentation and Latex workflows run Julia '1', so they need no extra step. All three -- the two `[sources]` entries and the CI step -- come out again once NeuralNetworkParameters is registered. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/CI.yml | 5 +++++ CHANGELOG.md | 5 +++++ Project.toml | 9 ++++++--- docs/Project.toml | 2 ++ 4 files changed, 18 insertions(+), 3 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 6c44a9398..5e8d846dc 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -60,6 +60,11 @@ jobs: version: ${{ matrix.version }} arch: ${{ matrix.arch }} - uses: julia-actions/cache@v1 + # `[sources]` in Project.toml is Pkg 1.11+, so the 1.10 jobs ignore it and fail to + # resolve the unregistered NeuralNetworkParameters. Adding it explicitly works on every + # version. Remove this step, and the `[sources]` entry, once it is registered. + - name: Add unregistered NeuralNetworkParameters + run: julia --project=. -e 'using Pkg; Pkg.add(url="https://github.com/JuliaGNI/NeuralNetworkParameters.jl", rev="main")' - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 - uses: julia-actions/julia-processcoverage@v1 diff --git a/CHANGELOG.md b/CHANGELOG.md index f0759e07d..e43eafb35 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -45,6 +45,11 @@ network covers a whole parameter range rather than a single problem instance. network input, and `flatten`/`unflatten` do that. `ParameterHandling` cannot: `GeometricOptimizers` defines `ParameterHandling.flatten(x)` with an unbound type parameter, and that method wins. + Until NeuralNetworkParameters is registered it is pinned to its `main` branch: a `[sources]` entry + in `Project.toml` and `docs/Project.toml`, plus an explicit `Pkg.add(url = …, rev = "main")` step + in `.github/workflows/CI.yml`, because `[sources]` is Pkg 1.11+ and the test matrix includes + Julia 1.10. All three go away on registration. + ### Changed - **`SymplecticEuler`, `SymplecticEulerA` and `SymplecticEulerB` are no longer exported.** The names diff --git a/Project.toml b/Project.toml index ce82c774a..0eebb290c 100644 --- a/Project.toml +++ b/Project.toml @@ -29,12 +29,15 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [weakdeps] HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" -[sources] -NeuralNetworkParameters = {path = "/Users/mkraus/Datashare/Julia/NeuralNetworkParameters"} - [extensions] HDF5Ext = "HDF5" +# NeuralNetworkParameters is not in the General registry yet. `[sources]` is Pkg 1.11+, so the +# Julia 1.10 CI jobs add it explicitly in `.github/workflows/CI.yml` instead. Both go away, along +# with this comment, once it is registered. +[sources] +NeuralNetworkParameters = {url = "https://github.com/JuliaGNI/NeuralNetworkParameters.jl", rev = "main"} + [compat] AbstractNeuralNetworks = "0.6.4" ChainRules = "1" diff --git a/docs/Project.toml b/docs/Project.toml index d1ca49661..8858db830 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -14,8 +14,10 @@ GeometricOptimizers = "fc236c15-5557-4942-aa65-b650f329279e" GeometricProblems = "18cb22b4-ad41-5c80-9c5f-710df63fbdc9" HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f" +NeuralNetworkParameters = "67f4d93a-60e9-472b-8cdd-1ccf6005724a" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [sources] BrenierTwoFluid = {rev = "main", url = "https://github.com/ToBlick/BrenierTwoFluids.git"} GeometricMachineLearning = {path = ".."} +NeuralNetworkParameters = {url = "https://github.com/JuliaGNI/NeuralNetworkParameters.jl", rev = "main"} From b0fc02a94aa1fdb44f1aac516a41d4ce1604882d Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Wed, 19 Aug 2026 23:18:07 +0900 Subject: [PATCH 3/3] Cover the parametric and forced pieces, and fix what that uncovered codecov put the patch at 55%, with six of the new files at 0%: nothing built 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 defects. `ForcedGeneralizedHamiltonianArchitecture` could not be evaluated at all. `(nn::NeuralNetwork{GT})(qp, problem_params)` and the `Optimizer` entry point were both written for `GT <: GeneralizedHamiltonianArchitecture`, and the forced architecture is a *sibling* of that under `HamiltonianArchitecture`, not a subtype. So `nn(x, mu)` fell through to AbstractNeuralNetworks' generic functor, which read the system parameters as the network parameters and reached the first layer as a `Float64`. Both methods are now defined for it too, in its own file; widening to `HamiltonianArchitecture` would be wrong, since `StandardHamiltonianArchitecture` takes no system parameters. `ParametricResNet(::DataLoader, n_blocks, width; parameters = ...)` accepted `parameters` and did not forward it, so that constructor always built a network with no parameter dependence. The test 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 `p` equation. It perturbs one coordinate at a time and checks the output only moves for a coordinate the layer is named after. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 7 + ..._generalized_hamiltonian_neural_network.jl | 17 ++- src/architectures/parametric_resnet.jl | 2 +- ...arametric_layers_and_architectures_test.jl | 129 ++++++++++++++++++ test/runtests.jl | 3 + 5 files changed, 156 insertions(+), 2 deletions(-) create mode 100644 test/generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index e43eafb35..edc29e4a9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -66,6 +66,13 @@ defines `ParameterHandling.flatten(x)` with an unbound type parameter, and that - `concatenate_array_with_parameters(::AbstractMatrix, ::AbstractVector)` concatenated a batch with `vcat` rather than `hcat`, collapsing it into a single long vector. +- **`ForcedGeneralizedHamiltonianArchitecture` could not be evaluated at all.** The + parameter-dependent `NeuralNetwork` functor and the `Optimizer` entry point were defined for + `GeneralizedHamiltonianArchitecture` only, and the two are siblings under `HamiltonianArchitecture` + rather than sub- and supertype, so `nn(x, μ)` fell through to the generic functor and read the + *system* parameters as the *network* parameters. +- `ParametricResNet(::DataLoader, n_blocks, width; parameters = …)` accepted `parameters` and then + dropped it, silently building a network with no parameter dependence. [nnp]: https://github.com/JuliaGNI/NeuralNetworkParameters.jl diff --git a/src/architectures/forced_generalized_hamiltonian_neural_network.jl b/src/architectures/forced_generalized_hamiltonian_neural_network.jl index 2bb97632e..9b417c7d6 100644 --- a/src/architectures/forced_generalized_hamiltonian_neural_network.jl +++ b/src/architectures/forced_generalized_hamiltonian_neural_network.jl @@ -32,4 +32,19 @@ function Chain(arch::ForcedGeneralizedHamiltonianArchitecture{FT}) where {FT} layers = (layers..., SymplecticEulerB(potential_energy; return_parameters = _return_parameters)) end Chain(layers...) -end \ No newline at end of file +end + +# `ForcedGeneralizedHamiltonianArchitecture` and `GeneralizedHamiltonianArchitecture` are siblings +# under `HamiltonianArchitecture`, so the parameter-dependent forward pass and training entry point +# defined for the latter do not cover this one. `HamiltonianArchitecture` itself is too wide: +# `StandardHamiltonianArchitecture` takes no system parameters. +function (nn::NeuralNetwork{<:ForcedGeneralizedHamiltonianArchitecture})(qp::QPTOAT2, + problem_params::OptionalParameters) + nn.model(qp, problem_params, params(nn)) +end + +function (o::Optimizer)(nn::NeuralNetwork{<:ForcedGeneralizedHamiltonianArchitecture}, + dl::ParametricDataLoader, batch::Batch{:FeedForward}, n_epochs::Integer = 1, + loss::NetworkLoss = ParametricLoss(); kwargs...) + o(nn, dl, batch, n_epochs, loss, ZygotePullback(loss); kwargs...) +end diff --git a/src/architectures/parametric_resnet.jl b/src/architectures/parametric_resnet.jl index 5652ef4cc..2a4ff022d 100644 --- a/src/architectures/parametric_resnet.jl +++ b/src/architectures/parametric_resnet.jl @@ -12,7 +12,7 @@ struct ParametricResNet{AT <: Activation, PT <: OptionalParameters} <: NeuralNet end function ParametricResNet(dl::DataLoader, n_blocks::Integer, width::Integer=dl.input_dim; activation=HNN_activation_default, parameters=NullParameters()) - ParametricResNet(dl.input_dim; width=width, n_blocks=n_blocks, activation) + ParametricResNet(dl.input_dim; width=width, n_blocks=n_blocks, activation=activation, parameters=parameters) end function ResNet(input_dim::Integer, n_blocks::Integer, width::Integer=input_dim; activation=HNN_activation_default, parameters=NullParameters()) diff --git a/test/generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl b/test/generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl new file mode 100644 index 000000000..eaa669edd --- /dev/null +++ b/test/generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl @@ -0,0 +1,129 @@ +# Construction and a forward pass for each parameter-dependent and forced piece. The other PGHNN +# tests build a plain `GeneralizedHamiltonianArchitecture`, so none of these were evaluated +# anywhere -- which is how `ForcedGeneralizedHamiltonianArchitecture`, exported and documented, +# came to have a forward pass that threw. + +using GeometricMachineLearning +using GeometricMachineLearning: ForcingLayerQ, ForcingLayerP, ForcingLayerQP, + ParametricResNetLayer, WideResNetLayer, ParametricResNet +using AbstractNeuralNetworks: params +using Random: seed! +using Test + +seed!(1234) + +const DIM = 4 +const HALF = DIM ÷ 2 +const WIDTH = 8 +const SYSTEM_PARAMETERS = (m = 1.0, ω = π / 2) + +finite(x::AbstractArray) = all(isfinite, x) +finite(qp::NamedTuple) = finite(qp.q) && finite(qp.p) + +# The `Q`/`P`/`QP` suffix names what the forcing *depends on*, not what it changes: a force enters +# the `ṗ` equation, so all three add to `p` and leave `q` alone. +@testset "ForcingLayer$name" for (name, Layer, depends_on) in ( + ("Q", ForcingLayerQ, (:q,)), ("P", ForcingLayerP, (:p,)), ("QP", ForcingLayerQP, (:q, :p))) + layer = Layer(DIM; parameters = SYSTEM_PARAMETERS) + nn = NeuralNetwork(layer) + @test parameterlength(nn) > 0 + + z = (q = rand(HALF), p = rand(HALF)) + out = layer(z, SYSTEM_PARAMETERS, params(nn)) + @test keys(out) == (:q, :p) + @test size(out.q) == size(z.q) && size(out.p) == size(z.p) + @test finite(out) + @test out.q == z.q + @test out.p != z.p + + # perturb one coordinate at a time: the forcing may only move when a coordinate it is named + # after does + for coordinate in (:q, :p) + perturbed = merge(z, NamedTuple{(coordinate,)}((z[coordinate] .+ 1.0,))) + moved = layer(perturbed, SYSTEM_PARAMETERS, params(nn)).p .- perturbed.p != + out.p .- z.p + @test moved == (coordinate in depends_on) + end + + # the same layer applied to the concatenated array form + array_out = layer(vcat(z.q, z.p), SYSTEM_PARAMETERS, params(nn)) + @test array_out ≈ vcat(out.q, out.p) +end + +@testset "WideResNetLayer" begin + layer = WideResNetLayer(DIM, WIDTH, tanh) + nn = NeuralNetwork(Chain(layer)) + ps = params(nn).L1 + @test parameterlength(layer) == WIDTH * (DIM + 1) + DIM * (WIDTH + 1) + + for input in (rand(DIM), rand(DIM, 3), rand(DIM, 3, 2)) + out = layer(input, ps) + @test size(out) == size(input) + @test finite(out) + end + + z = (q = rand(HALF), p = rand(HALF)) + out = layer(z, ps) + @test keys(out) == (:q, :p) + @test out ≈ (q = layer(vcat(z.q, z.p), ps)[1:HALF], p = layer(vcat(z.q, z.p), ps)[(HALF + 1):DIM]) +end + +@testset "ParametricResNetLayer" begin + layer = ParametricResNetLayer(DIM, WIDTH, tanh; + parameters = SYSTEM_PARAMETERS, return_parameters = false) + nn = NeuralNetwork(Chain(layer)) + ps = params(nn).L1 + + # one `NamedTuple` of system parameters describes one sample, so a matrix input is a single + # column; a batch is a vector of parameter sets, which is what `ParametricResNet` builds + for input in (rand(DIM), rand(DIM, 1)) + out = layer(input, SYSTEM_PARAMETERS, ps) + @test size(out) == size(input) + @test finite(out) + end + @test_throws AssertionError layer(rand(DIM, 3), SYSTEM_PARAMETERS, ps) + + z = (q = rand(HALF), p = rand(HALF)) + @test finite(layer(z, SYSTEM_PARAMETERS, ps)) +end + +@testset "ResNet with a width of its own" begin + # `sys_dim == width` keeps the plain `ResNetLayer`; a different width switches to + # `WideResNetLayer`, which is the path `ParametricResNet` compares against + narrow = NeuralNetwork(ResNet(DIM, 2, DIM)) + wide = NeuralNetwork(ResNet(DIM, 2, WIDTH)) + @test parameterlength(wide) > parameterlength(narrow) + @test size(wide(rand(DIM))) == (DIM,) + @test finite(wide(rand(DIM))) +end + +@testset "ParametricResNet" begin + arch = ParametricResNet(DIM; width = WIDTH, n_blocks = 2, parameters = SYSTEM_PARAMETERS) + nn = NeuralNetwork(arch) + out = nn.model(rand(DIM), SYSTEM_PARAMETERS, params(nn)) + @test size(out) == (DIM,) + @test finite(out) + + # the `DataLoader` constructor used to accept `parameters` and drop it + dl = DataLoader(rand(DIM, 20); suppress_info = true) + @test ParametricResNet(dl, 2, WIDTH; parameters = SYSTEM_PARAMETERS).parameters == + SYSTEM_PARAMETERS +end + +@testset "ForcedSympNet $forcing_type" for forcing_type in (:Q, :P, :QP) + nn = NeuralNetwork(ForcedSympNet(DIM; forcing_type = forcing_type)) + out = nn(rand(DIM)) + @test size(out) == (DIM,) + @test finite(out) +end + +@testset "ForcedGeneralizedHamiltonianArchitecture $forcing_type" for forcing_type in (:Q, :P, :QP) + arch = ForcedGeneralizedHamiltonianArchitecture(DIM; parameters = SYSTEM_PARAMETERS, + forcing_type = forcing_type) + nn = NeuralNetwork(arch) + out = nn(rand(DIM), SYSTEM_PARAMETERS) + @test size(out) == (DIM,) + @test finite(out) +end + +@test_throws ErrorException ForcedGeneralizedHamiltonianArchitecture(DIM; forcing_type = :X) diff --git a/test/runtests.jl b/test/runtests.jl index b697aea21..9539d9de8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -46,6 +46,9 @@ end @safetestset "PGHNN training on a ParametricDataLoader " begin include("generalized_hamiltonian_neural_networks/pghnn_training_test.jl") end +@safetestset "Parametric and forced layers and architectures " begin + include("generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl") +end @safetestset "Manifold Neural Network Layers " begin include("layers/manifold_layers.jl") end