How to monitor the progress of LinearSolve?

Question

I am trying to do something like

d = RandomReal[{0, 50}, 3];
U = ResourceFunction["RandomMatrix"]["Orthogonal", Real, 3];
m = Transpose[U] . DiagonalMatrix[d] . U;
b = {6, 4, 5};
LinearSolve[N@m, N@b, 
 Method -> {"Krylov", Method -> "ConjugateGradient"}]

and want to monitor the decline of residue in the solving process. It seems that Solve and LinearSolve do not have the StepMonitor and EvaluationMonitor options, which are only available for iterative built-in functions.

I'm guessing blindly that, LinearSolve is not eligible for these options because even when an iterative method is specified, the iterating procedure is handled by low-level linalg libraries, rather than within the Wolfram language abstraction. If this is the case, it seems also to rule out the possibility of achieving this with Sow and Reap.

Is there some other way I have not think of that may do the trick?

Related posts:

1. Monitoring the execution of Solve - @StarBucK
2. What does LinearSolve Preconditioner do? - Answer by @Henrik Schumacher

---

Indirect workaround: Since this is not a real world problem, one can always use CG FindMinimum on the corresponding variational minimization problem, and StepMonitor the process. I however still appreciate further insights on the internals of LinearSolve. From the second linked post I understand that there is internally the function SparseArray`KrylovLinearSolve; is it possible to probe its implementation?

Answer 1

I think I got it. MSE rules and this should be obvious had I been competent.
This is what @unlikely did in his question

Reap@LinearSolve[m, b, Method -> {"Krylov", 
    "ResidualNormFunction" -> ((Sow[#, "residue_vector"]; Sow[Norm[#, 2]]) &)
  }
]

This should intuitively suffice the purpose in question, however does not.

---

I finished my mock-up that I think is mathematically equivalent to @user293787 's

ClearAll[Gravifer`myKrylovLinearSolve]
Module[{symbol = Gravifer`myKrylovLinearSolve, $symbol = 
    SparseArray`KrylovLinearSolve},
  ((MessageName[Evaluate@symbol, #] = 
       MessageName[Evaluate@$symbol, #]) &) /@ ToExpression[
     FindList[
      FileNameJoin[{System`Private`$MessagesDir, $Language, 
        "Messages.m"}], ToString[$symbol]]
     , StandardForm, HoldForm][[All, 1, 1, -1]]
  ];
Options[Gravifer`myKrylovLinearSolve] ^= 
  Options[SparseArray`KrylovLinearSolve];
SyntaxInformation[
   Gravifer`myKrylovLinearSolve] ^= {"ArgumentsPattern" -> {_, _, 
     OptionsPattern[]}, 
   "OptionNames" -> 
    Union[First /@ Options[Gravifer`myKrylovLinearSolve], 
     First /@ Options[SparseArray`KrylovLinearSolve]]};
Gravifer`myKrylovLinearSolve[m_, b_, 
  opt : OptionsPattern[SparseArray`KrylovLinearSolve]] := 
 Module[{dim = Dimensions[m][[1]], 
    maxiter = OptionValue[MaxIterations], tol = OptionValue[Tolerance],
    x0 = OptionValue["StartingVector"], 
    resnorm = OptionValue["ResidualNormFunction"] /. Automatic -> Norm,
    CGfunc},
   
   CGfunc = 
    Function[{n, A, \[FormalB], \[FormalCapitalN], \[Epsilon], x0}, 
     Module[{x = x0, r = \[FormalB] - A . x0, 
       p = \[FormalB] - A . x0, \[Alpha], \[Beta]},
      Do[If[resnorm[r] < \[Epsilon], Break[]]; 
       If[p . A . p == 0, 
        Message[Gravifer`myKrylovLinearSolve::krystg]; Break[]];
       \[Alpha] = Norm[r]^2/p . A . p; 
       x += \[Alpha] p; \[Beta] = Norm[r - A . (\[Alpha] p)]^2/
        Norm[r]^2; r -= A . (\[Alpha] p); p = \[Beta] p + r,
       \[FormalCapitalN]]; 
      If[resnorm[r] > \[Epsilon], 
       Message[Gravifer`myKrylovLinearSolve::krymit, \
\[FormalCapitalN]]]; x
      ]];
   If[Precision[{m, b}] == \[Infinity], 
    Message[Gravifer`myKrylovLinearSolve::kryinfp]];
   (*
   If[Head[m]==SparseArray\[And]m["Background"]!=0,Message[
   Gravifer`myKrylovLinearSolve::kryinz,"m"]];
   If[Head[b]==SparseArray\[And]b["Background"]!=0,Message[
   Gravifer`myKrylovLinearSolve::kryinz,"b"]];
   *)
   Switch[OptionValue[Method],
    Automatic | "ConjugateGradient", 
    CGfunc[dim, m, b, maxiter /. Automatic -> n, 
     tol /. Automatic -> 10^-12, 
     x0 /. Automatic -> ConstantArray[0, dim]]
    ]
   ] /; And[CheckArguments[Gravifer`myKrylovLinearSolve[m, b, opt], 2],
   (If[MatrixQ[m], True, 
     Message[Gravifer`myKrylovLinearSolve::matrix, m, 1];
     If[MatrixQ[Evaluate@m], True, 
      Message[Gravifer`myKrylovLinearSolve::krynfa, m, 1]]; False]),
   (VectorQ[b] \[Or] 
     If[MatrixQ[b], True, 
      Message[Gravifer`myKrylovLinearSolve::matrix, b, 2]; False]),
   ((OptionValue[MaxIterations] == Automatic) \[Or] 
     If[IntegerQ@OptionValue[MaxIterations], True, 
      Message[Gravifer`myKrylovLinearSolve::kryit, 
       OptionValue[MaxIterations]]; False]),
   ((OptionValue["StartingVector"] == Automatic) \[Or] 
     If[VectorQ[OptionValue["StartingVector"], NumberQ], True, 
      Message[Gravifer`myKrylovLinearSolve::kryivec]; False]),
   Switch[OptionValue[Method],
    Automatic | "ConjugateGradient", True,
    "SteepestDescent" | "Lanczos" | "Arnoldi" | "MinRES" | "GMRES" | 
     "CGR" | "CGNR" | "CGLE" | "QMR" | "BiCG" | "BiCGStab", 
    Message[Gravifer`myKrylovLinearSolve::krynimp]; False,
    _, Message[Gravifer`myKrylovLinearSolve::kryme, 
     OptionValue["Method"]]; False
    ]
   ]

On the cases from his answer,

SeedRandom[1];
dim = 2000;
U = RandomVariate[CircularRealMatrixDistribution[dim]];
A = Transpose[U] . DiagonalMatrix[RandomReal[{1/100, 1}, dim]] . U;
b = RandomReal[{-1, 1}, dim];
{sol, reap} = 
  LinearSolve[A, b, 
    Method -> {"Krylov", Method -> "ConjugateGradient", 
      "ResidualNormFunction" -> ((Sow[Norm[#]]) &)}] // Reap;
{mysol, myreap} = 
 Gravifer`myKrylovLinearSolve[A, b, Method -> "ConjugateGradient", 
   "ResidualNormFunction" -> ((Sow[Norm[#]]) &)] // 
  Reap; 
ListLinePlot[{reap[[1]], myreap[[1]]}, 
 ScalingFunctions -> "Log"]
{sol, reap} = 
  LinearSolve[N@A, N@b, 
    Method -> {"Krylov", Method -> "ConjugateGradient", 
      "ResidualNormFunction" -> ((Sow[Norm[#]]) &)}] // Reap;
{mysol, myreap} = 
 Gravifer`myKrylovLinearSolve[N@A, N@b, Method -> "ConjugateGradient",
    "ResidualNormFunction" -> ((Sow[Norm[#]]) &)] // 
  Reap; 
ListLinePlot[{reap[[1]], myreap[[1]]}, 
 ScalingFunctions -> "Log"]

The results are

So maybe as @unlikely indeed suggested in his post, OptionValue["ResidualNormFunction"] is not called once per iteration for the built-in function.

Answer 2

Not directly an answer, but here is an implementation of conjugate gradient for a real, symmetric and positive definite matrix $A$. It Sows the norm of the residual at each step. This implementation uses more matrix-vector multiplications than necessary:

linearSolveConjugateGradient[A_,b_,tol_:10^(-10)] := Module[{p,x,r},
    p = {};
    x = 0.*b;
    Do[
       r = A.x-b;
       If[Sow[Norm[r]]<tol,Break[]];
       p = {r,First[p,Nothing]} // Developer`ToPackedArray;
       p[[1]] = LeastSquares[p.A.Transpose[p],-p.r].p;
       x += p[[1]];
    ,{Infinity}];
    x
];

Example. Create a large (whatever that means) matrix $A$ with condition number about 100:

SeedRandom[1];
dim = 2000;
U = RandomVariate[CircularRealMatrixDistribution[dim]];
A = Transpose[U].DiagonalMatrix[RandomReal[{1/100,1},dim]].U;
b = RandomReal[{-1,1},dim];

Monitor progress:

ListLogPlot[Last[Reap[linearSolveConjugateGradient[A,b]]],
            AxesLabel->{"step","Norm[residual]"}]

In this particular example, it is faster than the built-in conjugate gradient solver:

Norm[A.linearSolveConjugateGradient[A,b]-b] // AbsoluteTiming
(* {0.606085, 8.21629*10^-11} *)

Norm[A.LinearSolve[A,b,Method->{"Krylov",Method->"ConjugateGradient"}]-b] // AbsoluteTiming
(* {11.5928, 2.43374*10^-10} *)