Merge pull request #765 from SciML/myb/pretty_print

Pretty print and some fixes

Author: YingboMa

src/ModelingToolkit.jl

@@ -152,13 +152,13 @@ end
 tosymbolic(a::Num) = tosymbolic(value(a))
 tosymbolic(a::Sym) = tovar(a)
 tosymbolic(a) = a
-@num_method Base.isless isless(tosymbolic(a), tosymbolic(b)) (Real,)
-@num_method Base.:(<) (tosymbolic(a) < tosymbolic(b)) (Real,)
-@num_method Base.:(<=) (tosymbolic(a) <= tosymbolic(b)) (Real,)
-@num_method Base.:(>) (tosymbolic(a) > tosymbolic(b)) (Real,)
-@num_method Base.:(>=) (tosymbolic(a) >= tosymbolic(b)) (Real,)
+@num_method Base.isless  (val = isless(tosymbolic(a), tosymbolic(b)); val isa Bool ? val : Num(val)) (Real,)
+@num_method Base.:(<)    (val = tosymbolic(a) < tosymbolic(b)       ; val isa Bool ? val : Num(val)) (Real,)
+@num_method Base.:(<=)   (val = tosymbolic(a) <= tosymbolic(b)      ; val isa Bool ? val : Num(val)) (Real,)
+@num_method Base.:(>)    (val = tosymbolic(a) > tosymbolic(b)       ; val isa Bool ? val : Num(val)) (Real,)
+@num_method Base.:(>=)   (val = tosymbolic(a) >= tosymbolic(b)      ; val isa Bool ? val : Num(val)) (Real,)
+@num_method Base.:(==)   (val = tosymbolic(a) == tosymbolic(b)      ; val isa Bool ? val : Num(val)) (AbstractFloat,Number)
 @num_method Base.isequal isequal(tosymbolic(a), tosymbolic(b)) (AbstractFloat, Number, Symbolic)
-@num_method Base.:(==) tosymbolic(a) == tosymbolic(b) (AbstractFloat,Number)
 
 Base.hash(x::Num, h::UInt) = hash(value(x), h)
 

src/equations.jl

@@ -14,6 +14,7 @@ struct Equation
 end
 Base.:(==)(a::Equation, b::Equation) = all(isequal.((a.lhs, a.rhs), (b.lhs, b.rhs)))
 Base.hash(a::Equation, salt::UInt) = hash(a.lhs, hash(a.rhs, salt))
+
 Base.show(io::IO, eq::Equation) = print(io, eq.lhs, " ~ ", eq.rhs)
 
 SymbolicUtils.simplify(x::Equation; kw...) = simplify(x.lhs; kw...) ~ simplify(x.rhs; kw...)

src/extra_functions.jl

@@ -1,5 +1,5 @@
-@register Base.getindex(x,i::Integer)
-@register Base.getindex(x,i)
+@register Base.getindex(x,i::Integer) false
+@register Base.getindex(x,i) false
 @register Base.binomial(n,k)
 
 @register Base.signbit(x)
