aboutsummaryrefslogtreecommitdiff
path: root/solvers.py
blob: 38ac247c5040bc2e4f27cc3782183e97b24bd161 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
'''Routines for solving a linear system of equations.'''

import numpy as np
from numpy.typing import NDArray

def lu_factorization(
        mm: NDArray[np.float_],
        tolerance: float = 1e-6
) -> tuple[NDArray[np.float_], NDArray[np.float_], NDArray[np.float_]]:
    '''Computes the LUP factorization of a matrix

    Args:
        mm: Matrix to factorize
        tolerance: Greatest absolute value considered as 0

    Returns:
        ll: Lower triangular matrix
        uu: Upper triangular matrix
        pp: Permutation matrix
    '''
    nn = mm.shape[0]

    uu = mm.copy()
    ll = np.zeros((nn, nn), dtype=np.float_)
    pp = np.eye(nn, dtype=np.float_)

    i = 0
    j = 0

    while i < nn and j < nn:
        pivot_row = i + int(np.argmax(np.abs(uu[i:, j])))

        if np.abs(uu[pivot_row, j]) < tolerance:
            j = j + 1
            continue

        if i != pivot_row:
            # Swap rows
            uu[[i, pivot_row]] = uu[[pivot_row, i]]
            ll[[i, pivot_row]] = ll[[pivot_row, i]]
            pp[[i, pivot_row]] = pp[[pivot_row, i]]

        for k in range(i + 1, nn):
            q = uu[k, j] / uu[i, j]
            uu[k, :] = uu[k, :] - q * uu[i, :]
            ll[k, j] = q

        i = i + 1
        j = j + 1

    return ll + np.eye(nn), uu, pp

def forward_substitution(
        aa: NDArray[np.float_],
        bb: NDArray[np.float_]
) -> NDArray[np.float_]:
    '''Solves (T x = b), where T is a lower triangular matrix

    Args:
        aa: lower triangular matrix
        bb: vector

    Returns:
        x: solution of the system of equations
    '''
    nn = aa.shape[0]
    xx = np.zeros(nn, dtype=np.float_)
    for i in range(nn):
        xx[i] = (bb[i] - np.dot(aa[i, 0:i], xx[0:i])) / aa[i, i]
    return xx

def back_substitution(
        aa: NDArray[np.float_],
        bb: NDArray[np.float_]
) -> NDArray[np.float_]:
    '''Solves (T x = b), where T is a upper triangular matrix

    Args:
        aa: upper triangular matrix
        bb: vector

    Returns:
        x: solution of the system of equations
    '''
    nn = aa.shape[0]
    xx = np.zeros(nn, dtype=np.float_)
    for i in range(nn - 1, -1, -1):
        xx[i] = (bb[i] - np.dot(aa[i, i:], xx[i:nn])) / aa[i, i]
    return xx

def gaussian_eliminate(
        aa: NDArray[np.float_],
        bb: NDArray[np.float_],
        tolerance: float = 1e-6
) -> NDArray[np.float_]:
    '''Solves a linear system of equations (A x = b) by LUP factorization.

    Args:
        aa: upper triangular matrix
        bb: vector
        tolerance: Greatest absolute value considered as 0

    Returns:
        x: solution of the system of equations
    '''
    ll, uu, pp = lu_factorization(aa)

    nn = uu.shape[0]
    # Check if rank of matrix is lower than nn
    if np.abs(uu[nn - 1, nn - 1]) < tolerance:
        raise ValueError

    # L y = P @ b
    y = forward_substitution(ll, pp @ bb)
    # U @ x = y
    x = back_substitution(uu, y)
    return x