DiagonalMatrix.h

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2009 Gael Guennebaud <gael.guennebaud@inria.fr>
// Copyright (C) 2007-2009 Benoit Jacob <jacob.benoit.1@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_DIAGONALMATRIX_H
#define EIGEN_DIAGONALMATRIX_H

// IWYU pragma: private
#include "./InternalHeaderCheck.h"

namespace Eigen {

template <typename Derived>
class DiagonalBase : public EigenBase<Derived> {
 public:
  typedef typename internal::traits<Derived>::DiagonalVectorType DiagonalVectorType;
  typedef typename DiagonalVectorType::Scalar Scalar;
  typedef typename DiagonalVectorType::RealScalar RealScalar;
  typedef typename internal::traits<Derived>::StorageKind StorageKind;
  typedef typename internal::traits<Derived>::StorageIndex StorageIndex;

  enum {
    RowsAtCompileTime = DiagonalVectorType::SizeAtCompileTime,
    ColsAtCompileTime = DiagonalVectorType::SizeAtCompileTime,
    MaxRowsAtCompileTime = DiagonalVectorType::MaxSizeAtCompileTime,
    MaxColsAtCompileTime = DiagonalVectorType::MaxSizeAtCompileTime,
    IsVectorAtCompileTime = 0,
    Flags = NoPreferredStorageOrderBit
  };

  typedef Matrix<Scalar, RowsAtCompileTime, ColsAtCompileTime, 0, MaxRowsAtCompileTime, MaxColsAtCompileTime>
      DenseMatrixType;
  typedef DenseMatrixType DenseType;
  typedef DiagonalMatrix<Scalar, DiagonalVectorType::SizeAtCompileTime, DiagonalVectorType::MaxSizeAtCompileTime>
      PlainObject;

  EIGEN_DEVICE_FUNC inline const Derived& derived() const { return *static_cast<const Derived*>(this); }
  EIGEN_DEVICE_FUNC inline Derived& derived() { return *static_cast<Derived*>(this); }

  EIGEN_DEVICE_FUNC DenseMatrixType toDenseMatrix() const { return derived(); }

  EIGEN_DEVICE_FUNC inline const DiagonalVectorType& diagonal() const { return derived().diagonal(); }
  EIGEN_DEVICE_FUNC inline DiagonalVectorType& diagonal() { return derived().diagonal(); }

  EIGEN_DEVICE_FUNC inline Scalar coeff(Index row, Index col) const {
    eigen_assert(row >= 0 && col >= 0 && row < rows() && col <= cols());
    return row == col ? diagonal().coeff(row) : Scalar(0);
  }

  EIGEN_DEVICE_FUNC constexpr Index rows() const { return diagonal().size(); }
  EIGEN_DEVICE_FUNC constexpr Index cols() const { return diagonal().size(); }

  template <typename MatrixDerived>
  EIGEN_DEVICE_FUNC const Product<Derived, MatrixDerived, LazyProduct> operator*(
      const MatrixBase<MatrixDerived>& matrix) const {
    return Product<Derived, MatrixDerived, LazyProduct>(derived(), matrix.derived());
  }

  template <typename OtherDerived>
  using DiagonalProductReturnType = DiagonalWrapper<const EIGEN_CWISE_BINARY_RETURN_TYPE(
      DiagonalVectorType, typename OtherDerived::DiagonalVectorType, product)>;

  template <typename OtherDerived>
  EIGEN_DEVICE_FUNC const DiagonalProductReturnType<OtherDerived> operator*(
      const DiagonalBase<OtherDerived>& other) const {
    return diagonal().cwiseProduct(other.diagonal()).asDiagonal();
  }

  using DiagonalInverseReturnType =
      DiagonalWrapper<const CwiseUnaryOp<internal::scalar_inverse_op<Scalar>, const DiagonalVectorType>>;

  EIGEN_DEVICE_FUNC inline const DiagonalInverseReturnType inverse() const {
    return diagonal().cwiseInverse().asDiagonal();
  }

