1 #ifndef STAN_MATH_PRIM_FUNCTOR_APPLY_SCALAR_BINARY_HPP
2 #define STAN_MATH_PRIM_FUNCTOR_APPLY_SCALAR_BINARY_HPP
3 
4 #include <stan/math/prim/fun/as_column_vector_or_scalar.hpp>
5 #include <stan/math/prim/meta/is_stan_scalar.hpp>
6 #include <stan/math/prim/meta/is_container.hpp>
7 #include <stan/math/prim/meta/is_eigen.hpp>
8 #include <stan/math/prim/meta/require_generics.hpp>
9 #include <stan/math/prim/err/check_matching_dims.hpp>
10 #include <stan/math/prim/err/check_matching_sizes.hpp>
11 #include <stan/math/prim/fun/num_elements.hpp>
12 #include <vector>
13 
14 namespace stan {
15 namespace math {
16 
17 /**
18  * Base template function for vectorization of binary scalar functions
19  * defined by applying a functor to a combination of scalars,
20  * containers of matching sizes, or a combination of a scalar and a container.
21  * These containers can be a standard library vector, Eigen dense
22  * matrix expression template, or container of these. For each specialisation,
23  * the same type as the largest (dimension) input is returned.
24  *
25  * This base template function takes (and returns) scalars.
26  *
27  * @tparam T1 Type of first argument to which functor is applied.
28  * @tparam T2 Type of second argument to which functor is applied.
29  * @tparam F Type of functor to apply.
30  * @param x First input to which operation is applied.
31  * @param y Second input to which operation is applied.
32  * @param f functor to apply to inputs.
33  * @return Scalar with result of applying functor to input.
34  */
35 template <typename T1, typename T2, typename F,
36           require_all_stan_scalar_t<T1, T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)37 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
38   return f(x, y);
39 }
40 
41 /**
42  * Specialisation for use with two Eigen inputs. Eigen's binaryExpr framework
43  * is used for more efficient indexing of both row- and column-major inputs
44  * without separate loops.
45  *
46  * @tparam T1 Type of first argument to which functor is applied.
47  * @tparam T2 Type of second argument to which functor is applied.
48  * @tparam F Type of functor to apply.
49  * @param x First Eigen input to which operation is applied.
50  * @param y Second Eigen input to which operation is applied.
51  * @param f functor to apply to Eigen input.
52  * @return Eigen object with result of applying functor to inputs.
53  */
54 template <typename T1, typename T2, typename F,
55           require_all_eigen_t<T1, T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)56 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
57   check_matching_dims("Binary function", "x", x, "y", y);
58   return x.binaryExpr(y, f);
59 }
60 
61 /**
62  * Specialisation for use with one Eigen vector (row or column) and
63  * a one-dimensional std::vector of integer types
64  *
65  * @tparam T1 Type of first argument to which functor is applied.
66  * @tparam T2 Type of second argument to which functor is applied.
67  * @tparam F Type of functor to apply.
68  * @param x Eigen input to which operation is applied.
69  * @param y Integer std::vector input to which operation is applied.
70  * @param f functor to apply to inputs.
71  * @return Eigen object with result of applying functor to inputs.
72  */
73 template <typename T1, typename T2, typename F,
74           require_eigen_vector_vt<is_stan_scalar, T1>* = nullptr,
75           require_std_vector_vt<std::is_integral, T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)76 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
77   check_matching_sizes("Binary function", "x", x, "y", y);
78   using int_vec_t = promote_scalar_t<value_type_t<T2>, plain_type_t<T1>>;
79   Eigen::Map<const int_vec_t> y_map(y.data(), y.size());
80   return x.binaryExpr(y_map, f);
81 }
82 
83 /**
84  * Specialisation for use with a one-dimensional std::vector of integer types
85  * and one Eigen vector (row or column).
86  *
87  * @tparam T1 Type of first argument to which functor is applied.
88  * @tparam T2 Type of second argument to which functor is applied.
89  * @tparam F Type of functor to apply.
90  * @param x Integer std::vector input to which operation is applied.
91  * @param y Eigen input to which operation is applied.
92  * @param f functor to apply to inputs.
93  * @return Eigen object with result of applying functor to inputs.
94  */
95 template <typename T1, typename T2, typename F,
96           require_std_vector_vt<std::is_integral, T1>* = nullptr,
97           require_eigen_vector_vt<is_stan_scalar, T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)98 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
99   check_matching_sizes("Binary function", "x", x, "y", y);
100   using int_vec_t = promote_scalar_t<value_type_t<T1>, plain_type_t<T2>>;
101   Eigen::Map<const int_vec_t> x_map(x.data(), x.size());
102   return x_map.binaryExpr(y, f);
103 }
104 
105 /**
106  * Specialisation for use with one Eigen matrix and
107  * a two-dimensional std::vector of integer types
108  *
109  * @tparam T1 Type of first argument to which functor is applied.
110  * @tparam T2 Type of second argument to which functor is applied.
111  * @tparam F Type of functor to apply.
112  * @param x Eigen matrix input to which operation is applied.
113  * @param y Nested integer std::vector input to which operation is applied.
114  * @param f functor to apply to inputs.
115  * @return Eigen object with result of applying functor to inputs.
116  */
117 template <typename T1, typename T2, typename F,
118           require_eigen_matrix_dynamic_vt<is_stan_scalar, T1>* = nullptr,
119           require_std_vector_vt<is_std_vector, T2>* = nullptr,
120           require_std_vector_st<std::is_integral, T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)121 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
122   if (num_elements(x) != num_elements(y)) {
123     std::ostringstream msg;
124     msg << "Inputs to vectorized binary function must match in"
125         << " size if one is not a scalar";
126     throw std::invalid_argument(msg.str());
127   }
128   using return_st = decltype(f(x(0), y[0][0]));
129   Eigen::Matrix<return_st, Eigen::Dynamic, Eigen::Dynamic> result(x.rows(),
130                                                                   x.cols());
131   for (size_t i = 0; i < y.size(); ++i) {
132     result.row(i) = apply_scalar_binary(x.row(i).transpose(),
133                                         as_column_vector_or_scalar(y[i]), f);
134   }
135   return result;
136 }
137 
138 /**
139  * Specialisation for use with a two-dimensional std::vector of integer types
140  * and one Eigen matrix.
141  *
142  * @tparam T1 Type of first argument to which functor is applied.
143  * @tparam T2 Type of second argument to which functor is applied.
144  * @tparam F Type of functor to apply.
145  * @param x Nested integer std::vector input to which operation is applied.
146  * @param y Eigen matrix input to which operation is applied.
147  * @param f functor to apply to inputs.
148  * @return Eigen object with result of applying functor to inputs.
149  */
150 template <typename T1, typename T2, typename F,
151           require_std_vector_vt<is_std_vector, T1>* = nullptr,
152           require_std_vector_st<std::is_integral, T1>* = nullptr,
153           require_eigen_matrix_dynamic_vt<is_stan_scalar, T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)154 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
155   if (num_elements(x) != num_elements(y)) {
156     std::ostringstream msg;
157     msg << "Inputs to vectorized binary function must match in"
158         << " size if one is not a scalar";
159     throw std::invalid_argument(msg.str());
160   }
161   using return_st = decltype(f(x[0][0], y(0)));
162   Eigen::Matrix<return_st, Eigen::Dynamic, Eigen::Dynamic> result(y.rows(),
163                                                                   y.cols());
164   for (size_t i = 0; i < x.size(); ++i) {
165     result.row(i) = apply_scalar_binary(as_column_vector_or_scalar(x[i]),
166                                         y.row(i).transpose(), f);
167   }
168   return result;
169 }
170 
171 /**
172  * Specialisation for use when the first input is an Eigen type and the second
173  * is a scalar. Eigen's unaryExpr framework is used for more efficient indexing
174  * of both row- and column-major inputs. The unaryExpr framework also allows
175  * for the scalar to be captured and applied to each element in the Eigen
176  * object.
177  *
178  * @tparam T1 Type of Eigen object to which functor is applied.
179  * @tparam T2 Type of scalar to which functor is applied.
180  * @tparam F Type of functor to apply.
181  * @param x Eigen input to which operation is applied.
182  * @param y Scalar input to which operation is applied.
183  * @param f functor to apply to Eigen and scalar inputs.
184  * @return Eigen object with result of applying functor to inputs.
185  *
186  * Note: The return expresssion needs to be evaluated, otherwise the captured
187  *         function and scalar fall out of scope.
188  */
189 template <typename T1, typename T2, typename F, require_eigen_t<T1>* = nullptr,
190           require_stan_scalar_t<T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)191 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
192   return x.unaryExpr([&f, y](const auto& v) { return f(v, y); });
193 }
194 
195 /**
196  * Specialisation for use when the first input is an scalar and the second is
197  * an Eigen type. Eigen's unaryExpr framework is used for more efficient
198  * indexing of both row- and column-major inputs. The unaryExpr framework also
199  * allows for the scalar to be captured and applied to each element in the
200  * Eigen object.
201  *
202  * @tparam T1 Type of scalar to which functor is applied.
203  * @tparam T2 Type of Eigen object to which functor is applied.
204  * @tparam F Type of functor to apply.
205  * @param x Scalar input to which operation is applied.
206  * @param y Eigen input to which operation is applied.
207  * @param f functor to apply to Eigen and scalar inputs.
208  * @return Eigen object with result of applying functor to inputs.
209  *
210  * Note: The return expresssion needs to be evaluated, otherwise the captured
211  *         function and scalar fall out of scope.
212  */
213 template <typename T1, typename T2, typename F,
214           require_stan_scalar_t<T1>* = nullptr, require_eigen_t<T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)215 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
216   return y.unaryExpr([&f, x](const auto& v) { return f(x, v); });
217 }
218 
219 /**
220  * Specialisation for use with (non-nested) std::vectors. Inputs are mapped
221  * to Eigen column vectors and then the result is evaluated directly into the
222  * returned std::vector (via Eigen::Map).
223  *
224  * The returned scalar type is deduced to allow for cases where the input and
225  * return scalar types differ (e.g., functions implicitly promoting
226  * integers).
227  *
228  * @tparam T1 Type of first std::vector to which functor is applied.
229  * @tparam T2 Type of second std::vector to which functor is applied.
230  * @tparam F Type of functor to apply.
231  * @param x First std::vector input to which operation is applied.
232  * @param y Second std::vector input to which operation is applied.
233  * @param f functor to apply to std::vector inputs.
234  * @return std::vector with result of applying functor to inputs.
235  */
236 template <typename T1, typename T2, typename F,
237           require_all_std_vector_vt<is_stan_scalar, T1, T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)238 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
239   check_matching_sizes("Binary function", "x", x, "y", y);
240   decltype(auto) x_vec = as_column_vector_or_scalar(x);
241   decltype(auto) y_vec = as_column_vector_or_scalar(y);
242   using T_return = std::decay_t<decltype(f(x[0], y[0]))>;
243   std::vector<T_return> result(x.size());
244   Eigen::Map<Eigen::Matrix<T_return, -1, 1>>(result.data(), result.size())
245       = x_vec.binaryExpr(y_vec, f);
246   return result;
247 }
248 
249 /**
250  * Specialisation for use when the first input is a (non-nested) std::vector
251  * and the second is a scalar. The std::vector input is mapped to an Eigen
252  * column vector and then the result is evaluated directly into the returned
253  * std::vector (via Eigen::Map).
254  *
255  * The returned scalar type is deduced to allow for cases where the input and
256  * return scalar types differ (e.g., functions implicitly promoting
257  * integers).
258  *
259  * @tparam T1 Type of std::vector to which functor is applied.
260  * @tparam T2 Type of scalar to which functor is applied.
261  * @tparam F Type of functor to apply.
262  * @param x std::vector input to which operation is applied.
263  * @param y Scalar input to which operation is applied.
264  * @param f functor to apply to std::vector and scalar inputs.
265  * @return std::vector with result of applying functor to inputs.
266  */
267 template <typename T1, typename T2, typename F,
268           require_std_vector_vt<is_stan_scalar, T1>* = nullptr,
269           require_stan_scalar_t<T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)270 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
271   decltype(auto) x_vec = as_column_vector_or_scalar(x);
272   using T_return = std::decay_t<decltype(f(x[0], y))>;
273   std::vector<T_return> result(x.size());
274   Eigen::Map<Eigen::Matrix<T_return, -1, 1>>(result.data(), result.size())
275       = x_vec.unaryExpr([&f, &y](const auto& v) { return f(v, y); });
276   return result;
277 }
278 
279 /**
280  * Specialisation for use when the first input is a scalar and the second is a
281  * (non-nested) std::vector. The std::vector input is mapped to an Eigen
282  * column vector and then the result is evaluated directly into the returned
283  * std::vector (via Eigen::Map).
284  *
285  * The returned scalar type is deduced to allow for cases where the input and
286  * return scalar types differ (e.g., functions implicitly promoting
287  * integers).
288  *
289  * @tparam T1 Type of scalar to which functor is applied.
290  * @tparam T2 Type of std::vector to which functor is applied.
291  * @tparam F Type of functor to apply.
292  * @param x Scalar input to which operation is applied.
293  * @param y std::vector input to which operation is applied.
294  * @param f functor to apply to std::vector and scalar inputs.
295  * @return std::vector with result of applying functor to inputs.
296  */
297 template <typename T1, typename T2, typename F,
298           require_stan_scalar_t<T1>* = nullptr,
299           require_std_vector_vt<is_stan_scalar, T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)300 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
301   decltype(auto) y_vec = as_column_vector_or_scalar(y);
302   using T_return = std::decay_t<decltype(f(x, y[0]))>;
303   std::vector<T_return> result(y.size());
304   Eigen::Map<Eigen::Matrix<T_return, -1, 1>>(result.data(), result.size())
305       = y_vec.unaryExpr([&f, &x](const auto& v) { return f(x, v); });
306   return result;
307 }
308 
309 /**
310  * Specialisation for use with two nested containers (std::vectors).
311  * The returned scalar type is deduced to allow for cases where the input and
312  * return scalar types differ (e.g., functions implicitly promoting
313  * integers).
314  *
315  * @tparam T1 Type of first std::vector to which functor is applied.
316  * @tparam T2 Type of second std::vector to which functor is applied.
317  * @tparam F Type of functor to apply.
318  * @param x First std::vector input to which operation is applied.
319  * @param y Second std::vector input to which operation is applied.
320  * @param f functor to apply to std::vector inputs.
321  * @return std::vector with result of applying functor to inputs.
322  */
323 template <
324     typename T1, typename T2, typename F,
325     require_all_std_vector_vt<is_container_or_var_matrix, T1, T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)326 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
327   check_matching_sizes("Binary function", "x", x, "y", y);
328   using T_return = plain_type_t<decltype(apply_scalar_binary(x[0], y[0], f))>;
329   size_t y_size = y.size();
330   std::vector<T_return> result(y_size);
331   for (size_t i = 0; i < y_size; ++i) {
332     result[i] = apply_scalar_binary(x[i], y[i], f);
333   }
334   return result;
335 }
336 
337 /**
338  * Specialisation for use when the first input is a nested std::vector and the
339  * second is a scalar. The returned scalar type is deduced to allow for cases
340  * where the input and return scalar types differ (e.g., functions implicitly
341  * promoting integers).
342  *
343  * @tparam T1 Type of std::vector to which functor is applied.
344  * @tparam T2 Type of scalar to which functor is applied.
345  * @tparam F Type of functor to apply.
346  * @param x std::vector input to which operation is applied.
347  * @param y Scalar input to which operation is applied.
348  * @param f functor to apply to inputs.
349  * @return std::vector with result of applying functor to inputs.
350  */
351 template <typename T1, typename T2, typename F,
352           require_std_vector_vt<is_container_or_var_matrix, T1>* = nullptr,
353           require_stan_scalar_t<T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)354 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
355   using T_return = plain_type_t<decltype(apply_scalar_binary(x[0], y, f))>;
356   size_t x_size = x.size();
357   std::vector<T_return> result(x_size);
358   for (size_t i = 0; i < x_size; ++i) {
359     result[i] = apply_scalar_binary(x[i], y, f);
360   }
361   return result;
362 }
363 
364 /**
365  * Specialisation for use when the first input is a scalar and the second is a
366  * nested std::vector. The returned scalar type is deduced to allow for cases
367  * where the input and return scalar types differ (e.g., functions implicitly
368  * promoting integers).
369  *
370  * @tparam T1 Type of scalar to which functor is applied.
371  * @tparam T2 Type of std::vector to which functor is applied.
372  * @tparam F Type of functor to apply.
373  * @param x Scalar input to which operation is applied.
374  * @param y std::vector input to which operation is applied.
375  * @param f functor to apply to inputs.
376  * @return std::vector with result of applying functor to inputs.
377  */
378 template <typename T1, typename T2, typename F,
379           require_stan_scalar_t<T1>* = nullptr,
380           require_std_vector_vt<is_container_or_var_matrix, T2>* = nullptr>
apply_scalar_binary(const T1 & x,const T2 & y,const F & f)381 inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
382   using T_return = plain_type_t<decltype(apply_scalar_binary(x, y[0], f))>;
383   size_t y_size = y.size();
384   std::vector<T_return> result(y_size);
385   for (size_t i = 0; i < y_size; ++i) {
386     result[i] = apply_scalar_binary(x, y[i], f);
387   }
388   return result;
389 }
390 
391 }  // namespace math
392 }  // namespace stan
393 #endif
394