DiagonalMatrix.h

#ifndef theplu_yat_utility_diagonal_matrix
#define theplu_yat_utility_diagonal_matrix

// $Id: DiagonalMatrix.h 4207 2022-08-26 04:36:28Z peter $

/*
  Copyright (C) 2017, 2021, 2022 Peter Johansson

  This file is part of the yat library, http://dev.thep.lu.se/yat

  The yat library 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 3 of the
  License, or (at your option) any later version.

  The yat library 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 yat. If not, see <http://www.gnu.org/licenses/>.
*/

#include "BasicMatrix.h"
#include "BasicVector.h"
#include "BLAS_utility.h"

#include "MatrixExpression.h"
#include "VectorExpression.h"

#include "Vector.h"
#include "yat_assert.h"

namespace theplu {
namespace yat {
namespace utility {

  class MatrixBase;

  class DiagonalMatrix
  {
  public:
    DiagonalMatrix(void);

    DiagonalMatrix(size_t row, size_t column, double x=0.0);

    explicit DiagonalMatrix(const VectorBase& diagonal);

    explicit DiagonalMatrix(const MatrixBase& m);

    size_t rows(void) const;

    size_t columns(void) const;

    const double operator()(size_t row, size_t col) const;

    double& operator()(size_t i);

    DiagonalMatrix& operator*=(const DiagonalMatrix& rhs);

    DiagonalMatrix& operator+=(const DiagonalMatrix& rhs);

    DiagonalMatrix& operator-=(const DiagonalMatrix& rhs);
  private:
    Vector data_;
    size_t row_;
    size_t col_;
  };


  // expression classes
  namespace expression
  {
    template<typename T>
    class DiagonalMatrixMatrix :
      public MatrixExpression<DiagonalMatrixMatrix<T> >
    {
    public:
      DiagonalMatrixMatrix(const DiagonalMatrix& lhs,const BasicMatrix<T>& rhs)
        : lhs_(lhs), rhs_(rhs) {}
      size_t rows(void) const { return lhs_.rows(); }
      size_t columns(void) const { return rhs_.columns(); }
      double operator()(size_t i, size_t j) const
      { return lhs_(i, i) * rhs_(i, j); }

      void calculate_matrix(gsl_matrix*& result) const
      {
        detail::reallocate(result, rows(), columns());
        for (size_t i=0; i<rows(); ++i)
          for (size_t j=0; j<columns(); ++j)
            gsl_matrix_set(result, i, j, (*this)(i, j));
      }

    private:
      const DiagonalMatrix& lhs_;
      const BasicMatrix<T>& rhs_;
    };


    template<class T>
    class DiagonalMatrixVector
      : public VectorExpression<DiagonalMatrixVector<T> >
    {
    public:
      DiagonalMatrixVector(const DiagonalMatrix& lhs, const BasicVector<T>& rhs)
        : m_(lhs), vec_(rhs), size_(lhs.rows())
      {
      }

      DiagonalMatrixVector(const BasicVector<T>& lhs, const DiagonalMatrix& rhs)
        : m_(rhs), vec_(lhs), size_(rhs.columns())
      {
      }

      double operator()(size_t i) const
      {
        YAT_ASSERT(i < size());
        return m_(i, i) * vec_(i);
      }


      size_t size(void) const
      {
        return size_;
      }


      void calculate_gsl_vector_p(void) const
      {
        this->allocate_memory(size_);
        for (size_t i=0; i<size_; ++i)
          gsl_vector_set(this->v_, i, (*this)(i));
      }
    private:
      const DiagonalMatrix& m_;
      const BasicVector<T>& vec_;
      size_t size_;
    };



    template<typename T>
    class MatrixDiagonalMatrix :
      public MatrixExpression<MatrixDiagonalMatrix<T> >
    {
    public:
      MatrixDiagonalMatrix(const BasicMatrix<T>& lhs, const DiagonalMatrix& rhs)
        : lhs_(lhs), rhs_(rhs) {}
      size_t rows(void) const { return lhs_.rows(); }
      size_t columns(void) const { return rhs_.columns(); }
      double operator()(size_t i, size_t j) const
      { return lhs_(i, j) * rhs_(j, j); }

      void calculate_matrix(gsl_matrix*& result) const
      {
        detail::reallocate(result, rows(), columns());
        for (size_t i=0; i<rows(); ++i)
          for (size_t j=0; j<columns(); ++j)
            gsl_matrix_set(result, i, j, (*this)(i, j));
      }

    private:
      const BasicMatrix<T>& lhs_;
      const DiagonalMatrix& rhs_;
    };


    template<typename T>
    class VectorDiagonalMatrix :
      public BasicVector<VectorDiagonalMatrix<T> >
    {
    public:
      VectorDiagonalMatrix(const BasicVector<T>& rhs, const DiagonalMatrix& lhs)
        : lhs_(lhs), rhs_(rhs) {}

      double operator()(size_t i) const
      {
        YAT_ASSERT(i < size());
        return lhs_(i) * rhs_(i, i);
      }

      size_t size(void) const
      {
        return rhs_.columns();
      }


      const gsl_vector* gsl_vector_p(void) const
      {
        size_t n = size();
        gsl_vector* vec = detail::create_gsl_vector(n);
        for (size_t i=0; i<n; ++i)
          gsl_vector_set(vec, i, (*this)(i));
        return vec;
      }

    private:
      const BasicVector<T>& lhs_;
      const DiagonalMatrix& rhs_;
    };


