Skip to content

Move the beta, neg_binomial and student_t cdfs to the new grad_reg_inc_beta (part 3 of #3409) #3490

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
SteveBronder wants to merge 1 commit into from
+372 −528
Draft

Move the beta, neg_binomial and student_t cdfs to the new grad_reg_inc_beta (part 3 of #3409) #3490

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 33 additions & 53 deletions stan/math/prim/prob/beta_cdf.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,8 @@
#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/fun/digamma.hpp>
#include <stan/math/prim/fun/grad_reg_inc_beta.hpp>
#include <stan/math/prim/fun/inc_beta.hpp>
#include <stan/math/prim/fun/inc_beta_dda.hpp>
#include <stan/math/prim/fun/inc_beta_ddb.hpp>
#include <stan/math/prim/fun/inc_beta_ddz.hpp>
#include <stan/math/prim/fun/inv.hpp>
#include <stan/math/prim/fun/max_size.hpp>
Expand Down Expand Up @@ -59,9 +58,6 @@ inline return_type_t<T_y, T_scale_succ, T_scale_fail> beta_cdf(
scalar_seq_view<T_y_ref> y_vec(y_ref);
scalar_seq_view<T_alpha_ref> alpha_vec(alpha_ref);
scalar_seq_view<T_beta_ref> beta_vec(beta_ref);
size_t size_alpha = stan::math::size(alpha);
size_t size_beta = stan::math::size(beta);
size_t size_alpha_beta = max_size(alpha, beta);
size_t N = max_size(y, alpha, beta);

// Explicit return for extreme values
Expand All @@ -72,64 +68,48 @@ inline return_type_t<T_y, T_scale_succ, T_scale_fail> beta_cdf(
}
}

VectorBuilder<is_autodiff_v<T_scale_succ>, T_partials_return, T_scale_succ>
digamma_alpha(size_alpha);
if constexpr (is_autodiff_v<T_scale_succ>) {
for (size_t n = 0; n < size_alpha; n++) {
digamma_alpha[n] = digamma(alpha_vec.val(n));
}
}

VectorBuilder<is_autodiff_v<T_scale_fail>, T_partials_return, T_scale_fail>
digamma_beta(size_beta);
if constexpr (is_autodiff_v<T_scale_fail>) {
for (size_t n = 0; n < size_beta; n++) {
digamma_beta[n] = digamma(beta_vec.val(n));
}
}

VectorBuilder<is_any_autodiff_v<T_scale_succ, T_scale_fail>,
T_partials_return, T_scale_succ, T_scale_fail>
digamma_sum(size_alpha_beta);
if constexpr (is_any_autodiff_v<T_scale_succ, T_scale_fail>) {
for (size_t n = 0; n < size_alpha_beta; n++) {
digamma_sum[n] = digamma(alpha_vec.val(n) + beta_vec.val(n));
}
}

constexpr bool calc_alpha = is_autodiff_v<T_scale_succ>;
constexpr bool calc_beta = is_autodiff_v<T_scale_fail>;
for (size_t n = 0; n < N; n++) {
const T_partials_return y_dbl = y_vec.val(n);
const auto y_dbl = y_vec.val(n);

// Explicit results for extreme values
// The gradients are technically ill-defined, but treated as zero
if (y_dbl >= 1.0) {
continue;
}

const T_partials_return alpha_dbl = alpha_vec.val(n);
const T_partials_return beta_dbl = beta_vec.val(n);
const T_partials_return Pn = inc_beta(alpha_dbl, beta_dbl, y_dbl);
const T_partials_return inv_Pn
= is_constant_all<T_y, T_scale_succ, T_scale_fail>::value ? 0 : inv(Pn);
const auto alpha_dbl = alpha_vec.val(n);
const auto beta_dbl = beta_vec.val(n);
const auto Pn = inc_beta(alpha_dbl, beta_dbl, y_dbl);

P *= Pn;

if constexpr (is_autodiff_v<T_y>) {
partials<0>(ops_partials)[n]
+= inc_beta_ddz(alpha_dbl, beta_dbl, y_dbl) * inv_Pn;
}

if constexpr (is_autodiff_v<T_scale_succ>) {
partials<1>(ops_partials)[n]
+= inc_beta_dda(alpha_dbl, beta_dbl, y_dbl, digamma_alpha[n],
digamma_sum[n])
* inv_Pn;
}
if constexpr (is_autodiff_v<T_scale_fail>) {
partials<2>(ops_partials)[n]
+= inc_beta_ddb(alpha_dbl, beta_dbl, y_dbl, digamma_beta[n],
digamma_sum[n])
* inv_Pn;
if constexpr (is_any_autodiff_v<T_y, T_scale_succ, T_scale_fail>) {
const auto inv_Pn = inv(Pn);
if constexpr (is_autodiff_v<T_y>) {
partials<0>(ops_partials)[n]
+= inc_beta_ddz(alpha_dbl, beta_dbl, y_dbl) * inv_Pn;
}

if constexpr (calc_alpha && calc_beta) {
const auto [g1, g2] = grad_reg_inc_beta<GRAD_RIB::BOTH>(
alpha_dbl, beta_dbl, y_dbl, digamma(alpha_dbl), digamma(beta_dbl),
digamma(alpha_dbl + beta_dbl), 0.0);
partials<1>(ops_partials)[n] += g1 * inv_Pn;
partials<2>(ops_partials)[n] += g2 * inv_Pn;
} else if constexpr (calc_alpha) {
partials<1>(ops_partials)[n]
+= grad_reg_inc_beta<GRAD_RIB::G1>(
alpha_dbl, beta_dbl, y_dbl, digamma(alpha_dbl), 0.0,
digamma(alpha_dbl + beta_dbl), 0.0)
* inv_Pn;
} else if constexpr (calc_beta) {
partials<2>(ops_partials)[n]
+= grad_reg_inc_beta<GRAD_RIB::G2>(
alpha_dbl, beta_dbl, y_dbl, 0.0, digamma(beta_dbl),
digamma(alpha_dbl + beta_dbl), 0.0)
* inv_Pn;
}
}
}

Expand Down
65 changes: 24 additions & 41 deletions stan/math/prim/prob/beta_lccdf.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,40 +67,17 @@ inline return_type_t<T_y, T_scale_succ, T_scale_fail> beta_lccdf(
scalar_seq_view<T_y_ref> y_vec(y_ref);
scalar_seq_view<T_alpha_ref> alpha_vec(alpha_ref);
scalar_seq_view<T_beta_ref> beta_vec(beta_ref);
size_t size_alpha = stan::math::size(alpha);
size_t size_beta = stan::math::size(beta_param);
size_t size_alpha_beta = max_size(alpha, beta_param);
size_t N = max_size(y, alpha, beta_param);

VectorBuilder<is_any_autodiff_v<T_scale_succ, T_scale_fail>,
T_partials_return, T_scale_succ>
digamma_alpha(size_alpha);
VectorBuilder<is_any_autodiff_v<T_scale_succ, T_scale_fail>,
T_partials_return, T_scale_fail>
digamma_beta(size_beta);
VectorBuilder<is_any_autodiff_v<T_scale_succ, T_scale_fail>,
T_partials_return, T_scale_succ, T_scale_fail>
digamma_sum(size_alpha_beta);

if constexpr (is_any_autodiff_v<T_scale_succ, T_scale_fail>) {
for (size_t i = 0; i < size_alpha; i++) {
digamma_alpha[i] = digamma(alpha_vec.val(i));
}
for (size_t i = 0; i < size_beta; i++) {
digamma_beta[i] = digamma(beta_vec.val(i));
}
for (size_t i = 0; i < size_alpha_beta; i++) {
digamma_sum[i] = digamma(alpha_vec.val(i) + beta_vec.val(i));
}
}

constexpr bool calc_alpha = is_autodiff_v<T_scale_succ>;
constexpr bool calc_beta = is_autodiff_v<T_scale_fail>;
for (size_t n = 0; n < N; n++) {
const T_partials_return y_dbl = y_vec.val(n);
const T_partials_return alpha_dbl = alpha_vec.val(n);
const T_partials_return beta_dbl = beta_vec.val(n);
const T_partials_return betafunc_dbl = beta(alpha_dbl, beta_dbl);
const T_partials_return Pn = 1.0 - inc_beta(alpha_dbl, beta_dbl, y_dbl);
const T_partials_return inv_Pn
const auto y_dbl = y_vec.val(n);
const auto alpha_dbl = alpha_vec.val(n);
const auto beta_dbl = beta_vec.val(n);
const auto betafunc_dbl = beta(alpha_dbl, beta_dbl);
const auto Pn = 1.0 - inc_beta(alpha_dbl, beta_dbl, y_dbl);
const auto inv_Pn
= is_constant_all<T_y, T_scale_succ, T_scale_fail>::value ? 0 : inv(Pn);

ccdf_log += log(Pn);
Expand All @@ -111,18 +88,24 @@ inline return_type_t<T_y, T_scale_succ, T_scale_fail> beta_lccdf(
/ betafunc_dbl;
}

T_partials_return g1 = 0;
T_partials_return g2 = 0;

if constexpr (is_any_autodiff_v<T_scale_succ, T_scale_fail>) {
grad_reg_inc_beta(g1, g2, alpha_dbl, beta_dbl, y_dbl, digamma_alpha[n],
digamma_beta[n], digamma_sum[n], betafunc_dbl);
}
if constexpr (is_autodiff_v<T_scale_succ>) {
if constexpr (calc_alpha && calc_beta) {
const auto [g1, g2] = grad_reg_inc_beta<GRAD_RIB::BOTH>(
alpha_dbl, beta_dbl, y_dbl, digamma(alpha_dbl), digamma(beta_dbl),
digamma(alpha_dbl + beta_dbl), 0.0);
partials<1>(ops_partials)[n] -= g1 * inv_Pn;
}
if constexpr (is_autodiff_v<T_scale_fail>) {
partials<2>(ops_partials)[n] -= g2 * inv_Pn;
} else if constexpr (calc_alpha) {
partials<1>(ops_partials)[n]
-= grad_reg_inc_beta<GRAD_RIB::G1>(alpha_dbl, beta_dbl, y_dbl,
digamma(alpha_dbl), 0.0,
digamma(alpha_dbl + beta_dbl), 0.0)
* inv_Pn;
} else if constexpr (calc_beta) {
partials<2>(ops_partials)[n]
-= grad_reg_inc_beta<GRAD_RIB::G2>(alpha_dbl, beta_dbl, y_dbl, 0.0,
digamma(beta_dbl),
digamma(alpha_dbl + beta_dbl), 0.0)
* inv_Pn;
}
}
return ops_partials.build(ccdf_log);
Expand Down
67 changes: 26 additions & 41 deletions stan/math/prim/prob/beta_lcdf.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,40 +58,18 @@ inline return_type_t<T_y, T_scale_succ, T_scale_fail> beta_lcdf(
scalar_seq_view<T_alpha_ref> alpha_vec(alpha_ref);
scalar_seq_view<T_beta_ref> beta_vec(beta_ref);

const size_t size_alpha = stan::math::size(alpha);
const size_t size_beta = stan::math::size(beta_param);
const size_t size_alpha_beta = max_size(alpha, beta_param);
const size_t N = max_size(y, alpha, beta_param);

// Allocate digamma buffers only if alpha/beta contain any autodiff scalars.
constexpr bool need_digamma = is_any_autodiff_v<T_scale_succ, T_scale_fail>;
VectorBuilder<need_digamma, T_partials_return, T_scale_succ> digamma_alpha(
size_alpha);
VectorBuilder<need_digamma, T_partials_return, T_scale_fail> digamma_beta(
size_beta);
VectorBuilder<need_digamma, T_partials_return, T_scale_succ, T_scale_fail>
digamma_sum(size_alpha_beta);

if constexpr (is_any_autodiff_v<T_scale_succ, T_scale_fail>) {
for (size_t i = 0; i < size_alpha; ++i) {
digamma_alpha[i] = digamma(alpha_vec.val(i));
}
for (size_t i = 0; i < size_beta; ++i) {
digamma_beta[i] = digamma(beta_vec.val(i));
}
for (size_t i = 0; i < size_alpha_beta; ++i) {
digamma_sum[i] = digamma(alpha_vec.val(i) + beta_vec.val(i));
}
}

constexpr bool calc_alpha = is_autodiff_v<T_scale_succ>;
constexpr bool calc_beta = is_autodiff_v<T_scale_fail>;
for (size_t n = 0; n < N; ++n) {
const T_partials_return y_dbl = y_vec.val(n);
const T_partials_return alpha_dbl = alpha_vec.val(n);
const T_partials_return beta_dbl = beta_vec.val(n);
const T_partials_return betafunc_dbl = beta(alpha_dbl, beta_dbl);
const T_partials_return Pn = inc_beta(alpha_dbl, beta_dbl, y_dbl);
const auto y_dbl = y_vec.val(n);
const auto alpha_dbl = alpha_vec.val(n);
const auto beta_dbl = beta_vec.val(n);
const auto betafunc_dbl = beta(alpha_dbl, beta_dbl);
const auto Pn = inc_beta(alpha_dbl, beta_dbl, y_dbl);

const T_partials_return inv_Pn
const auto inv_Pn
= is_any_autodiff_v<T_y, T_scale_succ, T_scale_fail> ? inv(Pn) : 0;

cdf_log += log(Pn);
Expand All @@ -102,17 +80,24 @@ inline return_type_t<T_y, T_scale_succ, T_scale_fail> beta_lcdf(
/ betafunc_dbl;
}

if constexpr (is_any_autodiff_v<T_scale_succ, T_scale_fail>) {
T_partials_return g1 = 0;
T_partials_return g2 = 0;
grad_reg_inc_beta(g1, g2, alpha_dbl, beta_dbl, y_dbl, digamma_alpha[n],
digamma_beta[n], digamma_sum[n], betafunc_dbl);
if constexpr (is_any_autodiff_v<T_scale_succ>) {
partials<1>(ops_partials)[n] += g1 * inv_Pn;
}
if constexpr (is_any_autodiff_v<T_scale_fail>) {
partials<2>(ops_partials)[n] += g2 * inv_Pn;
}
if constexpr (calc_alpha && calc_beta) {
const auto [g1, g2] = grad_reg_inc_beta<GRAD_RIB::BOTH>(
alpha_dbl, beta_dbl, y_dbl, digamma(alpha_dbl), digamma(beta_dbl),
digamma(alpha_dbl + beta_dbl), 0.0);
partials<1>(ops_partials)[n] += g1 * inv_Pn;
partials<2>(ops_partials)[n] += g2 * inv_Pn;
} else if constexpr (calc_alpha) {
partials<1>(ops_partials)[n]
+= grad_reg_inc_beta<GRAD_RIB::G1>(alpha_dbl, beta_dbl, y_dbl,
digamma(alpha_dbl), 0.0,
digamma(alpha_dbl + beta_dbl), 0.0)
* inv_Pn;
} else if constexpr (calc_beta) {
partials<2>(ops_partials)[n]
+= grad_reg_inc_beta<GRAD_RIB::G2>(alpha_dbl, beta_dbl, y_dbl, 0.0,
digamma(beta_dbl),
digamma(alpha_dbl + beta_dbl), 0.0)
* inv_Pn;
}
}

Expand Down
48 changes: 11 additions & 37 deletions stan/math/prim/prob/beta_proportion_lccdf.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,32 +71,8 @@ inline return_type_t<T_y, T_loc, T_prec> beta_proportion_lccdf(
scalar_seq_view<T_y_ref> y_vec(y_ref);
scalar_seq_view<T_mu_ref> mu_vec(mu_ref);
scalar_seq_view<T_kappa_ref> kappa_vec(kappa_ref);
size_t size_kappa = stan::math::size(kappa);
size_t size_mu_kappa = max_size(mu, kappa);
size_t N = max_size(y, mu, kappa);

VectorBuilder<is_any_autodiff_v<T_loc, T_prec>, T_partials_return, T_loc,
T_prec>
digamma_mukappa(size_mu_kappa);
VectorBuilder<is_any_autodiff_v<T_loc, T_prec>, T_partials_return, T_loc,
T_prec>
digamma_kappa_mukappa(size_mu_kappa);
VectorBuilder<is_any_autodiff_v<T_loc, T_prec>, T_partials_return, T_prec>
digamma_kappa(size_kappa);

if constexpr (is_any_autodiff_v<T_loc, T_prec>) {
for (size_t i = 0; i < size_mu_kappa; i++) {
const T_partials_return kappa_dbl = kappa_vec.val(i);
const T_partials_return mukappa_dbl = mu_vec.val(i) * kappa_dbl;
digamma_mukappa[i] = digamma(mukappa_dbl);
digamma_kappa_mukappa[i] = digamma(kappa_dbl - mukappa_dbl);
}

for (size_t i = 0; i < size_kappa; i++) {
digamma_kappa[i] = digamma(kappa_vec.val(i));
}
}

for (size_t n = 0; n < N; n++) {
const T_partials_return y_dbl = y_vec.val(n);
const T_partials_return mu_dbl = mu_vec.val(n);
Expand All @@ -118,20 +94,18 @@ inline return_type_t<T_y, T_loc, T_prec> beta_proportion_lccdf(
/ betafunc_dbl;
}

T_partials_return g1 = 0;
T_partials_return g2 = 0;

// the mu and kappa partials each need both shape partials
if constexpr (is_any_autodiff_v<T_loc, T_prec>) {
grad_reg_inc_beta(g1, g2, mukappa_dbl, kappa_mukappa_dbl, y_dbl,
digamma_mukappa[n], digamma_kappa_mukappa[n],
digamma_kappa[n], betafunc_dbl);
}
if constexpr (is_autodiff_v<T_loc>) {
partials<1>(ops_partials)[n] -= kappa_dbl * (g1 - g2) * inv_Pn;
}
if constexpr (is_autodiff_v<T_prec>) {
partials<2>(ops_partials)[n]
-= (g1 * mu_dbl + g2 * (1 - mu_dbl)) * inv_Pn;
const auto [g1, g2] = grad_reg_inc_beta<GRAD_RIB::BOTH>(
mukappa_dbl, kappa_mukappa_dbl, y_dbl, digamma(mukappa_dbl),
digamma(kappa_mukappa_dbl), digamma(kappa_dbl), 0.0);
if constexpr (is_autodiff_v<T_loc>) {
partials<1>(ops_partials)[n] -= kappa_dbl * (g1 - g2) * inv_Pn;
}
if constexpr (is_autodiff_v<T_prec>) {
partials<2>(ops_partials)[n]
-= (g1 * mu_dbl + g2 * (1 - mu_dbl)) * inv_Pn;
}
}
}
return ops_partials.build(ccdf_log);
Expand Down
Loading
Loading