xref: /aosp_15_r20/external/eigen/Eigen/src/Core/Diagonal.h (revision bf2c37156dfe67e5dfebd6d394bad8b2ab5804d4)
1 // This file is part of Eigen, a lightweight C++ template library
2 // for linear algebra.
3 //
4 // Copyright (C) 2007-2009 Benoit Jacob <[email protected]>
5 // Copyright (C) 2009-2010 Gael Guennebaud <[email protected]>
6 //
7 // This Source Code Form is subject to the terms of the Mozilla
8 // Public License v. 2.0. If a copy of the MPL was not distributed
9 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
10 
11 #ifndef EIGEN_DIAGONAL_H
12 #define EIGEN_DIAGONAL_H
13 
14 namespace Eigen {
15 
16 /** \class Diagonal
17   * \ingroup Core_Module
18   *
19   * \brief Expression of a diagonal/subdiagonal/superdiagonal in a matrix
20   *
21   * \param MatrixType the type of the object in which we are taking a sub/main/super diagonal
22   * \param DiagIndex the index of the sub/super diagonal. The default is 0 and it means the main diagonal.
23   *              A positive value means a superdiagonal, a negative value means a subdiagonal.
24   *              You can also use DynamicIndex so the index can be set at runtime.
25   *
26   * The matrix is not required to be square.
27   *
28   * This class represents an expression of the main diagonal, or any sub/super diagonal
29   * of a square matrix. It is the return type of MatrixBase::diagonal() and MatrixBase::diagonal(Index) and most of the
30   * time this is the only way it is used.
31   *
32   * \sa MatrixBase::diagonal(), MatrixBase::diagonal(Index)
33   */
34 
35 namespace internal {
36 template<typename MatrixType, int DiagIndex>
37 struct traits<Diagonal<MatrixType,DiagIndex> >
38  : traits<MatrixType>
39 {
40   typedef typename ref_selector<MatrixType>::type MatrixTypeNested;
41   typedef typename remove_reference<MatrixTypeNested>::type _MatrixTypeNested;
42   typedef typename MatrixType::StorageKind StorageKind;
43   enum {
44     RowsAtCompileTime = (int(DiagIndex) == DynamicIndex || int(MatrixType::SizeAtCompileTime) == Dynamic) ? Dynamic
45                       : (EIGEN_PLAIN_ENUM_MIN(MatrixType::RowsAtCompileTime - EIGEN_PLAIN_ENUM_MAX(-DiagIndex, 0),
46                                               MatrixType::ColsAtCompileTime - EIGEN_PLAIN_ENUM_MAX( DiagIndex, 0))),
47     ColsAtCompileTime = 1,
48     MaxRowsAtCompileTime = int(MatrixType::MaxSizeAtCompileTime) == Dynamic ? Dynamic
49                          : DiagIndex == DynamicIndex ? EIGEN_SIZE_MIN_PREFER_FIXED(MatrixType::MaxRowsAtCompileTime,
50                                                                               MatrixType::MaxColsAtCompileTime)
51                          : (EIGEN_PLAIN_ENUM_MIN(MatrixType::MaxRowsAtCompileTime - EIGEN_PLAIN_ENUM_MAX(-DiagIndex, 0),
52                                                  MatrixType::MaxColsAtCompileTime - EIGEN_PLAIN_ENUM_MAX( DiagIndex, 0))),
53     MaxColsAtCompileTime = 1,
54     MaskLvalueBit = is_lvalue<MatrixType>::value ? LvalueBit : 0,
55     Flags = (unsigned int)_MatrixTypeNested::Flags & (RowMajorBit | MaskLvalueBit | DirectAccessBit) & ~RowMajorBit, // FIXME DirectAccessBit should not be handled by expressions
56     MatrixTypeOuterStride = outer_stride_at_compile_time<MatrixType>::ret,
57     InnerStrideAtCompileTime = MatrixTypeOuterStride == Dynamic ? Dynamic : MatrixTypeOuterStride+1,
58     OuterStrideAtCompileTime = 0
59   };
60 };
61 }
62 
63 template<typename MatrixType, int _DiagIndex> class Diagonal
64    : public internal::dense_xpr_base< Diagonal<MatrixType,_DiagIndex> >::type
65 {
66   public:
67 
68     enum { DiagIndex = _DiagIndex };
69     typedef typename internal::dense_xpr_base<Diagonal>::type Base;
70     EIGEN_DENSE_PUBLIC_INTERFACE(Diagonal)
71 
72     EIGEN_DEVICE_FUNC
73     explicit inline Diagonal(MatrixType& matrix, Index a_index = DiagIndex) : m_matrix(matrix), m_index(a_index)
74     {
75       eigen_assert( a_index <= m_matrix.cols() && -a_index <= m_matrix.rows() );
76     }
77 
78     EIGEN_INHERIT_ASSIGNMENT_OPERATORS(Diagonal)
79 
80     EIGEN_DEVICE_FUNC
81     inline Index rows() const
82     {
83       return m_index.value()<0 ? numext::mini<Index>(m_matrix.cols(),m_matrix.rows()+m_index.value())
84                                : numext::mini<Index>(m_matrix.rows(),m_matrix.cols()-m_index.value());
85     }
86 
87     EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR
88     inline Index cols() const EIGEN_NOEXCEPT { return 1; }
89 
90     EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR
91     inline Index innerStride() const EIGEN_NOEXCEPT {
92       return m_matrix.outerStride() + 1;
93     }
94 
95     EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR
96     inline Index outerStride() const EIGEN_NOEXCEPT { return 0; }
97 
98     typedef typename internal::conditional<
99                        internal::is_lvalue<MatrixType>::value,
100                        Scalar,
101                        const Scalar
102                      >::type ScalarWithConstIfNotLvalue;
103 
104     EIGEN_DEVICE_FUNC
105     inline ScalarWithConstIfNotLvalue* data() { return &(m_matrix.coeffRef(rowOffset(), colOffset())); }
106     EIGEN_DEVICE_FUNC
107     inline const Scalar* data() const { return &(m_matrix.coeffRef(rowOffset(), colOffset())); }
108 
109     EIGEN_DEVICE_FUNC
110     inline Scalar& coeffRef(Index row, Index)
111     {
112       EIGEN_STATIC_ASSERT_LVALUE(MatrixType)
113       return m_matrix.coeffRef(row+rowOffset(), row+colOffset());
114     }
115 
116     EIGEN_DEVICE_FUNC
117     inline const Scalar& coeffRef(Index row, Index) const
118     {
119       return m_matrix.coeffRef(row+rowOffset(), row+colOffset());
120     }
121 
122     EIGEN_DEVICE_FUNC
123     inline CoeffReturnType coeff(Index row, Index) const
124     {
125       return m_matrix.coeff(row+rowOffset(), row+colOffset());
126     }
127 
128     EIGEN_DEVICE_FUNC
129     inline Scalar& coeffRef(Index idx)
130     {
131       EIGEN_STATIC_ASSERT_LVALUE(MatrixType)
132       return m_matrix.coeffRef(idx+rowOffset(), idx+colOffset());
133     }
134 
135     EIGEN_DEVICE_FUNC
136     inline const Scalar& coeffRef(Index idx) const
137     {
138       return m_matrix.coeffRef(idx+rowOffset(), idx+colOffset());
139     }
140 
141     EIGEN_DEVICE_FUNC
142     inline CoeffReturnType coeff(Index idx) const
143     {
144       return m_matrix.coeff(idx+rowOffset(), idx+colOffset());
145     }
146 
147     EIGEN_DEVICE_FUNC
148     inline const typename internal::remove_all<typename MatrixType::Nested>::type&
149     nestedExpression() const
150     {
151       return m_matrix;
152     }
153 
154     EIGEN_DEVICE_FUNC
155     inline Index index() const
156     {
157       return m_index.value();
158     }
159 
160   protected:
161     typename internal::ref_selector<MatrixType>::non_const_type m_matrix;
162     const internal::variable_if_dynamicindex<Index, DiagIndex> m_index;
163 
164   private:
165     // some compilers may fail to optimize std::max etc in case of compile-time constants...
166     EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE EIGEN_CONSTEXPR
167     Index absDiagIndex() const EIGEN_NOEXCEPT { return m_index.value()>0 ? m_index.value() : -m_index.value(); }
168     EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE EIGEN_CONSTEXPR
169     Index rowOffset() const EIGEN_NOEXCEPT { return m_index.value()>0 ? 0 : -m_index.value(); }
170     EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE EIGEN_CONSTEXPR
171     Index colOffset() const EIGEN_NOEXCEPT { return m_index.value()>0 ? m_index.value() : 0; }
172     // trigger a compile-time error if someone try to call packet
173     template<int LoadMode> typename MatrixType::PacketReturnType packet(Index) const;
174     template<int LoadMode> typename MatrixType::PacketReturnType packet(Index,Index) const;
175 };
176 
177 /** \returns an expression of the main diagonal of the matrix \c *this
178   *
179   * \c *this is not required to be square.
180   *
181   * Example: \include MatrixBase_diagonal.cpp
182   * Output: \verbinclude MatrixBase_diagonal.out
183   *
184   * \sa class Diagonal */
185 template<typename Derived>
186 EIGEN_DEVICE_FUNC inline typename MatrixBase<Derived>::DiagonalReturnType
187 MatrixBase<Derived>::diagonal()
188 {
189   return DiagonalReturnType(derived());
190 }
191 
192 /** This is the const version of diagonal(). */
193 template<typename Derived>
194 EIGEN_DEVICE_FUNC inline typename MatrixBase<Derived>::ConstDiagonalReturnType
195 MatrixBase<Derived>::diagonal() const
196 {
197   return ConstDiagonalReturnType(derived());
198 }
199 
200 /** \returns an expression of the \a DiagIndex-th sub or super diagonal of the matrix \c *this
201   *
202   * \c *this is not required to be square.
203   *
204   * The template parameter \a DiagIndex represent a super diagonal if \a DiagIndex > 0
205   * and a sub diagonal otherwise. \a DiagIndex == 0 is equivalent to the main diagonal.
206   *
207   * Example: \include MatrixBase_diagonal_int.cpp
208   * Output: \verbinclude MatrixBase_diagonal_int.out
209   *
210   * \sa MatrixBase::diagonal(), class Diagonal */
211 template<typename Derived>
212 EIGEN_DEVICE_FUNC inline typename MatrixBase<Derived>::DiagonalDynamicIndexReturnType
213 MatrixBase<Derived>::diagonal(Index index)
214 {
215   return DiagonalDynamicIndexReturnType(derived(), index);
216 }
217 
218 /** This is the const version of diagonal(Index). */
219 template<typename Derived>
220 EIGEN_DEVICE_FUNC inline typename MatrixBase<Derived>::ConstDiagonalDynamicIndexReturnType
221 MatrixBase<Derived>::diagonal(Index index) const
222 {
223   return ConstDiagonalDynamicIndexReturnType(derived(), index);
224 }
225 
226 /** \returns an expression of the \a DiagIndex-th sub or super diagonal of the matrix \c *this
227   *
228   * \c *this is not required to be square.
229   *
230   * The template parameter \a DiagIndex represent a super diagonal if \a DiagIndex > 0
231   * and a sub diagonal otherwise. \a DiagIndex == 0 is equivalent to the main diagonal.
232   *
233   * Example: \include MatrixBase_diagonal_template_int.cpp
234   * Output: \verbinclude MatrixBase_diagonal_template_int.out
235   *
236   * \sa MatrixBase::diagonal(), class Diagonal */
237 template<typename Derived>
238 template<int Index_>
239 EIGEN_DEVICE_FUNC
240 inline typename MatrixBase<Derived>::template DiagonalIndexReturnType<Index_>::Type
241 MatrixBase<Derived>::diagonal()
242 {
243   return typename DiagonalIndexReturnType<Index_>::Type(derived());
244 }
245 
246 /** This is the const version of diagonal<int>(). */
247 template<typename Derived>
248 template<int Index_>
249 EIGEN_DEVICE_FUNC
250 inline typename MatrixBase<Derived>::template ConstDiagonalIndexReturnType<Index_>::Type
251 MatrixBase<Derived>::diagonal() const
252 {
253   return typename ConstDiagonalIndexReturnType<Index_>::Type(derived());
254 }
255 
256 } // end namespace Eigen
257 
258 #endif // EIGEN_DIAGONAL_H
259