Repository: TheAlgorithms/Python
File: linear_algebra/src/gaussian_elimination_pivoting.py (Python)

Code before:
    Example:
    >>> A = np.array([[2, 1, -1], [-3, -1, 2], [-2, 1, 2]], dtype=float)
    >>> B = np.array([8, -11, -3], dtype=float)
    >>> solution = solve_linear_system(np.column_stack((A, B)))
    >>> np.allclose(solution, np.array([2., 3., -1.]))
    True
    >>> solve_linear_system(np.array([[0, 0, 0], [0, 0, 0]],  dtype=float))
    Traceback (most recent call last):
        ...
    ValueError: Matrix is not correct
    """
    ab = np.copy(matrix)
    num_of_rows = ab.shape[0]
    num_of_columns = ab.shape[1] - 1
    x_lst: list[float] = []

    assert num_of_rows == num_of_columns

    for column_num in range(num_of_rows):
        # Lead element search
        for i in range(column_num, num_of_columns):
            if abs(ab[i][column_num]) > abs(ab[column_num][column_num]):
                ab[[column_num, i]] = ab[[i, column_num]]

        # Upper triangular matrix
        if ab[column_num, column_num] == 0.0:
            raise ValueError("Matrix is not correct")

        if column_num != 0:
            for i in range(column_num, num_of_rows):
                ab[i, :] -= (
                    ab[i, column_num - 1]
                    / ab[column_num - 1, column_num - 1]
                    * ab[column_num - 1, :]
                )

    # Find x vector (Back Substitution)
    for column_num in range(num_of_rows - 1, -1, -1):
        x = ab[column_num, -1] / ab[column_num, column_num]
        x_lst.insert(0, x)
        for i in range(column_num - 1, -1, -1):
            ab[i, -1] -= ab[i, column_num] * x

    # Return the solution vector
    return np.asarray(x_lst)


if __name__ == "__main__":
    from doctest import testmod

    testmod()

Review comment (bug):
Should there be a tolerance for this comparison? That is, is there a risk that floating-point errors can cause this check to incorrectly fail?

Code after:
    Example:
    >>> A = np.array([[2, 1, -1], [-3, -1, 2], [-2, 1, 2]], dtype=float)
    >>> B = np.array([8, -11, -3], dtype=float)
    >>> solution = solve_linear_system(np.column_stack((A, B)))
    >>> np.allclose(solution, np.array([2., 3., -1.]))
    True
    >>> solve_linear_system(np.array([[0, 0, 0]], dtype=float))
    Traceback (most recent call last):
        ...
    ValueError: Matrix is not square
    >>> solve_linear_system(np.array([[0, 0, 0], [0, 0, 0]], dtype=float))
    Traceback (most recent call last):
        ...
    ValueError: Matrix is singular
    """
    ab = np.copy(matrix)
    num_of_rows = ab.shape[0]
    num_of_columns = ab.shape[1] - 1
    x_lst: list[float] = []

    if num_of_rows != num_of_columns:
        raise ValueError("Matrix is not square")

    for column_num in range(num_of_rows):
        # Lead element search
        for i in range(column_num, num_of_columns):
            if abs(ab[i][column_num]) > abs(ab[column_num][column_num]):
                ab[[column_num, i]] = ab[[i, column_num]]

        # Upper triangular matrix
        if abs(ab[column_num, column_num]) < 1e-8:
            raise ValueError("Matrix is singular")

        if column_num != 0:
            for i in range(column_num, num_of_rows):
                ab[i, :] -= (
                    ab[i, column_num - 1]
                    / ab[column_num - 1, column_num - 1]
                    * ab[column_num - 1, :]
                )

    # Find x vector (Back Substitution)
    for column_num in range(num_of_rows - 1, -1, -1):
        x = ab[column_num, -1] / ab[column_num, column_num]
        x_lst.insert(0, x)
        for i in range(column_num - 1, -1, -1):
            ab[i, -1] -= ab[i, column_num] * x

    # Return the solution vector
    return np.asarray(x_lst)