Merge pull request #16 from SymbolicML/complex-numbers

Open OperatorEnum to 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.2"
+version = "0.4.3"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

docs/src/eval.md

@@ -4,7 +4,7 @@ Given an expression tree specified with a `Node` type, you may evaluate the expr
 over an array of data with the following command:
 
 ```@docs
-eval_tree_array(tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum) where {T<:Real}
+eval_tree_array(tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum) where {T<:Number}
 ```
 
 Assuming you are only using a single `OperatorEnum`, you can also use
@@ -45,15 +45,15 @@ all variables (or, all constants). Both use forward-mode automatic, but use
 `Zygote.jl` to compute derivatives of each operator, so this is very efficient.
 
 ```@docs
-eval_diff_tree_array(tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum, direction::Int) where {T<:Real}
-eval_grad_tree_array(tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum; variable::Bool=false) where {T<:Real}
+eval_diff_tree_array(tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum, direction::Int) where {T<:Number}
+eval_grad_tree_array(tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum; variable::Bool=false) where {T<:Number}
 ```
 
 Alternatively, you can compute higher-order derivatives by using `ForwardDiff` on
 the function `differentiable_eval_tree_array`, although this will be slower.
 
 ```@docs
-differentiable_eval_tree_array(tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum) where {T<:Real}
+differentiable_eval_tree_array(tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum) where {T<:Number}
 ```
 
 ## Printing

docs/src/types.md

@@ -16,7 +16,7 @@ Construct this operator specification as follows:
 OperatorEnum(; binary_operators=[], unary_operators=[], enable_autodiff::Bool=false, define_helper_functions::Bool=true)
 ```
 
-This is just for scalar real operators. However, you can use
+This is just for scalar operators. However, you can use
 the following for more general operators:
 
 ```@docs

src/EvaluateEquation.jl

@@ -63,7 +63,7 @@ which speed up evaluation significantly.
 """
 function eval_tree_array(
     tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum; turbo::Bool=false
-)::Tuple{AbstractVector{T},Bool} where {T<:Real}
+)::Tuple{AbstractVector{T},Bool} where {T<:Number}
     n = size(cX, 2)
     if turbo
         @assert T in (Float32, Float64)
@@ -77,7 +77,7 @@ function eval_tree_array(
 end
 function eval_tree_array(
     tree::Node{T1}, cX::AbstractMatrix{T2}, operators::OperatorEnum; turbo::Bool=false
-) where {T1<:Real,T2<:Real}
+) where {T1<:Number,T2<:Number}
     T = promote_type(T1, T2)
     @warn "Warning: eval_tree_array received mixed types: tree=$(T1) and data=$(T2)."
     tree = convert(Node{T}, tree)
@@ -87,7 +87,7 @@ end
 
 function _eval_tree_array(
     tree::Node{T}, cX::AbstractMatrix{T}, operators::OperatorEnum, ::Val{turbo}
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Number,turbo}
     n = size(cX, 2)
     # First, we see if there are only constants in the tree - meaning
     # we can just return the constant result.
@@ -148,7 +148,7 @@ end
 
 function deg2_eval(
     cumulator_l::AbstractVector{T}, cumulator_r::AbstractVector{T}, op::F, ::Val{turbo}
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Number,F,turbo}
     @maybe_turbo turbo for j in indices(cumulator_l)
         x = op(cumulator_l[j], cumulator_r[j])::T
         cumulator_l[j] = x
@@ -158,7 +158,7 @@ end
 
 function deg1_eval(
     cumulator::AbstractVector{T}, op::F, ::Val{turbo}
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Number,F,turbo}
     @maybe_turbo turbo for j in indices(cumulator)
         x = op(cumulator[j])::T
         cumulator[j] = x
@@ -168,7 +168,7 @@ end
 
 function deg0_eval(
     tree::Node{T}, cX::AbstractMatrix{T}
-)::Tuple{AbstractVector{T},Bool} where {T<:Real}
+)::Tuple{AbstractVector{T},Bool} where {T<:Number}
     if tree.constant
         n = size(cX, 2)
         return (fill(tree.val::T, n), true)
@@ -179,7 +179,7 @@ end
 
 function deg1_l2_ll0_lr0_eval(
     tree::Node{T}, cX::AbstractMatrix{T}, op::F, op_l::F2, ::Val{turbo}
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,F2,turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Number,F,F2,turbo}
     n = size(cX, 2)
     if tree.l.l.constant && tree.l.r.constant
         val_ll = tree.l.l.val::T
@@ -229,7 +229,7 @@ end
 # op(op2(x)) for x variable or constant
 function deg1_l1_ll0_eval(
     tree::Node{T}, cX::AbstractMatrix{T}, op::F, op_l::F2, ::Val{turbo}
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,F2,turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Number,F,F2,turbo}
     n = size(cX, 2)
     if tree.l.l.constant
         val_ll = tree.l.l.val::T
@@ -254,7 +254,7 @@ end
 # op(x, y) for x and y variable/constant
 function deg2_l0_r0_eval(
     tree::Node{T}, cX::AbstractMatrix{T}, op::F, ::Val{turbo}
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Number,F,turbo}
     n = size(cX, 2)
     if tree.l.constant && tree.r.constant
         val_l = tree.l.val::T
@@ -297,7 +297,7 @@ end
 # op(x, y) for x variable/constant, y arbitrary
 function deg2_l0_eval(
     tree::Node{T}, cumulator::AbstractVector{T}, cX::AbstractArray{T}, op::F, ::Val{turbo}
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Number,F,turbo}
     n = size(cX, 2)
     if tree.l.constant
         val = tree.l.val::T
@@ -319,7 +319,7 @@ end
 # op(x, y) for x arbitrary, y variable/constant
 function deg2_r0_eval(
     tree::Node{T}, cumulator::AbstractVector{T}, cX::AbstractArray{T}, op::F, ::Val{turbo}
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,turbo}
+)::Tuple{AbstractVector{T},Bool} where {T<:Number,F,turbo}
     n = size(cX, 2)
     if tree.r.constant
         val = tree.r.val::T
@@ -339,15 +339,15 @@ function deg2_r0_eval(
 end
 
 """
