Estimating the Region of Attraction

In this demonstration, we add awareness of the region of attraction (RoA) estimation task to our training.

We'll be examining the simple one-dimensional differential equation:

\[\frac{dx}{dt} = - x + x^3.\]

This system has a fixed point at $x = 0$ which has a RoA of $x \in (-1, 1)$, which we will attempt to identify.

We'll train in the larger domain $x \in [-2, 2]$.

Copy-Pastable Code

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

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

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

f(x, p, t) = -x .+ x .^ 3
lb = [-2.0];
ub = [ 2.0];
fixed_point = [0.0];

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

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

# Define training strategy
strategy = GridTraining(0.1)
discretization = PhysicsInformedNN(chain, strategy; init_params = ps)

# Define neural Lyapunov structure
structure = PositiveSemiDefiniteStructure(dim_output)
minimization_condition = DontCheckNonnegativity()

# Define Lyapunov decrease condition
decrease_condition = make_RoA_aware(AsymptoticStability())

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

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

@named pde_system = NeuralLyapunovPDESystem(
    f,
    lb,
    ub,
    spec
)

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

prob = discretize(pde_system, discretization)

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

res = Optimization.solve(prob, OptimizationOptimisers.Adam(); maxiters = 300)
prob = Optimization.remake(prob, u0 = res.u)
res = Optimization.solve(prob, OptimizationOptimJL.BFGS(); maxiters = 300)

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

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

