1 #ifndef STAN_MATH_OPENCL_PRIM_NEG_BINOMIAL_2_LPMF_HPP
2 #define STAN_MATH_OPENCL_PRIM_NEG_BINOMIAL_2_LPMF_HPP
3 #ifdef STAN_OPENCL
4
5 #include <stan/math/prim/meta.hpp>
6 #include <stan/math/prim/err.hpp>
7 #include <stan/math/prim/fun/constants.hpp>
8 #include <stan/math/prim/fun/elt_divide.hpp>
9 #include <stan/math/prim/fun/elt_multiply.hpp>
10 #include <stan/math/prim/fun/exp.hpp>
11 #include <stan/math/opencl/kernel_generator.hpp>
12 #include <stan/math/prim/functor/operands_and_partials.hpp>
13 #include <stan/math/prim/fun/multiply_log.hpp>
14
15 namespace stan {
16 namespace math {
17
18 /** \ingroup opencl
19 * The log of the negative binomial density for the specified scalars given
20 * the specified mean(s) and deviation(s). n, mu, or phi can
21 * each be either a scalar or a vector matrix_cl. Any vector inputs
22 * must be the same length.
23 *
24 * <p>The result log probability is defined to be the sum of the
25 * log probabilities for each observation/mean/deviation triple.
26 *
27 * @tparam T_n_cl type of scalar
28 * @tparam T_location_cl type of location parameter
29 * @tparam T_precision_cl type of precision parameter
30 * @param n (Sequence of) scalar(s).
31 * @param mu (Sequence of) location parameter(s)
32 * @param phi (Sequence of) precision parameters
33 * @return The log of the product of the densities.
34 * @throw std::domain_error if the scale is not positive.
35 */
36 template <bool propto, typename T_n_cl, typename T_location_cl,
37 typename T_precision_cl,
38 require_all_prim_or_rev_kernel_expression_t<
39 T_n_cl, T_location_cl, T_precision_cl>* = nullptr,
40 require_any_not_stan_scalar_t<T_n_cl, T_location_cl,
41 T_precision_cl>* = nullptr>
neg_binomial_2_lpmf(const T_n_cl & n,const T_location_cl & mu,const T_precision_cl & phi)42 inline return_type_t<T_n_cl, T_location_cl, T_precision_cl> neg_binomial_2_lpmf(
43 const T_n_cl& n, const T_location_cl& mu, const T_precision_cl& phi) {
44 static const char* function = "neg_binomial_2_lpmf(OpenCL)";
45 using T_partials_return
46 = partials_return_t<T_n_cl, T_location_cl, T_precision_cl>;
47 using std::isfinite;
48 using std::isnan;
49
50 check_consistent_sizes(function, "Failures variable", n, "Location parameter",
51 mu, "Precision parameter", phi);
52 const size_t N = max_size(n, mu, phi);
53 if (N == 0) {
54 return 0.0;
55 }
56 if (!include_summand<propto, T_n_cl, T_location_cl, T_precision_cl>::value) {
57 return 0.0;
58 }
59
60 const auto& mu_col = as_column_vector_or_scalar(mu);
61 const auto& phi_col = as_column_vector_or_scalar(phi);
62
63 const auto& mu_val = value_of(mu_col);
64 const auto& phi_val = value_of(phi_col);
65
66 auto check_n_nonnegative
67 = check_cl(function, "Failures variable", n, "nonnegative");
68 auto n_nonnegative = n >= 0;
69 auto check_mu_positive_finite
70 = check_cl(function, "Log location parameter", mu_val, "positive finite");
71 auto mu_positive_finite = 0 < mu_val && isfinite(mu_val);
72 auto check_phi_positive_finite
73 = check_cl(function, "Precision parameter", phi_val, "positive finite");
74 auto phi_positive_finite = 0 < phi_val && isfinite(phi_val);
75
76 auto log_phi = log(phi_val);
77 auto mu_plus_phi = mu_val + phi_val;
78 auto log_mu_plus_phi = log(mu_plus_phi);
79 auto n_plus_phi = n + phi_val;
80
81 auto logp1 = -elt_multiply(phi_val, log1p(elt_divide(mu_val, phi_val)))
82 - elt_multiply(n, log_mu_plus_phi);
83 auto logp2 = static_select<include_summand<propto, T_precision_cl>::value>(
84 logp1 + binomial_coefficient_log(n_plus_phi - 1, n), logp1);
85 auto logp_expr = colwise_sum(
86 static_select<include_summand<propto, T_location_cl>::value>(
87 logp2 + multiply_log(n, mu_val), logp2));
88
89 auto mu_deriv = elt_divide(n, mu_val) - elt_divide(n + phi_val, mu_plus_phi);
90 auto log_term
91 = select(mu_val < phi_val, log1p(-elt_divide(mu_val, mu_plus_phi)),
92 log_phi - log_mu_plus_phi);
93 auto phi_deriv = elt_divide(mu_val - n, mu_plus_phi) + log_term
94 - digamma(phi_val) + digamma(n_plus_phi);
95
96 matrix_cl<double> logp_cl;
97 matrix_cl<double> mu_deriv_cl;
98 matrix_cl<double> phi_deriv_cl;
99
100 results(check_n_nonnegative, check_mu_positive_finite,
101 check_phi_positive_finite, logp_cl, mu_deriv_cl, phi_deriv_cl)
102 = expressions(n_nonnegative, mu_positive_finite, phi_positive_finite,
103 logp_expr,
104 calc_if<!is_constant<T_location_cl>::value>(mu_deriv),
105 calc_if<!is_constant<T_precision_cl>::value>(phi_deriv));
106
107 T_partials_return logp = sum(from_matrix_cl(logp_cl));
108
109 operands_and_partials<decltype(mu_col), decltype(phi_col)> ops_partials(
110 mu_col, phi_col);
111
112 if (!is_constant<T_location_cl>::value) {
113 ops_partials.edge1_.partials_ = std::move(mu_deriv_cl);
114 }
115 if (!is_constant<T_precision_cl>::value) {
116 ops_partials.edge2_.partials_ = std::move(phi_deriv_cl);
117 }
118 return ops_partials.build(logp);
119 }
120
121 } // namespace math
122 } // namespace stan
123 #endif
124 #endif
125