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, T_true&& y_true, const T_false& y_false) {
  if (c) {
    return y_true;
  } else {
    return apply_scalar_binary(
        [](auto&& y_true_inner, auto&& y_false_inner) { return y_false_inner; },
        std::forward<T_true>(y_true), 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_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, T_false&& y_false) {
  if (c) {
    return apply_scalar_binary(
        [](auto&& y_true_inner, auto&& y_false_inner) { return y_true_inner; },
        y_true, std::forward<T_false>(y_false));
  } 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(T_bool&& c, const T_true& y_true, const T_false& y_false) {
  using ret_t = return_type_t<T_true, T_false>;
  return make_holder(
      [y_true, y_false](auto&& c_) {
        return std::forward<decltype(c_)>(c_).unaryExpr(
            [y_true, y_false](bool cond) {
              return cond ? ret_t(y_true) : ret_t(y_false);
            });
      },
      std::forward<T_bool>(c));
}
 
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(T_bool&& c, T_true&& y_true, 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>;
  if constexpr (!std::is_same_v<std::decay_t<T_true>, std::decay_t<T_false>>) {
    return make_holder(
        [](auto&& c_, auto&& y_true_, auto&& y_false_) {
          return std::forward<decltype(c_)>(c_).select(
              std::forward<decltype(y_true_)>(y_true_),
              std::forward<decltype(y_false_)>(y_false_));
        },
        std::forward<T_bool>(c), std::forward<T_true>(y_true),
        std::forward<T_false>(y_false));
  } else {
    return make_holder(
        [](auto&& c_, auto&& y_true_, auto&& y_false_) {
          return std::forward<decltype(c_)>(c_)
              .select(std::forward<decltype(y_true_)>(y_true_),
                      std::forward<decltype(y_false_)>(y_false_))
              .template cast<ret_t>();
        },
        std::forward<T_bool>(c), std::forward<T_true>(y_true),
        std::forward<T_false>(y_false));
  }
}
}  // namespace math
}  // namespace stan
 
#endif