  using DiagonalScaleReturnType =
      DiagonalWrapper<const EIGEN_EXPR_BINARYOP_SCALAR_RETURN_TYPE(DiagonalVectorType, Scalar, product)>;

  EIGEN_DEVICE_FUNC inline const DiagonalScaleReturnType operator*(const Scalar& scalar) const {
    return (diagonal() * scalar).asDiagonal();
  }

  using ScaleDiagonalReturnType =
      DiagonalWrapper<const EIGEN_SCALAR_BINARYOP_EXPR_RETURN_TYPE(Scalar, DiagonalVectorType, product)>;

  EIGEN_DEVICE_FUNC friend inline const ScaleDiagonalReturnType operator*(const Scalar& scalar,
                                                                          const DiagonalBase& other) {
    return (scalar * other.diagonal()).asDiagonal();
  }

  template <typename OtherDerived>
  using DiagonalSumReturnType = DiagonalWrapper<const EIGEN_CWISE_BINARY_RETURN_TYPE(
      DiagonalVectorType, typename OtherDerived::DiagonalVectorType, sum)>;

  template <typename OtherDerived>
  EIGEN_DEVICE_FUNC inline const DiagonalSumReturnType<OtherDerived> operator+(
      const DiagonalBase<OtherDerived>& other) const {
    return (diagonal() + other.diagonal()).asDiagonal();
  }

  template <typename OtherDerived>
  using DiagonalDifferenceReturnType = DiagonalWrapper<const EIGEN_CWISE_BINARY_RETURN_TYPE(
      DiagonalVectorType, typename OtherDerived::DiagonalVectorType, difference)>;

  template <typename OtherDerived>
  EIGEN_DEVICE_FUNC inline const DiagonalDifferenceReturnType<OtherDerived> operator-(
      const DiagonalBase<OtherDerived>& other) const {
    return (diagonal() - other.diagonal()).asDiagonal();
  }
};

namespace internal {
template <typename Scalar_, int SizeAtCompileTime, int MaxSizeAtCompileTime>
struct traits<DiagonalMatrix<Scalar_, SizeAtCompileTime, MaxSizeAtCompileTime>>
    : traits<Matrix<Scalar_, SizeAtCompileTime, SizeAtCompileTime, 0, MaxSizeAtCompileTime, MaxSizeAtCompileTime>> {
  typedef Matrix<Scalar_, SizeAtCompileTime, 1, 0, MaxSizeAtCompileTime, 1> DiagonalVectorType;
  typedef DiagonalShape StorageKind;
  enum { Flags = LvalueBit | NoPreferredStorageOrderBit | NestByRefBit };
};
}  // namespace internal
template <typename Scalar_, int SizeAtCompileTime, int MaxSizeAtCompileTime>
class DiagonalMatrix : public DiagonalBase<DiagonalMatrix<Scalar_, SizeAtCompileTime, MaxSizeAtCompileTime>> {
 public:
#ifndef EIGEN_PARSED_BY_DOXYGEN
  typedef typename internal::traits<DiagonalMatrix>::DiagonalVectorType DiagonalVectorType;
  typedef const DiagonalMatrix& Nested;
  typedef Scalar_ Scalar;
  typedef typename internal::traits<DiagonalMatrix>::StorageKind StorageKind;
  typedef typename internal::traits<DiagonalMatrix>::StorageIndex StorageIndex;
#endif

 protected:
  DiagonalVectorType m_diagonal;

 public:
  EIGEN_DEVICE_FUNC inline const DiagonalVectorType& diagonal() const { return m_diagonal; }
  EIGEN_DEVICE_FUNC inline DiagonalVectorType& diagonal() { return m_diagonal; }

  EIGEN_DEVICE_FUNC inline DiagonalMatrix() {}

  EIGEN_DEVICE_FUNC explicit inline DiagonalMatrix(Index dim) : m_diagonal(dim) {}

  EIGEN_DEVICE_FUNC inline DiagonalMatrix(const Scalar& x, const Scalar& y) : m_diagonal(x, y) {}

