Address review: trim comment, scipy-only test comparison, changelog

- Condense the inline comment in LSQR to a single line.
- Drop the dense diag(inv(A^H A)) comparison from the test and keep scipy's
  LSQR as the reference; LSQR is not guaranteed to match the dense variance.
- Use niter = 2 * nx (scipy's default) and note that niter is only a ceiling.
- Move the issue reference out of the docstring into the targeted comment.
- Add a CHANGELOG.md entry.

Author: gaoflow

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,11 +1248,7 @@ 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:
-            # ``var`` estimates the diagonal of (A^H A)^-1, so it must accumulate
-            # the *element-wise* square of ``dk`` (one variance per unknown), as in
-            # Paige & Saunders (1982) eq. on p.53 and scipy's lsqr. Using
-            # ``dot(dk, dk)`` instead adds the same scalar (sum of squares) to every
-            # entry, yielding identical, wrong variances (issue #639).
+            # accumulate the element-wise square of dk (one variance per unknown)
             self.var = self.var + to_numpy_conditional(
                 self.var, self.ncp.square(self.dk)
             )

pytests/test_solver.py

@@ -379,29 +379,26 @@ def test_lsqr_pylops_scipy(par):
 
 @pytest.mark.parametrize("par", [(par3), (par4)])
 def test_lsqr_calc_var(par):
-    """LSQR ``var`` estimates the diagonal of (A^H A)^-1 element-wise (issue #639).
+    """LSQR ``var`` estimates the diagonal of (A^H A)^-1 element-wise.
 
-    Each unknown must get its own variance; a regression here would make all
-    entries identical (the previous ``dot(dk, dk)`` bug). Compare against scipy's
-    LSQR and the dense ``diag(inv(A^H A))`` on a full-column-rank system.
+    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 = 10 * 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]
-    var_dense = np.diag(np.linalg.inv(np.conj(A.T) @ A))
+    var_sp = sp_lsqr(Aop, y, iter_lim=niter, atol=1e-8, btol=1e-8, calc_var=True)[9]
 
-    # not all identical (the bug produced a constant vector)
+    # 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)
-    assert_array_almost_equal(var, var_dense, decimal=6)
 
 
 @pytest.mark.parametrize(