Merge pull request #24 from SciML/interface

Various interface changes and minor performance improvements

Project.toml

@@ -34,6 +34,7 @@ OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+CSDP = "0a46da34-8e4b-519e-b418-48813639ff34"
 
 [targets]
-test = ["SafeTestsets", "Test", "Lux", "Optimization", "OptimizationOptimJL", "OptimizationOptimisers", "NLopt", "Random", "NeuralPDE", "DifferentialEquations"]
+test = ["SafeTestsets", "Test", "Lux", "Optimization", "OptimizationOptimJL", "OptimizationOptimisers", "NLopt", "Random", "NeuralPDE", "DifferentialEquations", "CSDP"]

src/NeuralLyapunov.jl

@@ -12,16 +12,31 @@ include("minimization_conditions.jl")
 include("decrease_conditions.jl")
 include("decrease_conditions_RoA_aware.jl")
 include("NeuralLyapunovPDESystem.jl")
-include("local_Lyapunov.jl")
+include("numerical_lyapunov_functions.jl")
+include("local_lyapunov.jl")
 include("policy_search.jl")
 
-export NeuralLyapunovPDESystem, NumericalNeuralLyapunovFunctions
-export local_Lyapunov
-export NeuralLyapunovSpecification, NeuralLyapunovStructure, UnstructuredNeuralLyapunov,
- NonnegativeNeuralLyapunov, PositiveSemiDefiniteStructure,
- LyapunovMinimizationCondition, StrictlyPositiveDefinite, PositiveSemiDefinite,
- DontCheckNonnegativity, LyapunovDecreaseCondition, AsymptoticDecrease,
- ExponentialDecrease, DontCheckDecrease, RoAAwareDecreaseCondition, make_RoA_aware,
- add_policy_search, get_policy
+# Lyapunov function structures
+export NeuralLyapunovStructure, UnstructuredNeuralLyapunov, NonnegativeNeuralLyapunov,
+ PositiveSemiDefiniteStructure, get_numerical_lyapunov_function
+
+# Minimization conditions
+export LyapunovMinimizationCondition, StrictlyPositiveDefinite, PositiveSemiDefinite,
+ DontCheckNonnegativity
+
+# Decrease conditions
+export LyapunovDecreaseCondition, AsymptoticDecrease, ExponentialDecrease, DontCheckDecrease
+
+# Setting up the PDESystem for NeuralPDE
+export NeuralLyapunovSpecification, NeuralLyapunovPDESystem
+
+# Region of attraction handling
+export RoAAwareDecreaseCondition, make_RoA_aware
+
+# Policy search
+export add_policy_search, get_policy
+
+# Local Lyapunov analysis
+export local_lyapunov
 
 end

src/NeuralLyapunovPDESystem.jl

@@ -2,12 +2,7 @@
  NeuralLyapunovPDESystem(dynamics::ODESystem, bounds, spec; <keyword_arguments>)
  NeuralLyapunovPDESystem(dynamics::Function, lb, ub, spec; <keyword_arguments>)
 
-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, then operates columnwise on the
-inputs.
+Construct a `ModelingToolkit.PDESystem` representing the specified neural Lyapunov problem.
 
 # Arguments
 - `dynamics`: the dynamical system being analyzed, represented as an `ODESystem` or the
@@ -37,6 +32,7 @@ inputs.
  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.
