Merge pull request #919 from anandijain/master

[wip] pretty print system

Author: YingboMa

Project.toml

@@ -4,6 +4,7 @@ authors = ["Chris Rackauckas <[email protected]>"]
 version = "5.13.6"
 
 [deps]
+AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
@@ -35,11 +36,11 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
-TreeViews = "a2a6695c-b41b-5b7d-aed9-dbfdeacea5d7"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [compat]
+AbstractTrees = "0.3"
 ArrayInterface = "2.8, 3.0"
 ConstructionBase = "1"
 DataStructures = "0.17, 0.18"
@@ -66,7 +67,6 @@ SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"
 SymbolicUtils = "0.8.3, 0.9"
 Symbolics = "0.1.8"
-TreeViews = "0.3"
 UnPack = "0.1, 1.0"
 Unitful = "1.1"
 julia = "1.2"

src/ModelingToolkit.jl

@@ -3,7 +3,7 @@ $(DocStringExtensions.README)
 """
 module ModelingToolkit
 using DocStringExtensions
-
+using AbstractTrees
 using DiffEqBase, SciMLBase, Reexport
 using Distributed
 using StaticArrays, LinearAlgebra, SparseArrays, LabelledArrays
@@ -54,8 +54,6 @@ import DiffEqBase: @add_kwonly
 
 import LightGraphs: SimpleDiGraph, add_edge!
 
-import TreeViews
-
 using Requires
 
 for fun in [:toexpr]

src/systems/abstractsystem.jl

@@ -539,3 +539,7 @@ struct InvalidSystemException <: Exception
     msg::String
 end
 Base.showerror(io::IO, e::InvalidSystemException) = print(io, "InvalidSystemException: ", e.msg)
+
+AbstractTrees.children(sys::ModelingToolkit.AbstractSystem) = ModelingToolkit.get_systems(sys)
+AbstractTrees.printnode(io::IO, sys::ModelingToolkit.AbstractSystem) = print(io, nameof(sys))
+AbstractTrees.nodetype(::ModelingToolkit.AbstractSystem) = ModelingToolkit.AbstractSystem

test/components.jl

@@ -1,86 +1,8 @@
 using Test
 using ModelingToolkit, OrdinaryDiffEq
 
-# Basic electric components
-const t = Sym{ModelingToolkit.Parameter{Real}}(:t)
-function Pin(;name)
-    @variables v(t) i(t)
-    ODESystem(Equation[], t, [v, i], [], name=name, defaults=[v=>1.0, i=>1.0])
-end
+include("rc_model.jl")
 
-function Ground(name)
-    @named g = Pin()
-    eqs = [g.v ~ 0]
-    ODESystem(eqs, t, [], [], systems=[g], name=name)
-end
-
-function ConstantVoltage(name; V = 1.0)
-    val = V
-    @named p = Pin()
-    @named n = Pin()
-    @parameters V
-    eqs = [
-           V ~ p.v - n.v
-           0 ~ p.i + n.i
-          ]
-    ODESystem(eqs, t, [], [V], systems=[p, n], defaults=Dict(V => val), name=name)
-end
-
-function Resistor(name; R = 1.0)
-    val = R
-    @named p = Pin()
-    @named n = Pin()
-    @variables v(t)
-    @parameters R
-    eqs = [
-           v ~ p.v - n.v
-           0 ~ p.i + n.i
-           v ~ p.i * R
-          ]
-    ODESystem(eqs, t, [v], [R], systems=[p, n], defaults=Dict(R => val), name=name)
-end
-
-function Capacitor(name; C = 1.0)
-    val = C
-    @named p = Pin()
-    @named n = Pin()
-    @variables v(t)
-    @parameters C
-    D = Differential(t)
-    eqs = [
-           v ~ p.v - n.v
-           0 ~ p.i + n.i
-           D(v) ~ p.i / C
-          ]
-    ODESystem(eqs, t, [v], [C], systems=[p, n], defaults=Dict(C => val), name=name)
-end
-
-R = 1.0
-C = 1.0
-V = 1.0
-resistor = Resistor(:resistor, R=R)
-capacitor = Capacitor(:capacitor, C=C)
-source = ConstantVoltage(:source, V=V)
-ground = Ground(:ground)
-
-function connect(ps...)
-    eqs = [
-           0 ~ sum(p->p.i, ps) # KCL
-          ]
-    # KVL
-    for i in 1:length(ps)-1
-        push!(eqs, ps[i].v ~ ps[i+1].v)
-    end
-
-    return eqs
-end
-rc_eqs = [
-          connect(source.p, resistor.p)
-          connect(resistor.n, capacitor.p)
-          connect(capacitor.n, source.n, ground.g)
-         ]
-
-rc_model = ODESystem(rc_eqs, t, systems=[resistor, capacitor, source, ground], name=:rc)
 sys = structural_simplify(rc_model)
 @test !isempty(ModelingToolkit.defaults(sys))
 u0 = [

test/print_tree.jl

@@ -0,0 +1,22 @@
+using ModelingToolkit, AbstractTrees, Test
+
+include("rc_model.jl")
+
+io = IOBuffer()
+print_tree(io, rc_model)
+ser = String(take!(io))
+str = 
+"""rc
+├─ resistor
+│  ├─ p
+│  └─ n
+├─ capacitor
+│  ├─ p
+│  └─ n
+├─ source
+│  ├─ p
+│  └─ n
+└─ ground
+   └─ g
+"""
+@test strip(ser) == strip(str)

test/rc_model.jl

@@ -0,0 +1,83 @@
+using Test
+using ModelingToolkit, OrdinaryDiffEq
+
+# Basic electric components
+const t = Sym{ModelingToolkit.Parameter{Real}}(:t)
+function Pin(;name)
+    @variables v(t) i(t)
+    ODESystem(Equation[], t, [v, i], [], name=name, defaults=[v=>1.0, i=>1.0])
+end
+
+function Ground(name)
+    @named g = Pin()
+    eqs = [g.v ~ 0]
+    ODESystem(eqs, t, [], [], systems=[g], name=name)
+end
+
+function ConstantVoltage(name; V = 1.0)
+    val = V
+    @named p = Pin()
+    @named n = Pin()
+    @parameters V
+    eqs = [
+           V ~ p.v - n.v
+           0 ~ p.i + n.i
+          ]
+    ODESystem(eqs, t, [], [V], systems=[p, n], defaults=Dict(V => val), name=name)
+end
+
+function Resistor(name; R = 1.0)
+    val = R
+    @named p = Pin()
+    @named n = Pin()
+    @variables v(t)
+    @parameters R
+    eqs = [
+           v ~ p.v - n.v
+           0 ~ p.i + n.i
+           v ~ p.i * R
+          ]
+    ODESystem(eqs, t, [v], [R], systems=[p, n], defaults=Dict(R => val), name=name)
+end
+
+function Capacitor(name; C = 1.0)
+    val = C
+    @named p = Pin()
+    @named n = Pin()
+    @variables v(t)
+    @parameters C
+    D = Differential(t)
+    eqs = [
+           v ~ p.v - n.v
+           0 ~ p.i + n.i
+           D(v) ~ p.i / C
+          ]
+    ODESystem(eqs, t, [v], [C], systems=[p, n], defaults=Dict(C => val), name=name)
+end
+
+R = 1.0
+C = 1.0
+V = 1.0
+resistor = Resistor(:resistor, R=R)
+capacitor = Capacitor(:capacitor, C=C)
+source = ConstantVoltage(:source, V=V)
+ground = Ground(:ground)
+
+function connect(ps...)
+    eqs = [
+           0 ~ sum(p->p.i, ps) # KCL
+          ]
+    # KVL
+    for i in 1:length(ps)-1
+        push!(eqs, ps[i].v ~ ps[i+1].v)
+    end
+
+    return eqs
+end
+rc_eqs = [
+          connect(source.p, resistor.p)
+          connect(resistor.n, capacitor.p)
+          connect(capacitor.n, source.n, ground.g)
+         ]
+
+rc_model = ODESystem(rc_eqs, t, systems=[resistor, capacitor, source, ground], name=:rc)

test/runtests.jl

@@ -35,3 +35,4 @@ println("Last test requires gcc available in the path!")
 @safetestset "Latexify recipes Test" begin include("latexify.jl") end
 @safetestset "StructuralTransformations" begin include("structural_transformation/runtests.jl") end
 @testset "Serialization" begin include("serialization.jl") end
+@safetestset "print_tree" begin include("print_tree.jl") end