Support ForwardDiff duals over BareInterval

ext/IntervalArithmeticDiffRulesExt.jl

@@ -3,6 +3,9 @@ module IntervalArithmeticDiffRulesExt
 using IntervalArithmetic
 import DiffRules
 
+DiffRules._abs_deriv(x::BareInterval{T}) where {T<:IntervalArithmetic.NumTypes} =
+    ifelse(isthinzero(x), bareinterval(-one(T), one(T)), sign(x))
+
 function DiffRules._abs_deriv(x::Interval{T}) where {T<:IntervalArithmetic.NumTypes}
     r = ifelse(isthinzero(x), bareinterval(-one(T), one(T)), sign(bareinterval(x)))
     d = decoration(x)

ext/IntervalArithmeticForwardDiffExt.jl

@@ -2,7 +2,7 @@ module IntervalArithmeticForwardDiffExt
 
 using IntervalArithmetic, ForwardDiff
 using IntervalArithmetic: isempty_domain, overlap_domain, intersect_domain, in_domain, leftof
-using ForwardDiff: Dual, Partials, ≺, value, partials
+using ForwardDiff: Dual, Partials, ≺, value, partials, DiffRules
 
 ForwardDiff.can_dual(::Type{ExactReal}) = true
 
@@ -39,6 +39,7 @@ const NestedInterval = let
     U
 end
 _isthinzero(x::Interval) = isthinzero(x)
+_isthinzero(x::BareInterval) = isthinzero(x)
 _isthinzero(x::Real) = iszero(x)
 _isthinzero(d::Dual) = _isthinzero(value(d)) && all(_isthinzero, partials(d))
 
@@ -172,4 +173,162 @@ end
 ForwardDiff.DiffRules._abs_deriv(x::Dual{T,<:Interval}) where {T} =
     Dual{T}(ForwardDiff.DiffRules._abs_deriv(value(x)), zero(partials(x)))
 