+- `name`: the name of the constructed `PDESystem`
 """
 function NeuralLyapunovPDESystem(
  dynamics::Function,
@@ -47,8 +43,9 @@ function NeuralLyapunovPDESystem(
  p = SciMLBase.NullParameters(),
  state_syms = [],
  parameter_syms = [],
- policy_search::Bool = false
-)::Tuple{PDESystem, Function}
+ policy_search::Bool = false,
+ name
+)::PDESystem
  ########################## Define state symbols ###########################
  state_dim = length(lb)
 
@@ -93,7 +90,8 @@ function NeuralLyapunovPDESystem(
  state,
  params,
  defaults,
- policy_search
+ policy_search,
+ name
  )
 end
 
@@ -103,15 +101,16 @@ function NeuralLyapunovPDESystem(
  ub,
  spec::NeuralLyapunovSpecification;
  fixed_point = zeros(length(lb)),
- p = SciMLBase.NullParameters()
-)::Tuple{PDESystem, Function}
+ p = SciMLBase.NullParameters(),
+ name
+)::PDESystem
  if dynamics.mass_matrix !== I
  throw(ErrorException("DAEs are not supported at this time. Please supply dynamics" *
  " without a mass matrix."))
  end
 
  s_syms, p_syms = if dynamics.sys isa ODESystem
- s_syms = Symbol.(operation.(states(dynamics.sys)))
+ s_syms = Symbol.(operation.(unknowns(dynamics.sys)))
  p_syms = Symbol.(parameters(dynamics.sys))
  (s_syms, p_syms)
  elseif dynamics.sys isa SciMLBase.SymbolCache
@@ -135,19 +134,21 @@ function NeuralLyapunovPDESystem(
  p = p,
  state_syms = s_syms,
  parameter_syms = p_syms,
- policy_search = false
+ policy_search = false,
+ name
  )
 end
 
 function NeuralLyapunovPDESystem(
  dynamics::ODESystem,
  bounds,
  spec::NeuralLyapunovSpecification;
- fixed_point = zeros(length(bounds))
-)::Tuple{PDESystem, Function}
+ fixed_point = zeros(length(bounds)),
+ name
+)::PDESystem
  ######################### Check for policy search #########################
  f, x, p, policy_search = if isempty(ModelingToolkit.unbound_inputs(dynamics))
- (ODEFunction(dynamics), states(dynamics), parameters(dynamics), false)
+ (ODEFunction(dynamics), unknowns(dynamics), parameters(dynamics), false)
  else
  (f, _), x, p = ModelingToolkit.generate_control_function(dynamics; simplify = true)
  (f, x, p, true)
@@ -172,15 +173,16 @@ function NeuralLyapunovPDESystem(
  end
 
  ########################### Construct PDESystem ###########################
- _NeuralLyapunovPDESystem(
+ return _NeuralLyapunovPDESystem(
  f,
  domains,
  spec,
  fixed_point,
  state,
  p,
  ModelingToolkit.get_defaults(dynamics),
- policy_search
+ policy_search,
+ name
  )
 end
 
@@ -192,8 +194,9 @@ function _NeuralLyapunovPDESystem(
  state,
  params,
  defaults,
- policy_search::Bool
-)::Tuple{PDESystem, Function}
+ policy_search::Bool,
+ name
+)::PDESystem
  ########################## Unpack specifications ##########################
  structure = spec.structure
  minimzation_condition = spec.minimzation_condition
@@ -266,129 +269,14 @@ function _NeuralLyapunovPDESystem(
  end
 
  ########################### Construct PDESystem ###########################
- @named lyapunov_pde_system = PDESystem(
+ return PDESystem(
  eqs,
  bcs,
  domains,
  state,
  φ(state),
  params;
- defaults = defaults
+ defaults = defaults,
+ name = name
  )
-
- ################### Return PDESystem and neural network ###################
- # φ_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, φ_func
-end
-
-"""
- NumericalNeuralLyapunovFunctions(phi, θ, network_func, structure, dynamics, fixed_point;
- jac, J_net)
-
-Returns the Lyapunov function, its time derivative, and its gradient: `V(state)`,
-`V̇(state)`, and `∇V(state)`
-
-These functions can operate on a state vector or columnwise on a matrix of state vectors.
-`phi` is the neural network with parameters `θ`. `network_func(phi, θ, state)` is an output
-of `NeuralLyapunovPDESystem`, which evaluates the neural network represented by `phi` with
-parameters `θ` at `state`.
-
-The Lyapunov function structure is specified in structure, which is a
-`NeuralLyapunovStructure`. The Jacobian of the network is either specified via
-`J_net(_phi, _θ, state)` or calculated using `jac`, which defaults to
-`ForwardDiff.jacobian`.
-"""
-function NumericalNeuralLyapunovFunctions(
- phi,
- θ,
- network_func::Function,
- structure::NeuralLyapunovStructure,
- dynamics::Function,
- fixed_point;
- p = SciMLBase.NullParameters(),
- jac = ForwardDiff.jacobian,
- J_net = (_phi, _θ, x) -> jac((y) -> network_func(_phi, _θ, y), x)
-)::Tuple{Function, Function, Function}
- # Make Network function
- _net_func = (x) -> network_func(phi, θ, x)
- _J_net = (x) -> J_net(phi, θ, x)
-
- # Numerical form of Lyapunov function
- V_func(state::AbstractVector) = structure.V(_net_func, state, fixed_point)
- V_func(state::AbstractMatrix) = mapslices(V_func, state, dims = [1])
-
- # Numerical gradient of Lyapunov function
- ∇V_func(state::AbstractVector) = structure.∇V(
- _net_func,
- _J_net,
- state,
- fixed_point
- )
- ∇V_func(state::AbstractMatrix) = mapslices(∇V_func, state, dims = [1])
-
- # Numerical time derivative of Lyapunov function
- function V̇_func(state::AbstractVector)
- structure.V̇(
- _net_func,
- _J_net,
- dynamics,
- state,
- p,
- 0.0,
- fixed_point
- )
- end
- V̇_func(state::AbstractMatrix) = mapslices(V̇_func, state, dims = [1])
-
- return V_func, V̇_func, ∇V_func
-end
-
-"""
- NumericalNeuralLyapunovFunctions(phi, θ, network_func, V_structure, dynamics,
- fixed_point, grad)
-
-Returns the Lyapunov function, its time derivative, and its gradient: `V(state)`,
-`V̇(state)`, and `∇V(state)`.
-
-These functions 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 Lyapunov function structure is defined by
- `V_structure(_network_func, state, fixed_point)`
-Its gradient is calculated using `grad`, which defaults to `ForwardDiff.gradient`.
-"""
-function NumericalNeuralLyapunovFunctions(
- phi,
- θ,
- network_func::Function,
- V_structure::Function,
- dynamics::Function,
- fixed_point;
- p = SciMLBase.NullParameters(),
- grad = ForwardDiff.gradient
-)::Tuple{Function, Function, Function}
- # Make network function
- _net_func = (x) -> network_func(phi, θ, x)
-
- # Numerical form of Lyapunov function
- V_func(state::AbstractVector) = V_structure(_net_func, state, fixed_point)
- V_func(state::AbstractMatrix) = mapslices(V_func, state, dims = [1])
-
- # Numerical gradient of Lyapunov function
- ∇V_func(state::AbstractVector) = grad(V_func, state)
- ∇V_func(state::AbstractMatrix) = mapslices(∇V_func, state, dims = [1])
-
- # Numerical time derivative of Lyapunov function
- V̇_func(state::AbstractVector) = dynamics(state, p, 0.0) ⋅ ∇V_func(state)
- V̇_func(state::AbstractMatrix) = mapslices(V̇_func, state, dims = [1])
-
- return V_func, V̇_func, ∇V_func
 end

src/conditions_specification.jl

@@ -57,10 +57,10 @@ struct NeuralLyapunovSpecification
 end
 
 """
