133 lines
2.8 KiB
Python
133 lines
2.8 KiB
Python
|
import numpy as np
|
|||
|
|
import scipy as sp
|
|||
|
|
|
|||
|
|
def decompose(a, c, e):
|
|||
|
|
"""
|
|||
|
|
Computes the LU decomposition of a tridiagonal matrix.
|
|||
|
|
"""
|
|||
|
|
α = np.zeros(len(c))
|
|||
|
|
β = a.copy()
|
|||
|
|
|
|||
|
|
for i in range(len(c)):
|
|||
|
|
α[i] = e[i] / β[i]
|
|||
|
|
β[i + 1] -= c[i] * α[i]
|
|||
|
|
|
|||
|
|
# Sanity check
|
|||
|
|
if len(a) <= 10:
|
|||
|
|
assert np.allclose(
|
|||
|
|
to_array(a, c, e),
|
|||
|
|
to_array(*from_lower(α)) @ to_array(*from_upper(β, c))
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return (α, β)
|
|||
|
|
|
|||
|
|
def to_csr(a, c, e):
|
|||
|
|
"""
|
|||
|
|
Converts a tridiagonal matrix into a scipy csr sparse matrix.
|
|||
|
|
"""
|
|||
|
|
n = len(c)
|
|||
|
|
|
|||
|
|
values = np.zeros(n * 3 + 1)
|
|||
|
|
values[::3] = a
|
|||
|
|
values[1::3] = c
|
|||
|
|
values[2::3] = e
|
|||
|
|
|
|||
|
|
col_indices = np.zeros_like(values)
|
|||
|
|
col_indices[1::3] = np.arange(1, n + 1)
|
|||
|
|
col_indices[2::3] = np.arange(0, n)
|
|||
|
|
col_indices[3::3] = np.arange(1, n + 1)
|
|||
|
|
|
|||
|
|
index_ptr = np.zeros(n + 2)
|
|||
|
|
index_ptr[1:n+1] = np.arange(2, n * 3 + 2, 3)
|
|||
|
|
index_ptr[n+1] = n * 3 + 1
|
|||
|
|
|
|||
|
|
return sp.sparse.csr_array((values, col_indices, index_ptr))
|
|||
|
|
|
|||
|
|
def to_array(a, c, e):
|
|||
|
|
"""
|
|||
|
|
Converts a tridiagonal matrix into a numpy matrix.
|
|||
|
|
"""
|
|||
|
|
return to_csr(a, c, e).toarray()
|
|||
|
|
|
|||
|
|
def from_lower(α):
|
|||
|
|
"""
|
|||
|
|
Turns the lower vector of a decomposition into a tridiagonal matrix.
|
|||
|
|
|
|||
|
|
Example ussage:
|
|||
|
|
```py
|
|||
|
|
α, β = decompose(m)
|
|||
|
|
print(from_lower(α))
|
|||
|
|
```
|
|||
|
|
"""
|
|||
|
|
return (np.ones(len(α) + 1), np.zeros(len(α)), α)
|
|||
|
|
|
|||
|
|
def from_upper(β, c):
|
|||
|
|
"""
|
|||
|
|
Turns the upper vectors of a decomposition into a tridiagonal matrix.
|
|||
|
|
|
|||
|
|
Example ussage:
|
|||
|
|
```py
|
|||
|
|
α, β = decompose((a, c, e))
|
|||
|
|
print(from_upper(β, c))
|
|||
|
|
```
|
|||
|
|
"""
|
|||
|
|
return (β, c, np.zeros(len(c)))
|
|||
|
|
|
|||
|
|
def solve_lower(α, rhs):
|
|||
|
|
"""
|
|||
|
|
Solve a linear system of equations Mx = v
|
|||
|
|
where M is a lower triangular matrix constructed
|
|||
|
|
by LU decomposing a tridiagonal matrix.
|
|||
|
|
"""
|
|||
|
|
x = np.zeros_like(rhs)
|
|||
|
|
|
|||
|
|
x[0] = rhs[0]
|
|||
|
|
|
|||
|
|
for i in range(1, len(rhs)):
|
|||
|
|
x[i] = rhs[i] - α[i - 1] * x[i - 1]
|
|||
|
|
|
|||
|
|
if len(α) <= 10:
|
|||
|
|
assert np.allclose(to_array(*from_lower(α)) @ x, rhs)
|
|||
|
|
|
|||
|
|
return x
|
|||
|
|
|
|||
|
|
def solve_upper(β, c, rhs):
|
|||
|
|
"""
|
|||
|
|
Solve a linear system of equations Mx = v
|
|||
|
|
where M is an upper triangular matrix constructed
|
|||
|
|
by LU decomposing a tridiagonal matrix.
|
|||
|
|
"""
|
|||
|
|
x = np.zeros_like(rhs)
|
|||
|
|
|
|||
|
|
x[-1] = rhs[-1] / β[-1]
|
|||
|
|
|
|||
|
|
for i in reversed(range(len(rhs) - 1)):
|
|||
|
|
x[i] = (rhs[i] - c[i] * x[i+1]) / β[i]
|
|||
|
|
|
|||
|
|
if len(β) <= 10:
|
|||
|
|
assert np.allclose(to_array(*from_upper(β, c)) @ x, rhs)
|
|||
|
|
|
|||
|
|
return x
|
|||
|
|
|
|||
|
|
def solve(a, c, e, rhs):
|
|||
|
|
α, β = decompose(a, c, e)
|
|||
|
|
x = solve_upper(β, c, solve_lower(α, rhs))
|
|||
|
|
|
|||
|
|
if len(α) <= 10:
|
|||
|
|
assert np.allclose(to_array(a, c, e)@x, rhs)
|
|||
|
|
|
|||
|
|
return x
|
|||
|
|
|
|||
|
|
# Small sanity check for the above code
|
|||
|
|
def main():
|
|||
|
|
a, c, e = (np.ones(4), 2*np.ones(3), 3*np.ones(3))
|
|||
|
|
|
|||
|
|
rhs = np.ones(4)
|
|||
|
|
result = solve(a, c, e, rhs)
|
|||
|
|
print(f"m={to_array(a, c, e)}")
|
|||
|
|
print(f"{rhs=}")
|
|||
|
|
print(f"{result=}")
|
|||
|
|
print(to_array(a, c, e) @ result)
|
|||
|
|
|
|||
|
|
main()
|