Skip to content

Commit fb6ae11

Browse files
bbbales2bbbales2
authored
Merge pull request #2140 from bstatcomp/generalize_binomial
Generalizes binomial
2 parents a258497 + 42b2a84 commit fb6ae11

File tree

7 files changed

+110
-162
lines changed

7 files changed

+110
-162
lines changed

‎stan/math/prim/prob/binomial_cdf.hpp

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -35,26 +35,33 @@ template <typename T_n, typename T_N, typename T_prob>
3535
return_type_t<T_prob> binomial_cdf(const T_n& n, const T_N& N,
3636
const T_prob& theta) {
3737
using T_partials_return = partials_return_t<T_n, T_N, T_prob>;
38+
using T_n_ref = ref_type_t<T_n>;
39+
using T_N_ref = ref_type_t<T_N>;
40+
using T_theta_ref = ref_type_t<T_prob>;
3841
using std::exp;
3942
using std::pow;
4043
static const char* function = "binomial_cdf";
41-
check_nonnegative(function, "Population size parameter", N);
42-
check_finite(function, "Probability parameter", theta);
43-
check_bounded(function, "Probability parameter", theta, 0.0, 1.0);
4444
check_consistent_sizes(function, "Successes variable", n,
4545
"Population size parameter", N,
4646
"Probability parameter", theta);
4747

48+
T_n_ref n_ref = n;
49+
T_N_ref N_ref = N;
50+
T_theta_ref theta_ref = theta;
51+
52+
check_nonnegative(function, "Population size parameter", N_ref);
53+
check_bounded(function, "Probability parameter", theta_ref, 0.0, 1.0);
54+
4855
if (size_zero(n, N, theta)) {
4956
return 1.0;
5057
}
5158

5259
T_partials_return P(1.0);
53-
operands_and_partials<T_prob> ops_partials(theta);
60+
operands_and_partials<T_theta_ref> ops_partials(theta_ref);
5461

55-
scalar_seq_view<T_n> n_vec(n);
56-
scalar_seq_view<T_N> N_vec(N);
57-
scalar_seq_view<T_prob> theta_vec(theta);
62+
scalar_seq_view<T_n_ref> n_vec(n_ref);
63+
scalar_seq_view<T_N_ref> N_vec(N_ref);
64+
scalar_seq_view<T_theta_ref> theta_vec(theta_ref);
5865
size_t max_size_seq_view = max_size(n, N, theta);
5966

6067
// Explicit return for extreme values

‎stan/math/prim/prob/binomial_lccdf.hpp

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -37,27 +37,34 @@ template <typename T_n, typename T_N, typename T_prob>
3737
return_type_t<T_prob> binomial_lccdf(const T_n& n, const T_N& N,
3838
const T_prob& theta) {
3939
using T_partials_return = partials_return_t<T_n, T_N, T_prob>;
40+
using T_n_ref = ref_type_t<T_n>;
41+
using T_N_ref = ref_type_t<T_N>;
42+
using T_theta_ref = ref_type_t<T_prob>;
4043
using std::exp;
4144
using std::log;
4245
using std::pow;
4346
static const char* function = "binomial_lccdf";
44-
check_nonnegative(function, "Population size parameter", N);
45-
check_finite(function, "Probability parameter", theta);
46-
check_bounded(function, "Probability parameter", theta, 0.0, 1.0);
4747
check_consistent_sizes(function, "Successes variable", n,
4848
"Population size parameter", N,
4949
"Probability parameter", theta);
5050

51+
T_n_ref n_ref = n;
52+
T_N_ref N_ref = N;
53+
T_theta_ref theta_ref = theta;
54+
55+
check_nonnegative(function, "Population size parameter", N_ref);
56+
check_bounded(function, "Probability parameter", theta_ref, 0.0, 1.0);
57+
5158
if (size_zero(n, N, theta)) {
5259
return 0;
5360
}
5461

5562
T_partials_return P(0.0);
56-
operands_and_partials<T_prob> ops_partials(theta);
63+
operands_and_partials<T_theta_ref> ops_partials(theta_ref);
5764

58-
scalar_seq_view<T_n> n_vec(n);
59-
scalar_seq_view<T_N> N_vec(N);
60-
scalar_seq_view<T_prob> theta_vec(theta);
65+
scalar_seq_view<T_n_ref> n_vec(n_ref);
66+
scalar_seq_view<T_N_ref> N_vec(N_ref);
67+
scalar_seq_view<T_theta_ref> theta_vec(theta_ref);
6168
size_t max_size_seq_view = max_size(n, N, theta);
6269

6370
// Explicit return for extreme values

‎stan/math/prim/prob/binomial_lcdf.hpp

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -37,27 +37,34 @@ template <typename T_n, typename T_N, typename T_prob>
3737
return_type_t<T_prob> binomial_lcdf(const T_n& n, const T_N& N,
3838
const T_prob& theta) {
3939
using T_partials_return = partials_return_t<T_n, T_N, T_prob>;
40+
using T_n_ref = ref_type_t<T_n>;
41+
using T_N_ref = ref_type_t<T_N>;
42+
using T_theta_ref = ref_type_t<T_prob>;
4043
using std::exp;
4144
using std::log;
4245
using std::pow;
4346
static const char* function = "binomial_lcdf";
44-
check_nonnegative(function, "Population size parameter", N);
45-
check_finite(function, "Probability parameter", theta);
46-
check_bounded(function, "Probability parameter", theta, 0.0, 1.0);
4747
check_consistent_sizes(function, "Successes variable", n,
4848
"Population size parameter", N,
4949
"Probability parameter", theta);
5050

51+
T_n_ref n_ref = n;
52+
T_N_ref N_ref = N;
53+
T_theta_ref theta_ref = theta;
54+
55+
check_nonnegative(function, "Population size parameter", N_ref);
56+
check_bounded(function, "Probability parameter", theta_ref, 0.0, 1.0);
57+
5158
if (size_zero(n, N, theta)) {
5259
return 0;
5360
}
5461

5562
T_partials_return P(0.0);
56-
operands_and_partials<T_prob> ops_partials(theta);
63+
operands_and_partials<T_theta_ref> ops_partials(theta_ref);
5764

58-
scalar_seq_view<T_n> n_vec(n);
59-
scalar_seq_view<T_N> N_vec(N);
60-
scalar_seq_view<T_prob> theta_vec(theta);
65+
scalar_seq_view<T_n_ref> n_vec(n_ref);
66+
scalar_seq_view<T_N_ref> N_vec(N_ref);
67+
scalar_seq_view<T_theta_ref> theta_vec(theta_ref);
6168
size_t max_size_seq_view = max_size(n, N, theta);
6269

6370
// Explicit return for extreme values

‎stan/math/prim/prob/binomial_logit_lpmf.hpp

Lines changed: 44 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
#include <stan/math/prim/fun/max_size.hpp>
1212
#include <stan/math/prim/fun/size.hpp>
1313
#include <stan/math/prim/fun/size_zero.hpp>
14+
#include <stan/math/prim/fun/to_ref.hpp>
1415
#include <stan/math/prim/fun/value_of.hpp>
1516
#include <stan/math/prim/functor/operands_and_partials.hpp>
1617

@@ -34,75 +35,65 @@ template <bool propto, typename T_n, typename T_N, typename T_prob>
3435
return_type_t<T_prob> binomial_logit_lpmf(const T_n& n, const T_N& N,
3536
const T_prob& alpha) {
3637
using T_partials_return = partials_return_t<T_n, T_N, T_prob>;
37-
using std::log;
38+
using T_n_ref = ref_type_if_t<!is_constant<T_n>::value, T_n>;
39+
using T_N_ref = ref_type_if_t<!is_constant<T_N>::value, T_N>;
40+
using T_alpha_ref = ref_type_if_t<!is_constant<T_prob>::value, T_prob>;
3841
static const char* function = "binomial_logit_lpmf";
39-
check_bounded(function, "Successes variable", n, 0, N);
40-
check_nonnegative(function, "Population size parameter", N);
41-
check_finite(function, "Probability parameter", alpha);
4242
check_consistent_sizes(function, "Successes variable", n,
4343
"Population size parameter", N,
4444
"Probability parameter", alpha);
4545

46-
if (size_zero(n, N, alpha)) {
47-
return 0.0;
48-
}
49-
if (!include_summand<propto, T_prob>::value) {
50-
return 0.0;
51-
}
46+
T_n_ref n_ref = n;
47+
T_N_ref N_ref = N;
48+
T_alpha_ref alpha_ref = alpha;
5249

53-
T_partials_return logp = 0;
54-
operands_and_partials<T_prob> ops_partials(alpha);
50+
const auto& n_col = as_column_vector_or_scalar(n_ref);
51+
const auto& N_col = as_column_vector_or_scalar(N_ref);
52+
const auto& alpha_col = as_column_vector_or_scalar(alpha_ref);
5553

56-
scalar_seq_view<T_n> n_vec(n);
57-
scalar_seq_view<T_N> N_vec(N);
58-
scalar_seq_view<T_prob> alpha_vec(alpha);
59-
size_t size_alpha = stan::math::size(alpha);
60-
size_t max_size_seq_view = max_size(n, N, alpha);
54+
const auto& n_arr = as_array_or_scalar(n_col);
55+
const auto& N_arr = as_array_or_scalar(N_col);
56+
const auto& alpha_arr = as_array_or_scalar(alpha_col);
6157

62-
if (include_summand<propto>::value) {
63-
for (size_t i = 0; i < max_size_seq_view; ++i) {
64-
logp += binomial_coefficient_log(N_vec[i], n_vec[i]);
65-
}
66-
}
58+
ref_type_t<decltype(value_of(n_arr))> n_val = value_of(n_arr);
59+
ref_type_t<decltype(value_of(N_arr))> N_val = value_of(N_arr);
60+
ref_type_t<decltype(value_of(alpha_arr))> alpha_val = value_of(alpha_arr);
6761

68-
VectorBuilder<true, T_partials_return, T_prob> inv_logit_alpha(size_alpha);
69-
VectorBuilder<true, T_partials_return, T_prob> inv_logit_neg_alpha(
70-
size_alpha);
71-
VectorBuilder<true, T_partials_return, T_prob> log_inv_logit_alpha(
72-
size_alpha);
73-
VectorBuilder<true, T_partials_return, T_prob> log_inv_logit_neg_alpha(
74-
size_alpha);
62+
check_bounded(function, "Successes variable", n_val, 0, N_val);
63+
check_nonnegative(function, "Population size parameter", N_val);
64+
check_finite(function, "Probability parameter", alpha_val);
7565

76-
for (size_t i = 0; i < size_alpha; ++i) {
77-
const T_partials_return alpha_dbl = value_of(alpha_vec[i]);
78-
inv_logit_alpha[i] = inv_logit(alpha_dbl);
79-
inv_logit_neg_alpha[i] = inv_logit(-alpha_dbl);
80-
log_inv_logit_alpha[i] = log(inv_logit_alpha[i]);
81-
log_inv_logit_neg_alpha[i] = log(inv_logit_neg_alpha[i]);
66+
if (size_zero(n, N, alpha)) {
67+
return 0.0;
68+
}
69+
if (!include_summand<propto, T_prob>::value) {
70+
return 0.0;
8271
}
72+
const auto& inv_logit_alpha
73+
= to_ref_if<!is_constant_all<T_prob>::value>(inv_logit(alpha_val));
74+
const auto& inv_logit_neg_alpha
75+
= to_ref_if<!is_constant_all<T_prob>::value>(inv_logit(-alpha_val));
8376

84-
for (size_t i = 0; i < max_size_seq_view; ++i) {
85-
logp += n_vec[i] * log_inv_logit_alpha[i]
86-
+ (N_vec[i] - n_vec[i]) * log_inv_logit_neg_alpha[i];
77+
size_t maximum_size = max_size(n, N, alpha);
78+
const auto& log_inv_logit_alpha = log(inv_logit_alpha);
79+
const auto& log_inv_logit_neg_alpha = log(inv_logit_neg_alpha);
80+
T_partials_return logp = sum(n_val * log_inv_logit_alpha
81+
+ (N_val - n_val) * log_inv_logit_neg_alpha);
82+
if (include_summand<propto, T_n, T_N>::value) {
83+
logp += sum(binomial_coefficient_log(N_val, n_val)) * maximum_size
84+
/ max_size(n, N);
8785
}
8886

87+
operands_and_partials<T_alpha_ref> ops_partials(alpha_ref);
8988
if (!is_constant_all<T_prob>::value) {
90-
if (size_alpha == 1) {
91-
T_partials_return sum_n = 0;
92-
T_partials_return sum_N = 0;
93-
for (size_t i = 0; i < max_size_seq_view; ++i) {
94-
sum_n += n_vec[i];
95-
sum_N += N_vec[i];
96-
}
97-
ops_partials.edge1_.partials_[0]
98-
+= sum_n * inv_logit_neg_alpha[0]
99-
- (sum_N - sum_n) * inv_logit_alpha[0];
89+
if (is_vector<T_prob>::value) {
90+
ops_partials.edge1_.partials_
91+
= n_val * inv_logit_neg_alpha - (N_val - n_val) * inv_logit_alpha;
10092
} else {
101-
for (size_t i = 0; i < max_size_seq_view; ++i) {
102-
ops_partials.edge1_.partials_[i]
103-
+= n_vec[i] * inv_logit_neg_alpha[i]
104-
- (N_vec[i] - n_vec[i]) * inv_logit_alpha[i];
105-
}
93+
T_partials_return sum_n = sum(n_val) * maximum_size / size(n);
94+
ops_partials.edge1_.partials_[0] = forward_as<T_partials_return>(
95+
sum_n * inv_logit_neg_alpha
96+
- (sum(N_val) * maximum_size / size(N) - sum_n) * inv_logit_alpha);
10697
}
10798
}
10899

‎stan/math/prim/prob/binomial_lpmf.hpp

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -36,15 +36,23 @@ template <bool propto, typename T_n, typename T_N, typename T_prob>
3636
return_type_t<T_prob> binomial_lpmf(const T_n& n, const T_N& N,
3737
const T_prob& theta) {
3838
using T_partials_return = partials_return_t<T_n, T_N, T_prob>;
39+
using T_n_ref = ref_type_t<T_n>;
40+
using T_N_ref = ref_type_t<T_N>;
41+
using T_theta_ref = ref_type_t<T_prob>;
3942
static const char* function = "binomial_lpmf";
40-
check_bounded(function, "Successes variable", n, 0, N);
41-
check_nonnegative(function, "Population size parameter", N);
42-
check_finite(function, "Probability parameter", theta);
43-
check_bounded(function, "Probability parameter", theta, 0.0, 1.0);
4443
check_consistent_sizes(function, "Successes variable", n,
4544
"Population size parameter", N,
4645
"Probability parameter", theta);
4746

47+
T_n_ref n_ref = n;
48+
T_N_ref N_ref = N;
49+
T_theta_ref theta_ref = theta;
50+
51+
check_bounded(function, "Successes variable", n_ref, 0, N_ref);
52+
check_nonnegative(function, "Population size parameter", N_ref);
53+
check_finite(function, "Probability parameter", theta_ref);
54+
check_bounded(function, "Probability parameter", theta_ref, 0.0, 1.0);
55+
4856
if (size_zero(n, N, theta)) {
4957
return 0.0;
5058
}
@@ -53,11 +61,11 @@ return_type_t<T_prob> binomial_lpmf(const T_n& n, const T_N& N,
5361
}
5462

5563
T_partials_return logp = 0;
56-
operands_and_partials<T_prob> ops_partials(theta);
64+
operands_and_partials<T_theta_ref> ops_partials(theta_ref);
5765

58-
scalar_seq_view<T_n> n_vec(n);
59-
scalar_seq_view<T_N> N_vec(N);
60-
scalar_seq_view<T_prob> theta_vec(theta);
66+
scalar_seq_view<T_n_ref> n_vec(n_ref);
67+
scalar_seq_view<T_N_ref> N_vec(N_ref);
68+
scalar_seq_view<T_theta_ref> theta_vec(theta_ref);
6169
size_t size_theta = stan::math::size(theta);
6270
size_t max_size_seq_view = max_size(n, N, theta);
6371

‎stan/math/prim/prob/binomial_rng.hpp

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,14 +32,18 @@ inline typename VectorBuilder<true, int, T_N, T_theta>::type binomial_rng(
3232
const T_N& N, const T_theta& theta, RNG& rng) {
3333
using boost::binomial_distribution;
3434
using boost::variate_generator;
35+
using T_N_ref = ref_type_t<T_N>;
36+
using T_theta_ref = ref_type_t<T_theta>;
3537
static const char* function = "binomial_rng";
36-
check_nonnegative(function, "Population size parameter", N);
37-
check_bounded(function, "Probability parameter", theta, 0.0, 1.0);
3838
check_consistent_sizes(function, "Population size parameter", N,
3939
"Probability Parameter", theta);
40+
T_N_ref N_ref = N;
41+
T_theta_ref theta_ref = theta;
42+
check_nonnegative(function, "Population size parameter", N_ref);
43+
check_bounded(function, "Probability parameter", theta_ref, 0.0, 1.0);
4044

41-
scalar_seq_view<T_N> N_vec(N);
42-
scalar_seq_view<T_theta> theta_vec(theta);
45+
scalar_seq_view<T_N_ref> N_vec(N_ref);
46+
scalar_seq_view<T_theta_ref> theta_vec(theta_ref);
4347
size_t M = max_size(N, theta);
4448
VectorBuilder<true, int, T_N, T_theta> output(M);
4549

0 commit comments

Comments
 (0)