Training integrator

scripts/data_problem.jl

@@ -2,7 +2,7 @@ using Plots
 using Zygote
 using  GeometricIntegrators
 using LinearAlgebra
-
+using GeometricMachineLearning
 
 
 #################################################################################################
@@ -14,13 +14,14 @@ dict_problem_H= Dict(
     #:three_body_system =>(,)
 )
 
-# Lagrangian Dictionary
+# Lagrangian DictionaryHH
 
 dict_problem_L = Dict(
     :pendulum => (x-> x[2]^2 / 2 - (1-cos(x[1])),1),
 )
 
 
+
 ###########################################################################
 # get data of a problem for HNN (results is (data = (q,p), target = (q̇,ṗ)))
 
@@ -153,4 +154,136 @@ function compute_phase_space(H_problem, q₀, p₀, tspan = (0., 100.), tstep =
     p = hcat([sol.p[:,i] for i in 1:n_dim]...)
 
     return (q, p)
+end
+
+
+###############################################################################
+# get structure data for multiple trajectory
+
+
+struct storing_data{T}
+    Δt::Float64
+    nb_trajectory::Int
+    data::T
+end
+
+
+function get_multiple_trajectory_structure(nameproblem; n_trajectory = 1, n_points = 10, tstep = 0.1, qmin = -1.2, pmin = -1.2, qmax = 1.2, pmax = 1.2)
+
+
+    # get the Hamiltonien corresponding to name_problem   
+    H_problem, n_dim = dict_problem_H[nameproblem] 
+
+    #define timespan
+    tspan=(0.,n_points*tstep)
+
+    #compute phase space for each trajectory staring from a random point
+    pre_data = NamedTuple()
+
+    for i in 1:n_trajectory
+
+        q₀ = [rand()*(qmax-qmin)+qmin for _ in 1:n_dim]
+        p₀ = [rand()*(pmax-pmin)+pmin for _ in 1:n_dim]
+        q, p = compute_phase_space(H_problem, q₀, p₀, tspan, tstep)
+
+        Data = [(q[n], p[n]) for n in 1:size(q,1)]
+
+
+        nt = NamedTuple{(Symbol("Trajectory_"*string(i)),)}(((data = Data, len = n_points+1),))
+
+        pre_data = merge(pre_data,nt)
+    end
+
+    data = storing_data(tstep, n_trajectory, pre_data)
+
+
+
+    return data #data_trajectory(data, Get_nb_trajectory, Get_length_trajectory, Get_q, Get_p, Get_Δt)
+end
+
+function get_multiple_trajectory_structure_with_target(nameproblem; n_trajectory = 1, n_points = 10, tstep = 0.1, qmin = -0.2, pmin = -0.2, qmax = 0.2, pmax = 0.2)
+
+
+    # get the Hamiltonien corresponding to name_problem   
+    H_problem, n_dim = dict_problem_H[nameproblem] 
+
+    # compute vector field
+    ∇H(x) = Zygote.gradient(H,x)[1]
+
+    I = Diagonal(ones(n_dim))
+    Z = zeros(n_dim,n_dim)
+    symplectic_matrix = [Z I;-I Z]
+
+    dH(x) = symplectic_matrix * ∇H(x)
+
+    #define timespan
+    tspan=(0.,n_points*tstep)
+
+    #compute phase space for each trajectory staring from a random point
+    pre_data = NamedTuple()
+    Target = NamedTuple()
+
+    for i in 1:n_trajectory
+
+        q₀ = [rand()*(qmax-qmin)+qmin for _ in 1:n_dim]
+        p₀ = [rand()*(pmax-pmin)+pmin for _ in 1:n_dim]
+        q, p = compute_phase_space(H_problem, q₀, p₀, tspan, tstep)
+
+        Data = [(q[n], p[n]) for n in 1:size(q,1)]
+        data_calc = [[q[n]..., p[n]...] for n in 1:size(q,1)]
+
+        nt = NamedTuple{(Symbol("Trajectory_"*string(i)),)}(((data = Data, len = n_points+1),))
+
+        pre_data = merge(pre_data,nt)
+
+        calc = dH.(data_calc)
+        pre_target = [(calc[n][1:n_dim], calc[n][1+n_dim:end]) for n in 1:size(q,1)]
+        ntt = NamedTuple{(Symbol("Trajectory_"*string(i)),)}(((target = pre_target, len = n_points+1),))
+        Target = merge(Target,ntt)
+
+    end
+
+    data = storing_data(tstep, n_trajectory, pre_data)
+
+    Get_Δt(Data) = Data.Δt
+    Get_nb_trajectory(Data) = Data.nb_trajectory
+    Get_length_trajectory(Data, i) = Data.data[Symbol("Trajectory_"*string(i))][:len]
+    Get_q(Data, i, n) = Data.data[Symbol("Trajectory_"*string(i))][:data][n][1]
+    Get_p(Data, i, n) = Data.data[Symbol("Trajectory_"*string(i))][:data][n][2]
+
+    Get_q̇(Target,i,n)= Target[Symbol("Trajectory_"*string(i))][:target][n][1][1]
+    Get_ṗ(Target,i,n)= Target[Symbol("Trajectory_"*string(i))][:target][n][2][1]
+
+    return dataTarget(data_trajectory(data, Get_nb_trajectory, Get_length_trajectory, Get_q, Get_p, Get_Δt), Target, Get_q̇, Get_ṗ)
+end
+
+
+function get_multiple_trajectory_structure_Lagrangian(nameproblem; n_trajectory = 1, n_points = 10, tstep = 0.1, qmin = -0.2, pmin = -0.2, qmax = 0.2, pmax = 0.2)
+
+
+    # get the Hamiltonien corresponding to name_problem   
+    H_problem, n_dim = dict_problem_L[nameproblem] 
+
+    #define timespan
+    tspan=(0.,n_points*tstep)
+
+    #compute phase space for each trajectory staring from a random point
+    pre_data = []
+    push!(pre_data,[tstep])
+    push!(pre_data,[n_trajectory])
+    push!(pre_data, [n_points+1])
+
+
+    for i in 1:n_trajectory
+
+        q₀ = [rand()*(qmax-qmin)+qmin for _ in 1:n_dim]
+        p₀ = [rand()*(pmax-pmin)+pmin for _ in 1:n_dim]
+        q, p = compute_phase_space(H_problem, q₀, p₀, tspan, tstep)
+
+        Data = [q[n] for n in 1:size(q,1)]
+
+        pre_data = push!(pre_data,Data)
+    end
+
+    return pre_data 
 end

scripts/hnn_scripts/hnn_script.jl

@@ -0,0 +1,60 @@
+# using Profile
+using GeometricMachineLearning
+
+# this contains the functions for generating the training data
+include("../data_problem.jl")
+
+
+function HNN(integrator::TrainingIntegrator{<:HnnTrainingIntegrator}, data::AbstractTrainingData, nameproblem::Symbol = :pendulum, opt =  MomentumOptimizer(1e-3,0.5))
+
+    _, n_dim = dict_problem_H[nameproblem]
+
+    # layer dimension/width
+    ld = 5
+
+    # hidden layers
+    ln = 3
+
+    # number of inputs/dimension of system
+    ninput = 2*n_dim
+
+    # number of training runs
+    nruns = 3
+
+    # create HNN
+    hnn = HamiltonianNeuralNetwork(ninput; nhidden = ln, width = ld)
+
+    # create Lux network
+    nn = NeuralNetwork(hnn, LuxBackend())
+
+    # perform training (returns array that contains the total loss for each training step)
+    total_loss = train!(nn, opt, data; ntraining = nruns, ti = integrator, showprogress = false)
+
+    return nn, total_loss
+end
+
+
+
+Data,Target = get_HNN_data(:pendulum)
+
+Get_Data = Dict(
+    :nb_points => Data -> length(Data),
+    :q => (Data,n) -> Data[n][1],
+    :p => (Data,n) -> Data[n][2]
+)
+pdata = DataSampled(Data, Get_Data)
+
+Get_Target = Dict(
+    :q̇ => (Target,n) -> Target[n][1],
+    :ṗ => (Target,n) -> Target[n][2],
+)
+
+data = DataTarget(pdata, Target, Get_Target)
+
+nn, total_loss = HNN(ExactHnn(), data,  :pendulum, MomentumOptimizer())
+
+include("../plots.jl")
+
+H,_ = dict_problem_H[:pendulum]
+
+plot_hnn(H, nn, total_loss; filename="hnn_pendulum.png", xmin=-1.2, xmax=+1.2, ymin=-1.2, ymax=+1.2)

scripts/hnn_scripts/tests_file.jl

@@ -0,0 +1,99 @@
+
+# using Profile
+using GeometricMachineLearning
+
+# this include the scripts using GeometricMachineLearning
+include("hnn_script.jl")
+include("../lnn_scripts/lnn_script.jl")
+include("../sympnet_script.jl")
+
+# this contains the functions for generating the training data
+include("../data_problem.jl")
+
+# this include the macro @testerror ans @testerror set for testing the code
+include("../macro_test.jl")
+
+
+#Data HNN with target
+Data,Target = get_HNN_data(:pendulum)
+
+Get_Data = Dict(
+    :nb_points => Data -> length(Data),
+    :q => (Data,n) -> Data[n][1],
+    :p => (Data,n) -> Data[n][2]
+)
+pdata = DataSampled(Data, Get_Data)
+
+Get_Target = Dict(
+    :q̇ => (Target,n) -> Target[n][1],
+    :ṗ => (Target,n) -> Target[n][2],
+)
+
+data = DataTarget(pdata, Target, Get_Target)
+
+#Data multiple trajectory 
+Data = get_multiple_trajectory_structure(:pendulum; n_trajectory = 2, n_points = 3, tstep = 0.1, qmin = -1.2, pmin = -1.2, qmax = 1.2, pmax = 1.2)
+
+Get_Data = Dict(
+    :Δt => Data -> Data.Δt,
+    :nb_trajectory => Data -> Data.nb_trajectory,
+    :length_trajectory => (Data,i) -> Data.data[Symbol("Trajectory_"*string(i))][:len],
+    :q => (Data,i,n) -> Data.data[Symbol("Trajectory_"*string(i))][:data][n][1],
+    :p => (Data,i,n) -> Data.data[Symbol("Trajectory_"*string(i))][:data][n][2],
+)
+data2 = DataTrajectory(Data, Get_Data)
+
+#Data LNN with target
+Data, Target = get_LNN_data(:pendulum)
+
+Get_Data = Dict(
+    :nb_points => Data -> length(Data),
+    :q => (Data,n) -> Data[n][1][1],
+    :q̇ => (Data,n) -> Data[n][2][1]
+)
+pdata = DataSampled(Data, Get_Data)
+
+Get_Target = Dict(
+    :q̈ => (Target,n) -> Target[n][1],
+)
+
+data3 = DataTarget(pdata, Target, Get_Target)
+
+#Data LNN without target
+Data = get_multiple_trajectory_structure(:pendulum; n_trajectory = 2, n_points = 3, tstep = 0.1, qmin = -1.2, pmin = -1.2, qmax = 1.2, pmax = 1.2)
+
+Get_Data = Dict(
+    :Δt => Data -> Data.Δt,
+    :nb_trajectory => Data -> Data.nb_trajectory,
+    :length_trajectory => (Data,i) -> Data.data[Symbol("Trajectory_"*string(i))][:len],
+    :q => (Data,i,n) -> Data.data[Symbol("Trajectory_"*string(i))][:data][n][1],
+)
+data4 = DataTrajectory(Data, Get_Data)
+
+
+@testseterrors begin
+
+    @testerror HNN ExactHnn() data :pendulum MomentumOptimizer()
+    @testerror HNN ExactHnn() data :pendulum AdamOptimizer()
+    @testerror HNN ExactHnn() data :pendulum GradientOptimizer()
+
+    @testerror HNN SEulerA() data2 :pendulum MomentumOptimizer()
+    @testerror HNN SEulerA() data2 :pendulum AdamOptimizer()
+    @testerror HNN SEulerA() data2 :pendulum GradientOptimizer()
+
+    @testerror SYMPNET BasicSympNet() data2 :pendulum MomentumOptimizer()
+    @testerror SYMPNET BasicSympNet() data2 :pendulum AdamOptimizer()
+    @testerror SYMPNET BasicSympNet() data2 :pendulum GradientOptimizer()
+
+    #@testerror LNN ExactLnn() data3 :pendulum MomentumOptimizer()
+    #@testerror LNN ExactLnn() data3 :pendulum AdamOptimizer()
+    #@testerror LNN ExactLnn() data3 :pendulum GradientOptimizer()
+
+    #@testerror LNN VariaMidPoint() data4 :pendulum MomentumOptimizer()
+    #@testerror LNN VariaMidPoint() data4 :pendulum AdamOptimizer()
+    #@testerror LNN VariaMidPoint() data4 :pendulum GradientOptimizer()
+
+end
+
+
+