Skip to content

Commit 40344e3

Browse files
committed
refactor: Move linear algebra algorithms outside of src/
1 parent 0b6f4f6 commit 40344e3

8 files changed

+1289
-0
lines changed

linear_algebra/conjugate_gradient.py

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

0 commit comments

Comments
 (0)