Merge pull request #204 from SciML/myb/chainrules

Add derivative frules

Author: YingboMa

Project.toml

@@ -4,13 +4,15 @@ authors = ["Chris Rackauckas and Julia Computing"]
 version = "2.1.0"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [compat]
+ChainRulesCore = "1"
 DiffEqBase = "6"
 IfElse = "0.1"
 ModelingToolkit = "8.50"

src/Blocks/Blocks.jl

@@ -3,7 +3,7 @@ The module `Blocks` contains common input-output components, referred to as bloc
 """
 module Blocks
 using ModelingToolkit, Symbolics
-using IfElse: ifelse
+import IfElse: ifelse
 import ..@symcheck
 using ModelingToolkit: getdefault
 

src/Blocks/sources.jl

@@ -1,4 +1,5 @@
 using DiffEqBase
+import ChainRulesCore
 
 # Define and register smooth functions
 # These are "smooth" aka differentiable and avoid Gibbs effect
@@ -496,19 +497,32 @@ Base.:+(x::Number, y::Parameter) = x + y.ref
 Base.:+(y::Parameter, x::Number) = Base.:+(x, y)
 Base.:+(x::Parameter, y::Parameter) = x.ref + y.ref
 
+Base.:-(y::Parameter) = -y.ref
 Base.:-(x::Number, y::Parameter) = x - y.ref
 Base.:-(y::Parameter, x::Number) = y.ref - x
 Base.:-(x::Parameter, y::Parameter) = x.ref - y.ref
 
-Base.:^(x::Number, y::Parameter) = Base.power_by_squaring(x, y.ref)
-Base.:^(y::Parameter, x::Number) = Base.power_by_squaring(y.ref, x)
-Base.:^(x::Parameter, y::Parameter) = Base.power_by_squaring(x.ref, y.ref)
+Base.:^(x::Number, y::Parameter) = Base.:^(x, y.ref)
+Base.:^(y::Parameter, x::Number) = Base.:^(y.ref, x)
+Base.:^(x::Parameter, y::Parameter) = Base.:^(x.ref, y.ref)
 
 Base.isless(x::Parameter, y::Number) = Base.isless(x.ref, y)
 Base.isless(y::Number, x::Parameter) = Base.isless(y, x.ref)
 
 Base.copy(x::Parameter{T}) where {T} = Parameter{T}(copy(x.data), x.ref)
 
+ifelse(c::Bool, x::Parameter, y::Parameter) = ifelse(c, x.ref, y.ref)
+ifelse(c::Bool, x::Parameter, y::Number) = ifelse(c, x.ref, y)
+ifelse(c::Bool, x::Number, y::Parameter) = ifelse(c, x, y.ref)
+
+Base.max(x::Number, y::Parameter) = max(x, y.ref)
+Base.max(x::Parameter, y::Number) = max(x.ref, y)
+Base.max(x::Parameter, y::Parameter) = max(x.ref, y.ref)
+
+Base.min(x::Number, y::Parameter) = min(x, y.ref)
+Base.min(x::Parameter, y::Number) = min(x.ref, y)
+Base.min(x::Parameter, y::Parameter) = min(x.ref, y.ref)
+
 function Base.show(io::IO, m::MIME"text/plain", p::Parameter)
     if !isempty(p.data)
         print(io, p.data)
@@ -575,6 +589,9 @@ function Symbolics.derivative(::typeof(get_sampled_data), args::NTuple{2, Any},
     memory = @inbounds args[2]
     first_order_backwards_difference(t, memory)
 end
+function ChainRulesCore.frule((_, ẋ, _), ::typeof(get_sampled_data), t, memory)
+    first_order_backwards_difference(t, memory) * ẋ
+end
 
 """
     SampledData(; name, buffer)
@@ -614,6 +631,9 @@ function SampledData(data::Vector{T}, dt::T, circular_buffer = true; name) where
 end
 
 Base.convert(::Type{T}, x::Parameter{T}) where {T <: Real} = x.ref
+function Base.convert(::Type{<:Parameter{T}}, x::Number) where {T <: Real}
+    Parameter{T}(T[], x, true)
+end
 
 # Beta Code for potential AE Hack ----------------------
 function set_sampled_data!(memory::Parameter{T}, t, x, Δt::Parameter{T}) where {T}
@@ -649,6 +669,14 @@ function Symbolics.derivative(::typeof(set_sampled_data!), args::NTuple{4, Any},
     first_order_backwards_difference(t, x, Δt, memory)
 end
 Symbolics.derivative(::typeof(set_sampled_data!), args::NTuple{4, Any}, ::Val{3}) = 1 #set_sampled_data returns x, therefore d/dx (x) = 1
+function ChainRulesCore.frule((_, _, ṫ, ẋ, _),
+    ::typeof(set_sampled_data!),
+    memory,
+    t,
+    x,
+    Δt)
+    first_order_backwards_difference(t, x, Δt, memory) * ṫ + ẋ
+end
 
 function first_order_backwards_difference(t, x, Δt, memory)
     x1 = set_sampled_data!(memory, t, x, Δt)

src/Hydraulic/IsothermalCompressible/utils.jl

@@ -1,3 +1,5 @@
+import ChainRulesCore
+
 regPow(x, a, delta = 0.01) = x * (x * x + delta * delta)^((a - 1) / 2);
 regRoot(x, delta = 0.01) = regPow(x, 0.5, delta)
 
@@ -116,6 +118,9 @@ end
 @register_symbolic friction_factor(dm, area, d_h, viscosity, shape_factor)
 Symbolics.derivative(::typeof(friction_factor), args, ::Val{1}) = 0
 Symbolics.derivative(::typeof(friction_factor), args, ::Val{4}) = 0
+function ChainRulesCore.frule(_, ::typeof(friction_factor), args...)
+    (friction_factor(args...), ChainRulesCore.ZeroTangent)
+end
 
 function transition(x1, x2, y1, y2, x)
     u = (x - x1) / (x2 - x1)

src/Mechanical/Translational/Translational.jl

@@ -18,7 +18,7 @@ include("utils.jl")
 export Mass, Spring, Damper, Fixed
 include("components.jl")
 
-export Force, Position
+export Force, Position, Velocity, Acceleration
 include("sources.jl")
 
 export ForceSensor, PositionSensor