Parametric Generalized Hamiltonian Neural Networks

CHANGELOG.md

@@ -108,6 +108,22 @@ breaking release).
   matrix cotangent through `_matrix_cotangent`, which fixes the rank *and* gives it the array type
   of the primal.
 
+- `concatenate_array_with_parameters(::AbstractMatrix, ::AbstractVector)` concatenated a batch with
+  `vcat` rather than `hcat`, collapsing it into a single long vector ([#207](https://github.com/JuliaGNI/GeometricMachineLearning.jl/pull/207)).
+- **`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.
+- `SymbolicPullback(nn, ::ParametricLoss, μ)` now throws for `n_integrators > 1` instead of appearing
+  to hang. The symbolic expression grows *multiplicatively* with the number of integrators — measured
+  at `dim = 4, width = 4, nhidden = 1`, the loss is 3.4 ⋅ 10⁵ characters at one integrator and
+  1.4 ⋅ 10⁹ at two, and the build never returns. One integrator builds in ≈1.4 s, and the result
+  evaluates about 100× faster than the `Zygote` pullback. See
+  [#245](https://github.com/JuliaGNI/GeometricMachineLearning.jl/issues/245).
+
 ### Changed
 
 - **The kernel `rrule`s honour the ChainRules interface for thunked cotangents.** Twelve pullbacks
@@ -139,8 +155,45 @@ breaking release).
   behind — the value of the innermost `h5save`, an implementation detail of the traversal. Returning
   the path is what `NeuralNetworkParameters.save(filename, ps)` does, so the two now agree.
 
+- **`SymplecticEuler`, `SymplecticEulerA` and `SymplecticEulerB` are no longer exported** ([#207](https://github.com/JuliaGNI/GeometricMachineLearning.jl/pull/207)). 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 ([#207](https://github.com/JuliaGNI/GeometricMachineLearning.jl/pull/207)).
+
 ### 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`. The system parameters reach the network as extra input components, flattened
+  with `NeuralNetworkParameters`' `flatten`/`unflatten`.
+- **`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.
+  Building that pullback refuses `n_integrators > 1`: the symbolic expression grows
+  *multiplicatively* with the number of integrators, so the build does not finish. See
+  [#245](https://github.com/JuliaGNI/GeometricMachineLearning.jl/issues/245).
+- **`ParametricResNet`** and a widened **`ResNet`**, which now takes a `width` separate from the
+  system dimension and uses `WideResNetLayer` when the two differ — 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.
+
 - **`load(NeuralNetwork, h5, arch, prototype)`** — a parameter set of the right shape to rebuild the
   structured leaves against. It is the form that needs no registration: `rebuild` has a prototype to
   take the non-differentiable fields from, so the file's type tags and

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},

docs/src/architectures/hamiltonian_neural_network.md

@@ -42,13 +42,33 @@ 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.apply_parametric
+GeometricMachineLearning.ParametricResNet
+GeometricMachineLearning.ParametricResNetLayer
+GeometricMachineLearning.WideResNetLayer
+GeometricMachineLearning.SymbolicPullback(::GeometricMachineLearning.NeuralNetwork, ::GeometricMachineLearning.ParametricLoss, ::GeometricMachineLearning.GeometricBase.OptionalParameters)
+GeometricMachineLearning._check_symbolic_pullback_is_tractable
+GeometricMachineLearning._flatten_system_parameters
+GeometricMachineLearning._unwrap_gradient
 GeometricMachineLearning._processing
 GeometricMachineLearning._get_contents
 GeometricMachineLearning._get_params

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

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
 ```

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

scripts/TimeDependentHarmonicOscillatorForcedGeneralizedHamiltonianNeuralNetwork.jl

@@ -0,0 +1,79 @@
+using HDF5
+using GeometricMachineLearning
+using GeometricMachineLearning: QPT, QPT2, Activation, ParametricLoss, SymbolicNeuralNetwork, SymbolicPullback
+using CairoMakie
+using NNlib: relu
+
+include(joinpath(@__DIR__, "parametric_data_helpers.jl"))
+
+# 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
+t  = collect(dt * range(0, nt, step=1))
+q, p = forced_harmonic_oscillator_solution(t, IC; omega = omega, Omega = Omega, F = F)
+
+# 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")

scripts/TimeDependentHarmonicOscillatorParametricResnet.jl

@@ -0,0 +1,68 @@
+using HDF5
+using GeometricMachineLearning
+using GeometricMachineLearning: QPT, QPT2, Activation, ParametricLoss, SymbolicNeuralNetwork, SymbolicPullback
+using CairoMakie
+using NNlib: relu
+
+include(joinpath(@__DIR__, "parametric_data_helpers.jl"))
+
+# 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
+t  = collect(dt * range(0, nt, step=1))
+q, p = forced_harmonic_oscillator_solution(t, IC; omega = omega, Omega = Omega, F = F)
+
+# 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")

scripts/Train_DampedOscillator_QP.jl

@@ -0,0 +1,79 @@
+using HDF5
+using GeometricMachineLearning
+using GeometricMachineLearning: QPT, QPT2
+using CairoMakie
+using JLD2
+using NNlib: relu
+
+include(joinpath(@__DIR__, "parametric_data_helpers.jl"))
+
+# 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
+
+# next to the script, unless GML_OUTPUT_DIR says otherwise -- an absolute path from whoever ran it
+# last is no use to anybody else
+path_out = joinpath(get(ENV, "GML_OUTPUT_DIR", @__DIR__), "damped_oscillator_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
+
+
+
+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)
+
+