TridiagonalMatrix.hpp

Go to the documentation of this file.

// -*- mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*-
// vi: set et ts=4 sw=4 sts=4:
/*
  This file is part of the Open Porous Media project (OPM).
 
  OPM is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 2 of the License, or
  (at your option) any later version.
 
  OPM is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
 
  You should have received a copy of the GNU General Public License
  along with OPM.  If not, see <http://www.gnu.org/licenses/>.
 
  Consult the COPYING file in the top-level source directory of this
  module for the precise wording of the license and the list of
  copyright holders.
*/
#ifndef OPM_TRIDIAGONAL_MATRIX_HH
#define OPM_TRIDIAGONAL_MATRIX_HH
 
#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
 
#include <assert.h>
 
namespace Opm {
 
template <class Scalar>
class TridiagonalMatrix
{
    struct TridiagRow_ {
        TridiagRow_(TridiagonalMatrix& m, size_t rowIdx)
            : matrix_(m)
            , rowIdx_(rowIdx)
        {}
 
        Scalar& operator[](size_t colIdx)
        { return matrix_.at(rowIdx_, colIdx); }
 
        Scalar operator[](size_t colIdx) const
        { return matrix_.at(rowIdx_, colIdx); }
 
        TridiagRow_& operator++()
        { ++ rowIdx_; return *this; }
 
        TridiagRow_& operator--()
        { -- rowIdx_; return *this; }
 
        bool operator==(const TridiagRow_& other) const
        { return other.rowIdx_ == rowIdx_ && &other.matrix_ == &matrix_; }
 
        bool operator!=(const TridiagRow_& other) const
        { return !operator==(other); }
 
        TridiagRow_& operator*()
        { return *this; }
 
        size_t index() const
        { return rowIdx_; }
 
    private:
        TridiagonalMatrix& matrix_;
        mutable size_t rowIdx_;
    };
 
public:
    typedef Scalar FieldType;
    typedef TridiagRow_ RowType;
    typedef size_t SizeType;
    typedef TridiagRow_ iterator;
    typedef TridiagRow_ const_iterator;
 
    explicit TridiagonalMatrix(size_t numRows = 0)
    {
        resize(numRows);
    }
 
    TridiagonalMatrix(size_t numRows, Scalar value)
    {
        resize(numRows);
        this->operator=(value);
    }
 
    TridiagonalMatrix(const TridiagonalMatrix& source)
    { *this = source; }
 
    size_t size() const
    { return diag_[0].size(); }
 
    size_t rows() const
    { return size(); }
 
    size_t cols() const
    { return size(); }
 
    void resize(size_t n)
    {
        if (n == size())
            return;
 
        for (int diagIdx = 0; diagIdx < 3; ++ diagIdx)
            diag_[diagIdx].resize(n);
    }
 
    Scalar& at(size_t rowIdx, size_t colIdx)
    {
        size_t n = size();
 
        // special cases
        if (n > 2) {
            if (rowIdx == 0 && colIdx == n - 1)
                return diag_[2][0];
            if (rowIdx == n - 1 && colIdx == 0)
                return diag_[0][n - 1];
        }
 
        size_t diagIdx = 1 + colIdx - rowIdx;
        // make sure that the requested column is in range
        assert(diagIdx < 3);
        return diag_[diagIdx][colIdx];
    }
 
    Scalar at(size_t rowIdx, size_t colIdx) const
    {
        size_t n = size();
 
        // special cases
        if (rowIdx == 0 && colIdx == n - 1)
            return diag_[2][0];
        if (rowIdx == n - 1 && colIdx == 0)
            return diag_[0][n - 1];
 
        size_t diagIdx = 1 + colIdx - rowIdx;
        // make sure that the requested column is in range
        assert(diagIdx < 3);
        return diag_[diagIdx][colIdx];
    }
 
