Damped Simple Harmonic Oscillator

Let's train a neural network to prove the exponential stability of the damped simple harmonic oscillator (SHO).

The damped SHO is represented by the system of differential equations

\[\begin{align*} \frac{dx}{dt} &= v \\ \frac{dv}{dt} &= -2 \zeta \omega_0 v - \omega_0^2 x \end{align*}\]

where $x$ is the position, $v$ is the velocity, $t$ is time, and $\zeta, \omega_0$ are parameters.

We'll consider just the box domain $x \in [-5, 5], v \in [-2, 2]$.

Copy-Pastable Code

using NeuralPDE, Lux, NeuralLyapunov
using Optimization, OptimizationOptimisers, OptimizationOptimJL
using StableRNGs, Random

rng = StableRNG(0)
Random.seed!(200)

######################### Define dynamics and domain ##########################

"Simple Harmonic Oscillator Dynamics"
function f(state, p, t)
    ζ, ω_0 = p
    pos = state[1]
    vel = state[2]
    vcat(vel, -2ζ * ω_0 * vel - ω_0^2 * pos)
end
lb = [-5.0, -2.0];
ub = [ 5.0,  2.0];
p = [0.5, 1.0];
fixed_point = [0.0, 0.0];
dynamics = ODEFunction(f; sys = SciMLBase.SymbolCache([:x, :v], [:ζ, :ω_0]))

####################### Specify neural Lyapunov problem #######################

# Define neural network discretization
dim_state = length(lb)
dim_hidden = 10
dim_output = 4
chain = [Lux.Chain(
             Dense(dim_state, dim_hidden, tanh),
             Dense(dim_hidden, dim_hidden, tanh),
             Dense(dim_hidden, 1)
         ) for _ in 1:dim_output]
ps = Lux.initialparameters(rng, chain)

# Define training strategy
strategy = QuasiRandomTraining(1000)
discretization = PhysicsInformedNN(chain, strategy; init_params = ps)

# Define neural Lyapunov structure
structure = NonnegativeStructure(
    dim_output;
    δ = 1e-6
)
minimization_condition = DontCheckNonnegativity(check_fixed_point = true)

# Define Lyapunov decrease condition
# Damped SHO has exponential stability at a rate of k = ζ * ω_0, so we train to certify that
decrease_condition = ExponentialStability(prod(p))

# Construct neural Lyapunov specification
spec = NeuralLyapunovSpecification(
    structure,
    minimization_condition,
    decrease_condition
)

############################# Construct PDESystem #############################

@named pde_system = NeuralLyapunovPDESystem(
    dynamics,
    lb,
    ub,
    spec;
    p = p
)

######################## Construct OptimizationProblem ########################

prob = discretize(pde_system, discretization)

########################## Solve OptimizationProblem ##########################

res = Optimization.solve(prob, OptimizationOptimisers.Adam(); maxiters = 500)

###################### Get numerical numerical functions ######################
net = discretization.phi
θ = res.u.depvar

V, V̇ = get_numerical_lyapunov_function(
    net,
    θ,
    structure,
    f,
    fixed_point;
    p = p
)

Detailed description

In this example, we set the dynamics up as an ODEFunction and use a SciMLBase.SymbolCache to tell the ultimate PDESystem what to call our state and parameter variables.

using SciMLBase # for ODEFunction and SciMLBase.SymbolCache

"Simple Harmonic Oscillator Dynamics"
function f(state, p, t)
    ζ, ω_0 = p
    pos = state[1]
    vel = state[2]
    vcat(vel, -2ζ * ω_0 * vel - ω_0^2 * pos)
end
lb = [-5.0, -2.0];
ub = [ 5.0,  2.0];
p = [0.5, 1.0];
fixed_point = [0.0, 0.0];
dynamics = ODEFunction(f; sys = SciMLBase.SymbolCache([:x, :v], [:ζ, :ω_0]))

Setting up the neural network using Lux and NeuralPDE training strategy is no different from any other physics-informed neural network problem. For more on that aspect, see the NeuralPDE documentation.

using Lux, StableRNGs

# Stable random number generator for doc stability
rng = StableRNG(0)

# Define neural network discretization
dim_state = length(lb)
dim_hidden = 10
dim_output = 3
chain = [Lux.Chain(
             Dense(dim_state, dim_hidden, tanh),
             Dense(dim_hidden, dim_hidden, tanh),
             Dense(dim_hidden, 1)
         ) for _ in 1:dim_output]
ps = Lux.initialparameters(rng, chain)

