Gradient.h

Go to the documentation of this file.

#pragma once

#include <Eigen/Core>
#include <Eigen/Dense>
#include <unsupported/Eigen/AutoDiff>

#include "drake/common/drake_assert.h"

namespace Drake {
// todo: recursive template to get arbitrary gradient order

// note: tried using template default values (e.g. Eigen::Dynamic), but they
// didn't seem to work on my mac clang
template <int num_vars>
using TaylorVard = Eigen::AutoDiffScalar<Eigen::Matrix<double, num_vars, 1> >;
template <int num_vars, int rows>
using TaylorVecd = Eigen::Matrix<TaylorVard<num_vars>, rows, 1>;
template <int num_vars, int rows, int cols>
using TaylorMatd = Eigen::Matrix<TaylorVard<num_vars>, rows, cols>;

typedef TaylorVard<Eigen::Dynamic> TaylorVarXd;
typedef TaylorVecd<Eigen::Dynamic, Eigen::Dynamic> TaylorVecXd;
typedef TaylorMatd<Eigen::Dynamic, Eigen::Dynamic, Eigen::Dynamic> TaylorMatXd;

template <typename Derived, int Nq>
using AutoDiffMatrixType = Eigen::Matrix<
    Eigen::AutoDiffScalar<Eigen::Matrix<typename Derived::Scalar, Nq, 1> >,
    Derived::RowsAtCompileTime, Derived::ColsAtCompileTime, 0,
    Derived::MaxRowsAtCompileTime, Derived::MaxColsAtCompileTime>;

template <typename Derived, typename DerivedGradient, typename DerivedAutoDiff>
void initializeAutoDiffGivenGradientMatrix(
    const Eigen::MatrixBase<Derived> &val,
    const Eigen::MatrixBase<DerivedGradient> &gradient,
    Eigen::MatrixBase<DerivedAutoDiff> &auto_diff_matrix) {
  static_assert(static_cast<int>(Derived::SizeAtCompileTime) ==
                    static_cast<int>(DerivedGradient::RowsAtCompileTime),
                "gradient has wrong number of rows at compile time");
  DRAKE_ASSERT(val.size() == gradient.rows() &&
               "gradient has wrong number of rows at runtime");
  typedef AutoDiffMatrixType<Derived, DerivedGradient::ColsAtCompileTime>
      ExpectedAutoDiffType;
  static_assert(static_cast<int>(ExpectedAutoDiffType::RowsAtCompileTime) ==
                    static_cast<int>(DerivedAutoDiff::RowsAtCompileTime),
                "auto diff matrix has wrong number of rows at compile time");
  static_assert(static_cast<int>(ExpectedAutoDiffType::ColsAtCompileTime) ==
                    static_cast<int>(DerivedAutoDiff::ColsAtCompileTime),
                "auto diff matrix has wrong number of columns at compile time");
  static_assert(std::is_same<typename DerivedAutoDiff::Scalar,
                             typename ExpectedAutoDiffType::Scalar>::value,
                "wrong auto diff scalar type");

  typedef typename Eigen::MatrixBase<DerivedGradient>::Index Index;
  auto_diff_matrix.resize(val.rows(), val.cols());
  auto num_derivs = gradient.cols();
  for (Index row = 0; row < auto_diff_matrix.size(); row++) {
    auto_diff_matrix(row).value() = val(row);
    auto_diff_matrix(row).derivatives().resize(num_derivs, 1);
    auto_diff_matrix(row).derivatives() = gradient.row(row).transpose();
  }
}

template <typename Derived, typename DerivedGradient>
AutoDiffMatrixType<Derived, DerivedGradient::ColsAtCompileTime>
initializeAutoDiffGivenGradientMatrix(
    const Eigen::MatrixBase<Derived> &val,
    const Eigen::MatrixBase<DerivedGradient> &gradient) {
  AutoDiffMatrixType<Derived, DerivedGradient::ColsAtCompileTime> ret(
      val.rows(), val.cols());
  initializeAutoDiffGivenGradientMatrix(val, gradient, ret);
  return ret;
}

template <typename Derived, typename DerivedAutoDiff>
void initializeAutoDiff(const Eigen::MatrixBase<Derived> &val,
                        Eigen::MatrixBase<DerivedAutoDiff> &auto_diff_matrix,
                        Eigen::DenseIndex num_derivatives = Eigen::Dynamic,
                        Eigen::DenseIndex deriv_num_start = 0) {
  using ADScalar = typename DerivedAutoDiff::Scalar;
  static_assert(static_cast<int>(Derived::RowsAtCompileTime) ==
                    static_cast<int>(DerivedAutoDiff::RowsAtCompileTime),
                "auto diff matrix has wrong number of rows at compile time");
  static_assert(static_cast<int>(Derived::ColsAtCompileTime) ==
                    static_cast<int>(DerivedAutoDiff::ColsAtCompileTime),
                "auto diff matrix has wrong number of columns at compile time");

  if (num_derivatives == Eigen::Dynamic) num_derivatives = val.size();

  auto_diff_matrix.resize(val.rows(), val.cols());
  Eigen::DenseIndex deriv_num = deriv_num_start;
  for (Eigen::DenseIndex i = 0; i < val.size(); i++) {
    auto_diff_matrix(i) = ADScalar(val(i), num_derivatives, deriv_num++);
  }
}

template <int Nq = Eigen::Dynamic, typename Derived>
AutoDiffMatrixType<Derived, Nq> initializeAutoDiff(
    const Eigen::MatrixBase<Derived> &mat,
    Eigen::DenseIndex num_derivatives = -1,
    Eigen::DenseIndex deriv_num_start = 0) {
  if (num_derivatives == -1) num_derivatives = mat.size();

  AutoDiffMatrixType<Derived, Nq> ret(mat.rows(), mat.cols());
  initializeAutoDiff(mat, ret, num_derivatives, deriv_num_start);
  return ret;
}

// TODO(tkoolen) begin: move to different location?
namespace internal {
template <typename Head, typename... Tail>
struct TotalSizeAtCompileTime {
  static constexpr int eval() {
    return Head::SizeAtCompileTime == Eigen::Dynamic ||
                   TotalSizeAtCompileTime<Tail...>::eval() == Eigen::Dynamic
               ? Eigen::Dynamic
               : Head::SizeAtCompileTime +
                     TotalSizeAtCompileTime<Tail...>::eval();
  }
};

template <typename Head>
struct TotalSizeAtCompileTime<Head> {
  static constexpr int eval() { return Head::SizeAtCompileTime; }
};
}

template <typename... Args>
constexpr int totalSizeAtCompileTime() {
  return internal::TotalSizeAtCompileTime<Args...>::eval();
}

constexpr Eigen::DenseIndex totalSizeAtRunTime() { return 0; }

template <typename Head, typename... Tail>
Eigen::DenseIndex totalSizeAtRunTime(const Eigen::MatrixBase<Head> &head,
                                     const Tail &... tail) {
  return head.size() + totalSizeAtRunTime(tail...);
}
// TODO(tkoolen) end: move to different location?

namespace internal {
template <size_t Index>
struct InitializeAutoDiffTupleHelper {
  template <typename... ValueTypes, typename... AutoDiffTypes>
  static void run(const std::tuple<ValueTypes...> &values,
                  std::tuple<AutoDiffTypes...> &auto_diffs,
                  Eigen::DenseIndex num_derivatives,
                  Eigen::DenseIndex deriv_num_start) {
    constexpr size_t tuple_index = sizeof...(AutoDiffTypes)-Index;
    const auto &value = std::get<tuple_index>(values);
    auto &auto_diff = std::get<tuple_index>(auto_diffs);
    auto_diff.resize(value.rows(), value.cols());
    initializeAutoDiff(value, auto_diff, num_derivatives, deriv_num_start);
    InitializeAutoDiffTupleHelper<Index - 1>::run(
        values, auto_diffs, num_derivatives, deriv_num_start + value.size());
  }
};

template <>
struct InitializeAutoDiffTupleHelper<0> {
  template <typename... ValueTypes, typename... AutoDiffTypes>
  static void run(const std::tuple<ValueTypes...> &values,
                  const std::tuple<AutoDiffTypes...> &auto_diffs,
                  Eigen::DenseIndex num_derivatives,
                  Eigen::DenseIndex deriv_num_start) {
    // empty
  }
};
}

template <typename... Deriveds>
std::tuple<
    AutoDiffMatrixType<Deriveds, totalSizeAtCompileTime<Deriveds...>()>...>
initializeAutoDiffTuple(const Eigen::MatrixBase<Deriveds> &... args) {
  Eigen::DenseIndex dynamic_num_derivs = totalSizeAtRunTime(args...);
  std::tuple<
      AutoDiffMatrixType<Deriveds, totalSizeAtCompileTime<Deriveds...>()>...>
      ret(AutoDiffMatrixType<Deriveds, totalSizeAtCompileTime<Deriveds...>()>(
          args.rows(), args.cols())...);
  auto values = std::forward_as_tuple(args...);
  internal::InitializeAutoDiffTupleHelper<sizeof...(args)>::run(
      values, ret, dynamic_num_derivs, 0);
  return ret;
}
}