    TridiagonalMatrix& operator=(const TridiagonalMatrix& source)
    {
        for (unsigned diagIdx = 0; diagIdx < 3; ++ diagIdx)
            diag_[diagIdx] = source.diag_[diagIdx];
 
        return *this;
    }
 
    TridiagonalMatrix& operator=(Scalar value)
    {
        for (unsigned diagIdx = 0; diagIdx < 3; ++ diagIdx)
            diag_[diagIdx].assign(size(), value);
 
        return *this;
    }
 
    iterator begin()
    { return TridiagRow_(*this, 0); }
 
    const_iterator begin() const
    { return TridiagRow_(const_cast<TridiagonalMatrix&>(*this), 0); }
 
    const_iterator end() const
    { return TridiagRow_(const_cast<TridiagonalMatrix&>(*this), size()); }
 
    TridiagRow_ operator[](size_t rowIdx)
    { return TridiagRow_(*this, rowIdx); }
 
    const TridiagRow_ operator[](size_t rowIdx) const
    { return TridiagRow_(*this, rowIdx); }
 
    TridiagonalMatrix& operator*=(Scalar alpha)
    {
        unsigned n = size();
        for (unsigned diagIdx = 0; diagIdx < 3; ++ diagIdx) {
            for (unsigned i = 0; i < n; ++i) {
                diag_[diagIdx][i] *= alpha;
            }
        }
 
        return *this;
    }
 
    TridiagonalMatrix& operator/=(Scalar alpha)
    {
        alpha = 1.0/alpha;
        unsigned n = size();
        for (unsigned diagIdx = 0; diagIdx < 3; ++ diagIdx) {
            for (unsigned i = 0; i < n; ++i) {
                diag_[diagIdx][i] *= alpha;
            }
        }
 
        return *this;
    }
 
    TridiagonalMatrix& operator-=(const TridiagonalMatrix& other)
    { return axpy(-1.0, other); }
 
    TridiagonalMatrix& operator+=(const TridiagonalMatrix& other)
    { return axpy(1.0, other); }
 
 
    TridiagonalMatrix& axpy(Scalar alpha, const TridiagonalMatrix& other)
    {
        assert(size() == other.size());
 
        unsigned n = size();
        for (unsigned diagIdx = 0; diagIdx < 3; ++ diagIdx)
            for (unsigned i = 0; i < n; ++ i)
                diag_[diagIdx][i] += alpha * other[diagIdx][i];
 
        return *this;
    }
 
    template<class Vector>
    void mv(const Vector& source, Vector& dest) const
    {
        assert(source.size() == size());
        assert(dest.size() == size());
        assert(size() > 1);
 
        // deal with rows 1 .. n-2
        unsigned n = size();
        for (unsigned i = 1; i < n - 1; ++ i) {
            dest[i] =
                diag_[0][i - 1]*source[i-1] +
                diag_[1][i]*source[i] +
                diag_[2][i + 1]*source[i + 1];
        }
 
        // rows 0 and n-1
        dest[0] =
            diag_[1][0]*source[0] +
            diag_[2][1]*source[1] +
            diag_[2][0]*source[n - 1];
 
        dest[n - 1] =
            diag_[0][n-1]*source[0] +
            diag_[0][n-2]*source[n-2] +
            diag_[1][n-1]*source[n-1];
    }
 
    template<class Vector>
    void umv(const Vector& source, Vector& dest) const
    {
        assert(source.size() == size());
        assert(dest.size() == size());
        assert(size() > 1);
 
        // deal with rows 1 .. n-2
        unsigned n = size();
        for (unsigned i = 1; i < n - 1; ++ i) {
            dest[i] +=
                diag_[0][i - 1]*source[i-1] +
                diag_[1][i]*source[i] +
                diag_[2][i + 1]*source[i + 1];
        }
 
        // rows 0 and n-1
        dest[0] +=
            diag_[1][0]*source[0] +
            diag_[2][1]*source[1] +
            diag_[2][0]*source[n - 1];
 
        dest[n - 1] +=
            diag_[0][n-1]*source[0] +
            diag_[0][n-2]*source[n-2] +
            diag_[1][n-1]*source[n-1];
    }
 