  EIGEN_DEVICE_FUNC inline DiagonalMatrix(const Scalar& x, const Scalar& y, const Scalar& z) : m_diagonal(x, y, z) {}

  template <typename... ArgTypes>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DiagonalMatrix(const Scalar& a0, const Scalar& a1, const Scalar& a2,
                                                       const ArgTypes&... args)
      : m_diagonal(a0, a1, a2, args...) {}

  EIGEN_DEVICE_FUNC explicit EIGEN_STRONG_INLINE DiagonalMatrix(
      const std::initializer_list<std::initializer_list<Scalar>>& list)
      : m_diagonal(list) {}

  EIGEN_DEVICE_FUNC explicit inline DiagonalMatrix(DiagonalVectorType&& diag) : m_diagonal(std::move(diag)) {}

  template <typename OtherDerived>
  EIGEN_DEVICE_FUNC inline DiagonalMatrix(const DiagonalBase<OtherDerived>& other) : m_diagonal(other.diagonal()) {}

#ifndef EIGEN_PARSED_BY_DOXYGEN

  inline DiagonalMatrix(const DiagonalMatrix& other) : m_diagonal(other.diagonal()) {}
#endif

  template <typename OtherDerived>
  EIGEN_DEVICE_FUNC explicit inline DiagonalMatrix(const MatrixBase<OtherDerived>& other) : m_diagonal(other) {}

  template <typename OtherDerived>
  EIGEN_DEVICE_FUNC DiagonalMatrix& operator=(const DiagonalBase<OtherDerived>& other) {
    m_diagonal = other.diagonal();
    return *this;
  }

#ifndef EIGEN_PARSED_BY_DOXYGEN

  EIGEN_DEVICE_FUNC DiagonalMatrix& operator=(const DiagonalMatrix& other) {
    m_diagonal = other.diagonal();
    return *this;
  }
#endif

  typedef DiagonalWrapper<const CwiseNullaryOp<internal::scalar_constant_op<Scalar>, DiagonalVectorType>>
      InitializeReturnType;

  typedef DiagonalWrapper<const CwiseNullaryOp<internal::scalar_zero_op<Scalar>, DiagonalVectorType>>
      ZeroInitializeReturnType;

  EIGEN_DEVICE_FUNC static const ZeroInitializeReturnType Zero() { return DiagonalVectorType::Zero().asDiagonal(); }
  EIGEN_DEVICE_FUNC static const ZeroInitializeReturnType Zero(Index size) {
    return DiagonalVectorType::Zero(size).asDiagonal();
  }
  EIGEN_DEVICE_FUNC static const InitializeReturnType Identity() { return DiagonalVectorType::Ones().asDiagonal(); }
  EIGEN_DEVICE_FUNC static const InitializeReturnType Identity(Index size) {
    return DiagonalVectorType::Ones(size).asDiagonal();
  }

  EIGEN_DEVICE_FUNC inline void resize(Index size) { m_diagonal.resize(size); }
  EIGEN_DEVICE_FUNC inline void setZero() { m_diagonal.setZero(); }
  EIGEN_DEVICE_FUNC inline void setZero(Index size) { m_diagonal.setZero(size); }
  EIGEN_DEVICE_FUNC inline void setIdentity() { m_diagonal.setOnes(); }
  EIGEN_DEVICE_FUNC inline void setIdentity(Index size) { m_diagonal.setOnes(size); }
};

namespace internal {
template <typename DiagonalVectorType_>
struct traits<DiagonalWrapper<DiagonalVectorType_>> {
  typedef DiagonalVectorType_ DiagonalVectorType;
  typedef typename DiagonalVectorType::Scalar Scalar;
  typedef typename DiagonalVectorType::StorageIndex StorageIndex;
  typedef DiagonalShape StorageKind;
  typedef typename traits<DiagonalVectorType>::XprKind XprKind;
  enum {
    RowsAtCompileTime = DiagonalVectorType::SizeAtCompileTime,
    ColsAtCompileTime = DiagonalVectorType::SizeAtCompileTime,
    MaxRowsAtCompileTime = DiagonalVectorType::MaxSizeAtCompileTime,
    MaxColsAtCompileTime = DiagonalVectorType::MaxSizeAtCompileTime,
    Flags = (traits<DiagonalVectorType>::Flags & LvalueBit) | NoPreferredStorageOrderBit
  };
};
}  // namespace internal

template <typename DiagonalVectorType_>
class DiagonalWrapper : public DiagonalBase<DiagonalWrapper<DiagonalVectorType_>>, internal::no_assignment_operator {
 public:
#ifndef EIGEN_PARSED_BY_DOXYGEN
  typedef DiagonalVectorType_ DiagonalVectorType;
  typedef DiagonalWrapper Nested;
#endif

