diff --git a/README.md b/README.md index 124be64f3e..258123d3a9 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,7 @@ and parameters. Therefore we label them as follows: using ModelingToolkit # Define some variables -@parameters t σ ρ β +@parameters t() σ() ρ() β() @variables x(t) y(t) z(t) @derivatives D'~t ``` @@ -80,22 +80,22 @@ state of the previous ODE. This is the nonlinear system defined by where the derivatives are zero. We use (unknown) variables for our nonlinear system. ```julia -@variables x y z -@parameters σ ρ β +@variables x() y() z() +@parameters σ() ρ() β() # Define a nonlinear system eqs = [0 ~ σ*(y-x), 0 ~ x*(ρ-z)-y, 0 ~ x*y - β*z] ns = NonlinearSystem(eqs, [x,y,z]) -nlsys_func = generate_function(ns, [x,y,z], [ρ,σ,β]) +nlsys_func = generate_function(ns, [x,y,z], [σ,ρ,β]) ``` which generates: ```julia :((##364, u, p)->begin - let (x, z, y, ρ, σ, β) = (u[1], u[2], u[3], p[1], p[2], p[3]) + let (x, y, z, σ, ρ, β) = (u[1], u[2], u[3], p[1], p[2], p[3]) ##364[1] = σ * (y - x) ##364[2] = x * (ρ - z) - y ##364[3] = x * y - β * z diff --git a/src/systems/diffeqs/diffeqsystem.jl b/src/systems/diffeqs/diffeqsystem.jl index 2c78cffb0c..215276224f 100644 --- a/src/systems/diffeqs/diffeqsystem.jl +++ b/src/systems/diffeqs/diffeqsystem.jl @@ -94,11 +94,6 @@ function calculate_jacobian(sys::ODESystem) return jac end -function generate_jacobian(sys::ODESystem; version::FunctionVersion = ArrayFunction) - jac = calculate_jacobian(sys) - return build_function(jac, sys.dvs, sys.ps, (sys.iv.name,); version = version) -end - struct ODEToExpr sys::ODESystem end @@ -113,6 +108,11 @@ function (f::ODEToExpr)(O::Operation) end (f::ODEToExpr)(x) = convert(Expr, x) +function generate_jacobian(sys::ODESystem; version::FunctionVersion = ArrayFunction) + jac = calculate_jacobian(sys) + return build_function(jac, sys.dvs, sys.ps, (sys.iv.name,), ODEToExpr(sys); version = version) +end + function generate_function(sys::ODESystem, dvs, ps; version::FunctionVersion = ArrayFunction) rhss = [deq.rhs for deq ∈ sys.eqs] dvs′ = [clean(dv) for dv ∈ dvs] diff --git a/src/systems/nonlinear/nonlinear_system.jl b/src/systems/nonlinear/nonlinear_system.jl index 3940dfed0a..fda2626fe9 100644 --- a/src/systems/nonlinear/nonlinear_system.jl +++ b/src/systems/nonlinear/nonlinear_system.jl @@ -45,12 +45,26 @@ end function generate_jacobian(sys::NonlinearSystem; version::FunctionVersion = ArrayFunction) jac = calculate_jacobian(sys) - return build_function(jac, clean.(sys.vs), sys.ps; version = version) + return build_function(jac, clean.(sys.vs), sys.ps, (), NLSysToExpr(sys); version = version) end +struct NLSysToExpr + sys::NonlinearSystem +end +function (f::NLSysToExpr)(O::Operation) + any(isequal(O), f.sys.vs) && return O.op.name # variables + if isa(O.op, Variable) + isempty(O.args) && return O.op.name # 0-ary parameters + return build_expr(:call, Any[O.op.name; f.(O.args)]) + end + return build_expr(:call, Any[O.op; f.(O.args)]) +end +(f::NLSysToExpr)(x) = convert(Expr, x) + + function generate_function(sys::NonlinearSystem, vs, ps; version::FunctionVersion = ArrayFunction) rhss = [eq.rhs for eq ∈ sys.eqs] vs′ = [clean(v) for v ∈ vs] ps′ = [clean(p) for p ∈ ps] - return build_function(rhss, vs′, ps′; version = version) + return build_function(rhss, vs′, ps′, (), NLSysToExpr(sys); version = version) end diff --git a/test/system_construction.jl b/test/system_construction.jl index 2a91b4c994..2126a2b5b2 100644 --- a/test/system_construction.jl +++ b/test/system_construction.jl @@ -119,8 +119,12 @@ eqs = [0 ~ σ*(y-x), 0 ~ x*y - β*z] ns = NonlinearSystem(eqs, [x,y,z]) test_nlsys_inference("standard", ns, (x, y, z), (σ, ρ, β)) - -generate_function(ns, [x,y,z], [σ,ρ,β]) +@test begin + f = eval(generate_function(ns, [x,y,z], [σ,ρ,β])) + du = [0.0, 0.0, 0.0] + f(du, [1,2,3], [1,2,3]) + du ≈ [1, -3, -7] +end @derivatives D'~t @parameters A() B() C()