    template<class Vector>
    void mmv(const Vector& source, Vector& dest) const
    {
        assert(source.size() == size());
        assert(dest.size() == size());
        assert(size() > 1);
 
        // deal with rows 1 .. n-2
        unsigned n = size();
        for (unsigned i = 1; i < n - 1; ++ i) {
            dest[i] -=
                diag_[0][i - 1]*source[i-1] +
                diag_[1][i]*source[i] +
                diag_[2][i + 1]*source[i + 1];
        }
 
        // rows 0 and n-1
        dest[0] -=
            diag_[1][0]*source[0] +
            diag_[2][1]*source[1] +
            diag_[2][0]*source[n - 1];
 
        dest[n - 1] -=
            diag_[0][n-1]*source[0] +
            diag_[0][n-2]*source[n-2] +
            diag_[1][n-1]*source[n-1];
    }
 
    template<class Vector>
    void usmv(Scalar alpha, const Vector& source, Vector& dest) const
    {
        assert(source.size() == size());
        assert(dest.size() == size());
        assert(size() > 1);
 
        // deal with rows 1 .. n-2
        unsigned n = size();
        for (unsigned i = 1; i < n - 1; ++ i) {
            dest[i] +=
                alpha*(
                    diag_[0][i - 1]*source[i-1] +
                    diag_[1][i]*source[i] +
                    diag_[2][i + 1]*source[i + 1]);
        }
 
        // rows 0 and n-1
        dest[0] +=
            alpha*(
                diag_[1][0]*source[0] +
                diag_[2][1]*source[1] +
                diag_[2][0]*source[n - 1]);
 
        dest[n - 1] +=
            alpha*(
                diag_[0][n-1]*source[0] +
                diag_[0][n-2]*source[n-2] +
                diag_[1][n-1]*source[n-1]);
    }
 
    template<class Vector>
    void mtv(const Vector& source, Vector& dest) const
    {
        assert(source.size() == size());
        assert(dest.size() == size());
        assert(size() > 1);
 
        // deal with rows 1 .. n-2
        unsigned n = size();
        for (unsigned i = 1; i < n - 1; ++ i) {
            dest[i] =
                diag_[2][i + 1]*source[i-1] +
                diag_[1][i]*source[i] +
                diag_[0][i - 1]*source[i + 1];
        }
 
        // rows 0 and n-1
        dest[0] =
            diag_[1][0]*source[0] +
            diag_[0][1]*source[1] +
            diag_[0][n-1]*source[n - 1];
 
        dest[n - 1] =
            diag_[2][0]*source[0] +
            diag_[2][n-1]*source[n-2] +
            diag_[1][n-1]*source[n-1];
    }
 
    template<class Vector>
    void umtv(const Vector& source, Vector& dest) const
    {
        assert(source.size() == size());
        assert(dest.size() == size());
        assert(size() > 1);
 
        // deal with rows 1 .. n-2
        unsigned n = size();
        for (unsigned i = 1; i < n - 1; ++ i) {
            dest[i] +=
                diag_[2][i + 1]*source[i-1] +
                diag_[1][i]*source[i] +
                diag_[0][i - 1]*source[i + 1];
        }
 
        // rows 0 and n-1
        dest[0] +=
            diag_[1][0]*source[0] +
            diag_[0][1]*source[1] +
            diag_[0][n-1]*source[n - 1];
 
        dest[n - 1] +=
            diag_[2][0]*source[0] +
            diag_[2][n-1]*source[n-2] +
            diag_[1][n-1]*source[n-1];
    }
 
