Create function for pretty printing of graphs

Author: MilesCranmer

src/DynamicExpressions.jl

@@ -36,7 +36,7 @@ import .EquationModule: constructorof, preserve_sharing
     has_constants,
     get_constants,
     set_constants!
-@reexport import .StringsModule: string_tree, print_tree
+@reexport import .StringsModule: string_tree, print_tree, pretty_string_graph
 @reexport import .OperatorEnumModule: AbstractOperatorEnum
 @reexport import .OperatorEnumConstructionModule:
     OperatorEnum, GenericOperatorEnum, @extend_operators, set_default_variable_names!

src/Strings.jl

@@ -1,5 +1,7 @@
 module StringsModule
 
+import Compat: Returns
+
 import ..UtilsModule: deprecate_varmap
 import ..OperatorEnumModule: AbstractOperatorEnum
 import ..EquationModule: AbstractExpressionNode, tree_mapreduce
@@ -181,4 +183,131 @@ for io in ((), (:(io::IO),))
     end
 end
 
+function pretty_string_graph(
+    tree::AbstractExpressionNode{T},
+    operators::Union{AbstractOperatorEnum,Nothing}=nothing;
+    f_variable::F1=string_variable,
+    f_constant::F2=string_constant,
+    variable_names::Union{Vector{String},Nothing}=nothing,
+) where {T,F1,F2}
+    output = Char[]
+
+    # First, we build a mapping of shared nodes. We skip varible nodes
+    # as it wouldn't simplify printing.
+    shared_nodes = _build_shared_node_dict(tree, node -> node.degree != 0 || node.constant)
+    # NOTE THAT THIS IS ASSUMED BY _leaf_string_or_shared_variable
+
+    # We collect the shared nodes in order of appearance, and will print
+    # them in that order (deepest first)
+    iter = collect(values(shared_nodes))
+    sort!(iter; by=x -> x.index)
+
+    # We also want to print the final expression:
+    push!(iter, (node=tree, index=length(iter) + 1))
+
+    for (; node, index) in iter
+        raw_output, _ = tree_mapreduce(
+            leaf -> _leaf_string_or_shared_variable(
+                index, leaf, shared_nodes; f_variable, f_constant, variable_names
+            ),
+            branch ->
+                _branch_string_or_shared_variable(index, branch, shared_nodes, operators),
+            _combine_op_with_inputs_or_shared_variable,
+            node,
+            @NamedTuple{chars::Vector{Char}, shared::Bool};
+            break_sharing=Val(true),
+        )
+        is_output = index == length(iter)
+        if is_output
+            if !isempty(shared_nodes)
+                append!(output, ('│', '\n'))
+            end
+            append!(output, ('=', ' '))
+        else
+            append!(output, (index == 1 ? '┌' : '├', '─', '─', '─', ' '))
+            append!(output, _get_z_name(index))
+            append!(output, (' ', '=', ' '))
+        end
+        append!(output, strip_brackets(raw_output))
+        append!(output, ('\n',))
+    end
+    return String(output)
+end
+
+function _build_shared_node_dict(tree, filter::F=Returns(true)) where {F}
+    i = Ref(0)
+    node_counts = Dict{UInt,@NamedTuple{node::typeof(tree), index::Int}}()
+    tree_mapreduce(
+        p -> p,
+        (p, _...) -> p,
+        tree,
+        typeof(tree);
+        f_on_shared=(node, is_shared) -> begin
+            if is_shared && filter(node)
+                oid = objectid(node)
+                if !haskey(node_counts, oid)
+                    node_counts[oid] = (node=node, index=(i[] += 1))
+                end
+            end
+            node
+        end,
+    )
+    return node_counts
+end
+
+function _get_z_name(j::Int)
+    out = ['z']
+    for k in digits(j)
+        push!(out, Char('0' + k))
+    end
+    return out
+end
+
+function _combine_op_with_inputs_or_shared_variable(p, c...)
+    if p.shared
+        # We assume that this is an intermediate variable, so don't wish to expand it!
+        return (chars=p.chars, shared=false)
+    else
+        return (
+            chars=combine_op_with_inputs(p.chars, (ci -> ci.chars).(c)...), shared=false
+        )
+    end
+end
+
+function _leaf_string_or_shared_variable(
+    cur_intermediate_variable,
+    leaf::AbstractExpressionNode{T},
+    shared_nodes;
+    f_variable::F1,
+    f_constant::F2,
+    variable_names,
+) where {T,F1,F2}
+    if leaf.constant
+        oid = objectid(leaf)
+        if haskey(shared_nodes, oid) &&
+            (cur_i = @inbounds(shared_nodes[oid][2])) < cur_intermediate_variable
+            return (chars=_get_z_name(cur_i), shared=true)
+        else
+            return (chars=collect(f_constant(leaf.val::T)), shared=false)
+        end
+    else
+        return (chars=collect(f_variable(leaf.feature, variable_names)), shared=false)
+    end
+end
+function _branch_string_or_shared_variable(
+    cur_intermediate_variable, branch, shared_nodes, operators
+)
+    oid = objectid(branch)
+    if haskey(shared_nodes, oid) &&
+        (cur_i = @inbounds(shared_nodes[oid][2])) < cur_intermediate_variable
+        return (chars=_get_z_name(cur_i), shared=true)
+    else
+        if branch.degree == 1
+            return (chars=dispatch_op_name(Val(1), operators, branch.op), shared=false)
+        else
+            return (chars=dispatch_op_name(Val(2), operators, branch.op), shared=false)
+        end
+    end
+end
+
 end