3-element Vector{@NamedTuple{layer_1::@NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}, layer_2::@NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}, layer_3::@NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}}}:
 (layer_1 = (weight = [-0.3237479 -0.5437387; -0.048560407 -1.9027755; … ; -1.6571343 -1.0880196; 0.11297749 -0.21115682], bias = [0.12009778, 0.17294402, 0.29379776, 0.5655657, -0.108922966, 0.04379597, 0.12652467, 0.31780428, -0.2874046, 0.30533487]), layer_2 = (weight = [-0.72350246 -0.3401721 … -0.7510324 -0.069471754; 0.045605414 0.15956083 … 0.21771742 0.19644302; … ; 0.26822245 0.57399637 … -0.81163055 0.012115666; -0.78351724 0.8305831 … 0.5208219 -0.55054986], bias = [0.04156058, 0.02013397, 0.101262905, 0.26914015, -0.2074199, -0.13821417, 0.15174063, -0.14683884, -0.1337309, 0.19964091]), layer_3 = (weight = [-0.35277173 0.4447264 … 0.27290577 -0.103856936], bias = [0.22611074]))
 (layer_1 = (weight = [-1.0708913 -0.2169331; -0.36197582 1.5401223; … ; 0.15410937 -1.0339351; 1.2886711 -0.7981994], bias = [0.48097476, -0.14829622, 0.40842858, -0.15696919, -0.08328807, -0.5449018, -0.091243714, -0.44998, -0.5328614, -0.3034612]), layer_2 = (weight = [0.06931635 0.70472336 … -0.555685 0.06449507; -0.2095592 -0.2052483 … -0.58327854 0.41810396; … ; 0.16986024 -0.6749801 … -0.83834386 -0.3477046; -0.3656068 0.76393116 … -0.6311579 0.52442175], bias = [0.14191626, 0.2844518, 0.07772283, 0.15210773, 0.23963133, -0.0060114358, 0.2544337, -0.013675445, 0.112837784, -0.2950333]), layer_3 = (weight = [-0.49402636 0.35983545 … 0.28786662 0.4767699], bias = [0.074471064]))
 (layer_1 = (weight = [1.2439798 -0.31668776; 0.4865117 -1.0685825; … ; -1.3613875 -1.053776; 1.840911 0.17835672], bias = [-0.089505404, 0.371908, 0.59805584, 0.36353645, -0.499529, 0.6626108, -0.13184647, 0.3982164, -0.2829247, 0.25867644]), layer_2 = (weight = [-0.7837689 0.10970415 … -0.04149061 -0.3317203; -0.55738413 -0.2639527 … 0.73763 0.69223243; … ; 0.04197509 -0.22453102 … 0.47352976 0.5388029; 0.4078881 -0.3659533 … 0.5233756 -0.20981494], bias = [-0.21433939, 0.15363123, 0.02154136, 0.09674127, 0.10775468, 0.033996847, -0.2623025, 0.21038397, -0.036707506, -0.15219104]), layer_3 = (weight = [-0.5051298 -0.4015753 … 0.14219329 -0.5399407], bias = [-0.2908838]))

using NeuralPDE

# Define training strategy
strategy = QuasiRandomTraining(1000)
discretization = PhysicsInformedNN(chain, strategy; init_params = ps)

We now define our Lyapunov candidate structure along with the form of the Lyapunov conditions we'll be using.

For this example, let's use a Lyapunov candidate

\[V(x) = \lVert \phi(x) \rVert^2 + \delta \log \left( 1 + \lVert x \rVert^2 \right),\]

which structurally enforces nonnegativity, but doesn't guarantee $V([0, 0]) = 0$. We therefore don't need a term in the loss function enforcing $V(x) > 0 \, \forall x \ne 0$, but we do need something enforcing $V([0, 0]) = 0$. So, we use DontCheckNonnegativity(check_fixed_point = true).

To train for exponential stability we use ExponentialStability, but we must specify the rate of exponential decrease, which we know in this case to be $\zeta \omega_0$.

using NeuralLyapunov

# Define neural Lyapunov structure
structure = NonnegativeStructure(
    dim_output;
    δ = 1e-6
)
minimization_condition = DontCheckNonnegativity(check_fixed_point = true)

# Define Lyapunov decrease condition
# Damped SHO has exponential stability at a rate of k = ζ * ω_0, so we train to certify that
decrease_condition = ExponentialStability(prod(p))

# Construct neural Lyapunov specification
spec = NeuralLyapunovSpecification(
    structure,
    minimization_condition,
    decrease_condition
)

# Construct PDESystem
@named pde_system = NeuralLyapunovPDESystem(
    dynamics,
    lb,
    ub,
    spec;
    p = p
)