    template<class Vector>
    void mmtv (const Vector& source, Vector& dest) const
    {
        assert(source.size() == size());
        assert(dest.size() == size());
        assert(size() > 1);
 
        // deal with rows 1 .. n-2
        unsigned n = size();
        for (unsigned i = 1; i < n - 1; ++ i) {
            dest[i] -=
                diag_[2][i + 1]*source[i-1] +
                diag_[1][i]*source[i] +
                diag_[0][i - 1]*source[i + 1];
        }
 
        // rows 0 and n-1
        dest[0] -=
            diag_[1][0]*source[0] +
            diag_[0][1]*source[1] +
            diag_[0][n-1]*source[n - 1];
 
        dest[n - 1] -=
            diag_[2][0]*source[0] +
            diag_[2][n-1]*source[n-2] +
            diag_[1][n-1]*source[n-1];
    }
 
    template<class Vector>
    void usmtv(Scalar alpha, const Vector& source, Vector& dest) const
    {
        assert(source.size() == size());
        assert(dest.size() == size());
        assert(size() > 1);
 
        // deal with rows 1 .. n-2
        unsigned n = size();
        for (unsigned i = 1; i < n - 1; ++ i) {
            dest[i] +=
                alpha*(
                    diag_[2][i + 1]*source[i-1] +
                    diag_[1][i]*source[i] +
                    diag_[0][i - 1]*source[i + 1]);
        }
 
        // rows 0 and n-1
        dest[0] +=
            alpha*(
                diag_[1][0]*source[0] +
                diag_[0][1]*source[1] +
                diag_[0][n-1]*source[n - 1]);
 
        dest[n - 1] +=
            alpha*(
                diag_[2][0]*source[0] +
                diag_[2][n-1]*source[n-2] +
                diag_[1][n-1]*source[n-1]);
    }
 
    Scalar frobeniusNorm() const
    { return std::sqrt(frobeniusNormSquared()); }
 
    Scalar frobeniusNormSquared() const
    {
        Scalar result = 0;
 
        unsigned n = size();
        for (unsigned i = 0; i < n; ++ i)
            for (unsigned diagIdx = 0; diagIdx < 3; ++ diagIdx)
                result += diag_[diagIdx][i];
 
        return result;
    }
 
    Scalar infinityNorm() const
    {
        Scalar result=0;
 
        // deal with rows 1 .. n-2
        unsigned n = size();
        for (unsigned i = 1; i < n - 1; ++ i) {
            result = std::max(result,
                              std::abs(diag_[0][i - 1]) +
                              std::abs(diag_[1][i]) +
                              std::abs(diag_[2][i + 1]));
        }
 
        // rows 0 and n-1
        result = std::max(result,
                          std::abs(diag_[1][0]) +
                          std::abs(diag_[2][1]) +
                          std::abs(diag_[2][0]));
 
 
        result = std::max(result,
                          std::abs(diag_[0][n-1]) +
                          std::abs(diag_[0][n-2]) +
                          std::abs(diag_[1][n-2]));
 
        return result;
    }
 
    template <class XVector, class BVector>
    void solve(XVector& x, const BVector& b) const
    {
        if (size() > 2 && std::abs(diag_[2][0]) < 1e-30)
            solveWithUpperRight_(x, b);
        else
            solveWithoutUpperRight_(x, b);
    }
 