-check_nonnegativity(cond::AbstractLyapunovMinimizationCondition)
+ check_nonnegativity(cond::AbstractLyapunovMinimizationCondition)
 
-`true` if `cond` specifies training to meet the Lyapunov minimization condition, `false` if
-`cond` specifies no training to meet this condition.
+Return `true` if `cond` specifies training to meet the Lyapunov minimization condition, and
+`false` if `cond` specifies no training to meet this condition.
 """
 function check_nonnegativity(cond::AbstractLyapunovMinimizationCondition)::Bool
  error("check_nonnegativity not implemented for AbstractLyapunovMinimizationCondition " *
@@ -70,8 +70,8 @@ end
 """
  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.
+Return `true` if `cond` specifies training for the Lyapunov function to equal zero at the
+fixed point, and `false` if `cond` specifies no training to meet this condition.
 """
 function check_minimal_fixed_point(cond::AbstractLyapunovMinimizationCondition)::Bool
  error("check_minimal_fixed_point not implemented for " *
@@ -81,8 +81,8 @@ end
 """
  get_minimization_condition(cond::AbstractLyapunovMinimizationCondition)
 
-Returns a function of `V`, `state`, and `fixed_point` that equals zero when the
-Lyapunov minimization condition is met and greater than zero when it's violated.
+Return a function of `V`, `state`, and `fixed_point` that equals zero when the Lyapunov
+minimization condition is met and is greater than zero when it's violated.
 """
 function get_minimization_condition(cond::AbstractLyapunovMinimizationCondition)
  error("get_condition not implemented for AbstractLyapunovMinimizationCondition of " *
@@ -92,8 +92,8 @@ end
 """
  check_decrease(cond::AbstractLyapunovDecreaseCondition)
 
-`true` if `cond` specifies training to meet the Lyapunov decrease condition, `false`
-if `cond` specifies no training to meet this condition.
+Return `true` if `cond` specifies training to meet the Lyapunov decrease condition, and
+`false` if `cond` specifies no training to meet this condition.
 """
 function check_decrease(cond::AbstractLyapunovDecreaseCondition)::Bool
  error("check_decrease not implemented for AbstractLyapunovDecreaseCondition of type " *
@@ -103,8 +103,8 @@ end
 """
  get_decrease_condition(cond::AbstractLyapunovDecreaseCondition)
 
-Returns a function of `V`, `dVdt`, `state`, and `fixed_point` that is equal to zero
-when the Lyapunov decrease condition is met and greater than zero when it is violated.
+Return a function of `V`, `dVdt`, `state`, and `fixed_point` that is equal to zero
+when the Lyapunov decrease condition is met and is greater than zero when it is violated.
 """
 function get_decrease_condition(cond::AbstractLyapunovDecreaseCondition)
  error("get_condition not implemented for AbstractLyapunovDecreaseCondition of type " *