Chris@49: // Copyright (C) 2008-2012 NICTA (www.nicta.com.au)
Chris@49: // Copyright (C) 2008-2012 Conrad Sanderson
Chris@49: // Copyright (C) 2012 Ryan Curtin
Chris@49: // 
Chris@49: // This Source Code Form is subject to the terms of the Mozilla Public
Chris@49: // License, v. 2.0. If a copy of the MPL was not distributed with this
Chris@49: // file, You can obtain one at http://mozilla.org/MPL/2.0/.
Chris@49: 
Chris@49: 
Chris@49: //! \addtogroup fn_trace
Chris@49: //! @{
Chris@49: 
Chris@49: 
Chris@49: //! Immediate trace (sum of diagonal elements) of a square dense matrix
Chris@49: template<typename T1>
Chris@49: arma_hot
Chris@49: arma_warn_unused
Chris@49: inline
Chris@49: typename enable_if2<is_arma_type<T1>::value, typename T1::elem_type>::result
Chris@49: trace(const T1& X)
Chris@49:   {
Chris@49:   arma_extra_debug_sigprint();
Chris@49:   
Chris@49:   typedef typename T1::elem_type eT;
Chris@49:   
Chris@49:   const Proxy<T1> A(X);
Chris@49:   
Chris@49:   arma_debug_check( (A.get_n_rows() != A.get_n_cols()), "trace(): matrix must be square sized" );
Chris@49:   
Chris@49:   const uword N = A.get_n_rows();
Chris@49:   
Chris@49:   eT val1 = eT(0);
Chris@49:   eT val2 = eT(0);
Chris@49:   
Chris@49:   uword i,j;
Chris@49:   for(i=0, j=1; j<N; i+=2, j+=2)
Chris@49:     {
Chris@49:     val1 += A.at(i,i);
Chris@49:     val2 += A.at(j,j);
Chris@49:     }
Chris@49:   
Chris@49:   if(i < N)
Chris@49:     {
Chris@49:     val1 += A.at(i,i);
Chris@49:     }
Chris@49:   
Chris@49:   return val1 + val2;
Chris@49:   }
Chris@49: 
Chris@49: 
Chris@49: 
Chris@49: template<typename T1>
Chris@49: arma_hot
Chris@49: arma_warn_unused
Chris@49: inline
Chris@49: typename T1::elem_type
Chris@49: trace(const Op<T1, op_diagmat>& X)
Chris@49:   {
Chris@49:   arma_extra_debug_sigprint();
Chris@49:   
Chris@49:   typedef typename T1::elem_type eT;
Chris@49:   
Chris@49:   const diagmat_proxy<T1> A(X.m);
Chris@49:   
Chris@49:   const uword N = A.n_elem;
Chris@49:   
Chris@49:   eT val = eT(0);
Chris@49:   
Chris@49:   for(uword i=0; i<N; ++i)
Chris@49:     {
Chris@49:     val += A[i];
Chris@49:     }
Chris@49:   
Chris@49:   return val;
Chris@49:   }
Chris@49: 
Chris@49: 
Chris@49: 
Chris@49: //! speedup for trace(A*B), where the result of A*B is a square sized matrix
Chris@49: template<typename T1, typename T2>
Chris@49: arma_hot
Chris@49: inline
Chris@49: typename T1::elem_type
Chris@49: trace_mul_unwrap(const T1& XA, const T2& XB)
Chris@49:   {
Chris@49:   arma_extra_debug_sigprint();
Chris@49:   
Chris@49:   typedef typename T1::elem_type eT;
Chris@49:   
Chris@49:   const Proxy<T1>    PA(XA);
Chris@49:   const unwrap<T2> tmpB(XB);
Chris@49:   
Chris@49:   const Mat<eT>& B = tmpB.M;
Chris@49:   
Chris@49:   arma_debug_assert_mul_size(PA.get_n_rows(), PA.get_n_cols(), B.n_rows, B.n_cols, "matrix multiplication");
Chris@49:   
Chris@49:   arma_debug_check( (PA.get_n_rows() != B.n_cols), "trace(): matrix must be square sized" );
Chris@49:   
Chris@49:   const uword N1 = PA.get_n_rows();   // equivalent to B.n_cols, due to square size requirements
Chris@49:   const uword N2 = PA.get_n_cols();   // equivalent to B.n_rows, due to matrix multiplication requirements
Chris@49:   
Chris@49:   eT val = eT(0);
Chris@49:   
Chris@49:   for(uword i=0; i<N1; ++i)
Chris@49:     {
Chris@49:     const eT* B_colmem = B.colptr(i);
Chris@49:     
Chris@49:     eT acc1 = eT(0);
Chris@49:     eT acc2 = eT(0);
Chris@49:     
Chris@49:     uword j,k;
Chris@49:     for(j=0, k=1; k < N2; j+=2, k+=2)
Chris@49:       {
Chris@49:       const eT tmp_j = B_colmem[j];
Chris@49:       const eT tmp_k = B_colmem[k];
Chris@49:       
Chris@49:       acc1 += PA.at(i,j) * tmp_j;
Chris@49:       acc2 += PA.at(i,k) * tmp_k;
Chris@49:       }
Chris@49:     
Chris@49:     if(j < N2)
Chris@49:       {
Chris@49:       acc1 += PA.at(i,j) * B_colmem[j];
Chris@49:       }
Chris@49:     
Chris@49:     val += (acc1 + acc2);
Chris@49:     }
Chris@49:   
Chris@49:   return val;
Chris@49:   }
Chris@49: 
Chris@49: 
Chris@49: 
Chris@49: //! speedup for trace(A*B), where the result of A*B is a square sized matrix
Chris@49: template<typename T1, typename T2>
Chris@49: arma_hot
Chris@49: inline
Chris@49: typename T1::elem_type
Chris@49: trace_mul_proxy(const T1& XA, const T2& XB)
Chris@49:   {
Chris@49:   arma_extra_debug_sigprint();
Chris@49:   
Chris@49:   typedef typename T1::elem_type eT;
Chris@49:   
Chris@49:   const Proxy<T1> PA(XA);
Chris@49:   const Proxy<T2> PB(XB);
Chris@49:   
Chris@49:   if(is_Mat<typename Proxy<T2>::stored_type>::value == true)
Chris@49:     {
Chris@49:     return trace_mul_unwrap(PA.Q, PB.Q);
Chris@49:     }
Chris@49:   
Chris@49:   arma_debug_assert_mul_size(PA.get_n_rows(), PA.get_n_cols(), PB.get_n_rows(), PB.get_n_cols(), "matrix multiplication");
Chris@49:   
Chris@49:   arma_debug_check( (PA.get_n_rows() != PB.get_n_cols()), "trace(): matrix must be square sized" );
Chris@49:   
Chris@49:   const uword N1 = PA.get_n_rows();   // equivalent to PB.get_n_cols(), due to square size requirements
Chris@49:   const uword N2 = PA.get_n_cols();   // equivalent to PB.get_n_rows(), due to matrix multiplication requirements
Chris@49:   
Chris@49:   eT val = eT(0);
Chris@49:   
Chris@49:   for(uword i=0; i<N1; ++i)
Chris@49:     {
Chris@49:     eT acc1 = eT(0);
Chris@49:     eT acc2 = eT(0);
Chris@49:     
Chris@49:     uword j,k;
Chris@49:     for(j=0, k=1; k < N2; j+=2, k+=2)
Chris@49:       {
Chris@49:       const eT tmp_j = PB.at(j,i);
Chris@49:       const eT tmp_k = PB.at(k,i);
Chris@49:       
Chris@49:       acc1 += PA.at(i,j) * tmp_j;
Chris@49:       acc2 += PA.at(i,k) * tmp_k;
Chris@49:       }
Chris@49:     
Chris@49:     if(j < N2)
Chris@49:       {
Chris@49:       acc1 += PA.at(i,j) * PB.at(j,i);
Chris@49:       }
Chris@49:     
Chris@49:     val += (acc1 + acc2);
Chris@49:     }
Chris@49:   
Chris@49:   return val;
Chris@49:   }
Chris@49: 
Chris@49: 
Chris@49: 
Chris@49: //! speedup for trace(A*B), where the result of A*B is a square sized matrix
Chris@49: template<typename T1, typename T2>
Chris@49: arma_hot
Chris@49: arma_warn_unused
Chris@49: inline
Chris@49: typename T1::elem_type
Chris@49: trace(const Glue<T1, T2, glue_times>& X)
Chris@49:   {
Chris@49:   arma_extra_debug_sigprint();
Chris@49:   
Chris@49:   return (is_Mat<T2>::value) ? trace_mul_unwrap(X.A, X.B) : trace_mul_proxy(X.A, X.B);
Chris@49:   }
Chris@49: 
Chris@49: 
Chris@49: 
Chris@49: //! trace of sparse object
Chris@49: template<typename T1>
Chris@49: arma_hot
Chris@49: arma_warn_unused
Chris@49: inline
Chris@49: typename enable_if2<is_arma_sparse_type<T1>::value, typename T1::elem_type>::result
Chris@49: trace(const T1& x)
Chris@49:   {
Chris@49:   arma_extra_debug_sigprint();
Chris@49:   
Chris@49:   const SpProxy<T1> p(x);
Chris@49:   
Chris@49:   arma_debug_check( (p.get_n_rows() != p.get_n_cols()), "trace(): matrix must be square sized" );
Chris@49:   
Chris@49:   typedef typename T1::elem_type eT;
Chris@49:   
Chris@49:   eT result = eT(0);
Chris@49:   
Chris@49:   typename SpProxy<T1>::const_iterator_type it     = p.begin();
Chris@49:   typename SpProxy<T1>::const_iterator_type it_end = p.end();
Chris@49:   
Chris@49:   while(it != it_end)
Chris@49:     {
Chris@49:     if(it.row() == it.col())
Chris@49:       {
Chris@49:       result += (*it);
Chris@49:       }
Chris@49:     
Chris@49:     ++it;
Chris@49:     }
Chris@49:   
Chris@49:   return result;
Chris@49:   }
Chris@49: 
Chris@49: 
Chris@49: 
Chris@49: //! @}