Merge pull request #876 from huanglangwen/simplify_dutils

[WIP] Simplify `derivative_utils.jl`

Author: devmotion

src/derivative_utils.jl

@@ -51,6 +51,8 @@ function calc_J(integrator, cache::OrdinaryDiffEqConstantCache, is_compos)
   if DiffEqBase.has_jac(f)
     J = f.jac(uprev, p, t)
   else
+    cache.uf.t = t
+    cache.uf.p = p
     J = jacobian(cache.uf,uprev,integrator)
   end
   integrator.destats.njacs += 1
@@ -63,6 +65,8 @@ function calc_J(nlsolver, integrator, cache::OrdinaryDiffEqConstantCache, is_com
   if DiffEqBase.has_jac(f)
     J = f.jac(uprev, p, t)
   else
+    nlsolver.uf.t = t
+    nlsolver.uf.p = p
     J = jacobian(nlsolver.uf,uprev,integrator)
   end
   integrator.destats.njacs += 1
@@ -81,7 +85,7 @@ jacobian update function, then it will be called for the update. Otherwise,
 either ForwardDiff or finite difference will be used depending on the
 `jac_config` of the cache.
 """
-function calc_J!(integrator, cache::OrdinaryDiffEqMutableCache, is_compos)
+function calc_J!(integrator, cache::OrdinaryDiffEqCache, is_compos)
   if isdefined(cache, :nlsolver)
     calc_J!(cache.nlsolver, integrator, cache, is_compos)
   elseif isdefined(cache, :J)
@@ -106,6 +110,10 @@ function calc_J!(nlsolver::NLSolver, integrator, cache::OrdinaryDiffEqMutableCac
   is_compos && (integrator.eigen_est = opnorm(J, Inf))
 end
 
+function calc_J!(nlsolver::NLSolver, integrator, cache::OrdinaryDiffEqConstantCache, is_compos)
+  nlsolver.cache.J = calc_J(nlsolver,integrator,cache,is_compos)
+end
+
 function calc_J_in_cache!(integrator, cache::OrdinaryDiffEqMutableCache, is_compos)
   @unpack t,dt,uprev,u,f,p = integrator
   J = cache.J
@@ -121,6 +129,10 @@ function calc_J_in_cache!(integrator, cache::OrdinaryDiffEqMutableCache, is_comp
   is_compos && (integrator.eigen_est = opnorm(J, Inf))
 end
 
+function calc_J_in_cache!(integrator, cache::OrdinaryDiffEqConstantCache, is_compos)
+  cache.J = calc_J(integrator,cache,is_compos)
+end
+
 """
     WOperator(mass_matrix,gamma,J[;transform=false])
 
@@ -312,6 +324,7 @@ end
 
 @noinline _throwWJerror(W, J) = throw(DimensionMismatch("W: $(axes(W)), J: $(axes(J))"))
 @noinline _throwWMerror(W, mass_matrix) = throw(DimensionMismatch("W: $(axes(W)), mass matrix: $(axes(mass_matrix))"))
+@noinline _throwJMerror(J, mass_matrix) = throw(DimensionMismatch("J: $(axes(J)), mass matrix: $(axes(mass_matrix))"))
 
 @inline function jacobian2W!(W::AbstractMatrix, mass_matrix::MT, dtgamma::Number, J::AbstractMatrix, W_transform::Bool)::Nothing where MT
   # check size and dimension
@@ -341,6 +354,28 @@ end
   return nothing
 end
 
+@inline function jacobian2W(mass_matrix::MT, dtgamma::Number, J::AbstractMatrix, W_transform::Bool)::Nothing where MT
+  # check size and dimension
+  mass_matrix isa UniformScaling || @boundscheck axes(mass_matrix) === axes(J) || _throwJMerror(J, mass_matrix)
+  @inbounds if W_transform
+    invdtgamma = inv(dtgamma)
+    if MT <: UniformScaling
+      λ = -mass_matrix.λ
+      W = J + (λ * invdtgamma)*I
+    else
+      W = muladd(-mass_matrix, invdtgamma, J)
+    end
+  else
+    if MT <: UniformScaling
+      λ = -mass_matrix.λ
+      W = dtgamma*J + λ*I
+    else
+      W = muladd(dtgamma, J, -mass_matrix)
+    end
+  end
+  return W
+end
+
 function calc_W!(integrator, cache::OrdinaryDiffEqMutableCache, dtgamma, repeat_step, W_transform=false)
   @unpack t,dt,uprev,u,f,p = integrator
   @unpack J,W = cache