feat: integrate SciMLJacobianOperators into NonlinearSolve

Author: avik-pal

docs/src/devdocs/operators.md

@@ -6,21 +6,6 @@
 NonlinearSolve.AbstractNonlinearSolveOperator
 ```
 
-## Jacobian Operators
-
-```@docs
-NonlinearSolve.JacobianOperator
-NonlinearSolve.VecJacOperator
-NonlinearSolve.JacVecOperator
-```
-
-### Stateful Jacobian Operators
-
-```@docs
-NonlinearSolve.StatefulJacobianOperator
-NonlinearSolve.StatefulJacobianNormalFormOperator
-```
-
 ## Low-Rank Jacobian Operators
 
 ```@docs

lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl

@@ -5,6 +5,7 @@ using ConcreteStructs: @concrete
 using ConstructionBase: ConstructionBase
 using DifferentiationInterface: DifferentiationInterface
 using FastClosures: @closure
+using LinearAlgebra: LinearAlgebra
 using SciMLBase: SciMLBase, AbstractNonlinearProblem, AbstractNonlinearFunction
 using SciMLOperators: AbstractSciMLOperator
 using Setfield: @set!
@@ -23,6 +24,57 @@ struct JVP <: AbstractMode end
 flip_mode(::VJP) = JVP()
 flip_mode(::JVP) = VJP()
 
+"""
+    JacobianOperator{iip, T} <: AbstractJacobianOperator{T} <: AbstractSciMLOperator{T}
+
+A Jacobian Operator Provides both JVP and VJP without materializing either (if possible).
+
+### Constructor
+
+```julia
+JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = nothing,
+    vjp_autodiff = nothing, skip_vjp::Val = Val(false), skip_jvp::Val = Val(false))
+```
+
+By default, the `JacobianOperator` will compute `JVP`. Use `Base.adjoint` or
+`Base.transpose` to switch to `VJP`.
+
+### Computing the VJP
+
+Computing the VJP is done according to the following rules:
+
+  - If `f` has a `vjp` method, then we use that.
+  - If `f` has a `jac` method and no `vjp_autodiff` is provided, then we use `jac * v`.
+  - If `vjp_autodiff` is provided we using DifferentiationInterface.jl to compute the VJP.
+
+### Computing the JVP
+
+Computing the JVP is done according to the following rules:
+
+  - If `f` has a `jvp` method, then we use that.
+  - If `f` has a `jac` method and no `jvp_autodiff` is provided, then we use `v * jac`.
+  - If `jvp_autodiff` is provided we using DifferentiationInterface.jl to compute the JVP.
+
+### Special Case (Number)
+
+For Number inputs, VJP and JVP are not distinct. Hence, if either `vjp` or `jvp` is
+provided, then we use that. If neither is provided, then we use `v * jac` if `jac` is
+provided. Finally, we use the respective autodiff methods to compute the derivative
+using DifferentiationInterface.jl and multiply by `v`.
+
+### Methods Provided
+
+!!! warning
+
+    Currently it is expected that `p` during problem construction is same as `p` during
+    operator evaluation. This restriction will be lifted in the future.
+
+  - `(op::JacobianOperator)(v, u, p)`: Computes `∂f(u, p)/∂u * v` or `∂f(u, p)/∂uᵀ * v`.
+  - `(op::JacobianOperator)(res, v, u, p)`: Computes `∂f(u, p)/∂u * v` or `∂f(u, p)/∂uᵀ * v`
+    and stores the result in `res`.
+
+See also [`VecJacOperator`](@ref) and [`JacVecOperator`](@ref).
+"""
 @concrete struct JacobianOperator{iip, T <: Real} <: AbstractJacobianOperator{T}
     mode <: AbstractMode
 
@@ -65,8 +117,8 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff =
     vjp_op = prepare_vjp(skip_vjp, prob, f, u, fu; autodiff = vjp_autodiff)
     jvp_op = prepare_jvp(skip_jvp, prob, f, u, fu; autodiff = jvp_autodiff)
 
-    output_cache = iip ? similar(fu, T) : nothing
-    input_cache = iip ? similar(u, T) : nothing
+    output_cache = similar(fu, T)
+    input_cache = similar(u, T)
 
     return JacobianOperator{iip, T}(
         JVP(), jvp_op, vjp_op, (length(fu), length(u)), output_cache, input_cache)
@@ -112,14 +164,106 @@ function (op::JacobianOperator)(Jv, v, u, p)
     end
 end
 
+"""
+    VecJacOperator(args...; autodiff = nothing, kwargs...)
+
+Constructs a [`JacobianOperator`](@ref) which only provides the VJP using the
+`vjp_autodiff = autodiff`.
+"""
 function VecJacOperator(args...; autodiff = nothing, kwargs...)
     return JacobianOperator(args...; kwargs..., skip_jvp = True, vjp_autodiff = autodiff)'
 end
 
