Merge pull request #754 from SciML/nonlinearsystem

Complete NonlinearSystem

Author: ChrisRackauckas

Project.toml

@@ -36,7 +36,7 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 [compat]
 ArrayInterface = "2.8"
 DataStructures = "0.17, 0.18"
-DiffEqBase = "6.48.1"
+DiffEqBase = "6.54.0"
 DiffEqJump = "6.7.5"
 DiffRules = "0.1, 1.0"
 Distributions = "0.24"
@@ -64,6 +64,7 @@ julia = "1.2"
 Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GalacticOptim = "a75be94c-b780-496d-a8a9-0878b188d577"
+NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -72,4 +73,4 @@ StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Dagger", "ForwardDiff", "GalacticOptim", "OrdinaryDiffEq", "Optim", "Random", "SteadyStateDiffEq", "Test", "StochasticDiffEq"]
+test = ["Dagger", "ForwardDiff", "GalacticOptim", "NonlinearSolve", "OrdinaryDiffEq", "Optim", "Random", "SteadyStateDiffEq", "Test", "StochasticDiffEq"]

src/systems/nonlinear/nonlinearsystem.jl

@@ -87,6 +87,132 @@ jacobian_sparsity(sys::NonlinearSystem) =
     jacobian_sparsity([eq.rhs for eq ∈ equations(sys)],
                       states(sys))
 
+function DiffEqBase.NonlinearFunction(sys::NonlinearSystem, args...; kwargs...)
+    NonlinearFunction{true}(sys, args...; kwargs...)
+end
+
+"""
+```julia
+function DiffEqBase.NonlinearFunction{iip}(sys::NonlinearSystem, dvs = states(sys),
+                                     ps = parameters(sys);
+                                     version = nothing,
+                                     jac = false,
+                                     sparse = false,
+                                     kwargs...) where {iip}
+```
+
+Create an `NonlinearFunction` from the [`NonlinearSystem`](@ref). The arguments
+`dvs` and `ps` are used to set the order of the dependent variable and parameter
+vectors, respectively.
+"""
+function DiffEqBase.NonlinearFunction{iip}(sys::NonlinearSystem, dvs = states(sys),
+                                     ps = parameters(sys), u0 = nothing;
+                                     version = nothing,
+                                     jac = false,
+                                     eval_expression = true,
+                                     sparse = false, simplify=false,
+                                     kwargs...) where {iip}
+
+    f_gen = generate_function(sys, dvs, ps; expression=Val{eval_expression}, kwargs...)
+    f_oop,f_iip = eval_expression ? (@RuntimeGeneratedFunction(ex) for ex in f_gen) : f_gen
+    f(u,p) = f_oop(u,p)
+    f(du,u,p) = f_iip(du,u,p)
+
+    if jac
+        jac_gen = generate_jacobian(sys, dvs, ps;
+                                    simplify=simplify, sparse = sparse,
+                                    expression=Val{eval_expression}, kwargs...)
+        jac_oop,jac_iip = eval_expression ? (@RuntimeGeneratedFunction(ex) for ex in jac_gen) : jac_gen
+        _jac(u,p) = jac_oop(u,p)
+        _jac(J,u,p) = jac_iip(J,u,p)
+    else
+        _jac = nothing
+    end
+
+    NonlinearFunction{iip}(f,
+                     jac = _jac === nothing ? nothing : _jac,
+                     jac_prototype = sparse ? similar(sys.jac[],Float64) : nothing,
+                     syms = Symbol.(states(sys)))
+end
+
+"""
+```julia
+function DiffEqBase.NonlinearFunctionExpr{iip}(sys::NonlinearSystem, dvs = states(sys),
+                                     ps = parameters(sys);
+                                     version = nothing,
+                                     jac = false,
+                                     sparse = false,
+                                     kwargs...) where {iip}
+```
+
+Create a Julia expression for an `ODEFunction` from the [`ODESystem`](@ref).
+The arguments `dvs` and `ps` are used to set the order of the dependent
+variable and parameter vectors, respectively.
+"""
+struct NonlinearFunctionExpr{iip} end
+
+function NonlinearFunctionExpr{iip}(sys::NonlinearSystem, dvs = states(sys),
+                                     ps = parameters(sys), u0 = nothing;
+                                     version = nothing, tgrad=false,
+                                     jac = false,
+                                     linenumbers = false,
+                                     sparse = false, simplify=false,
+                                     kwargs...) where {iip}
+
+    idx = iip ? 2 : 1
+    f = generate_function(sys, dvs, ps; expression=Val{true}, kwargs...)[idx]
+
+    if jac
+        _jac = generate_jacobian(sys, dvs, ps;
+                                 sparse=sparse, simplify=simplify,
+                                 expression=Val{true}, kwargs...)[idx]
+    else
+        _jac = :nothing
+    end
+
+    jp_expr = sparse ? :(similar($(sys.jac[]),Float64)) : :nothing
+
+    ex = quote
+        f = $f
+        jac = $_jac
+        NonlinearFunction{$iip}(f,
+                         jac = jac,
+                         jac_prototype = $jp_expr,
+                         syms = $(Symbol.(states(sys))))
+    end
+    !linenumbers ? striplines(ex) : ex
+end
+
+function process_NonlinearProblem(constructor, sys::NonlinearSystem,u0map,parammap;
+                           version = nothing,
+                           jac = false,
+                           checkbounds = false, sparse = false,
+                           simplify=false,
+                           linenumbers = true, parallel=SerialForm(),
+                           eval_expression = true,
+                           kwargs...)
+    dvs = states(sys)
+    ps = parameters(sys)
+    u0map′ = lower_mapnames(u0map)
+    u0 = varmap_to_vars(u0map′,dvs; defaults=get_default_u0(sys))
+
+    if !(parammap isa DiffEqBase.NullParameters)
+        parammap′ = lower_mapnames(parammap)
+        p = varmap_to_vars(parammap′,ps; defaults=get_default_p(sys))
+    else
+        p = ps
+    end
+
+    f = constructor(sys,dvs,ps,u0;jac=jac,checkbounds=checkbounds,
+                    linenumbers=linenumbers,parallel=parallel,simplify=simplify,
+                    sparse=sparse,eval_expression=eval_expression,kwargs...)
+    return f, u0, p
+end
+
+function DiffEqBase.NonlinearProblem(sys::NonlinearSystem, args...; kwargs...)
+    NonlinearProblem{true}(sys, args...; kwargs...)
+end
+
 """
 ```julia
 function DiffEqBase.NonlinearProblem{iip}(sys::NonlinearSystem,u0map,
@@ -101,19 +227,9 @@ Generates an NonlinearProblem from a NonlinearSystem and allows for automaticall
 symbolically calculating numerical enhancements.
 """
 function DiffEqBase.NonlinearProblem{iip}(sys::NonlinearSystem,u0map,
-                                          parammap=DiffEqBase.NullParameters();
-                                          jac = false, sparse=false,
-                                          checkbounds = false,
-                                          linenumbers = true, parallel=SerialForm(),
-                                          kwargs...) where iip
-    dvs = states(sys)
-    ps = parameters(sys)
-
-    f = generate_function(sys;checkbounds=checkbounds,linenumbers=linenumbers,
-                              parallel=parallel,sparse=sparse,expression=Val{false})
-    u0 = varmap_to_vars(u0map,dvs; defaults=get_default_u0(sys))
-    p = varmap_to_vars(parammap,ps; defaults=get_default_p(sys))
-    NonlinearProblem(f,u0,p;kwargs...)
+                                    parammap=DiffEqBase.NullParameters();kwargs...) where iip
+    f, u0, p = process_NonlinearProblem(NonlinearFunction{iip}, sys, u0map, parammap; kwargs...)
+    NonlinearProblem{iip}(f,u0,p;kwargs...)
 end
 
 """
@@ -132,23 +248,22 @@ numerical enhancements.
 """
 struct NonlinearProblemExpr{iip} end
 
