Fix elementwise variance accumulation in lsqr (closes #639)

CHANGELOG.md

@@ -17,6 +17,9 @@ Changelog
   and `pylops.waveeqprocessing.Marchenko`
 * Improved typing annotations across all operators (and enforced use of
   Literal for parameters with multiple options)
+* 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.6.0
 * Added `pylops.medical` module and `pylops.medical.CT2D` operator

pylops/optimization/cls_basic.py

@@ -1248,8 +1248,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

@@ -377,6 +377,33 @@ 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):
+    """LSQR ``var`` estimates the diagonal of (A^H A)^-1 element-wise.
+
+    Each unknown must get its own variance; a regression would make all entries
+    identical. scipy's LSQR is used as the reference implementation.
+    """
+    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]
+
+    # the dot(dk, dk) bug produced a constant vector instead (issue #639)
+    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)]
 )