Add tests for batched broyden

Author: avik-pal

.gitignore

@@ -0,0 +1,3 @@
+Manifest.toml
+
+wip

ext/SimpleBatchedNonlinearSolveExt.jl

@@ -1,6 +1,6 @@
 module SimpleBatchedNonlinearSolveExt
 
-using SimpleNonlinearSolve, SciMLBase
+using ArrayInterfaceCore, LinearAlgebra, SimpleNonlinearSolve, SciMLBase
 isdefined(Base, :get_extension) ? (using NNlib) : (using ..NNlib)
 
 _batch_transpose(x) = reshape(x, 1, size(x)...)
@@ -53,15 +53,15 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::Broyden{true}, args...;
     xₙ = x
     xₙ₋₁ = x
     fₙ₋₁ = fₙ
-    for _ in 1:maxiters
-        xₙ = xₙ₋₁ .- _batched_mul(J⁻¹, fₙ₋₁, batch)
+    for i in 1:maxiters
+        xₙ = xₙ₋₁ .- _batched_mul(J⁻¹, fₙ₋₁)
         fₙ = f(xₙ)
         Δ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)
+        J⁻¹Δfₙ = _batched_mul(J⁻¹, Δfₙ)
+        J⁻¹ += _batched_mul(((Δxₙ .- J⁻¹Δfₙ) ./
+                             (_batched_mul(_batch_transpose(Δxₙ), J⁻¹Δfₙ) .+ T(1e-5))),
+                            _batched_mul(_batch_transpose(Δxₙ), J⁻¹))
 
         iszero(fₙ) &&
             return SciMLBase.build_solution(prob, alg, xₙ, fₙ;
@@ -78,5 +78,4 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::Broyden{true}, args...;
     return SciMLBase.build_solution(prob, alg, xₙ, fₙ; retcode = ReturnCode.MaxIters)
 end
 
-
-end
+end

src/SimpleNonlinearSolve.jl

@@ -16,9 +16,7 @@ end
 
 function __init__()
     @static if !isdefined(Base, :get_extension)
-        @require NNlib="872c559c-99b0-510c-b3b7-b6c96a88d5cd" begin
-            include("../ext/SimpleBatchedNonlinearSolveExt.jl")
-        end
+        @require NNlib="872c559c-99b0-510c-b3b7-b6c96a88d5cd" begin include("../ext/SimpleBatchedNonlinearSolveExt.jl") end
     end
 end
 

src/broyden.jl

@@ -3,6 +3,11 @@
 
 A low-overhead implementation of Broyden. This method is non-allocating on scalar
 and static array problems.
+
+!!! note
+
+    To use the `batched` version, remember to load `NNlib`, i.e., `using NNlib` or
+    `import NNlib` must be present in your code.
 """
 struct Broyden{batched} <: AbstractSimpleNonlinearSolveAlgorithm
     Broyden(; batched = false) = new{batched}()

test/basictests.jl

@@ -370,3 +370,15 @@ for options in list_of_options
     sol = solve(probN, alg)
     @test all(abs.(f(u, p)) .< 1e-10)
 end
+
+# Batched Broyden
+using NNlib
+
+f, u0 = (u, p) -> u .* u .- p, randn(1, 3)
+
+p = [2.0 1.0 5.0];
+probN = NonlinearProblem{false}(f, u0, p);
+
+sol = solve(probN, Broyden(batched = true))
+
+@test abs.(sol.u) ≈ sqrt.(p)