|
| 1 | +""" |
| 2 | +Resources: |
| 3 | +- https://en.wikipedia.org/wiki/Conjugate_gradient_method |
| 4 | +- https://en.wikipedia.org/wiki/Definite_symmetric_matrix |
| 5 | +""" |
| 6 | +import numpy as np |
| 7 | + |
| 8 | + |
| 9 | +def _is_matrix_spd(matrix: np.array) -> bool: |
| 10 | + """ |
| 11 | + Returns True if input matrix is symmetric positive definite. |
| 12 | + Returns False otherwise. |
| 13 | +
|
| 14 | + For a matrix to be SPD, all eigenvalues must be positive. |
| 15 | +
|
| 16 | + >>> import numpy as np |
| 17 | + >>> matrix = np.array([ |
| 18 | + ... [4.12401784, -5.01453636, -0.63865857], |
| 19 | + ... [-5.01453636, 12.33347422, -3.40493586], |
| 20 | + ... [-0.63865857, -3.40493586, 5.78591885]]) |
| 21 | + >>> _is_matrix_spd(matrix) |
| 22 | + True |
| 23 | + >>> matrix = np.array([ |
| 24 | + ... [0.34634879, 1.96165514, 2.18277744], |
| 25 | + ... [0.74074469, -1.19648894, -1.34223498], |
| 26 | + ... [-0.7687067 , 0.06018373, -1.16315631]]) |
| 27 | + >>> _is_matrix_spd(matrix) |
| 28 | + False |
| 29 | + """ |
| 30 | + # Ensure matrix is square. |
| 31 | + assert np.shape(matrix)[0] == np.shape(matrix)[1] |
| 32 | + |
| 33 | + # If matrix not symmetric, exit right away. |
| 34 | + if np.allclose(matrix, matrix.T) is False: |
| 35 | + return False |
| 36 | + |
| 37 | + # Get eigenvalues and eignevectors for a symmetric matrix. |
| 38 | + eigen_values, _ = np.linalg.eigh(matrix) |
| 39 | + |
| 40 | + # Check sign of all eigenvalues. |
| 41 | + return np.all(eigen_values > 0) |
| 42 | + |
| 43 | + |
| 44 | +def _create_spd_matrix(dimension: np.int64) -> np.array: |
| 45 | + """ |
| 46 | + Returns a symmetric positive definite matrix given a dimension. |
| 47 | +
|
| 48 | + Input: |
| 49 | + dimension gives the square matrix dimension. |
| 50 | +
|
| 51 | + Output: |
| 52 | + spd_matrix is an diminesion x dimensions symmetric positive definite (SPD) matrix. |
| 53 | +
|
| 54 | + >>> import numpy as np |
| 55 | + >>> dimension = 3 |
| 56 | + >>> spd_matrix = _create_spd_matrix(dimension) |
| 57 | + >>> _is_matrix_spd(spd_matrix) |
| 58 | + True |
| 59 | + """ |
| 60 | + random_matrix = np.random.randn(dimension, dimension) |
| 61 | + spd_matrix = np.dot(random_matrix, random_matrix.T) |
| 62 | + assert _is_matrix_spd(spd_matrix) |
| 63 | + return spd_matrix |
| 64 | + |
| 65 | + |
| 66 | +def conjugate_gradient( |
| 67 | + spd_matrix: np.array, |
| 68 | + load_vector: np.array, |
| 69 | + max_iterations: int = 1000, |
| 70 | + tol: float = 1e-8, |
| 71 | +) -> np.array: |
| 72 | + """ |
| 73 | + Returns solution to the linear system np.dot(spd_matrix, x) = b. |
| 74 | +
|
| 75 | + Input: |
| 76 | + spd_matrix is an NxN Symmetric Positive Definite (SPD) matrix. |
| 77 | + load_vector is an Nx1 vector. |
| 78 | +
|
| 79 | + Output: |
| 80 | + x is an Nx1 vector that is the solution vector. |
| 81 | +
|
| 82 | + >>> import numpy as np |
| 83 | + >>> spd_matrix = np.array([ |
| 84 | + ... [8.73256573, -5.02034289, -2.68709226], |
| 85 | + ... [-5.02034289, 3.78188322, 0.91980451], |
| 86 | + ... [-2.68709226, 0.91980451, 1.94746467]]) |
| 87 | + >>> b = np.array([ |
| 88 | + ... [-5.80872761], |
| 89 | + ... [ 3.23807431], |
| 90 | + ... [ 1.95381422]]) |
| 91 | + >>> conjugate_gradient(spd_matrix, b) |
| 92 | + array([[-0.63114139], |
| 93 | + [-0.01561498], |
| 94 | + [ 0.13979294]]) |
| 95 | + """ |
| 96 | + # Ensure proper dimensionality. |
| 97 | + assert np.shape(spd_matrix)[0] == np.shape(spd_matrix)[1] |
| 98 | + assert np.shape(load_vector)[0] == np.shape(spd_matrix)[0] |
| 99 | + assert _is_matrix_spd(spd_matrix) |
| 100 | + |
| 101 | + # Initialize solution guess, residual, search direction. |
| 102 | + x0 = np.zeros((np.shape(load_vector)[0], 1)) |
| 103 | + r0 = np.copy(load_vector) |
| 104 | + p0 = np.copy(r0) |
| 105 | + |
| 106 | + # Set initial errors in solution guess and residual. |
| 107 | + error_residual = 1e9 |
| 108 | + error_x_solution = 1e9 |
| 109 | + error = 1e9 |
| 110 | + |
| 111 | + # Set iteration counter to threshold number of iterations. |
| 112 | + iterations = 0 |
| 113 | + |
| 114 | + while error > tol: |
| 115 | + |
| 116 | + # Save this value so we only calculate the matrix-vector product once. |
| 117 | + w = np.dot(spd_matrix, p0) |
| 118 | + |
| 119 | + # The main algorithm. |
| 120 | + |
| 121 | + # Update search direction magnitude. |
| 122 | + alpha = np.dot(r0.T, r0) / np.dot(p0.T, w) |
| 123 | + # Update solution guess. |
| 124 | + x = x0 + alpha * p0 |
| 125 | + # Calculate new residual. |
| 126 | + r = r0 - alpha * w |
| 127 | + # Calculate new Krylov subspace scale. |
| 128 | + beta = np.dot(r.T, r) / np.dot(r0.T, r0) |
| 129 | + # Calculate new A conjuage search direction. |
| 130 | + p = r + beta * p0 |
| 131 | + |
| 132 | + # Calculate errors. |
| 133 | + error_residual = np.linalg.norm(r - r0) |
| 134 | + error_x_solution = np.linalg.norm(x - x0) |
| 135 | + error = np.maximum(error_residual, error_x_solution) |
| 136 | + |
| 137 | + # Update variables. |
| 138 | + x0 = np.copy(x) |
| 139 | + r0 = np.copy(r) |
| 140 | + p0 = np.copy(p) |
| 141 | + |
| 142 | + # Update number of iterations. |
| 143 | + iterations += 1 |
| 144 | + |
| 145 | + return x |
| 146 | + |
| 147 | + |
| 148 | +def test_conjugate_gradient() -> None: |
| 149 | + """ |
| 150 | + >>> test_conjugate_gradient() # self running tests |
| 151 | + """ |
| 152 | + # Create linear system with SPD matrix and known solution x_true. |
| 153 | + dimension = 3 |
| 154 | + spd_matrix = _create_spd_matrix(dimension) |
| 155 | + x_true = np.random.randn(dimension, 1) |
| 156 | + b = np.dot(spd_matrix, x_true) |
| 157 | + |
| 158 | + # Numpy solution. |
| 159 | + x_numpy = np.linalg.solve(spd_matrix, b) |
| 160 | + |
| 161 | + # Our implementation. |
| 162 | + x_conjugate_gradient = conjugate_gradient(spd_matrix, b) |
| 163 | + |
| 164 | + # Ensure both solutions are close to x_true (and therefore one another). |
| 165 | + assert np.linalg.norm(x_numpy - x_true) <= 1e-6 |
| 166 | + assert np.linalg.norm(x_conjugate_gradient - x_true) <= 1e-6 |
| 167 | + |
| 168 | + |
| 169 | +if __name__ == "__main__": |
| 170 | + import doctest |
| 171 | + |
| 172 | + doctest.testmod() |
| 173 | + test_conjugate_gradient() |
0 commit comments