Automatic Differentiation
 
Loading...
Searching...
No Matches
dirichlet_lpdf.hpp
Go to the documentation of this file.
1#ifndef STAN_MATH_OPENCL_PRIM_DIRICHLET_LPDF_HPP
2#define STAN_MATH_OPENCL_PRIM_DIRICHLET_LPDF_HPP
3#ifdef STAN_OPENCL
4
5#include <stan/math/prim/fun/Eigen.hpp>
6#include <stan/math/prim/meta.hpp>
7#include <stan/math/prim/err.hpp>
8#include <stan/math/prim/fun/constants.hpp>
9#include <stan/math/prim/fun/elt_divide.hpp>
10#include <stan/math/prim/fun/elt_multiply.hpp>
11#include <stan/math/prim/fun/lgamma.hpp>
12#include <stan/math/opencl/kernel_generator.hpp>
13#include <stan/math/prim/functor/partials_propagator.hpp>
14#include <stan/math/prim/err/constraint_tolerance.hpp>
15
16namespace stan {
17namespace math {
18
58template <bool propto, typename T_prob_cl, typename T_prior_size_cl,
59 require_all_prim_or_rev_kernel_expression_t<
60 T_prob_cl, T_prior_size_cl>* = nullptr,
61 require_any_not_stan_scalar_t<T_prob_cl, T_prior_size_cl>* = nullptr>
62inline return_type_t<T_prob_cl, T_prior_size_cl> dirichlet_lpdf(
63 const T_prob_cl& theta, const T_prior_size_cl& alpha) {
64 static constexpr const char* function = "dirichlet_lpdf(OpenCL)";
65
66 check_consistent_sizes(function, "probabilities", theta, "prior sample sizes",
67 alpha);
68
69 if (max_size(theta, alpha) == 0) {
70 return 0.0;
71 }
72 if constexpr (!include_summand<propto, T_prob_cl, T_prior_size_cl>::value) {
73 return 0.0;
74 }
75
76 const auto& theta_val = value_of(theta);
77 const auto& alpha_val = value_of(alpha);
78
79 auto check_alpha_positive
80 = check_cl(function, "prior sample sizes", alpha_val, "positive");
81 auto alpha_positive = alpha_val > 0.0;
82 auto check_theta_nonnegative
83 = check_cl(function, "probabilities", theta_val, "nonnegative");
84 auto theta_nonnegative = theta_val >= 0.0;
85
86 auto theta_csum = colwise_sum(theta_val);
87 auto alpha_m_1 = rowwise_optional_broadcast(alpha_val - 1.0);
88 auto theta_log = rowwise_optional_broadcast(log(theta_val));
89
90 matrix_cl<double> theta_csum_cl;
91 matrix_cl<double> alpha_csum_cl;
92 matrix_cl<double> lgamma_alpha_csum_cl;
93 matrix_cl<double> theta_log_alpha_m_1_sum_cl;
94 matrix_cl<double> theta_deriv_cl;
95 matrix_cl<double> alpha_deriv_cl;
96 if (theta.cols() > alpha.cols()) {
97 auto alpha_csum = colwise_sum(alpha_val);
98 auto lgamma_alpha_csum = colwise_sum(lgamma(alpha_val));
99 matrix_cl<double> digamma_alpha_cl(alpha.rows(), alpha.cols());
100 results(check_alpha_positive, alpha_csum_cl, lgamma_alpha_csum_cl,
101 digamma_alpha_cl)
102 = expressions(
103 alpha_positive,
104 calc_if<include_summand<propto, T_prior_size_cl>::value>(
105 alpha_csum),
106 calc_if<include_summand<propto, T_prior_size_cl>::value>(
107 lgamma_alpha_csum),
108 calc_if<is_autodiff_v<T_prior_size_cl>>(digamma(alpha_val)));
109
110 auto theta_log_alpha_m_1_sum = sum_2d(elt_multiply(theta_log, alpha_m_1));
111
112 auto theta_deriv
113 = elt_divide(alpha_m_1, rowwise_optional_broadcast(theta_val));
114 auto alpha_deriv = theta_log - rowwise_optional_broadcast(digamma_alpha_cl);
115
116 results(check_theta_nonnegative, theta_csum_cl, theta_log_alpha_m_1_sum_cl,
117 theta_deriv_cl, alpha_deriv_cl)
118 = expressions(
119 theta_nonnegative, theta_csum,
120 calc_if<include_summand<propto, T_prob_cl, T_prior_size_cl>::value>(
121 theta_log_alpha_m_1_sum),
122 calc_if<is_autodiff_v<T_prob_cl>>(theta_deriv),
123 calc_if<is_autodiff_v<T_prior_size_cl>>(alpha_deriv));
124
125 if constexpr (include_summand<propto, T_prior_size_cl>::value) {
126 matrix_cl<double> alpha_csum_cl2;
127 matrix_cl<double> lgamma_alpha_csum_cl2;
128 while (alpha_csum_cl.rows() > 1) {
129 results(alpha_csum_cl2, lgamma_alpha_csum_cl2) = expressions(
130 calc_if<include_summand<propto, T_prior_size_cl>::value>(
131 colwise_sum(alpha_csum_cl)),
132 calc_if<include_summand<propto, T_prior_size_cl>::value>(
133 colwise_sum(lgamma_alpha_csum_cl)));
134 alpha_csum_cl = std::move(alpha_csum_cl2);
135 lgamma_alpha_csum_cl = std::move(lgamma_alpha_csum_cl2);
136 }
137 }
138 while (theta_csum_cl.rows() > 1) {
139 theta_csum_cl = colwise_sum(theta_csum_cl).eval();
140 }
141 } else {
142 auto alpha_csum = colwise_sum(alpha_val);
143 auto lgamma_alpha_csum = colwise_sum(lgamma(alpha_val));
144 if (alpha.cols() > theta.cols()) {
145 matrix_cl<double> log_theta_cl(theta.rows(), theta.cols());
146 results(check_theta_nonnegative, theta_csum_cl, log_theta_cl)
147 = expressions(
148 theta_nonnegative, theta_csum,
149 calc_if<
150 include_summand<propto, T_prob_cl, T_prior_size_cl>::value>(
151 theta_log));
152
153 auto log_theta_bc = rowwise_optional_broadcast(log_theta_cl);
154 auto theta_log_alpha_m_1_sum
155 = sum_2d(elt_multiply(log_theta_bc, alpha_m_1));
156
157 auto theta_deriv
158 = elt_divide(alpha_m_1, rowwise_optional_broadcast(theta_val));
159 auto alpha_deriv
160 = log_theta_bc - rowwise_optional_broadcast(digamma(alpha_val));
161
162 results(check_alpha_positive, alpha_csum_cl, lgamma_alpha_csum_cl,
163 theta_log_alpha_m_1_sum_cl, theta_deriv_cl, alpha_deriv_cl)
164 = expressions(
165 alpha_positive,
166 calc_if<include_summand<propto, T_prior_size_cl>::value>(
167 alpha_csum),
168 calc_if<include_summand<propto, T_prior_size_cl>::value>(
169 lgamma_alpha_csum),
170 calc_if<
171 include_summand<propto, T_prob_cl, T_prior_size_cl>::value>(
172 theta_log_alpha_m_1_sum),
173 calc_if<is_autodiff_v<T_prob_cl>>(theta_deriv),
174 calc_if<is_autodiff_v<T_prior_size_cl>>(alpha_deriv));
175
176 while (alpha_csum_cl.rows() > 1) {
177 matrix_cl<double> alpha_csum_cl2;
178 matrix_cl<double> lgamma_alpha_csum_cl2;
179 results(alpha_csum_cl2, lgamma_alpha_csum_cl2) = expressions(
180 calc_if<include_summand<propto, T_prior_size_cl>::value>(
181 colwise_sum(alpha_csum_cl)),
182 calc_if<include_summand<propto, T_prior_size_cl>::value>(
183 colwise_sum(lgamma_alpha_csum_cl)));
184 if constexpr (include_summand<propto, T_prior_size_cl>::value) {
185 alpha_csum_cl = std::move(alpha_csum_cl2);
186 lgamma_alpha_csum_cl = std::move(lgamma_alpha_csum_cl2);
187 }
188 }
189 double theta_sum = sum(from_matrix_cl(theta_csum));
190 check_cl(function, "sum of probabilities", theta_sum, "equal to 1")
191 = (fabs(theta_sum - 1.0) <= CONSTRAINT_TOLERANCE);
192 } else {
193 auto theta_log_alpha_m_1_sum = sum_2d(elt_multiply(theta_log, alpha_m_1));
194
195 auto theta_deriv
196 = elt_divide(alpha_m_1, rowwise_optional_broadcast(theta_val));
197 auto alpha_deriv
198 = theta_log - rowwise_optional_broadcast(digamma(alpha_val));
199
200 results(check_alpha_positive, check_theta_nonnegative, theta_csum_cl,
201 alpha_csum_cl, lgamma_alpha_csum_cl, theta_log_alpha_m_1_sum_cl,
202 theta_deriv_cl, alpha_deriv_cl)
203 = expressions(
204 alpha_positive, theta_nonnegative, theta_csum,
205 calc_if<include_summand<propto, T_prior_size_cl>::value>(
206 alpha_csum),
207 calc_if<include_summand<propto, T_prior_size_cl>::value>(
208 lgamma_alpha_csum),
209 calc_if<
210 include_summand<propto, T_prob_cl, T_prior_size_cl>::value>(
211 theta_log_alpha_m_1_sum),
212 calc_if<is_autodiff_v<T_prob_cl>>(theta_deriv),
213 calc_if<is_autodiff_v<T_prior_size_cl>>(alpha_deriv));
214
215 while (theta_csum_cl.rows() > 1) {
216 matrix_cl<double> theta_csum_cl2;
217 matrix_cl<double> alpha_csum_cl2;
218 matrix_cl<double> lgamma_alpha_csum_cl2;
219 results(theta_csum_cl2, alpha_csum_cl2, lgamma_alpha_csum_cl2)
220 = expressions(
221 colwise_sum(theta_csum_cl),
222 calc_if<include_summand<propto, T_prior_size_cl>::value>(
223 colwise_sum(alpha_csum_cl)),
224 calc_if<include_summand<propto, T_prior_size_cl>::value>(
225 colwise_sum(lgamma_alpha_csum_cl)));
226 theta_csum_cl = std::move(theta_csum_cl2);
227 if constexpr (include_summand<propto, T_prior_size_cl>::value) {
228 alpha_csum_cl = std::move(alpha_csum_cl2);
229 lgamma_alpha_csum_cl = std::move(lgamma_alpha_csum_cl2);
230 }
231 }
232 }
233 }
234
235 if (theta.cols() >= alpha.cols()) {
236 // transpose is there just because working on col vectors is more efficient
237 // than on row vectors with kernel generator
238 check_cl(function, "sum of probabilities", transpose(theta_csum_cl),
239 "equal to 1")
240 = (fabs(transpose(theta_csum_cl) - 1.0) <= CONSTRAINT_TOLERANCE);
241 }
242
243 double lp = 0.0;
244
245 if constexpr (include_summand<propto, T_prior_size_cl>::value) {
246 if (theta.cols() > alpha.cols()) {
247 lp += (lgamma(from_matrix_cl(alpha_csum_cl)) * theta.cols()
248 - from_matrix_cl(lgamma_alpha_csum_cl) * theta.cols())
249 .sum();
250 } else {
251 lp += (lgamma(from_matrix_cl(alpha_csum_cl))
252 - from_matrix_cl(lgamma_alpha_csum_cl))
253 .sum();
254 }
255 }
256 if constexpr (include_summand<propto, T_prob_cl, T_prior_size_cl>::value) {
257 lp += from_matrix_cl(theta_log_alpha_m_1_sum_cl).sum();
258 }
259
260 auto ops_partials = make_partials_propagator(theta, alpha);
261
262 if constexpr (is_autodiff_v<T_prob_cl>) {
263 if (theta.cols() < alpha.cols()) {
264 partials<0>(ops_partials) = rowwise_sum(theta_deriv_cl);
265 } else {
266 partials<0>(ops_partials) = std::move(theta_deriv_cl);
267 }
268 }
269 if constexpr (is_autodiff_v<T_prior_size_cl>) {
270 if (theta.cols() > alpha.cols()) {
271 matrix_cl<double> tmp_cl
272 = digamma(alpha_csum_cl) * static_cast<double>(theta.cols());
273 partials<1>(ops_partials)
274 = colwise_broadcast(tmp_cl) + rowwise_sum(alpha_deriv_cl);
275 } else {
276 matrix_cl<double> tmp_cl = digamma(alpha_csum_cl);
277 partials<1>(ops_partials) = colwise_broadcast(tmp_cl) + alpha_deriv_cl;
278 }
279 }
280 return ops_partials.build(lp);
281}
282
283} // namespace math
284} // namespace stan
285#endif
286#endif
stan::math::matrix_cl::rows
int rows() const
Definition matrix_cl.hpp:64
stan::math::matrix_cl
Represents an arithmetic matrix on the OpenCL device.
Definition matrix_cl.hpp:47
constraint_tolerance.hpp
stan::math::elt_multiply
elt_multiply_< as_operation_cl_t< T_a >, as_operation_cl_t< T_b > > elt_multiply(T_a &&a, T_b &&b)
Definition binary_operation.hpp:204
stan::math::sum_2d
auto sum_2d(T &&a)
Two dimensional sum - reduction of a kernel generator expression.
Definition reduction_2d.hpp:233
stan::math::check_cl
auto check_cl(const char *function, const char *var_name, T &&y, const char *must_be)
Constructs a check on opencl matrix or expression.
Definition check_cl.hpp:219
stan::math::results
results_cl< T_results... > results(T_results &&... results)
Deduces types for constructing results_cl object.
Definition multi_result_kernel.hpp:668
stan::math::transpose
auto transpose(Arg &&a)
Transposes a kernel generator expression.
Definition transpose.hpp:139
stan::math::elt_divide
elt_divide_< as_operation_cl_t< T_a >, as_operation_cl_t< T_b > > elt_divide(T_a &&a, T_b &&b)
Definition binary_operation.hpp:209
stan::math::rowwise_sum
auto rowwise_sum(T &&a)
Rowwise sum reduction of a kernel generator expression.
Definition rowwise_reduction.hpp:330
stan::math::calc_if
calc_if_< true, as_operation_cl_t< T > > calc_if(T &&a)
Definition calc_if.hpp:121
stan::math::colwise_broadcast
auto colwise_broadcast(T &&a)
Broadcast an expression in colwise dimmension.
Definition broadcast.hpp:163
stan::math::colwise_sum
auto colwise_sum(T &&a)
Column wise sum - reduction of a kernel generator expression.
Definition colwise_reduction.hpp:224
stan::math::expressions
expressions_cl< T_expressions... > expressions(T_expressions &&... expressions)
Deduces types for constructing expressions_cl object.
Definition multi_result_kernel.hpp:289
stan::math::dirichlet_lpdf
return_type_t< T_prob_cl, T_prior_size_cl > dirichlet_lpdf(const T_prob_cl &theta, const T_prior_size_cl &alpha)
The log of the Dirichlet density for the given theta and a vector of prior sample sizes,...
Definition dirichlet_lpdf.hpp:62
stan::math::from_matrix_cl
auto from_matrix_cl(const T &src)
Copies the source matrix that is stored on the OpenCL device to the destination Eigen matrix.
Definition copy.hpp:61
stan::require_all_prim_or_rev_kernel_expression_t
require_all_t< is_prim_or_rev_kernel_expression< std::decay_t< Types > >... > require_all_prim_or_rev_kernel_expression_t
Require type satisfies is_prim_or_rev_kernel_expression.
Definition is_kernel_expression.hpp:148
stan::return_type_t
typename return_type< Ts... >::type return_type_t
Convenience type for the return type of the specified template parameters.
Definition return_type.hpp:218
kernel_generator.hpp
stan::math::value_of
T value_of(const fvar< T > &v)
Return the value of the specified variable.
Definition value_of.hpp:18
stan::math::log
fvar< T > log(const fvar< T > &x)
Definition log.hpp:18
stan::math::rowwise_optional_broadcast
auto rowwise_optional_broadcast(T &&a)
Broadcast an expression in rowwise dimmension if the number of columns equals to 1.
Definition optional_broadcast.hpp:223
stan::math::check_consistent_sizes
void check_consistent_sizes(const char *)
Trivial no input case, this function is a no-op.
Definition check_consistent_sizes.hpp:15
stan::math::lgamma
fvar< T > lgamma(const fvar< T > &x)
Return the natural logarithm of the gamma function applied to the specified argument.
Definition lgamma.hpp:21
stan::math::sum
auto sum(const std::vector< T > &m)
Return the sum of the entries of the specified standard vector.
Definition sum.hpp:23
stan::math::max_size
int64_t max_size(const T1 &x1, const Ts &... xs)
Calculate the size of the largest input.
Definition max_size.hpp:20
stan::math::CONSTRAINT_TOLERANCE
const double CONSTRAINT_TOLERANCE
The tolerance for checking arithmetic bounds in rank and in simplexes.
Definition constraint_tolerance.hpp:18
stan::math::make_partials_propagator
auto make_partials_propagator(Ops &&... ops)
Construct an partials_propagator.
Definition partials_propagator.hpp:119
stan::math::digamma
fvar< T > digamma(const fvar< T > &x)
Return the derivative of the log gamma function at the specified argument.
Definition digamma.hpp:23
stan::math::fabs
fvar< T > fabs(const fvar< T > &x)
Definition fabs.hpp:16
stan
The lgamma implementation in stan-math is based on either the reentrant safe lgamma_r implementation ...
Definition unit_vector_constrain.hpp:15
elt_divide.hpp
elt_multiply.hpp
partials_propagator.hpp
stan::math::include_summand
Template metaprogram to calculate whether a summand needs to be included in a proportional (log) prob...
Definition include_summand.hpp:39