Merge pull request #107 from SymbolicML/fix-safe-sqrt

feat: allow expression algebra for safe aliases

Author: MilesCranmer

.github/workflows/CI.yml

@@ -62,7 +62,7 @@ jobs:
   additional_tests:
     name: test ${{ matrix.test_name }} - ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
-    timeout-minutes: 60
+    timeout-minutes: 120
     strategy:
       fail-fast: false
       matrix:
@@ -71,7 +71,7 @@ jobs:
         julia-version:
           - "1"
         test_name:
-          - "enzyme"
+          # - "enzyme"  # flaky; seems to infinitely compile and fail the CI
           - "jet"
     steps:
       - uses: actions/checkout@v2

Project.toml

@@ -1,7 +1,7 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <[email protected]>"]
-version = "1.3.0"
+version = "1.4.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/DynamicExpressions.jl

@@ -90,6 +90,7 @@ import .StringsModule: get_op_name
 import .ExpressionModule:
     get_operators, get_variable_names, Metadata, default_node_type, node_type
 @reexport import .ExpressionAlgebraModule: @declare_expression_operator
+import .ExpressionAlgebraModule: declare_operator_alias
 @reexport import .ParseModule: @parse_expression, parse_expression
 import .ParseModule: parse_leaf
 @reexport import .ParametricExpressionModule: ParametricExpression, ParametricNode

src/ExpressionAlgebra.jl

@@ -34,9 +34,31 @@ function Base.showerror(io::IO, e::MissingOperatorError)
     return print(io, e.msg)
 end
 
+"""
+    declare_operator_alias(op::Function, ::Val{arity})::Function
+
+Define how an internal operator should be matched against user-provided operators in expression trees.
+
+By default, operators match themselves. Override this method to specify that an internal operator
+should match a different operator when searching the operator lists in expressions.
+
+For example, to make `safe_sqrt` match `sqrt` user-space:
+
+```julia
+DynamicExpressions.declare_operator_alias(safe_sqrt, Val(1)) = sqrt
+```
+
+Which would allow a user to write `sqrt(x::Expression)`
+and have it match the operator `safe_sqrt` stored in the binary operators
+of the expression.
+"""
+declare_operator_alias(op::F, _) where {F<:Function} = op
+
 function apply_operator(op::F, l::AbstractExpression) where {F<:Function}
     operators = get_operators(l, nothing)
-    op_idx = findfirst(==(op), operators.unaops)
+    op_idx = findfirst(
+        ==(op), map(Base.Fix2(declare_operator_alias, Val(1)), operators.unaops)
+    )
     if op_idx === nothing
         throw(
             MissingOperatorError(
@@ -56,7 +78,9 @@ function apply_operator(op::F, l, r) where {F<:Function}
         r::AbstractExpression
         (get_operators(r, nothing), r)
     end
-    op_idx = findfirst(==(op), operators.binops)
+    op_idx = findfirst(
+        ==(op), map(Base.Fix2(declare_operator_alias, Val(2)), operators.binops)
+    )
     if op_idx === nothing
         throw(
             MissingOperatorError(

test/test_expression_math.jl

@@ -145,3 +145,40 @@ end
         )
     end
 end
+@testitem "Custom operators and aliases" begin
+    using DynamicExpressions
+
+    # Define a custom safe sqrt that avoids negative numbers
+    safe_sqrt(x) = x < 0 ? zero(x) : sqrt(x)
+    # And a custom function that squares its input
+    my_func(x) = x^2
+
+    # Define that safe_sqrt should match sqrt in expressions, with correct type!
+    DynamicExpressions.declare_operator_alias(::typeof(safe_sqrt), ::Val{1}) = sqrt
+
+    # Declare my_func as a new operator
+    @declare_expression_operator my_func 1
+
+    # Create an expression with just safe_sqrt:
+    ex = parse_expression(
+        :(x);
+        expression_type=Expression{Float64},
+        unary_operators=[safe_sqrt, my_func],
+        variable_names=["x"],
+    )
+
+    # Test that sqrt(ex) maps to safe_sqrt through the alias:
+    ex_sqrt = sqrt(ex)
+    ex_my = my_func(ex)
+
+    shower(ex) = sprint((io, e) -> show(io, MIME"text/plain"(), e), ex)
+
+    @test shower(ex_sqrt) == "safe_sqrt(x)"
+    @test shower(ex_my) == "my_func(x)"
+
+    # Test evaluation:
+    X = [4.0 -4.0]
+
+    @test ex_sqrt(X) ≈ [2.0; 0.0]
+    @test ex_my(X) ≈ [16.0; 16.0]
+end