+# ForwardDiff support for `BareInterval`, which is not a subtype of `Number`.
+
+ForwardDiff.can_dual(::Type{<:BareInterval}) = true
+
+function ForwardDiff.derivative(f::F, x::BareInterval) where {F}
+    T = typeof(ForwardDiff.Tag(f, typeof(x)))
+    return ForwardDiff.extract_derivative(T, f(Dual{T}(x, Partials((one(x),)))))
+end
+
+# Bare intervals nested in up to four `Dual` layers.
+const NestedBareInterval = let
+    U = BareInterval
+    for _ in 1:4
+        U = Union{U, Dual{T,<:U} where {T}}
+    end
+    U
+end
+
+# numtype of the bare interval at the center of a nest of `Dual` layers
+_numtype_bare(x::BareInterval) = numtype(x)
+_numtype_bare(d::Dual) = _numtype_bare(value(d))
+
+Base.:+(d::Dual{T,<:NestedBareInterval}, x::BareInterval) where {T} = Dual{T}(value(d) + x, partials(d))
+Base.:+(x::BareInterval, d::Dual{<:Any,<:NestedBareInterval}) = d + x
+Base.:-(d::Dual{T,<:NestedBareInterval}, x::BareInterval) where {T} = Dual{T}(value(d) - x, partials(d))
+Base.:-(x::BareInterval, d::Dual{T,<:NestedBareInterval}) where {T} = Dual{T}(x - value(d), -partials(d))
+Base.:*(d::Dual{T,<:NestedBareInterval}, x::BareInterval) where {T} = Dual{T}(value(d) * x, partials(d) * x)
+Base.:*(x::BareInterval, d::Dual{<:Any,<:NestedBareInterval}) = d * x
+Base.:/(d::Dual{T,<:NestedBareInterval}, x::BareInterval) where {T} = Dual{T}(value(d) / x, partials(d) / x)
+Base.:/(x::BareInterval, d::Dual{<:Any,<:NestedBareInterval}) = x * inv(d)
+
+# Extend ForwardDiff's partial scaling to bare intervals.
+Base.:*(p::Partials, x::BareInterval) = Partials(ForwardDiff.scale_tuple(p.values, x))
+Base.:*(x::BareInterval, p::Partials) = p * x
+Base.:/(p::Partials, x::BareInterval) = Partials(ForwardDiff.div_tuple_by_scalar(p.values, x))
+ForwardDiff._mul_partial(p::BareInterval, x::BareInterval) = p * x
+ForwardDiff._mul_partial(p::BareInterval, x::ExactReal) = p * x
+ForwardDiff._mul_partial(p::Dual, x::BareInterval) = p * x
+ForwardDiff._div_partial(p::BareInterval, x::BareInterval) = p / x
+ForwardDiff._div_partial(p::BareInterval, x::ExactReal) = p / x
+ForwardDiff._div_partial(p::Dual, x::BareInterval) = p / x
+
+# Derivative rules may return either bare intervals or real coefficients.
+ForwardDiff.dual_definition_retval(::Val{T}, val::BareInterval, deriv, partial::Partials) where {T} =
+    Dual{T}(val, deriv * partial)
+ForwardDiff.dual_definition_retval(::Val{T}, val::BareInterval, deriv1, partial1::Partials,
+                                   deriv2, partial2::Partials) where {T} =
+    Dual{T}(val, ForwardDiff._mul_partials(partial1, partial2, deriv1, deriv2))
+
+# Use general integer powers to scale partials via `_mul_partial`.
+# Match ForwardDiff's `Val` specializations to avoid ambiguities, and use a
+# variable exponent to avoid recursing through `literal_pow`.
+for p in 0:3
+    @eval Base.literal_pow(::typeof(^), x::Dual{T,<:NestedBareInterval}, ::Val{$p}) where {T} =
+        (n = $p; x^n)
+end
+Base.literal_pow(::typeof(^), x::Dual{T,<:NestedBareInterval}, ::Val{p}) where {T,p} = (n = p; x^n)
+
+# Integer and rational exponents are exactly representable, and keep their
+# `IntervalArithmetic` semantics for a base straddling zero.
+for R in (:Integer, :Rational)
+    @eval function Base.:(^)(x::Dual{Tx,<:NestedBareInterval}, y::$R) where {Tx}
+        v = value(x)
+        expv = v^y
+        if iszero(y) || all(_isthinzero, values(partials(x)))
+            return Dual{Tx}(expv, zero(partials(x)))
+        else
+            return Dual{Tx}(expv, partials(x) * exact(y) * v^(y - 1))
+        end
+    end
+end
+
+function Base.:(^)(x::Dual{Tx,<:NestedBareInterval}, y::BareInterval) where {Tx}
+    v = value(x)
+    expv = v^y
+    if isthinzero(y) || all(_isthinzero, values(partials(x)))
+        return Dual{Tx}(expv, zero(partials(x)))
+    else
+        return Dual{Tx}(expv, partials(x) * y * v^(y - bareinterval(1)))
+    end
+end
+
+# `true` and `false` scale a dual without naming an interval of their own.
+# ForwardDiff's `(Dual, Bool)` method inspects `signbit(value(d))`; an interval
+# has no single sign bit and does not distinguish ±0.
+Base.:*(d::Dual{T,<:NestedBareInterval}, x::Bool) where {T} = x ? d : zero(d)
+Base.:*(x::Bool, d::Dual{T,<:NestedBareInterval}) where {T} = d * x
+
+DiffRules._abs_deriv(x::Dual{T,<:NestedBareInterval}) where {T} =
+    Dual{T}(DiffRules._abs_deriv(value(x)), zero(partials(x)))
+
+# Derivative rules for bare intervals
+#
+# The rules tabulated by `DiffRules` embed bare numbers, and ForwardDiff splices
+# them into its `Dual` methods unchanged. Vendor them with `ExactReal` or enclosing
+# bare interval.
+
+_rigorous(x::Integer) = :(exact($x))
+_rigorous(c::AbstractIrrational) = :(bareinterval(_numtype_bare(x), $c, $c))
+_rigorous(x::Number) = error("no rigorous `BareInterval` form for the constant `$x`")
+function _rigorous(s::Symbol)
+    s === :π && return _rigorous(π)
+    s ∈ (:x, :y, :vx, :vy) && return s
+    return error("unrecognized symbol `$s` in a derivative rule")
+end
+function _rigorous(ex::Expr)
+    ex.head === :call || error("unsupported expression head `$(ex.head)` in a derivative rule")
+    return Expr(:call, ex.args[1], map(_rigorous, ex.args[2:end])...)
+end
+
+# The functions of one argument that `DiffRules` tabulates and
+# `IntervalArithmetic` implements on `BareInterval`. `+`, `-`, `sin` and `cos`
+# are absent because ForwardDiff differentiates them without consulting the
+# rule table.
+const BARE_UNARY_RULES = (
+    :abs, :abs2, :acos, :acosh, :acot, :acoth, :asin, :asinh, :atan, :atanh,
+    :cbrt, :cosd, :cosh, :cospi, :cot, :coth, :csc, :csch, :deg2rad, :exp,
+    :exp10, :exp2, :expm1, :inv, :log, :log10, :log1p, :log2, :rad2deg, :sec,
+    :sech, :sind, :sinh, :sinpi, :sqrt, :tan, :tanh,
+)
+
+for f in BARE_UNARY_RULES
+    rule = _rigorous(DiffRules.diffrule(:Base, f, :x))
+    # Share subexpressions between the value and the derivative, as ForwardDiff
+    # does for its own rules: a derivative that repeats the value (`exp`, `inv`,
+    # `tan`, …) otherwise evaluates it twice per `Dual` layer.
+    work = ForwardDiff.qualified_cse!(quote
+        val = $f(x)
+        deriv = $rule
+    end)
+    @eval function Base.$f(d::Dual{T,<:NestedBareInterval}) where {T}
+        x = value(d)
+        $work
+        return ForwardDiff.dual_definition_retval(Val{T}(), val, deriv, partials(d))
+    end
+end
+
+# Binary rules with a bare interval on one side. ForwardDiff's `Dual`-`Dual`
+# methods carry bare intervals through unchanged; its `Dual`-`Real` methods do
+# not apply, because a `BareInterval` is not a `Real`.
+# `^` and the field operations are defined above.
+for f in (:atan, :hypot)
+    dvx, _ = DiffRules.diffrule(:Base, f, :vx, :y)
+    _, dvy = DiffRules.diffrule(:Base, f, :x, :vy)
+    @eval begin
+        function Base.$f(x::Dual{Tx,<:NestedBareInterval}, y::BareInterval) where {Tx}
+            vx = value(x)
+            val = Base.$f(vx, y)
+            return ForwardDiff.dual_definition_retval(Val{Tx}(), val, $(_rigorous(dvx)), partials(x))
+        end
+        function Base.$f(x::BareInterval, y::Dual{Ty,<:NestedBareInterval}) where {Ty}
+            vy = value(y)
+            val = Base.$f(x, vy)
+            return ForwardDiff.dual_definition_retval(Val{Ty}(), val, $(_rigorous(dvy)), partials(y))
+        end
+    end
+end
+
 end

