From 78632d08d340842d3b505d3f5c6006e1519f4216 Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Fri, 15 Mar 2024 14:57:42 -0400
Subject: [PATCH 01/16] Added RoA estimation test to runtests

---
 test/runtests.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/test/runtests.jl b/test/runtests.jl
index 4887b76..ce449d9 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -7,4 +7,7 @@ using SafeTestsets
     @time @safetestset "Damped pendulum" begin
         include("damped_pendulum.jl")
     end
+    @time @safetestset "Region of attraction estimation" begin
+        include("roa_estimation.jl")
+    end
 end

From e990bd756b47d867337db0950b4aa89738148f19 Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Fri, 15 Mar 2024 15:07:11 -0400
Subject: [PATCH 02/16] Ceded form of dynamics function to
 NeuralLyapunovStructure

---
 src/NeuralLyapunovPDESystem.jl  |  4 +++-
 src/conditions_specification.jl |  8 ++++----
 src/structure_specification.jl  | 26 +++++++++++++++-----------
 3 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/src/NeuralLyapunovPDESystem.jl b/src/NeuralLyapunovPDESystem.jl
index 349ae2d..a60859a 100644
--- a/src/NeuralLyapunovPDESystem.jl
+++ b/src/NeuralLyapunovPDESystem.jl
@@ -204,8 +204,10 @@ function _NeuralLyapunovPDESystem(
         structure.V̇(
             u,
             y -> Symbolics.jacobian(u(y), y),
-            y -> dynamics(y, params, 0.0),
+            dynamics,
             x,
+            params,
+            0.0,
             fixed_point
         )
     end
diff --git a/src/conditions_specification.jl b/src/conditions_specification.jl
index 8e30305..31c1c95 100644
--- a/src/conditions_specification.jl
+++ b/src/conditions_specification.jl
@@ -8,10 +8,10 @@ enforce Lyapunov conditions.
 `network_dim` is the dimension of the output of the neural network.
 `V(phi::Function, state, fixed_point)` takes in the neural network, the state, and the fixed
 point, and outputs the value of the Lyapunov function at `state`.
-`V̇(phi::Function, J_phi::Function, f::Function, state, fixed_point)` takes in the neural
-network, the jacobian of the neural network, the dynamics (as a function of the state
-alone), the state, and the fixed point, and outputs the time derivative of the Lyapunov
-function at `state`.
+`V̇(phi::Function, J_phi::Function, f::Function, state, params, t, fixed_point)` takes in the
+neural network, jacobian of the neural network, dynamics, state, parameters and time (for
+calling the dynamics, when relevant), and fixed point, and outputs the time derivative of
+the Lyapunov function at `state`.
 """
 struct NeuralLyapunovStructure
     V::Function
diff --git a/src/structure_specification.jl b/src/structure_specification.jl
index e225fb4..20c23db 100644
--- a/src/structure_specification.jl
+++ b/src/structure_specification.jl
@@ -8,7 +8,8 @@ function UnstructuredNeuralLyapunov()::NeuralLyapunovStructure
     NeuralLyapunovStructure(
         (net, state, fixed_point) -> net(state),
         (net, grad_net, state, fixed_point) -> grad_net(state),
-        (net, grad_net, f, state, fixed_point) -> grad_net(state) ⋅ f(state),
+        (net, grad_net, f, state, params, t, fixed_point) -> grad_net(state) ⋅
+                                                             f(state, params, t),
         1
     )
 end
@@ -46,8 +47,9 @@ function NonnegativeNeuralLyapunov(
         NeuralLyapunovStructure(
             (net, state, fixed_point) -> net(state) ⋅ net(state),
             (net, J_net, state, fixed_point) -> 2 * transpose(net(state)) * J_net(state),
-            (net, J_net, f, state, fixed_point) -> 2 *
-                                                   dot(net(state), J_net(state), f(state)),
+            (net, J_net, f, state, params, t, fixed_point) -> 2 *
+                                                              dot(
+                net(state), J_net(state), f(state, params, t)),
             network_dim
         )
     else
@@ -61,11 +63,11 @@ function NonnegativeNeuralLyapunov(
                                          δ * pos_def(state, fixed_point),
             (net, J_net, state, fixed_point) -> 2 * transpose(net(state)) * J_net(state) +
                                                 δ * grad_pos_def(state, fixed_point),
-            (net, J_net, f, state, fixed_point) -> 2 * dot(
+            (net, J_net, f, state, params, t, fixed_point) -> 2 * dot(
                 net(state),
                 J_net(state),
-                f(state)
-            ) + δ * grad_pos_def(state, fixed_point) ⋅ f(state),
+                f(state, params, t)
+            ) + δ * grad_pos_def(state, fixed_point) ⋅ f(state, params, t),
             network_dim
         )
     end
@@ -123,11 +125,13 @@ function PositiveSemiDefiniteStructure(
                                             non_neg(net, state, fixed_point) +
                                             pos_def(state, fixed_point) *
                                             grad_non_neg(net, J_net, state, fixed_point),
-        (net, J_net, f, state, fixed_point) -> (f(state) ⋅
-                                                grad_pos_def(state, fixed_point)) *
-                                               non_neg(net, state, fixed_point) +
-                                               pos_def(state, fixed_point) *
-                                               (f(state) ⋅ grad_non_neg(
+        (net, J_net, f, state, params, t, fixed_point) -> (f(state, params, t) ⋅
+                                                           grad_pos_def(
+            state, fixed_point)) *
+                                                          non_neg(net, state, fixed_point) +
+                                                          pos_def(state, fixed_point) *
+                                                          (f(state, params, t) ⋅
+                                                           grad_non_neg(
             net, J_net, state, fixed_point)),
         network_dim
     )

From 657d609a9d62c2c48c00240daccaf7dce258a834 Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Fri, 15 Mar 2024 18:02:24 -0400
Subject: [PATCH 03/16] Set up NeuralPDESystem(f::Function, ...) to handle
 f(x,u,p,t) when spec uses that form

---
 src/NeuralLyapunovPDESystem.jl | 78 ++++++++++++++++++++--------------
 1 file changed, 47 insertions(+), 31 deletions(-)

diff --git a/src/NeuralLyapunovPDESystem.jl b/src/NeuralLyapunovPDESystem.jl
index a60859a..a9f1fd8 100644
--- a/src/NeuralLyapunovPDESystem.jl
+++ b/src/NeuralLyapunovPDESystem.jl
@@ -9,10 +9,11 @@ columnwise.
 The neural Lyapunov function will only be trained for `{ x : lb .≤ x .≤ ub }`. The Lyapunov
 function will be for the dynamical system represented by `dynamics`. If `dynamics` is an
 `ODEProblem` or `ODEFunction`, then the corresponding ODE; if `dynamics` is a function, then
-the ODE is `ẋ = dynamics(x, p, t)`. This ODE should not depend on `t` (time `t=0.0` alone
-will be used) and should have a fixed point at `x = fixed_point`. The particular Lyapunov
-conditions to be used and structure of the neural Lyapunov function are specified through
-`spec`, which is a `NeuralLyapunovSpecification`.
+the ODE is typically `ẋ = dynamics(x, p, t)`, but `spec` may specify a different form. This
+ODE should not depend on `t` (time `t=0.0` alone will be used) and should have a fixed point
+at `x = fixed_point`. The particular Lyapunov conditions to be used and structure of the
+neural Lyapunov function are specified through `spec`, which is a
+`NeuralLyapunovSpecification`.
 
 The returned neural network function takes three inputs: the neural network structure `phi`,
 the trained network parameters, and a matrix of inputs to operate on columnwise.
@@ -20,36 +21,26 @@ the trained network parameters, and a matrix of inputs to operate on columnwise.
 If `dynamics` requires parameters, their values can be supplied through the Vector `p`, or
 through the parameters of `dynamics` if `dynamics isa ODEProblem` (in which case, let
 the other be `SciMLBase.NullParameters()`). If `dynamics` is an `ODEFunction` and
-`dynamics.paramsyms` is defined, then `p` should have the same order.
+`dynamics.sys.parameters` is defined, then `p` should have the same order.
+
+To use specific symbols for the state and parameters in the resulting `PDESystem`, supply
+them through the `state_syms` and `parameter_syms` keywords (if not using an `ODEFunction`,
+`ODEProblem`, or `ODESystem`, which each have their own ways of defining symbols.)
 """
 function NeuralLyapunovPDESystem(
-        dynamics::ODEFunction,
+        dynamics::Function,
         lb,
         ub,
         spec::NeuralLyapunovSpecification;
         fixed_point = zeros(length(lb)),
-        p = SciMLBase.NullParameters()
+        p = SciMLBase.NullParameters(),
+        state_syms = [],
+        parameter_syms = []
 )::Tuple{PDESystem, Function}
-    if dynamics.mass_matrix !== I
-        throw(ErrorException("DAEs are not supported at this time."))
-    end
-
-    if dynamics.sys isa ODESystem
-        return NeuralLyapunovPDESystem(
-            dynamics.sys,
-            lb,
-            ub,
-            spec;
-            fixed_point = fixed_point,
-            p = p
-        )
-    end
-
     ########################## Define state symbols ###########################
     state_dim = length(lb)
 
     # Define state symbols, if not already defined
-    state_syms = SciMLBase.variable_symbols(dynamics.sys)
     state_syms = if isempty(state_syms)
         [Symbol(:state, i) for i in 1:state_dim]
     else
@@ -62,10 +53,10 @@ function NeuralLyapunovPDESystem(
     param_syms = if p == SciMLBase.NullParameters()
         []
     else
-        if isnothing(dynamics.sys.parameters)
+        if isempty(parameter_syms)
             [Symbol(:param, i) for i in 1:length(p)]
         else
-            dynamics.sys.parameters
+            parameter_syms
         end
     end
 
@@ -92,20 +83,43 @@ function NeuralLyapunovPDESystem(
 end
 
 function NeuralLyapunovPDESystem(
-        dynamics::Function,
+        dynamics::ODEFunction,
         lb,
         ub,
         spec::NeuralLyapunovSpecification;
         fixed_point = zeros(length(lb)),
         p = SciMLBase.NullParameters()
 )::Tuple{PDESystem, Function}
+    if dynamics.mass_matrix !== I
+        throw(ErrorException("DAEs are not supported at this time."))
+    end
+
+    if dynamics.sys isa ODESystem
+        return NeuralLyapunovPDESystem(
+            dynamics.sys,
+            lb,
+            ub,
+            spec;
+            fixed_point = fixed_point,
+            p = p
+        )
+    end
+
+    p_syms = if isnothing(dynamics.sys.parameters)
+        []
+    else
+        dynamics.sys.parameters
+    end
+
     return NeuralLyapunovPDESystem(
-        ODEFunction(dynamics),
+        dynamics.f,
         lb,
         ub,
         spec;
         fixed_point = fixed_point,
-        p = p
+        p = p,
+        state_syms = SciMLBase.variable_symbols(dynamics.sys),
+        parameter_syms = p_syms
     )
 end
 
@@ -294,7 +308,7 @@ function NumericalNeuralLyapunovFunctions(
         structure::NeuralLyapunovStructure,
         dynamics::Function,
         fixed_point;
-        p,
+        p = SciMLBase.NullParameters(),
         jac = ForwardDiff.jacobian,
         J_net = (_phi, _θ, x) -> jac((y) -> network_func(_phi, _θ, y), x)
 )::Tuple{Function, Function, Function}
@@ -320,8 +334,10 @@ function NumericalNeuralLyapunovFunctions(
         structure.V̇(
             _net_func,
             _J_net,
-            y -> dynamics(y, p, 0.0),
+            dynamics,
             state,
+            p,
+            0.0,
             fixed_point
         )
     end
@@ -352,7 +368,7 @@ function NumericalNeuralLyapunovFunctions(
         V_structure::Function,
         dynamics::Function,
         fixed_point;
-        p = SciMLBase.NullParameters,
+        p = SciMLBase.NullParameters(),
         grad = ForwardDiff.gradient
 )::Tuple{Function, Function, Function}
     # Make network function

From 2d8b84c44299b3718c354305c01f27b82068b599 Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Fri, 15 Mar 2024 18:02:58 -0400
Subject: [PATCH 04/16] Modified RoA estimation test for more code coverage

---
 test/roa_estimation.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/roa_estimation.jl b/test/roa_estimation.jl
index 43fb96e..8c341d4 100644
--- a/test/roa_estimation.jl
+++ b/test/roa_estimation.jl
@@ -70,8 +70,8 @@ V_func, V̇_func, ∇V_func = NumericalNeuralLyapunovFunctions(
     discretization.phi,
     res.u,
     network_func,
-    structure.V,
-    ODEFunction(f),
+    structure,
+    f,
     zeros(length(lb))
 )
 

From 8b56f31e16a97368a5c1dbcf28af1bcda19c1f77 Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Sat, 16 Mar 2024 14:42:27 -0400
Subject: [PATCH 05/16] Renamed check_fixed_point to check_minimal_fixed_point

---
 src/NeuralLyapunovPDESystem.jl  | 2 +-
 src/conditions_specification.jl | 8 ++++----
 src/minimization_conditions.jl  | 2 +-
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/NeuralLyapunovPDESystem.jl b/src/NeuralLyapunovPDESystem.jl
index a9f1fd8..2cdb73b 100644
--- a/src/NeuralLyapunovPDESystem.jl
+++ b/src/NeuralLyapunovPDESystem.jl
@@ -241,7 +241,7 @@ function _NeuralLyapunovPDESystem(
 
     bcs = []
 
-    if check_fixed_point(minimzation_condition)
+    if check_minimal_fixed_point(minimzation_condition)
         push!(bcs, V_sym(fixed_point) ~ 0.0)
     end
     if check_stationary_fixed_point(decrease_condition)
diff --git a/src/conditions_specification.jl b/src/conditions_specification.jl
index 31c1c95..7498665 100644
--- a/src/conditions_specification.jl
+++ b/src/conditions_specification.jl
@@ -63,14 +63,14 @@ function check_nonnegativity(cond::AbstractLyapunovMinimizationCondition)::Bool
 end
 
 """
-    check_fixed_point(cond::AbstractLyapunovMinimizationCondition)
+    check_minimal_fixed_point(cond::AbstractLyapunovMinimizationCondition)
 
 `true` if `cond` specifies training for the Lyapunov function to equal zero at the
 fixed point, `false` if `cond` specifies no training to meet this condition.
 """
-function check_fixed_point(cond::AbstractLyapunovMinimizationCondition)::Bool
-    error("check_fixed_point not implemented for AbstractLyapunovMinimizationCondition of" *
-          " type $(typeof(cond))")
+function check_minimal_fixed_point(cond::AbstractLyapunovMinimizationCondition)::Bool
+    error("check_minimal_fixed_point not implemented for " *
+          "AbstractLyapunovMinimizationCondition of type $(typeof(cond))")
 end
 
 """
diff --git a/src/minimization_conditions.jl b/src/minimization_conditions.jl
index f19b53b..e11af3e 100644
--- a/src/minimization_conditions.jl
+++ b/src/minimization_conditions.jl
@@ -38,7 +38,7 @@ function check_nonnegativity(cond::LyapunovMinimizationCondition)::Bool
     cond.check_nonnegativity
 end
 
-function check_fixed_point(cond::LyapunovMinimizationCondition)::Bool
+function check_minimal_fixed_point(cond::LyapunovMinimizationCondition)::Bool
     cond.check_fixed_point
 end
 

From 6b127b49356b3260a0f265e3120c39be3b2900da Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Sat, 16 Mar 2024 15:34:42 -0400
Subject: [PATCH 06/16] Renamed network variables in PDESystem

---
 src/NeuralLyapunovPDESystem.jl | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/NeuralLyapunovPDESystem.jl b/src/NeuralLyapunovPDESystem.jl
index 2cdb73b..a15218b 100644
--- a/src/NeuralLyapunovPDESystem.jl
+++ b/src/NeuralLyapunovPDESystem.jl
@@ -204,20 +204,20 @@ function _NeuralLyapunovPDESystem(
 
     ################## Define Lyapunov function & derivative ##################
     output_dim = structure.network_dim
-    net_syms = [Symbol(:u, i) for i in 1:output_dim]
+    net_syms = [Symbol(:φ, i) for i in 1:output_dim]
     net = [first(@variables $s(..)) for s in net_syms]
 
-    # u(x) is the symbolic form of neural network output
-    u(x) = Num.([ui(x...) for ui in net])
+    # φ(x) is the symbolic form of neural network output
+    φ(x) = Num.([φi(x...) for φi in net])
 
     # V_sym(x) is the symobolic form of the Lyapunov function
-    V_sym(x) = structure.V(u, x, fixed_point)
+    V_sym(x) = structure.V(φ, x, fixed_point)
 
     # V̇_sym(x) is the symbolic time derivative of the Lyapunov function
     function V̇_sym(x)
         structure.V̇(
-            u,
-            y -> Symbolics.jacobian(u(y), y),
+            φ,
+            y -> Symbolics.jacobian(φ(y), y),
             dynamics,
             x,
             params,
@@ -267,21 +267,21 @@ function _NeuralLyapunovPDESystem(
         bcs,
         domains,
         state,
-        u(state),
+        φ(state),
         params;
         defaults = defaults
     )
 
     ################### Return PDESystem and neural network ###################
-    # u_func is the numerical form of neural network output
-    function u_func(phi, θ, x)
+    # φ_func is the numerical form of neural network output
+    function φ_func(phi, θ, x)
         reduce(
             vcat,
             Array(phi[i](x, θ.depvar[net_syms[i]])) for i in 1:output_dim
         )
     end
 
-    return lyapunov_pde_system, u_func
+    return lyapunov_pde_system, φ_func
 end
 
 """

From 9e06f1c3ede0fe38b27df0527c7f50866f3a493a Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Sat, 16 Mar 2024 16:08:16 -0400
Subject: [PATCH 07/16] Removed check_stationary_fixed_point

---
 src/NeuralLyapunovPDESystem.jl       |  3 ---
 src/conditions_specification.jl      | 13 +-----------
 src/decrease_conditions.jl           | 30 ++++++++--------------------
 src/decrease_conditions_RoA_aware.jl | 13 ++----------
 4 files changed, 11 insertions(+), 48 deletions(-)

diff --git a/src/NeuralLyapunovPDESystem.jl b/src/NeuralLyapunovPDESystem.jl
index a15218b..7096e90 100644
--- a/src/NeuralLyapunovPDESystem.jl
+++ b/src/NeuralLyapunovPDESystem.jl
@@ -244,9 +244,6 @@ function _NeuralLyapunovPDESystem(
     if check_minimal_fixed_point(minimzation_condition)
         push!(bcs, V_sym(fixed_point) ~ 0.0)
     end
-    if check_stationary_fixed_point(decrease_condition)
-        push!(bcs, V̇_sym(fixed_point) ~ 0.0)
-    end
 
     if isempty(eqs) && isempty(bcs)
         error("No training conditions specified.")
diff --git a/src/conditions_specification.jl b/src/conditions_specification.jl
index 7498665..e0044e7 100644
--- a/src/conditions_specification.jl
+++ b/src/conditions_specification.jl
@@ -36,7 +36,7 @@ abstract type AbstractLyapunovMinimizationCondition end
 Represents the decrease condition in a neural Lyapunov problem
 
 All concrete `AbstractLyapunovDecreaseCondition` subtypes should define the
-`check_decrease`, `check_stationary_fixed_point`, and `get_decrease_condition` functions.
+`check_decrease` and `get_decrease_condition` functions.
 """
 abstract type AbstractLyapunovDecreaseCondition end
 
@@ -95,17 +95,6 @@ function check_decrease(cond::AbstractLyapunovDecreaseCondition)::Bool
           string(typeof(cond)))
 end
 
-"""
-    check_stationary_fixed_point(cond::AbstractLyapunovDecreaseCondition)
-
-`true` if `cond` specifies training for the Lyapunov function not to change at the
-fixed point, `false` if `cond` specifies no training to meet this condition.
-"""
-function check_stationary_fixed_point(cond::AbstractLyapunovDecreaseCondition)::Bool
-    error("check_fixed_point not implemented for AbstractLyapunovDecreaseCondition of " *
-          "type $(typeof(cond))")
-end
-
 """
     get_decrease_condition(cond::AbstractLyapunovDecreaseCondition)
 
diff --git a/src/decrease_conditions.jl b/src/decrease_conditions.jl
index 48353d4..f173092 100644
--- a/src/decrease_conditions.jl
+++ b/src/decrease_conditions.jl
@@ -1,13 +1,11 @@
 """
-    LyapunovDecreaseCondition(check_decrease, decrease, strength, relu, check_fixed_point)
+    LyapunovDecreaseCondition(check_decrease, decrease, strength, relu)
 
 Specifies the form of the Lyapunov conditions to be used; if `check_decrease`, training will
 enforce `decrease(V, dVdt) ≤ strength(state, fixed_point)`.
 
 The inequality will be approximated by the equation
     `relu(decrease(V, dVdt) - strength(state, fixed_point)) = 0.0`.
-If `check_fixed_point` is `false`, then training assumes `dVdt(fixed_point) = 0`, but
-if `check_fixed_point` is `true`, then training will enforce `dVdt(fixed_point) = 0`.
 
 If the dynamics truly have a fixed point at `fixed_point` and `dVdt` has been defined
 properly in terms of the dynamics, then `dVdt(fixed_point)` will be `0` and there is no need
@@ -30,17 +28,12 @@ struct LyapunovDecreaseCondition <: AbstractLyapunovDecreaseCondition
     decrease::Function
     strength::Function
     relu::Function
-    check_fixed_point::Bool
 end
 
 function check_decrease(cond::LyapunovDecreaseCondition)::Bool
     cond.check_decrease
 end
 
-function check_stationary_fixed_point(cond::LyapunovDecreaseCondition)::Bool
-    cond.check_fixed_point
-end
-
 function get_decrease_condition(cond::LyapunovDecreaseCondition)
     if cond.check_decrease
         return (V, dVdt, x, fixed_point) -> cond.relu(
@@ -52,7 +45,7 @@ function get_decrease_condition(cond::LyapunovDecreaseCondition)
 end
 
 """
-    AsymptoticDecrease(strict; check_fixed_point, C)
+    AsymptoticDecrease(; strict, C, relu)
 
 Constructs a `LyapunovDecreaseCondition` corresponding to asymptotic decrease.
 
@@ -63,7 +56,6 @@ The inequality is represented by `a ≥ b` <==> `relu(b-a) = 0.0`.
 """
 function AsymptoticDecrease(;
         strict::Bool = false,
-        check_fixed_point::Bool = false,
         C::Real = 1e-6,
         relu = (t) -> max(0.0, t)
 )::LyapunovDecreaseCondition
@@ -77,13 +69,12 @@ function AsymptoticDecrease(;
         true,
         (V, dVdt) -> dVdt,
         strength,
-        relu,
-        check_fixed_point
+        relu
     )
 end
 
 """
-    ExponentialDecrease(k, strict; check_fixed_point, C)
+    ExponentialDecrease(k; strict, C, relu)
 
 Constructs a `LyapunovDecreaseCondition` corresponding to exponential decrease of rate `k`.
 
@@ -95,7 +86,6 @@ The inequality is represented by `a ≥ b` <==> `relu(b-a) = 0.0`.
 function ExponentialDecrease(
         k::Real;
         strict::Bool = false,
-        check_fixed_point::Bool = false,
         C::Real = 1e-6,
         relu = (t) -> max(0.0, t)
 )::LyapunovDecreaseCondition
@@ -109,26 +99,22 @@ function ExponentialDecrease(
         true,
         (V, dVdt) -> dVdt + k * V,
         strength,
-        relu,
-        check_fixed_point
+        relu
     )
 end
 
 """
-    DontCheckDecrease(check_fixed_point = false)
+    DontCheckDecrease()
 
 Constructs a `LyapunovDecreaseCondition` which represents not checking for
 decrease of the Lyapunov function along system trajectories. This is appropriate
 in cases when the Lyapunov decrease condition has been structurally enforced.
-
-It is still possible to check for `dV/dt = 0` at `fixed_point`, even in this case.
 """
-function DontCheckDecrease(check_fixed_point::Bool = false)::LyapunovDecreaseCondition
+function DontCheckDecrease()::LyapunovDecreaseCondition
     return LyapunovDecreaseCondition(
         false,
         (V, dVdt) -> 0.0,
         (state, fixed_point) -> 0.0,
-        (t) -> 0.0,
-        check_fixed_point
+        (t) -> 0.0
     )
 end
diff --git a/src/decrease_conditions_RoA_aware.jl b/src/decrease_conditions_RoA_aware.jl
index d6a9de3..fba7e1f 100644
--- a/src/decrease_conditions_RoA_aware.jl
+++ b/src/decrease_conditions_RoA_aware.jl
@@ -1,6 +1,6 @@
 """
-    RoAAwareDecreaseCondition(check_decrease, decrease, strength, relu, check_fixed_point,
-                              ρ, out_of_RoA_penalty)
+    RoAAwareDecreaseCondition(check_decrease, decrease, strength, relu, ρ,
+                              out_of_RoA_penalty)
 
 Specifies the form of the Lyapunov conditions to be used, training for a region of
 attraction estimate of `{ x : V(x) ≤ ρ }`
@@ -12,9 +12,6 @@ and will instead apply
 
 The inequality will be approximated by the equation
     `relu(decrease(V, dVdt) - strength(state, fixed_point)) = 0.0`.
-If `check_fixed_point` is `false`, then training assumes `dVdt(fixed_point) = 0`, but
-if `check_fixed_point` is `true`, then training will attempt to enforce
-`dVdt(fixed_point) = 0`.
 
 If the dynamics truly have a fixed point at `fixed_point` and `dVdt` has been defined
 properly in terms of the dynamics, then `dVdt(fixed_point)` will be `0` and there is no need
@@ -44,7 +41,6 @@ struct RoAAwareDecreaseCondition <: AbstractLyapunovDecreaseCondition
     decrease::Function
     strength::Function
     relu::Function
-    check_fixed_point::Bool
     ρ::Real
     out_of_RoA_penalty::Function
 end
@@ -53,10 +49,6 @@ function check_decrease(cond::RoAAwareDecreaseCondition)::Bool
     cond.check_decrease
 end
 
-function check_stationary_fixed_point(cond::RoAAwareDecreaseCondition)::Bool
-    cond.check_fixed_point
-end
-
 function get_decrease_condition(cond::RoAAwareDecreaseCondition)
     if cond.check_decrease
         return function (V, dVdt, x, fixed_point)
@@ -92,7 +84,6 @@ function make_RoA_aware(
         cond.decrease,
         cond.strength,
         cond.relu,
-        cond.check_fixed_point,
         ρ,
         out_of_RoA_penalty
     )

From c49364c84bdda0e1b94210a1749848b47981e744 Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Sat, 16 Mar 2024 19:47:41 -0400
Subject: [PATCH 08/16] Added policy search capabilities

---
 Project.toml                    |   3 +-
 src/NeuralLyapunov.jl           |   4 +-
 src/NeuralLyapunovPDESystem.jl  |  20 ++-
 src/conditions_specification.jl |   7 +-
 src/policy_search.jl            |  65 +++++++++
 src/structure_specification.jl  |  13 ++
 test/inverted_pendulum.jl       | 246 ++++++++++++++++++++++++++++++++
 test/runtests.jl                |   3 +
 8 files changed, 353 insertions(+), 8 deletions(-)
 create mode 100644 src/policy_search.jl
 create mode 100644 test/inverted_pendulum.jl

diff --git a/Project.toml b/Project.toml
index fe72351..e950ba6 100644
--- a/Project.toml
+++ b/Project.toml
@@ -24,6 +24,7 @@ SciMLBase = "2"
 julia = "1.10"
 
 [extras]
+DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 NLopt = "76087f3c-5699-56af-9a33-bf431cd00edd"
 NeuralPDE = "315f7962-48a3-4962-8226-d0f33b1235f0"
@@ -35,4 +36,4 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["SafeTestsets", "Test", "Lux", "Optimization", "OptimizationOptimJL", "OptimizationOptimisers", "NLopt", "Random", "NeuralPDE"]
+test = ["SafeTestsets", "Test", "Lux", "Optimization", "OptimizationOptimJL", "OptimizationOptimisers", "NLopt", "Random", "NeuralPDE", "DifferentialEquations"]
diff --git a/src/NeuralLyapunov.jl b/src/NeuralLyapunov.jl
index 78171b2..b6e41b7 100644
--- a/src/NeuralLyapunov.jl
+++ b/src/NeuralLyapunov.jl
@@ -13,6 +13,7 @@ include("decrease_conditions.jl")
 include("decrease_conditions_RoA_aware.jl")
 include("NeuralLyapunovPDESystem.jl")
 include("local_Lyapunov.jl")
+include("policy_search.jl")
 
 export NeuralLyapunovPDESystem, NumericalNeuralLyapunovFunctions
 export local_Lyapunov
@@ -20,6 +21,7 @@ export NeuralLyapunovSpecification, NeuralLyapunovStructure, UnstructuredNeuralL
        NonnegativeNeuralLyapunov, PositiveSemiDefiniteStructure,
        LyapunovMinimizationCondition, StrictlyPositiveDefinite, PositiveSemiDefinite,
        DontCheckNonnegativity, LyapunovDecreaseCondition, AsymptoticDecrease,
-       ExponentialDecrease, DontCheckDecrease, RoAAwareDecreaseCondition, make_RoA_aware
+       ExponentialDecrease, DontCheckDecrease, RoAAwareDecreaseCondition, make_RoA_aware,
+       add_policy_search, get_policy
 
 end
diff --git a/src/NeuralLyapunovPDESystem.jl b/src/NeuralLyapunovPDESystem.jl
index 7096e90..51391e8 100644
--- a/src/NeuralLyapunovPDESystem.jl
+++ b/src/NeuralLyapunovPDESystem.jl
@@ -35,7 +35,8 @@ function NeuralLyapunovPDESystem(
         fixed_point = zeros(length(lb)),
         p = SciMLBase.NullParameters(),
         state_syms = [],
-        parameter_syms = []
+        parameter_syms = [],
+        policy_search::Bool = false
 )::Tuple{PDESystem, Function}
     ########################## Define state symbols ###########################
     state_dim = length(lb)
@@ -78,7 +79,8 @@ function NeuralLyapunovPDESystem(
         fixed_point,
         state,
         params,
-        defaults
+        defaults,
+        policy_search
     )
 end
 
@@ -119,7 +121,8 @@ function NeuralLyapunovPDESystem(
         fixed_point = fixed_point,
         p = p,
         state_syms = SciMLBase.variable_symbols(dynamics.sys),
-        parameter_syms = p_syms
+        parameter_syms = p_syms,
+        policy_search = false
     )
 end
 
@@ -179,7 +182,8 @@ function NeuralLyapunovPDESystem(
         fixed_point,
         state,
         Num.(parameters(dynamics)),
-        ModelingToolkit.get_defaults(dynamics)
+        ModelingToolkit.get_defaults(dynamics),
+        false
     )
 end
 
@@ -191,12 +195,14 @@ function _NeuralLyapunovPDESystem(
         fixed_point,
         state,
         params,
-        defaults
+        defaults,
+        policy_search::Bool
 )::Tuple{PDESystem, Function}
     ########################## Unpack specifications ##########################
     structure = spec.structure
     minimzation_condition = spec.minimzation_condition
     decrease_condition = spec.decrease_condition
+    f_call = structure.f_call
 
     ############################# Define domains ##############################
     state_dim = length(lb)
@@ -245,6 +251,10 @@ function _NeuralLyapunovPDESystem(
         push!(bcs, V_sym(fixed_point) ~ 0.0)
     end
 
+    if policy_search
+        append!(bcs, f_call(dynamics, φ, fixed_point, params, 0.0) .~ zeros(state_dim))
+    end
+
     if isempty(eqs) && isempty(bcs)
         error("No training conditions specified.")
     end
diff --git a/src/conditions_specification.jl b/src/conditions_specification.jl
index e0044e7..010e7b8 100644
--- a/src/conditions_specification.jl
+++ b/src/conditions_specification.jl
@@ -5,18 +5,23 @@ Specifies the structure of the neural Lyapunov function and its derivative.
 
 Allows the user to define the Lyapunov in terms of the neural network to structurally
 enforce Lyapunov conditions.
-`network_dim` is the dimension of the output of the neural network.
 `V(phi::Function, state, fixed_point)` takes in the neural network, the state, and the fixed
 point, and outputs the value of the Lyapunov function at `state`.
 `V̇(phi::Function, J_phi::Function, f::Function, state, params, t, fixed_point)` takes in the
 neural network, jacobian of the neural network, dynamics, state, parameters and time (for
 calling the dynamics, when relevant), and fixed point, and outputs the time derivative of
 the Lyapunov function at `state`.
+`f_call(dynamics::Function, phi::Function, state, p, t)` takes in the dynamics, the neural
+network, the state, the parameters of the dynamics, and time, and outputs the derivative of
+the state; this is useful for making closed-loop dynamics which depend on the neural
+network, such as in the policy search case
+`network_dim` is the dimension of the output of the neural network.
 """
 struct NeuralLyapunovStructure
     V::Function
     ∇V::Function
     V̇::Function
+    f_call::Function
     network_dim::Integer
 end
 
diff --git a/src/policy_search.jl b/src/policy_search.jl
new file mode 100644
index 0000000..2c81869
--- /dev/null
+++ b/src/policy_search.jl
@@ -0,0 +1,65 @@
+"""
+    add_policy_search(lyapunov_structure, new_dims, control_structure)
+
+Adds dependence on the neural network to the dynamics in a `NeuralLyapunovStructure`
+
+Adds `new_dims` outputs to the neural network and feeds them through `control_structure` to
+calculatethe contribution of the neural network to the dynamics.
+The existing `lyapunov_structure.network_dim` dimensions are used as in `lyapunov_structure`
+to calculate the Lyapunov function.
+
+`lyapunov_structure` should assume in its `V̇` that the dynamics take a form `f(x, p, t)`.
+The returned `NeuralLyapunovStructure` will assume instead `f(x, u, p, t)`, where `u` is the
+contribution from the neural network. Therefore, this structure cannot be used with a
+`NeuralLyapunovPDESystem` method that requires an `ODEFunction`, `ODESystem`, or
+`ODEProblem`.
+"""
+function add_policy_search(
+        lyapunov_structure::NeuralLyapunovStructure,
+        new_dims::Integer;
+        control_structure::Function = identity
+)::NeuralLyapunovStructure
+    let V = lyapunov_structure.V, ∇V = lyapunov_structure.∇V, V̇ = lyapunov_structure.V̇,
+        V_dim = lyapunov_structure.network_dim, nd = new_dims, u = control_structure
+
+        NeuralLyapunovStructure(
+            function (net, state, fixed_point)
+                if length(size(state)) == 1
+                    if V_dim == 1
+                        V(st -> net(st)[1], state, fixed_point)
+                    else
+                        V(st -> net(st)[1:V_dim], state, fixed_point)
+                    end
+                else
+                    V(st -> net(st)[1:V_dim, :], state, fixed_point)
+                end
+            end,
+            function (net, J_net, state, fixed_point)
+                ∇V(st -> net(st)[1:V_dim], st -> J_net(st)[1:V_dim, :], state, fixed_point)
+            end,
+            function (net, J_net, f, state, params, t, fixed_point)
+                V̇(st -> net(st)[1:V_dim], st -> J_net(st)[1:V_dim, :],
+                    (st, p, t) -> f(st, u(net(st)[(V_dim + 1):end]), p, t), state, params,
+                    t, fixed_point)
+            end,
+            (f, net, state, p, t) -> f(state, u(net(state)[(V_dim + 1):end]), p, t),
+            V_dim + nd
+        )
+    end
+end
+
+function get_policy(
+        phi,
+        θ,
+        network_func::Function,
+        dim_u::Integer;
+        u_func::Function = identity
+)
+    function policy(state::AbstractVector)
+        u_func(network_func(phi, θ, state)[(end - dim_u + 1):end])
+    end
+
+    policy(state::AbstractMatrix) = mapslices(policy, state, dims = [1])
+
+    return policy
+end
diff --git a/src/structure_specification.jl b/src/structure_specification.jl
index 20c23db..b1f6af5 100644
--- a/src/structure_specification.jl
+++ b/src/structure_specification.jl
@@ -3,6 +3,9 @@
 
 Creates a `NeuralLyapunovStructure` where the Lyapunov function is the neural network
 evaluated at the state. This does not structurally enforce any Lyapunov conditions.
+
+Dynamics are assumed to be in `f(state, p, t)` form, as in an `ODEFunction`. For
+`f(state, input, p, t)`, consider using `add_policy_search`.
 """
 function UnstructuredNeuralLyapunov()::NeuralLyapunovStructure
     NeuralLyapunovStructure(
@@ -10,6 +13,7 @@ function UnstructuredNeuralLyapunov()::NeuralLyapunovStructure
         (net, grad_net, state, fixed_point) -> grad_net(state),
         (net, grad_net, f, state, params, t, fixed_point) -> grad_net(state) ⋅
                                                              f(state, params, t),
+        (f, net, state, p, t) -> f(state, p, t),
         1
     )
 end
@@ -33,6 +37,9 @@ Lyapunov loss function.
 defaults to `ForwardDiff.gradient`.
 
 The neural network output has dimension `network_dim`.
+
+Dynamics are assumed to be in `f(state, p, t)` form, as in an `ODEFunction`. For
+`f(state, input, p, t)`, consider using `add_policy_search`.
 """
 function NonnegativeNeuralLyapunov(
         network_dim::Integer;
@@ -50,6 +57,7 @@ function NonnegativeNeuralLyapunov(
             (net, J_net, f, state, params, t, fixed_point) -> 2 *
                                                               dot(
                 net(state), J_net(state), f(state, params, t)),
+            (f, net, state, p, t) -> f(state, p, t),
             network_dim
         )
     else
@@ -68,6 +76,7 @@ function NonnegativeNeuralLyapunov(
                 J_net(state),
                 f(state, params, t)
             ) + δ * grad_pos_def(state, fixed_point) ⋅ f(state, params, t),
+            (f, net, state, p, t) -> f(state, p, t),
             network_dim
         )
     end
@@ -94,6 +103,9 @@ gradient of `non_neg(net, state, fixed_point)` with respect to `state` at `state
 defaults to `ForwardDiff.gradient`.
 
 The neural network output has dimension `network_dim`.
+
+Dynamics are assumed to be in `f(state, p, t)` form, as in an `ODEFunction`. For
+`f(state, input, p, t)`, consider using `add_policy_search`.
 """
 function PositiveSemiDefiniteStructure(
         network_dim::Integer;
@@ -133,6 +145,7 @@ function PositiveSemiDefiniteStructure(
                                                           (f(state, params, t) ⋅
                                                            grad_non_neg(
             net, J_net, state, fixed_point)),
+        (f, net, state, p, t) -> f(state, p, t),
         network_dim
     )
 end
diff --git a/test/inverted_pendulum.jl b/test/inverted_pendulum.jl
new file mode 100644
index 0000000..f35a87e
--- /dev/null
+++ b/test/inverted_pendulum.jl
@@ -0,0 +1,246 @@
+using LinearAlgebra
+using NeuralPDE, Lux, ModelingToolkit
+using Optimization, OptimizationOptimisers, OptimizationOptimJL, NLopt
+using NeuralLyapunov
+using Random
+using Test
+
+Random.seed!(200)
+
+println("Inverted Pendulum - Policy Search")
+
+######################### Define dynamics and domain ##########################
+
+function controlled_pendulum_dynamics(x, u, p, t)
+    θ, ω = x
+    ζ, ω_0 = p
+    τ = u[]
+    return [ω
+            -2ζ * ω - ω_0^2 * sin(θ) + τ]
+end
+
+lb = [0.0, -10.0];
+ub = [2π, 10.0];
+upright_equilibrium = [π, 0.0]
+p = [0.5, 1.0]
+state_syms = [:θ, :ω]
+parameter_syms = [:ζ, :ω_0]
+
+####################### Specify neural Lyapunov problem #######################
+
+# Define neural network discretization
+# We use an input layer that is periodic with period 2π with respect to θ
+dim_state = length(lb)
+dim_hidden = 15
+dim_phi = 2
+dim_u = 1
+dim_output = dim_phi + dim_u
+chain = [Lux.Chain(
+             Lux.WrappedFunction(x -> vcat(
+                 transpose(sin.(x[1, :])),
+                 transpose(cos.(x[1, :])),
+                 transpose(x[2, :])
+             )),
+             Dense(3, dim_hidden, tanh),
+             Dense(dim_hidden, dim_hidden, tanh),
+             Dense(dim_hidden, 1, use_bias = false)
+         ) for _ in 1:dim_output]
+
+# Define neural network discretization
+strategy = GridTraining(0.1)
+discretization = PhysicsInformedNN(chain, strategy)
+
+# Define neural Lyapunov structure
+structure = PositiveSemiDefiniteStructure(
+    dim_phi;
+    pos_def = function (state, fixed_point)
+        θ, ω = state
+        θ_eq, ω_eq = fixed_point
+        log(1.0 + (sin(θ) - sin(θ_eq))^2 + (cos(θ) - cos(θ_eq))^2 + (ω - ω_eq)^2)
+    end
+)
+structure = add_policy_search(
+    structure,
+    dim_u
+)
+minimization_condition = DontCheckNonnegativity(check_fixed_point = false)
+
+# Define Lyapunov decrease condition
+decrease_condition = AsymptoticDecrease(strict = true)
+
+# Construct neural Lyapunov specification
+spec = NeuralLyapunovSpecification(
+    structure,
+    minimization_condition,
+    decrease_condition
+)
+
+############################# Construct PDESystem #############################
+
+pde_system, network_func = NeuralLyapunovPDESystem(
+    controlled_pendulum_dynamics,
+    lb,
+    ub,
+    spec;
+    fixed_point = upright_equilibrium,
+    p = p,
+    state_syms = state_syms,
+    parameter_syms = parameter_syms,
+    policy_search = true
+)
+
+######################## Construct OptimizationProblem ########################
+
+sym_prob = symbolic_discretize(pde_system, discretization)
+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, BFGS(); maxiters = 300)
+
+###################### Get numerical numerical functions ######################
+
+V_func, V̇_func, ∇V_func = NumericalNeuralLyapunovFunctions(
+    discretization.phi,
+    res.u,
+    network_func,
+    structure,
+    controlled_pendulum_dynamics,
+    upright_equilibrium;
+    p = p
+)
+
+u = get_policy(discretization.phi, res.u, network_func, dim_u)
+
+################################## Simulate ###################################
+
+xs = (-2π):0.02:(2π)
+ys = lb[2]:0.02:ub[2]
+states = Iterators.map(collect, Iterators.product(xs, ys))
+V_predict = vec(V_func(hcat(states...)))
+dVdt_predict = vec(V̇_func(hcat(states...)))
+
+#################################### Tests ####################################
+
+# Network structure should enforce positive definiteness
+@test V_func(upright_equilibrium) == 0.0
+@test min(V_func(upright_equilibrium), minimum(V_predict)) ≥ 0.0
+
+# Network structure should enforce periodicity in θ
+x0 = (ub .- lb) .* rand(2, 100) .+ lb
+@test all(isapprox.(V_func(x0), V_func(x0 .+ [2π, 0.0]); rtol = 1e-3))
+
+# Training should result in a fixed point at the upright equilibrium
+@test all(isapprox.(
+    controlled_pendulum_dynamics(upright_equilibrium, u(upright_equilibrium), p, 0.0),
+    0.0; atol = 1e-8))
+@test V̇_func(upright_equilibrium) == 0.0
+
+# V̇ should be negative almost everywhere
+@test sum(dVdt_predict .> 0) / length(dVdt_predict) < 1e-3
+
+################################## Simulate ###################################
+
+using DifferentialEquations
+
+closed_loop_dynamics = ODEFunction(
+    (x, p, t) -> controlled_pendulum_dynamics(x, u(x), p, t);
+    sys = SciMLBase.SymbolCache(state_syms, parameter_syms)
+)
+
+# Starting still at bottom
+downward_equilibrium = zeros(2)
+ode_prob = ODEProblem(closed_loop_dynamics, downward_equilibrium, [0.0, 35.0], p)
+sol = solve(ode_prob, Tsit5())
+# plot(sol)
+
+# Should make it to the top
+θ_end, ω_end = sol.u[end]
+x_end, y_end = sin(θ_end), -cos(θ_end)
+@test all(isapprox.([x_end, y_end, ω_end], [0.0, 1.0, 0.0]; atol = 1e-3))
+
+# Starting at a random point
+x0 = lb .+ rand(2) .* (ub .- lb)
+ode_prob = ODEProblem(closed_loop_dynamics, x0, [0.0, 20.0], p)
+sol = solve(ode_prob, Tsit5())
+# plot(sol)
+
+# Should make it to the top
+θ_end, ω_end = sol.u[end]
+x_end, y_end = sin(θ_end), -cos(θ_end)
+@test all(isapprox.([x_end, y_end, ω_end], [0.0, 1.0, 0.0]; atol = 1e-3))
+
+#=
+# Print statistics
+println("V(π, 0) = ", V_func(upright_equilibrium))
+println(
+    "f([π, 0], u([π, 0])) = ",
+    controlled_pendulum_dynamics(upright_equilibrium, u(upright_equilibrium), p, 0.0)
+)
+println(
+    "V ∋ [",
+    min(V_func(upright_equilibrium),
+    minimum(V_predict)),
+    ", ",
+    maximum(V_predict),
+    "]"
+)
+println(
+    "V̇ ∋ [",
+    minimum(dVdt_predict),
+    ", ",
+    max(V̇_func(upright_equilibrium), maximum(dVdt_predict)),
+    "]"
+)
+
+# Plot results
+using Plots
+
+p1 = plot(
+    xs / pi,
+    ys,
+    V_predict,
+    linetype =
+    :contourf,
+    title = "V",
+    xlabel = "θ/π",
+    ylabel = "ω",
+    c = :bone_1
+);
+p1 = scatter!([-2 * pi, 0, 2 * pi] / pi, [0, 0, 0],
+    label = "Downward Equilibria", color = :green, markershape = :+);
+p1 = scatter!(
+    [-pi, pi] / pi, [0, 0], label = "Upward Equilibria", color = :red, markershape = :x);
+p2 = plot(
+    xs / pi,
+    ys,
+    dVdt_predict,
+    linetype = :contourf,
+    title = "dV/dt",
+    xlabel = "θ/π",
+    ylabel = "ω",
+    c = :binary
+);
+p2 = scatter!([-2 * pi, 0, 2 * pi] / pi, [0, 0, 0],
+    label = "Downward Equilibria", color = :green, markershape = :+);
+p2 = scatter!([-pi, pi] / pi, [0, 0], label = "Upward Equilibria",
+    color = :red, markershape = :x, legend = false);
+p3 = plot(
+    xs / pi,
+    ys,
+    dVdt_predict .< 0,
+    linetype = :contourf,
+    title = "dV/dt < 0",
+    xlabel = "θ/π",
+    ylabel = "ω",
+    colorbar = false,
+    linewidth = 0
+);
+p3 = scatter!([-2 * pi, 0, 2 * pi] / pi, [0, 0, 0],
+    label = "Downward Equilibria", color = :green, markershape = :+);
+p3 = scatter!([-pi, pi] / pi, [0, 0], label = "Upward Equilibria",
+    color = :red, markershape = :x, legend = false);
+plot(p1, p2, p3)
+=#
diff --git a/test/runtests.jl b/test/runtests.jl
index ce449d9..fae3962 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -10,4 +10,7 @@ using SafeTestsets
     @time @safetestset "Region of attraction estimation" begin
         include("roa_estimation.jl")
     end
+    @time @safetestset "Policy search - inverted pendulum" begin
+        include("inverted_pendulum.jl")
+    end
 end

From dd9618960027b9bf4b6fa56b53439458f46eec00 Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Sat, 16 Mar 2024 19:57:48 -0400
Subject: [PATCH 09/16] Added documentation for get_policy

---
 src/policy_search.jl | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/policy_search.jl b/src/policy_search.jl
index 2c81869..80b055e 100644
--- a/src/policy_search.jl
+++ b/src/policy_search.jl
@@ -48,15 +48,28 @@ function add_policy_search(
     end
 end
 
+"""
+    get_policy(phi, θ, network_func, dim; control_structure)
+
+Returns the control policy as a function of the state
+
+The returned function can operate on a state vector or columnwise on a matrix of state
+vectors.
+
+`phi` is the neural network with parameters `θ`. `network_func` is an output of
+`NeuralLyapunovPDESystem`.
+The control policy is `control_structure` composed with the last `dim` outputs of the
+neural network, as set up by `add_policy_search`.
+"""
 function get_policy(
         phi,
         θ,
         network_func::Function,
-        dim_u::Integer;
-        u_func::Function = identity
+        dim::Integer;
+        control_structure::Function = identity
 )
     function policy(state::AbstractVector)
-        u_func(network_func(phi, θ, state)[(end - dim_u + 1):end])
+        control_structure(network_func(phi, θ, state)[(end - dim + 1):end])
     end
 
     policy(state::AbstractMatrix) = mapslices(policy, state, dims = [1])

From b2453d1703229db16cebcd4ff62e3b7ec24a8fb1 Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Sat, 16 Mar 2024 20:03:37 -0400
Subject: [PATCH 10/16] Updated documenation on NeuralLyapunovPDESystem

---
 src/NeuralLyapunovPDESystem.jl | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/NeuralLyapunovPDESystem.jl b/src/NeuralLyapunovPDESystem.jl
index 51391e8..0974cee 100644
--- a/src/NeuralLyapunovPDESystem.jl
+++ b/src/NeuralLyapunovPDESystem.jl
@@ -9,11 +9,13 @@ columnwise.
 The neural Lyapunov function will only be trained for `{ x : lb .≤ x .≤ ub }`. The Lyapunov
 function will be for the dynamical system represented by `dynamics`. If `dynamics` is an
 `ODEProblem` or `ODEFunction`, then the corresponding ODE; if `dynamics` is a function, then
-the ODE is typically `ẋ = dynamics(x, p, t)`, but `spec` may specify a different form. This
-ODE should not depend on `t` (time `t=0.0` alone will be used) and should have a fixed point
-at `x = fixed_point`. The particular Lyapunov conditions to be used and structure of the
-neural Lyapunov function are specified through `spec`, which is a
-`NeuralLyapunovSpecification`.
+the ODE is typically `ẋ = dynamics(x, p, t)`, but `spec` may specify a different form. In
+particular, if `policy_search`, then the ODE should have a control input:
+`ẋ = dynamics(x, input, p, t)`. In that case, the loss function will include a term
+penalizing nonzero `ẋ` at the fixed point. Otherwise, the ODE should have a fixed point
+at `x = fixed_point`. Additionally, the ODE should not depend on `t` (time `t=0.0` alone
+will be used). The particular Lyapunov conditions to be used and structure of the neural
+Lyapunov function are specified through `spec`, which is a `NeuralLyapunovSpecification`.
 
 The returned neural network function takes three inputs: the neural network structure `phi`,
 the trained network parameters, and a matrix of inputs to operate on columnwise.

From dbd4956a52d963dbf8ba9d275f0b3185f69b4333 Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Mon, 18 Mar 2024 17:42:35 -0400
Subject: [PATCH 11/16] Renamed inverted pendulum dynamics

---
 test/inverted_pendulum.jl | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/test/inverted_pendulum.jl b/test/inverted_pendulum.jl
index f35a87e..54dfbf0 100644
--- a/test/inverted_pendulum.jl
+++ b/test/inverted_pendulum.jl
@@ -11,7 +11,7 @@ println("Inverted Pendulum - Policy Search")
 
 ######################### Define dynamics and domain ##########################
 
-function controlled_pendulum_dynamics(x, u, p, t)
+function open_loop_pendulum_dynamics(x, u, p, t)
     θ, ω = x
     ζ, ω_0 = p
     τ = u[]
@@ -78,7 +78,7 @@ spec = NeuralLyapunovSpecification(
 ############################# Construct PDESystem #############################
 
 pde_system, network_func = NeuralLyapunovPDESystem(
-    controlled_pendulum_dynamics,
+    open_loop_pendulum_dynamics,
     lb,
     ub,
     spec;
@@ -107,7 +107,7 @@ V_func, V̇_func, ∇V_func = NumericalNeuralLyapunovFunctions(
     res.u,
     network_func,
     structure,
-    controlled_pendulum_dynamics,
+    open_loop_pendulum_dynamics,
     upright_equilibrium;
     p = p
 )
@@ -134,7 +134,7 @@ x0 = (ub .- lb) .* rand(2, 100) .+ lb
 
 # Training should result in a fixed point at the upright equilibrium
 @test all(isapprox.(
-    controlled_pendulum_dynamics(upright_equilibrium, u(upright_equilibrium), p, 0.0),
+    open_loop_pendulum_dynamics(upright_equilibrium, u(upright_equilibrium), p, 0.0),
     0.0; atol = 1e-8))
 @test V̇_func(upright_equilibrium) == 0.0
 
@@ -146,7 +146,7 @@ x0 = (ub .- lb) .* rand(2, 100) .+ lb
 using DifferentialEquations
 
 closed_loop_dynamics = ODEFunction(
-    (x, p, t) -> controlled_pendulum_dynamics(x, u(x), p, t);
+    (x, p, t) -> open_loop_pendulum_dynamics(x, u(x), p, t);
     sys = SciMLBase.SymbolCache(state_syms, parameter_syms)
 )
 
@@ -177,7 +177,7 @@ x_end, y_end = sin(θ_end), -cos(θ_end)
 println("V(π, 0) = ", V_func(upright_equilibrium))
 println(
     "f([π, 0], u([π, 0])) = ",
-    controlled_pendulum_dynamics(upright_equilibrium, u(upright_equilibrium), p, 0.0)
+    open_loop_pendulum_dynamics(upright_equilibrium, u(upright_equilibrium), p, 0.0)
 )
 println(
     "V ∋ [",

From 18a859d9fbead3ebb1fb35f7bee5e31445302415 Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Wed, 20 Mar 2024 15:28:10 -0400
Subject: [PATCH 12/16] Added policy search support for ODESystems with unbound
 inputs

---
 src/NeuralLyapunovPDESystem.jl      |  23 ++-
 test/inverted_pendulum_ODESystem.jl | 255 ++++++++++++++++++++++++++++
 test/runtests.jl                    |   3 +
 3 files changed, 275 insertions(+), 6 deletions(-)
 create mode 100644 test/inverted_pendulum_ODESystem.jl

diff --git a/src/NeuralLyapunovPDESystem.jl b/src/NeuralLyapunovPDESystem.jl
index 0974cee..dca6d5e 100644
--- a/src/NeuralLyapunovPDESystem.jl
+++ b/src/NeuralLyapunovPDESystem.jl
@@ -167,25 +167,32 @@ function NeuralLyapunovPDESystem(
         spec::NeuralLyapunovSpecification;
         fixed_point = zeros(length(lb))
 )::Tuple{PDESystem, Function}
+    ######################### Check for policy search #########################
+    f, x, p, policy_search = if isempty(ModelingToolkit.unbound_inputs(dynamics))
+        (ODEFunction(dynamics), states(dynamics), parameters(dynamics), false)
+    else
+        (f, _), x, p = ModelingToolkit.generate_control_function(dynamics; simplify = true)
+        (f, x, p, true)
+    end
+
     ########################## Define state symbols ###########################
-    state = states(dynamics)
     # States should all be functions of time, but we just want the symbol
     # e.g., if the state is ω(t), we just want ω
-    state = map(st -> istree(st) ? operation(st) : st, state)
+    state = map(st -> istree(st) ? operation(st) : st, x)
     state_syms = Symbol.(state)
     state = [first(@parameters $s) for s in state_syms]
 
     ########################### Construct PDESystem ###########################
     _NeuralLyapunovPDESystem(
-        ODEFunction(dynamics),
+        f,
         lb,
         ub,
         spec,
         fixed_point,
         state,
-        Num.(parameters(dynamics)),
+        Num.(p),
         ModelingToolkit.get_defaults(dynamics),
-        false
+        policy_search
     )
 end
 
@@ -262,7 +269,11 @@ function _NeuralLyapunovPDESystem(
     end
 
     # NeuralPDE requires an equation and a boundary condition, even if they are
-    # trivial like 0.0 == 0.0
+    # trivial like 0.0 == 0.0, so we remove those trivial equations when they showed up
+    # naturally alongside other equations and them in we're left with no other DifferentialEquations
+    eqs = filter(eq -> eq != (0.0 ~ 0.0), eqs)
+    bcs = filter(eq -> eq != (0.0 ~ 0.0), bcs)
+
     if isempty(eqs)
         push!(eqs, 0.0 ~ 0.0)
     end
diff --git a/test/inverted_pendulum_ODESystem.jl b/test/inverted_pendulum_ODESystem.jl
new file mode 100644
index 0000000..6605513
--- /dev/null
+++ b/test/inverted_pendulum_ODESystem.jl
@@ -0,0 +1,255 @@
+using LinearAlgebra
+using NeuralPDE, Lux, ModelingToolkit
+using Optimization, OptimizationOptimisers, OptimizationOptimJL, NLopt
+using NeuralLyapunov
+using Random
+using Test
+
+Random.seed!(200)
+
+println("Inverted Pendulum - Policy Search 2")
+
+######################### Define dynamics and domain ##########################
+
+@parameters ζ ω_0
+defaults = Dict([ζ => 0.5, ω_0 => 1.0])
+
+@variables t θ(t) τ(t) [input = true]
+Dt = Differential(t)
+DDt = Dt^2
+
+eqs = [DDt(θ) + 2ζ * Dt(θ) + ω_0^2 * sin(θ) ~ τ]
+
+@named driven_pendulum = ODESystem(
+    eqs,
+    t,
+    [θ, τ],
+    [ζ, ω_0];
+    defaults = defaults
+)
+
+(open_loop_pendulum_dynamics, _), state_order, p_order = ModelingToolkit.generate_control_function(
+    driven_pendulum; simplify = true)
+
+upright_equilibrium = [π, 0.0]
+lb = [0.0, -10.0];
+ub = [2π, 10.0];
+p = [defaults[param] for param in p_order]
+
+####################### Specify neural Lyapunov problem #######################
+
+# Define neural network discretization
+# We use an input layer that is periodic with period 2π with respect to θ
+dim_state = length(lb)
+dim_hidden = 15
+dim_phi = 2
+dim_u = 1
+dim_output = dim_phi + dim_u
+chain = [Lux.Chain(
+             Lux.WrappedFunction(x -> vcat(
+                 transpose(sin.(x[1, :])),
+                 transpose(cos.(x[1, :])),
+                 transpose(x[2, :])
+             )),
+             Dense(3, dim_hidden, tanh),
+             Dense(dim_hidden, dim_hidden, tanh),
+             Dense(dim_hidden, 1, use_bias = false)
+         ) for _ in 1:dim_output]
+
+# Define neural network discretization
+strategy = GridTraining(0.1)
+discretization = PhysicsInformedNN(chain, strategy)
+
+# Define neural Lyapunov structure
+structure = PositiveSemiDefiniteStructure(
+    dim_phi;
+    pos_def = function (state, fixed_point)
+        θ, ω = state
+        θ_eq, ω_eq = fixed_point
+        log(1.0 + (sin(θ) - sin(θ_eq))^2 + (cos(θ) - cos(θ_eq))^2 + (ω - ω_eq)^2)
+    end
+)
+structure = add_policy_search(
+    structure,
+    dim_u
+)
+minimization_condition = DontCheckNonnegativity(check_fixed_point = false)
+
+# Define Lyapunov decrease condition
+decrease_condition = AsymptoticDecrease(strict = true)
+
+# Construct neural Lyapunov specification
+spec = NeuralLyapunovSpecification(
+    structure,
+    minimization_condition,
+    decrease_condition
+)
+
+############################# Construct PDESystem #############################
+
+pde_system, network_func = NeuralLyapunovPDESystem(
+    driven_pendulum,
+    lb,
+    ub,
+    spec;
+    fixed_point = upright_equilibrium
+)
+
+######################## Construct OptimizationProblem ########################
+
+sym_prob = symbolic_discretize(pde_system, discretization)
+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, BFGS(); maxiters = 300)
+
+###################### Get numerical numerical functions ######################
+
+V_func, V̇_func, ∇V_func = NumericalNeuralLyapunovFunctions(
+    discretization.phi,
+    res.u,
+    network_func,
+    structure,
+    open_loop_pendulum_dynamics,
+    upright_equilibrium;
+    p = p
+)
+
+u = get_policy(discretization.phi, res.u, network_func, dim_u)
+
+################################## Simulate ###################################
+
+xs = (-2π):0.02:(2π)
+ys = lb[2]:0.02:ub[2]
+states = Iterators.map(collect, Iterators.product(xs, ys))
+V_predict = vec(V_func(hcat(states...)))
+dVdt_predict = vec(V̇_func(hcat(states...)))
+
+#################################### Tests ####################################
+
+# Network structure should enforce positive definiteness
+@test V_func(upright_equilibrium) == 0.0
+@test min(V_func(upright_equilibrium), minimum(V_predict)) ≥ 0.0
+
+# Network structure should enforce periodicity in θ
+x0 = (ub .- lb) .* rand(2, 100) .+ lb
+@test all(isapprox.(V_func(x0), V_func(x0 .+ [2π, 0.0]); rtol = 1e-3))
+
+# Training should result in a fixed point at the upright equilibrium
+@test all(isapprox.(
+    open_loop_pendulum_dynamics(upright_equilibrium, u(upright_equilibrium), p, 0.0),
+    0.0; atol = 1e-8))
+@test V̇_func(upright_equilibrium) == 0.0
+
+# V̇ should be negative almost everywhere
+@test sum(dVdt_predict .> 0) / length(dVdt_predict) < 1.04e-3
+
+################################## Simulate ###################################
+
+using DifferentialEquations
+
+state_order = map(st -> istree(st) ? operation(st) : st, state_order)
+state_syms = Symbol.(state_order)
+
+closed_loop_dynamics = ODEFunction(
+    (x, p, t) -> open_loop_pendulum_dynamics(x, u(x), p, t);
+    sys = SciMLBase.SymbolCache(state_syms, Symbol.(p_order))
+)
+
+# Starting still at bottom
+downward_equilibrium = zeros(2)
+ode_prob = ODEProblem(closed_loop_dynamics, downward_equilibrium, [0.0, 70.0], p)
+sol = solve(ode_prob, Tsit5())
+# plot(sol)
+
+# Should make it to the top
+θ_end, ω_end = sol.u[end]
+x_end, y_end = sin(θ_end), -cos(θ_end)
+@test all(isapprox.([x_end, y_end, ω_end], [0.0, 1.0, 0.0]; atol = 1e-3))
+
+# Starting at a random point
+x0 = lb .+ rand(2) .* (ub .- lb)
+ode_prob = ODEProblem(closed_loop_dynamics, x0, [0.0, 100.0], p)
+sol = solve(ode_prob, Tsit5())
+# plot(sol)
+
+# Should make it to the top
+θ_end, ω_end = sol.u[end]
+x_end, y_end = sin(θ_end), -cos(θ_end)
+@test all(isapprox.([x_end, y_end, ω_end], [0.0, 1.0, 0.0]; atol = 1e-3))
+
+#=
+# Print statistics
+println("V(π, 0) = ", V_func(upright_equilibrium))
+println(
+    "f([π, 0], u([π, 0])) = ",
+    open_loop_pendulum_dynamics(upright_equilibrium, u(upright_equilibrium), p, 0.0)
+)
+println(
+    "V ∋ [",
+    min(V_func(upright_equilibrium),
+    minimum(V_predict)),
+    ", ",
+    maximum(V_predict),
+    "]"
+)
+println(
+    "V̇ ∋ [",
+    minimum(dVdt_predict),
+    ", ",
+    max(V̇_func(upright_equilibrium), maximum(dVdt_predict)),
+    "]"
+)
+
+# Plot results
+using Plots
+
+p1 = plot(
+    xs / pi,
+    ys,
+    V_predict,
+    linetype =
+    :contourf,
+    title = "V",
+    xlabel = "θ/π",
+    ylabel = "ω",
+    c = :bone_1
+);
+p1 = scatter!([-2 * pi, 0, 2 * pi] / pi, [0, 0, 0],
+    label = "Downward Equilibria", color = :green, markershape = :+);
+p1 = scatter!(
+    [-pi, pi] / pi, [0, 0], label = "Upward Equilibria", color = :red, markershape = :x);
+p2 = plot(
+    xs / pi,
+    ys,
+    dVdt_predict,
+    linetype = :contourf,
+    title = "dV/dt",
+    xlabel = "θ/π",
+    ylabel = "ω",
+    c = :binary
+);
+p2 = scatter!([-2 * pi, 0, 2 * pi] / pi, [0, 0, 0],
+    label = "Downward Equilibria", color = :green, markershape = :+);
+p2 = scatter!([-pi, pi] / pi, [0, 0], label = "Upward Equilibria",
+    color = :red, markershape = :x, legend = false);
+p3 = plot(
+    xs / pi,
+    ys,
+    dVdt_predict .< 0,
+    linetype = :contourf,
+    title = "dV/dt < 0",
+    xlabel = "θ/π",
+    ylabel = "ω",
+    colorbar = false,
+    linewidth = 0
+);
+p3 = scatter!([-2 * pi, 0, 2 * pi] / pi, [0, 0, 0],
+    label = "Downward Equilibria", color = :green, markershape = :+);
+p3 = scatter!([-pi, pi] / pi, [0, 0], label = "Upward Equilibria",
+    color = :red, markershape = :x, legend = false);
+plot(p1, p2, p3)
+=#
diff --git a/test/runtests.jl b/test/runtests.jl
index fae3962..ef1120d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -13,4 +13,7 @@ using SafeTestsets
     @time @safetestset "Policy search - inverted pendulum" begin
         include("inverted_pendulum.jl")
     end
+    @time @safetestset "Policy search - inverted pendulum 2" begin
+        include("inverted_pendulum_ODESystem.jl")
+    end
 end

From eb0ca776bb682b73ba695b45f19ddd71e11db0ce Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Wed, 20 Mar 2024 18:16:32 -0400
Subject: [PATCH 13/16] Replaced `lb` and `ub` with `bounds` when `dynamics isa
 ODESystem`; removed ODEProblem method

---
 src/NeuralLyapunovPDESystem.jl      | 131 ++++++++++++----------------
 test/damped_pendulum.jl             |  18 ++--
 test/inverted_pendulum_ODESystem.jl |  18 ++--
 3 files changed, 79 insertions(+), 88 deletions(-)

diff --git a/src/NeuralLyapunovPDESystem.jl b/src/NeuralLyapunovPDESystem.jl
index dca6d5e..e6ec167 100644
--- a/src/NeuralLyapunovPDESystem.jl
+++ b/src/NeuralLyapunovPDESystem.jl
@@ -1,33 +1,42 @@
 """
-    NeuralLyapunovPDESystem(dynamics, lb, ub, spec; fixed_point, ps)
+    NeuralLyapunovPDESystem(dynamics::ODESystem, bounds, spec; <keyword_arguments>)
+    NeuralLyapunovPDESystem(dynamics::Function, lb, ub, spec; <keyword_arguments>)
 
-Constructs a ModelingToolkit `PDESystem` to train a neural Lyapunov function
-
-Returns the `PDESystem` and a function representing the neural network, which operates
-columnwise.
-
-The neural Lyapunov function will only be trained for `{ x : lb .≤ x .≤ ub }`. The Lyapunov
-function will be for the dynamical system represented by `dynamics`. If `dynamics` is an
-`ODEProblem` or `ODEFunction`, then the corresponding ODE; if `dynamics` is a function, then
-the ODE is typically `ẋ = dynamics(x, p, t)`, but `spec` may specify a different form. In
-particular, if `policy_search`, then the ODE should have a control input:
-`ẋ = dynamics(x, input, p, t)`. In that case, the loss function will include a term
-penalizing nonzero `ẋ` at the fixed point. Otherwise, the ODE should have a fixed point
-at `x = fixed_point`. Additionally, the ODE should not depend on `t` (time `t=0.0` alone
-will be used). The particular Lyapunov conditions to be used and structure of the neural
-Lyapunov function are specified through `spec`, which is a `NeuralLyapunovSpecification`.
+Construct and return a `PDESystem` representing the specified neural Lyapunov problem, along
+with a function representing the neural network.
 
 The returned neural network function takes three inputs: the neural network structure `phi`,
-the trained network parameters, and a matrix of inputs to operate on columnwise.
-
-If `dynamics` requires parameters, their values can be supplied through the Vector `p`, or
-through the parameters of `dynamics` if `dynamics isa ODEProblem` (in which case, let
-the other be `SciMLBase.NullParameters()`). If `dynamics` is an `ODEFunction` and
-`dynamics.sys.parameters` is defined, then `p` should have the same order.
-
-To use specific symbols for the state and parameters in the resulting `PDESystem`, supply
-them through the `state_syms` and `parameter_syms` keywords (if not using an `ODEFunction`,
-`ODEProblem`, or `ODESystem`, which each have their own ways of defining symbols.)
+the trained network parameters, and a matrix of inputs, then operates columnwise on the
+inputs.
+
+# Arguments
+- `dynamics`: the dynamical system being analyzed, represented as an `ODESystem` or the
+        function `f` such that `ẋ = f(x[, u], p, t)`; either way, the ODE should not depend
+        on time and only `t = 0.0` will be used
+- `bounds`: an array of domains, defining the training domain by bounding the states (and
+        derivatives, when applicable) of `dynamics`; only used when `dynamics isa
+        ODESystem`, otherwise use `lb` and `ub`.
+- `lb` and `ub`: the training domain will be ``[lb_1, ub_1]×[lb_2, ub_2]×...``; not used
+        when `dynamics isa ODESystem`, then use `bounds`.
+- `spec::NeuralLyapunovSpecification`: defines the Lyapunov function structure, as well as
+        the minimization and decrease conditions.
+- `fixed_point`: the equilibrium being analyzed; defaults to the origin.
+- `p`: the values of the parameters of the dynamical system being analyzed; defaults to
+        `SciMLBase.NullParameters()`; not used when `dynamics isa ODESystem`, then use the
+        default parameter values of `dynamics`.
+- `state_syms`: an array of the `Symbol` representing each state; not used when `dynamics
+        isa ODESystem`, then the symbols from `dynamics` are used; if `dynamics isa
+        ODEFunction`, symbols stored there are used, unless overridden here; if not provided
+        here and cannot be inferred, `[:state1, :state2, ...]` will be used.
+- `parameter_syms`: an array of the `Symbol` representing each parameter; not used when
+        `dynamics isa ODESystem`, then the symbols from `dynamics` are used; if `dynamics
+        isa ODEFunction`, symbols stored there are used, unless overridden here; if not
+        provided here and cannot be inferred, `[:param1, :param2, ...]` will be used.
+- `policy_search::Bool`: whether or not to include a loss term enforcing `fixed_point` to
+        actually be a fixed point; defaults to `false`; only used when `dynamics isa
+        Function && !(dynamics isa ODEFunction)`; when `dynamics isa ODEFunction`,
+        `policy_search` must be `false`, so should not be supplied; when `dynamics isa
+        ODESystem`, value inferred by the presence of unbound inputs.
 """
 function NeuralLyapunovPDESystem(
         dynamics::Function,
@@ -72,11 +81,13 @@ function NeuralLyapunovPDESystem(
         Dict([param => param_val for (param, param_val) in zip(params, p)])
     end
 
+    ############################# Define domains ##############################
+    domains = [state[i] ∈ (lb[i], ub[i]) for i in 1:state_dim]
+
     ########################### Construct PDESystem ###########################
     return _NeuralLyapunovPDESystem(
         dynamics,
-        lb,
-        ub,
+        domains,
         spec,
         fixed_point,
         state,
@@ -99,6 +110,7 @@ function NeuralLyapunovPDESystem(
     end
 
     if dynamics.sys isa ODESystem
+        ##### TODO: This will cause a MethodError; we need to set up bounds to make this work
         return NeuralLyapunovPDESystem(
             dynamics.sys,
             lb,
@@ -128,44 +140,11 @@ function NeuralLyapunovPDESystem(
     )
 end
 
-function NeuralLyapunovPDESystem(
-        dynamics::ODEProblem,
-        lb,
-        ub,
-        spec::NeuralLyapunovSpecification;
-        fixed_point = zeros(length(lb)),
-        p = SciMLBase.NullParameters()
-)::Tuple{PDESystem, Function}
-    f = dynamics.f
-
-    p = if dynamics.p == SciMLBase.NullParameters()
-        p
-    elseif p == SciMLBase.NullParameters()
-        dynamics.p
-    elseif dynamics.p == p
-        p
-    else
-        throw(ErrorException("Conflicting parameter definitions. Please define parameters" *
-                             " only through p or dynamics.p; the other should be " *
-                             "SciMLBase.NullParameters()"))
-    end
-
-    return NeuralLyapunovPDESystem(
-        f,
-        lb,
-        ub,
-        spec;
-        fixed_point = fixed_point,
-        p = p
-    )
-end
-
 function NeuralLyapunovPDESystem(
         dynamics::ODESystem,
-        lb,
-        ub,
+        bounds,
         spec::NeuralLyapunovSpecification;
-        fixed_point = zeros(length(lb))
+        fixed_point = zeros(length(bounds))
 )::Tuple{PDESystem, Function}
     ######################### Check for policy search #########################
     f, x, p, policy_search = if isempty(ModelingToolkit.unbound_inputs(dynamics))
@@ -178,15 +157,25 @@ function NeuralLyapunovPDESystem(
     ########################## Define state symbols ###########################
     # States should all be functions of time, but we just want the symbol
     # e.g., if the state is ω(t), we just want ω
-    state = map(st -> istree(st) ? operation(st) : st, x)
-    state_syms = Symbol.(state)
+    _state = map(st -> istree(st) ? operation(st) : st, x)
+    state_syms = Symbol.(_state)
     state = [first(@parameters $s) for s in state_syms]
 
+    ###################### Remove derivatives in domains ######################
+    domains = map(
+        d -> Num(operation(Symbolics.diff2term(Symbolics.value(d.variables)))) ∈ d.domain,
+        bounds
+    )
+    domain_vars = map(d -> d.variables, domains)
+    if Set(_state) != Set(domain_vars)
+        error("Domain variables from `domains` do not match those extracted from " *
+              "`dynamics`. Got $_state from `dynamics` and $domain_vars from `domains`.")
+    end
+
     ########################### Construct PDESystem ###########################
     _NeuralLyapunovPDESystem(
         f,
-        lb,
-        ub,
+        domains,
         spec,
         fixed_point,
         state,
@@ -198,8 +187,7 @@ end
 
 function _NeuralLyapunovPDESystem(
         dynamics::Function,
-        lb,
-        ub,
+        domains,
         spec::NeuralLyapunovSpecification,
         fixed_point,
         state,
@@ -212,10 +200,7 @@ function _NeuralLyapunovPDESystem(
     minimzation_condition = spec.minimzation_condition
     decrease_condition = spec.decrease_condition
     f_call = structure.f_call
-
-    ############################# Define domains ##############################
-    state_dim = length(lb)
-    domains = [state[i] ∈ (lb[i], ub[i]) for i in 1:state_dim]
+    state_dim = length(domains)
 
     ################## Define Lyapunov function & derivative ##################
     output_dim = structure.network_dim
diff --git a/test/damped_pendulum.jl b/test/damped_pendulum.jl
index 6632021..020167e 100644
--- a/test/damped_pendulum.jl
+++ b/test/damped_pendulum.jl
@@ -29,16 +29,16 @@ eqs = [DDt(θ) + 2ζ * Dt(θ) + ω_0^2 * sin(θ) ~ 0.0]
 )
 
 dynamics = structural_simplify(dynamics)
-
-lb = [-pi, -10.0];
-ub = [pi, 10.0];
-p = [defaults[param] for param in parameters(dynamics)]
+bounds = [
+    θ ∈ (-π, π),
+    Dt(θ) ∈ (-10.0, 10.0)
+]
 
 ####################### Specify neural Lyapunov problem #######################
 
 # Define neural network discretization
 # We use an input layer that is periodic with period 2π with respect to θ
-dim_state = length(lb)
+dim_state = length(bounds)
 dim_hidden = 15
 dim_output = 2
 chain = [Lux.Chain(
@@ -81,8 +81,7 @@ spec = NeuralLyapunovSpecification(
 
 pde_system, network_func = NeuralLyapunovPDESystem(
     dynamics,
-    lb,
-    ub,
+    bounds,
     spec
 )
 
@@ -99,18 +98,21 @@ res = Optimization.solve(prob, BFGS(); maxiters = 300)
 
 ###################### Get numerical numerical functions ######################
 
+p = [defaults[param] for param in parameters(dynamics)]
 V_func, V̇_func, ∇V_func = NumericalNeuralLyapunovFunctions(
     discretization.phi,
     res.u,
     network_func,
     structure.V,
     ODEFunction(dynamics),
-    zeros(length(lb));
+    zeros(length(bounds));
     p = p
 )
 
 ################################## Simulate ###################################
 
+lb = [-pi, -10.0];
+ub = [pi, 10.0];
 xs = (2 * lb[1]):0.02:(2 * ub[1])
 ys = lb[2]:0.02:ub[2]
 states = Iterators.map(collect, Iterators.product(xs, ys))
diff --git a/test/inverted_pendulum_ODESystem.jl b/test/inverted_pendulum_ODESystem.jl
index 6605513..5738c82 100644
--- a/test/inverted_pendulum_ODESystem.jl
+++ b/test/inverted_pendulum_ODESystem.jl
@@ -28,19 +28,22 @@ eqs = [DDt(θ) + 2ζ * Dt(θ) + ω_0^2 * sin(θ) ~ τ]
     defaults = defaults
 )
 
+bounds = [
+    θ ∈ (0, 2π),
+    Dt(θ) ∈ (-10.0, 10.0)
+]
+
 (open_loop_pendulum_dynamics, _), state_order, p_order = ModelingToolkit.generate_control_function(
     driven_pendulum; simplify = true)
 
 upright_equilibrium = [π, 0.0]
-lb = [0.0, -10.0];
-ub = [2π, 10.0];
 p = [defaults[param] for param in p_order]
 
 ####################### Specify neural Lyapunov problem #######################
 
 # Define neural network discretization
 # We use an input layer that is periodic with period 2π with respect to θ
-dim_state = length(lb)
+dim_state = length(bounds)
 dim_hidden = 15
 dim_phi = 2
 dim_u = 1
@@ -89,8 +92,7 @@ spec = NeuralLyapunovSpecification(
 
 pde_system, network_func = NeuralLyapunovPDESystem(
     driven_pendulum,
-    lb,
-    ub,
+    bounds,
     spec;
     fixed_point = upright_equilibrium
 )
@@ -122,8 +124,10 @@ u = get_policy(discretization.phi, res.u, network_func, dim_u)
 
 ################################## Simulate ###################################
 
-xs = (-2π):0.02:(2π)
-ys = lb[2]:0.02:ub[2]
+lb = [0.0, -10.0];
+ub = [2π, 10.0];
+xs = (-2π):0.1:(2π)
+ys = lb[2]:0.1:ub[2]
 states = Iterators.map(collect, Iterators.product(xs, ys))
 V_predict = vec(V_func(hcat(states...)))
 dVdt_predict = vec(V̇_func(hcat(states...)))

From 651eb52a3a09bb02b55013db30a6933e82f51a13 Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Thu, 21 Mar 2024 12:42:11 -0400
Subject: [PATCH 14/16] Fixed `MethodError` when an `ODEFunction` was made from
 an `ODESystem`

---
 src/NeuralLyapunovPDESystem.jl | 33 ++++++++++++++++-----------------
 test/damped_pendulum.jl        | 14 ++++++++------
 2 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/src/NeuralLyapunovPDESystem.jl b/src/NeuralLyapunovPDESystem.jl
index e6ec167..e7dc851 100644
--- a/src/NeuralLyapunovPDESystem.jl
+++ b/src/NeuralLyapunovPDESystem.jl
@@ -106,25 +106,24 @@ function NeuralLyapunovPDESystem(
         p = SciMLBase.NullParameters()
 )::Tuple{PDESystem, Function}
     if dynamics.mass_matrix !== I
-        throw(ErrorException("DAEs are not supported at this time."))
+        throw(ErrorException("DAEs are not supported at this time. Please supply dynamics" *
+                             " without a mass matrix."))
     end
 
-    if dynamics.sys isa ODESystem
-        ##### TODO: This will cause a MethodError; we need to set up bounds to make this work
-        return NeuralLyapunovPDESystem(
-            dynamics.sys,
-            lb,
-            ub,
-            spec;
-            fixed_point = fixed_point,
-            p = p
-        )
-    end
-
-    p_syms = if isnothing(dynamics.sys.parameters)
-        []
+    s_syms, p_syms = if dynamics.sys isa ODESystem
+        s_syms = Symbol.(operation.(states(dynamics.sys)))
+        p_syms = Symbol.(parameters(dynamics.sys))
+        (s_syms, p_syms)
+    elseif dynamics.sys isa SciMLBase.SymbolCache
+        s_syms = SciMLBase.variable_symbols(dynamics.sys)
+        p_syms = if isnothing(dynamics.sys.parameters)
+            []
+        else
+            dynamics.sys.parameters
+        end
+        (s_syms, p_syms)
     else
-        dynamics.sys.parameters
+        ([], [])
     end
 
     return NeuralLyapunovPDESystem(
@@ -134,7 +133,7 @@ function NeuralLyapunovPDESystem(
         spec;
         fixed_point = fixed_point,
         p = p,
-        state_syms = SciMLBase.variable_symbols(dynamics.sys),
+        state_syms = s_syms,
         parameter_syms = p_syms,
         policy_search = false
     )
diff --git a/test/damped_pendulum.jl b/test/damped_pendulum.jl
index 020167e..fcda6e9 100644
--- a/test/damped_pendulum.jl
+++ b/test/damped_pendulum.jl
@@ -33,6 +33,9 @@ bounds = [
     θ ∈ (-π, π),
     Dt(θ) ∈ (-10.0, 10.0)
 ]
+lb = [-π, -10.0];
+ub = [π, 10.0];
+p = [defaults[param] for param in parameters(dynamics)]
 
 ####################### Specify neural Lyapunov problem #######################
 
@@ -80,9 +83,11 @@ spec = NeuralLyapunovSpecification(
 ############################# Construct PDESystem #############################
 
 pde_system, network_func = NeuralLyapunovPDESystem(
-    dynamics,
-    bounds,
-    spec
+    ODEFunction(dynamics),
+    lb,
+    ub,
+    spec;
+    p = p
 )
 
 ######################## Construct OptimizationProblem ########################
@@ -98,7 +103,6 @@ res = Optimization.solve(prob, BFGS(); maxiters = 300)
 
 ###################### Get numerical numerical functions ######################
 
-p = [defaults[param] for param in parameters(dynamics)]
 V_func, V̇_func, ∇V_func = NumericalNeuralLyapunovFunctions(
     discretization.phi,
     res.u,
@@ -111,8 +115,6 @@ V_func, V̇_func, ∇V_func = NumericalNeuralLyapunovFunctions(
 
 ################################## Simulate ###################################
 
-lb = [-pi, -10.0];
-ub = [pi, 10.0];
 xs = (2 * lb[1]):0.02:(2 * ub[1])
 ys = lb[2]:0.02:ub[2]
 states = Iterators.map(collect, Iterators.product(xs, ys))

From 6e9732bd18d101bc42e147d910985b2edb82779d Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Thu, 21 Mar 2024 12:52:38 -0400
Subject: [PATCH 15/16] Cleaned up handling of ODESystem symbols

---
 src/NeuralLyapunovPDESystem.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/NeuralLyapunovPDESystem.jl b/src/NeuralLyapunovPDESystem.jl
index e7dc851..b65a9ae 100644
--- a/src/NeuralLyapunovPDESystem.jl
+++ b/src/NeuralLyapunovPDESystem.jl
@@ -156,7 +156,7 @@ function NeuralLyapunovPDESystem(
     ########################## Define state symbols ###########################
     # States should all be functions of time, but we just want the symbol
     # e.g., if the state is ω(t), we just want ω
-    _state = map(st -> istree(st) ? operation(st) : st, x)
+    _state = operation.(x)
     state_syms = Symbol.(_state)
     state = [first(@parameters $s) for s in state_syms]
 
@@ -178,7 +178,7 @@ function NeuralLyapunovPDESystem(
         spec,
         fixed_point,
         state,
-        Num.(p),
+        p,
         ModelingToolkit.get_defaults(dynamics),
         policy_search
     )

From b5f9226f8fd47429faa337248f68c15005dc797c Mon Sep 17 00:00:00 2001
From: Nicholas Klugman <13633349+nicholaskl97@users.noreply.github.com>
Date: Fri, 22 Mar 2024 14:52:48 -0400
Subject: [PATCH 16/16] Fixed typo in comment

---
 src/NeuralLyapunovPDESystem.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/NeuralLyapunovPDESystem.jl b/src/NeuralLyapunovPDESystem.jl
index b65a9ae..81a059a 100644
--- a/src/NeuralLyapunovPDESystem.jl
+++ b/src/NeuralLyapunovPDESystem.jl
@@ -253,8 +253,8 @@ function _NeuralLyapunovPDESystem(
     end
 
     # NeuralPDE requires an equation and a boundary condition, even if they are
-    # trivial like 0.0 == 0.0, so we remove those trivial equations when they showed up
-    # naturally alongside other equations and them in we're left with no other DifferentialEquations
+    # trivial like 0.0 == 0.0, so we remove those trivial equations if they showed up
+    # naturally alongside other equations and add them in if we have no other equations
     eqs = filter(eq -> eq != (0.0 ~ 0.0), eqs)
     bcs = filter(eq -> eq != (0.0 ~ 0.0), bcs)