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