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
import Optimization, OptimizationOptimisers
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]
    return [vel, -2ζ * ω_0 * vel - ω_0^2 * pos]
end
lb = Float32[-5.0, -2.0];
ub = Float32[ 5.0,  2.0];
p = Float32[0.5, 1.0];
fixed_point = Float32[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 = [Chain(
             Dense(dim_state, dim_hidden, tanh),
             Dense(dim_hidden, dim_hidden, tanh),
             Dense(dim_hidden, 1)
         ) for _ in 1:dim_output]
ps, st = Lux.setup(rng, chain)

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

# Define neural Lyapunov structure and corresponding minimization condition
structure = NonnegativeStructure(dim_output; δ = 1.0f-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)

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

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]
    return [vel, -2ζ * ω_0 * vel - ω_0^2 * pos]
end
lb = Float32[-5.0, -2.0];
ub = Float32[ 5.0,  2.0];
p = Float32[0.5, 1.0];
fixed_point = Float32[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 = [Chain(
             Dense(dim_state, dim_hidden, tanh),
             Dense(dim_hidden, dim_hidden, tanh),
             Dense(dim_hidden, 1)
         ) for _ in 1:dim_output]
ps, st = Lux.setup(rng, chain)

Since Lux.setup defaults to Float32 parameters for Dense layers, we set up the bounds and parameters using Float32 as well. To use Float64 parameters instead, add the following lines:

using ComponentArrays
ps = ps |> ComponentArray |> f64
st = st |> f64

To train the model, we'll sample 1000 points from the domain using NeuralPDE.QuasiRandomTraining. See the NeuralPDE.jl docs for more on how the PDESystem will be converted into an OptimizationProblem.

using NeuralPDE

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

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).\]

using NeuralLyapunov

# Define neural Lyapunov structure and corresponding minimization condition
structure = NonnegativeStructure(dim_output; δ = 1.0f-6)

NeuralLyapunovStructure
    Network dimension: 3
    V(x) = 1.0e-6log(1.0 + (x - x_0)²) + ||φ(x)||²
    V̇(x) = 2(φ(x))⋅(f(x, p, t)*Jφ(x)) + (2.0e-6(x - x_0)*f(x, p, t)) / (1.0 + (x - x_0)²)
    f_call(x) = f(x, p, t)

This structure 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).

minimization_condition = DontCheckNonnegativity(check_fixed_point = true)

LyapunovMinimizationCondition
    Does not train for nonnegativity of V(x)
    Trains for V(x_0) = 0

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$.

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

LyapunovDecreaseCondition
    Trains for 0.5V(x) + V̇(x) ≤ 0
    with approximation a ≤ 0 => max(0, a) ≈ 0

We package these in a NeuralLyapunovSpecification and use it to construct a PDESystem.

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

NeuralLyapunovSpecification
    Structure:
        NeuralLyapunovStructure
            Network dimension: 3
            V(x) = 1.0e-6log(1.0 + (x - x_0)²) + ||φ(x)||²
            V̇(x) = 2(φ(x))⋅(f(x, p, t)*Jφ(x)) + (2.0e-6(x - x_0)*f(x, p, t)) / (1.0 + (x - x_0)²)
            f_call(x) = f(x, p, t)
    Minimization Condition:
        LyapunovMinimizationCondition
            Does not train for nonnegativity of V(x)
            Trains for V(x_0) = 0
    Decrease Condition:
        LyapunovDecreaseCondition
            Trains for 0.5V(x) + V̇(x) ≤ 0
            with approximation a ≤ 0 => max(0, a) ≈ 0

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

PDESystem
Equations: Symbolics.Equation[max(0, (2.0e-6v*(-2v*ζ*ω_0 - x*(ω_0^2))) / (1.0 + v^2 + x^2) + (2.0e-6v*x) / (1.0 + v^2 + x^2) + 0.5(1.0e-6log(1.0 + v^2 + x^2) + φ3(x, v)^2 + φ2(x, v)^2 + φ1(x, v)^2) + 2(v*(φ3(x, v)*Differential(x, 1)(φ3(x, v)) + φ2(x, v)*Differential(x, 1)(φ2(x, v)) + Differential(x, 1)(φ1(x, v))*φ1(x, v)) + (-2v*ζ*ω_0 - x*(ω_0^2))*(φ3(x, v)*Differential(v, 1)(φ3(x, v)) + φ2(x, v)*Differential(v, 1)(φ2(x, v)) + Differential(v, 1)(φ1(x, v))*φ1(x, v)))) ~ 0.0]
Boundary Conditions: Symbolics.Equation[φ1(0.0, 0.0)^2 + φ3(0.0, 0.0)^2 + φ2(0.0, 0.0)^2 ~ 0.0]
Domain: Symbolics.VarDomainPairing[Symbolics.VarDomainPairing(x, -5.0 .. 5.0), Symbolics.VarDomainPairing(v, -2.0 .. 2.0)]
Dependent Variables: Symbolics.Num[φ1(x, v), φ2(x, v), φ3(x, v)]
Independent Variables: Symbolics.Num[x, v]
Parameters: Symbolics.Num[ζ, ω_0]
Default Parameter ValuesModelingToolkitBase.AtomicArrayDict{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymbolicUtils.SymReal}, Dict{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymbolicUtils.SymReal}, SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymbolicUtils.SymReal}}}(ζ => 0.5, ω_0 => 1.0)

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

