select.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_PRIM_FUN_SELECT_HPP
#define STAN_MATH_PRIM_FUN_SELECT_HPP
 
#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/functor/apply_scalar_binary.hpp>
 
namespace stan {
namespace math {
 
template <typename T_true, typename T_false,
          typename ReturnT = return_type_t<T_true, T_false>,
          require_all_stan_scalar_t<T_true, T_false>* = nullptr>
inline ReturnT select(const bool c, const T_true y_true,
                      const T_false y_false) {
  return c ? ReturnT(y_true) : ReturnT(y_false);
}
 
template <
    typename T_true, typename T_false,
    typename T_return = return_type_t<T_true, T_false>,
    typename T_true_plain = promote_scalar_t<T_return, plain_type_t<T_true>>,
    typename T_false_plain = promote_scalar_t<T_return, plain_type_t<T_false>>,
    require_all_container_t<T_true, T_false>* = nullptr,
    require_all_same_t<T_true_plain, T_false_plain>* = nullptr>
inline T_true_plain select(const bool c, T_true&& y_true, T_false&& y_false) {
  check_matching_dims("select", "left hand side", y_true, "right hand side",
                      y_false);
  return c ? T_true_plain(std::forward<T_true>(y_true))
           : T_true_plain(std::forward<T_false>(y_false));
}
 
template <typename T_true, typename T_false,
          typename ReturnT = promote_scalar_t<return_type_t<T_true, T_false>,
                                              plain_type_t<T_true>>,
          require_container_t<T_true>* = nullptr,
          require_stan_scalar_t<T_false>* = nullptr>
inline ReturnT select(const bool c, const T_true& y_true,
                      const T_false& y_false) {
  if (c) {
    return y_true;
  } else {
    return apply_scalar_binary(
        y_true, y_false,
        [](const auto& y_true_inner, const auto& y_false_inner) {
          return y_false_inner;
        });
  }
}
 
template <typename T_true, typename T_false,
          typename ReturnT = promote_scalar_t<return_type_t<T_true, T_false>,
                                              plain_type_t<T_false>>,
          require_stan_scalar_t<T_true>* = nullptr,
          require_container_t<T_false>* = nullptr>
inline ReturnT select(const bool c, const T_true y_true,
                      const T_false y_false) {
  if (c) {
    return apply_scalar_binary(
        y_true, y_false,
        [](const auto& y_true_inner, const auto& y_false_inner) {
          return y_true_inner;
        });
  } else {
    return y_false;
  }
}
 
template <typename T_bool, typename T_true, typename T_false,
          require_eigen_array_vt<std::is_integral, T_bool>* = nullptr,
          require_all_stan_scalar_t<T_true, T_false>* = nullptr>
inline auto select(const T_bool c, const T_true y_true, const T_false y_false) {
  using ret_t = return_type_t<T_true, T_false>;
  return c
      .unaryExpr([y_true, y_false](bool cond) {
        return cond ? ret_t(y_true) : ret_t(y_false);
      })
      .eval();
}
 
template <typename T_bool, typename T_true, typename T_false,
          require_eigen_array_t<T_bool>* = nullptr,
          require_any_eigen_array_t<T_true, T_false>* = nullptr>
inline auto select(const T_bool c, const T_true y_true, const T_false y_false) {
  check_consistent_sizes("select", "boolean", c, "left hand side", y_true,
                         "right hand side", y_false);
  using ret_t = return_type_t<T_true, T_false>;
  return c.select(y_true, y_false).template cast<ret_t>().eval();
}
 
}  // namespace math
}  // namespace stan
 
#endif