Merge pull request #66 from SymbolicML/statictools

(BREAKING) Refactor code to be more idiomatic
SymbolicML · Feb 21, 2024 · 4bd83e9 · 4bd83e9 · MilesCranmer · Mar 9, 2024
2 parents 1f1ad6c + 7397885
commit 4bd83e9
Show file tree

Hide file tree

Showing 18 changed files with 477 additions and 363 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,15 +1,13 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <[email protected]>"]
-version = "0.15.0"
+version = "0.16.0"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
-LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
-Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
@@ -34,8 +32,8 @@ Bumper = "0.6"
 Compat = "3.37, 4"
 Enzyme = "^0.11.12"
 LoopVectorization = "0.12"
-Optim = "0.19, 1"
 MacroTools = "0.4, 0.5"
+Optim = "0.19, 1"
 PackageExtensionCompat = "1"
 PrecompileTools = "1"
 Reexport = "1"
@@ -48,6 +46,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -58,4 +57,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "SafeTestsets", "Aqua", "Bumper", "Enzyme", "ForwardDiff", "LoopVectorization", "Optim", "SpecialFunctions", "StaticArrays", "SymbolicUtils", "Zygote"]
+test = ["Test", "SafeTestsets", "Aqua", "Bumper", "Enzyme", "ForwardDiff", "LinearAlgebra", "LoopVectorization", "Optim", "SpecialFunctions", "StaticArrays", "SymbolicUtils", "Zygote"]
diff --git a/ext/DynamicExpressionsBumperExt.jl b/ext/DynamicExpressionsBumperExt.jl
@@ -23,7 +23,7 @@ function bumper_eval_tree_array(
             leaf_node -> begin
                 ar = @alloc(T, n)
                 ok = if leaf_node.constant
-                    v = leaf_node.val::T
+                    v = leaf_node.val
                     ar .= v
                     isfinite(v)
                 else

diff --git a/ext/DynamicExpressionsLoopVectorizationExt.jl b/ext/DynamicExpressionsLoopVectorizationExt.jl
@@ -41,8 +41,8 @@ function deg1_l2_ll0_lr0_eval(
     tree::AbstractExpressionNode{T}, cX::AbstractMatrix{T}, op::F, op_l::F2, ::Val{true}
 ) where {T<:Number,F,F2}
     if tree.l.l.constant && tree.l.r.constant
-        val_ll = tree.l.l.val::T
-        val_lr = tree.l.r.val::T
+        val_ll = tree.l.l.val
+        val_lr = tree.l.r.val
         @return_on_check val_ll cX
         @return_on_check val_lr cX
         x_l = op_l(val_ll, val_lr)::T
@@ -51,7 +51,7 @@ function deg1_l2_ll0_lr0_eval(
         @return_on_check x cX
         return ResultOk(fill_similar(x, cX, axes(cX, 2)), true)
     elseif tree.l.l.constant
-        val_ll = tree.l.l.val::T
+        val_ll = tree.l.l.val
         @return_on_check val_ll cX
         feature_lr = tree.l.r.feature
         cumulator = similar(cX, axes(cX, 2))
@@ -63,7 +63,7 @@ function deg1_l2_ll0_lr0_eval(
         return ResultOk(cumulator, true)
     elseif tree.l.r.constant
         feature_ll = tree.l.l.feature
-        val_lr = tree.l.r.val::T
+        val_lr = tree.l.r.val
         @return_on_check val_lr cX
         cumulator = similar(cX, axes(cX, 2))
         @turbo for j in axes(cX, 2)
@@ -89,7 +89,7 @@ function deg1_l1_ll0_eval(
     tree::AbstractExpressionNode{T}, cX::AbstractMatrix{T}, op::F, op_l::F2, ::Val{true}
 ) where {T<:Number,F,F2}
     if tree.l.l.constant
-        val_ll = tree.l.l.val::T
+        val_ll = tree.l.l.val
         @return_on_check val_ll cX
         x_l = op_l(val_ll)::T
         @return_on_check x_l cX
@@ -112,16 +112,16 @@ function deg2_l0_r0_eval(
     tree::AbstractExpressionNode{T}, cX::AbstractMatrix{T}, op::F, ::Val{true}
 ) where {T<:Number,F}
     if tree.l.constant && tree.r.constant
-        val_l = tree.l.val::T
+        val_l = tree.l.val
         @return_on_check val_l cX
-        val_r = tree.r.val::T
+        val_r = tree.r.val
         @return_on_check val_r cX
         x = op(val_l, val_r)::T
         @return_on_check x cX
         return ResultOk(fill_similar(x, cX, axes(cX, 2)), true)
     elseif tree.l.constant
         cumulator = similar(cX, axes(cX, 2))
-        val_l = tree.l.val::T
+        val_l = tree.l.val
         @return_on_check val_l cX
         feature_r = tree.r.feature
         @turbo for j in axes(cX, 2)
@@ -132,7 +132,7 @@ function deg2_l0_r0_eval(
     elseif tree.r.constant
         cumulator = similar(cX, axes(cX, 2))
         feature_l = tree.l.feature
-        val_r = tree.r.val::T
+        val_r = tree.r.val
         @return_on_check val_r cX
         @turbo for j in axes(cX, 2)
             x = op(cX[feature_l, j], val_r)
@@ -160,7 +160,7 @@ function deg2_l0_eval(
     ::Val{true},
 ) where {T<:Number,F}
     if tree.l.constant
-        val = tree.l.val::T
+        val = tree.l.val
         @return_on_check val cX
         @turbo for j in eachindex(cumulator)
             x = op(val, cumulator[j])
@@ -185,7 +185,7 @@ function deg2_r0_eval(
     ::Val{true},
 ) where {T<:Number,F}
     if tree.r.constant
-        val = tree.r.val::T
+        val = tree.r.val
         @return_on_check val cX
         @turbo for j in eachindex(cumulator)
             x = op(cumulator[j], val)

diff --git a/ext/DynamicExpressionsSymbolicUtilsExt.jl b/ext/DynamicExpressionsSymbolicUtilsExt.jl
@@ -34,7 +34,7 @@ function parse_tree_to_eqs(
 ) where {T}
     if tree.degree == 0
         # Return constant if needed
-        tree.constant && return subs_bad(tree.val::T)
+        tree.constant && return subs_bad(tree.val)
         return SymbolicUtils.Sym{LiteralReal}(Symbol("x$(tree.feature)"))
     end
     # Collect the next children
@@ -93,10 +93,10 @@ function split_eq(
     else
         ind = findoperation(op, operators.binops)
     end
-    return constructorof(N)(
-        ind,
-        convert(N, args[1], operators; variable_names=variable_names),
-        convert(N, op(args[2:end]...), operators; variable_names=variable_names),
+    return constructorof(N)(;
+        op=ind,
+        l=convert(N, args[1], operators; variable_names=variable_names),
+        r=convert(N, op(args[2:end]...), operators; variable_names=variable_names),
     )
 end
 
@@ -157,8 +157,9 @@ function Base.convert(
         findoperation(op, operators.unaops)
     end
 
-    return constructorof(N)(
-        ind, map(x -> convert(N, x, operators; variable_names=variable_names), args)...
+    return constructorof(N)(;
+        op=ind,
+        children=map(x -> convert(N, x, operators; variable_names=variable_names), args),
     )
 end
 

diff --git a/src/DynamicExpressions.jl b/src/DynamicExpressions.jl
@@ -5,6 +5,7 @@ include("ExtensionInterface.jl")
 include("OperatorEnum.jl")
 include("Equation.jl")
 include("EquationUtils.jl")
+include("Strings.jl")
 include("EvaluateEquation.jl")
 include("EvaluateEquationDerivative.jl")
 include("EvaluationHelpers.jl")
@@ -19,8 +20,6 @@ import Reexport: @reexport
     AbstractExpressionNode,
     GraphNode,
     Node,
-    string_tree,
-    print_tree,
     copy_node,
     set_node!,
     tree_mapreduce,
@@ -39,6 +38,7 @@ import .EquationModule: constructorof, preserve_sharing
     set_constants!,
     get_constant_refs,
     set_constant_refs!
+@reexport import .StringsModule: string_tree, print_tree
 @reexport import .OperatorEnumModule: AbstractOperatorEnum
 @reexport import .OperatorEnumConstructionModule:
     OperatorEnum, GenericOperatorEnum, @extend_operators, set_default_variable_names!