Make broyden batched

Author: avik-pal

Project.toml

@@ -1,14 +1,15 @@
 name = "SimpleNonlinearSolve"
 uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 authors = ["SciML"]
-version = "0.1.11"
+version = "0.1.12"
 
 [deps]
 ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
@@ -19,6 +20,7 @@ ArrayInterfaceCore = "0.1.1"
 DiffEqBase = "6.114"
 FiniteDiff = "2"
 ForwardDiff = "0.10.3"
+NNlib = "0.8"
 Reexport = "0.2, 1"
 SciMLBase = "1.73"
 SnoopPrecompile = "1"

src/SimpleNonlinearSolve.jl

@@ -7,6 +7,7 @@ using StaticArraysCore
 using LinearAlgebra
 import ArrayInterfaceCore
 using DiffEqBase
+using NNlib  # Batched Matrix Multiplication
 
 @reexport using SciMLBase
 

src/broyden.jl

@@ -8,15 +8,18 @@ and static array problems.
 """
 struct Broyden <: AbstractSimpleNonlinearSolveAlgorithm end
 
-function SciMLBase.__solve(prob::NonlinearProblem,
-                           alg::Broyden, args...; abstol = nothing,
-                           reltol = nothing,
-                           maxiters = 1000, kwargs...)
+function SciMLBase.__solve(prob::NonlinearProblem, alg::Broyden, args...; abstol = nothing,
+                           reltol = nothing, maxiters = 1000, batch = false, kwargs...)
     f = Base.Fix2(prob.f, prob.p)
     x = float(prob.u0)
+
+    if batch && ndims(x) != 2
+        error("`batch` mode works only if `ndims(prob.u0) == 2`")
+    end
+
     fₙ = f(x)
     T = eltype(x)
-    J⁻¹ = init_J(x)
+    J⁻¹ = init_J(x; batch)
 
     if SciMLBase.isinplace(prob)
         error("Broyden currently only supports out-of-place nonlinear problems")
@@ -30,11 +33,14 @@ function SciMLBase.__solve(prob::NonlinearProblem,
     xₙ₋₁ = x
     fₙ₋₁ = fₙ
     for _ in 1:maxiters
-        xₙ = xₙ₋₁ - J⁻¹ * fₙ₋₁
+        xₙ = xₙ₋₁ .- _batched_mul(J⁻¹, fₙ₋₁, batch)
         fₙ = f(xₙ)
-        Δxₙ = xₙ - xₙ₋₁
-        Δfₙ = fₙ - fₙ₋₁
-        J⁻¹ += ((Δxₙ - J⁻¹ * Δfₙ) ./ (Δxₙ' * J⁻¹ * Δfₙ)) * (Δxₙ' * J⁻¹)
+        Δxₙ = xₙ .- xₙ₋₁
+        Δfₙ = fₙ .- fₙ₋₁
+        J⁻¹Δfₙ = _batched_mul(J⁻¹, Δfₙ, batch)
+        J⁻¹ += _batched_mul(((Δxₙ .- J⁻¹Δfₙ, batch) ./
+                             (_batched_mul(_batch_transpose(Δxₙ, batch), J⁻¹Δfₙ, batch))),
+                            _batched_mul(_batch_transpose(Δxₙ, batch), J⁻¹, batch), batch)
 
         iszero(fₙ) &&
             return SciMLBase.build_solution(prob, alg, xₙ, fₙ;
@@ -50,3 +56,8 @@ function SciMLBase.__solve(prob::NonlinearProblem,
 
     return SciMLBase.build_solution(prob, alg, xₙ, fₙ; retcode = ReturnCode.MaxIters)
 end
+
+function _batch_transpose(x, batch)
+    !batch && return x'
+    return reshape(x, 1, size(x)...)
+end

src/lbroyden.jl

@@ -8,10 +8,9 @@ Base.@kwdef struct LBroyden <: AbstractSimpleNonlinearSolveAlgorithm
     threshold::Int = 27
 end
 
-@views function SciMLBase.__solve(prob::NonlinearProblem,
-                                  alg::LBroyden, args...; abstol = nothing,
-                                  reltol = nothing,
-                                  maxiters = 1000, kwargs...)
+@views function SciMLBase.__solve(prob::NonlinearProblem, alg::LBroyden, args...;
+                                  abstol = nothing, reltol = nothing, maxiters = 1000,
+                                  batch = false, kwargs...)
     threshold = min(maxiters, alg.threshold)
     x = float(prob.u0)
 

src/utils.jl

@@ -34,14 +34,17 @@ value(x) = x
 value(x::Dual) = ForwardDiff.value(x)
 value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x)
 
-function init_J(x)
-    J = ArrayInterfaceCore.zeromatrix(x)
-    if ismutable(x)
-        J[diagind(J)] .= one(eltype(x))
+function init_J(x; batch = false)
+    x_ = batch ? x[:, 1] : x
+
+    J = ArrayInterfaceCore.zeromatrix(x_)
+    if ismutable(x_)
+        J[diagind(J)] .= one(eltype(x_))
     else
         J += I
     end
-    return J
+
+    return batch ? repeat(J, 1, 1, size(x, 2)) : J
 end
 
 function dogleg_method(H, g, Δ)
@@ -68,3 +71,18 @@ function dogleg_method(H, g, Δ)
     tau = (-dot_δsd_δN_δsd + sqrt(fact)) / dot_δN_δsd
     return δsd + tau * δN_δsd
 end
+
+_batched_mul(x, y, batch) = x * y
+function _batched_mul(x::AbstractArray{T, 3}, y::AbstractMatrix, batch) where {T}
+    !batch && return x * y
+    return dropdims(batched_mul(x, reshape(y, size(y, 1), 1, size(y, 2))); dims = 2)
+end
+function _batched_mul(x::AbstractMatrix, y::AbstractArray{T, 3}, batch) where {T}
+    !batch && return x * y
+    return batched_mul(reshape(x, size(x, 1), 1, size(x, 2)), y)
+end
+function _batched_mul(x::AbstractArray{T1, 3}, y::AbstractArray{T2, 3},
+                      batch) where {T1, T2}
+    !batch && return x * y
+    return batched_mul(x, y)
+end