\[ \begin{align} max\left( 0, \frac{2 \cdot 10^{-6} v x}{1 + v^{2} + x^{2}} + \frac{2 \cdot 10^{-6} v \left( - 2 v \zeta \mathtt{\omega\_0} - \mathtt{\omega\_0}^{2} x \right)}{1 + v^{2} + x^{2}} + 0.5 \left( 1 \cdot 10^{-6} \log\left( 1 + v^{2} + x^{2} \right) + \left( \mathtt{{\varphi}3}\left( x, v \right) \right)^{2} + \left( \mathtt{{\varphi}2}\left( x, v \right) \right)^{2} + \left( \mathtt{{\varphi}1}\left( x, v \right) \right)^{2} \right) + 2 \left( v \left( \mathtt{{\varphi}3}\left( x, v \right) \frac{\mathrm{d}}{\mathrm{d}x} \mathtt{{\varphi}3}\left( x, v \right) + \mathtt{{\varphi}2}\left( x, v \right) \frac{\mathrm{d}}{\mathrm{d}x} \mathtt{{\varphi}2}\left( x, v \right) + \mathtt{{\varphi}1}\left( x, v \right) \frac{\mathrm{d}}{\mathrm{d}x} \mathtt{{\varphi}1}\left( x, v \right) \right) + \left( - 2 v \zeta \mathtt{\omega\_0} - \mathtt{\omega\_0}^{2} x \right) \left( \mathtt{{\varphi}3}\left( x, v \right) \frac{\mathrm{d}}{\mathrm{d}v} \mathtt{{\varphi}3}\left( x, v \right) + \mathtt{{\varphi}2}\left( x, v \right) \frac{\mathrm{d}}{\mathrm{d}v} \mathtt{{\varphi}2}\left( x, v \right) + \mathtt{{\varphi}1}\left( x, v \right) \frac{\mathrm{d}}{\mathrm{d}v} \mathtt{{\varphi}1}\left( x, v \right) \right) \right) \right) &= 0 \end{align} \]

Now, we solve the PDESystem using NeuralPDE the same way we would any PINN problem.

prob = discretize(pde_system, discretization)

using Optimization, OptimizationOptimisers, OptimizationOptimJL

res = Optimization.solve(prob, OptimizationOptimisers.Adam(); maxiters = 500)

net = discretization.phi
θ = res.u.depvar

ComponentVector{Float32,SubArray...}(φ1 = (layer_1 = (weight = Float32[-0.37360314 -0.42321053; -0.07675064 -1.9219619; … ; -1.7319822 -1.0410416; 0.15594046 -0.11850441], bias = Float32[0.10652527, 0.14863712, 0.36818013, 0.38835937, -0.14100409, 0.23546612, 0.13478608, 0.34779385, -0.19842929, 0.45498586]), layer_2 = (weight = Float32[-0.67359906 -0.27110812 … -0.77882844 -0.053122725; 0.13309379 0.21480861 … 0.27924326 0.12886639; … ; 0.29979512 0.6047719 … -0.7822111 -0.020179873; -0.81878585 0.7986151 … 0.4857474 -0.5157928], bias = Float32[0.069324724, -0.1168963, 0.18361771, 0.20625809, -0.14732324, -0.14842212, 0.18981332, -0.07319037, -0.18152443, 0.25321078]), layer_3 = (weight = Float32[-0.38729748 0.41377982 … 0.26129347 -0.10148953], bias = Float32[0.16644306])), φ2 = (layer_1 = (weight = Float32[-0.9556224 -0.2728692; -0.34532735 1.5248199; … ; 0.19478579 -0.98960227; 1.2491342 -0.727222], bias = Float32[0.5878323, -0.12998076, 0.25378463, -0.28239173, 0.02366038, -0.5436345, -0.02280118, -0.10832134, -0.47408903, -0.2520412]), layer_2 = (weight = Float32[0.093578435 0.6653792 … -0.51887465 0.050949916; -0.12537558 -0.16435142 … -0.6068842 0.33788154; … ; 0.21162704 -0.64659154 … -0.794251 -0.39211708; -0.30772236 0.805525 … -0.5743718 0.46409097], bias = Float32[0.1012345, 0.19101354, 0.0964599, 0.14057213, 0.2832215, 0.061200857, 0.2766938, -0.006369291, 0.069844544, -0.3462718]), layer_3 = (weight = Float32[-0.43539968 0.33881748 … 0.3672816 0.39426452], bias = Float32[0.037504934])), φ3 = (layer_1 = (weight = Float32[1.3151463 -0.3988785; 0.40486404 -1.1303412; … ; -1.2699504 -1.1549482; 1.7450848 0.26064458], bias = Float32[-0.041738775, 0.34145245, 0.6395982, 0.28324446, -0.53772694, 0.6036485, -0.16615228, 0.33531722, -0.36993554, 0.17587061]), layer_2 = (weight = Float32[-0.7667081 0.07233707 … 0.006496981 -0.27740625; -0.48995423 -0.1941515 … 0.6966831 0.763854; … ; -0.000984827 -0.25361356 … 0.31768247 0.49123415; 0.36626056 -0.41563198 … 0.55952495 -0.24075067], bias = Float32[-0.2727152, 0.0783829, 0.07052464, 0.16007066, 0.024844714, -0.05241387, -0.25464392, 0.26658413, 0.033414543, -0.19929643]), layer_3 = (weight = Float32[-0.45096064 -0.41024217 … 0.21109617 -0.5734637], bias = Float32[-0.21359639])))

