apply_scalar_ternary.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_PRIM_FUNCTOR_APPLY_SCALAR_TERNARY_HPP
#define STAN_MATH_PRIM_FUNCTOR_APPLY_SCALAR_TERNARY_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.hpp>
#include <stan/math/prim/meta/require_generics.hpp>
#include <stan/math/prim/err/check_matching_dims.hpp>
#include <stan/math/prim/err/check_matching_sizes.hpp>
#include <stan/math/prim/fun/num_elements.hpp>
#include <vector>
 
namespace stan {
namespace math {
 
template <typename F, typename T1, typename T2, typename T3,
          require_all_stan_scalar_t<T1, T2, T3>* = nullptr>
inline auto apply_scalar_ternary(F&& f, T1&& x, T2&& y, T3&& z) {
  return std::forward<F>(f)(std::forward<T1>(x), std::forward<T2>(y),
                            std::forward<T3>(z));
}
 
template <typename F, typename T1, typename T2, typename T3,
          require_all_eigen_t<T1, T2, T3>* = nullptr>
inline auto apply_scalar_ternary(F&& f, T1&& x, T2&& y, T3&& z) {
  check_matching_dims("Ternary function", "x", x, "y", y);
  check_matching_dims("Ternary function", "y", y, "z", z);
  return make_holder(
      [](auto&& f_inner, auto&& x_inner, auto&& y_inner, auto&& z_inner) {
        return Eigen::CwiseTernaryOp<
            std::decay_t<decltype(f_inner)>, std::decay_t<decltype(x_inner)>,
            std::decay_t<decltype(y_inner)>, std::decay_t<decltype(z_inner)>>(
            x_inner, y_inner, z_inner, f_inner);
      },
      std::forward<F>(f), std::forward<T1>(x), std::forward<T2>(y),
      std::forward<T3>(z));
}
 
template <typename F, typename T1, typename T2, typename T3,
          require_all_std_vector_vt<is_stan_scalar, T1, T2, T3>* = nullptr>
inline auto apply_scalar_ternary(F&& f, T1&& x, T2&& y, T3&& z) {
  check_matching_sizes("Ternary function", "x", x, "y", y);
  check_matching_sizes("Ternary function", "y", y, "z", z);
  using T_return = std::decay_t<decltype(f(x[0], y[0], z[0]))>;
  decltype(auto) x_vec = as_column_vector_or_scalar(std::forward<T1>(x));
  decltype(auto) y_vec = as_column_vector_or_scalar(std::forward<T2>(y));
  decltype(auto) z_vec = as_column_vector_or_scalar(std::forward<T3>(z));
  std::vector<T_return> result(x_vec.size());
  Eigen::Map<Eigen::Matrix<T_return, -1, 1>>(result.data(), result.size())
      = apply_scalar_ternary(std::forward<F>(f),
                             std::forward<decltype(x_vec)>(x_vec),
                             std::forward<decltype(y_vec)>(y_vec),
                             std::forward<decltype(z_vec)>(z_vec));
  return result;
}
 
template <typename F, typename T1, typename T2, typename T3,
          require_all_std_vector_vt<is_container_or_var_matrix, T1, T2,
                                    T3>* = nullptr>
inline auto apply_scalar_ternary(F&& f, T1&& x, T2&& y, T3&& z) {
  check_matching_sizes("Ternary function", "x", x, "y", y);
  check_matching_sizes("Ternary function", "y", y, "z", z);
  using T_return
      = plain_type_t<decltype(apply_scalar_ternary(f, x[0], y[0], z[0]))>;
  size_t y_size = y.size();
  std::vector<T_return> result(y_size);
  for (size_t i = 0; i < y_size; ++i) {
    result[i] = apply_scalar_ternary(f, x[i], y[i], z[i]);
  }
  return result;
}
 
template <typename F, typename T1, typename T2, typename T3,
          require_any_container_t<T1, T2>* = nullptr,
          require_stan_scalar_t<T3>* = nullptr>
inline auto apply_scalar_ternary(F&& f, T1&& x, T2&& y, T3&& z) {
  return apply_scalar_binary(
      [f_ = std::forward<F>(f), z](auto&& a, auto&& b) {
        return f_(std::forward<decltype(a)>(a), std::forward<decltype(b)>(b),
                  z);
      },
      std::forward<T1>(x), std::forward<T2>(y));
}
 
template <typename F, typename T1, typename T2, typename T3,
          require_all_container_t<T1, T3>* = nullptr,
          require_stan_scalar_t<T2>* = nullptr>
inline auto apply_scalar_ternary(F&& f, T1&& x, T2&& y, T3&& z) {
  return apply_scalar_binary(
      [f_ = std::forward<F>(f), y](auto&& a, auto&& c) {
        return f_(std::forward<decltype(a)>(a), y,
                  std::forward<decltype(c)>(c));
      },
      std::forward<T1>(x), std::forward<T3>(z));
}
 
template <typename F, typename T1, typename T2, typename T3,
          require_container_t<T3>* = nullptr,
          require_stan_scalar_t<T1>* = nullptr>
inline auto apply_scalar_ternary(F&& f, T1&& x, T2&& y, T3&& z) {
  return apply_scalar_binary(
      [f_ = std::forward<F>(f), x](auto&& b, auto&& c) {
        return f_(x, std::forward<decltype(b)>(b),
                  std::forward<decltype(c)>(c));
      },
      std::forward<T2>(y), std::forward<T3>(z));
}
 
}  // namespace math
}  // namespace stan
#endif