Merge pull request #751 from gaoflow/fix/lsqr-calc-var-elementwise

Fix elementwise variance accumulation in lsqr (closes #639)

Author: mrava87

CHANGELOG.md

@@ -1,6 +1,11 @@
 Changelog
 =========
 
+# 2.8.0
+* Fixed element-wise variance accumulation in `pylops.optimization.cls_basic.LSQR`
+  and `pylops.lsqr`, which previously assigned the same value to every entry of
+  `var` (#639)
+
 # 2.7.0
 * Added cubic spline interpolation operator via
   `pylops.signalprocessing.interpspline.InterpCubicSpline` (also interfaceable via

pylops/optimization/cls_basic.py

@@ -1271,8 +1271,9 @@ def step(self, x: NDArray, show: bool = False) -> NDArray:
             self.ncp.add(self.v, self.w, out=self.w)
         self.ddnorm = self.ddnorm + self.ncp.linalg.norm(self.dk) ** 2.0
         if self.calc_var:
+            # accumulate the element-wise square of dk (one variance per unknown)
             self.var = self.var + to_numpy_conditional(
-                self.var, self.ncp.dot(self.dk, self.dk)
+                self.var, self.ncp.square(self.dk)
             )
 
         # use a plane rotation on the right to eliminate the

pytests/test_solver.py

@@ -410,6 +410,29 @@ def test_lsqr_pylops_scipy(par):
         assert_array_almost_equal(xinv_sp, x, decimal=4)
 
 
+@pytest.mark.skipif(
+    int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 1, reason="Not CuPy enabled"
+)
+@pytest.mark.parametrize("par", [(par3), (par4)])
+def test_lsqr_calc_var(par):
+    """Compare PyLops and scipy LSQR variance computation for the diagonal of
+    (A^H A)^-1 (issue #639)."""
+    np.random.seed(10)
+
+    A = np.random.normal(0, 1, (par["ny"], par["nx"]))
+    Aop = MatrixMult(A, dtype=par["dtype"])
+    y = Aop * (np.ones(par["nx"]))
+
+    # niter is only a ceiling (atol/btol stop earlier); use scipy's default of
+    # 2 * nx so both implementations run for the same number of iterations.
+    niter = 2 * par["nx"]
+    var = lsqr(Aop, y, x0=None, niter=niter, atol=1e-8, btol=1e-8)[9]
+    var_sp = sp_lsqr(Aop, y, iter_lim=niter, atol=1e-8, btol=1e-8, calc_var=True)[9]
+
+    assert not np.allclose(var, var[0])
+    assert_array_almost_equal(var, var_sp, decimal=6)
+
+
 @pytest.mark.parametrize(
     "par", [(par1), (par2), (par3), (par4), (par1j), (par2j), (par3j), (par3j)]
 )