test_lsqr.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153
  1. import numpy as np
  2. from numpy.testing import assert_allclose, assert_array_equal, assert_equal
  3. import pytest
  4. import scipy.sparse
  5. import scipy.sparse.linalg
  6. from scipy.sparse.linalg import lsqr
  7. from time import time
  8. # Set up a test problem
  9. n = 35
  10. G = np.eye(n)
  11. normal = np.random.normal
  12. norm = np.linalg.norm
  13. for jj in range(5):
  14. gg = normal(size=n)
  15. hh = gg * gg.T
  16. G += (hh + hh.T) * 0.5
  17. G += normal(size=n) * normal(size=n)
  18. b = normal(size=n)
  19. # tolerance for atol/btol keywords of lsqr()
  20. tol = 2e-10
  21. # tolerances for testing the results of the lsqr() call with assert_allclose
  22. # These tolerances are a bit fragile - see discussion in gh-15301.
  23. atol_test = 4e-10
  24. rtol_test = 2e-8
  25. show = False
  26. maxit = None
  27. def test_lsqr_basic():
  28. b_copy = b.copy()
  29. xo, *_ = lsqr(G, b, show=show, atol=tol, btol=tol, iter_lim=maxit)
  30. assert_array_equal(b_copy, b)
  31. svx = np.linalg.solve(G, b)
  32. assert_allclose(xo, svx, atol=atol_test, rtol=rtol_test)
  33. # Now the same but with damp > 0.
  34. # This is equivalent to solving the extented system:
  35. # ( G ) @ x = ( b )
  36. # ( damp*I ) ( 0 )
  37. damp = 1.5
  38. xo, *_ = lsqr(
  39. G, b, damp=damp, show=show, atol=tol, btol=tol, iter_lim=maxit)
  40. Gext = np.r_[G, damp * np.eye(G.shape[1])]
  41. bext = np.r_[b, np.zeros(G.shape[1])]
  42. svx, *_ = np.linalg.lstsq(Gext, bext, rcond=None)
  43. assert_allclose(xo, svx, atol=atol_test, rtol=rtol_test)
  44. def test_gh_2466():
  45. row = np.array([0, 0])
  46. col = np.array([0, 1])
  47. val = np.array([1, -1])
  48. A = scipy.sparse.coo_matrix((val, (row, col)), shape=(1, 2))
  49. b = np.asarray([4])
  50. lsqr(A, b)
  51. def test_well_conditioned_problems():
  52. # Test that sparse the lsqr solver returns the right solution
  53. # on various problems with different random seeds.
  54. # This is a non-regression test for a potential ZeroDivisionError
  55. # raised when computing the `test2` & `test3` convergence conditions.
  56. n = 10
  57. A_sparse = scipy.sparse.eye(n, n)
  58. A_dense = A_sparse.toarray()
  59. with np.errstate(invalid='raise'):
  60. for seed in range(30):
  61. rng = np.random.RandomState(seed + 10)
  62. beta = rng.rand(n)
  63. beta[beta == 0] = 0.00001 # ensure that all the betas are not null
  64. b = A_sparse @ beta[:, np.newaxis]
  65. output = lsqr(A_sparse, b, show=show)
  66. # Check that the termination condition corresponds to an approximate
  67. # solution to Ax = b
  68. assert_equal(output[1], 1)
  69. solution = output[0]
  70. # Check that we recover the ground truth solution
  71. assert_allclose(solution, beta)
  72. # Sanity check: compare to the dense array solver
  73. reference_solution = np.linalg.solve(A_dense, b).ravel()
  74. assert_allclose(solution, reference_solution)
  75. def test_b_shapes():
  76. # Test b being a scalar.
  77. A = np.array([[1.0, 2.0]])
  78. b = 3.0
  79. x = lsqr(A, b)[0]
  80. assert norm(A.dot(x) - b) == pytest.approx(0)
  81. # Test b being a column vector.
  82. A = np.eye(10)
  83. b = np.ones((10, 1))
  84. x = lsqr(A, b)[0]
  85. assert norm(A.dot(x) - b.ravel()) == pytest.approx(0)
  86. def test_initialization():
  87. # Test the default setting is the same as zeros
  88. b_copy = b.copy()
  89. x_ref = lsqr(G, b, show=show, atol=tol, btol=tol, iter_lim=maxit)
  90. x0 = np.zeros(x_ref[0].shape)
  91. x = lsqr(G, b, show=show, atol=tol, btol=tol, iter_lim=maxit, x0=x0)
  92. assert_array_equal(b_copy, b)
  93. assert_allclose(x_ref[0], x[0])
  94. # Test warm-start with single iteration
  95. x0 = lsqr(G, b, show=show, atol=tol, btol=tol, iter_lim=1)[0]
  96. x = lsqr(G, b, show=show, atol=tol, btol=tol, iter_lim=maxit, x0=x0)
  97. assert_allclose(x_ref[0], x[0])
  98. assert_array_equal(b_copy, b)
  99. if __name__ == "__main__":
  100. svx = np.linalg.solve(G, b)
  101. tic = time()
  102. X = lsqr(G, b, show=show, atol=tol, btol=tol, iter_lim=maxit)
  103. xo = X[0]
  104. phio = X[3]
  105. psio = X[7]
  106. k = X[2]
  107. chio = X[8]
  108. mg = np.amax(G - G.T)
  109. if mg > 1e-14:
  110. sym = 'No'
  111. else:
  112. sym = 'Yes'
  113. print('LSQR')
  114. print("Is linear operator symmetric? " + sym)
  115. print("n: %3g iterations: %3g" % (n, k))
  116. print("Norms computed in %.2fs by LSQR" % (time() - tic))
  117. print(" ||x|| %9.4e ||r|| %9.4e ||Ar|| %9.4e " % (chio, phio, psio))
  118. print("Residual norms computed directly:")
  119. print(" ||x|| %9.4e ||r|| %9.4e ||Ar|| %9.4e" % (norm(xo),
  120. norm(G*xo - b),
  121. norm(G.T*(G*xo-b))))
  122. print("Direct solution norms:")
  123. print(" ||x|| %9.4e ||r|| %9.4e " % (norm(svx), norm(G*svx - b)))
  124. print("")
  125. print(" || x_{direct} - x_{LSQR}|| %9.4e " % norm(svx-xo))
  126. print("")