linalg.h

Go to the documentation of this file.

#pragma once
 
#include <cassert>
#include <cmath>
#include <iostream>
 
#include "../config.h"
 
namespace alfi::util::linalg {
    template <typename Number = DefaultNumber, template <typename, typename...> class Container = DefaultContainer>
    Container<Number> lup_solve(Container<Container<Number>>&& A, Container<Number>&& B, const Number& epsilon = std::numeric_limits<Number>::epsilon()) {
        const auto n = B.size();
        assert(n == A.size());
        for (SizeT i = 0; i < n; ++i) {
            assert(n == A[i].size());
            Number max_a = std::abs(A[i][i]);
            SizeT i_max = i;
            for (SizeT k = i + 1; k < n; ++k) {
                const auto cur = std::abs(A[k][i]);
                if (cur > max_a) {
                    max_a = cur;
                    i_max = k;
                }
            }
            if (max_a <= epsilon) {
                std::cerr << "Error in function " << __FUNCTION__ << ": Matrix A is degenerate. Returning an empty array..." << std::endl;
                return {};
            }
            if (i_max != i) {
                std::swap(A[i], A[i_max]);
                std::swap(B[i], B[i_max]);
            }
            for (SizeT j = i + 1; j < n; ++j) {
                A[j][i] /= A[i][i];
                for (SizeT k = i + 1; k < n; ++k) {
                    A[j][k] -= A[j][i] * A[i][k];
                }
            }
        }
        Container<Number> X(n);
        for (SizeT i = 0; i < n; ++i) {
            X[i] = B[i];
            for (SizeT k = 0; k < i; ++k) {
                X[i] -= A[i][k] * X[k];
            }
        }
        for (SizeT iter = 0; iter < n; ++iter) {
            const auto i = n - 1 - iter;
            for (SizeT k = i + 1; k < n; ++k) {
                X[i] -= A[i][k] * X[k];
            }
            X[i] /= A[i][i];
        }
        return X;
    }

    template <typename Number = DefaultNumber, template <typename, typename...> class Container = DefaultContainer>
    Container<Number> tridiag_solve_unstable(
            const Container<Number>& lower,
            Container<Number>&& diag,
            const Container<Number>& upper,
            Container<Number>&& right
    ) {
        const auto n = right.size();
        assert(n == lower.size());
        assert(n == diag.size());
        assert(n == upper.size());
        if (n == 0) {
            return {};
        }
        for (SizeT i = 0; i < n - 1; ++i) {
            const auto m = lower[i+1] / diag[i];
            diag[i+1] -= m * upper[i];
            right[i+1] -= m * right[i];
        }
        Container<Number> X(n);
        X[n-1] = right[n-1] / diag[n-1];
        for (SizeT iter = 2; iter <= n; ++iter) {
            const auto i = n - iter;
            X[i] = (right[i] - upper[i] * X[i+1]) / diag[i];
        }
        return X;
    }

    template <typename Number = DefaultNumber, template <typename, typename...> class Container = DefaultContainer>
    Container<Number> tridiag_solve(
            Container<Number>&& lower,
            Container<Number>&& diag,
            Container<Number>&& upper,
            Container<Number>&& right
    ) {
        const auto n = right.size();
        assert(n == lower.size());
        assert(n == diag.size());
        assert(n == upper.size());
        if (n == 0) {
            return {};
        }
        if (n == 1) {
            return {right[0]/diag[0]};
        }
        for (SizeT i = 0; i < n - 1; ++i) {
            if (std::abs(diag[i]) >= std::abs(lower[i+1])) {
                const auto m = lower[i+1] / diag[i];
                diag[i+1] -= m * upper[i];
                right[i+1] -= m * right[i];
                lower[i+1] = 0;
            } else {
                // Swap rows i and (i+1).
                // Eliminate the lower[i+1] element by reducing row (i+1) by row i.
                // Use lower[i+1] as a buffer for the non-tridiagonal element (i,i+2) (hack, used below).
                const auto m = diag[i] / lower[i+1];
                diag[i] = lower[i+1];
                lower[i+1] = upper[i+1];
                upper[i+1] *= -m;
                std::swap(upper[i], diag[i+1]);
                diag[i+1] -= m * upper[i];
                std::swap(right[i], right[i+1]);
                right[i+1] -= m * right[i];
            }
        }
        Container<Number> X(n);
        X[n-1] = right[n-1] / diag[n-1];
        X[n-2] = (right[n-2] - upper[n-2] * X[n-1]) / diag[n-2];
        for (SizeT iter = 3; iter <= n; ++iter) {
            const auto i = n - iter;
            X[i] = (right[i] - upper[i] * X[i+1] - lower[i+1] * X[i+2]) / diag[i];
        }
        return X;
    }
}