@@ -22,7 +22,7 @@ function ModelingToolkit.derivative(::typeof(max), args::NTuple{2,Any}, ::Val{2}
     IfElse.ifelse(x > y, zero(y), one(y))
 end
 
-@register IfElse.ifelse(x,y,z::Any)
+IfElse.ifelse(x::Num,y,z) = Num(Term{Real}(IfElse.ifelse, [value(x), value(y), value(z)]))
 ModelingToolkit.derivative(::typeof(IfElse.ifelse), args::NTuple{3,Any}, ::Val{1}) = 0
 ModelingToolkit.derivative(::typeof(IfElse.ifelse), args::NTuple{3,Any}, ::Val{2}) = IfElse.ifelse(args[1],1,0)
 ModelingToolkit.derivative(::typeof(IfElse.ifelse), args::NTuple{3,Any}, ::Val{3}) = IfElse.ifelse(args[1],0,1)
@@ -36,8 +36,8 @@ ModelingToolkit.@register Distributions.cdf(dist,x)
 ModelingToolkit.@register Distributions.logcdf(dist,x)
 ModelingToolkit.@register Distributions.quantile(dist,x)
 
-ModelingToolkit.@register Distributions.Uniform(mu,sigma)
-ModelingToolkit.@register Distributions.Normal(mu,sigma)
+ModelingToolkit.@register Distributions.Uniform(mu,sigma) false
+ModelingToolkit.@register Distributions.Normal(mu,sigma) false
 
 @register ∈(x::Num, y::AbstractArray)
 @register ∪(x, y)
@@ -74,4 +74,4 @@ function LinearAlgebra.det(A::AbstractMatrix{<:Num}; laplace=true)
         end
         return det(lu(A; check = false))
     end
-end
+end

src/register_function.jl

@@ -1,5 +1,5 @@
 """
-    @register(expr, Ts = [Num, Symbolic, Real])
+    @register(expr, define_promotion, Ts = [Num, Symbolic, Real])
 
 Overload approperate methods such that ModelingToolkit can stop tracing into the
 registered function.
@@ -11,7 +11,7 @@ registered function.
 @register hoo(x, y)::Int # `hoo` returns `Int`
 ```
 """
-macro register(expr, Ts = [Num, Symbolic, Real])
+macro register(expr, define_promotion = true, Ts = [Num, Symbolic, Real])
     if expr.head === :(::)
         ret_type = expr.args[2]
         expr = expr.args[1]
@@ -48,7 +48,7 @@ macro register(expr, Ts = [Num, Symbolic, Real])
     push!(
           ex.args,
           quote
-              if $!($hasmethod($promote_symtype, $Tuple{$typeof($f), $Vararg}))
+              if $define_promotion
                   (::$typeof($promote_symtype))(::$typeof($f), args...) = $ret_type
               end
           end

src/solve.jl

@@ -82,17 +82,23 @@ Assumes `length(eqs) == length(vars)`
 
 Currently only works if all equations are linear.
 """
-function solve_for(eqs, vars)
+function solve_for(eqs, vars; simplify=true)
     A, b = A_b(eqs, vars)
-    _solve(A, b)
+    #TODO: we need to make sure that `solve_for(eqs, vars)` contains no `vars`
+    _solve(A, b, simplify)
 end
 
-function _solve(A::AbstractMatrix, b::AbstractArray)
+function _solve(A::AbstractMatrix, b::AbstractArray, do_simplify)
     A = SymbolicUtils.simplify.(Num.(A), polynorm=true)
     b = SymbolicUtils.simplify.(Num.(b), polynorm=true)
-    value.(SymbolicUtils.simplify.(sym_lu(A) \ b))
+    sol = value.(sym_lu(A) \ b)
+    do_simplify ? SymbolicUtils.simplify.(sol, polynorm=true) : sol
+end
+
+function _solve(a, b, do_simplify)
+    sol = value(b/a)
+    do_simplify ? SymbolicUtils.simplify(sol, polynorm=true) : sol
 end
-_solve(a, b) = value(SymbolicUtils.simplify(b/a, polynorm=true))
 
 # ldiv below
 

src/systems/abstractsystem.jl

@@ -316,3 +316,53 @@ function (f::AbstractSysToExpr)(O)
     end
     return build_expr(:call, Any[operation(O); f.(arguments(O))])
 end
+
+function Base.show(io::IO, sys::AbstractSystem)
+    eqs = equations(sys)
+    Base.printstyled(io, "Equations ($(length(eqs))):\n"; bold=true)
+    Base.print_matrix(io, eqs)
+    println(io)
+
+    rows = first(displaysize(io)) ÷ 3
+    limit = get(io, :limit, false)
+
+    vars = states(sys); nvars = length(vars)
+    Base.printstyled(io, "States ($nvars):"; bold=true)
+    nrows = min(nvars, limit ? rows : nvars)
+    limited = nrows < length(vars)
+    d_u0 = default_u0(sys)
+    for i in 1:nrows
+        s = vars[i]
+        print(io, "\n  ", s)
+
+        val = get(d_u0, s, nothing)
+        if val !== nothing
+            print(io, " [defaults to $val]")
+        end
+    end
+    limited && print(io, "\n⋮")
+    println(io)
+
+    vars = parameters(sys); nvars = length(vars)
+    Base.printstyled(io, "Parameters ($nvars):"; bold=true)
+    nrows = min(nvars, limit ? rows : nvars)
+    limited = nrows < length(vars)
+    d_p = default_p(sys)
+    for i in 1:nrows
+        s = vars[i]
+        print(io, "\n  ", s)
+
+        val = get(d_p, s, nothing)
+        if val !== nothing
+            print(io, " [defaults to $val]")
+        end
+    end
+    limited && print(io, "\n⋮")
+
+    s = get_structure(sys)
+    if s !== nothing
+        Base.printstyled(io, "\nIncidence matrix:"; color=:magenta)
+        show(io, incidence_matrix(s.graph, Num(Sym{Real}(:×))))
+    end
+    return nothing
+end

src/systems/diffeqs/odesystem.jl

@@ -138,6 +138,7 @@ function ODESystem(eqs, iv=nothing; kwargs...)
             end
         end
     end
+    iv = value(iv)
     iv === nothing && throw(ArgumentError("Please pass in independent variables."))
     for eq in eqs
         collect_vars!(allstates, ps, eq.lhs, iv)

src/systems/systemstructure.jl

@@ -95,7 +95,6 @@ function init_graph(sys)
     algvar_offset = 2dxvar_offset
 
     fullvars = [xvars; dxvars; algvars]
-    sys = reordersys(sys, dxvar_offset, fullvars)
     eqs = equations(sys)
     idxmap = Dict(fullvars .=> 1:length(fullvars))
     graph = BipartiteGraph(length(eqs), length(fullvars))
@@ -122,21 +121,3 @@ function init_graph(sys)
     varassoc = Int[(1:dxvar_offset) .+ dxvar_offset; zeros(Int, length(fullvars) - dxvar_offset)] # variable association list
     sys, dxvar_offset, fullvars, varassoc, graph, solvable_graph
 end
-
-function reordersys(sys, dxvar_offset, fullvars)
-    eqs = equations(sys)
-    neweqs = similar(eqs, Equation)
-    eqidxmap = Dict(@view(fullvars[dxvar_offset+1:2dxvar_offset]) .=> (1:dxvar_offset))
-    varidxmap = Dict([@view(fullvars[1:dxvar_offset]); @view(fullvars[2dxvar_offset+1:end])] .=> (1:length(fullvars)-dxvar_offset))
-    algidx = dxvar_offset
-    for eq in eqs
-        if isdiffeq(eq)
-            neweqs[eqidxmap[eq.lhs]] = eq
-        else
-            neweqs[algidx+=1] = eq
-        end
-    end
-    sts = states(sys)
-    @set! sys.eqs = neweqs
-    @set! sys.states = sts[map(s->varidxmap[s], sts)]
-end

test/operation_overloads.jl

@@ -25,7 +25,7 @@ F = lu(X)
 R = simplify.(F.L * F.U - X[F.p, :], polynorm=true)
 @test iszero(R)
 @test simplify.(F \ X) == I
-@test ModelingToolkit._solve(X, X) == I
+@test ModelingToolkit._solve(X, X, true) == I
 inv(X)
 qr(X)
 

test/reduction.jl

@@ -97,13 +97,13 @@ reduced_system = alias_elimination(connected; conservative=false)
        ]) |> isempty
 
 reduced_eqs = [
+               0 ~ a + lorenz1.x - (lorenz2.y)
                D(lorenz1.x) ~ lorenz2.x + lorenz2.y + lorenz1.σ*((lorenz1.y) - (lorenz1.x)) - (lorenz2.z)
                D(lorenz1.y) ~ lorenz1.x*(lorenz1.ρ - (lorenz1.z)) - ((lorenz1.x) + (lorenz1.y) - (lorenz1.z))
                D(lorenz1.z) ~ lorenz1.x*lorenz1.y - (lorenz1.β*(lorenz1.z))
                D(lorenz2.x) ~ lorenz1.x + lorenz1.y + lorenz2.σ*((lorenz2.y) - (lorenz2.x)) - (lorenz1.z)
                D(lorenz2.y) ~ lorenz2.x*(lorenz2.ρ - (lorenz2.z)) - ((lorenz2.x) + (lorenz2.y) - (lorenz2.z))
                D(lorenz2.z) ~ lorenz2.x*lorenz2.y - (lorenz2.β*(lorenz2.z))
-               0 ~ a + lorenz1.x - (lorenz2.y)
               ]
 
 test_equal.(equations(reduced_system), reduced_eqs)