Lazy W operator support

REQUIRE

@@ -1,5 +1,6 @@
 julia 0.7-beta2
 DiffEqBase 3.8.0
+DiffEqOperators 3.2.0
 Parameters 0.5.0
 ForwardDiff 0.7.0
 GenericSVD 0.0.2

src/OrdinaryDiffEq.jl

@@ -17,6 +17,8 @@ module OrdinaryDiffEq
   # Internal utils
   import DiffEqBase: ODE_DEFAULT_NORM, ODE_DEFAULT_ISOUTOFDOMAIN, ODE_DEFAULT_PROG_MESSAGE, ODE_DEFAULT_UNSTABLE_CHECK
 
+  using DiffEqOperators: DiffEqArrayOperator
+
   import RecursiveArrayTools: chain, recursivecopy!
 
   using Parameters, GenericSVD, ForwardDiff, RecursiveArrayTools,

src/caches/sdirk_caches.jl

@@ -1,4 +1,4 @@
-mutable struct ImplicitEulerCache{uType,rateType,uNoUnitsType,J,UF,JC,uToltype,F} <: OrdinaryDiffEqMutableCache
+mutable struct ImplicitEulerCache{uType,rateType,uNoUnitsType,J,W,UF,JC,uToltype,F} <: OrdinaryDiffEqMutableCache
   u::uType
   uprev::uType
   uprev2::uType