-    _eval_constant_tree(tree::Node{T}, operators::OperatorEnum)::Tuple{T,Bool} where {T<:Real}
+    _eval_constant_tree(tree::Node{T}, operators::OperatorEnum)::Tuple{T,Bool} where {T<:Number}
 
 Evaluate a tree which is assumed to not contain any variable nodes. This
 gives better performance, as we do not need to perform computation
 over an entire array when the values are all the same.
 """
 function _eval_constant_tree(
     tree::Node{T}, operators::OperatorEnum
-)::Tuple{T,Bool} where {T<:Real}
+)::Tuple{T,Bool} where {T<:Number}
     if tree.degree == 0
         return deg0_eval_constant(tree)
     elseif tree.degree == 1
@@ -357,13 +357,13 @@ function _eval_constant_tree(
     end
 end
 
-@inline function deg0_eval_constant(tree::Node{T})::Tuple{T,Bool} where {T<:Real}
+@inline function deg0_eval_constant(tree::Node{T})::Tuple{T,Bool} where {T<:Number}
     return tree.val::T, true
 end
 
 function deg1_eval_constant(
     tree::Node{T}, op::F, operators::OperatorEnum
-)::Tuple{T,Bool} where {T<:Real,F}
+)::Tuple{T,Bool} where {T<:Number,F}
     (cumulator, complete) = _eval_constant_tree(tree.l, operators)
     !complete && return zero(T), false
     output = op(cumulator)::T
@@ -372,7 +372,7 @@ end
 
 function deg2_eval_constant(
     tree::Node{T}, op::F, operators::OperatorEnum
-)::Tuple{T,Bool} where {T<:Real,F}
+)::Tuple{T,Bool} where {T<:Number,F}
     (cumulator, complete) = _eval_constant_tree(tree.l, operators)
     !complete && return zero(T), false
     (cumulator2, complete2) = _eval_constant_tree(tree.r, operators)
@@ -388,7 +388,7 @@ Evaluate an expression tree in a way that can be auto-differentiated.
 """
 function differentiable_eval_tree_array(
     tree::Node{T1}, cX::AbstractMatrix{T}, operators::OperatorEnum
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,T1}
+)::Tuple{AbstractVector{T},Bool} where {T<:Number,T1}
     n = size(cX, 2)
     if tree.degree == 0
         if tree.constant
