Merge pull request #40 from SymbolicML/custom-printing

Allow user-defined printing

Author: MilesCranmer

benchmark/benchmarks.jl

@@ -98,14 +98,15 @@ function benchmark_utilities()
         :is_constant,
         :get_set_constants!,
         :index_constants,
+        :string_tree,
     )
 
     operators = OperatorEnum(; binary_operators=[+, -, /, *], unary_operators=[cos, exp])
-
     for func_k in all_funcs
         suite[func_k] = let s = BenchmarkGroup()
             for k in (:break_sharing, :preserve_sharing)
-                k == :preserve_sharing && !(func_k in (:copy, :convert)) && continue
+                has_both_modes = func_k in (:copy, :convert)
+                k == :preserve_sharing && !has_both_modes && continue
 
                 f = if func_k == :copy
                     tree -> _copy_node(tree; preserve_sharing=(k == :preserve_sharing))
@@ -115,7 +116,7 @@ function benchmark_utilities()
                         tree;
                         preserve_sharing=(k == :preserve_sharing),
                     )
-                elseif func_k in (:simplify_tree, :combine_operators)
+                elseif func_k in (:simplify_tree, :combine_operators, :string_tree)
                     g = getfield(@__MODULE__, func_k)
                     tree -> f_tree_op(g, tree, operators)
                 else
@@ -133,6 +134,9 @@ function benchmark_utilities()
                         trees=[gen_random_tree_fixed_size(n, $operators, 5, Float32) for _ in 1:ntrees]
                     )
                 )
+                if !has_both_modes
+                    s = s[k]
+                end
                 #! format: on
             end
             s

src/Equation.jl

@@ -190,110 +190,148 @@ const OP_NAMES = Dict(
     "safe_pow" => "^",
 )
 
-function get_op_name(op::String)
-    return get(OP_NAMES, op, op)
+@generated function get_op_name(op::F) where {F}
+    try
+        # Bit faster to just cache the name of the operator:
+        op_s = string(F.instance)
+        out = get(OP_NAMES, op_s, op_s)
+        return :($out)
+    catch
+    end
+    return quote
+        op_s = string(op)
+        out = get(OP_NAMES, op_s, op_s)
+        return out
+    end
 end
 
 function string_op(
-    op::F,
-    tree::Node,
-    operators::AbstractOperatorEnum;
-    bracketed::Bool=false,
-    variable_names::Union{Array{String,1},Nothing}=nothing,
-    # Deprecated
-    varMap=nothing,
+    ::Val{2}, op::F, tree::Node, args...; bracketed, kws...
 )::String where {F}
-    variable_names = deprecate_varmap(variable_names, varMap, :string_op)
-    op_name = get_op_name(string(op))
+    op_name = get_op_name(op)
     if op_name in ["+", "-", "*", "/", "^"]
-        l = string_tree(tree.l, operators; bracketed=false, variable_names=variable_names)
-        r = string_tree(tree.r, operators; bracketed=false, variable_names=variable_names)
+        l = string_tree(tree.l, args...; bracketed=false, kws...)
+        r = string_tree(tree.r, args...; bracketed=false, kws...)
         if bracketed
-            return "$l $op_name $r"
+            return l * " " * op_name * " " * r
         else
-            return "($l $op_name $r)"
+            return "(" * l * " " * op_name * " " * r * ")"
         end
     else
-        l = string_tree(tree.l, operators; bracketed=true, variable_names=variable_names)
-        r = string_tree(tree.r, operators; bracketed=true, variable_names=variable_names)
-        return "$op_name($l, $r)"
+        l = string_tree(tree.l, args...; bracketed=true, kws...)
+        r = string_tree(tree.r, args...; bracketed=true, kws...)
+        # return "$op_name($l, $r)"
+        return op_name * "(" * l * ", " * r * ")"
+    end
+end
+function string_op(
+    ::Val{1}, op::F, tree::Node, args...; bracketed, kws...
+)::String where {F}
+    op_name = get_op_name(op)
+    l = string_tree(tree.l, args...; bracketed=true, kws...)
+    return op_name * "(" * l * ")"
+end
+
+function string_constant(val, bracketed::Bool)
+    does_not_need_brackets = (typeof(val) <: Union{Real,AbstractArray})
+    if does_not_need_brackets || bracketed
+        string(val)
+    else
+        "(" * string(val) * ")"
+    end
+end
+
+function string_variable(feature, variable_names)
+    if variable_names === nothing
+        return "x" * string(feature)
+    else
+        return variable_names[feature]
     end
 end
 
 """
-    string_tree(tree::Node, operators::AbstractOperatorEnum; kws...)
+    string_tree(tree::Node, operators::AbstractOperatorEnum[; bracketed, variable_names, f_variable, f_constant])
 
 Convert an equation to a string.
 
 # Arguments
-
-- `variable_names::Union{Array{String, 1}, Nothing}=nothing`: what variables
-    to print for each feature.
+- `tree`: the tree to convert to a string
+- `operators`: the operators used to define the tree
+
+# Keyword Arguments
+- `bracketed`: (optional) whether to put brackets around the outside.
+- `f_variable`: (optional) function to convert a variable to a string, of the form `(feature::Int, variable_names)`.
+- `f_constant`: (optional) function to convert a constant to a string, of the form `(val, bracketed::Bool)`
+- `variable_names::Union{Array{String, 1}, Nothing}=nothing`: (optional) what variables to print for each feature.
 """
 function string_tree(
     tree::Node{T},
     operators::AbstractOperatorEnum;
     bracketed::Bool=false,
+    f_variable::F1=string_variable,
+    f_constant::F2=string_constant,
     variable_names::Union{Array{String,1},Nothing}=nothing,
     # Deprecated
     varMap=nothing,
-)::String where {T}
+)::String where {T,F1<:Function,F2<:Function}
     variable_names = deprecate_varmap(variable_names, varMap, :string_tree)
     if tree.degree == 0
