From f43a52d75eadcbcae46de5903e63dd80bfa50dd7 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 15 Oct 2023 18:40:46 -0400
Subject: [PATCH] Use symmetric linear solve if possible

---
 src/gaussnewton.jl |  8 ++++----
 src/jacobian.jl    | 12 ++++++++----
 src/levenberg.jl   |  8 ++++----
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/src/gaussnewton.jl b/src/gaussnewton.jl
index f8c0c7fbe..55a3fa27f 100644
--- a/src/gaussnewton.jl
+++ b/src/gaussnewton.jl
@@ -97,8 +97,8 @@ function perform_step!(cache::GaussNewtonCache{true})
     __matmul!(Jᵀf, J', fu1)
 
     # u = u - J \ fu
-    linres = dolinsolve(alg.precs, linsolve; A = JᵀJ, b = _vec(Jᵀf), linu = _vec(du),
-        p, reltol = cache.abstol)
+    linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(JᵀJ), b = _vec(Jᵀf),
+        linu = _vec(du), p, reltol = cache.abstol)
     cache.linsolve = linres.cache
     @. u = u - du
     f(cache.fu_new, u, p)
@@ -125,8 +125,8 @@ function perform_step!(cache::GaussNewtonCache{false})
     if linsolve === nothing
         cache.du = fu1 / cache.J
     else
-        linres = dolinsolve(alg.precs, linsolve; A = cache.JᵀJ, b = _vec(cache.Jᵀf),
-            linu = _vec(cache.du), p, reltol = cache.abstol)
+        linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(cache.JᵀJ),
+            b = _vec(cache.Jᵀf), linu = _vec(cache.du), p, reltol = cache.abstol)
         cache.linsolve = linres.cache
     end
     cache.u = @. u - cache.du  # `u` might not support mutation
diff --git a/src/jacobian.jl b/src/jacobian.jl
index 10ac69ec6..58d18f1d4 100644
--- a/src/jacobian.jl
+++ b/src/jacobian.jl
@@ -95,14 +95,14 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p, ::Val{ii
         Jᵀfu = J' * fu
     end
 
-    linprob = LinearProblem(needsJᵀJ ? JᵀJ : J, needsJᵀJ ? _vec(Jᵀfu) : _vec(fu);
-        u0 = _vec(du))
+    linprob = LinearProblem(needsJᵀJ ? __maybe_symmetric(JᵀJ) : J,
+        needsJᵀJ ? _vec(Jᵀfu) : _vec(fu); u0 = _vec(du))
 
     weight = similar(u)
     recursivefill!(weight, true)
 
-    Pl, Pr = wrapprecs(alg.precs(J, nothing, u, p, nothing, nothing, nothing, nothing,
-            nothing)..., weight)
+    Pl, Pr = wrapprecs(alg.precs(needsJᵀJ ? __maybe_symmetric(JᵀJ) : J, nothing, u, p,
+            nothing, nothing, nothing, nothing, nothing)..., weight)
     linsolve = init(linprob, alg.linsolve; alias_A = true, alias_b = true, Pl, Pr,
         linsolve_kwargs...)
 
@@ -119,6 +119,10 @@ __init_JᵀJ(J::Number) = zero(J)
 __init_JᵀJ(J::AbstractArray) = J' * J
 __init_JᵀJ(J::StaticArray) = MArray{Tuple{size(J, 2), size(J, 2)}, eltype(J)}(undef)
 
+__maybe_symmetric(x) = Symmetric(x)
+__maybe_symmetric(x::Number) = x
+__maybe_symmetric(x::SparseArrays.AbstractSparseMatrix) = x
+
 ## Special Handling for Scalars
 function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u::Number, p,
     ::Val{false}; linsolve_with_JᵀJ::Val{needsJᵀJ} = Val(false),
diff --git a/src/levenberg.jl b/src/levenberg.jl
index 8b0e5728e..87384137a 100644
--- a/src/levenberg.jl
+++ b/src/levenberg.jl
@@ -203,8 +203,8 @@ function perform_step!(cache::LevenbergMarquardtCache{true})
     # The following lines do: cache.v = -cache.mat_tmp \ cache.u_tmp
     mul!(cache.u_tmp, J', fu1)
     @. cache.mat_tmp = JᵀJ + λ * DᵀD
-    linres = dolinsolve(alg.precs, linsolve; A = cache.mat_tmp, b = _vec(cache.u_tmp),
-        linu = _vec(cache.du), p = p, reltol = cache.abstol)
+    linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(cache.mat_tmp),
+        b = _vec(cache.u_tmp), linu = _vec(cache.du), p = p, reltol = cache.abstol)
     cache.linsolve = linres.cache
     @. cache.v = -cache.du
 
@@ -280,8 +280,8 @@ function perform_step!(cache::LevenbergMarquardtCache{false})
     if linsolve === nothing
         cache.v = -cache.mat_tmp \ (J' * fu1)
     else
-        linres = dolinsolve(alg.precs, linsolve; A = -cache.mat_tmp, b = _vec(J' * fu1),
-            linu = _vec(cache.v), p, reltol = cache.abstol)
+        linres = dolinsolve(alg.precs, linsolve; A = -__maybe_symmetric(cache.mat_tmp),
+            b = _vec(J' * fu1), linu = _vec(cache.v), p, reltol = cache.abstol)
         cache.linsolve = linres.cache
     end