@@ -405,7 +405,7 @@ end
 
 function deg1_diff_eval(
     tree::Node{T1}, cX::AbstractMatrix{T}, op::F, operators::OperatorEnum
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,T1}
+)::Tuple{AbstractVector{T},Bool} where {T<:Number,F,T1}
     (left, complete) = differentiable_eval_tree_array(tree.l, cX, operators)
     @return_on_false complete left
     out = op.(left)
@@ -415,7 +415,7 @@ end
 
 function deg2_diff_eval(
     tree::Node{T1}, cX::AbstractMatrix{T}, op::F, operators::OperatorEnum
-)::Tuple{AbstractVector{T},Bool} where {T<:Real,F,T1}
+)::Tuple{AbstractVector{T},Bool} where {T<:Number,F,T1}
     (left, complete) = differentiable_eval_tree_array(tree.l, cX, operators)
     @return_on_false complete left
     (right, complete2) = differentiable_eval_tree_array(tree.r, cX, operators)

src/EvaluateEquationDerivative.jl

@@ -43,7 +43,7 @@ function eval_diff_tree_array(
     operators::OperatorEnum,
     direction::Int;
     turbo::Bool=false,
-)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Real}
+)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Number}
     assert_autodiff_enabled(operators)
     # TODO: Implement quick check for whether the variable is actually used
     # in this tree. Otherwise, return zero.
@@ -57,7 +57,7 @@ function eval_diff_tree_array(
     operators::OperatorEnum,
     direction::Int;
     turbo::Bool=false,
-) where {T1<:Real,T2<:Real}
+) where {T1<:Number,T2<:Number}
     T = promote_type(T1, T2)
     @warn "Warning: eval_diff_tree_array received mixed types: tree=$(T1) and data=$(T2)."
     tree = convert(Node{T}, tree)
@@ -71,7 +71,7 @@ function _eval_diff_tree_array(
     operators::OperatorEnum,
     direction::Int,
     ::Val{turbo},
-)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Real,turbo}
+)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Number,turbo}
     evaluation, derivative, complete = if tree.degree == 0
         diff_deg0_eval(tree, cX, direction)
     elseif tree.degree == 1
@@ -101,7 +101,7 @@ end
 
 function diff_deg0_eval(
     tree::Node{T}, cX::AbstractMatrix{T}, direction::Int
-)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Real}
+)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Number}
     n = size(cX, 2)
     const_part = deg0_eval(tree, cX)[1]
     derivative_part =
@@ -117,7 +117,7 @@ function diff_deg1_eval(
     operators::OperatorEnum,
     direction::Int,
     ::Val{turbo},
-)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Real,F,dF,turbo}
+)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Number,F,dF,turbo}
     n = size(cX, 2)
     (cumulator, dcumulator, complete) = _eval_diff_tree_array(
         tree.l, cX, operators, direction, Val(turbo)
@@ -143,7 +143,7 @@ function diff_deg2_eval(
     operators::OperatorEnum,
     direction::Int,
     ::Val{turbo},
-)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Real,F,dF,turbo}
+)::Tuple{AbstractVector{T},AbstractVector{T},Bool} where {T<:Number,F,dF,turbo}
     (cumulator, dcumulator, complete) = _eval_diff_tree_array(
         tree.l, cX, operators, direction, Val(turbo)
     )
@@ -194,7 +194,7 @@ function eval_grad_tree_array(
     operators::OperatorEnum;
     variable::Bool=false,
     turbo::Bool=false,
-)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real}
+)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Number}
     assert_autodiff_enabled(operators)
     n = size(cX, 2)
     if variable
@@ -224,7 +224,7 @@ function eval_grad_tree_array(
     operators::OperatorEnum,
     ::Val{variable},
     ::Val{turbo},
-)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real,variable,turbo}
+)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Number,variable,turbo}
     evaluation, gradient, complete = _eval_grad_tree_array(
         tree, n, n_gradients, index_tree, cX, operators, Val(variable), Val(turbo)
     )