+"""
+    JacVecOperator(args...; autodiff = nothing, kwargs...)
+
+Constructs a [`JacobianOperator`](@ref) which only provides the JVP using the
+`jvp_autodiff = autodiff`.
+"""
 function JacVecOperator(args...; autodiff = nothing, kwargs...)
     return JacobianOperator(args...; kwargs..., skip_vjp = True, jvp_autodiff = autodiff)
 end
 
+"""
+    StatefulJacobianOperator(jac_op::JacobianOperator, u, p)
+
+Wrapper over a [`JacobianOperator`](@ref) which stores the input `u` and `p` and defines
+`mul!` and `*` for computing VJPs and JVPs.
+"""
+@concrete struct StatefulJacobianOperator{M <: AbstractMode, T} <:
+                 AbstractJacobianOperator{T}
+    mode::M
+    jac_op <: JacobianOperator
+    u
+    p
+
+    function StatefulJacobianOperator(jac_op::JacobianOperator, u, p)
+        return new{
+            typeof(jac_op.mode), eltype(jac_op), typeof(jac_op), typeof(u), typeof(p)}(
+            jac_op.mode, jac_op, u, p)
+    end
+end
+
+Base.size(J::StatefulJacobianOperator) = size(J.jac_op)
+Base.size(J::StatefulJacobianOperator, d::Integer) = size(J.jac_op, d)
+
+for op in (:adjoint, :transpose)
+    @eval function Base.$(op)(operator::StatefulJacobianOperator)
+        return StatefulJacobianOperator($(op)(operator.jac_op), operator.u, operator.p)
+    end
+end
+
+Base.:*(J::StatefulJacobianOperator, v::AbstractArray) = J.jac_op(v, J.u, J.p)
+
+function LinearAlgebra.mul!(
+        Jv::AbstractArray, J::StatefulJacobianOperator, v::AbstractArray)
+    J.jac_op(Jv, v, J.u, J.p)
+    return Jv
+end
+
+"""
+    StatefulJacobianNormalFormOperator(vjp_operator, jvp_operator, cache)
+
+This constructs a Normal Form Jacobian Operator, i.e. it constructs the operator
+corresponding to `JᵀJ` where `J` is the Jacobian Operator. This is not meant to be directly
+constructed, rather it is constructed with `*` on two [`StatefulJacobianOperator`](@ref)s.
+"""
+@concrete mutable struct StatefulJacobianNormalFormOperator{T} <:
+                         AbstractJacobianOperator{T}
+    vjp_operator <: StatefulJacobianOperator{VJP}
+    jvp_operator <: StatefulJacobianOperator{JVP}
+    cache
+end
+
+function Base.size(J::StatefulJacobianNormalFormOperator)
+    return size(J.vjp_operator, 1), size(J.jvp_operator, 2)
+end
+
+function Base.:*(J1::StatefulJacobianOperator{VJP}, J2::StatefulJacobianOperator{JVP})
+    cache = J2 * J2.jac_op.input_cache
+    T = promote_type(eltype(J1), eltype(J2))
+    return StatefulJacobianNormalFormOperator{T}(J1, J2, cache)
+end
+
+function LinearAlgebra.mul!(C::StatefulJacobianNormalFormOperator,
+        A::StatefulJacobianOperator{VJP}, B::StatefulJacobianOperator{JVP})
+    C.vjp_operator = A
+    C.jvp_operator = B
+    return C
+end
+
+function Base.:*(JᵀJ::StatefulJacobianNormalFormOperator, x::AbstractArray)
+    return JᵀJ.vjp_operator * (JᵀJ.jvp_operator * x)
+end
+
+function LinearAlgebra.mul!(
+        JᵀJx::AbstractArray, JᵀJ::StatefulJacobianNormalFormOperator, x::AbstractArray)
+    mul!(JᵀJ.cache, JᵀJ.jvp_operator, x)
+    mul!(JᵀJx, JᵀJ.vjp_operator, JᵀJ.cache)
+    return JᵀJx
+end
+
+# Helper Functions
 prepare_vjp(::Val{true}, args...; kwargs...) = nothing
 
 function prepare_vjp(::Val{false}, prob::AbstractNonlinearProblem,
@@ -163,8 +307,8 @@ function prepare_vjp(::Val{false}, prob::AbstractNonlinearProblem,
             DI.pullback!(fₚ, fu_cache, reshape(vJ, size(u)), autodiff, u, v, di_extras)
         end
     else
-        di_extras = DI.prepare_pullback(f, autodiff, u, fu)
-        return @closure (v, u, p) -> DI.pullback(f, autodiff, u, v, di_extras)
+        di_extras = DI.prepare_pullback(fₚ, autodiff, u, fu)
+        return @closure (v, u, p) -> DI.pullback(fₚ, autodiff, u, v, di_extras)
     end
 end
 
@@ -206,8 +350,8 @@ function prepare_jvp(::Val{false}, prob::AbstractNonlinearProblem,
         fu_cache = copy(fu)
         di_extras = DI.prepare_pushforward(fₚ, fu_cache, autodiff, u, u)
         return @closure (Jv, v, u, p) -> begin
-            DI.pushforward!(fₚ, fu_cache, reshape(Jv, size(fu_cache)), autodiff, u, v,
-                di_extras)
+            DI.pushforward!(
+                fₚ, fu_cache, reshape(Jv, size(fu_cache)), autodiff, u, v, di_extras)
             return
         end
     else
@@ -231,5 +375,7 @@ function prepare_scalar_op(::Val{false}, prob::AbstractNonlinearProblem,
 end
 
 export JacobianOperator, VecJacOperator, JacVecOperator
+export StatefulJacobianOperator
+export StatefulJacobianNormalFormOperator
 
 end

src/NonlinearSolve.jl

@@ -40,6 +40,8 @@ using Preferences: Preferences, @load_preference, @set_preferences!
 using RecursiveArrayTools: recursivecopy!, recursivefill!
 using SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem,
                  AbstractSciMLOperator, _unwrap_val, has_jac, isinplace, NLStats
+using SciMLJacobianOperators: JacobianOperator, VecJacOperator, JacVecOperator,
+                              StatefulJacobianOperator, StatefulJacobianNormalFormOperator
 using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC
 using SparseDiffTools: SparseDiffTools, AbstractSparsityDetection,
                        ApproximateJacobianSparsity, JacPrototypeSparsityDetection,
@@ -72,7 +74,6 @@ include("descent/dogleg.jl")
 include("descent/damped_newton.jl")
 include("descent/geodesic_acceleration.jl")
 
-include("internal/operators.jl")
 include("internal/jacobian.jl")
 include("internal/forward_diff.jl")
 include("internal/linear_solve.jl")

src/globalization/trust_region.jl

@@ -386,11 +386,13 @@ function __internal_init(
     p1, p2, p3, p4 = __get_parameters(T, alg.method)
     ϵ = T(1e-8)
 
+    reverse_ad = get_concrete_reverse_ad(alg.reverse_ad, prob; check_reverse_mode = true)
     vjp_operator = alg.method isa RUS.__Yuan || alg.method isa RUS.__Bastin ?
-                   VecJacOperator(prob, fu, u; autodiff = alg.reverse_ad) : nothing
+                   VecJacOperator(prob, fu, u; autodiff = reverse_ad) : nothing
 
+    forward_ad = get_concrete_forward_ad(alg.forward_ad, prob; check_forward_mode = true)
     jvp_operator = alg.method isa RUS.__Bastin ?
-                   JacVecOperator(prob, fu, u; autodiff = alg.forward_ad) : nothing
+                   JacVecOperator(prob, fu, u; autodiff = forward_ad) : nothing
 
     if alg.method isa RUS.__Yuan
         Jᵀfu_cache = StatefulJacobianOperator(vjp_operator, u, prob.p) * _vec(fu)

src/internal/jacobian.jl

@@ -25,7 +25,8 @@ Construct a cache for the Jacobian of `f` w.r.t. `u`.
   - `jvp_autodiff`: Automatic Differentiation or Finite Differencing backend for computing
     the Jacobian-vector product.
   - `linsolve`: Linear Solver Algorithm used to determine if we need a concrete jacobian
-    or if possible we can just use a [`NonlinearSolve.JacobianOperator`](@ref) instead.
+    or if possible we can just use a [`SciMLJacobianOperators.JacobianOperator`](@ref)
+    instead.
 """
 @concrete mutable struct JacobianCache{iip} <: AbstractNonlinearSolveJacobianCache{iip}
     J
@@ -85,8 +86,7 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
             __similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) :
             copy(f.jac_prototype)
         elseif f.jac_prototype === nothing
-            zero(init_jacobian(
-                jac_cache; preserve_immutable = Val(true)))
+            zero(init_jacobian(jac_cache; preserve_immutable = Val(true)))
         else
             f.jac_prototype
         end
@@ -114,9 +114,9 @@ end
 
 @inline (cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p)
 @inline function (cache::JacobianCache)(::Nothing)
-    J = cache.J
-    J isa JacobianOperator && return StatefulJacobianOperator(J, cache.u, cache.p)
-    return J
+    cache.J isa JacobianOperator &&
+        return StatefulJacobianOperator(cache.J, cache.u, cache.p)
+    return cache.J
 end
 
 function (cache::JacobianCache)(J::JacobianOperator, u, p = cache.p)