refactor: update to new derivative rule syntax

Author: AayushSabharwal

src/diff.jl

@@ -225,15 +225,15 @@ function Base.showerror(io::IO, err::DerivativeNotDefinedError)
     # and code fences
     err_str = Markdown.parse("""
         Derivative of `$(err.expr)` with respect to its $(err.i)-th argument is not defined.
-        Define a derivative by adding a method to `Symbolics.derivative`:
+        Define a derivative by using `@register_derivative`:
 
         ```julia
-        function Symbolics.derivative(::typeof($op), args::NTuple{$nargs, Any}, ::Val{$(err.i)})
+        @register_derivative $op(args...) $(err.i) begin
             # ...
         end
         ```
 
-        Refer to the documentation for `Symbolics.derivative` and the
+        Refer to the documentation for `@register_derivative` and the
         "[Adding Analytical Derivatives](@ref)" section of the docs for further information.
         """)
     show(io, MIME"text/plain"(), err_str)
@@ -321,12 +321,12 @@ function executediff(D::Differential, arg::BasicSymbolic{VartypeT}; simplify=fal
                 # We know `D.x` is in `arg`, so the derivative is not identically zero.
                 # `arg` cannot be `D.x` since, that would have also early exited. 
                 for (i, a) in enumerate(inner_args)
-                    der = derivative_idx(arr, i)
+                    der = derivative_idx(arr, i)::Union{Nothing, SymbolicT}
                     if isequal(a, D.x)
-                        der isa NoDeriv && return D(arg)
+                        der === nothing && return D(arg)
                         push!(summed_args, der[idx])
                         continue
-                    elseif der isa NoDeriv
+                    elseif der === nothing
                         push!(summed_args, Differential(a)(arg) * executediff(D, a))
                     else
                         push!(summed_args, der[idx] * executediff(D, a))
@@ -381,8 +381,8 @@ function executediff(D::Differential, arg::BasicSymbolic{VartypeT}; simplify=fal
                 for (i, iarg) in enumerate(inner_args)
                     t2 = executediff(D, iarg; simplify, throw_no_derivative)::SymbolicT
                     _iszero(t2) && continue
-                    t = derivative_idx(arg, i)::Union{NoDeriv, SymbolicT}
-                    if t isa NoDeriv
+                    t = derivative_idx(arg, i)::Union{Nothing, SymbolicT}
+                    if t === nothing
                         throw_no_derivative && throw(DerivativeNotDefinedError(arg, i))
                         t = D(arg)
                     end
@@ -506,89 +506,6 @@ function expand_derivatives(n::Complex{Num}, simplify=false; kwargs...)
 end
 expand_derivatives(x, simplify=false; kwargs...) = x
 
-# Don't specialize on the function here
-"""
-$(SIGNATURES)
-
-Calculate the derivative of the op `O` with respect to its argument with index
-`idx`.
-
-# Examples
-
-```jldoctest label1
-julia> using Symbolics
-
-julia> @variables x y;
-
-julia> Symbolics.derivative_idx(Symbolics.value(sin(x)), 1)
-cos(x)
-```
-
-Note that the function does not recurse into the operation's arguments, i.e., the
-chain rule is not applied:
-
-```jldoctest label1
-julia> myop = Symbolics.value(sin(x) * y^2)
-sin(x)*(y^2)
-
-julia> typeof(Symbolics.operation(myop))  # Op is multiplication function
-typeof(*)
-
-julia> Symbolics.derivative_idx(myop, 1)  # wrt. sin(x)
-y^2
-
-julia> Symbolics.derivative_idx(myop, 2)  # wrt. y^2
-sin(x)
-```
-"""
-derivative_idx(O::Any, ::Any) = COMMON_ZERO
-function derivative_idx(O::BasicSymbolic, idx)
-    iscall(O) || return COMMON_ZERO
-    res = derivative(operation(O), (arguments(O)...,), Val(idx))
-    if res isa NoDeriv
-        return res
-    else
-        return Const{VartypeT}(res)
-    end
-end
-
-# Indicate that no derivative is defined.
-struct NoDeriv
-end
-
-"""
-    Symbolics.derivative(::typeof(f), args::NTuple{N, Any}, ::Val{i})
-
-Return the derivative of `f(args...)` with respect to `args[i]`. `N` should be the number
-of arguments that `f` takes and `i` is the argument with respect to which the derivative
-is taken. The result can be a numeric value (if the derivative is constant) or a symbolic
-expression. This function is useful for defining derivatives of custom functions registered
-via `@register_symbolic`, to be used when calling `expand_derivatives`.
-"""
-derivative(f, args, v) = NoDeriv()
-
-# Pre-defined derivatives
-import DiffRules
-for (modu, fun, arity) ∈ DiffRules.diffrules(; filter_modules=(:Base, :SpecialFunctions, :NaNMath))
-    fun in [:*, :+, :abs, :mod, :rem, :max, :min] && continue # special
-    for i ∈ 1:arity
-
-        expr = if arity == 1
-            DiffRules.diffrule(modu, fun, :(args[1]))
-        else
-            DiffRules.diffrule(modu, fun, ntuple(k->:(args[$k]), arity)...)[i]
-        end
-        @eval derivative(::typeof($modu.$fun), args::NTuple{$arity,Any}, ::Val{$i}) = $expr
-    end
-end
-
-derivative(::typeof(+), args::NTuple{N,Any}, ::Val) where {N} = 1
-derivative(::typeof(*), args::NTuple{N,Any}, ::Val{i}) where {N,i} = *(deleteat!(collect(args), i)...)
-derivative(::typeof(one), args::Tuple{<:Any}, ::Val) = 0
-
-derivative(f::Function, x::Union{Num, <:BasicSymbolic}) = derivative(f(x), x)
-derivative(::Function, x::Any) = TypeError(:derivative, "2nd argument", Union{Num, <:BasicSymbolic}, x) |> throw
-
 function count_order(x)
     @assert !(x isa Symbol) "The variable $x must have an order of differentiation that is greater or equal to 1!"
     n = 1
@@ -666,6 +583,8 @@ function derivative(O, var; simplify=false, kwargs...)
         Num(expand_derivatives(Differential(var)(unwrap(O)), simplify; kwargs...))
     end
 end
+derivative(f::Function, var::Union{SymbolicT, Num}) = derivative(f(var), var)
+derivative(::Function, x::Any) = throw(TypeError(:derivative, "2nd argument", Union{Num, SymbolicT}, x))
 
 """
 $(SIGNATURES)

src/extra_functions.jl

@@ -7,31 +7,16 @@ for (T1, T2) in Iterators.product([Number, BasicSymbolic{VartypeT}, Num], [Integ
     end
 end
 
-derivative(::typeof(sign), args::NTuple{1,Any}, ::Val{1}) = 0
-
-derivative(::typeof(signbit), args::NTuple{1,Any}, ::Val{1}) = 0
-derivative(::typeof(abs), args::NTuple{1,Any}, ::Val{1}) = ifelse(signbit(args[1]),-one(args[1]),one(args[1]))
-
-function derivative(::typeof(min), args::NTuple{2,Any}, ::Val{1})
-    x, y = args
-    ifelse(x < y, one(x), zero(x))
-end
-function derivative(::typeof(min), args::NTuple{2,Any}, ::Val{2})
-    x, y = args
-    ifelse(x < y, zero(y), one(y))
-end
-function derivative(::typeof(max), args::NTuple{2,Any}, ::Val{1})
-    x, y = args
-    ifelse(x > y, one(x), zero(x))
-end
-function derivative(::typeof(max), args::NTuple{2,Any}, ::Val{2})
-    x, y = args
-    ifelse(x > y, zero(y), one(y))
-end
-
-function derivative(::Union{typeof(ceil),typeof(floor),typeof(factorial)}, args::NTuple{1,Any}, ::Val{1})
-    zero(args[1])
-end
+@register_derivative sign(x) 1 COMMON_ZERO
+@register_derivative signbit(x) 1 COMMON_ZERO
+@register_derivative abs(x) 1 ifelse(signbit(x),-one(x),one(x))
+@register_derivative min(x, y) 1 ifelse(x < y, one(x), zero(x))
+@register_derivative min(x, y) 2 ifelse(x < y, zero(y), one(y))
+@register_derivative max(x, y) 1 ifelse(x > y, one(x), zero(x))
+@register_derivative max(x, y) 2 ifelse(x > y, zero(y), one(y))
+@register_derivative ceil(x) 1 COMMON_ZERO
+@register_derivative floor(x) 1 COMMON_ZERO
+@register_derivative factorial(x) 1 COMMON_ZERO
 
 @register_symbolic Base.rand(x)
 @register_symbolic Base.randn(x)
@@ -45,13 +30,7 @@ for (T1, T2, T3) in Iterators.product(Iterators.repeated((Num, BasicSymbolic{Var
     end
 end
 
-function derivative(::typeof(Base.clamp), args::NTuple{3, Any}, ::Val{1})
-    x, l, h = args
-    T = promote_type(symtype(x), symtype(l), symtype(h))
-    z = zero(T)
-    o = one(T)
-    ifelse(x<l, z, ifelse(x>h, z, o))
-end
+@register_derivative clamp(x, l, h) 1 ifelse(x < l, COMMON_ZERO, ifelse(x > h, COMMON_ZERO, COMMON_ONE))
 
 for T1 in [Real, Num, BasicSymbolic{VartypeT}], T2 in [AbstractArray, Arr, BasicSymbolic{VartypeT}]
     if T1 != Num && T2 != Arr
@@ -78,11 +57,10 @@ end
 
 LinearAlgebra.norm(x::Num, p::Real) = abs(x)
 
-derivative(::typeof(<), ::NTuple{2, Any}, ::Val{i}) where {i} = 0
-derivative(::typeof(<=), ::NTuple{2, Any}, ::Val{i}) where {i} = 0
-derivative(::typeof(>), ::NTuple{2, Any}, ::Val{i}) where {i} = 0
-derivative(::typeof(>=), ::NTuple{2, Any}, ::Val{i}) where {i} = 0
-derivative(::typeof(==), ::NTuple{2, Any}, ::Val{i}) where {i} = 0
-derivative(::typeof(!=), ::NTuple{2, Any}, ::Val{i}) where {i} = 0
-
-derivative(::typeof(expinti), args::NTuple{1,Any}, ::Val{1}) = exp(args[1])/args[1]
+@register_derivative <(x, y) I COMMON_ZERO
+@register_derivative <=(x, y) I COMMON_ZERO
+@register_derivative >(x, y) I COMMON_ZERO
+@register_derivative >=(x, y) I COMMON_ZERO
+@register_derivative ==(x, y) I COMMON_ZERO
+@register_derivative !=(x, y) I COMMON_ZERO
+@register_derivative expinti(x) 1 exp(x) / x

src/register_derivatives.jl

@@ -123,6 +123,9 @@ macro register_derivative(f::Expr, I::Union{Symbol, Int}, body)
     return esc(Expr(:function, derhead, Expr(:block, unpack, body)))
 end
 
+# Fallback
+@register_derivative (f::Any)(args...) I nothing
+
 # Pre-defined derivatives
 import DiffRules
 for (modu, fun, arity) ∈ DiffRules.diffrules(; filter_modules=(:Base, :SpecialFunctions, :NaNMath))
@@ -136,7 +139,7 @@ for (modu, fun, arity) ∈ DiffRules.diffrules(; filter_modules=(:Base, :Special
         end
 
         # Using the macro here doesn't work somehow.
-        @eval function derivative_rule(::typeof($modu.$fun), ::Val{$arity}, args::SymbolicUtils.ArgsT{VartypeT}, ::Val{$i})
+        @eval function derivative_rule(::typeof($modu.$fun), ::Val{$arity}, args::ROArgsT{VartypeT}, ::Val{$i})
             $SConst($expr)
         end
     end
@@ -285,7 +288,7 @@ sin(x)
 ```
 """
 @inline derivative_idx(::Any, ::Any) = COMMON_ZERO
-function derivative_idx(O::VartypeT, idx::Int)
+function derivative_idx(O::SymbolicT, idx::Int)
     iscall(O) || return COMMON_ZERO
     f = operation(O)
     args = arguments(O)

src/solver/solve_helpers.jl

@@ -41,7 +41,9 @@ end
 SymbolicUtils.promote_symtype(::typeof(ssqrt), ::Type{T}) where {T} = T
 SymbolicUtils.promote_shape(::typeof(ssqrt), @nospecialize(sh::SymbolicUtils.ShapeT)) = sh
 
-derivative(::typeof(ssqrt), args...) = substitute(derivative(sqrt, args...), sqrt => ssqrt)
+@register_derivative ssqrt(x) I begin
+    substitute(@derivative_rule(sqrt(x), I), sqrt => ssqrt)
+end
 
 function scbrt(n)
     n = unwrap(n)
@@ -62,7 +64,9 @@ end
 
 SymbolicUtils.promote_symtype(::typeof(scbrt), ::Type{T}) where {T} = T
 SymbolicUtils.promote_shape(::typeof(scbrt), @nospecialize(sh::SymbolicUtils.ShapeT)) = sh
-derivative(::typeof(scbrt), args...) = substitute(derivative(cbrt, args...), cbrt => scbrt)
+@register_derivative scbrt(x) I begin
+    substitute(@derivative_rule(cbrt(x), I), cbrt => scbrt)
+end
 
 function slog(n)
     n = unwrap(n)
@@ -82,7 +86,9 @@ end
 SymbolicUtils.promote_symtype(::typeof(slog), ::Type{T}) where {T} = T
 SymbolicUtils.promote_shape(::typeof(slog), @nospecialize(sh::SymbolicUtils.ShapeT)) = sh
 
-derivative(::typeof(slog), args...) = substitute(derivative(log, args...), log => slog)
+@register_derivative slog(x) I begin
+    substitute(@derivative_rule(log(x), I), log => slog)
+end
 
 const RootsOf = (SymbolicUtils.@syms roots_of(poly,var))[1]
 

test/diff.jl

@@ -630,25 +630,25 @@ end
     test_equal.(Symbolics.unwrap.(Symbolics.jacobian([f], vp) .- Symbolics.jacobian([f], p)), 0)
 end
 
+@register_array_symbolic foobar(x, y) begin
+    size = (2,)
+    eltype = Real
+    ndims = 1
+end
+@register_array_symbolic dfoobar1(x, y) begin
+    size = (2,)
+    eltype = Real
+    ndims = 1
+end
+@register_array_symbolic dfoobar2(x, y) begin
+    size = (2,)
+    eltype = Real
+    ndims = 1
+end
+@register_derivative foobar(x, y) 1 dfoobar1(x, y)
+@register_derivative foobar(x, y) 2 dfoobar2(x, y)
+
 @testset "Derivatives of indexed array expressions" begin
-    @register_array_symbolic foobar(x, y) begin
-        size = (2,)
-        eltype = Real
-        ndims = 1
-    end
-    @register_array_symbolic dfoobar1(x, y) begin
-        size = (2,)
-        eltype = Real
-        ndims = 1
-    end
-    @register_array_symbolic dfoobar2(x, y) begin
-        size = (2,)
-        eltype = Real
-        ndims = 1
-    end
-    Symbolics.derivative(::typeof(foobar), args::NTuple{2, Any}, ::Val{1}) = dfoobar1(args...)
-    Symbolics.derivative(::typeof(foobar), args::NTuple{2, Any}, ::Val{2}) = dfoobar2(args...)
-    
     @variables x y
     ex = foobar(x + 2y, y)
     @test ex isa Symbolics.Arr