@@ -11,7 +11,7 @@ mutable struct ImplicitEulerCache{uType,rateType,uNoUnitsType,J,UF,JC,uToltype,F
   tmp::uType
   atmp::uNoUnitsType
   J::J
-  W::J
+  W::W
   uf::UF
   jac_config::JC
   linsolve::F
@@ -28,8 +28,13 @@ function alg_cache(alg::ImplicitEuler,u,rate_prototype,uEltypeNoUnits,uBottomElt
                    tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
 
   du1 = zero(rate_prototype)
-  J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
-  W = similar(J)
+  if DiffEqBase.has_jac(f)
+    W = WOperator(f, dt)
+    J = nothing # is J = W.J better?
+  else
+    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
+    W = similar(J)
+  end
   z = similar(u,axes(u))
   dz = similar(u,axes(u)); tmp = similar(u,axes(u)); b = similar(u,axes(u))
   fsalfirst = zero(rate_prototype)

src/derivative_utils.jl

@@ -54,10 +54,161 @@ function calc_J!(integrator, cache::OrdinaryDiffEqMutableCache, is_compos)
   is_compos && (integrator.eigen_est = opnorm(J, Inf))
 end
 
+"""
+    WOperator(mass_matrix,gamma,J[;transform=false])
+
+A linear operator that represents the W matrix of an ODEProblem, defined as
+
+```math
+W = MM - \\gamma J
+```
+
+or, if `transform=true`:
+
+```math
+W = \\frac{1}{\\gamma}MM - J
+```
+
+where `MM` is the mass matrix (a regular `AbstractMatrix` or a `UniformScaling`),
+`γ` is a real number proportional to the time step, and `J` is the Jacobian
+operator (must be a `AbstractDiffEqLinearOperator`). A `WOperator` can also be
+constructed using a `*DEFunction` directly as
+
+    WOperator(f,gamma[;transform=false])
+
+`f` needs to have a jacobian and `jac_prototype`, but the prototype does not need
+to be a diffeq operator --- it will automatically be converted to one.
+
+`WOperator` supports lazy `*` and `mul!` operations, the latter utilizing an
+internal cache (can be specified in the constructor; default to regular `Vector`).
+It supports all of `AbstractDiffEqLinearOperator`'s interface.
+"""
+mutable struct WOperator{T,
+  MType <: Union{UniformScaling,AbstractMatrix},
+  GType <: Real,
+  JType <: DiffEqBase.AbstractDiffEqLinearOperator{T}
+  } <: DiffEqBase.AbstractDiffEqLinearOperator{T}
+  mass_matrix::MType
+  gamma::GType
+  J::JType
+  transform::Bool       # true => W = mm/gamma - J; false => W = mm - gamma*J
+  _func_cache           # cache used in `mul!`
+  _concrete_form         # non-lazy form (matrix/number) of the operator
+  WOperator(mass_matrix, gamma, J; transform=false) = new{eltype(J),typeof(mass_matrix),
+    typeof(gamma),typeof(J)}(mass_matrix,gamma,J,transform,nothing,nothing)
+end
+function WOperator(f::DiffEqBase.AbstractODEFunction, gamma; transform=false)
+  @assert DiffEqBase.has_jac(f) "f needs to have an associated jacobian"
+  if isa(f, Union{SplitFunction, DynamicalODEFunction})
+    error("WOperator does not support $(typeof(f)) yet")
+  end
+  # Convert mass matrix, if needed
+  mass_matrix = f.mass_matrix
+  if !isa(mass_matrix, Union{AbstractMatrix,UniformScaling})
+    mass_matrix = convert(AbstractMatrix, mass_matrix)
+  end
+  # Convert jacobian, if needed
+  J = deepcopy(f.jac_prototype)
+  if !isa(J, DiffEqBase.AbstractDiffEqLinearOperator)
+    J = DiffEqArrayOperator(J; update_func=f.jac)
+  end
+  return WOperator(mass_matrix, gamma, J; transform=transform)
+end
+
+set_gamma!(W::WOperator, gamma) = (W.gamma = gamma; W)
+DiffEqBase.update_coefficients!(W::WOperator,u,p,t) = (update_coefficients!(W.J,u,p,t); W)
+function Base.convert(::Type{AbstractMatrix}, W::WOperator)
+  if W._concrete_form == nothing
+    # Allocating
+    if W.transform
+      W._concrete_form = W.mass_matrix / W.gamma - convert(AbstractMatrix,W.J)
+    else
+      W._concrete_form = W.mass_matrix - W.gamma * convert(AbstractMatrix,W.J)
+    end
+  else
+    # Non-allocating
+    if W.transform
+      rmul!(copyto!(W._concrete_form, W.mass_matrix), 1/W.gamma)
+      axpy!(-1, convert(AbstractMatrix,W.J), W._concrete_form)
+    else
+      copyto!(W._concrete_form, W.mass_matrix)
+      axpy!(-W.gamma, convert(AbstractMatrix,W.J), W._concrete_form)
+    end
+  end
+  W._concrete_form
+end
+function Base.convert(::Type{Number}, W::WOperator)
+  if W.transform
+    W._concrete_form = W.mass_matrix / W.gamma - convert(Number,W.J)
+  else
+    W._concrete_form = W.mass_matrix - W.gamma * convert(Number,W.J)
+  end
+  W._concrete_form
+end
+Base.size(W::WOperator, args...) = size(W.J, args...)
+function Base.getindex(W::WOperator, i::Int)
+  if W.transform
+    W.mass_matrix[i] / W.gamma - W.J[i]
+  else
+    W.mass_matrix[i] - W.gamma * W.J[i]
+  end
+end
+function Base.getindex(W::WOperator, I::Vararg{Int,N}) where {N}
+  if W.transform
+    W.mass_matrix[I...] / W.gamma - W.J[I...]
+  else
+    W.mass_matrix[I...] - W.gamma * W.J[I...]
+  end
+end
+function Base.:*(W::WOperator, x::Union{AbstractVecOrMat,Number})
+  if W.transform
+    (W.mass_matrix*x) / W.gamma - W.J*x
+  else
+    W.mass_matrix*x - W.gamma * (W.J*x)
+  end
+end
+function Base.:\(W::WOperator, x::Union{AbstractVecOrMat,Number})
+  if size(W) == () # scalar operator
+    convert(Number,W) \ x
+  else
+    convert(AbstractMatrix,W) \ x
+  end
+end
+function LinearAlgebra.mul!(Y::AbstractVecOrMat, W::WOperator, B::AbstractVecOrMat)
+  if W._func_cache == nothing
+    # Allocate cache only if needed
+    W._func_cache = Vector{eltype(W)}(undef, size(Y, 1))
+  end
+  if W.transform
+    # Compute mass_matrix * B
+    if isa(W.mass_matrix, UniformScaling)
+      a = W.mass_matrix.λ / W.gamma
+      @. Y = a * B
+    else
+      mul!(Y, W.mass_matrix, B)
+      lmul!(1/W.gamma, Y)
+    end
+    # Compute J * B and subtract
+    mul!(W._func_cache, W.J, B)
+    Y .-= W._func_cache
+  else
+    # Compute mass_matrix * B
+    if isa(W.mass_matrix, UniformScaling)
+      @. Y = W.mass_matrix.λ * B
+    else
+      mul!(Y, W.mass_matrix, B)
+    end
+    # Compute J * B
+    mul!(W._func_cache, W.J, B)
+    # Subtract result
+    axpy!(-W.gamma, W._func_cache, Y)
+  end
+end
+
 function calc_W!(integrator, cache::OrdinaryDiffEqMutableCache, dtgamma, repeat_step, W_transform=false)
   @inbounds begin
     @unpack t,dt,uprev,u,f,p = integrator
-    @unpack J,W,jac_config = cache
+    @unpack J,W = cache
     mass_matrix = integrator.f.mass_matrix
     is_compos = typeof(integrator.alg) <: CompositeAlgorithm
     alg = unwrap_alg(integrator, true)
@@ -78,21 +229,14 @@ function calc_W!(integrator, cache::OrdinaryDiffEqMutableCache, dtgamma, repeat_
         new_jac = false
       else
         new_jac = true
-        calc_J!(integrator, cache, is_compos)
+        DiffEqBase.update_coefficients!(W,uprev,p,t)
       end
       # skip calculation of W if step is repeated
       if !repeat_step && (!alg_can_repeat_jac(alg) ||
                           (integrator.iter < 1 || new_jac ||
                            abs(dt - (t-integrator.tprev)) > 100eps(typeof(integrator.t))))
-        if W_transform
-          for j in 1:length(u), i in 1:length(u)
-              W[i,j] = mass_matrix[i,j]/dtgamma - J[i,j]
-          end
-        else
-          for j in 1:length(u), i in 1:length(u)
-              W[i,j] = mass_matrix[i,j] - dtgamma*J[i,j]
-          end
-        end
+        W.transform = W_transform
+        set_gamma!(W, dtgamma)
       else
         new_W = false
       end
@@ -102,27 +246,42 @@ function calc_W!(integrator, cache::OrdinaryDiffEqMutableCache, dtgamma, repeat_
 end
 
 function calc_W!(integrator, cache::OrdinaryDiffEqConstantCache, dtgamma, repeat_step, W_transform=false)
-  @unpack t,uprev,f = integrator
+  @unpack t,uprev,p,f = integrator
   @unpack uf = cache
+  mass_matrix = integrator.f.mass_matrix
   # calculate W
   uf.t = t
   isarray = typeof(uprev) <: AbstractArray
   iscompo = typeof(integrator.alg) <: CompositeAlgorithm
   if !W_transform
-    if isarray
-      J = ForwardDiff.jacobian(uf,uprev)
-      W = I - dtgamma*J
+    if DiffEqBase.has_jac(f)
+      J = f.jac(uprev, p, t)
+      if !isa(f, DiffEqBase.AbstractDiffEqLinearOperator)
+        J = DiffEqArrayOperator(J)
+      end
+      W = WOperator(mass_matrix, dtgamma, J; transform=false)
     else
-      J = ForwardDiff.derivative(uf,uprev)
-      W = 1 - dtgamma*J
+      if isarray
+        J = ForwardDiff.jacobian(uf,uprev)
+      else
+        J = ForwardDiff.derivative(uf,uprev)
+      end
+      W = mass_matrix - dtgamma*J
     end
   else
-    if isarray
-      J = ForwardDiff.jacobian(uf,uprev)
-      W = I*inv(dtgamma) - J
+    if DiffEqBase.has_jac(f)
+      J = f.jac(uprev, p, t)
+      if !isa(f, DiffEqBase.AbstractDiffEqLinearOperator)
+        J = DiffEqArrayOperator(J)
+      end
+      W = WOperator(mass_matrix, dtgamma, J; transform=true)
     else
-      J = ForwardDiff.derivative(uf,uprev)
-      W = inv(dtgamma) - J
+      if isarray
+        J = ForwardDiff.jacobian(uf,uprev)
+      else
+        J = ForwardDiff.derivative(uf,uprev)
+      end
+      W = mass_matrix*inv(dtgamma) - J
     end
   end
   iscompo && (integrator.eigen_est = isarray ? opnorm(J, Inf) : J)

src/perform_step/sdirk_perform_step.jl

@@ -68,7 +68,7 @@ end
     c = 7/12 # default correction factor in SPICE (LTE overestimated by DD)
     r = c*dt^2 # by mean value theorem 2nd DD equals y''(s)/2 for some s
 
-    tmp = r*abs((u - uprev)/dt1 - (uprev - uprev2)/dt2)
+    tmp = r*abs.((u - uprev)/dt1 - (uprev - uprev2)/dt2)
     atmp = calculate_residuals(tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm)
     integrator.EEst = integrator.opts.internalnorm(atmp)
   else

test/utility_tests.jl

@@ -1,5 +1,6 @@
 using OrdinaryDiffEq: phi, phi, phiv, phiv_timestep, expv, expv_timestep, arnoldi, getH, getV
-using LinearAlgebra, SparseArrays, Random, Test
+using OrdinaryDiffEq: WOperator, set_gamma!, calc_W!
+using OrdinaryDiffEq, LinearAlgebra, SparseArrays, Random, Test, DiffEqOperators
 
 @testset "Exponential Utilities" begin
   # Scalar phi
@@ -67,3 +68,42 @@ using LinearAlgebra, SparseArrays, Random, Test
   u = expv_timestep(t, A, B[:, 1]; adaptive=true, tol=tol)
   @test norm(u - u_exact) / norm(u_exact) < tol
 end
+
+@testset "Derivative Utilities" begin
+  @testset "WOperator" begin
+    srand(0); y = zeros(2); b = rand(2)
+    mm = I; _J = rand(2,2)
+    W = WOperator(mm, 1.0, DiffEqArrayOperator(_J))
+    set_gamma!(W, 2.0)
+    _W = mm - 2.0 * _J
+    @test convert(AbstractMatrix,W) ≈ _W
+    @test W * b ≈ _W * b
+    mul!(y, W, b); @test y ≈ _W * b
+  end
+
+  @testset "calc_W!" begin
+    A = [-1.0 0.0; 0.0 -0.5]; mm = [2.0 0.0; 0.0 1.0]
+    u0 = [1.0, 1.0]; tmp = zeros(2)
+    tspan = (0.0,1.0); dt = 0.01; dtgamma = 0.5dt
+    concrete_W = mm - dtgamma * A
+    concrete_Wt = mm/dtgamma - A
+
+    # Out-of-place
+    fun = ODEFunction((u,p,t) -> A*u;
+                      mass_matrix=mm,
+                      jac=(u,p,t) -> A)
+    integrator = init(ODEProblem(fun,u0,tspan), ImplicitEuler(); adaptive=false, dt=dt)
+    W = calc_W!(integrator, integrator.cache, dtgamma, false)
+    @test convert(AbstractMatrix, W) ≈ concrete_W
+    @test W \ u0 ≈ concrete_W \ u0
+
+    # In-place
+    fun = ODEFunction((du,u,p,t) -> mul!(du,A,u);
+                      mass_matrix=mm,
+                      jac_prototype=DiffEqArrayOperator(A))
+    integrator = init(ODEProblem(fun,u0,tspan), ImplicitEuler(); adaptive=false, dt=dt)
+    calc_W!(integrator, integrator.cache, dtgamma, false)
+    @test convert(AbstractMatrix, integrator.cache.W) ≈ concrete_W
+    ldiv!(tmp, lu!(W), u0); @test tmp ≈ concrete_W \ u0
+  end
+end