@@ -238,7 +238,7 @@ function eval_grad_tree_array(
     operators::OperatorEnum;
     variable::Bool=false,
     turbo::Bool=false,
-) where {T1<:Real,T2<:Real}
+) where {T1<:Number,T2<:Number}
     T = promote_type(T1, T2)
     return eval_grad_tree_array(
         convert(Node{T}, tree),
@@ -258,7 +258,7 @@ function _eval_grad_tree_array(
     operators::OperatorEnum,
     ::Val{variable},
     ::Val{turbo},
-)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real,variable,turbo}
+)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Number,variable,turbo}
     if tree.degree == 0
         grad_deg0_eval(tree, n, n_gradients, index_tree, cX, Val(variable))
     elseif tree.degree == 1
@@ -297,7 +297,7 @@ function grad_deg0_eval(
     index_tree::NodeIndex,
     cX::AbstractMatrix{T},
     ::Val{variable},
-)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real,variable}
+)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Number,variable}
     const_part = deg0_eval(tree, cX)[1]
 
     if variable == tree.constant
@@ -321,7 +321,7 @@ function grad_deg1_eval(
     operators::OperatorEnum,
     ::Val{variable},
     ::Val{turbo},
-)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real,F,dF,variable,turbo}
+)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Number,F,dF,variable,turbo}
     (cumulator, dcumulator, complete) = eval_grad_tree_array(
         tree.l, n, n_gradients, index_tree.l, cX, operators, Val(variable), Val(turbo)
     )
@@ -350,7 +350,7 @@ function grad_deg2_eval(
     operators::OperatorEnum,
     ::Val{variable},
     ::Val{turbo},
-)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Real,F,dF,variable,turbo}
+)::Tuple{AbstractVector{T},AbstractMatrix{T},Bool} where {T<:Number,F,dF,variable,turbo}
     (cumulator1, dcumulator1, complete) = eval_grad_tree_array(
         tree.l, n, n_gradients, index_tree.l, cX, operators, Val(variable), Val(turbo)
     )

src/OperatorEnum.jl

@@ -7,8 +7,8 @@ abstract type AbstractOperatorEnum end
 
 Defines an enum over operators, along with their derivatives.
 # Fields
-- `binops`: A tuple of binary operators. Real scalar input type.
-- `unaops`: A tuple of unary operators. Real scalar input type.
+- `binops`: A tuple of binary operators. Scalar input type.
+- `unaops`: A tuple of unary operators. Scalar input type.
 - `diff_binops`: A tuple of Zygote-computed derivatives of the binary operators.
 - `diff_unaops`: A tuple of Zygote-computed derivatives of the unary operators.
 """

src/OperatorEnumConstruction.jl

@@ -126,7 +126,7 @@ function _extend_operators(operators, skip_user_operators, __module__::Module)
         local type_requirements
         local build_converters
         if isa($operators, OperatorEnum)
-            type_requirements = Real
+            type_requirements = Number
             build_converters = true
         else
             type_requirements = Any
@@ -257,9 +257,9 @@ and `(::Node)(X)`.
 
 # Arguments
 - `binary_operators::Vector{Function}`: A vector of functions, each of which is a binary
-  operator on real scalars.
+  operator.
 - `unary_operators::Vector{Function}`: A vector of functions, each of which is a unary
-  operator on real scalars.
+  operator.
 - `define_helper_functions::Bool=true`: Whether to define helper functions for creating
    and evaluating node types. Turn this off when doing precompilation. Note that these
    are *not* needed for the package to work; they are purely for convenience.

src/SimplifyEquation.jl

@@ -9,7 +9,7 @@ function combine_operators(
     tree::Node{T},
     operators::AbstractOperatorEnum,
     id_map::IdDict{Node{T},Node{T}}=IdDict{Node{T},Node{T}}(),
-)::Node{T} where {T<:Real}
+)::Node{T} where {T}
     # NOTE: (const (+*-) const) already accounted for. Call simplify_tree before.
     # ((const + var) + const) => (const + var)
     # ((const * var) * const) => (const * var)