Merge pull request #978 from SciML/myb/print

Round-trip system printing

Author: YingboMa

src/systems/abstractsystem.jl

@@ -411,6 +411,98 @@ function (f::AbstractSysToExpr)(O)
     return build_expr(:call, Any[operation(O); f.(arguments(O))])
 end
 
+###
+### System utils
+###
+function push_vars!(stmt, name, typ, vars)
+    isempty(vars) && return
+    vars_expr = Expr(:macrocall, typ, nothing)
+    for s in vars
+        if istree(s)
+            f = nameof(operation(s))
+            args = arguments(s)
+            ex = :($f($(args...)))
+        else
+            ex = nameof(s)
+        end
+        push!(vars_expr.args, ex)
+    end
+    push!(stmt, :($name = $collect($vars_expr)))
+    return
+end
+
+function round_trip_expr(t, var2name)
+    name = get(var2name, t, nothing)
+    name !== nothing && return name
+    t isa Sym && return nameof(t)
+    istree(t) || return t
+    f = round_trip_expr(operation(t), var2name)
+    args = map(Base.Fix2(round_trip_expr, var2name), arguments(t))
+    return :($f($(args...)))
+end
+round_trip_eq(eq, var2name) = Expr(:call, :~, round_trip_expr(eq.lhs, var2name), round_trip_expr(eq.rhs, var2name))
+
+function push_eqs!(stmt, eqs, var2name)
+    eqs_name = gensym(:eqs)
+    eqs_expr = Expr(:vcat)
+    eqs_blk = Expr(:(=), eqs_name, eqs_expr)
+    for eq in eqs
+        push!(eqs_expr.args, round_trip_eq(eq, var2name))
+    end
+
+    push!(stmt, eqs_blk)
+    return eqs_name
+end
+
+function push_defaults!(stmt, defs, var2name)
+    defs_name = gensym(:defs)
+    defs_expr = Expr(:call, Dict)
+    defs_blk = Expr(:(=), defs_name, defs_expr)
+    for d in defs
+        n = round_trip_expr(d.first, var2name)
+        v = round_trip_expr(d.second, var2name)
+        push!(defs_expr.args, :($(=>)($n, $v)))
+    end
+
+    push!(stmt, defs_blk)
+    return defs_name
+end
+
+function toexpr(sys::AbstractSystem)
+    sys = flatten(sys)
+    expr = Expr(:block)
+    stmt = expr.args
+
+    iv = independent_variable(sys)
+    ivname = gensym(:iv)
+    if iv !== nothing
+        push!(stmt, :($ivname = (@variables $(getname(iv)))[1]))
+    end
+
+    stsname = gensym(:sts)
+    sts = states(sys)
+    push_vars!(stmt, stsname, Symbol("@variables"), sts)
+    psname = gensym(:ps)
+    ps = parameters(sys)
+    push_vars!(stmt, psname, Symbol("@parameters"), ps)
+
+    var2name = Dict{Any,Symbol}()
+    for v in Iterators.flatten((sts, ps))
+        var2name[v] = getname(v)
+    end
+
+    eqs_name = push_eqs!(stmt, equations(sys), var2name)
+    defs_name = push_defaults!(stmt, defaults(sys), var2name)
+
+    if sys isa ODESystem
+        push!(stmt, :($ODESystem($eqs_name, $ivname, $stsname, $psname; defaults=$defs_name)))
+    elseif sys isa NonlinearSystem
+        push!(stmt, :($NonlinearSystem($eqs_name, $stsname, $psname; defaults=$defs_name)))
+    end
+
+    striplines(expr) # keeping the line numbers is never helpful
+end
+
 function Base.show(io::IO, ::MIME"text/plain", sys::AbstractSystem)
     eqs = equations(sys)
     if eqs isa AbstractArray
@@ -581,6 +673,10 @@ function check_eqs_u0(eqs, dvs, u0)
     return nothing
 end
 
+###
+### Connectors
+###
+
 function with_connection_type(expr)
     @assert expr isa Expr && (expr.head == :function || (expr.head == :(=) &&
                                        expr.args[1] isa Expr &&

src/systems/nonlinear/nonlinearsystem.jl

@@ -308,3 +308,10 @@ function flatten(sys::NonlinearSystem)
                               )
     end
 end
+
+function Base.:(==)(sys1::NonlinearSystem, sys2::NonlinearSystem)
+    _eq_unordered(get_eqs(sys1), get_eqs(sys2)) &&
+    _eq_unordered(get_states(sys1), get_states(sys2)) &&
+    _eq_unordered(get_ps(sys1), get_ps(sys2)) &&
+    all(s1 == s2 for (s1, s2) in zip(get_systems(sys1), get_systems(sys2)))
+end

test/nonlinearsystem.jl

@@ -21,7 +21,8 @@ end
 eqs = [0 ~ σ*(y-x),
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
-ns = NonlinearSystem(eqs, [x,y,z], [σ,ρ,β])
+ns = NonlinearSystem(eqs, [x,y,z], [σ,ρ,β], defaults = Dict(x => 2))
+@test eval(toexpr(ns)) == ns
 test_nlsys_inference("standard", ns, (x, y, z), (σ, ρ, β))
 @test begin
     f = eval(generate_function(ns, [x,y,z], [σ,ρ,β])[2])

test/odesystem.jl

@@ -17,8 +17,8 @@ eqs = [D(x) ~ σ*(y-x),
        D(z) ~ x*y - β*z]
 
 ModelingToolkit.toexpr.(eqs)[1]
-:(derivative(x(t), t) = σ * (y(t) - x(t))).args
-de = ODESystem(eqs)
+de = ODESystem(eqs; defaults=Dict(x => 1))
+@test eval(toexpr(de)) == de
 
 generate_function(de)