    void print(std::ostream& os = std::cout) const
    {
        size_t n = size();
 
        // row 0
        os << at(0, 0) << "\t"
           << at(0, 1) << "\t";
 
        if (n > 3)
            os << "\t";
        if (n > 2)
            os << at(0, n-1);
        os << "\n";
 
        // row 1 .. n - 2
        for (unsigned rowIdx = 1; rowIdx < n-1; ++rowIdx) {
            if (rowIdx > 1)
                os << "\t";
            if (rowIdx == n - 2)
                os << "\t";
 
            os << at(rowIdx, rowIdx - 1) << "\t"
               << at(rowIdx, rowIdx) << "\t"
               << at(rowIdx, rowIdx + 1) << "\n";
        }
 
        // row n - 1
        if (n > 2)
            os << at(n-1, 0) << "\t";
        if (n > 3)
            os << "\t";
        if (n > 4)
            os << "\t";
        os << at(n-1, n-2) << "\t"
           << at(n-1, n-1) << "\n";
    }
 
private:
    template <class XVector, class BVector>
    void solveWithUpperRight_(XVector& x, const BVector& b) const
    {
        size_t n = size();
 
        std::vector<Scalar> lowerDiag(diag_[0]), mainDiag(diag_[1]), upperDiag(diag_[2]), lastColumn(n);
        std::vector<Scalar> bStar(n);
        std::copy(b.begin(), b.end(), bStar.begin());
 
        lastColumn[0] = upperDiag[0];
 
        // forward elimination
        for (size_t i = 1; i < n; ++i) {
            Scalar alpha = lowerDiag[i - 1]/mainDiag[i - 1];
 
            lowerDiag[i - 1] -= alpha * mainDiag[i - 1];
            mainDiag[i] -= alpha * upperDiag[i];
 
            bStar[i] -= alpha * bStar[i - 1];
        };
 
        // deal with the last row if the entry on the lower left is not zero
        if (lowerDiag[n - 1] != 0.0 && size() > 2) {
            Scalar lastRow = lowerDiag[n - 1];
            for (size_t i = 0; i < n - 1; ++i) {
                Scalar alpha = lastRow/mainDiag[i];
                lastRow = - alpha*upperDiag[i + 1];
                bStar[n - 1] -= alpha * bStar[i];
            }
 
            mainDiag[n-1] += lastRow;
        }
 
        // backward elimination
        x[n - 1] = bStar[n - 1]/mainDiag[n-1];
        for (int i = static_cast<int>(n) - 2; i >= 0; --i) {
            unsigned iu = static_cast<unsigned>(i);
            x[iu] = (bStar[iu] - x[iu + 1]*upperDiag[iu+1] - x[n-1]*lastColumn[iu])/mainDiag[iu];
        }
    }
 
    template <class XVector, class BVector>
    void solveWithoutUpperRight_(XVector& x, const BVector& b) const
    {
        size_t n = size();
 
        std::vector<Scalar> lowerDiag(diag_[0]), mainDiag(diag_[1]), upperDiag(diag_[2]);
        std::vector<Scalar> bStar(n);
        std::copy(b.begin(), b.end(), bStar.begin());
 
        // forward elimination
        for (size_t i = 1; i < n; ++i) {
            Scalar alpha = lowerDiag[i - 1]/mainDiag[i - 1];
 
            lowerDiag[i - 1] -= alpha * mainDiag[i - 1];
            mainDiag[i] -= alpha * upperDiag[i];
 
            bStar[i] -= alpha * bStar[i - 1];
        };
 
        // deal with the last row if the entry on the lower left is not zero
        if (lowerDiag[n - 1] != 0.0 && size() > 2) {
            Scalar lastRow = lowerDiag[n - 1];
            for (size_t i = 0; i < n - 1; ++i) {
                Scalar alpha = lastRow/mainDiag[i];
                lastRow = - alpha*upperDiag[i + 1];
                bStar[n - 1] -= alpha * bStar[i];
            }
 
            mainDiag[n-1] += lastRow;
        }
 
        // backward elimination
        x[n - 1] = bStar[n - 1]/mainDiag[n-1];
        for (int i = static_cast<int>(n) - 2; i >= 0; --i) {
            unsigned iu = static_cast<unsigned>(i);
            x[iu] = (bStar[iu] - x[iu + 1]*upperDiag[iu+1])/mainDiag[iu];
        }
    }
 
    mutable std::vector<Scalar> diag_[3];
};
 
} // namespace Opm
 
template <class Scalar>
std::ostream& operator<<(std::ostream& os, const Opm::TridiagonalMatrix<Scalar>& mat)
{
    mat.print(os);
    return os;
}
 
#endif