Parametric Generalized Hamiltonian Neural Networks

Project.toml

@@ -13,13 +13,15 @@ Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GeometricBase = "9a0b12b7-583b-4f04-aa1f-d8551b6addc9"
 GeometricEquations = "c85262ba-a08a-430a-b926-d29770767bf2"
+GeometricProblems = "18cb22b4-ad41-5c80-9c5f-710df63fbdc9"
 GeometricSolutions = "7843afe4-64f4-4df4-9231-049495c56661"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+ParameterHandling = "2412ca09-6db7-441c-8e3a-88d5709968c5"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -48,6 +50,7 @@ HDF5 = "0.16, 0.17"
 KernelAbstractions = "0.9"
 LazyArrays = "=2.3.2"
 NNlib = "0.8, 0.9"
+ParameterHandling = "0.5.0"
 ProgressMeter = "1"
 SafeTestsets = "0.1"
 StatsBase = "0.33, 0.34"

docs/src/GeometricMachineLearning.bib

@@ -751,9 +751,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,
@@ -928,6 +929,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,12 +42,18 @@ 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.GeneralizedHamiltonianArchitecture
+GeometricMachineLearning.ForcingLayer
+GeometricMachineLearning.ParametricDataLoader
 GeometricMachineLearning._processing
 ```

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")

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")