Merge pull request #19 from SymbolicML/fix-complex-printing

Fix printing of complex numbers

Author: MilesCranmer

Project.toml

@@ -1,7 +1,7 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <miles.cranmer@gmail.com>"]
-version = "0.4.3"
+version = "0.5.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/Equation.jl

@@ -344,7 +344,7 @@ function string_tree(
 )::String where {T}
     if tree.degree == 0
         if tree.constant
-            return string(tree.val::T)
+            return string_constant(tree.val::T; bracketed=bracketed)
         else
             if varMap === nothing
                 return "x$(tree.feature)"
@@ -362,6 +362,15 @@ function string_tree(
     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,

test/test_print.jl

@@ -38,3 +38,20 @@ for binop in [safe_pow, ^]
     minitree = Node(5, Node("x1"), Node("x2"))
     @test string_tree(minitree, opts) == "(x1 ^ x2)"
 end
+
+@testset "Test printing of complex numbers" begin
+    @eval my_custom_op(x, y) = x + y
+    operators = OperatorEnum(;
+        default_params...,
+        binary_operators=(+, *, /, -, my_custom_op),
+        unary_operators=(cos, sin),
+    )
+    @extend_operators operators
+    x1, x2, x3 = [Node(; feature=i) for i in 1:3]
+    tree = sin(x1 * 1.0)
+    @test string_tree(tree, operators) == "sin(x1 * 1.0)"
+    tree = sin(x1 * (1.0 + 2.0im))
+    @test string_tree(tree, operators) == "sin(x1 * (1.0 + 2.0im))"
+    tree = my_custom_op(x1, 1.0 + 2.0im)
+    @test string_tree(tree, operators) == "my_custom_op(x1, 1.0 + 2.0im)"
+end