Count statistics inside calls and not in individual algorithms

avik-pal · avik-pal · commit ee042973bf6f · 2023-12-05T13:49:30.000-05:00
diff --git a/src/gaussnewton.jl b/src/gaussnewton.jl
@@ -123,7 +123,7 @@ function perform_step!(cache::GaussNewtonCache{iip}) where {iip}
         A, b = cache.J, _vec(cache.fu)
     end
 
-    linres = dolinsolve(cache.alg.precs, cache.linsolve; A, b, linu = _vec(cache.du),
+    linres = dolinsolve(cache, cache.alg.precs, cache.linsolve; A, b, linu = _vec(cache.du),
         cache.p, reltol = cache.abstol)
     cache.linsolve = linres.cache
     cache.du = _restructure(cache.du, linres.u)
@@ -142,9 +142,6 @@ function perform_step!(cache::GaussNewtonCache{iip}) where {iip}
     @bb copyto!(cache.u_cache, cache.u)
     @bb copyto!(cache.dfu, cache.fu)
 
-    cache.stats.njacs += 1
-    cache.stats.nsolve += 1
-    cache.stats.nfactors += 1
     return nothing
 end
 
diff --git a/src/klement.jl b/src/klement.jl
@@ -123,7 +123,7 @@ function perform_step!(cache::GeneralKlementCache{iip}) where {iip}
     A = ifelse(cache.J isa SMatrix || cache.J isa Number || !fact_done, cache.J, nothing)
 
     # u = u - J \ fu
-    linres = dolinsolve(alg.precs, cache.linsolve; A,
+    linres = dolinsolve(cache, alg.precs, cache.linsolve; A,
         b = _vec(cache.fu), linu = _vec(cache.du), cache.p, reltol = cache.abstol)
     cache.linsolve = linres.cache
     cache.du = _restructure(cache.du, linres.u)
@@ -139,9 +139,6 @@ function perform_step!(cache::GeneralKlementCache{iip}) where {iip}
 
     @bb copyto!(cache.u_cache, cache.u)
 
-    cache.stats.nsolve += 1
-    cache.stats.nfactors += 1
-
     cache.force_stop && return nothing
 
     # Update the Jacobian
diff --git a/src/levenberg.jl b/src/levenberg.jl
@@ -263,13 +263,14 @@ function perform_step!(cache::LevenbergMarquardtCache{iip, fastls}) where {iip,
         else
             cache.rhs_tmp = _vcat(_vec(cache.fu), zero(_vec(cache.u)))
         end
-        linres = dolinsolve(alg.precs, linsolve; A = cache.mat_tmp,
+        linres = dolinsolve(cache, alg.precs, linsolve; A = cache.mat_tmp,
             b = cache.rhs_tmp, linu = _vec(cache.v), cache.p, reltol = cache.abstol)
     else
         @bb cache.u_cache_2 = transpose(cache.J) × cache.fu
         @bb @. cache.mat_tmp = cache.JᵀJ + cache.λ * cache.DᵀD
-        linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(cache.mat_tmp),
-            b = _vec(cache.u_cache_2), linu = _vec(cache.v), cache.p, reltol = cache.abstol)
+        linres = dolinsolve(cache, alg.precs, linsolve;
+            A = __maybe_symmetric(cache.mat_tmp), b = _vec(cache.u_cache_2),
+            linu = _vec(cache.v), cache.p, reltol = cache.abstol)
     end
     cache.linsolve = linres.cache
     linu = _restructure(cache.v, linres.u)
@@ -293,20 +294,17 @@ function perform_step!(cache::LevenbergMarquardtCache{iip, fastls}) where {iip,
         else
             cache.rhs_tmp = _vcat(_vec(cache.fu_cache_2), zero(_vec(cache.u)))
         end
-        linres = dolinsolve(alg.precs, linsolve; b = cache.rhs_tmp, linu = _vec(cache.a),
-            cache.p, reltol = cache.abstol)
+        linres = dolinsolve(cache, alg.precs, linsolve; b = cache.rhs_tmp,
+            linu = _vec(cache.a), cache.p, reltol = cache.abstol)
     else
         @bb cache.u_cache_2 = transpose(cache.J) × cache.fu_cache_2
-        linres = dolinsolve(alg.precs, linsolve; b = _vec(cache.u_cache_2),
+        linres = dolinsolve(cache, alg.precs, linsolve; b = _vec(cache.u_cache_2),
             linu = _vec(cache.a), cache.p, reltol = cache.abstol)
     end
     cache.linsolve = linres.cache
     linu = _restructure(cache.a, linres.u)
     @bb @. cache.a = -linu
 
-    cache.stats.nsolve += 2
-    cache.stats.nfactors += 2
-
     # Require acceptable steps to satisfy the following condition.
     norm_v = cache.internalnorm(cache.v)
     if 2 * cache.internalnorm(cache.a) ≤ cache.α_geodesic * norm_v
diff --git a/src/pseudotransient.jl b/src/pseudotransient.jl
@@ -127,7 +127,7 @@ function perform_step!(cache::PseudoTransientCache{iip}) where {iip}
     end
 
     # u = u - J \ fu
-    linres = dolinsolve(alg.precs, cache.linsolve; A, b = _vec(cache.fu),
+    linres = dolinsolve(cache, alg.precs, cache.linsolve; A, b = _vec(cache.fu),
         linu = _vec(cache.du), cache.p, reltol = cache.abstol)
     cache.linsolve = linres.cache
     cache.du = _restructure(cache.du, linres.u)
@@ -145,9 +145,6 @@ function perform_step!(cache::PseudoTransientCache{iip}) where {iip}
     check_and_update!(cache, cache.fu, cache.u, cache.u_cache)
 
     @bb copyto!(cache.u_cache, cache.u)
-    cache.stats.njacs += 1
-    cache.stats.nsolve += 1
-    cache.stats.nfactors += 1
     return nothing
 end
 
diff --git a/src/raphson.jl b/src/raphson.jl
@@ -104,7 +104,7 @@ function perform_step!(cache::NewtonRaphsonCache{iip}) where {iip}
     cache.J = jacobian!!(cache.J, cache)
 
     # u = u - J \ fu
-    linres = dolinsolve(alg.precs, cache.linsolve; A = cache.J, b = _vec(cache.fu),
+    linres = dolinsolve(cache, alg.precs, cache.linsolve; A = cache.J, b = _vec(cache.fu),
         linu = _vec(cache.du), cache.p, reltol = cache.abstol)
     cache.linsolve = linres.cache
     cache.du = _restructure(cache.du, linres.u)
@@ -119,8 +119,5 @@ function perform_step!(cache::NewtonRaphsonCache{iip}) where {iip}
     check_and_update!(cache, cache.fu, cache.u, cache.u_cache)
 
     @bb copyto!(cache.u_cache, cache.u)
-    cache.stats.njacs += 1
-    cache.stats.nsolve += 1
-    cache.stats.nfactors += 1
     return nothing
 end
diff --git a/src/trustRegion.jl b/src/trustRegion.jl
@@ -360,7 +360,7 @@ function perform_step!(cache::TrustRegionCache{iip}) where {iip}
 
         # do not use A = cache.H, b = _vec(cache.g) since it is equivalent
         # to  A = cache.J, b = _vec(fu) as long as the Jacobian is non-singular
-        linres = dolinsolve(cache.alg.precs, cache.linsolve, A = cache.J,
+        linres = dolinsolve(cache, cache.alg.precs, cache.linsolve, A = cache.J,
             b = _vec(cache.fu), linu = _vec(cache.u_gauss_newton), p = cache.p,
             reltol = cache.abstol)
         cache.linsolve = linres.cache
@@ -375,8 +375,6 @@ function perform_step!(cache::TrustRegionCache{iip}) where {iip}
     @bb @. cache.u_cache_2 = cache.u + cache.du
     evaluate_f(cache, cache.u_cache_2, cache.p, Val{:fu_cache_2}())
     trust_region_step!(cache)
-    cache.stats.nsolve += 1
-    cache.stats.nfactors += 1
     return nothing
 end
 
diff --git a/src/utils.jl b/src/utils.jl
@@ -88,18 +88,26 @@ end
 
 DEFAULT_PRECS(W, du, u, p, t, newW, Plprev, Prprev, cachedata) = nothing, nothing
 
-function dolinsolve(precs::P, linsolve::FakeLinearSolveJLCache; A = nothing,
+function dolinsolve(cache, precs::P, linsolve::FakeLinearSolveJLCache; A = nothing,
         linu = nothing, b = nothing, du = nothing, p = nothing, weight = nothing,
         cachedata = nothing, reltol = nothing, reuse_A_if_factorization = false) where {P}
+    # Update Statistics
+    cache.stats.nsolve += 1
+    cache.stats.nfactors += !(A isa Number)
+
     A !== nothing && (linsolve.A = A)
     b !== nothing && (linsolve.b = b)
     linres = linsolve.A \ linsolve.b
     return FakeLinearSolveJLResult(linsolve, linres)
 end
 
-function dolinsolve(precs::P, linsolve; A = nothing, linu = nothing, b = nothing,
+function dolinsolve(cache, precs::P, linsolve; A = nothing, linu = nothing, b = nothing,
         du = nothing, p = nothing, weight = nothing, cachedata = nothing, reltol = nothing,
         reuse_A_if_factorization = false) where {P}
+    # Update Statistics
+    cache.stats.nsolve += 1
+    cache.stats.nfactors += 1
+
     # Some Algorithms would reuse factorization but it causes the cache to not reset in
     # certain cases
     if A !== nothing
@@ -108,10 +116,16 @@ function dolinsolve(precs::P, linsolve; A = nothing, linu = nothing, b = nothing
            (alg isa LinearSolve.DefaultLinearSolver && !(alg ==
               LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.KrylovJL_GMRES)))
             # Factorization Algorithm
-            !reuse_A_if_factorization && (linsolve.A = A)
+            if reuse_A_if_factorization
+                cache.stats.nfactors -= 1
+            else
+                linsolve.A = A
+            end
         else
             linsolve.A = A
         end
+    else
+        cache.stats.nfactors -= 1
     end
     b !== nothing && (linsolve.b = b)
     linu !== nothing && (linsolve.u = linu)
