From 92983d888709c9a24177620dafe55bc24be82448 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Tue, 21 Feb 2023 15:13:37 +0100
Subject: [PATCH] Update cpr.jl

---
 src/cpr.jl | 30 ++++++++++++++++++------------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/src/cpr.jl b/src/cpr.jl
index 609a50a8..95507bf5 100644
--- a/src/cpr.jl
+++ b/src/cpr.jl
@@ -188,18 +188,8 @@ function apply!(x, cpr::CPRPreconditioner, r, arg...)
         # Apply preconditioner to pressure part
         @tic "p apply" begin
             p_rtol = cpr.p_rtol
-            apply!(Δp, cpr.pressure_precond, r_p)
-            if !isnothing(p_rtol)
-                A_p = cpr.A_p
-                if isnothing(cpr.psolver)
-                    cpr.psolver = FgmresSolver(A_p, r_p)
-                end
-                psolve = cpr.psolver
-                warm_start!(psolve, Δp)
-                M = Jutul.PrecondWrapper(linear_operator(cpr.pressure_precond))
-                fgmres!(psolve, A_p, r_p, M = M, rtol = p_rtol, itmax = 20)
-                @. Δp = psolve.x
-            end
+            p_precond = cpr.pressure_precond
+            cpr_p_apply!(Δp, cpr, p_precond, r_p, p_rtol)
         end
         @tic "r update" correct_residual_for_dp!(y, x, Δp, bz, cpr.buf, cpr.A_ps)
         @tic "s apply" apply!(x, cpr.system_precond, y)
@@ -207,6 +197,22 @@ function apply!(x, cpr::CPRPreconditioner, r, arg...)
     end
 end
 
+
+function cpr_p_apply!(Δp, cpr, p_precond, r_p, p_rtol)
+    apply!(Δp, p_precond, r_p)
+    if !isnothing(p_rtol)
+        A_p = cpr.A_p
+        if isnothing(cpr.psolver)
+            cpr.psolver = FgmresSolver(A_p, r_p)
+        end
+        psolve = cpr.psolver
+        warm_start!(psolve, Δp)
+        M = Jutul.PrecondWrapper(linear_operator(p_precond))
+        fgmres!(psolve, A_p, r_p, M = M, rtol = p_rtol, itmax = 20)
+        @. Δp = psolve.x
+    end
+end
+
 reservoir_residual(lsys) = lsys.r
 reservoir_jacobian(lsys) = lsys.jac