Merge pull request #266 from yonatanwesen/yd/pt-caching

Removing Allocation for Inexact Jacobian

Project.toml

@@ -9,6 +9,7 @@ ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
+FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
@@ -35,13 +36,14 @@ NonlinearSolveFastLevenbergMarquardtExt = "FastLevenbergMarquardt"
 NonlinearSolveLeastSquaresOptimExt = "LeastSquaresOptim"
 
 [compat]
-BandedMatrices = "1"
 ADTypes = "0.2"
 ArrayInterface = "6.0.24, 7"
+BandedMatrices = "1"
 ConcreteStructs = "0.2"
 DiffEqBase = "6.130"
 EnumX = "1"
 Enzyme = "0.11"
+FastBroadcast = "0.1.9, 0.2"
 FastLevenbergMarquardt = "0.1"
 FiniteDiff = "2"
 ForwardDiff = "0.10.3"
@@ -63,6 +65,7 @@ julia = "1.9"
 [extras]
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -79,7 +82,6 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
-DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 
 [targets]
 test = ["Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase"]

src/NonlinearSolve.jl

@@ -9,11 +9,12 @@ import PrecompileTools
 
 PrecompileTools.@recompile_invalidations begin
     using DiffEqBase, LinearAlgebra, LinearSolve, SparseArrays, SparseDiffTools
+    using FastBroadcast: @..
     import ArrayInterface: restructure
 
     import ADTypes: AbstractFiniteDifferencesMode
     import ArrayInterface: undefmatrix,
-        matrix_colors, parameterless_type, ismutable, issingular
+        matrix_colors, parameterless_type, ismutable, issingular,fast_scalar_indexing
     import ConcreteStructs: @concrete
     import EnumX: @enumx
     import ForwardDiff

src/pseudotransient.jl

@@ -114,12 +114,25 @@ end
 function perform_step!(cache::PseudoTransientCache{true})
     @unpack u, u_prev, fu1, f, p, alg, J, linsolve, du, alpha, tc_storage = cache
     jacobian!!(J, cache)
-    J_new = J - (1 / alpha) * I
+    inv_alpha = inv(alpha)
+
+    if J isa SciMLBase.AbstractSciMLOperator
+        J = J - inv_alpha * I
+    else
+        idxs = diagind(J)
+        if fast_scalar_indexing(J)
+            @inbounds for i in axes(J, 1)
+                J[i, i] = J[i, i] - inv_alpha
+            end
+        else
+            @.. broadcast=false @view(J[idxs])=@view(J[idxs]) - inv_alpha
+        end
+    end
 
     termination_condition = cache.termination_condition(tc_storage)
 
     # u = u - J \ fu
-    linres = dolinsolve(alg.precs, linsolve; A = J_new, b = _vec(fu1), linu = _vec(du),
+    linres = dolinsolve(alg.precs, linsolve; A = J, b = _vec(fu1), linu = _vec(du),
         p, reltol = cache.abstol)
     cache.linsolve = linres.cache
     @. u = u - du
@@ -147,13 +160,14 @@ function perform_step!(cache::PseudoTransientCache{false})
     termination_condition = cache.termination_condition(tc_storage)
 
     cache.J = jacobian!!(cache.J, cache)
+    inv_alpha = inv(alpha)
+
+    cache.J = cache.J - inv_alpha * I
     # u = u - J \ fu
     if linsolve === nothing
-        cache.du = fu1 / (cache.J - (1 / alpha) * I)
+        cache.du = fu1 / cache.J
     else
-        linres = dolinsolve(alg.precs, linsolve; A = cache.J - (1 / alpha) * I,
-            b = _vec(fu1),
-            linu = _vec(cache.du), p, reltol = cache.abstol)
+        linres = dolinsolve(alg.precs, linsolve; A = cache.J,b = _vec(fu1),linu = _vec(cache.du), p, reltol = cache.abstol)
         cache.linsolve = linres.cache
     end
     cache.u = @. u - cache.du  # `u` might not support mutation