|
| 1 | +#ifndef STAN_MATH_OPENCL_PRIM_BERNOULLI_LCCDF_HPP |
| 2 | +#define STAN_MATH_OPENCL_PRIM_BERNOULLI_LCCDF_HPP |
| 3 | +#ifdef STAN_OPENCL |
| 4 | + |
| 5 | +#include <stan/math/prim/meta.hpp> |
| 6 | +#include <stan/math/prim/err.hpp> |
| 7 | +#include <stan/math/opencl/kernel_generator.hpp> |
| 8 | +#include <stan/math/prim/fun/constants.hpp> |
| 9 | +#include <stan/math/prim/fun/elt_divide.hpp> |
| 10 | +#include <stan/math/prim/functor/operands_and_partials.hpp> |
| 11 | + |
| 12 | +namespace stan { |
| 13 | +namespace math { |
| 14 | + |
| 15 | +/** \ingroup prob_dists |
| 16 | + * Returns the log CCDF of the Bernoulli distribution. If containers are |
| 17 | + * supplied, returns the log sum of the probabilities. |
| 18 | + * |
| 19 | + * @tparam T_n_cl type of integer parameter |
| 20 | + * @tparam T_prob_cl type of chance of success parameter |
| 21 | + * @param n integer parameter |
| 22 | + * @param theta logit-transformed chance of success parameter |
| 23 | + * @return log probability or log sum of probabilities |
| 24 | + * @throw std::domain_error if theta is not a valid probability |
| 25 | + * @throw std::invalid_argument if container sizes mismatch. |
| 26 | + */ |
| 27 | +template < |
| 28 | + typename T_n_cl, typename T_prob_cl, |
| 29 | + require_all_prim_or_rev_kernel_expression_t<T_n_cl, T_prob_cl>* = nullptr, |
| 30 | + require_any_not_stan_scalar_t<T_n_cl, T_prob_cl>* = nullptr> |
| 31 | +return_type_t<T_prob_cl> bernoulli_lccdf(const T_n_cl& n, |
| 32 | + const T_prob_cl& theta) { |
| 33 | + static const char* function = "bernoulli_lccdf(OpenCL)"; |
| 34 | + using T_partials_return = partials_return_t<T_prob_cl>; |
| 35 | + using std::isnan; |
| 36 | + constexpr bool is_n_vector = !is_stan_scalar<T_n_cl>::value; |
| 37 | + constexpr bool is_theta_vector = !is_stan_scalar<T_prob_cl>::value; |
| 38 | + |
| 39 | + check_consistent_sizes(function, "Random variable", n, |
| 40 | + "Probability parameter", theta); |
| 41 | + const size_t N = is_n_vector ? size(n) : size(theta); |
| 42 | + if (N == 0) { |
| 43 | + return 0.0; |
| 44 | + } |
| 45 | + |
| 46 | + const auto& theta_col = as_column_vector_or_scalar(theta); |
| 47 | + const auto& theta_val = value_of(theta_col); |
| 48 | + |
| 49 | + auto check_theta_bounded = check_cl(function, "Probability parameter", |
| 50 | + theta_val, "in the interval [0, 1]"); |
| 51 | + auto theta_bounded_expr = 0.0 <= theta_val && theta_val <= 1.0; |
| 52 | + |
| 53 | + auto any_n_negative = colwise_max(0 + (n < 0)); |
| 54 | + auto any_n_over_one = colwise_max(constant(0, N, 1) + (n >= 1)); |
| 55 | + auto P_expr = colwise_sum(log(theta_val)); |
| 56 | + auto deriv = elt_divide(1.0, theta_val); |
| 57 | + |
| 58 | + matrix_cl<double> any_n_negative_cl; |
| 59 | + matrix_cl<double> any_n_over_one_cl; |
| 60 | + matrix_cl<double> P_cl; |
| 61 | + matrix_cl<double> deriv_cl; |
| 62 | + |
| 63 | + results(check_theta_bounded, any_n_negative_cl, any_n_over_one_cl, P_cl, |
| 64 | + deriv_cl) |
| 65 | + = expressions(theta_bounded_expr, any_n_negative, any_n_over_one, P_expr, |
| 66 | + calc_if<(!is_constant_all<T_prob_cl>::value)>(deriv)); |
| 67 | + |
| 68 | + if (from_matrix_cl(any_n_negative_cl).maxCoeff()) { |
| 69 | + return 0.0; |
| 70 | + } |
| 71 | + if (from_matrix_cl(any_n_over_one_cl).maxCoeff()) { |
| 72 | + return NEGATIVE_INFTY; |
| 73 | + } |
| 74 | + |
| 75 | + T_partials_return P = from_matrix_cl(P_cl).sum(); |
| 76 | + operands_and_partials<decltype(theta_col)> ops_partials(theta_col); |
| 77 | + |
| 78 | + if (!is_constant_all<T_prob_cl>::value) { |
| 79 | + ops_partials.edge1_.partials_ = std::move(deriv_cl); |
| 80 | + } |
| 81 | + |
| 82 | + return ops_partials.build(P); |
| 83 | +} |
| 84 | + |
| 85 | +} // namespace math |
| 86 | +} // namespace stan |
| 87 | +#endif |
| 88 | +#endif |
0 commit comments