aboutsummaryrefslogtreecommitdiff
path: root/solvers.py
blob: da38e2be4e27c34b16a9759bd4ed75c64403f9da (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
"""Routines for solving a linear system of equations."""

import numpy as np
from numpy.typing import NDArray


def gaussian_eliminate(
        aa: NDArray[np.float_],
        bb: NDArray[np.float_],
        tolerance: float = 1e-6
) -> NDArray[np.float_] | None:
    """Solves a linear system of equations (Ax = b) by Gauss-elimination

    Args:
        aa: Matrix with the coefficients. Shape: (n, n).
        bb: Right hand side of the equation. Shape: (n,)
        tolerance: Greatest absolute value considered as 0

    Returns:
        Vector xx with the solution of the linear equation or None
        if the equations are linearly dependent.
    """
    nn = aa.shape[0]

    ee = np.zeros((nn, nn + 1))
    ee[:, 0:nn] = aa
    ee[:, nn] = bb

    i = 0
    j = 0

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

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

        if i != pivot_row:
            # Swap rows
            ee[[i, pivot_row]] = ee[[pivot_row, i]]

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

        i = i + 1
        j = j + 1

    # Check if rank of matrix is lower than nn
    if np.abs(ee[nn - 1, nn - 1]) < tolerance:
        return None

    # Back substitution
    xx = np.zeros((nn, 1))
    for i in range(nn - 1, -1, -1):
        xx[i] = (ee[i, nn] - np.dot(ee[i, i:nn], xx[i:nn])) / ee[i, i]

    return xx