src/intervals/arithmetic/power.jl

@@ -81,6 +81,10 @@ end
 Base.:^(x::Interval, n::Integer) = ^(x, n//one(n))
 Base.:^(x::Interval, y::Rational) = ^(x, convert(Interval{typeof(y)}, y))
 
+# `BareInterval` has no `convert` method for scalar exponents.
+Base.:^(x::BareInterval, n::Integer) = ^(x, n//one(n))
+Base.:^(x::BareInterval, y::Rational) = ^(x, bareinterval(typeof(y), y))
+
 # power for complex intervals
 # See Improve power of complex intervals #745 for a description of the algorithm
 
@@ -125,6 +129,7 @@ Base.:^(x::Complex{<:Interval}, y::Interval) = ^(promote(x, y)...)
 Base.:^(x::Interval, y::Complex{<:Interval}) = ^(promote(x, y)...)
 
 # overwrite behaviour for small integer powers from https://github.com/JuliaLang/julia/pull/24240
+Base.literal_pow(::typeof(^), x::BareInterval, ::Val{n}) where {n} = _select_pown(x, n)
 Base.literal_pow(::typeof(^), x::Interval, ::Val{n}) where {n} = _select_pown(x, n)
 Base.literal_pow(::typeof(^), x::Complex{<:Interval}, ::Val{n}) where {n} = fastpown(x, n)
 

test/ext/IntervalArithmeticDiffRulesExt.jl

@@ -40,7 +40,15 @@ end
 end
 
 @testset "BareInterval" begin
-    @test_throws MethodError ForwardDiff.derivative(abs, bareinterval(-1, 1))
+    @test DiffRules._abs_deriv(bareinterval(-2, -1)) === bareinterval(-1.0, -1.0)
+    @test DiffRules._abs_deriv(bareinterval(0)) === bareinterval(-1.0, 1.0)
+    @test DiffRules._abs_deriv(bareinterval(BigFloat, 1, 2)) isa BareInterval{BigFloat}
+
+    @test ForwardDiff.derivative(abs, bareinterval(-2, -1)) === bareinterval(-1.0, -1.0)
+    @test ForwardDiff.derivative(abs, bareinterval(1, 2)) === bareinterval(1.0, 1.0)
+    @test ForwardDiff.derivative(abs, bareinterval(0)) === bareinterval(-1.0, 1.0)
+    @test ForwardDiff.derivative(abs, bareinterval(-1, 0)) === bareinterval(-1.0, 0.0)
+    @test ForwardDiff.derivative(abs, bareinterval(-2, 2)) === bareinterval(-1.0, 1.0)
 end
 
 @testset "ForwardDiff abs" begin