prob = discretize(pde_system, discretization)

import Optimization, OptimizationOptimisers

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

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

ComponentVector{Float32,SubArray...}(φ1 = (layer_1 = (weight = Float32[-0.37360153 -0.4232104; -0.07675064 -1.9219621; … ; -1.7319793 -1.0410423; 0.15594552 -0.11850377], bias = Float32[0.1065257, 0.14863703, 0.36818168, 0.3883605, -0.14100441, 0.23546763, 0.13478595, 0.3477938, -0.19842929, 0.4549732]), layer_2 = (weight = Float32[-0.6736002 -0.27111223 … -0.7788277 -0.053122614; 0.13309428 0.21480964 … 0.2792428 0.12886752; … ; 0.2997937 0.6047715 … -0.78221124 -0.020179374; -0.81878567 0.7986161 … 0.48574686 -0.51579344], bias = Float32[0.06932491, -0.11689407, 0.18362087, 0.20625578, -0.14732526, -0.14842188, 0.18981409, -0.073184915, -0.18152446, 0.25321007]), layer_3 = (weight = Float32[-0.3872982 0.41377953 … 0.2612935 -0.101489775], bias = Float32[0.16644332])), φ2 = (layer_1 = (weight = Float32[-0.9556231 -0.27286774; -0.34532633 1.5248189; … ; 0.19478573 -0.98960245; 1.2491341 -0.72722256], bias = Float32[0.5878332, -0.12997817, 0.25378886, -0.28239143, 0.023660073, -0.5436362, -0.022800812, -0.10832287, -0.47408974, -0.25204107]), layer_2 = (weight = Float32[0.093578495 0.66537964 … -0.518875 0.050950184; -0.1253755 -0.16435154 … -0.606884 0.33788073; … ; 0.2116271 -0.6465917 … -0.79425156 -0.39211693; -0.30772206 0.80552447 … -0.57437325 0.4640908], bias = Float32[0.10123479, 0.1910134, 0.0964589, 0.14056684, 0.28322223, 0.061201483, 0.27669352, -0.006370619, 0.06984444, -0.3462714]), layer_3 = (weight = Float32[-0.4353998 0.33881843 … 0.36728048 0.39426765], bias = Float32[0.037505213])), φ3 = (layer_1 = (weight = Float32[1.315149 -0.39887452; 0.4048633 -1.1303413; … ; -1.2699504 -1.1549482; 1.7450985 0.26064488], bias = Float32[-0.04173923, 0.34145322, 0.639603, 0.28324494, -0.53772825, 0.6036485, -0.16615383, 0.33531725, -0.36993498, 0.17586939]), layer_2 = (weight = Float32[-0.7667076 0.07233688 … 0.006497935 -0.27741084; -0.4899542 -0.1941513 … 0.696684 0.76385444; … ; -0.0009849019 -0.25361285 … 0.31768337 0.49123487; 0.366261 -0.41563177 … 0.55952454 -0.24075013], bias = Float32[-0.27271354, 0.07838221, 0.07052567, 0.16007534, 0.024843711, -0.05241839, -0.2546432, 0.26658788, 0.033414897, -0.19929594]), layer_3 = (weight = Float32[-0.45096052 -0.4102418 … 0.21109319 -0.57346463], bias = Float32[-0.21359646])))

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)

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(reduce(hcat, states)))
V̇_samples = vec(V̇(reduce(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)),
    "]",
)

V(0.,0.) = 0.00892636552453041
V ∋ [0.00419643414211518, 1.098724439623576]
Minimal sample of V is at Float32[0.0, -0.08]
V̇ ∋ [-3.411007987503139, 1.118394802796075]

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.