We can use the result of the optimization problem to build the Lyapunov candidate as a Julia function.

V, V̇ = get_numerical_lyapunov_function(
    net,
    θ,
    structure,
    f,
    fixed_point;
    p = p
)

Now let's see how we did. We'll evaluate both $V$ and $\dot{V}$ on a $101 \times 101$ grid:

Δx = (ub[1] - lb[1]) / 100
Δv = (ub[2] - lb[2]) / 100
xs = lb[1]:Δx:ub[1]
vs = lb[2]:Δv:ub[2]
states = Iterators.map(collect, Iterators.product(xs, vs))
V_samples = vec(V(hcat(states...)))
V̇_samples = vec(V̇(hcat(states...)))

# Print statistics
V_min, i_min = findmin(V_samples)
state_min = collect(states)[i_min]
V_min, state_min = if V(fixed_point) ≤ V_min
        V(fixed_point), fixed_point
    else
        V_min, state_min
    end

println("V(0.,0.) = ", V(fixed_point))
println("V ∋ [", V_min, ", ", maximum(V_samples), "]")
println("Minimal sample of V is at ", state_min)
println(
    "V̇ ∋ [",
    minimum(V̇_samples),
    ", ",
    max(V̇(fixed_point), maximum(V̇_samples)),
    "]",
)

┌ Warning: Mixed-Precision `matmul_cpu_fallback!` detected and Octavian.jl cannot be used for this set of inputs (C [Matrix{Float64}]: A [Base.ReshapedArray{Float32, 2, SubArray{Float32, 1, Vector{Float32}, Tuple{UnitRange{Int64}}, true}, Tuple{}}] x B [Matrix{Float64}]). Falling back to generic implementation. This may be slow.
└ @ LuxLib.Impl ~/.julia/packages/LuxLib/Kj0os/src/impl/matmul.jl:190
┌ Warning: Mixed-Precision `matmul_cpu_fallback!` detected and Octavian.jl cannot be used for this set of inputs (C [Matrix{Float64}]: A [Base.ReshapedArray{Float32, 2, SubArray{Float32, 1, Vector{Float32}, Tuple{UnitRange{Int64}}, true}, Tuple{}}] x B [Matrix{Float64}]). Falling back to generic implementation. This may be slow.
└ @ LuxLib.Impl ~/.julia/packages/LuxLib/Kj0os/src/impl/matmul.jl:190
┌ Warning: Mixed-Precision `matmul_cpu_fallback!` detected and Octavian.jl cannot be used for this set of inputs (C [Matrix{Float64}]: A [Base.ReshapedArray{Float32, 2, SubArray{Float32, 1, Vector{Float32}, Tuple{UnitRange{Int64}}, true}, Tuple{}}] x B [Matrix{Float64}]). Falling back to generic implementation. This may be slow.
└ @ LuxLib.Impl ~/.julia/packages/LuxLib/Kj0os/src/impl/matmul.jl:190
V(0.,0.) = 0.00892606277834081
V ∋ [0.0041962811003427615, 1.0987211792145484]
Minimal sample of V is at [0.0, -0.08]
V̇ ∋ [-3.4110034188891234, 1.1183980954619652]

At least at these validation samples, the conditions that $\dot{V}$ be negative semi-definite and $V$ be minimized at the origin are nearly satisfied.

using Plots

p1 = plot(xs, vs, V_samples, linetype = :contourf, title = "V", xlabel = "x", ylabel = "v");
p1 = scatter!([0], [0], label = "Equilibrium");
p2 = plot(
    xs,
    vs,
    V̇_samples,
    linetype = :contourf,
    title = "V̇",
    xlabel = "x",
    ylabel = "v",
);
p2 = scatter!([0], [0], label = "Equilibrium");
plot(p1, p2)

Each sublevel set of $V$ completely contained in the plot above has been verified as a subset of the region of attraction.