trace_quad_form.hpp

Go to the documentation of this file.

#ifndef STAN__AGRAD__REV__MATRIX__TRACE_QUAD_FORM_HPP
#define STAN__AGRAD__REV__MATRIX__TRACE_QUAD_FORM_HPP

#include <boost/utility/enable_if.hpp>
#include <boost/type_traits.hpp>
#include <stan/math/matrix/Eigen.hpp>
#include <stan/math/matrix/typedefs.hpp>
#include <stan/agrad/rev/matrix/value_of.hpp>
#include <stan/agrad/rev/var.hpp>
#include <stan/agrad/rev/vari.hpp>
#include <stan/agrad/rev/matrix/typedefs.hpp>
#include <stan/math/matrix/trace_quad_form.hpp>
#include <stan/math/error_handling/matrix/check_multiplicable.hpp>
#include <stan/math/error_handling/matrix/check_square.hpp>

namespace stan {
  namespace agrad {
    namespace {
      template<typename TA,int RA,int CA,typename TB,int RB,int CB>
      class trace_quad_form_vari_alloc : public chainable_alloc {
      public:
        trace_quad_form_vari_alloc(const Eigen::Matrix<TA,RA,CA> &A,
                                   const Eigen::Matrix<TB,RB,CB> &B)
        : A_(A), B_(B)
        { }
        
        double compute() {
          return stan::math::trace_quad_form(value_of(A_),
                                             value_of(B_));
        }
        
        Eigen::Matrix<TA,RA,CA>  A_;
        Eigen::Matrix<TB,RB,CB>  B_;
      };
      
      template<typename TA,int RA,int CA,typename TB,int RB,int CB>
      class trace_quad_form_vari : public vari {
      protected:
        static inline void chainA(Eigen::Matrix<double,RA,CA> &A, 
                                  const Eigen::Matrix<double,RB,CB> &Bd,
                                  const double &adjC) {}
        static inline void chainB(Eigen::Matrix<double,RB,CB> &B, 
                                  const Eigen::Matrix<double,RA,CA> &Ad,
                                  const Eigen::Matrix<double,RB,CB> &Bd,
                                  const double &adjC) {}
        
        static inline void chainA(Eigen::Matrix<var,RA,CA> &A, 
                                  const Eigen::Matrix<double,RB,CB> &Bd,
                                  const double &adjC)
        {
          int i,j;
          Eigen::Matrix<double,RA,CA>     adjA(adjC*Bd*Bd.transpose());
          for (j = 0; j < A.cols(); j++)
            for (i = 0; i < A.rows(); i++)
              A(i,j).vi_->adj_ += adjA(i,j);
        }
        static inline void chainB(Eigen::Matrix<var,RB,CB> &B, 
                                  const Eigen::Matrix<double,RA,CA> &Ad,
                                  const Eigen::Matrix<double,RB,CB> &Bd,
                                  const double &adjC)
        {
          int i,j;
          Eigen::Matrix<double,RA,CA>     adjB(adjC*(Ad + Ad.transpose())*Bd);
          for (j = 0; j < B.cols(); j++)
            for (i = 0; i < B.rows(); i++)
              B(i,j).vi_->adj_ += adjB(i,j);
        }
        
        inline void chainAB(Eigen::Matrix<TA,RA,CA> &A,
                            Eigen::Matrix<TB,RB,CB> &B,
                            const Eigen::Matrix<double,RA,CA> &Ad,
                            const Eigen::Matrix<double,RB,CB> &Bd,
                            const double &adjC)
        {
          chainA(A,Bd,adjC);
          chainB(B,Ad,Bd,adjC);
        }
        

      public:
        trace_quad_form_vari(trace_quad_form_vari_alloc<TA,RA,CA,TB,RB,CB> *impl)
        : vari(impl->compute()), _impl(impl) { }
        
        virtual void chain() {
          chainAB(_impl->A_, _impl->B_,
                  value_of(_impl->A_), value_of(_impl->B_),
                  adj_);
        };

        trace_quad_form_vari_alloc<TA,RA,CA,TB,RB,CB> *_impl;
      };
    }
    
    template<typename TA,int RA,int CA,typename TB,int RB,int CB>
    inline typename
    boost::enable_if_c< boost::is_same<TA,var>::value ||
                        boost::is_same<TB,var>::value,
                        var >::type
    trace_quad_form(const Eigen::Matrix<TA,RA,CA> &A,
                    const Eigen::Matrix<TB,RB,CB> &B)
    {
      stan::math::check_square("trace_quad_form(%1%)",A,"A",(double*)0);
      stan::math::check_multiplicable("trace_quad_form(%1%)",A,"A",
                                      B,"B",(double*)0);
      
      trace_quad_form_vari_alloc<TA,RA,CA,TB,RB,CB> *baseVari = new trace_quad_form_vari_alloc<TA,RA,CA,TB,RB,CB>(A,B);
      
      return var(new trace_quad_form_vari<TA,RA,CA,TB,RB,CB>(baseVari));
    }
  }
}

#endif