-        if tree.constant
-            return string_constant(tree.val::T; bracketed=bracketed)
+        if !tree.constant
+            return f_variable(tree.feature, variable_names)
         else
-            if variable_names === nothing
-                return "x$(tree.feature)"
-            else
-                return variable_names[tree.feature]
-            end
+            return f_constant(tree.val::T, bracketed)
         end
     elseif tree.degree == 1
-        op_name = get_op_name(string(operators.unaops[tree.op]))
-        return "$(op_name)($(string_tree(tree.l, operators, bracketed=true, variable_names=variable_names)))"
+        return string_op(
+            Val(1),
+            operators.unaops[tree.op],
+            tree,
+            operators;
+            bracketed,
+            f_variable,
+            f_constant,
+            variable_names,
+        )
     else
         return string_op(
+            Val(2),
             operators.binops[tree.op],
             tree,
             operators;
-            bracketed=bracketed,
-            variable_names=variable_names,
+            bracketed,
+            f_variable,
+            f_constant,
+            variable_names,
         )
     end
 end
 
-string_constant(val::T; bracketed::Bool) where {T<:Union{Real,AbstractArray}} = string(val)
-function string_constant(val; bracketed::Bool)
-    if bracketed
-        string(val)
-    else
-        "(" * string(val) * ")"
-    end
-end
-
 # Print an equation
 function print_tree(
     io::IO,
     tree::Node,
     operators::AbstractOperatorEnum;
+    f_variable::F1=string_variable,
+    f_constant::F2=string_constant,
     variable_names::Union{Array{String,1},Nothing}=nothing,
     # Deprecated
     varMap=nothing,
-)
+) where {F1<:Function,F2<:Function}
     variable_names = deprecate_varmap(variable_names, varMap, :print_tree)
-    return println(io, string_tree(tree, operators; variable_names=variable_names))
+    return println(io, string_tree(tree, operators; f_variable, f_constant, variable_names))
 end
 
 function print_tree(
     tree::Node,
     operators::AbstractOperatorEnum;
+    f_variable::F1=string_variable,
+    f_constant::F2=string_constant,
     variable_names::Union{Array{String,1},Nothing}=nothing,
     # Deprecated
     varMap=nothing,
-)
+) where {F1<:Function,F2<:Function}
     variable_names = deprecate_varmap(variable_names, varMap, :print_tree)
-    return println(string_tree(tree, operators; variable_names=variable_names))
+    return println(string_tree(tree, operators; f_variable, f_constant, variable_names))
 end
 
 end

test/test_print.jl

@@ -1,5 +1,6 @@
 using Test
 using DynamicExpressions
+import Compat: Returns
 
 include("test_params.jl")
 
@@ -39,6 +40,24 @@ for binop in [safe_pow, ^]
     @test string_tree(minitree, opts) == "(x1 ^ x2)"
 end
 
+@testset "Test print_tree function" begin
+    if VERSION > v"1.8"
+        operators = OperatorEnum(;
+            binary_operators=(+, *, /, -), unary_operators=(cos, sin)
+        )
+        x1, x2, x3 = [Node(Float64; feature=i) for i in 1:3]
+        tree = x1 * x1 + 0.5
+        # Capture stdout to variable:
+        pipe = Pipe()
+        redirect_stdout(pipe) do
+            print_tree(tree, operators)
+        end
+        close(pipe.in)
+        s = read(pipe.out, String)
+        @test s == "((x1 * x1) + 0.5)\n"
+    end
+end
+
 @testset "Test printing of complex numbers" begin
     @eval my_custom_op(x, y) = x + y
     operators = OperatorEnum(;
@@ -55,3 +74,25 @@ end
     tree = my_custom_op(x1, 1.0 + 2.0im)
     @test string_tree(tree, operators) == "my_custom_op(x1, 1.0 + 2.0im)"
 end
+
+@testset "Test user-define printing" begin
+    operators = OperatorEnum(;
+        default_params..., binary_operators=(+, *, /, -), unary_operators=(cos, sin)
+    )
+    @extend_operators operators
+    x1, x2, x3 = [Node(Float64; feature=i) for i in 1:3]
+    tree = x1 * x1 + 0.5
+    @test string_tree(tree, operators; f_constant=Returns("TEST")) == "((x1 * x1) + TEST)"
+    @test string_tree(tree, operators; f_variable=Returns("TEST")) ==
+        "((TEST * TEST) + 0.5)"
+    @test string_tree(
+        tree, operators; f_variable=Returns("TEST"), f_constant=Returns("TEST2")
+    ) == "((TEST * TEST) + TEST2)"
+
+    # Try printing with a precision:
+    tree = x1 * x1 + π
+    f_constant(val::Float64, args...) = string(round(val; digits=2))
+    @test string_tree(tree, operators; f_constant=f_constant) == "((x1 * x1) + 3.14)"
+    f_constant(val::Float64, args...) = string(round(val; digits=4))
+    @test string_tree(tree, operators; f_constant=f_constant) == "((x1 * x1) + 3.1416)"
+end