Merge pull request #481 from avik-pal/ap/vecjac

VecJacOperator

Author: ChrisRackauckas

.github/workflows/CI.yml

@@ -9,6 +9,7 @@ on:
 jobs:
   test:
     runs-on: ubuntu-latest
+    continue-on-error: ${{ matrix.group == 'MOLFiniteDifference' }}
     strategy:
       matrix:
         group:

Project.toml

@@ -16,9 +16,11 @@ LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
@@ -37,6 +39,7 @@ NNlib = "0.6, 0.7"
 NonlinearSolve = "0.3.7"
 RuntimeGeneratedFunctions = "0.4, 0.5"
 SciMLBase = "1.11"
+SparseDiffTools = "1.17"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"
 SymbolicUtils = "0.11, 0.12, 0.13, 0.14, 0.15, 0.16, 0.17"
 julia = "1"
@@ -49,6 +52,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["FillArrays", "OrdinaryDiffEq", "Parameters", "Random", "SafeTestsets", "SpecialFunctions", "Test"]
+test = ["FillArrays", "OrdinaryDiffEq", "Parameters", "Random", "SafeTestsets", "SpecialFunctions", "Test", "Zygote"]

docs/make.jl

@@ -22,6 +22,7 @@ makedocs(
             "operators/operator_overview.md",
             "operators/derivative_operators.md",
             "operators/vector_calculus_operators.md",
+            "operators/vector_jacobian_product.md",
             "operators/jacobian_vector_product.md",
             "operators/matrix_free_operators.md"
         ],

docs/src/operators/vector_jacobian_product.md

@@ -0,0 +1,12 @@
+# Vector-Jacobian Product Operators
+
+```julia
+VecJacOperator{T}(f,u::AbstractArray,p=nothing,t::Union{Nothing,Number}=nothing;autodiff=true,ishermitian=false,opnorm=true)
+```
+
+The `VecJacOperator` is a linear operator `J'*v` where `J` acts like `df/du`
+for some function `f(u,p,t)`. For in-place operations `mul!(w,J,v)`, `f`
+is an in-place function `f(du,u,p,t)`.
+
+!!! note
+    This operator is available when `Zygote` is imported.

src/DiffEqOperators.jl

@@ -7,9 +7,11 @@ import LinearAlgebra: mul!, ldiv!, lmul!, rmul!, axpy!, opnorm, factorize, I
 import DiffEqBase: update_coefficients!, isconstant
 using SciMLBase: AbstractDiffEqLinearOperator, AbstractDiffEqCompositeOperator, DiffEqScaledOperator
 import SciMLBase: getops
+using SparseDiffTools
 using SparseArrays, ForwardDiff, BandedMatrices, NNlib, LazyArrays, BlockBandedMatrices, LoopVectorization
 using LazyBandedMatrices, ModelingToolkit
 using RuntimeGeneratedFunctions
+using Requires
 RuntimeGeneratedFunctions.init(@__MODULE__)
 
 abstract type AbstractDiffEqAffineOperator{T} end
@@ -63,6 +65,14 @@ for T in [DiffEqScaledOperator, DiffEqOperatorCombination, DiffEqOperatorComposi
   (L::T)(du,u,p,t) = (update_coefficients!(L,u,p,t); mul!(du,L,u))
 end
 
+function __init__()
+  @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" begin
+    export VecJacOperator
+
+    include("vecjac_operators.jl")
+  end
+end
+
 export MatrixFreeOperator
 export AnalyticalJacVecOperator, JacVecOperator, getops
 export AbstractDerivativeOperator, DerivativeOperator,