apply_vector_unary.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_PRIM_FUNCTOR_APPLY_VECTOR_UNARY_HPP
#define STAN_MATH_PRIM_FUNCTOR_APPLY_VECTOR_UNARY_HPP
 
#include <stan/math/prim/fun/Eigen.hpp>
#include <stan/math/prim/fun/as_column_vector_or_scalar.hpp>
#include <stan/math/prim/meta/is_stan_scalar.hpp>
#include <stan/math/prim/meta/is_container.hpp>
#include <stan/math/prim/meta/is_eigen_matrix_base.hpp>
#include <stan/math/prim/meta/plain_type.hpp>
#include <stan/math/prim/meta/require_generics.hpp>
#include <vector>
 
namespace stan {
namespace math {
// Forward declaration to allow specializations
template <typename T, typename Enable = void>
struct apply_vector_unary {};
 
template <typename T>
struct apply_vector_unary<T, require_eigen_t<T>> {
  template <typename F, typename T2 = T,
            require_t<is_eigen_matrix_base<plain_type_t<T2>>>* = nullptr>
  static inline auto apply(T2&& x, F&& f) {
    if constexpr (is_eigen_array<decltype(f(x))>::value) {
      return make_holder(
          [](auto&& xx) { return std::forward<decltype(xx)>(xx).matrix(); },
          make_holder(std::forward<F>(f), std::forward<T2>(x)));
    } else {
      return make_holder(std::forward<F>(f), std::forward<T2>(x));
    }
  }
 
  template <typename F, typename T2 = T,
            require_t<is_eigen_array<plain_type_t<T2>>>* = nullptr>
  static inline auto apply(T2&& x, F&& f) {
    if constexpr (is_eigen_array<decltype(f(x))>::value) {
      return make_holder(std::forward<F>(f), std::forward<T2>(x));
    } else {
      return make_holder(
          [](auto&& xx) { return std::forward<decltype(xx)>(xx).array(); },
          make_holder(std::forward<F>(f), std::forward<T2>(x)));
    }
  }
 
  template <typename F, typename T2 = T,
            require_t<is_eigen_matrix_base<plain_type_t<T2>>>* = nullptr>
  static inline auto apply_no_holder(T2& x, F&& f) {
    return std::forward<F>(f)(std::forward<T2>(x));
  }
 
  template <typename F, typename T2 = T,
            require_t<is_eigen_array<plain_type_t<T2>>>* = nullptr>
  static inline auto apply_no_holder(T2&& x, F&& f) {
    return std::forward<F>(f)(std::forward<T2>(x));
  }
 
  template <typename T2, typename F>
  static inline auto reduce(T2&& x, F&& f) {
    return make_holder(std::forward<F>(f), std::forward<T2>(x));
  }
};
 
template <typename T>
struct apply_vector_unary<T, require_std_vector_vt<is_stan_scalar, T>> {
  using T_vt = value_type_t<T>;
  using T_map = typename Eigen::Map<const Eigen::Matrix<T_vt, -1, 1>>;
 
  template <typename T2, typename F>
  static inline auto apply(T2&& x, F&& f) {
    using T_return = value_type_t<decltype(
        f(as_column_vector_or_scalar(std::forward<T2>(x))))>;
    std::vector<T_return> result(x.size());
    Eigen::Map<Eigen::Matrix<T_return, -1, 1>>(result.data(), result.size())
        = std::forward<F>(f)(as_column_vector_or_scalar(std::forward<T2>(x)))
              .matrix();
    return result;
  }
 
  template <typename T2, typename F>
  static inline auto apply_no_holder(T2&& x, F&& f) {
    return apply(std::forward<T2>(x), std::forward<F>(f));
  }
 
  template <typename T2, typename F>
  static inline auto reduce(T2&& x, F&& f) {
    return apply_vector_unary<T_map>::reduce(
        as_column_vector_or_scalar(std::forward<T2>(x)), std::forward<F>(f));
  }
};
 
template <typename T>
struct apply_vector_unary<
    T, require_std_vector_vt<is_container_or_var_matrix, T>> {
  using T_vt = value_type_t<T>;
 
  template <typename F>
  static inline auto apply(const T& x, const F& f) {
    using T_return
        = plain_type_t<decltype(apply_vector_unary<T_vt>::apply(x[0], f))>;
    std::vector<T_return> result(x.size());
    std::transform(x.begin(), x.end(), result.begin(), [&f](auto&& xx) {
      return apply_vector_unary<T_vt>::apply(xx, f);
    });
    return result;
  }
 
  template <typename F>
  static inline auto apply_no_holder(const T& x, const F& f) {
    return apply(x, f);
  }
 
  template <typename F>
  static inline auto reduce(const T& x, const F& f) {
    size_t x_size = x.size();
    using T_return = decltype(apply_vector_unary<T_vt>::reduce(x[0], f));
    std::vector<T_return> result(x_size);
    for (size_t i = 0; i < x_size; ++i)
      result[i] = apply_vector_unary<T_vt>::reduce(x[i], f);
    return result;
  }
};
 
}  // namespace math
}  // namespace stan
#endif