[wip] pretty print system

Author: anandijain

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/tree_print.jl

@@ -0,0 +1,108 @@
+
+using ModelingToolkit, AbstractTrees, Test
+
+@parameters t
+
+# Basic electric components
+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 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
+
+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
+
+R = 1.0
+C = 1.0
+V = 1.0
+@named resistor = Resistor(R=R)
+@named capacitor = Capacitor(C=C)
+@named source = ConstantVoltage(V=V)
+@named ground = Ground()
+
+function connect_pins(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_pins(source.p, resistor.p)
+          connect_pins(resistor.n, capacitor.p)
+          connect_pins(capacitor.n, source.n, ground.g)
+         ]
+
+@named rc_model = ODESystem(rc_eqs, t,
+                            systems=[resistor, capacitor, source, ground])
+
+open("rc_tree.txt", "w") do io
+    print_tree(io, rc_model)
+end;
+
+str = 
+"""rc_model
+├─ resistor
+│  ├─ p
+│  └─ n
+├─ capacitor
+│  ├─ p
+│  └─ n
+├─ source
+│  ├─ p
+│  └─ n
+└─ ground
+   └─ g
+"""
+@test read("rc_tree.txt", String) == str
+
+rm("rc_tree.txt")