  EIGEN_DEVICE_FUNC explicit inline DiagonalWrapper(DiagonalVectorType& a_diagonal) : m_diagonal(a_diagonal) {}

  EIGEN_DEVICE_FUNC const DiagonalVectorType& diagonal() const { return m_diagonal; }

 protected:
  typename DiagonalVectorType::Nested m_diagonal;
};

template <typename Derived>
EIGEN_DEVICE_FUNC inline const DiagonalWrapper<const Derived> MatrixBase<Derived>::asDiagonal() const {
  return DiagonalWrapper<const Derived>(derived());
}

template <typename Derived>
bool MatrixBase<Derived>::isDiagonal(const RealScalar& prec) const {
  if (cols() != rows()) return false;
  RealScalar maxAbsOnDiagonal = static_cast<RealScalar>(-1);
  for (Index j = 0; j < cols(); ++j) {
    RealScalar absOnDiagonal = numext::abs(coeff(j, j));
    if (absOnDiagonal > maxAbsOnDiagonal) maxAbsOnDiagonal = absOnDiagonal;
  }
  for (Index j = 0; j < cols(); ++j)
    for (Index i = 0; i < j; ++i) {
      if (!internal::isMuchSmallerThan(coeff(i, j), maxAbsOnDiagonal, prec)) return false;
      if (!internal::isMuchSmallerThan(coeff(j, i), maxAbsOnDiagonal, prec)) return false;
    }
  return true;
}

namespace internal {

template <>
struct storage_kind_to_shape<DiagonalShape> {
  typedef DiagonalShape Shape;
};

struct Diagonal2Dense {};

template <>
struct AssignmentKind<DenseShape, DiagonalShape> {
  typedef Diagonal2Dense Kind;
};

// Diagonal matrix to Dense assignment
template <typename DstXprType, typename SrcXprType, typename Functor>
struct Assignment<DstXprType, SrcXprType, Functor, Diagonal2Dense> {
  static EIGEN_DEVICE_FUNC void run(
      DstXprType& dst, const SrcXprType& src,
      const internal::assign_op<typename DstXprType::Scalar, typename SrcXprType::Scalar>& /*func*/) {
    Index dstRows = src.rows();
    Index dstCols = src.cols();
    if ((dst.rows() != dstRows) || (dst.cols() != dstCols)) dst.resize(dstRows, dstCols);

    dst.setZero();
    dst.diagonal() = src.diagonal();
  }

  static EIGEN_DEVICE_FUNC void run(
      DstXprType& dst, const SrcXprType& src,
      const internal::add_assign_op<typename DstXprType::Scalar, typename SrcXprType::Scalar>& /*func*/) {
    dst.diagonal() += src.diagonal();
  }

  static EIGEN_DEVICE_FUNC void run(
      DstXprType& dst, const SrcXprType& src,
      const internal::sub_assign_op<typename DstXprType::Scalar, typename SrcXprType::Scalar>& /*func*/) {
    dst.diagonal() -= src.diagonal();
  }
};

}  // namespace internal

}  // end namespace Eigen

#endif  // EIGEN_DIAGONALMATRIX_H