Add derivative rules in ChainRulesCore

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/sources.jl

@@ -1,4 +1,5 @@
 using DiffEqBase
+import ChainRulesCore
 
 # Define and register smooth functions
 # These are "smooth" aka differentiable and avoid Gibbs effect
@@ -575,6 +576,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)
@@ -649,6 +653,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,7 @@ 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
+ChainRulesCore.frule(_, ::typeof(friction_factor), _...) = 0
 
 function transition(x1, x2, y1, y2, x)
     u = (x - x1) / (x2 - x1)