xref: /aosp_15_r20/external/eigen/Eigen/src/Core/Select.h (revision bf2c37156dfe67e5dfebd6d394bad8b2ab5804d4)
1 // This file is part of Eigen, a lightweight C++ template library
2 // for linear algebra.
3 //
4 // Copyright (C) 2008-2010 Gael Guennebaud <[email protected]>
5 //
6 // This Source Code Form is subject to the terms of the Mozilla
7 // Public License v. 2.0. If a copy of the MPL was not distributed
8 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
9 
10 #ifndef EIGEN_SELECT_H
11 #define EIGEN_SELECT_H
12 
13 namespace Eigen {
14 
15 /** \class Select
16   * \ingroup Core_Module
17   *
18   * \brief Expression of a coefficient wise version of the C++ ternary operator ?:
19   *
20   * \param ConditionMatrixType the type of the \em condition expression which must be a boolean matrix
21   * \param ThenMatrixType the type of the \em then expression
22   * \param ElseMatrixType the type of the \em else expression
23   *
24   * This class represents an expression of a coefficient wise version of the C++ ternary operator ?:.
25   * It is the return type of DenseBase::select() and most of the time this is the only way it is used.
26   *
27   * \sa DenseBase::select(const DenseBase<ThenDerived>&, const DenseBase<ElseDerived>&) const
28   */
29 
30 namespace internal {
31 template<typename ConditionMatrixType, typename ThenMatrixType, typename ElseMatrixType>
32 struct traits<Select<ConditionMatrixType, ThenMatrixType, ElseMatrixType> >
33  : traits<ThenMatrixType>
34 {
35   typedef typename traits<ThenMatrixType>::Scalar Scalar;
36   typedef Dense StorageKind;
37   typedef typename traits<ThenMatrixType>::XprKind XprKind;
38   typedef typename ConditionMatrixType::Nested ConditionMatrixNested;
39   typedef typename ThenMatrixType::Nested ThenMatrixNested;
40   typedef typename ElseMatrixType::Nested ElseMatrixNested;
41   enum {
42     RowsAtCompileTime = ConditionMatrixType::RowsAtCompileTime,
43     ColsAtCompileTime = ConditionMatrixType::ColsAtCompileTime,
44     MaxRowsAtCompileTime = ConditionMatrixType::MaxRowsAtCompileTime,
45     MaxColsAtCompileTime = ConditionMatrixType::MaxColsAtCompileTime,
46     Flags = (unsigned int)ThenMatrixType::Flags & ElseMatrixType::Flags & RowMajorBit
47   };
48 };
49 }
50 
51 template<typename ConditionMatrixType, typename ThenMatrixType, typename ElseMatrixType>
52 class Select : public internal::dense_xpr_base< Select<ConditionMatrixType, ThenMatrixType, ElseMatrixType> >::type,
53                internal::no_assignment_operator
54 {
55   public:
56 
57     typedef typename internal::dense_xpr_base<Select>::type Base;
58     EIGEN_DENSE_PUBLIC_INTERFACE(Select)
59 
60     inline EIGEN_DEVICE_FUNC
61     Select(const ConditionMatrixType& a_conditionMatrix,
62            const ThenMatrixType& a_thenMatrix,
63            const ElseMatrixType& a_elseMatrix)
64       : m_condition(a_conditionMatrix), m_then(a_thenMatrix), m_else(a_elseMatrix)
65     {
66       eigen_assert(m_condition.rows() == m_then.rows() && m_condition.rows() == m_else.rows());
67       eigen_assert(m_condition.cols() == m_then.cols() && m_condition.cols() == m_else.cols());
68     }
69 
70     inline EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR
71     Index rows() const EIGEN_NOEXCEPT { return m_condition.rows(); }
72     inline EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR
73     Index cols() const EIGEN_NOEXCEPT { return m_condition.cols(); }
74 
75     inline EIGEN_DEVICE_FUNC
76     const Scalar coeff(Index i, Index j) const
77     {
78       if (m_condition.coeff(i,j))
79         return m_then.coeff(i,j);
80       else
81         return m_else.coeff(i,j);
82     }
83 
84     inline EIGEN_DEVICE_FUNC
85     const Scalar coeff(Index i) const
86     {
87       if (m_condition.coeff(i))
88         return m_then.coeff(i);
89       else
90         return m_else.coeff(i);
91     }
92 
93     inline EIGEN_DEVICE_FUNC const ConditionMatrixType& conditionMatrix() const
94     {
95       return m_condition;
96     }
97 
98     inline EIGEN_DEVICE_FUNC const ThenMatrixType& thenMatrix() const
99     {
100       return m_then;
101     }
102 
103     inline EIGEN_DEVICE_FUNC const ElseMatrixType& elseMatrix() const
104     {
105       return m_else;
106     }
107 
108   protected:
109     typename ConditionMatrixType::Nested m_condition;
110     typename ThenMatrixType::Nested m_then;
111     typename ElseMatrixType::Nested m_else;
112 };
113 
114 
115 /** \returns a matrix where each coefficient (i,j) is equal to \a thenMatrix(i,j)
116   * if \c *this(i,j), and \a elseMatrix(i,j) otherwise.
117   *
118   * Example: \include MatrixBase_select.cpp
119   * Output: \verbinclude MatrixBase_select.out
120   *
121   * \sa class Select
122   */
123 template<typename Derived>
124 template<typename ThenDerived,typename ElseDerived>
125 inline EIGEN_DEVICE_FUNC const Select<Derived,ThenDerived,ElseDerived>
126 DenseBase<Derived>::select(const DenseBase<ThenDerived>& thenMatrix,
127                             const DenseBase<ElseDerived>& elseMatrix) const
128 {
129   return Select<Derived,ThenDerived,ElseDerived>(derived(), thenMatrix.derived(), elseMatrix.derived());
130 }
131 
132 /** Version of DenseBase::select(const DenseBase&, const DenseBase&) with
133   * the \em else expression being a scalar value.
134   *
135   * \sa DenseBase::select(const DenseBase<ThenDerived>&, const DenseBase<ElseDerived>&) const, class Select
136   */
137 template<typename Derived>
138 template<typename ThenDerived>
139 inline EIGEN_DEVICE_FUNC const Select<Derived,ThenDerived, typename ThenDerived::ConstantReturnType>
140 DenseBase<Derived>::select(const DenseBase<ThenDerived>& thenMatrix,
141                            const typename ThenDerived::Scalar& elseScalar) const
142 {
143   return Select<Derived,ThenDerived,typename ThenDerived::ConstantReturnType>(
144     derived(), thenMatrix.derived(), ThenDerived::Constant(rows(),cols(),elseScalar));
145 }
146 
147 /** Version of DenseBase::select(const DenseBase&, const DenseBase&) with
148   * the \em then expression being a scalar value.
149   *
150   * \sa DenseBase::select(const DenseBase<ThenDerived>&, const DenseBase<ElseDerived>&) const, class Select
151   */
152 template<typename Derived>
153 template<typename ElseDerived>
154 inline EIGEN_DEVICE_FUNC const Select<Derived, typename ElseDerived::ConstantReturnType, ElseDerived >
155 DenseBase<Derived>::select(const typename ElseDerived::Scalar& thenScalar,
156                            const DenseBase<ElseDerived>& elseMatrix) const
157 {
158   return Select<Derived,typename ElseDerived::ConstantReturnType,ElseDerived>(
159     derived(), ElseDerived::Constant(rows(),cols(),thenScalar), elseMatrix.derived());
160 }
161 
162 } // end namespace Eigen
163 
164 #endif // EIGEN_SELECT_H
165