+function NonlinearProblemExpr(sys::NonlinearSystem, args...; kwargs...)
+    NonlinearProblemExpr{true}(sys, args...; kwargs...)
+end
+
 function NonlinearProblemExpr{iip}(sys::NonlinearSystem,u0map,
-                                          parammap=DiffEqBase.NullParameters();
-                                          jac = false, sparse=false,
-                                          checkbounds = false,
-                                          linenumbers = false, parallel=SerialForm(),
-                                          kwargs...) where iip
-    dvs = states(sys)
-    ps = parameters(sys)
+                             parammap=DiffEqBase.NullParameters();
+                             kwargs...) where iip
+
+    f, u0, p = process_NonlinearProblem(NonlinearFunctionExpr{iip}, sys, u0map, parammap; kwargs...)
+    linenumbers = get(kwargs, :linenumbers, true)
 
-    f = generate_function(sys;checkbounds=checkbounds,linenumbers=linenumbers,
-                              parallel=parallel,sparse=sparse,expression=Val{true})
-    u0 = varmap_to_vars(u0map,dvs; defaults=get_default_u0(sys))
-    p = varmap_to_vars(parammap,ps; defaults=get_default_p(sys))
-    quote
+    ex = quote
         f = $f
         u0 = $u0
         p = $p
-        NonlinearProblem(f,u0,p;kwargs...)
+        NonlinearProblem(f,u0,p;$(kwargs...))
     end
+    !linenumbers ? striplines(ex) : ex
 end

test/nonlinearsystem.jl

@@ -1,6 +1,7 @@
 using ModelingToolkit, StaticArrays, LinearAlgebra
 using DiffEqBase, SparseArrays
 using Test
+using NonlinearSolve
 using ModelingToolkit: value
 
 canonequal(a, b) = isequal(simplify(a), simplify(b))
@@ -62,9 +63,14 @@ eqs = [0 ~ σ*a,
        0 ~ x*y - β*z]
 ns = NonlinearSystem(eqs, [x,y,z], [σ,ρ,β])
 nlsys_func = generate_function(ns, [x,y,z], [σ,ρ,β])
+nf = NonlinearFunction(ns)
 jac = calculate_jacobian(ns)
 
 @test ModelingToolkit.jacobian_sparsity(ns).colptr == sparse(jac).colptr
 @test ModelingToolkit.jacobian_sparsity(ns).rowval == sparse(jac).rowval
 
 jac = generate_jacobian(ns)
+
+prob = NonlinearProblem(ns,ones(3),ones(3))
+sol = solve(prob,NewtonRaphson())
+@test sol.u[1] ≈ sol.u[2]