    template<class M1, class M2, class OP>
    class DiagonalMatrix
      : public MatrixExpression<DiagonalMatrix<M1, M2, OP> >
    {
    public:
      DiagonalMatrix(const M1& lhs, const M2& rhs)
        : lhs_(lhs), rhs_(rhs)
      {
        YAT_ASSERT(lhs.rows() == rhs.rows());
        YAT_ASSERT(lhs.columns() == rhs.columns());
      }
      double operator()(size_t i, size_t j) const
      {
        return get(i, j, OP());
      }
      size_t rows(void) const { return lhs_.rows(); }
      size_t columns(void) const { return rhs_.columns(); }
      void calculate_matrix(gsl_matrix*& m) const
      {
        detail::reallocate(m, rows(), columns());
        for (size_t i=0; i<rows(); ++i)
          for (size_t j=0; j<columns(); ++j)
            gsl_matrix_set(m, i, j, (*this)(i,j));
      }

    private:
      const M1& lhs_;
      const M2& rhs_;

      double get(size_t i, size_t j, Plus) const
      { return lhs_(i, j) + rhs_(i, j); }

      double get(size_t i, size_t j, Minus) const
      { return lhs_(i, j) - rhs_(i, j); }
    };

  } // end of namespace expression



  DiagonalMatrix operator*(const DiagonalMatrix& lhs,
                           const DiagonalMatrix& rhs);

  template<class T>
  expression::DiagonalMatrixMatrix<T>
  operator*(const DiagonalMatrix& lhs, const BasicMatrix<T>& rhs)
  {
    YAT_ASSERT(lhs.columns() == rhs.rows());
    return expression::DiagonalMatrixMatrix<T>(lhs, rhs);
  }


  template<class T>
  expression::MatrixDiagonalMatrix<T>
  operator*(const BasicMatrix<T>& lhs, const DiagonalMatrix& rhs)
  {
    YAT_ASSERT(lhs.columns() == rhs.rows());
    return expression::MatrixDiagonalMatrix<T>(lhs, rhs);
  }


  DiagonalMatrix
  operator+(const DiagonalMatrix& lhs, const DiagonalMatrix& rhs);


  template<class T>
  expression::DiagonalMatrix<DiagonalMatrix, BasicMatrix<T>, expression::Plus>
  operator+(const DiagonalMatrix& lhs, const BasicMatrix<T>& rhs)
  {
    YAT_ASSERT(lhs.rows() == rhs.rows());
    YAT_ASSERT(lhs.columns() == rhs.columns());
    return expression::DiagonalMatrix<DiagonalMatrix, BasicMatrix<T>,
                                      expression::Plus>(lhs, rhs);
  }


  template<class T>
  expression::DiagonalMatrix<BasicMatrix<T>, DiagonalMatrix, expression::Plus>
  operator+(const BasicMatrix<T>& lhs, const DiagonalMatrix& rhs)
  {
    YAT_ASSERT(lhs.rows() == rhs.rows());
    YAT_ASSERT(lhs.columns() == rhs.columns());
    return expression::DiagonalMatrix<BasicMatrix<T>, DiagonalMatrix,
                                      expression::Plus>(lhs, rhs);
  }


  DiagonalMatrix
  operator-(const DiagonalMatrix& lhs, const DiagonalMatrix& rhs);


  template<class T>
  expression::DiagonalMatrix<DiagonalMatrix, BasicMatrix<T>, expression::Minus>
  operator-(const DiagonalMatrix& lhs, const BasicMatrix<T>& rhs)
  {
    YAT_ASSERT(lhs.rows() == rhs.rows());
    YAT_ASSERT(lhs.columns() == rhs.columns());
    return expression::DiagonalMatrix<DiagonalMatrix, BasicMatrix<T>,
                                      expression::Minus>(lhs, rhs);
  }


  template<class T>
  expression::DiagonalMatrix<BasicMatrix<T>, DiagonalMatrix, expression::Minus>
  operator-(const BasicMatrix<T>& lhs, const DiagonalMatrix& rhs)
  {
    YAT_ASSERT(lhs.rows() == rhs.rows());
    YAT_ASSERT(lhs.columns() == rhs.columns());
    return expression::DiagonalMatrix<BasicMatrix<T>, DiagonalMatrix,
                                      expression::Minus>(lhs, rhs);
  }


  template<class T>
  expression::DiagonalMatrixVector<T>
  operator*(const DiagonalMatrix& lhs, const BasicVector<T>& rhs)
  {
    YAT_ASSERT(lhs.columns() == rhs.size());
    return expression::DiagonalMatrixVector<T>(lhs, rhs);
  }


  template<class T>
  expression::DiagonalMatrixVector<T>
  operator*(const BasicVector<T>& lhs, const DiagonalMatrix& rhs)
  {
    YAT_ASSERT(lhs.size() == rhs.rows());
    return  expression::DiagonalMatrixVector<T>(lhs, rhs);
  }


}}}
#endif