Merge pull request #531 from vpuri3/ad

overload _alg_autodiff

Author: ChrisRackauckas

Project.toml

@@ -25,6 +25,7 @@ RandomNumbers = "e6cf234a-135c-5ec9-84dd-332b85af5143"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -44,7 +45,7 @@ JumpProcesses = "9"
 LevyArea = "1.0.0"
 MuladdMacro = "0.2.1"
 NLsolve = "4"
-OrdinaryDiffEq = "6.4"
+OrdinaryDiffEq = "6.52"
 RandomNumbers = "1.5.3"
 RecursiveArrayTools = "2"
 Reexport = "0.2, 1.0"

src/StochasticDiffEq.jl

@@ -24,7 +24,7 @@ using DocStringExtensions
   import DiffEqBase: ODE_DEFAULT_NORM, ODE_DEFAULT_ISOUTOFDOMAIN,
          ODE_DEFAULT_PROG_MESSAGE, ODE_DEFAULT_UNSTABLE_CHECK
 
-  using DiffEqBase: DiffEqArrayOperator
+  using SciMLOperators: MatrixOperator
 
   using DiffEqBase: TimeGradientWrapper, UJacobianWrapper, TimeDerivativeWrapper, UDerivativeWrapper
 

src/alg_utils.jl

@@ -260,10 +260,10 @@ alg_needs_extra_process(alg::PL1WM) = true
 alg_needs_extra_process(alg::NON) = true
 alg_needs_extra_process(alg::NON2) = true
 
-OrdinaryDiffEq.alg_autodiff(alg::StochasticDiffEqNewtonAlgorithm{CS,AD,FDT,ST,CJ,Controller}) where {CS,AD,FDT,ST,CJ,Controller} = AD
-OrdinaryDiffEq.alg_autodiff(alg::StochasticDiffEqNewtonAdaptiveAlgorithm{CS,AD,FDT,ST,CJ,Controller}) where {CS,AD,FDT,ST,CJ,Controller} = AD
-OrdinaryDiffEq.alg_autodiff(alg::StochasticDiffEqJumpNewtonAdaptiveAlgorithm{CS,AD,FDT,ST,CJ,Controller}) where {CS,AD,FDT,ST,CJ,Controller} = AD
-OrdinaryDiffEq.alg_autodiff(alg::StochasticDiffEqJumpNewtonDiffusionAdaptiveAlgorithm{CS,AD,FDT,ST,CJ,Controller}) where {CS,AD,FDT,ST,CJ,Controller} = AD
+OrdinaryDiffEq._alg_autodiff(alg::StochasticDiffEqNewtonAlgorithm{CS,AD,FDT,ST,CJ,Controller}) where {CS,AD,FDT,ST,CJ,Controller} = Val{AD}()
+OrdinaryDiffEq._alg_autodiff(alg::StochasticDiffEqNewtonAdaptiveAlgorithm{CS,AD,FDT,ST,CJ,Controller}) where {CS,AD,FDT,ST,CJ,Controller} = Val{AD}()
+OrdinaryDiffEq._alg_autodiff(alg::StochasticDiffEqJumpNewtonAdaptiveAlgorithm{CS,AD,FDT,ST,CJ,Controller}) where {CS,AD,FDT,ST,CJ,Controller} = Val{AD}()
+OrdinaryDiffEq._alg_autodiff(alg::StochasticDiffEqJumpNewtonDiffusionAdaptiveAlgorithm{CS,AD,FDT,ST,CJ,Controller}) where {CS,AD,FDT,ST,CJ,Controller} = Val{AD}()
 
 OrdinaryDiffEq.get_current_alg_autodiff(alg::StochasticDiffEqCompositeAlgorithm, cache) = OrdinaryDiffEq.alg_autodiff(alg.algs[cache.current])
 

src/integrators/integrator_utils.jl

@@ -418,7 +418,7 @@ function oop_generate_W(alg,u,uprev,p,t,dt,f,uEltypeNoUnits)
     # get the operator
     J = islin ? nf.f : f.jac(uprev, p, t)
     if !isa(J, DiffEqBase.AbstractDiffEqLinearOperator)
-      J = DiffEqArrayOperator(J)
+      J = MatrixOperator(J)
     end
     W = WOperator{false}(f.mass_matrix, dt, J, u)
   else

test/utility_tests.jl

@@ -1,5 +1,6 @@
 using StochasticDiffEq, LinearAlgebra, SparseArrays, Random, LinearSolve, Test
 using StochasticDiffEq.OrdinaryDiffEq: WOperator, set_gamma!, calc_W!
+using StochasticDiffEq.SciMLOperators: MatrixOperator
 
 @testset "Derivative Utilities" begin
   @testset "calc_W!" begin
@@ -25,7 +26,7 @@ using StochasticDiffEq.OrdinaryDiffEq: WOperator, set_gamma!, calc_W!
     _f = (du,u,p,t) -> mul!(du,A,u); _g = (du,u,p,t) -> mul!(du,σ,u)
     fun = SDEFunction(_f, _g;
                       mass_matrix=mm,
-                      jac_prototype=DiffEqArrayOperator(A))
+                      jac_prototype=MatrixOperator(A))
     prob = SDEProblem(fun, _g, u0, tspan)
     integrator = init(prob, ImplicitEM(theta=1); adaptive=false, dt=dt)
     W = integrator.cache.nlsolver.cache.W
@@ -52,7 +53,7 @@ using StochasticDiffEq.OrdinaryDiffEq: WOperator, set_gamma!, calc_W!
     prob1 = SDEProblem(SDEFunction(_f, _g; mass_matrix=mm), _g, u0, tspan)
     prob2 = SDEProblem(SDEFunction(_f, _g; mass_matrix=mm, jac=(u,p,t) -> t*A), _g, u0, tspan)
     prob1_ip = SDEProblem(SDEFunction(_f_ip, _g_ip; mass_matrix=mm), _g_ip, u0, tspan)
-    jac_prototype=DiffEqArrayOperator(similar(A); update_func=(J,u,p,t) -> (J .= t .* A; J))
+    jac_prototype=MatrixOperator(similar(A); update_func! = (J,u,p,t) -> (J .= t .* A; J))
     prob2_ip = SDEProblem(SDEFunction(_f_ip, _g_ip; mass_matrix=mm, jac_prototype=jac_prototype), _g_ip, u0, tspan)
 
     for Alg in [ImplicitEM, ISSEM]