################################## Simulate ###################################
states = lb[]:0.001:ub[]
V_samples = vec(V(states'))
V̇_samples = vec(V̇(states'))

# Calculated RoA estimate
ρ = decrease_condition.ρ
RoA_states = states[vec(V(transpose(states))) .≤ ρ]
RoA = (first(RoA_states), last(RoA_states))

Detailed description

In this example, we set up the dynamics as a Julia function and don't bother specifying the symbols for the variables (so $x$ will be called the default state1 in the PDESystem).

f(x, p, t) = -x .+ x .^ 3
lb = [-2.0];
ub = [ 2.0];
fixed_point = [0.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. Since we're only considering one dimension, training on a grid isn't so bad in this case.

using Lux, StableRNGs

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

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

2-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}}}}:
 (layer_1 = (weight = [-0.45784864; -0.06867478; … ; 1.64411; 0.23668337;;], bias = [0.3418486, -0.978009, -0.04572296, -0.8118267, 0.055347443]), layer_2 = (weight = [-0.34389058 0.9279804 … 0.07996013 -1.0231869; -1.2034209 0.28792876 … 0.23100138 0.064495794; … ; -0.50130814 -0.688124 … -0.52472657 0.8011582; -0.9508799 -0.13354729 … 0.55746263 0.5449771], bias = [0.27231756, 0.33674875, -0.11881286, 0.13140164, -0.3838435]), layer_3 = (weight = [-0.2886456 0.13539186 … -0.14496835 -0.17473423],))
 (layer_1 = (weight = [2.2938614; -2.1249437; … ; 1.815136; 2.6265342;;], bias = [-0.9846766, -0.63826203, 0.51425433, -0.46349335, -0.036490917]), layer_2 = (weight = [-0.02004436 -1.2301451 … -1.0701928 0.8329293; 0.69851124 1.0166358 … -1.0654731 -0.87174803; … ; -1.1370962 1.029001 … -1.1567752 -0.55754757; -0.8275013 0.5202121 … -0.15313427 -0.98145705], bias = [-0.15087044, -0.12097594, -0.40972832, -0.246295, -0.043452278]), layer_3 = (weight = [-0.13673632 0.20193696 … -0.49816325 -0.41002852],))

using NeuralPDE

# Define training strategy
strategy = GridTraining(0.1)
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 the default Lyapunov candidate from PositiveSemiDefiniteStructure:

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

which structurally enforces positive definiteness. We therefore use DontCheckNonnegativity().

We only require asymptotic stability in this example, but we use make_RoA_aware to only penalize positive values of $\dot{V}(x)$ when $V(x) \le 1$.

using NeuralLyapunov

# Define neural Lyapunov structure
structure = PositiveSemiDefiniteStructure(dim_output)
minimization_condition = DontCheckNonnegativity()

# Define Lyapunov decrease condition
decrease_condition = make_RoA_aware(AsymptoticStability())

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

# Construct PDESystem
@named pde_system = NeuralLyapunovPDESystem(
    f,
    lb,
    ub,
    spec
)

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

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

prob = discretize(pde_system, discretization)

import Optimization, OptimizationOptimisers, OptimizationOptimJL

res = Optimization.solve(prob, OptimizationOptimisers.Adam(); maxiters = 300)
prob = Optimization.remake(prob, u0 = res.u)
res = Optimization.solve(prob, OptimizationOptimJL.BFGS(); maxiters = 300)

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

ComponentVector{Float32,SubArray...}(φ1 = (layer_1 = (weight = Float32[-0.45784864; -0.06867478; … ; 1.64411; 0.23668337;;], bias = Float32[0.3418486, -0.978009, -0.04572296, -0.8118267, 0.055347443]), layer_2 = (weight = Float32[-0.34389058 0.9279804 … 0.07996013 -1.0231869; -1.2034209 0.28792876 … 0.23100138 0.064495794; … ; -0.50130814 -0.688124 … -0.52472657 0.8011582; -0.9508799 -0.13354729 … 0.55746263 0.5449771], bias = Float32[0.27231756, 0.33674875, -0.11881286, 0.13140164, -0.3838435]), layer_3 = (weight = Float32[-0.2886456 0.13539186 … -0.14496835 -0.17473423])), φ2 = (layer_1 = (weight = Float32[2.2938614; -2.1249437; … ; 1.815136; 2.6265342;;], bias = Float32[-0.9846766, -0.63826203, 0.51425433, -0.46349335, -0.036490917]), layer_2 = (weight = Float32[-0.02004436 -1.2301451 … -1.0701928 0.8329293; 0.69851124 1.0166358 … -1.0654731 -0.87174803; … ; -1.1370962 1.029001 … -1.1567752 -0.55754757; -0.8275013 0.5202121 … -0.15313427 -0.98145705], bias = Float32[-0.15087044, -0.12097594, -0.40972832, -0.246295, -0.043452278]), layer_3 = (weight = Float32[-0.13673632 0.20193696 … -0.49816325 -0.41002852])))

We can use the result of the optimization problem to build the Lyapunov candidate as a Julia function, then sample on a finer grid than we trained on to find the estimated region of attraction.

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

# Sample
states = lb[]:0.001:ub[]
V_samples = vec(V(states'))
V̇_samples = vec(V̇(states'))

# Calculate RoA estimate
ρ = decrease_condition.ρ
RoA_states = states[vec(V(transpose(states))) .≤ ρ]
RoA = (first(RoA_states), last(RoA_states))

# Print statistics
println("V(0.,0.) = ", V(fixed_point))
println("V ∋ [", min(V(fixed_point), minimum(V_samples)), ", ", maximum(V_samples), "]")
println(
    "V̇ ∋ [",
    minimum(V̇_samples),
    ", ",
    max(V̇(fixed_point), maximum(V̇_samples)),
    "]",
)
println("True region of attraction: (-1, 1)")
println("Estimated region of attraction: ", RoA)

┌ 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.0
┌ 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, 3.5251504663249973]
V̇ ∋ [-0.721374428820775, 11.087438260235606]
True region of attraction: (-1, 1)
Estimated region of attraction: (-0.788, 0.821)

The estimated region of attraction is within the true region of attraction.

using Plots

p_V = plot(states, V_samples, label = "V", xlabel = "x", linewidth=2);
p_V = hline!([ρ], label = "V = ρ", legend = :top);
p_V = vspan!(collect(RoA); label = "Estimated Region of Attraction", color = :gray, fillstyle = :/);
p_V = vspan!([-1, 1]; label = "True Region of Attraction", opacity = 0.2, color = :green);

p_V̇ = plot(states, V̇_samples, label = "dV/dt", xlabel = "x", linewidth=2);
p_V̇ = hline!([0.0], label = "dV/dt = 0", legend = :top);
p_V̇ = vspan!(collect(RoA); label = "Estimated Region of Attraction", color = :gray, fillstyle = :/);
p_V̇ = vspan!([-1, 1]; label = "True Region of Attraction", opacity = 0.2, color = :green);

plt = plot(p_V, p_V̇)