From 7ae8e42dacaa3b179a9b6eaf593bd266bb2f75a2 Mon Sep 17 00:00:00 2001 From: Aki Vehtari Date: Sun, 20 Sep 2026 22:44:29 +0300 Subject: [PATCH 01/10] Compute the incomplete beta shape gradients by continued fraction --- licenses/tensorflow-probability-license.txt | 206 ++++++++++++ stan/math/prim/fun/grad_reg_inc_beta.hpp | 347 ++++++++++++++++++-- 2 files changed, 528 insertions(+), 25 deletions(-) create mode 100644 licenses/tensorflow-probability-license.txt diff --git a/licenses/tensorflow-probability-license.txt b/licenses/tensorflow-probability-license.txt new file mode 100644 index 00000000000..b92e39eaa05 --- /dev/null +++ b/licenses/tensorflow-probability-license.txt @@ -0,0 +1,206 @@ +License of the code that stan/math/prim/fun/grad_reg_inc_beta.hpp is derived +from: TensorFlow Probability, tensorflow_probability/python/math/special.py, +Copyright 2022 The TensorFlow Probability Authors. +https://github.com/tensorflow/probability + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/stan/math/prim/fun/grad_reg_inc_beta.hpp b/stan/math/prim/fun/grad_reg_inc_beta.hpp index f4fa1b58178..83049f76975 100644 --- a/stan/math/prim/fun/grad_reg_inc_beta.hpp +++ b/stan/math/prim/fun/grad_reg_inc_beta.hpp @@ -1,27 +1,321 @@ +// Original code derived from TensorFlow Probability and is distributed here +// under the Apache License 2.0 (licenses/tensorflow-probability-license.txt) +// Copyright 2022 The TensorFlow Probability Authors. +// tensorflow_probability/python/math/special.py, betainc partials +// Secondary code copyright by its author and is distributed here +// under the BSD-3 license (LICENSE.md) + #ifndef STAN_MATH_PRIM_FUN_GRAD_REG_INC_BETA_HPP #define STAN_MATH_PRIM_FUN_GRAD_REG_INC_BETA_HPP #include -#include -#include +#include +#include +#include +#include +#include +#include +#include +#include +#include #include +#include namespace stan { namespace math { +namespace internal { + +/** + * Partial derivatives of the regularized incomplete beta function + * I_x(a, b) with respect to a and b, by the power series + * + * I_x(a, b) = x^a 2F1(a, 1 - b; a + 1; x) / (a B(a, b)), + * + * differentiated term by term. Used where at least one of these holds: + * x < a / (a + b), b x <= 1, x <= 0.95 + * x >= a / (a + b), a (1 - x) <= 1, x >= 0.05 + * For x >= a / (a + b) the symmetry relation I_x(a, b) = 1 - I_{1-x}(b, a) + * is applied first. Stops when a term is below eps / a in magnitude, or + * after 600 terms. + * + * @tparam T scalar type + * @param a first shape + * @param b second shape + * @param x upper limit, in (0, 1) + * @param digamma_a digamma(a) + * @param digamma_b digamma(b) + * @param digamma_ab digamma(a + b) + * @param[out] g1 d/da I_x(a, b) + * @param[out] g2 d/db I_x(a, b) + */ +template +inline void inc_beta_partials_power_series(T a, T b, T x, T digamma_a, + T digamma_b, T digamma_ab, T& g1, + T& g2) { + using std::exp; + using std::log; + const bool use_symmetry_relation + = value_of_rec(x) >= value_of_rec(a) / value_of_rec(a + b); + if (use_symmetry_relation) { + const T a_orig = a; + a = b; + b = a_orig; + x = 1.0 - x; + const T digamma_a_orig = digamma_a; + digamma_a = digamma_b; + digamma_b = digamma_a_orig; + } + + constexpr int max_iterations = 600; + const double tolerance + = std::numeric_limits::epsilon() / value_of_rec(a); + + // series for 2F1(a, 1 - b; a + 1; x) / a and its a and b derivatives + T product = 1.0; + T series_sum = inv(a); + T product_grad_b = 0.0; + T da = -inv(square(a)); + T db = 0.0; + for (int n = 1; n <= max_iterations; ++n) { + const T x_div_n = x / n; + const T factor = (n - b) * x_div_n; + const T apn = a + n; + const T new_product = product * factor; + const T term = new_product / apn; + product_grad_b = factor * product_grad_b - product * x_div_n; + const T db_increment = product_grad_b / apn; + da -= new_product / square(apn); + db += db_increment; + product = new_product; + series_sum += term; + // Differs from the TensorFlow Probability source: the stop test also + // requires the d/db increment to be small. For a positive integer b the + // value series is a polynomial whose term is exactly 0 at n = b, while + // the d/db series continues. + if (fabs(value_of_rec(term)) <= tolerance + && fabs(value_of_rec(db_increment)) <= tolerance) { + break; + } + } + + const T normalization = exp(a * log(x) - lbeta(a, b)); + const T grad_a + = normalization * (da + series_sum * (digamma_ab - digamma_a + log(x))); + const T grad_b = normalization * (db + series_sum * (digamma_ab - digamma_b)); + + if (use_symmetry_relation) { + g1 = -grad_b; + g2 = -grad_a; + } else { + g1 = grad_a; + g2 = grad_b; + } +} + +/** + * One step of the modified Lentz method for the continued fraction of the + * incomplete beta function, carrying the derivatives of the ratios with + * respect to a and b. Thompson and Barnett (1986), appendix. + * + * @tparam T scalar type + * @param d partial numerator d_n + * @param dd_a d/da of d_n + * @param dd_b d/db of d_n + * @param[in,out] C ratio of successive numerators A_n / A_{n-1} + * @param[in,out] D ratio of successive denominators B_{n-1} / B_n + * @param[in,out] h convergent A_n / B_n + * @param[in,out] dC_a d/da of C + * @param[in,out] dC_b d/db of C + * @param[in,out] dD_a d/da of D + * @param[in,out] dD_b d/db of D + * @param[in,out] dh_a d/da of h + * @param[in,out] dh_b d/db of h + * @return the multiplicative update of h at this step + */ +template +inline T inc_beta_lentz_step(const T& d, const T& dd_a, const T& dd_b, T& C, + T& D, T& h, T& dC_a, T& dC_b, T& dD_a, T& dD_b, + T& dh_a, T& dh_b) { + const double small = std::sqrt(std::numeric_limits::min()); + + T C_new = 1.0 + d / C; + if (fabs(value_of_rec(C_new)) < small) { + C_new = small; + } + T D_new = 1.0 + d * D; + if (fabs(value_of_rec(D_new)) < small) { + D_new = small; + } + D_new = inv(D_new); + const T delta = C_new * D_new; + const T h_new = h * delta; + + const T dC_a_new = (dd_a * C - d * dC_a) / square(C); + const T dC_b_new = (dd_b * C - d * dC_b) / square(C); + const T dD_a_new = -(dd_a * D + d * dD_a) * square(D_new); + const T dD_b_new = -(dd_b * D + d * dD_b) * square(D_new); + const T dh_a_new = dh_a * delta + h * dC_a_new * D_new + h * dD_a_new * C_new; + const T dh_b_new = dh_b * delta + h * dC_b_new * D_new + h * dD_b_new * C_new; + + C = C_new; + D = D_new; + h = h_new; + dC_a = dC_a_new; + dC_b = dC_b_new; + dD_a = dD_a_new; + dD_b = dD_b_new; + dh_a = dh_a_new; + dh_b = dh_b_new; + return delta; +} + +/** + * Partial derivatives of the regularized incomplete beta function + * I_x(a, b) with respect to a and b, by the continued fraction + * https://dlmf.nist.gov/8.17.E22 evaluated with the modified Lentz method, + * with the partial numerators https://dlmf.nist.gov/8.17.E23 differentiated + * in a and b and the derivatives carried through the recurrence. Used + * where the power series region does not apply. The continued fraction + * converges rapidly for x < (a - 1) / (a + b - 2); otherwise the symmetry + * relation I_x(a, b) = 1 - I_{1-x}(b, a) is applied first. Stops when the + * multiplicative update is within 3 eps of 1, or after 300 double steps. + * + * @tparam T scalar type + * @param a first shape + * @param b second shape + * @param x upper limit, in (0, 1) + * @param digamma_a digamma(a) + * @param digamma_b digamma(b) + * @param digamma_ab digamma(a + b) + * @param[out] g1 d/da I_x(a, b) + * @param[out] g2 d/db I_x(a, b) + */ +template +inline void inc_beta_partials_continued_fraction(T a, T b, T x, T digamma_a, + T digamma_b, T digamma_ab, + T& g1, T& g2) { + using std::exp; + using std::log; + const double a_d = value_of_rec(a); + const double b_d = value_of_rec(b); + // IEEE division: (a - 1) / 0 is +-inf and the comparison is well defined + const bool use_symmetry_relation + = value_of_rec(x) >= (a_d - 1.0) / (a_d + b_d - 2.0); + if (use_symmetry_relation) { + const T a_orig = a; + a = b; + b = a_orig; + x = 1.0 - x; + const T digamma_a_orig = digamma_a; + digamma_a = digamma_b; + digamma_b = digamma_a_orig; + } + + constexpr int max_iterations = 300; + const double tolerance = 3.0 * std::numeric_limits::epsilon(); + const double small = std::sqrt(std::numeric_limits::min()); + + const T apb = a + b; + const T ap1 = a + 1.0; + + // initialization and first step of the modified Lentz method + T C = 1.0; + T D = 1.0 - apb * x / ap1; + if (fabs(value_of_rec(D)) < small) { + D = small; + } + D = inv(D); + T h = D; + const T dD_denom = square(x * apb - ap1); + T dD_a = (1.0 - b) * x / dD_denom; + T dD_b = ap1 * x / dD_denom; + T dC_a = 0.0; + T dC_b = 0.0; + T dh_a = dD_a; + T dh_b = dD_b; + + for (int m = 1; m <= max_iterations; ++m) { + const double dh_a_prev = value_of_rec(dh_a); + const double dh_b_prev = value_of_rec(dh_b); + // even partial numerator d_{2m}, https://dlmf.nist.gov/8.17.E23 + { + const T a_plus_2m = a + 2.0 * m; + const T a_plus_2m_minus_one = a_plus_2m - 1.0; + const T denominator = a_plus_2m * a_plus_2m_minus_one; + const T dd_b = m * x / denominator; + const T d = dd_b * (b - m); + const T dd_a = -d * (a_plus_2m + a_plus_2m_minus_one) / denominator; + inc_beta_lentz_step(d, dd_a, dd_b, C, D, h, dC_a, dC_b, dD_a, dD_b, dh_a, + dh_b); + } + // odd partial numerator d_{2m+1} + T delta; + { + const T a_plus_m = a + m; + const T a_plus_2m = a_plus_m + m; + const T a_plus_2m_plus_one = a_plus_2m + 1.0; + const T a_plus_b_plus_m = a_plus_m + b; + const T denominator = a_plus_2m * a_plus_2m_plus_one; + const T dd_b = -a_plus_m * x / denominator; + const T d = dd_b * a_plus_b_plus_m; + const T dd_a = -d * ((a_plus_2m + a_plus_2m_plus_one) / denominator) + - x * (2.0 * a_plus_m + b) / denominator; + delta = inc_beta_lentz_step(d, dd_a, dd_b, C, D, h, dC_a, dC_b, dD_a, + dD_b, dh_a, dh_b); + } + // Differs from the TensorFlow Probability source: the stop test also + // requires the two derivatives to have converged. For a positive + // integer b the partial numerator d_{2m} is exactly 0 at m = b; the + // convergent is then exact and delta is 1, but the derivatives still + // need the tail of the fraction. + const double dh_a_now = value_of_rec(dh_a); + const double dh_b_now = value_of_rec(dh_b); + if (fabs(value_of_rec(delta) - 1.0) < tolerance + && fabs(dh_a_now - dh_a_prev) <= tolerance * fabs(dh_a_now) + && fabs(dh_b_now - dh_b_prev) <= tolerance * fabs(dh_b_now)) { + break; + } + } + + const T normalization + = exp(a * log(x) + b * log1p(-x) - log(a) - lbeta(a, b)); + const T grad_a + = normalization * (dh_a + h * (log(x) - inv(a) + digamma_ab - digamma_a)); + const T grad_b + = normalization * (dh_b + h * (log1p(-x) + digamma_ab - digamma_b)); + + if (use_symmetry_relation) { + g1 = -grad_b; + g2 = -grad_a; + } else { + g1 = grad_a; + g2 = grad_b; + } +} + +} // namespace internal + /** * Computes the gradients of the regularized incomplete beta * function. Specifically, this function computes gradients of - * ibeta(a, b, z), with respect to the arguments + * inc_beta(a, b, z), with respect to the arguments * a and b. * - * Uses the equivalence to a hypergeometric function. See - * http://dlmf.nist.gov/8.17#ii + * The algorithm is the one of Boik and Robinson-Cox (1998), "Derivatives + * of the Incomplete Beta Function", Journal of Statistical Software 3(1), + * in the form implemented by TensorFlow Probability: a power series near + * the endpoints and a continued fraction elsewhere, both differentiated + * term by term, with the prefactor evaluated in log space and the + * symmetry relation I_z(a, b) = 1 - I_{1-z}(b, a) applied where the + * expansion converges faster on the other side. + * + * The gradients are 0 at z = 0 and z = 1. * * @tparam T type of arguments - * @param[out] g1 partial derivative of ibeta(a, b, z) + * @param[out] g1 partial derivative of inc_beta(a, b, z) * with respect to a - * @param[out] g2 partial derivative of ibeta(a, b, + * @param[out] g2 partial derivative of inc_beta(a, b, * z) with respect to b * @param[in] a a * @param[in] b b @@ -29,32 +323,35 @@ namespace math { * @param[in] digammaA the value of digamma(a) * @param[in] digammaB the value of digamma(b) * @param[in] digammaSum the value of digamma(a + b) - * @param[in] betaAB the value of beta(a, b) + * @param[in] betaAB the value of beta(a, b); accepted for + * interface compatibility and not used, the normalization is evaluated in + * log space */ template inline void grad_reg_inc_beta(T& g1, T& g2, const T& a, const T& b, const T& z, const T& digammaA, const T& digammaB, const T& digammaSum, const T& betaAB) { - using std::exp; + const double a_d = value_of_rec(a); + const double b_d = value_of_rec(b); + const double z_d = value_of_rec(z); - T c1 = log(z); - T c2 = log1m(z); - T c3 = betaAB * inc_beta(a, b, z); - T C = exp(a * c1 + b * c2) / a; - T dF1 = 0; - T dF2 = 0; - T dF3 = 0; - T dFz = 0; - if (value_of_rec(C)) { - std::forward_as_tuple(dF1, dF2, dF3, dFz) - = grad_2F1(a + b, 1.0, a + 1, z); + if (z_d == 0.0 || z_d == 1.0) { + g1 = 0.0; + g2 = 0.0; + return; } - T dBda = (c1 - 1.0 / a) * c3 + C * (dF1 + dF3); - T dBdb = c2 * c3 + C * dF1; - - g1 = (dBda - c3 * (digammaA - digammaSum)) / betaAB; - g2 = (dBdb - c3 * (digammaB - digammaSum)) / betaAB; + const double mean = a_d / (a_d + b_d); + const bool use_power_series + = (z_d < mean && b_d * z_d <= 1.0 && z_d <= 0.95) + || (z_d >= mean && a_d * (1.0 - z_d) <= 1.0 && z_d >= 0.05); + if (use_power_series) { + internal::inc_beta_partials_power_series(a, b, z, digammaA, digammaB, + digammaSum, g1, g2); + } else { + internal::inc_beta_partials_continued_fraction(a, b, z, digammaA, digammaB, + digammaSum, g1, g2); + } } } // namespace math From 138e257c062c76e78a041749470850beb760a968 Mon Sep 17 00:00:00 2001 From: Aki Vehtari Date: Sun, 20 Sep 2026 22:44:34 +0300 Subject: [PATCH 02/10] Route inc_beta_dda and inc_beta_ddb through grad_reg_inc_beta --- stan/math/prim/fun.hpp | 3 ++ stan/math/prim/fun/inc_beta_dda.hpp | 83 ++++------------------------- stan/math/prim/fun/inc_beta_ddb.hpp | 75 ++++---------------------- 3 files changed, 23 insertions(+), 138 deletions(-) diff --git a/stan/math/prim/fun.hpp b/stan/math/prim/fun.hpp index 8e0b01d2093..14589a7cee1 100644 --- a/stan/math/prim/fun.hpp +++ b/stan/math/prim/fun.hpp @@ -133,6 +133,9 @@ #include #include #include +#include +#include +#include #include #include #include diff --git a/stan/math/prim/fun/inc_beta_dda.hpp b/stan/math/prim/fun/inc_beta_dda.hpp index 4ed2c0e4758..f45ab7cc5ac 100644 --- a/stan/math/prim/fun/inc_beta_dda.hpp +++ b/stan/math/prim/fun/inc_beta_dda.hpp @@ -2,13 +2,8 @@ #define STAN_MATH_PRIM_FUN_INC_BETA_DDA_HPP #include -#include -#include -#include -#include -#include -#include -#include +#include +#include namespace stan { namespace math { @@ -16,19 +11,16 @@ namespace math { /** * Returns the partial derivative of the regularized * incomplete beta function, I_{z}(a, b) with respect to a. - * The power series used to compute the derivative tends to - * converge slowly when a and b are large, especially if z - * approaches 1. The implementation will throw an exception - * if the series have not converged within 100,000 iterations. - * The current implementation has been tested for values - * of a and b up to 12500 and z = 0.999. + * + * Evaluated by grad_reg_inc_beta, which computes both shape + * derivatives; call that function directly when both are needed. * * @tparam T scalar types of arguments * @param a first argument * @param b second argument * @param z upper bound of the integral * @param digamma_a value of digamma(a) - * @param digamma_ab value of digamma(b) + * @param digamma_ab value of digamma(a + b) * @return partial derivative of the incomplete beta with respect to a * * @pre a >= 0 @@ -37,65 +29,10 @@ namespace math { */ template inline T inc_beta_dda(T a, T b, T z, T digamma_a, T digamma_ab) { - using std::fabs; - using std::log; - using std::pow; - - if (b > a) { - if ((0.1 < z && z <= 0.75 && b > 500) || (0.01 < z && z <= 0.1 && b > 2500) - || (0.001 < z && z <= 0.01 && b > 1e5)) { - return -inc_beta_ddb(b, a, 1 - z, digamma_a, digamma_ab); - } - } - - if (z > 0.75 && a < 500) { - return -inc_beta_ddb(b, a, 1 - z, digamma_a, digamma_ab); - } - if (z > 0.9 && a < 2500) { - return -inc_beta_ddb(b, a, 1 - z, digamma_a, digamma_ab); - } - if (z > 0.99 && a < 1e5) { - return -inc_beta_ddb(b, a, 1 - z, digamma_a, digamma_ab); - } - if (z > 0.999) { - return -inc_beta_ddb(b, a, 1 - z, digamma_a, digamma_ab); - } - - double threshold = 1e-10; - - const T a_plus_b = a + b; - const T a_plus_1 = a + 1; - - digamma_a += inv(a); // Need digamma(a + 1), not digamma(a); - - // Common prefactor to regularize numerator and denominator - T prefactor = pow(a_plus_1 / a_plus_b, 3); - - T sum_numer = (digamma_ab - digamma_a) * prefactor; - T sum_denom = prefactor; - - T summand = prefactor * z * a_plus_b / a_plus_1; - - T k = 1; - digamma_ab += inv(a_plus_b); - digamma_a += inv(a_plus_1); - - while (fabs(summand) > threshold) { - sum_numer += (digamma_ab - digamma_a) * summand; - sum_denom += summand; - - summand *= (1 + (a_plus_b) / k) * (1 + k) / (1 + a_plus_1 / k); - digamma_ab += inv(a_plus_b + k); - digamma_a += inv(a_plus_1 + k); - ++k; - summand *= z / k; - - if (k > 1e5) { - throw_domain_error("inc_beta_dda", - "did not converge within 10000 iterations", "", ""); - } - } - return inc_beta(a, b, z) * (log(z) + sum_numer / sum_denom); + T g1 = 0; + T g2 = 0; + grad_reg_inc_beta(g1, g2, a, b, z, digamma_a, digamma(b), digamma_ab, T(0)); + return g1; } } // namespace math diff --git a/stan/math/prim/fun/inc_beta_ddb.hpp b/stan/math/prim/fun/inc_beta_ddb.hpp index 5514743f3e9..a4d440da80a 100644 --- a/stan/math/prim/fun/inc_beta_ddb.hpp +++ b/stan/math/prim/fun/inc_beta_ddb.hpp @@ -2,36 +2,25 @@ #define STAN_MATH_PRIM_FUN_INC_BETA_DDB_HPP #include -#include -#include -#include -#include -#include -#include -#include +#include +#include namespace stan { namespace math { -template -inline T inc_beta_dda(T a, T b, T z, T digamma_a, T digamma_ab); - /** * Returns the partial derivative of the regularized * incomplete beta function, I_{z}(a, b) with respect to b. - * The power series used to compute the derivative tends to - * converge slowly when a and b are large, especially if z - * approaches 1. The implementation will throw an exception - * if the series have not converged within 100,000 iterations. - * The current implementation has been tested for values - * of a and b up to 12500 and z = 0.999. + * + * Evaluated by grad_reg_inc_beta, which computes both shape + * derivatives; call that function directly when both are needed. * * @tparam T scalar types of arguments * @param a first argument * @param b second argument * @param z upper bound of the integral * @param digamma_b value of digamma(b) - * @param digamma_ab value of digamma(b) + * @param digamma_ab value of digamma(a + b) * @return partial derivative of the incomplete beta with respect to b * * @pre a >= 0 @@ -40,54 +29,10 @@ inline T inc_beta_dda(T a, T b, T z, T digamma_a, T digamma_ab); */ template inline T inc_beta_ddb(T a, T b, T z, T digamma_b, T digamma_ab) { - using std::fabs; - using std::log; - using std::pow; - - if (b > a) { - if ((0.1 < z && z <= 0.75 && b > 500) || (0.01 < z && z <= 0.1 && b > 2500) - || (0.001 < z && z <= 0.01 && b > 1e5)) { - return -inc_beta_dda(b, a, 1 - z, digamma_b, digamma_ab); - } - } - - if ((z > 0.75 && a < 500) || (z > 0.9 && a < 2500) || (z > 0.99 && a < 1e5) - || (z > 0.999)) { - return -inc_beta_dda(b, a, 1 - z, digamma_b, digamma_ab); - } - - double threshold = 1e-10; - - const T a_plus_b = a + b; - const T a_plus_1 = a + 1; - - // Common prefactor to regularize numerator and denominator - T prefactor = pow(a_plus_1 / a_plus_b, 3); - - T sum_numer = digamma_ab * prefactor; - T sum_denom = prefactor; - - T summand = prefactor * z * a_plus_b / a_plus_1; - - T k = 1; - digamma_ab += inv(a_plus_b); - - while (fabs(summand) > threshold) { - sum_numer += digamma_ab * summand; - sum_denom += summand; - - summand *= (1 + (a_plus_b) / k) * (1 + k) / (1 + a_plus_1 / k); - digamma_ab += inv(a_plus_b + k); - ++k; - summand *= z / k; - - if (k > 1e5) { - throw_domain_error("inc_beta_ddb", - "did not converge within 100000 iterations", "", ""); - } - } - - return inc_beta(a, b, z) * (log1m(z) - digamma_b + sum_numer / sum_denom); + T g1 = 0; + T g2 = 0; + grad_reg_inc_beta(g1, g2, a, b, z, digamma(a), digamma_b, digamma_ab, T(0)); + return g2; } } // namespace math From eb89df5051bab27b8e4a81ec6a64f7950720aeba Mon Sep 17 00:00:00 2001 From: Aki Vehtari Date: Sun, 20 Sep 2026 22:44:40 +0300 Subject: [PATCH 03/10] Compute both beta_cdf shape partials in one root call --- stan/math/prim/prob/beta_cdf.hpp | 49 +++++++++++++++----------------- 1 file changed, 23 insertions(+), 26 deletions(-) diff --git a/stan/math/prim/prob/beta_cdf.hpp b/stan/math/prim/prob/beta_cdf.hpp index b099c8df6ab..77abbe6b3a6 100644 --- a/stan/math/prim/prob/beta_cdf.hpp +++ b/stan/math/prim/prob/beta_cdf.hpp @@ -4,9 +4,8 @@ #include #include #include +#include #include -#include -#include #include #include #include @@ -72,26 +71,24 @@ inline return_type_t beta_cdf( } } - VectorBuilder, T_partials_return, T_scale_succ> + // grad_reg_inc_beta computes both shape derivatives in one pass and needs + // all three digamma values, so they are built when either shape is + // autodiff. + constexpr bool need_shape_grads + = is_any_autodiff_v; + VectorBuilder digamma_alpha(size_alpha); - if constexpr (is_autodiff_v) { + VectorBuilder digamma_beta( + size_beta); + VectorBuilder + digamma_sum(size_alpha_beta); + if constexpr (need_shape_grads) { for (size_t n = 0; n < size_alpha; n++) { digamma_alpha[n] = digamma(alpha_vec.val(n)); } - } - - VectorBuilder, T_partials_return, T_scale_fail> - digamma_beta(size_beta); - if constexpr (is_autodiff_v) { for (size_t n = 0; n < size_beta; n++) { digamma_beta[n] = digamma(beta_vec.val(n)); } - } - - VectorBuilder, - T_partials_return, T_scale_succ, T_scale_fail> - digamma_sum(size_alpha_beta); - if constexpr (is_any_autodiff_v) { for (size_t n = 0; n < size_alpha_beta; n++) { digamma_sum[n] = digamma(alpha_vec.val(n) + beta_vec.val(n)); } @@ -119,17 +116,17 @@ inline return_type_t beta_cdf( += inc_beta_ddz(alpha_dbl, beta_dbl, y_dbl) * inv_Pn; } - if constexpr (is_autodiff_v) { - 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) { - partials<2>(ops_partials)[n] - += inc_beta_ddb(alpha_dbl, beta_dbl, y_dbl, digamma_beta[n], - digamma_sum[n]) - * inv_Pn; + if constexpr (need_shape_grads) { + 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], T_partials_return(0)); + if constexpr (is_autodiff_v) { + partials<1>(ops_partials)[n] += g1 * inv_Pn; + } + if constexpr (is_autodiff_v) { + partials<2>(ops_partials)[n] += g2 * inv_Pn; + } } } From 671d5dedd3e683e2193390b58758f3860f2219d8 Mon Sep 17 00:00:00 2001 From: Aki Vehtari Date: Sun, 20 Sep 2026 22:44:47 +0300 Subject: [PATCH 04/10] Test the incomplete beta shape gradients against fixed references --- .../math/prim/fun/grad_reg_inc_beta_test.cpp | 80 +++++++++++++--- test/unit/math/prim/fun/inc_beta_dda_test.cpp | 84 +++++++++------- test/unit/math/prim/fun/inc_beta_ddb_test.cpp | 96 ++++++++++--------- test/unit/math/rev/fun/inc_beta_test.cpp | 43 +++++++++ test/unit/math/rev/prob/beta_cdf_log_test.cpp | 44 +++++++++ test/unit/math/rev/prob/beta_cdf_test.cpp | 38 ++++++++ 6 files changed, 290 insertions(+), 95 deletions(-) create mode 100644 test/unit/math/rev/fun/inc_beta_test.cpp create mode 100644 test/unit/math/rev/prob/beta_cdf_log_test.cpp create mode 100644 test/unit/math/rev/prob/beta_cdf_test.cpp diff --git a/test/unit/math/prim/fun/grad_reg_inc_beta_test.cpp b/test/unit/math/prim/fun/grad_reg_inc_beta_test.cpp index 129028fbab7..881057a5cf6 100644 --- a/test/unit/math/prim/fun/grad_reg_inc_beta_test.cpp +++ b/test/unit/math/prim/fun/grad_reg_inc_beta_test.cpp @@ -1,8 +1,33 @@ #include #include +#include #include +namespace { + +// Both shape derivatives of I_z(a, b) against a high-precision reference +// (mpmath, 80 digits, density-form quadrature checked against numerical +// differentiation at 110 digits). +void expect_grads(double a, double b, double z, double dda, double ddb, + double rtol) { + using stan::math::digamma; + using stan::math::lbeta; + double g1 = 0; + double g2 = 0; + stan::math::grad_reg_inc_beta(g1, g2, a, b, z, digamma(a), digamma(b), + digamma(a + b), std::exp(lbeta(a, b))); + EXPECT_TRUE(std::isfinite(g1)) << "a=" << a << " b=" << b << " z=" << z; + EXPECT_TRUE(std::isfinite(g2)) << "a=" << a << " b=" << b << " z=" << z; + EXPECT_NEAR(dda, g1, rtol * std::fabs(dda)) + << "a=" << a << " b=" << b << " z=" << z; + EXPECT_NEAR(ddb, g2, rtol * std::fabs(ddb)) + << "a=" << a << " b=" << b << " z=" << z; +} + +} // namespace + TEST(grad_reg_inc_beta, 1) { + // at y = 1 the function is identically 1 and both derivatives are 0 double alpha = 1.0; double beta = 1.0; double y = 1.0; @@ -16,22 +41,49 @@ TEST(grad_reg_inc_beta, 1) { stan::math::grad_reg_inc_beta(g1, g2, alpha, beta, y, digamma_alpha, digamma_beta, digamma_sum, betafunc); EXPECT_FLOAT_EQ(0, g1); - EXPECT_FLOAT_EQ(-std::numeric_limits::infinity(), g2); + EXPECT_FLOAT_EQ(0, g2); } TEST(grad_reg_inc_beta, 2) { - double alpha = 1.0; - double beta = 1.0; - double y = 0.4; - double digamma_alpha = stan::math::digamma(alpha); - double digamma_beta = stan::math::digamma(beta); - double digamma_sum = stan::math::digamma(alpha + beta); - double betafunc = std::exp(stan::math::lbeta(alpha, beta)); + expect_grads(1.0, 1.0, 0.4, -0.36651629274966202, 0.30649537425959442, 1e-12); +} - double g1 = 0; - double g2 = 0; - stan::math::grad_reg_inc_beta(g1, g2, alpha, beta, y, digamma_alpha, - digamma_beta, digamma_sum, betafunc); - EXPECT_NEAR(-0.36651629, g1, 1e-6); - EXPECT_NEAR(0.30649537, g2, 1e-6); +TEST(grad_reg_inc_beta, moderate) { + expect_grads(2, 3, 0.25, -0.22805360232434637, 0.10692153005229139, 1e-12); + expect_grads(0.5, 0.5, 0.5, -0.58312180806163756, 0.58312180806163756, 1e-12); + expect_grads(20, 20, 0.5, -0.063740468830439859, 0.063740468830439859, 1e-12); + expect_grads(1.5, 1.25, 0.6, -0.23806756196382869, 0.32279574755409757, + 1e-12); +} + +TEST(grad_reg_inc_beta, beta_ab_underflow) { + // beta(a, b) underflows for a + b above about 1100 + expect_grads(600, 600, 0.5, -0.01152047147304755, 0.01152047147304755, 1e-12); +} + +TEST(grad_reg_inc_beta, prefactor_underflow) { + // z^a (1 - z)^b underflows while I is 1 and the derivatives are 1e-138 + expect_grads(352, 590, 0.758, -3.3772717797123242e-138, + 4.5390666074606649e-138, 1e-12); +} + +TEST(grad_reg_inc_beta, integer_b) { + // for integer b the power series of the value terminates but the + // derivative series does not; the stop test must watch both + expect_grads(10, 1, 0.5, -0.00067690154351557159, 0.00226574342615196, 1e-12); + expect_grads(1, 100, 0.01, -0.42882413065860014, 0.0036787479630394138, + 1e-12); + expect_grads(2500, 2, 0.999, -0.00020509953516671709, 0.24850317856904407, + 1e-12); +} + +TEST(grad_reg_inc_beta, tails) { + expect_grads(1.145, 6786, 1.8e-4, -0.38621356822721375, 5.8439109664263782e-5, + 1e-10); + expect_grads(446, 5, 0.76, -1.0644675765185877e-47, 1.2763067943072555e-46, + 1e-12); + expect_grads(15000, 1.25, 0.999, -6.5954322569556673e-10, + 2.008497927784773e-6, 1e-10); + expect_grads(15000, 12500, 0.5, -8.7210244281937301e-53, + 9.5528279245811057e-53, 1e-10); } diff --git a/test/unit/math/prim/fun/inc_beta_dda_test.cpp b/test/unit/math/prim/fun/inc_beta_dda_test.cpp index ffe17385a6d..10c0f250c08 100644 --- a/test/unit/math/prim/fun/inc_beta_dda_test.cpp +++ b/test/unit/math/prim/fun/inc_beta_dda_test.cpp @@ -1,5 +1,21 @@ #include #include +#include + +namespace { + +// d/da of I_z(a, b) against a high-precision reference (mpmath, 80 digits, +// density-form quadrature checked against numerical differentiation). +void expect_dda(double a, double b, double z, double expected, double rtol) { + using stan::math::digamma; + using stan::math::inc_beta_dda; + const double got = inc_beta_dda(a, b, z, digamma(a), digamma(a + b)); + EXPECT_TRUE(std::isfinite(got)) << "a=" << a << " b=" << b << " z=" << z; + EXPECT_NEAR(expected, got, rtol * std::fabs(expected)) + << "a=" << a << " b=" << b << " z=" << z; +} + +} // namespace TEST(MathFunctions, inc_beta_dda) { using stan::math::digamma; @@ -15,49 +31,43 @@ TEST(MathFunctions, inc_beta_dda) { double mid_z = 0.6; double large_z = 0.999; - EXPECT_FLOAT_EQ(-0.00028665637, - inc_beta_dda(small_a, small_b, small_z, digamma(small_a), - digamma(small_a + small_b))) - << "reasonable values for a, b, x"; - EXPECT_FLOAT_EQ(-0.23806757, - inc_beta_dda(small_a, small_b, mid_z, digamma(small_a), - digamma(small_a + small_b))) - << "reasonable values for a, b, x"; - EXPECT_FLOAT_EQ(-0.00022264493, - inc_beta_dda(small_a, small_b, large_z, digamma(small_a), - digamma(small_a + small_b))) - << "reasonable values for a, b, x"; + expect_dda(small_a, small_b, small_z, -0.00028665636570426377, 1e-10); + expect_dda(small_a, small_b, mid_z, -0.23806756196382869, 1e-10); + expect_dda(small_a, small_b, large_z, -0.00022264492614285471, 1e-10); + // I underflows to 0 in double; the derivative is below 1e-3000 EXPECT_FLOAT_EQ(0.0, inc_beta_dda(large_a, small_b, small_z, digamma(large_a), - digamma(large_a + small_b))) - << "reasonable values for a, b, x"; + digamma(large_a + small_b))); EXPECT_FLOAT_EQ(0.0, inc_beta_dda(large_a, small_b, mid_z, digamma(large_a), - digamma(large_a + small_b))) - << "reasonable values for a, b, x"; - EXPECT_FLOAT_EQ(0.00091716705, - inc_beta_dda(large_a, small_b, large_z, digamma(large_a), - digamma(large_a + small_b))) - << "reasonable values for a, b, x"; - - EXPECT_FLOAT_EQ(1.8226241, - inc_beta_dda(small_a, large_b, small_z, digamma(small_a), - digamma(small_a + large_b))) - << "reasonable values for a, b, x"; + digamma(large_a + small_b))); + expect_dda(large_a, small_b, large_z, -6.5954322569556673e-10, 1e-10); + + expect_dda(small_a, large_b, small_z, -3.937563719959171e-5, 1e-10); EXPECT_FLOAT_EQ(0.0, inc_beta_dda(small_a, large_b, mid_z, digamma(small_a), - digamma(small_a + large_b))) - << "reasonable values for a, b, x"; + digamma(small_a + large_b))); EXPECT_FLOAT_EQ(0.0, inc_beta_dda(small_a, large_b, large_z, digamma(small_a), - digamma(small_a + large_b))) - << "reasonable values for a, b, x"; + digamma(small_a + large_b))); EXPECT_FLOAT_EQ(0.0, inc_beta_dda(large_a, large_b, small_z, digamma(large_a), - digamma(large_a + large_b))) - << "reasonable values for a, b, x"; - EXPECT_FLOAT_EQ(-4.2743586e-14, - inc_beta_dda(large_a, large_b, mid_z, digamma(large_a), - digamma(large_a + large_b))) - << "reasonable values for a, b, x"; + digamma(large_a + large_b))); + expect_dda(large_a, large_b, mid_z, -1.755811372560583e-76, 1e-10); EXPECT_FLOAT_EQ(0.0, inc_beta_dda(large_a, large_b, large_z, digamma(large_a), - digamma(large_a + large_b))) - << "reasonable values for a, b, x"; + digamma(large_a + large_b))); +} + +TEST(MathFunctions, inc_beta_dda_regressions) { + // b >> a with small z + expect_dda(1.145, 6786, 1.8e-4, -0.38621356822721375, 1e-10); + // deep lower tail, I of order 1e-47 + expect_dda(446, 5, 0.76, -1.0644675765185877e-47, 1e-10); + // large a with z near 1 + expect_dda(2500, 2, 0.999, -0.00020509953516671709, 1e-10); + expect_dda(1, 100, 0.01, -0.42882413065860014, 1e-10); + // integer b: the value series terminates but the derivative does not + expect_dda(2, 3, 0.25, -0.22805360232434637, 1e-10); + expect_dda(10, 1, 0.5, -0.00067690154351557159, 1e-10); + // a + b above the beta(a, b) underflow limit + expect_dda(600, 600, 0.5, -0.01152047147304755, 1e-10); + expect_dda(0.5, 0.5, 0.5, -0.58312180806163756, 1e-10); + expect_dda(20, 20, 0.5, -0.063740468830439859, 1e-10); } diff --git a/test/unit/math/prim/fun/inc_beta_ddb_test.cpp b/test/unit/math/prim/fun/inc_beta_ddb_test.cpp index d6116dd437a..dfc8d50cfb0 100644 --- a/test/unit/math/prim/fun/inc_beta_ddb_test.cpp +++ b/test/unit/math/prim/fun/inc_beta_ddb_test.cpp @@ -1,5 +1,22 @@ #include #include +#include + +namespace { + +// d/db of I_z(a, b) against a high-precision reference (mpmath, 80 digits, +// density-form quadrature checked against numerical differentiation). +// The fourth argument of inc_beta_ddb is digamma(b). +void expect_ddb(double a, double b, double z, double expected, double rtol) { + using stan::math::digamma; + using stan::math::inc_beta_ddb; + const double got = inc_beta_ddb(a, b, z, digamma(b), digamma(a + b)); + EXPECT_TRUE(std::isfinite(got)) << "a=" << a << " b=" << b << " z=" << z; + EXPECT_NEAR(expected, got, rtol * std::fabs(expected)) + << "a=" << a << " b=" << b << " z=" << z; +} + +} // namespace TEST(MathFunctions, inc_beta_ddb) { using stan::math::digamma; @@ -15,48 +32,39 @@ TEST(MathFunctions, inc_beta_ddb) { double mid_z = 0.5; double large_z = 0.999; - EXPECT_FLOAT_EQ(3.2996082e-05, - inc_beta_ddb(small_a, small_b, small_z, digamma(small_a), - digamma(small_a + small_b))) - << "reasonable values for a, b, x"; - EXPECT_FLOAT_EQ(0.17945045, - inc_beta_ddb(small_a, small_b, mid_z, digamma(small_a), - digamma(small_a + small_b))) - << "reasonable values for a, b, x"; - EXPECT_FLOAT_EQ(0.0019721228, - inc_beta_ddb(small_a, small_b, large_z, digamma(small_a), - digamma(small_a + small_b))) - << "reasonable values for a, b, x"; - - EXPECT_FLOAT_EQ(0.0, inc_beta_ddb(large_a, small_b, small_z, digamma(large_a), - digamma(large_a + small_b))) - << "reasonable values for a, b, x"; - EXPECT_FLOAT_EQ(0.0, inc_beta_ddb(large_a, small_b, mid_z, digamma(large_a), - digamma(large_a + small_b))) - << "reasonable values for a, b, x"; - EXPECT_FLOAT_EQ(7.7076669, - inc_beta_ddb(large_a, small_b, large_z, digamma(large_a), - digamma(large_a + small_b))) - << "reasonable values for a, b, x"; - - EXPECT_FLOAT_EQ(9.3959293, - inc_beta_ddb(small_a, large_b, small_z, digamma(small_a), - digamma(small_a + large_b))) - << "reasonable values for a, b, x"; - EXPECT_FLOAT_EQ(0.0, inc_beta_ddb(small_a, large_b, mid_z, digamma(small_a), - digamma(small_a + large_b))) - << "reasonable values for a, b, x"; - EXPECT_FLOAT_EQ(0.0, inc_beta_ddb(small_a, large_b, large_z, digamma(small_a), - digamma(small_a + large_b))) - << "reasonable values for a, b, x"; - - EXPECT_FLOAT_EQ(0.0, inc_beta_ddb(large_a, large_b, small_z, digamma(large_a), - digamma(large_a + large_b))) - << "reasonable values for a, b, x"; - EXPECT_FLOAT_EQ(0.0, inc_beta_ddb(large_a, large_b, mid_z, digamma(large_a), - digamma(large_a + large_b))) - << "reasonable values for a, b, x"; - EXPECT_FLOAT_EQ(0.0, inc_beta_ddb(large_a, large_b, large_z, digamma(large_a), - digamma(large_a + large_b))) - << "reasonable values for a, b, x"; + expect_ddb(small_a, small_b, small_z, 4.4135732827365698e-5, 1e-10); + expect_ddb(small_a, small_b, mid_z, 0.29301795311347039, 1e-10); + expect_ddb(small_a, small_b, large_z, 0.0018969609780283825, 1e-10); + + // I underflows to 0 in double; the derivative is below 1e-3000 + EXPECT_FLOAT_EQ(0.0, inc_beta_ddb(large_a, small_b, small_z, digamma(small_b), + digamma(large_a + small_b))); + EXPECT_FLOAT_EQ(0.0, inc_beta_ddb(large_a, small_b, mid_z, digamma(small_b), + digamma(large_a + small_b))); + expect_ddb(large_a, small_b, large_z, 2.008497927784773e-6, 1e-10); + + expect_ddb(small_a, large_b, small_z, 1.4782043391986094e-8, 1e-10); + EXPECT_FLOAT_EQ(0.0, inc_beta_ddb(small_a, large_b, mid_z, digamma(large_b), + digamma(small_a + large_b))); + EXPECT_FLOAT_EQ(0.0, inc_beta_ddb(small_a, large_b, large_z, digamma(large_b), + digamma(small_a + large_b))); + + EXPECT_FLOAT_EQ(0.0, inc_beta_ddb(large_a, large_b, small_z, digamma(large_b), + digamma(large_a + large_b))); + expect_ddb(large_a, large_b, mid_z, 9.5528279245811057e-53, 1e-10); + EXPECT_FLOAT_EQ(0.0, inc_beta_ddb(large_a, large_b, large_z, digamma(large_b), + digamma(large_a + large_b))); +} + +TEST(MathFunctions, inc_beta_ddb_regressions) { + expect_ddb(1.145, 6786, 1.8e-4, 5.8439109664263782e-5, 1e-10); + expect_ddb(446, 5, 0.76, 1.2763067943072555e-46, 1e-10); + expect_ddb(2500, 2, 0.999, 0.24850317856904407, 1e-10); + expect_ddb(1, 100, 0.01, 0.0036787479630394138, 1e-10); + // integer b: the value series terminates but the derivative does not + expect_ddb(2, 3, 0.25, 0.10692153005229139, 1e-10); + expect_ddb(10, 1, 0.5, 0.00226574342615196, 1e-10); + expect_ddb(600, 600, 0.5, 0.01152047147304755, 1e-10); + expect_ddb(0.5, 0.5, 0.5, 0.58312180806163756, 1e-10); + expect_ddb(20, 20, 0.5, 0.063740468830439859, 1e-10); } diff --git a/test/unit/math/rev/fun/inc_beta_test.cpp b/test/unit/math/rev/fun/inc_beta_test.cpp new file mode 100644 index 00000000000..9b2aefb1272 --- /dev/null +++ b/test/unit/math/rev/fun/inc_beta_test.cpp @@ -0,0 +1,43 @@ +#include +#include +#include +#include +#include + +namespace { + +// Shape gradients of inc_beta against a high-precision reference (mpmath, +// 80 digits), at points where a finite-difference check is not informative. +void expect_inc_beta_grads(double a, double b, double z, double dda, double ddb, + double rtol) { + using stan::math::var; + var av = a; + var bv = b; + var zv = z; + var f = stan::math::inc_beta(av, bv, zv); + std::vector x{av, bv}; + std::vector g; + f.grad(x, g); + EXPECT_TRUE(std::isfinite(g[0])) << "a=" << a << " b=" << b << " z=" << z; + EXPECT_TRUE(std::isfinite(g[1])) << "a=" << a << " b=" << b << " z=" << z; + EXPECT_NEAR(dda, g[0], rtol * std::fabs(dda)) + << "a=" << a << " b=" << b << " z=" << z; + EXPECT_NEAR(ddb, g[1], rtol * std::fabs(ddb)) + << "a=" << a << " b=" << b << " z=" << z; +} + +} // namespace + +TEST_F(AgradRev, inc_beta_shape_gradients) { + expect_inc_beta_grads(2, 3, 0.25, -0.22805360232434637, 0.10692153005229139, + 1e-12); + // a + b above the beta(a, b) underflow limit + expect_inc_beta_grads(600, 600, 0.5, -0.01152047147304755, + 0.01152047147304755, 1e-12); + // z^a (1 - z)^b underflows while I is 1 + expect_inc_beta_grads(352, 590, 0.758, -3.3772717797123242e-138, + 4.5390666074606649e-138, 1e-12); + // integer b + expect_inc_beta_grads(10, 1, 0.5, -0.00067690154351557159, + 0.00226574342615196, 1e-12); +} diff --git a/test/unit/math/rev/prob/beta_cdf_log_test.cpp b/test/unit/math/rev/prob/beta_cdf_log_test.cpp new file mode 100644 index 00000000000..f4ab7ed85c3 --- /dev/null +++ b/test/unit/math/rev/prob/beta_cdf_log_test.cpp @@ -0,0 +1,44 @@ +#include +#include +#include +#include +#include + +namespace { + +// Shape gradients of beta_lcdf against a high-precision reference: the +// reference is (d/da I) / I and (d/db I) / I from mpmath at 80 digits. +void expect_beta_lcdf_grads(double y, double alpha, double beta, double dalpha, + double dbeta, double rtol) { + using stan::math::var; + var a = alpha; + var b = beta; + var lp = stan::math::beta_lcdf(y, a, b); + std::vector x{a, b}; + std::vector g; + lp.grad(x, g); + EXPECT_TRUE(std::isfinite(g[0])) + << "y=" << y << " a=" << alpha << " b=" << beta; + EXPECT_TRUE(std::isfinite(g[1])) + << "y=" << y << " a=" << alpha << " b=" << beta; + EXPECT_NEAR(dalpha, g[0], rtol * std::fabs(dalpha)) + << "y=" << y << " a=" << alpha << " b=" << beta; + EXPECT_NEAR(dbeta, g[1], rtol * std::fabs(dbeta)) + << "y=" << y << " a=" << alpha << " b=" << beta; +} + +} // namespace + +TEST_F(AgradRev, ProbDistributionsBeta_lcdf_shape_gradients) { + expect_beta_lcdf_grads(0.25, 2, 3, -0.87136898798556226, 0.40853599542368053, + 1e-12); + // large shapes: beta(alpha, beta) underflows + expect_beta_lcdf_grads(0.5, 600, 600, -0.0230409429460951, 0.0230409429460951, + 1e-12); + // y far above the mean: the cdf is 1 and the gradients are 1e-138 + expect_beta_lcdf_grads(0.758, 352, 590, -3.3772717797123241e-138, + 4.5390666074606647e-138, 1e-12); + // integer beta + expect_beta_lcdf_grads(0.5, 10, 1, -0.69314718055994531, 2.3201212683796071, + 1e-12); +} diff --git a/test/unit/math/rev/prob/beta_cdf_test.cpp b/test/unit/math/rev/prob/beta_cdf_test.cpp new file mode 100644 index 00000000000..ad01be2862f --- /dev/null +++ b/test/unit/math/rev/prob/beta_cdf_test.cpp @@ -0,0 +1,38 @@ +#include +#include +#include +#include +#include + +TEST_F(AgradRev, ProbDistributionsBeta_cdf_shape_gradients) { + // For one observation the shape partials of beta_cdf are d/da I and + // d/db I. Reference: mpmath, 80 digits. The point has b >> a with a + // small y. + using stan::math::var; + var a = 1.145; + var b = 6786; + var p = stan::math::beta_cdf(1.8e-4, a, b); + std::vector x{a, b}; + std::vector g; + p.grad(x, g); + EXPECT_NEAR(0.64944133941212122, p.val(), 1e-12); + EXPECT_NEAR(-0.38621356822721375, g[0], 1e-10); + EXPECT_NEAR(5.8439109664263782e-5, g[1], 1e-10 * 5.8e-5); +} + +TEST_F(AgradRev, ProbDistributionsBeta_cdf_mixed_shape_types) { + // one autodiff shape with the other constant must give the same partial + using stan::math::var; + { + var a = 1.145; + var p = stan::math::beta_cdf(1.8e-4, a, 6786.0); + p.grad(); + EXPECT_NEAR(-0.38621356822721375, a.adj(), 1e-10); + } + { + var b = 6786; + var p = stan::math::beta_cdf(1.8e-4, 1.145, b); + p.grad(); + EXPECT_NEAR(5.8439109664263782e-5, b.adj(), 1e-10 * 5.8e-5); + } +} From d2eeacd67664ad7e59c7711997d5551b8a9b516a Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Wed, 21 Oct 2026 17:59:16 -0400 Subject: [PATCH 05/10] Reverse-mode grad_reg_inc_beta with a Hessian-vector product reverse pass Add stan/math/rev/fun/grad_reg_inc_beta.hpp. For var arguments the values are computed in double by the prim overload, and the reverse pass computes the a and b adjoints directly as H v (H the Hessian of I_z(a, b) in (a, b)) from the power series or continued fraction, instead of taping through the loop. The z adjoint is closed form. One-sided requests compute only H_aa or H_bb unless the cross term is needed. Also: - simplify the prim grad_reg_inc_beta kernels: one Lentz step for even and odd terms, log-derivatives carried through the steps, the sign folded into the prefactor - rewrite rev and fwd inc_beta with make_callback_var / if constexpr; fvar and var mixes are not supported - move the beta, beta_proportion, neg_binomial and student_t cdfs to the GRAD_RIB interfaces, computing digamma inside the loops instead of with VectorBuilder - fix neg_binomial_cdf (missing return, P / P_i, size_zero, &&) - vectorized inc_beta_dda and inc_beta_ddz, with tests - remove the TensorFlow Probability license file --- licenses/tensorflow-probability-license.txt | 206 ----- stan/math/fwd/fun/inc_beta.hpp | 87 +- stan/math/prim/fun/grad_reg_inc_beta.hpp | 583 +++++++------ stan/math/prim/fun/inc_beta_dda.hpp | 68 +- stan/math/prim/fun/inc_beta_ddb.hpp | 23 +- stan/math/prim/fun/inc_beta_ddz.hpp | 46 +- stan/math/prim/prob/beta_cdf.hpp | 73 +- stan/math/prim/prob/beta_lccdf.hpp | 65 +- stan/math/prim/prob/beta_lcdf.hpp | 67 +- stan/math/prim/prob/beta_proportion_lccdf.hpp | 48 +- stan/math/prim/prob/beta_proportion_lcdf.hpp | 47 +- stan/math/prim/prob/neg_binomial_2_cdf.hpp | 41 +- stan/math/prim/prob/neg_binomial_cdf.hpp | 161 ++-- stan/math/prim/prob/neg_binomial_lccdf.hpp | 52 +- stan/math/prim/prob/neg_binomial_lcdf.hpp | 36 +- stan/math/prim/prob/student_t_cdf.hpp | 69 +- stan/math/prim/prob/student_t_lccdf.hpp | 71 +- stan/math/prim/prob/student_t_lcdf.hpp | 71 +- stan/math/rev/fun.hpp | 1 + stan/math/rev/fun/grad_reg_inc_beta.hpp | 774 ++++++++++++++++++ stan/math/rev/fun/inc_beta.hpp | 77 +- .../math/mix/fun/grad_reg_inc_beta_test.cpp | 206 ++--- .../math/prim/fun/grad_reg_inc_beta_test.cpp | 16 + test/unit/math/prim/fun/inc_beta_dda_test.cpp | 44 + test/unit/math/prim/fun/inc_beta_ddz_test.cpp | 37 + 25 files changed, 1777 insertions(+), 1192 deletions(-) delete mode 100644 licenses/tensorflow-probability-license.txt create mode 100644 stan/math/rev/fun/grad_reg_inc_beta.hpp diff --git a/licenses/tensorflow-probability-license.txt b/licenses/tensorflow-probability-license.txt deleted file mode 100644 index b92e39eaa05..00000000000 --- a/licenses/tensorflow-probability-license.txt +++ /dev/null @@ -1,206 +0,0 @@ -License of the code that stan/math/prim/fun/grad_reg_inc_beta.hpp is derived -from: TensorFlow Probability, tensorflow_probability/python/math/special.py, -Copyright 2022 The TensorFlow Probability Authors. -https://github.com/tensorflow/probability - - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/stan/math/fwd/fun/inc_beta.hpp b/stan/math/fwd/fun/inc_beta.hpp index d0b89e35e06..1c928880ac0 100644 --- a/stan/math/fwd/fun/inc_beta.hpp +++ b/stan/math/fwd/fun/inc_beta.hpp @@ -3,54 +3,59 @@ #include #include -#include #include -#include -#include -#include #include #include -#include +#include +#include namespace stan { namespace math { -template -inline fvar inc_beta(const fvar& a, const fvar& b, const fvar& x) { - T d_a; - T d_b; - const T beta_ab = beta(a.val_, b.val_); - grad_reg_inc_beta(d_a, d_b, a.val_, b.val_, x.val_, digamma(a.val_), - digamma(b.val_), digamma(a.val_ + b.val_), beta_ab); - T d_x = pow((1 - x.val_), b.val_ - 1) * pow(x.val_, a.val_ - 1) / beta_ab; - return fvar(inc_beta(a.val_, b.val_, x.val_), - a.d_ * d_a + b.d_ * d_b + x.d_ * d_x); -} - -template -inline fvar inc_beta(double a, const fvar& b, const fvar& x) { - return inc_beta(fvar(a), b, x); -} -template -inline fvar inc_beta(const fvar& a, double b, const fvar& x) { - return inc_beta(a, fvar(b), x); -} -template -inline fvar inc_beta(const fvar& a, const fvar& b, double x) { - return inc_beta(a, b, fvar(x)); -} - -template -inline fvar inc_beta(double a, double b, const fvar& x) { - return inc_beta(fvar(a), fvar(b), x); -} -template -inline fvar inc_beta(const fvar& a, double b, double x) { - return inc_beta(a, fvar(b), fvar(x)); -} -template -inline fvar inc_beta(double a, const fvar& b, double x) { - return inc_beta(fvar(a), b, fvar(x)); +/** + * The regularized incomplete beta function I_x(a, b) for fvar + * arguments. Only the partials of the fvar arguments are + * computed. Each argument is an fvar or arithmetic; mixing with + * var is not supported. + * + * @tparam T1 type of a + * @tparam T2 type of b + * @tparam T3 type of x + * @param a first shape, positive + * @param b second shape, positive + * @param x upper limit, in [0, 1] + * @return I_x(a, b) + */ +template * = nullptr, + require_all_not_var_t* = nullptr, + require_any_fvar_t* = nullptr> +inline fvar> inc_beta(const T1& a, const T2& b, + const T3& x) { + using T_return = partials_return_t; + const auto a_val = value_of(a); + const auto b_val = value_of(b); + const auto x_val = value_of(x); + T_return d(0); + if constexpr (is_autodiff_v && is_autodiff_v) { + const auto [d_a, d_b] = grad_reg_inc_beta( + a_val, b_val, x_val, digamma(a_val), digamma(b_val), + digamma(a_val + b_val), 0.0); + d += a.d_ * d_a + b.d_ * d_b; + } else if constexpr (is_autodiff_v) { + d += a.d_ + * grad_reg_inc_beta(a_val, b_val, x_val, digamma(a_val), + 0.0, digamma(a_val + b_val), 0.0); + } else if constexpr (is_autodiff_v) { + d += b.d_ + * grad_reg_inc_beta(a_val, b_val, x_val, 0.0, + digamma(b_val), + digamma(a_val + b_val), 0.0); + } + if constexpr (is_autodiff_v) { + d += x.d_ * inc_beta_ddz(a_val, b_val, x_val); + } + return fvar(inc_beta(a_val, b_val, x_val), d); } } // namespace math diff --git a/stan/math/prim/fun/grad_reg_inc_beta.hpp b/stan/math/prim/fun/grad_reg_inc_beta.hpp index 83049f76975..37c3c4426e8 100644 --- a/stan/math/prim/fun/grad_reg_inc_beta.hpp +++ b/stan/math/prim/fun/grad_reg_inc_beta.hpp @@ -1,10 +1,3 @@ -// Original code derived from TensorFlow Probability and is distributed here -// under the Apache License 2.0 (licenses/tensorflow-probability-license.txt) -// Copyright 2022 The TensorFlow Probability Authors. -// tensorflow_probability/python/math/special.py, betainc partials -// Secondary code copyright by its author and is distributed here -// under the BSD-3 license (LICENSE.md) - #ifndef STAN_MATH_PRIM_FUN_GRAD_REG_INC_BETA_HPP #define STAN_MATH_PRIM_FUN_GRAD_REG_INC_BETA_HPP @@ -19,155 +12,148 @@ #include #include #include +#include #include +#include namespace stan { namespace math { +/** + * Selects which partial derivatives of the regularized incomplete beta + * function grad_reg_inc_beta computes: G1 for the + * derivative with respect to a, G2 for the derivative with + * respect to b, BOTH for both. + */ +enum class GRAD_RIB : uint8_t { G1 = 0, G2 = 1, BOTH = 2 }; + namespace internal { +/** + * The partials requested after the symmetry relation + * I_z(a, b) = 1 - I_{1-z}(b, a) swaps the roles of a and b. + */ +inline constexpr GRAD_RIB swap_grad_rib(GRAD_RIB grads_requested) { + return grads_requested == GRAD_RIB::G1 + ? GRAD_RIB::G2 + : grads_requested == GRAD_RIB::G2 ? GRAD_RIB::G1 : GRAD_RIB::BOTH; +} + +/** + * True for the arguments of grad_reg_inc_beta handled by the var overload in + * stan/math/rev/fun/grad_reg_inc_beta.hpp: every argument is a var or + * arithmetic and one of a, b, z or the digamma values is a var. The prim + * overload takes every other combination. + */ +template +inline constexpr bool is_grad_reg_inc_beta_rev_v + = is_var_or_arithmetic::value + && (is_var::value || is_var::value || is_var::value + || is_var::value || is_var::value || is_var::value); + /** * Partial derivatives of the regularized incomplete beta function * I_x(a, b) with respect to a and b, by the power series * * I_x(a, b) = x^a 2F1(a, 1 - b; a + 1; x) / (a B(a, b)), * - * differentiated term by term. Used where at least one of these holds: - * x < a / (a + b), b x <= 1, x <= 0.95 - * x >= a / (a + b), a (1 - x) <= 1, x >= 0.05 - * For x >= a / (a + b) the symmetry relation I_x(a, b) = 1 - I_{1-x}(b, a) - * is applied first. Stops when a term is below eps / a in magnitude, or - * after 600 terms. + * differentiated term by term. Stops when a term is below eps / a in + * magnitude, or after 600 terms. Only the series for the requested partials + * are accumulated. * - * @tparam T scalar type + * @tparam grads_requested which partials to compute + * @tparam flip_sign true when the caller passes (b, a, 1 - x) to evaluate + * its I_x(a, b) through the symmetry relation I_x(a, b) = 1 - I_{1-x}(b, a), + * so the results are negated and returned in the caller's order of a and b + * @tparam T1 type of a + * @tparam T2 type of b + * @tparam T3 type of x + * @tparam T4 type of digamma_a + * @tparam T5 type of digamma_b + * @tparam T6 type of digamma_ab * @param a first shape * @param b second shape * @param x upper limit, in (0, 1) * @param digamma_a digamma(a) * @param digamma_b digamma(b) * @param digamma_ab digamma(a + b) - * @param[out] g1 d/da I_x(a, b) - * @param[out] g2 d/db I_x(a, b) + * @return d/da I_x(a, b) for G1, d/db for G2, the + * pair (d/da, d/db) for BOTH; with flip_sign + * these are negated and the pair is returned as (-d/db, -d/da) */ -template -inline void inc_beta_partials_power_series(T a, T b, T x, T digamma_a, - T digamma_b, T digamma_ab, T& g1, - T& g2) { - using std::exp; - using std::log; - const bool use_symmetry_relation - = value_of_rec(x) >= value_of_rec(a) / value_of_rec(a + b); - if (use_symmetry_relation) { - const T a_orig = a; - a = b; - b = a_orig; - x = 1.0 - x; - const T digamma_a_orig = digamma_a; - digamma_a = digamma_b; - digamma_b = digamma_a_orig; - } - +template +inline auto inc_beta_power_series_kernel(const T1& a, const T2& b, const T3& x, + const T4& digamma_a, + const T5& digamma_b, + const T6& digamma_ab) { + using ret_t = return_type_t; + constexpr bool calc_a = grads_requested != GRAD_RIB::G2; + constexpr bool calc_b = grads_requested != GRAD_RIB::G1; + // db_increment is d(term)/db. For a positive integer b the term is exactly + // 0 at n = b while d(term)/db is not, so the series must also run until + // d(term)/db is small: when d/db is requested, and when b is an autodiff + // type, because autodiff of the d/da series in b needs the same tail. + constexpr bool track_b = calc_b || is_autodiff_v; constexpr int max_iterations = 600; const double tolerance = std::numeric_limits::epsilon() / value_of_rec(a); - // series for 2F1(a, 1 - b; a + 1; x) / a and its a and b derivatives - T product = 1.0; - T series_sum = inv(a); - T product_grad_b = 0.0; - T da = -inv(square(a)); - T db = 0.0; + // series for 2F1(a, 1 - b; a + 1; x) / a and its a and b derivatives. The + // n-th term is product / (a + n), so its a derivative is -term / (a + n) + // and its b derivative is product_grad_b / (a + n). + ret_t product = 1.0; + ret_t series_sum = inv(a); + ret_t da = -series_sum / a; + ret_t db = 0.0; + ret_t product_grad_b = 0.0; for (int n = 1; n <= max_iterations; ++n) { - const T x_div_n = x / n; - const T factor = (n - b) * x_div_n; - const T apn = a + n; - const T new_product = product * factor; - const T term = new_product / apn; - product_grad_b = factor * product_grad_b - product * x_div_n; - const T db_increment = product_grad_b / apn; - da -= new_product / square(apn); - db += db_increment; + const auto x_div_n = x / n; + const auto factor = (n - b) * x_div_n; + const auto apn = a + n; + const auto new_product = product * factor; + const auto term = new_product / apn; + bool db_converged = true; + if constexpr (track_b) { + product_grad_b = factor * product_grad_b - product * x_div_n; + const auto db_increment = product_grad_b / apn; + if constexpr (calc_b) { + db += db_increment; + } + db_converged = fabs(value_of_rec(db_increment)) <= tolerance; + } + if constexpr (calc_a) { + da -= term / apn; + } product = new_product; series_sum += term; - // Differs from the TensorFlow Probability source: the stop test also - // requires the d/db increment to be small. For a positive integer b the - // value series is a polynomial whose term is exactly 0 at n = b, while - // the d/db series continues. - if (fabs(value_of_rec(term)) <= tolerance - && fabs(value_of_rec(db_increment)) <= tolerance) { + if (fabs(value_of_rec(term)) <= tolerance && db_converged) { break; } } - const T normalization = exp(a * log(x) - lbeta(a, b)); - const T grad_a - = normalization * (da + series_sum * (digamma_ab - digamma_a + log(x))); - const T grad_b = normalization * (db + series_sum * (digamma_ab - digamma_b)); - - if (use_symmetry_relation) { - g1 = -grad_b; - g2 = -grad_a; + // the symmetry relation negates the partials + constexpr double sign = flip_sign ? -1.0 : 1.0; + const auto log_x = log(x); + const ret_t normalization = sign * exp(a * log_x - lbeta(a, b)); + if constexpr (grads_requested == GRAD_RIB::G1) { + return ret_t(normalization + * (da + series_sum * (digamma_ab - digamma_a + log_x))); + } else if constexpr (grads_requested == GRAD_RIB::G2) { + return ret_t(normalization * (db + series_sum * (digamma_ab - digamma_b))); } else { - g1 = grad_a; - g2 = grad_b; - } -} - -/** - * One step of the modified Lentz method for the continued fraction of the - * incomplete beta function, carrying the derivatives of the ratios with - * respect to a and b. Thompson and Barnett (1986), appendix. - * - * @tparam T scalar type - * @param d partial numerator d_n - * @param dd_a d/da of d_n - * @param dd_b d/db of d_n - * @param[in,out] C ratio of successive numerators A_n / A_{n-1} - * @param[in,out] D ratio of successive denominators B_{n-1} / B_n - * @param[in,out] h convergent A_n / B_n - * @param[in,out] dC_a d/da of C - * @param[in,out] dC_b d/db of C - * @param[in,out] dD_a d/da of D - * @param[in,out] dD_b d/db of D - * @param[in,out] dh_a d/da of h - * @param[in,out] dh_b d/db of h - * @return the multiplicative update of h at this step - */ -template -inline T inc_beta_lentz_step(const T& d, const T& dd_a, const T& dd_b, T& C, - T& D, T& h, T& dC_a, T& dC_b, T& dD_a, T& dD_b, - T& dh_a, T& dh_b) { - const double small = std::sqrt(std::numeric_limits::min()); - - T C_new = 1.0 + d / C; - if (fabs(value_of_rec(C_new)) < small) { - C_new = small; - } - T D_new = 1.0 + d * D; - if (fabs(value_of_rec(D_new)) < small) { - D_new = small; + const ret_t grad_a + = normalization * (da + series_sum * (digamma_ab - digamma_a + log_x)); + const ret_t grad_b + = normalization * (db + series_sum * (digamma_ab - digamma_b)); + if constexpr (flip_sign) { + // the caller swapped a and b, so its (d/da, d/db) are (grad_b, grad_a) + return std::pair{grad_b, grad_a}; + } else { + return std::pair{grad_a, grad_b}; + } } - D_new = inv(D_new); - const T delta = C_new * D_new; - const T h_new = h * delta; - - const T dC_a_new = (dd_a * C - d * dC_a) / square(C); - const T dC_b_new = (dd_b * C - d * dC_b) / square(C); - const T dD_a_new = -(dd_a * D + d * dD_a) * square(D_new); - const T dD_b_new = -(dd_b * D + d * dD_b) * square(D_new); - const T dh_a_new = dh_a * delta + h * dC_a_new * D_new + h * dD_a_new * C_new; - const T dh_b_new = dh_b * delta + h * dC_b_new * D_new + h * dD_b_new * C_new; - - C = C_new; - D = D_new; - h = h_new; - dC_a = dC_a_new; - dC_b = dC_b_new; - dD_a = dD_a_new; - dD_b = dD_b_new; - dh_a = dh_a_new; - dh_b = dh_b_new; - return delta; } /** @@ -175,131 +161,181 @@ inline T inc_beta_lentz_step(const T& d, const T& dd_a, const T& dd_b, T& C, * I_x(a, b) with respect to a and b, by the continued fraction * https://dlmf.nist.gov/8.17.E22 evaluated with the modified Lentz method, * with the partial numerators https://dlmf.nist.gov/8.17.E23 differentiated - * in a and b and the derivatives carried through the recurrence. Used - * where the power series region does not apply. The continued fraction - * converges rapidly for x < (a - 1) / (a + b - 2); otherwise the symmetry - * relation I_x(a, b) = 1 - I_{1-x}(b, a) is applied first. Stops when the - * multiplicative update is within 3 eps of 1, or after 300 double steps. + * in a and b and the derivatives carried through the recurrence (Thompson + * and Barnett (1986), appendix). Only the requested derivatives are carried. + * Stops when + * the multiplicative update is within 3 eps of 1 and the requested + * derivatives have converged, or after 300 double steps. * - * @tparam T scalar type + * @tparam grads_requested which partials to compute + * @tparam flip_sign true when the caller passes (b, a, 1 - x) to evaluate + * its I_x(a, b) through the symmetry relation I_x(a, b) = 1 - I_{1-x}(b, a), + * so the results are negated and returned in the caller's order of a and b + * @tparam T1 type of a + * @tparam T2 type of b + * @tparam T3 type of x + * @tparam T4 type of digamma_a + * @tparam T5 type of digamma_b + * @tparam T6 type of digamma_ab * @param a first shape * @param b second shape * @param x upper limit, in (0, 1) * @param digamma_a digamma(a) * @param digamma_b digamma(b) * @param digamma_ab digamma(a + b) - * @param[out] g1 d/da I_x(a, b) - * @param[out] g2 d/db I_x(a, b) + * @return d/da I_x(a, b) for G1, d/db for G2, the + * pair (d/da, d/db) for BOTH; with flip_sign + * these are negated and the pair is returned as (-d/db, -d/da) */ -template -inline void inc_beta_partials_continued_fraction(T a, T b, T x, T digamma_a, - T digamma_b, T digamma_ab, - T& g1, T& g2) { - using std::exp; - using std::log; - const double a_d = value_of_rec(a); - const double b_d = value_of_rec(b); - // IEEE division: (a - 1) / 0 is +-inf and the comparison is well defined - const bool use_symmetry_relation - = value_of_rec(x) >= (a_d - 1.0) / (a_d + b_d - 2.0); - if (use_symmetry_relation) { - const T a_orig = a; - a = b; - b = a_orig; - x = 1.0 - x; - const T digamma_a_orig = digamma_a; - digamma_a = digamma_b; - digamma_b = digamma_a_orig; - } - +template +inline auto inc_beta_continued_fraction_kernel(const T1& a, const T2& b, + const T3& x, const T4& digamma_a, + const T5& digamma_b, + const T6& digamma_ab) { + using ret_t = return_type_t; + constexpr bool calc_a = grads_requested != GRAD_RIB::G2; + constexpr bool calc_b = grads_requested != GRAD_RIB::G1; constexpr int max_iterations = 300; const double tolerance = 3.0 * std::numeric_limits::epsilon(); const double small = std::sqrt(std::numeric_limits::min()); - const T apb = a + b; - const T ap1 = a + 1.0; + // Modified Lentz state: C = A_n / A_{n-1} is the ratio of successive + // numerators, D = B_{n-1} / B_n the ratio of successive denominators and + // h = A_n / B_n the convergent, the product of every step's C * D. The + // derivatives are carried as logarithmic derivatives, c = (dC/da) / C, + // e = (dD/da) / D and r = (dh/da) / h for a and likewise for b, so each step + // adds c + e to r. - // initialization and first step of the modified Lentz method - T C = 1.0; - T D = 1.0 - apb * x / ap1; + // first step, with partial numerator d_1 = -(a + b) x / (a + 1) and C held + // at 1: D = 1 / (1 + d_1) and e = -(d d_1) D + const auto ap1 = a + 1.0; + const auto d_1 = -(a + b) * x / ap1; + ret_t C = 1.0; + ret_t D = 1.0 + d_1; if (fabs(value_of_rec(D)) < small) { D = small; } D = inv(D); - T h = D; - const T dD_denom = square(x * apb - ap1); - T dD_a = (1.0 - b) * x / dD_denom; - T dD_b = ap1 * x / dD_denom; - T dC_a = 0.0; - T dC_b = 0.0; - T dh_a = dD_a; - T dh_b = dD_b; + ret_t h = D; + ret_t c_a = 0.0; + ret_t e_a = 0.0; + ret_t r_a = 0.0; + ret_t c_b = 0.0; + ret_t e_b = 0.0; + ret_t r_b = 0.0; + if constexpr (calc_a) { + // d d_1 / da = -(x + d_1) / (a + 1) + e_a = (x + d_1) / ap1 * D; + r_a = e_a; + } + if constexpr (calc_b) { + // d d_1 / db = -x / (a + 1) + e_b = x / ap1 * D; + r_b = e_b; + } - for (int m = 1; m <= max_iterations; ++m) { - const double dh_a_prev = value_of_rec(dh_a); - const double dh_b_prev = value_of_rec(dh_b); - // even partial numerator d_{2m}, https://dlmf.nist.gov/8.17.E23 - { - const T a_plus_2m = a + 2.0 * m; - const T a_plus_2m_minus_one = a_plus_2m - 1.0; - const T denominator = a_plus_2m * a_plus_2m_minus_one; - const T dd_b = m * x / denominator; - const T d = dd_b * (b - m); - const T dd_a = -d * (a_plus_2m + a_plus_2m_minus_one) / denominator; - inc_beta_lentz_step(d, dd_a, dd_b, C, D, h, dC_a, dC_b, dD_a, dD_b, dh_a, - dh_b); + // the stop test compares dh = h r after each odd step + double dh_a_prev = value_of_rec(h) * value_of_rec(r_a); + double dh_b_prev = value_of_rec(h) * value_of_rec(r_b); + for (int n = 2; n <= 2 * max_iterations + 1; ++n) { + // partial numerators https://dlmf.nist.gov/8.17.E23 over the shared + // denominator (a + n - 1)(a + n): + // d_{2m} = m (b - m) x / ((a + 2m - 1)(a + 2m)) + // d_{2m+1} = -(a + m)(a + b + m) x / ((a + 2m)(a + 2m + 1)) + // d log(denominator) / da = (2 (a + n) - 1) / denominator, and for odd n + // the numerator's a derivative is -(2 (a + m) + b) x + const int m = n / 2; + const bool odd = n % 2 == 1; + const auto a_plus_n = a + n; + const auto denominator = a_plus_n * (a_plus_n - 1.0); + ret_t dd_b; + ret_t d; + if (odd) { + dd_b = -(a + m) * x / denominator; + d = dd_b * (a + b + m); + } else { + dd_b = m * x / denominator; + d = dd_b * (b - m); } - // odd partial numerator d_{2m+1} - T delta; - { - const T a_plus_m = a + m; - const T a_plus_2m = a_plus_m + m; - const T a_plus_2m_plus_one = a_plus_2m + 1.0; - const T a_plus_b_plus_m = a_plus_m + b; - const T denominator = a_plus_2m * a_plus_2m_plus_one; - const T dd_b = -a_plus_m * x / denominator; - const T d = dd_b * a_plus_b_plus_m; - const T dd_a = -d * ((a_plus_2m + a_plus_2m_plus_one) / denominator) - - x * (2.0 * a_plus_m + b) / denominator; - delta = inc_beta_lentz_step(d, dd_a, dd_b, C, D, h, dC_a, dC_b, dD_a, - dD_b, dh_a, dh_b); + // one step of the modified Lentz method + auto C_new = 1.0 + d / C; + // C C_new = C + d, which does not wait on the division for C_new + ret_t C_C_new = C + d; + if (fabs(value_of_rec(C_new)) < small) { + C_new = small; + C_C_new = C * small; } - // Differs from the TensorFlow Probability source: the stop test also - // requires the two derivatives to have converged. For a positive - // integer b the partial numerator d_{2m} is exactly 0 at m = b; the - // convergent is then exact and delta is 1, but the derivatives still - // need the tail of the fraction. - const double dh_a_now = value_of_rec(dh_a); - const double dh_b_now = value_of_rec(dh_b); - if (fabs(value_of_rec(delta) - 1.0) < tolerance - && fabs(dh_a_now - dh_a_prev) <= tolerance * fabs(dh_a_now) - && fabs(dh_b_now - dh_b_prev) <= tolerance * fabs(dh_b_now)) { - break; + auto D_new = 1.0 + d * D; + if (fabs(value_of_rec(D_new)) < small) { + D_new = small; + } + D_new = inv(D_new); + // c_new = (dd - d c) / (C C_new) and e_new = -(dd + d e) D D_new + const auto D_D_new = D * D_new; + if constexpr (calc_a) { + auto dd_a = -d * (2.0 * a_plus_n - 1.0) / denominator; + if (odd) { + dd_a -= x * (2.0 * (a + m) + b) / denominator; + } + c_a = (dd_a - d * c_a) / C_C_new; + e_a = -(dd_a + d * e_a) * D_D_new; + r_a += c_a + e_a; + } + if constexpr (calc_b) { + c_b = (dd_b - d * c_b) / C_C_new; + e_b = -(dd_b + d * e_b) * D_D_new; + r_b += c_b + e_b; + } + C = C_new; + D = D_new; + const auto delta = C_new * D_new; + h *= delta; + if (odd) { + const double h_d = value_of_rec(h); + const double dh_a_now = h_d * value_of_rec(r_a); + const double dh_b_now = h_d * value_of_rec(r_b); + if (fabs(value_of_rec(delta) - 1.0) < tolerance + && (!calc_a + || fabs(dh_a_now - dh_a_prev) <= tolerance * fabs(dh_a_now)) + && (!calc_b + || fabs(dh_b_now - dh_b_prev) <= tolerance * fabs(dh_b_now))) { + break; + } + dh_a_prev = dh_a_now; + dh_b_prev = dh_b_now; } } - const T normalization - = exp(a * log(x) + b * log1p(-x) - log(a) - lbeta(a, b)); - const T grad_a - = normalization * (dh_a + h * (log(x) - inv(a) + digamma_ab - digamma_a)); - const T grad_b - = normalization * (dh_b + h * (log1p(-x) + digamma_ab - digamma_b)); - - if (use_symmetry_relation) { - g1 = -grad_b; - g2 = -grad_a; + // the symmetry relation negates the partials + constexpr double sign = flip_sign ? -1.0 : 1.0; + const auto log_x = log(x); + const auto log1m_x = log1p(-x); + // h times the normalization is I_x(a, b), shared by both partials + const auto prefactor + = sign * h * exp(a * log_x + b * log1m_x - log(a) - lbeta(a, b)); + if constexpr (grads_requested == GRAD_RIB::G1) { + return ret_t(prefactor * (r_a + log_x - inv(a) + digamma_ab - digamma_a)); + } else if constexpr (grads_requested == GRAD_RIB::G2) { + return ret_t(prefactor * (r_b + log1m_x + digamma_ab - digamma_b)); } else { - g1 = grad_a; - g2 = grad_b; + const ret_t grad_a + = prefactor * (r_a + log_x - inv(a) + digamma_ab - digamma_a); + const ret_t grad_b = prefactor * (r_b + log1m_x + digamma_ab - digamma_b); + if constexpr (flip_sign) { + // the caller swapped a and b, so its (d/da, d/db) are (grad_b, grad_a) + return std::pair{grad_b, grad_a}; + } else { + return std::pair{grad_a, grad_b}; + } } } } // namespace internal /** - * Computes the gradients of the regularized incomplete beta - * function. Specifically, this function computes gradients of - * inc_beta(a, b, z), with respect to the arguments + * Computes the requested gradients of the regularized incomplete beta + * function inc_beta(a, b, z) with respect to the arguments * a and b. * * The algorithm is the one of Boik and Robinson-Cox (1998), "Derivatives @@ -308,10 +344,96 @@ inline void inc_beta_partials_continued_fraction(T a, T b, T x, T digamma_a, * the endpoints and a continued fraction elsewhere, both differentiated * term by term, with the prefactor evaluated in log space and the * symmetry relation I_z(a, b) = 1 - I_{1-z}(b, a) applied where the - * expansion converges faster on the other side. + * expansion converges faster on the other side. Only the requested + * derivatives are carried through the expansion. * * The gradients are 0 at z = 0 and z = 1. * + * @tparam grads_requested G1 for the derivative with respect to + * a, G2 for the derivative with respect to b, + * BOTH for both + * @tparam T1 type of a + * @tparam T2 type of b + * @tparam T3 type of z + * @tparam T4 type of digammaA + * @tparam T5 type of digammaB + * @tparam T6 type of digammaSum + * @tparam T7 type of betaAB + * @param a a + * @param b b + * @param z z + * @param digammaA the value of digamma(a); not used for + * G2 + * @param digammaB the value of digamma(b); not used for + * G1 + * @param digammaSum the value of digamma(a + b) + * @param betaAB accepted for interface compatibility and not used, the + * normalization is evaluated in log space + * @return d/da I_z(a, b) for G1, d/db for G2, the + * pair (d/da, d/db) for BOTH + */ +template >>* = nullptr> +inline auto grad_reg_inc_beta(const T1& a, const T2& b, const T3& z, + const T4& digammaA, const T5& digammaB, + const T6& digammaSum, const T7& /* betaAB */) { + // T7 is left out so that every branch returns the same type + using ret_t = return_type_t; + // the partials requested after the symmetry relation swaps a and b + constexpr GRAD_RIB swapped_grads = internal::swap_grad_rib(grads_requested); + const double z_d = value_of_rec(z); + if (z_d == 0.0 || z_d == 1.0) { + if constexpr (grads_requested == GRAD_RIB::BOTH) { + return std::pair{0.0, 0.0}; + } else { + return ret_t{0.0}; + } + } + + const double a_d = value_of_rec(a); + const double b_d = value_of_rec(b); + const double mean = a_d / (a_d + b_d); + const bool use_power_series + = (z_d < mean && b_d * z_d <= 1.0 && z_d <= 0.95) + || (z_d >= mean && a_d * (1.0 - z_d) <= 1.0 && z_d >= 0.05); + if (use_power_series) { + // For z >= a / (a + b) the series is evaluated for I_{1-z}(b, a) through + // the symmetry relation; its (d/da, d/db) are the negated (d/db, d/da) + // of I_z(a, b). + if (z_d >= mean) { + return internal::inc_beta_power_series_kernel( + b, a, 1.0 - z, digammaB, digammaA, digammaSum); + } else { + return internal::inc_beta_power_series_kernel( + a, b, z, digammaA, digammaB, digammaSum); + } + } else { + // The continued fraction converges rapidly for z < (a - 1) / (a + b - 2); + // otherwise it is evaluated for I_{1-z}(b, a) through the symmetry + // relation. IEEE division: (a - 1) / 0 is +-inf and the comparison is + // well defined. + if (z_d >= (a_d - 1.0) / (a_d + b_d - 2.0)) { + return internal::inc_beta_continued_fraction_kernel( + b, a, 1.0 - z, digammaB, digammaA, digammaSum); + } else { + return internal::inc_beta_continued_fraction_kernel( + a, b, z, digammaA, digammaB, digammaSum); + } + } +} + +/** + * Computes the gradients of the regularized incomplete beta + * function. Specifically, this function computes gradients of + * inc_beta(a, b, z), with respect to the arguments + * a and b. + * + * Writes both partials; see the overload templated on + * GRAD_RIB to compute only one. + * * @tparam T type of arguments * @param[out] g1 partial derivative of inc_beta(a, b, z) * with respect to a @@ -331,27 +453,10 @@ template inline void grad_reg_inc_beta(T& g1, T& g2, const T& a, const T& b, const T& z, const T& digammaA, const T& digammaB, const T& digammaSum, const T& betaAB) { - const double a_d = value_of_rec(a); - const double b_d = value_of_rec(b); - const double z_d = value_of_rec(z); - - if (z_d == 0.0 || z_d == 1.0) { - g1 = 0.0; - g2 = 0.0; - return; - } - - const double mean = a_d / (a_d + b_d); - const bool use_power_series - = (z_d < mean && b_d * z_d <= 1.0 && z_d <= 0.95) - || (z_d >= mean && a_d * (1.0 - z_d) <= 1.0 && z_d >= 0.05); - if (use_power_series) { - internal::inc_beta_partials_power_series(a, b, z, digammaA, digammaB, - digammaSum, g1, g2); - } else { - internal::inc_beta_partials_continued_fraction(a, b, z, digammaA, digammaB, - digammaSum, g1, g2); - } + auto grads = grad_reg_inc_beta(a, b, z, digammaA, digammaB, + digammaSum, betaAB); + g1 = grads.first; + g2 = grads.second; } } // namespace math diff --git a/stan/math/prim/fun/inc_beta_dda.hpp b/stan/math/prim/fun/inc_beta_dda.hpp index f45ab7cc5ac..f5c7d4efad1 100644 --- a/stan/math/prim/fun/inc_beta_dda.hpp +++ b/stan/math/prim/fun/inc_beta_dda.hpp @@ -2,8 +2,11 @@ #define STAN_MATH_PRIM_FUN_INC_BETA_DDA_HPP #include -#include +#include +#include #include +#include +#include namespace stan { namespace math { @@ -12,10 +15,14 @@ namespace math { * Returns the partial derivative of the regularized * incomplete beta function, I_{z}(a, b) with respect to a. * - * Evaluated by grad_reg_inc_beta, which computes both shape - * derivatives; call that function directly when both are needed. + * Evaluated by grad_reg_inc_beta, computing only the derivative + * with respect to a. * - * @tparam T scalar types of arguments + * @tparam T1 type of a + * @tparam T2 type of b + * @tparam T3 type of z + * @tparam T4 type of digamma_a + * @tparam T5 type of digamma_ab * @param a first argument * @param b second argument * @param z upper bound of the integral @@ -27,12 +34,53 @@ namespace math { * @pre b >= 0 * @pre 0 <= z <= 1 */ -template -inline T inc_beta_dda(T a, T b, T z, T digamma_a, T digamma_ab) { - T g1 = 0; - T g2 = 0; - grad_reg_inc_beta(g1, g2, a, b, z, digamma_a, digamma(b), digamma_ab, T(0)); - return g1; +template * = nullptr> +inline auto inc_beta_dda(const T1& a, const T2& b, const T3& z, + const T4& digamma_a, const T5& digamma_ab) { + return grad_reg_inc_beta(a, b, z, digamma_a, 0.0, digamma_ab, + 0.0); +} + +/** + * Enables the vectorized application of the inc_beta_dda function, when + * any arguments are vectors. Scalar arguments are broadcast against the + * vectors, which must all have the same size. + * + * @tparam T1 type of a + * @tparam T2 type of b + * @tparam T3 type of z + * @tparam T4 type of digamma_a + * @tparam T5 type of digamma_ab + * @param a first argument + * @param b second argument + * @param z upper bound of the integral + * @param digamma_a value of digamma(a) + * @param digamma_ab value of digamma(a + b) + * @return column array of the partial derivatives of the incomplete beta + * with respect to a + * @throw std::invalid_argument if the vector arguments differ in size + */ +template * = nullptr, + require_all_not_var_matrix_t* = nullptr> +inline auto inc_beta_dda(const T1& a, const T2& b, const T3& z, + const T4& digamma_a, const T5& digamma_ab) { + check_consistent_sizes("inc_beta_dda", "a", a, "b", b, "z", z, "digamma_a", + digamma_a, "digamma_ab", digamma_ab); + using ret_t = return_type_t; + scalar_seq_view a_vec(a); + scalar_seq_view b_vec(b); + scalar_seq_view z_vec(z); + scalar_seq_view digamma_a_vec(digamma_a); + scalar_seq_view digamma_ab_vec(digamma_ab); + const size_t N = max_size(a, b, z, digamma_a, digamma_ab); + Eigen::Array result(N); + for (size_t i = 0; i < N; ++i) { + result.coeffRef(i) = inc_beta_dda(a_vec[i], b_vec[i], z_vec[i], + digamma_a_vec[i], digamma_ab_vec[i]); + } + return result; } } // namespace math diff --git a/stan/math/prim/fun/inc_beta_ddb.hpp b/stan/math/prim/fun/inc_beta_ddb.hpp index a4d440da80a..ee42baa91e3 100644 --- a/stan/math/prim/fun/inc_beta_ddb.hpp +++ b/stan/math/prim/fun/inc_beta_ddb.hpp @@ -2,7 +2,6 @@ #define STAN_MATH_PRIM_FUN_INC_BETA_DDB_HPP #include -#include #include namespace stan { @@ -12,10 +11,14 @@ namespace math { * Returns the partial derivative of the regularized * incomplete beta function, I_{z}(a, b) with respect to b. * - * Evaluated by grad_reg_inc_beta, which computes both shape - * derivatives; call that function directly when both are needed. + * Evaluated by grad_reg_inc_beta, computing only the derivative + * with respect to b. * - * @tparam T scalar types of arguments + * @tparam T1 type of a + * @tparam T2 type of b + * @tparam T3 type of z + * @tparam T4 type of digamma_b + * @tparam T5 type of digamma_ab * @param a first argument * @param b second argument * @param z upper bound of the integral @@ -27,12 +30,12 @@ namespace math { * @pre b >= 0 * @pre 0 <= z <= 1 */ -template -inline T inc_beta_ddb(T a, T b, T z, T digamma_b, T digamma_ab) { - T g1 = 0; - T g2 = 0; - grad_reg_inc_beta(g1, g2, a, b, z, digamma(a), digamma_b, digamma_ab, T(0)); - return g2; +template * = nullptr> +inline auto inc_beta_ddb(const T1& a, const T2& b, const T3& z, + const T4& digamma_b, const T5& digamma_ab) { + return grad_reg_inc_beta(a, b, z, 0.0, digamma_b, digamma_ab, + 0.0); } } // namespace math diff --git a/stan/math/prim/fun/inc_beta_ddz.hpp b/stan/math/prim/fun/inc_beta_ddz.hpp index 4543df47880..2700b0d34b5 100644 --- a/stan/math/prim/fun/inc_beta_ddz.hpp +++ b/stan/math/prim/fun/inc_beta_ddz.hpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -26,20 +27,55 @@ namespace math { * @pre b > 0 * @pre 0 < z <= 1 */ -template -inline T inc_beta_ddz(T a, T b, T z) { - using std::exp; - using std::log; +template * = nullptr> +inline auto inc_beta_ddz(T1 a, T2 b, T3 z) { return exp((b - 1) * log1m(z) + (a - 1) * log(z) + lgamma(a + b) - lgamma(a) - lgamma(b)); } -template <> +/** + * Returns the partial derivative of the regularized + * incomplete beta function, I_{z}(a, b) with respect to z. + * + * Overload for double arguments, evaluated by Boost. + * + * @param a first argument + * @param b second argument + * @param z upper bound of the integral + * @return partial derivative of the incomplete beta with respect to z + */ inline double inc_beta_ddz(double a, double b, double z) { using boost::math::ibeta_derivative; return ibeta_derivative(a, b, z); } +/** + * Enables the vectorized application of the inc_beta_ddz function, when + * any arguments are containers. + * + * @tparam T1 type of first input + * @tparam T2 type of second input + * @tparam T3 type of third input + * @param a First input + * @param b Second input + * @param z Third input + * @return Partial derivative of the incomplete beta function with respect + * to z, applied to the three inputs. + */ +template * = nullptr, + require_all_not_var_matrix_t* = nullptr> +inline auto inc_beta_ddz(T1&& a, T2&& b, T3&& z) { + return apply_scalar_ternary( + [](auto&& d, auto&& e, auto&& f) { + return inc_beta_ddz(std::forward(d), + std::forward(e), + std::forward(f)); + }, + std::forward(a), std::forward(b), std::forward(z)); +} + } // namespace math } // namespace stan #endif diff --git a/stan/math/prim/prob/beta_cdf.hpp b/stan/math/prim/prob/beta_cdf.hpp index 77abbe6b3a6..2e51fbdf842 100644 --- a/stan/math/prim/prob/beta_cdf.hpp +++ b/stan/math/prim/prob/beta_cdf.hpp @@ -58,9 +58,6 @@ inline return_type_t beta_cdf( scalar_seq_view y_vec(y_ref); scalar_seq_view alpha_vec(alpha_ref); scalar_seq_view 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 @@ -71,31 +68,10 @@ inline return_type_t beta_cdf( } } - // grad_reg_inc_beta computes both shape derivatives in one pass and needs - // all three digamma values, so they are built when either shape is - // autodiff. - constexpr bool need_shape_grads - = is_any_autodiff_v; - VectorBuilder - digamma_alpha(size_alpha); - VectorBuilder digamma_beta( - size_beta); - VectorBuilder - digamma_sum(size_alpha_beta); - if constexpr (need_shape_grads) { - for (size_t n = 0; n < size_alpha; n++) { - digamma_alpha[n] = digamma(alpha_vec.val(n)); - } - for (size_t n = 0; n < size_beta; n++) { - digamma_beta[n] = digamma(beta_vec.val(n)); - } - 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; + constexpr bool calc_beta = is_autodiff_v; 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 @@ -103,29 +79,36 @@ inline return_type_t beta_cdf( 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::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_any_autodiff_v) { + const auto inv_Pn = inv(Pn); + if constexpr (is_autodiff_v) { + partials<0>(ops_partials)[n] + += inc_beta_ddz(alpha_dbl, beta_dbl, y_dbl) * inv_Pn; + } - if constexpr (is_autodiff_v) { - partials<0>(ops_partials)[n] - += inc_beta_ddz(alpha_dbl, beta_dbl, y_dbl) * inv_Pn; - } - - if constexpr (need_shape_grads) { - 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], T_partials_return(0)); - if constexpr (is_autodiff_v) { + if constexpr (calc_alpha && calc_beta) { + const auto [g1, g2] = grad_reg_inc_beta( + 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) { partials<2>(ops_partials)[n] += g2 * inv_Pn; + } else if constexpr (calc_alpha) { + partials<1>(ops_partials)[n] + += grad_reg_inc_beta( + 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( + alpha_dbl, beta_dbl, y_dbl, 0.0, digamma(beta_dbl), + digamma(alpha_dbl + beta_dbl), 0.0) + * inv_Pn; } } } diff --git a/stan/math/prim/prob/beta_lccdf.hpp b/stan/math/prim/prob/beta_lccdf.hpp index 66232243893..ed4234a6f5a 100644 --- a/stan/math/prim/prob/beta_lccdf.hpp +++ b/stan/math/prim/prob/beta_lccdf.hpp @@ -67,40 +67,17 @@ inline return_type_t beta_lccdf( scalar_seq_view y_vec(y_ref); scalar_seq_view alpha_vec(alpha_ref); scalar_seq_view 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, - T_partials_return, T_scale_succ> - digamma_alpha(size_alpha); - VectorBuilder, - T_partials_return, T_scale_fail> - digamma_beta(size_beta); - VectorBuilder, - T_partials_return, T_scale_succ, T_scale_fail> - digamma_sum(size_alpha_beta); - - if constexpr (is_any_autodiff_v) { - 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; + constexpr bool calc_beta = is_autodiff_v; 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::value ? 0 : inv(Pn); ccdf_log += log(Pn); @@ -111,18 +88,24 @@ inline return_type_t beta_lccdf( / betafunc_dbl; } - T_partials_return g1 = 0; - T_partials_return g2 = 0; - - if constexpr (is_any_autodiff_v) { - 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) { + if constexpr (calc_alpha && calc_beta) { + const auto [g1, g2] = grad_reg_inc_beta( + 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) { partials<2>(ops_partials)[n] -= g2 * inv_Pn; + } else if constexpr (calc_alpha) { + partials<1>(ops_partials)[n] + -= grad_reg_inc_beta(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(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); diff --git a/stan/math/prim/prob/beta_lcdf.hpp b/stan/math/prim/prob/beta_lcdf.hpp index c2a514a766c..22517788ed4 100644 --- a/stan/math/prim/prob/beta_lcdf.hpp +++ b/stan/math/prim/prob/beta_lcdf.hpp @@ -58,40 +58,18 @@ inline return_type_t beta_lcdf( scalar_seq_view alpha_vec(alpha_ref); scalar_seq_view 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; - VectorBuilder digamma_alpha( - size_alpha); - VectorBuilder digamma_beta( - size_beta); - VectorBuilder - digamma_sum(size_alpha_beta); - - if constexpr (is_any_autodiff_v) { - 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; + constexpr bool calc_beta = is_autodiff_v; 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 ? inv(Pn) : 0; cdf_log += log(Pn); @@ -102,17 +80,24 @@ inline return_type_t beta_lcdf( / betafunc_dbl; } - if constexpr (is_any_autodiff_v) { - 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) { - partials<1>(ops_partials)[n] += g1 * inv_Pn; - } - if constexpr (is_any_autodiff_v) { - partials<2>(ops_partials)[n] += g2 * inv_Pn; - } + if constexpr (calc_alpha && calc_beta) { + const auto [g1, g2] = grad_reg_inc_beta( + 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(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(alpha_dbl, beta_dbl, y_dbl, 0.0, + digamma(beta_dbl), + digamma(alpha_dbl + beta_dbl), 0.0) + * inv_Pn; } } diff --git a/stan/math/prim/prob/beta_proportion_lccdf.hpp b/stan/math/prim/prob/beta_proportion_lccdf.hpp index c615d23b0df..04e606fd3c8 100644 --- a/stan/math/prim/prob/beta_proportion_lccdf.hpp +++ b/stan/math/prim/prob/beta_proportion_lccdf.hpp @@ -71,32 +71,8 @@ inline return_type_t beta_proportion_lccdf( scalar_seq_view y_vec(y_ref); scalar_seq_view mu_vec(mu_ref); scalar_seq_view 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, T_partials_return, T_loc, - T_prec> - digamma_mukappa(size_mu_kappa); - VectorBuilder, T_partials_return, T_loc, - T_prec> - digamma_kappa_mukappa(size_mu_kappa); - VectorBuilder, T_partials_return, T_prec> - digamma_kappa(size_kappa); - - if constexpr (is_any_autodiff_v) { - 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); @@ -118,20 +94,18 @@ inline return_type_t 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) { - 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) { - partials<1>(ops_partials)[n] -= kappa_dbl * (g1 - g2) * inv_Pn; - } - if constexpr (is_autodiff_v) { - partials<2>(ops_partials)[n] - -= (g1 * mu_dbl + g2 * (1 - mu_dbl)) * inv_Pn; + const auto [g1, g2] = grad_reg_inc_beta( + mukappa_dbl, kappa_mukappa_dbl, y_dbl, digamma(mukappa_dbl), + digamma(kappa_mukappa_dbl), digamma(kappa_dbl), 0.0); + if constexpr (is_autodiff_v) { + partials<1>(ops_partials)[n] -= kappa_dbl * (g1 - g2) * inv_Pn; + } + if constexpr (is_autodiff_v) { + partials<2>(ops_partials)[n] + -= (g1 * mu_dbl + g2 * (1 - mu_dbl)) * inv_Pn; + } } } return ops_partials.build(ccdf_log); diff --git a/stan/math/prim/prob/beta_proportion_lcdf.hpp b/stan/math/prim/prob/beta_proportion_lcdf.hpp index 5566b8e5fe2..e3e11ee3534 100644 --- a/stan/math/prim/prob/beta_proportion_lcdf.hpp +++ b/stan/math/prim/prob/beta_proportion_lcdf.hpp @@ -72,31 +72,8 @@ inline return_type_t beta_proportion_lcdf( scalar_seq_view y_vec(y_ref); scalar_seq_view mu_vec(mu_ref); scalar_seq_view 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, T_partials_return, T_loc, - T_prec> - digamma_mukappa(size_mu_kappa); - VectorBuilder, T_partials_return, T_loc, - T_prec> - digamma_kappa_mukappa(size_mu_kappa); - VectorBuilder, T_partials_return, T_prec> - digamma_kappa(size_kappa); - - if constexpr (is_any_autodiff_v) { - 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); @@ -118,20 +95,18 @@ inline return_type_t beta_proportion_lcdf( / 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) { - 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) { - partials<1>(ops_partials)[n] += kappa_dbl * (g1 - g2) * inv_Pn; - } - if constexpr (is_autodiff_v) { - partials<2>(ops_partials)[n] - += (g1 * mu_dbl + g2 * (1 - mu_dbl)) * inv_Pn; + const auto [g1, g2] = grad_reg_inc_beta( + mukappa_dbl, kappa_mukappa_dbl, y_dbl, digamma(mukappa_dbl), + digamma(kappa_mukappa_dbl), digamma(kappa_dbl), 0.0); + if constexpr (is_autodiff_v) { + partials<1>(ops_partials)[n] += kappa_dbl * (g1 - g2) * inv_Pn; + } + if constexpr (is_autodiff_v) { + partials<2>(ops_partials)[n] + += (g1 * mu_dbl + g2 * (1 - mu_dbl)) * inv_Pn; + } } } diff --git a/stan/math/prim/prob/neg_binomial_2_cdf.hpp b/stan/math/prim/prob/neg_binomial_2_cdf.hpp index 47cab11d76d..3890f5149b4 100644 --- a/stan/math/prim/prob/neg_binomial_2_cdf.hpp +++ b/stan/math/prim/prob/neg_binomial_2_cdf.hpp @@ -84,30 +84,27 @@ inline return_type_t neg_binomial_2_cdf( return ops_partials.build(1.0); } - const T_partials_return n_dbl_p1 = n_vec.val(i) + 1; - const T_partials_return mu_dbl = mu_vec.val(i); - const T_partials_return phi_dbl = phi_vec.val(i); - const T_partials_return inv_mu_plus_phi = inv(mu_dbl + phi_dbl); - const T_partials_return p_dbl = phi_dbl * inv_mu_plus_phi; - const T_partials_return d_dbl = square(inv_mu_plus_phi); - const T_partials_return P_i = inc_beta(phi_dbl, n_dbl_p1, p_dbl); - const T_partials_return inc_beta_ddz_i - = is_constant_all::value - ? 0 - : inc_beta_ddz(phi_dbl, n_dbl_p1, p_dbl) * d_dbl / P_i; + const auto n_dbl_p1 = n_vec.val(i) + 1; + const auto mu_dbl = mu_vec.val(i); + const auto phi_dbl = phi_vec.val(i); + const auto inv_mu_plus_phi = inv(mu_dbl + phi_dbl); + const auto p_dbl = phi_dbl * inv_mu_plus_phi; + const auto P_i = inc_beta(phi_dbl, n_dbl_p1, p_dbl); P *= P_i; - - if constexpr (is_autodiff_v) { - partials<0>(ops_partials)[i] -= inc_beta_ddz_i * phi_dbl; - } - - if constexpr (is_autodiff_v) { - partials<1>(ops_partials)[i] - += inc_beta_dda(phi_dbl, n_dbl_p1, p_dbl, digamma_phi_vec[i], - digamma_sum_vec[i]) - / P_i - + inc_beta_ddz_i * mu_dbl; + if constexpr (is_any_autodiff_v) { + const auto d_dbl = square(inv_mu_plus_phi); + const auto inc_beta_ddz_i = inc_beta_ddz(phi_dbl, n_dbl_p1, p_dbl) * d_dbl / P_i; + if constexpr (is_autodiff_v) { + partials<0>(ops_partials)[i] -= inc_beta_ddz_i * phi_dbl; + } + if constexpr (is_autodiff_v) { + partials<1>(ops_partials)[i] + += inc_beta_dda(phi_dbl, n_dbl_p1, p_dbl, digamma_phi_vec[i], + digamma_sum_vec[i]) + / P_i + + inc_beta_ddz_i * mu_dbl; + } } } diff --git a/stan/math/prim/prob/neg_binomial_cdf.hpp b/stan/math/prim/prob/neg_binomial_cdf.hpp index 31163102be1..b4e7c2b436a 100644 --- a/stan/math/prim/prob/neg_binomial_cdf.hpp +++ b/stan/math/prim/prob/neg_binomial_cdf.hpp @@ -3,122 +3,115 @@ #include #include +#include +#include +#include #include #include #include #include #include -#include -#include -#include +#include +#include #include #include -#include +#include #include -#include #include namespace stan { namespace math { -template +/** \ingroup prob_dists + * Returns the CDF of the negative binomial distribution. Given containers of + * matching sizes, returns the product of probabilities. + * + * With p = beta / (1 + beta), the CDF at n is the regularized incomplete beta + * function I_p(alpha, n + 1). + * + * @tparam T_n type of integer parameter + * @tparam T_shape type of shape parameter + * @tparam T_inv_scale type of inverse scale parameter + * @param n failures variable + * @param alpha shape parameter + * @param beta inverse scale parameter + * @return probability or product of probabilities + * @throw std::domain_error if alpha or beta is not positive and finite + * @throw std::invalid_argument if container sizes mismatch + */ +template * = nullptr> inline return_type_t neg_binomial_cdf( - const T_n& n, const T_shape& alpha, const T_inv_scale& beta) { + T_n&& n, T_shape&& alpha, T_inv_scale&& beta) { using T_partials_return = partials_return_t; using T_n_ref = ref_type_t; - using T_alpha_ref = ref_type_t; - using T_beta_ref = ref_type_t; + using T_alpha_ref = ref_type_if_not_constant_t; + using T_beta_ref = ref_type_if_not_constant_t; static constexpr const char* function = "neg_binomial_cdf"; check_consistent_sizes(function, "Failures variable", n, "Shape parameter", alpha, "Inverse scale parameter", beta); - T_n_ref n_ref = n; - T_alpha_ref alpha_ref = alpha; - T_beta_ref beta_ref = beta; - check_positive_finite(function, "Shape parameter", alpha_ref); - check_positive_finite(function, "Inverse scale parameter", beta_ref); + T_n_ref n_ref = std::forward(n); + T_alpha_ref alpha_ref = std::forward(alpha); + T_beta_ref beta_ref = std::forward(beta); - if (size_zero(n, alpha, beta)) { + decltype(auto) n_val = to_ref(as_value_column_array_or_scalar(n_ref)); + decltype(auto) alpha_val = to_ref(as_value_column_array_or_scalar(alpha_ref)); + decltype(auto) beta_val = to_ref(as_value_column_array_or_scalar(beta_ref)); + + check_positive_finite(function, "Shape parameter", alpha_val); + check_positive_finite(function, "Inverse scale parameter", beta_val); + + if (size_zero(n_ref, alpha_ref, beta_ref)) { return 1.0; } - T_partials_return P(1.0); auto ops_partials = make_partials_propagator(alpha_ref, beta_ref); - scalar_seq_view n_vec(n_ref); - scalar_seq_view alpha_vec(alpha_ref); - scalar_seq_view beta_vec(beta_ref); - size_t size_alpha = stan::math::size(alpha); - size_t size_n_alpha = max_size(n, alpha); - size_t max_size_seq_view = max_size(n, alpha, beta); - // Explicit return for extreme values // The gradients are technically ill-defined, but treated as zero - for (size_t i = 0; i < stan::math::size(n); i++) { - if (n_vec.val(i) < 0) { - return ops_partials.build(0.0); - } - } - - VectorBuilder, T_partials_return, T_shape> - digamma_alpha_vec(size_alpha); - VectorBuilder, T_partials_return, T_n, T_shape> - digamma_sum_vec(size_n_alpha); - - if constexpr (is_autodiff_v) { - for (size_t i = 0; i < size_alpha; i++) { - digamma_alpha_vec[i] = digamma(alpha_vec.val(i)); - } - for (size_t i = 0; i < size_n_alpha; i++) { - const T_partials_return n_dbl = n_vec.val(i); - const T_partials_return alpha_dbl = alpha_vec.val(i); - digamma_sum_vec[i] = digamma(n_dbl + alpha_dbl + 1); - } - } - - for (size_t i = 0; i < max_size_seq_view; i++) { - // Explicit results for extreme values - // The gradients are technically ill-defined, but treated as zero - if (n_vec.val(i) == std::numeric_limits::max()) { - return ops_partials.build(1.0); - } - - const T_partials_return n_dbl = n_vec.val(i); - const T_partials_return alpha_dbl = alpha_vec.val(i); - const T_partials_return beta_dbl = beta_vec.val(i); - const T_partials_return inv_beta_p1 = inv(beta_dbl + 1); - const T_partials_return p_dbl = beta_dbl * inv_beta_p1; - const T_partials_return d_dbl = square(inv_beta_p1); - - const T_partials_return P_i = inc_beta(alpha_dbl, n_dbl + 1.0, p_dbl); - - P *= P_i; - - if constexpr (is_autodiff_v) { - partials<0>(ops_partials)[i] - += inc_beta_dda(alpha_dbl, n_dbl + 1, p_dbl, digamma_alpha_vec[i], - digamma_sum_vec[i]) - / P_i; - } - - if constexpr (is_autodiff_v) { - partials<1>(ops_partials)[i] - += inc_beta_ddz(alpha_dbl, n_dbl + 1.0, p_dbl) * d_dbl / P_i; - } + if (any(n_val < 0)) { + return ops_partials.build(0.0); } - - if constexpr (is_autodiff_v) { - for (size_t i = 0; i < size_alpha; ++i) { - partials<0>(ops_partials)[i] *= P; - } + if (any(n_val == std::numeric_limits::max())) { + return ops_partials.build(1.0); } - if constexpr (is_autodiff_v) { - for (size_t i = 0; i < stan::math::size(beta); ++i) { - partials<1>(ops_partials)[i] *= P; + if constexpr (is_autodiff_v || is_autodiff_v) { + auto inv_beta_p1 = to_ref(inv(beta_val + 1.0)); + auto p_val = to_ref(beta_val * inv_beta_p1); + auto n_p1 = to_ref(promote_scalar(n_val + 1.0)); + auto P_i = to_ref(as_array_or_scalar(inc_beta(alpha_val, n_p1, p_val))); + const T_partials_return P = prod(P_i); + // d/dx prod(P_i) = sum(dP_i/dx * P / P_i) + if constexpr (is_autodiff_v && is_autodiff_v) { + auto P_div_P_i = to_ref(P / P_i); + partials<0>(ops_partials) + = inc_beta_dda(alpha_val, n_p1, p_val, to_ref(digamma(alpha_val)), + to_ref(digamma(alpha_val + n_p1))) + * P_div_P_i; + partials<1>(ops_partials) + = as_array_or_scalar(inc_beta_ddz(alpha_val, n_p1, p_val)) + * square(inv_beta_p1) * P_div_P_i; + } else if constexpr (is_autodiff_v) { + partials<0>(ops_partials) + = inc_beta_dda(alpha_val, n_p1, p_val, to_ref(digamma(alpha_val)), + to_ref(digamma(alpha_val + n_p1))) + / P_i * P; + } else if constexpr (is_autodiff_v) { + partials<1>(ops_partials) + = as_array_or_scalar(inc_beta_ddz(alpha_val, n_p1, p_val)) + * square(inv_beta_p1) / P_i * P; } + return ops_partials.build(P); + } else { + auto inv_beta_p1 = to_ref(inv(beta_val + 1.0)); + auto p_val = to_ref(beta_val * inv_beta_p1); + auto n_p1 = to_ref(promote_scalar(n_val + 1.0)); + const T_partials_return P + = prod(as_array_or_scalar(inc_beta(alpha_val, n_p1, p_val))); + return ops_partials.build(P); } - - return ops_partials.build(P); } } // namespace math diff --git a/stan/math/prim/prob/neg_binomial_lccdf.hpp b/stan/math/prim/prob/neg_binomial_lccdf.hpp index cc24a6b79a7..3de84a34589 100644 --- a/stan/math/prim/prob/neg_binomial_lccdf.hpp +++ b/stan/math/prim/prob/neg_binomial_lccdf.hpp @@ -53,8 +53,6 @@ inline return_type_t neg_binomial_lccdf( scalar_seq_view alpha_vec(alpha_ref); scalar_seq_view beta_vec(beta_ref); size_t size_n = stan::math::size(n); - size_t size_alpha = stan::math::size(alpha); - size_t size_n_alpha = max_size(n, alpha); size_t max_size_seq_view = max_size(n, alpha, beta_param); // Explicit return for extreme values @@ -65,27 +63,6 @@ inline return_type_t neg_binomial_lccdf( } } - VectorBuilder, T_partials_return, T_n> digammaN_vec( - size_n); - VectorBuilder, T_partials_return, T_shape> - digammaAlpha_vec(size_alpha); - VectorBuilder, T_partials_return, T_n, T_shape> - digammaSum_vec(size_n_alpha); - - if constexpr (is_autodiff_v) { - for (size_t i = 0; i < size_n; i++) { - digammaN_vec[i] = digamma(n_vec.val(i) + 1); - } - for (size_t i = 0; i < size_alpha; i++) { - digammaAlpha_vec[i] = digamma(alpha_vec.val(i)); - } - for (size_t i = 0; i < size_n_alpha; i++) { - const T_partials_return n_dbl = n_vec.val(i); - const T_partials_return alpha_dbl = alpha_vec.val(i); - digammaSum_vec[i] = digamma(n_dbl + alpha_dbl + 1); - } - } - for (size_t i = 0; i < max_size_seq_view; i++) { // Explicit results for extreme values // The gradients are technically ill-defined, but treated as zero @@ -93,25 +70,24 @@ inline return_type_t neg_binomial_lccdf( return ops_partials.build(negative_infinity()); } - const T_partials_return n_dbl = n_vec.val(i); - const T_partials_return alpha_dbl = alpha_vec.val(i); - const T_partials_return beta_dbl = beta_vec.val(i); - const T_partials_return inv_beta_p1 = inv(beta_dbl + 1); - const T_partials_return p_dbl = beta_dbl * inv_beta_p1; - const T_partials_return d_dbl = square(inv_beta_p1); - const T_partials_return Pi = 1.0 - inc_beta(alpha_dbl, n_dbl + 1.0, p_dbl); - const T_partials_return beta_func = beta(n_dbl + 1, alpha_dbl); + const auto n_dbl = n_vec.val(i); + const auto alpha_dbl = alpha_vec.val(i); + const auto beta_dbl = beta_vec.val(i); + const auto inv_beta_p1 = inv(beta_dbl + 1); + const auto p_dbl = beta_dbl * inv_beta_p1; + const auto d_dbl = square(inv_beta_p1); + const auto Pi = 1.0 - inc_beta(alpha_dbl, n_dbl + 1.0, p_dbl); + const auto beta_func = beta(n_dbl + 1, alpha_dbl); P += log(Pi); + // alpha is the first shape of inc_beta, so only its partial is needed if constexpr (is_autodiff_v) { - T_partials_return g1 = 0; - T_partials_return g2 = 0; - - grad_reg_inc_beta(g1, g2, alpha_dbl, n_dbl + 1, p_dbl, - digammaAlpha_vec[i], digammaN_vec[i], digammaSum_vec[i], - beta_func); - partials<0>(ops_partials)[i] -= g1 / Pi; + partials<0>(ops_partials)[i] + -= grad_reg_inc_beta( + alpha_dbl, n_dbl + 1, p_dbl, digamma(alpha_dbl), 0.0, + digamma(n_dbl + alpha_dbl + 1), 0.0) + / Pi; } if constexpr (is_autodiff_v) { partials<1>(ops_partials)[i] -= d_dbl * pow(1 - p_dbl, n_dbl) diff --git a/stan/math/prim/prob/neg_binomial_lcdf.hpp b/stan/math/prim/prob/neg_binomial_lcdf.hpp index ba85ba41d0d..7296870e58e 100644 --- a/stan/math/prim/prob/neg_binomial_lcdf.hpp +++ b/stan/math/prim/prob/neg_binomial_lcdf.hpp @@ -53,8 +53,6 @@ inline return_type_t neg_binomial_lcdf( scalar_seq_view alpha_vec(alpha_ref); scalar_seq_view beta_vec(beta_ref); size_t size_n = stan::math::size(n); - size_t size_alpha = stan::math::size(alpha); - size_t size_n_alpha = max_size(n, alpha); size_t max_size_seq_view = max_size(n, alpha, beta_param); // Explicit return for extreme values @@ -65,27 +63,6 @@ inline return_type_t neg_binomial_lcdf( } } - VectorBuilder, T_partials_return, T_n> digammaN_vec( - size_n); - VectorBuilder, T_partials_return, T_shape> - digammaAlpha_vec(size_alpha); - VectorBuilder, T_partials_return, T_n, T_shape> - digammaSum_vec(size_n_alpha); - - if constexpr (is_autodiff_v) { - for (size_t i = 0; i < size_n; i++) { - digammaN_vec[i] = digamma(n_vec.val(i) + 1); - } - for (size_t i = 0; i < size_alpha; i++) { - digammaAlpha_vec[i] = digamma(alpha_vec.val(i)); - } - for (size_t i = 0; i < size_n_alpha; i++) { - const T_partials_return n_dbl = n_vec.val(i); - const T_partials_return alpha_dbl = alpha_vec.val(i); - digammaSum_vec[i] = digamma(n_dbl + alpha_dbl + 1); - } - } - for (size_t i = 0; i < max_size_seq_view; i++) { // Explicit results for extreme values // The gradients are technically ill-defined, but treated as zero @@ -104,14 +81,13 @@ inline return_type_t neg_binomial_lcdf( P += log(Pi); + // alpha is the first shape of inc_beta, so only its partial is needed if constexpr (is_autodiff_v) { - T_partials_return g1 = 0; - T_partials_return g2 = 0; - - grad_reg_inc_beta(g1, g2, alpha_dbl, n_dbl + 1, p_dbl, - digammaAlpha_vec[i], digammaN_vec[i], digammaSum_vec[i], - beta_func); - partials<0>(ops_partials)[i] += g1 / Pi; + partials<0>(ops_partials)[i] + += grad_reg_inc_beta( + alpha_dbl, n_dbl + 1, p_dbl, digamma(alpha_dbl), 0.0, + digamma(n_dbl + alpha_dbl + 1), 0.0) + / Pi; } if constexpr (is_autodiff_v) { partials<1>(ops_partials)[i] += d_dbl * pow(1 - p_dbl, n_dbl) diff --git a/stan/math/prim/prob/student_t_cdf.hpp b/stan/math/prim/prob/student_t_cdf.hpp index f9720c61209..78cfa8f5ee3 100644 --- a/stan/math/prim/prob/student_t_cdf.hpp +++ b/stan/math/prim/prob/student_t_cdf.hpp @@ -60,26 +60,6 @@ inline return_type_t student_t_cdf( } } - T_partials_return digammaHalf = 0; - - VectorBuilder, T_partials_return, T_dof> digamma_vec( - math::size(nu)); - VectorBuilder, T_partials_return, T_dof> digammaNu_vec( - math::size(nu)); - VectorBuilder, T_partials_return, T_dof> - digammaNuPlusHalf_vec(math::size(nu)); - - if constexpr (is_autodiff_v) { - digammaHalf = digamma(0.5); - - for (size_t i = 0; i < stan::math::size(nu); i++) { - const T_partials_return nu_dbl = nu_vec.val(i); - - digammaNu_vec[i] = digamma(0.5 * nu_dbl); - digammaNuPlusHalf_vec[i] = digamma(0.5 + 0.5 * nu_dbl); - } - } - for (size_t n = 0; n < N; n++) { // Explicit results for extreme values // The gradients are technically ill-defined, but treated as zero @@ -87,20 +67,19 @@ inline return_type_t student_t_cdf( continue; } - const T_partials_return sigma_inv = 1.0 / sigma_vec.val(n); - const T_partials_return t = (y_vec.val(n) - mu_vec.val(n)) * sigma_inv; - const T_partials_return nu_dbl = nu_vec.val(n); - const T_partials_return q = nu_dbl / (t * t); - const T_partials_return r = 1.0 / (1.0 + q); - const T_partials_return J = 2 * r * r * q / t; - const T_partials_return betaNuHalf = beta(0.5, 0.5 * nu_dbl); + const auto sigma_inv = 1.0 / sigma_vec.val(n); + const auto t = (y_vec.val(n) - mu_vec.val(n)) * sigma_inv; + const auto nu_dbl = nu_vec.val(n); + const auto q = nu_dbl / (t * t); + const auto r = 1.0 / (1.0 + q); + const auto J = 2 * r * r * q / t; + const auto betaNuHalf = beta(0.5, 0.5 * nu_dbl); double zJacobian = t > 0 ? -0.5 : 0.5; if (q < 2) { - T_partials_return z - = inc_beta(0.5 * nu_dbl, (T_partials_return)0.5, 1.0 - r); - const T_partials_return Pn = t > 0 ? 1.0 - 0.5 * z : 0.5 * z; - const T_partials_return d_ibeta + auto z = inc_beta(0.5 * nu_dbl, 0.5, 1.0 - r); + const auto Pn = t > 0 ? 1.0 - 0.5 * z : 0.5 * z; + const auto d_ibeta = pow(r, -0.5) * pow(1.0 - r, 0.5 * nu_dbl - 1) / betaNuHalf; P *= Pn; @@ -110,12 +89,10 @@ inline return_type_t student_t_cdf( += -zJacobian * d_ibeta * J * sigma_inv / Pn; } if constexpr (is_autodiff_v) { - T_partials_return g1 = 0; - T_partials_return g2 = 0; - - grad_reg_inc_beta(g1, g2, 0.5 * nu_dbl, (T_partials_return)0.5, 1.0 - r, - digammaNu_vec[n], digammaHalf, - digammaNuPlusHalf_vec[n], betaNuHalf); + // nu / 2 is the first shape of inc_beta here and 1 / 2 the second + const auto g1 = grad_reg_inc_beta( + 0.5 * nu_dbl, 0.5, 1.0 - r, digamma(0.5 * nu_dbl), 0.0, + digamma(0.5 + 0.5 * nu_dbl), 0.0); partials<1>(ops_partials)[n] += zJacobian * (d_ibeta * (r / t) * (r / t) + 0.5 * g1) / Pn; @@ -131,15 +108,13 @@ inline return_type_t student_t_cdf( } } else { - T_partials_return z - = 1.0 - inc_beta((T_partials_return)0.5, 0.5 * nu_dbl, r); + auto z = 1.0 - inc_beta(0.5, 0.5 * nu_dbl, r); zJacobian *= -1; - const T_partials_return Pn = t > 0 ? 1.0 - 0.5 * z : 0.5 * z; + const auto Pn = t > 0 ? 1.0 - 0.5 * z : 0.5 * z; - T_partials_return d_ibeta - = pow(1.0 - r, 0.5 * nu_dbl - 1) * pow(r, -0.5) / betaNuHalf; + auto d_ibeta = pow(1.0 - r, 0.5 * nu_dbl - 1) * pow(r, -0.5) / betaNuHalf; P *= Pn; @@ -148,12 +123,10 @@ inline return_type_t student_t_cdf( += zJacobian * d_ibeta * J * sigma_inv / Pn; } if constexpr (is_autodiff_v) { - T_partials_return g1 = 0; - T_partials_return g2 = 0; - - grad_reg_inc_beta(g1, g2, (T_partials_return)0.5, 0.5 * nu_dbl, r, - digammaHalf, digammaNu_vec[n], - digammaNuPlusHalf_vec[n], betaNuHalf); + // nu / 2 is the second shape of inc_beta here and 1 / 2 the first + const auto g2 = grad_reg_inc_beta( + 0.5, 0.5 * nu_dbl, r, 0.0, digamma(0.5 * nu_dbl), + digamma(0.5 + 0.5 * nu_dbl), 0.0); partials<1>(ops_partials)[n] += zJacobian * (-d_ibeta * (r / t) * (r / t) + 0.5 * g2) / Pn; diff --git a/stan/math/prim/prob/student_t_lccdf.hpp b/stan/math/prim/prob/student_t_lccdf.hpp index af8cb6208b0..224a3165f50 100644 --- a/stan/math/prim/prob/student_t_lccdf.hpp +++ b/stan/math/prim/prob/student_t_lccdf.hpp @@ -62,26 +62,6 @@ inline return_type_t student_t_lccdf( } } - T_partials_return digammaHalf = 0; - - VectorBuilder, T_partials_return, T_dof> digamma_vec( - math::size(nu)); - VectorBuilder, T_partials_return, T_dof> digammaNu_vec( - math::size(nu)); - VectorBuilder, T_partials_return, T_dof> - digammaNuPlusHalf_vec(math::size(nu)); - - if constexpr (is_autodiff_v) { - digammaHalf = digamma(0.5); - - for (size_t i = 0; i < stan::math::size(nu); i++) { - const T_partials_return nu_dbl = nu_vec.val(i); - - digammaNu_vec[i] = digamma(0.5 * nu_dbl); - digammaNuPlusHalf_vec[i] = digamma(0.5 + 0.5 * nu_dbl); - } - } - for (size_t n = 0; n < N; n++) { // Explicit results for extreme values // The gradients are technically ill-defined, but treated as zero @@ -89,20 +69,19 @@ inline return_type_t student_t_lccdf( return ops_partials.build(negative_infinity()); } - const T_partials_return sigma_inv = 1.0 / sigma_vec.val(n); - const T_partials_return t = (y_vec.val(n) - mu_vec.val(n)) * sigma_inv; - const T_partials_return nu_dbl = nu_vec.val(n); - const T_partials_return q = nu_dbl / (t * t); - const T_partials_return r = 1.0 / (1.0 + q); - const T_partials_return J = 2 * r * r * q / t; - const T_partials_return betaNuHalf = beta(0.5, 0.5 * nu_dbl); - T_partials_return zJacobian = t > 0 ? -0.5 : 0.5; + const auto sigma_inv = 1.0 / sigma_vec.val(n); + const auto t = (y_vec.val(n) - mu_vec.val(n)) * sigma_inv; + const auto nu_dbl = nu_vec.val(n); + const auto q = nu_dbl / (t * t); + const auto r = 1.0 / (1.0 + q); + const auto J = 2 * r * r * q / t; + const auto betaNuHalf = beta(0.5, 0.5 * nu_dbl); + auto zJacobian = t > 0 ? -0.5 : 0.5; if (q < 2) { - T_partials_return z - = inc_beta(0.5 * nu_dbl, (T_partials_return)0.5, 1.0 - r); - const T_partials_return Pn = t > 0 ? 0.5 * z : 1.0 - 0.5 * z; - const T_partials_return d_ibeta + auto z = inc_beta(0.5 * nu_dbl, 0.5, 1.0 - r); + const auto Pn = t > 0 ? 0.5 * z : 1.0 - 0.5 * z; + const auto d_ibeta = pow(r, -0.5) * pow(1.0 - r, 0.5 * nu_dbl - 1) / betaNuHalf; P += log(Pn); @@ -113,12 +92,10 @@ inline return_type_t student_t_lccdf( } if constexpr (is_autodiff_v) { - T_partials_return g1 = 0; - T_partials_return g2 = 0; - - grad_reg_inc_beta(g1, g2, 0.5 * nu_dbl, (T_partials_return)0.5, 1.0 - r, - digammaNu_vec[n], digammaHalf, - digammaNuPlusHalf_vec[n], betaNuHalf); + // nu / 2 is the first shape of inc_beta here and 1 / 2 the second + const auto g1 = grad_reg_inc_beta( + 0.5 * nu_dbl, 0.5, 1.0 - r, digamma(0.5 * nu_dbl), 0.0, + digamma(0.5 + 0.5 * nu_dbl), 0.0); partials<1>(ops_partials)[n] -= zJacobian * (d_ibeta * (r / t) * (r / t) + 0.5 * g1) / Pn; @@ -134,14 +111,12 @@ inline return_type_t student_t_lccdf( } } else { - T_partials_return z - = 1.0 - inc_beta((T_partials_return)0.5, 0.5 * nu_dbl, r); + auto z = 1.0 - inc_beta(0.5, 0.5 * nu_dbl, r); zJacobian *= -1; - const T_partials_return Pn = t > 0 ? 0.5 * z : 1.0 - 0.5 * z; + const auto Pn = t > 0 ? 0.5 * z : 1.0 - 0.5 * z; - T_partials_return d_ibeta - = pow(1.0 - r, 0.5 * nu_dbl - 1) * pow(r, -0.5) / betaNuHalf; + auto d_ibeta = pow(1.0 - r, 0.5 * nu_dbl - 1) * pow(r, -0.5) / betaNuHalf; P += log(Pn); @@ -151,12 +126,10 @@ inline return_type_t student_t_lccdf( } if constexpr (is_autodiff_v) { - T_partials_return g1 = 0; - T_partials_return g2 = 0; - - grad_reg_inc_beta(g1, g2, (T_partials_return)0.5, 0.5 * nu_dbl, r, - digammaHalf, digammaNu_vec[n], - digammaNuPlusHalf_vec[n], betaNuHalf); + // nu / 2 is the second shape of inc_beta here and 1 / 2 the first + const auto g2 = grad_reg_inc_beta( + 0.5, 0.5 * nu_dbl, r, 0.0, digamma(0.5 * nu_dbl), + digamma(0.5 + 0.5 * nu_dbl), 0.0); partials<1>(ops_partials)[n] -= zJacobian * (-d_ibeta * (r / t) * (r / t) + 0.5 * g2) / Pn; diff --git a/stan/math/prim/prob/student_t_lcdf.hpp b/stan/math/prim/prob/student_t_lcdf.hpp index 793b23273ed..be791924977 100644 --- a/stan/math/prim/prob/student_t_lcdf.hpp +++ b/stan/math/prim/prob/student_t_lcdf.hpp @@ -62,26 +62,6 @@ inline return_type_t student_t_lcdf( } } - T_partials_return digammaHalf = 0; - - VectorBuilder, T_partials_return, T_dof> digamma_vec( - math::size(nu)); - VectorBuilder, T_partials_return, T_dof> digammaNu_vec( - math::size(nu)); - VectorBuilder, T_partials_return, T_dof> - digammaNuPlusHalf_vec(math::size(nu)); - - if constexpr (is_autodiff_v) { - digammaHalf = digamma(0.5); - - for (size_t i = 0; i < stan::math::size(nu); i++) { - const T_partials_return nu_dbl = nu_vec.val(i); - - digammaNu_vec[i] = digamma(0.5 * nu_dbl); - digammaNuPlusHalf_vec[i] = digamma(0.5 + 0.5 * nu_dbl); - } - } - for (size_t n = 0; n < N; n++) { // Explicit results for extreme values // The gradients are technically ill-defined, but treated as zero @@ -89,20 +69,19 @@ inline return_type_t student_t_lcdf( continue; } - const T_partials_return sigma_inv = 1.0 / sigma_vec.val(n); - const T_partials_return t = (y_vec.val(n) - mu_vec.val(n)) * sigma_inv; - const T_partials_return nu_dbl = nu_vec.val(n); - const T_partials_return q = nu_dbl / (t * t); - const T_partials_return r = 1.0 / (1.0 + q); - const T_partials_return J = 2 * r * r * q / t; - const T_partials_return betaNuHalf = beta(0.5, 0.5 * nu_dbl); - T_partials_return zJacobian = t > 0 ? -0.5 : 0.5; + const auto sigma_inv = 1.0 / sigma_vec.val(n); + const auto t = (y_vec.val(n) - mu_vec.val(n)) * sigma_inv; + const auto nu_dbl = nu_vec.val(n); + const auto q = nu_dbl / (t * t); + const auto r = 1.0 / (1.0 + q); + const auto J = 2 * r * r * q / t; + const auto betaNuHalf = beta(0.5, 0.5 * nu_dbl); + auto zJacobian = t > 0 ? -0.5 : 0.5; if (q < 2) { - T_partials_return z - = inc_beta(0.5 * nu_dbl, (T_partials_return)0.5, 1.0 - r); - const T_partials_return Pn = t > 0 ? 1.0 - 0.5 * z : 0.5 * z; - const T_partials_return d_ibeta + auto z = inc_beta(0.5 * nu_dbl, 0.5, 1.0 - r); + const auto Pn = t > 0 ? 1.0 - 0.5 * z : 0.5 * z; + const auto d_ibeta = pow(r, -0.5) * pow(1.0 - r, 0.5 * nu_dbl - 1) / betaNuHalf; P += log(Pn); @@ -113,12 +92,10 @@ inline return_type_t student_t_lcdf( } if constexpr (is_autodiff_v) { - T_partials_return g1 = 0; - T_partials_return g2 = 0; - - grad_reg_inc_beta(g1, g2, 0.5 * nu_dbl, (T_partials_return)0.5, 1.0 - r, - digammaNu_vec[n], digammaHalf, - digammaNuPlusHalf_vec[n], betaNuHalf); + // nu / 2 is the first shape of inc_beta here and 1 / 2 the second + const auto g1 = grad_reg_inc_beta( + 0.5 * nu_dbl, 0.5, 1.0 - r, digamma(0.5 * nu_dbl), 0.0, + digamma(0.5 + 0.5 * nu_dbl), 0.0); partials<1>(ops_partials)[n] += zJacobian * (d_ibeta * (r / t) * (r / t) + 0.5 * g1) / Pn; @@ -134,14 +111,12 @@ inline return_type_t student_t_lcdf( } } else { - T_partials_return z - = 1.0 - inc_beta((T_partials_return)0.5, 0.5 * nu_dbl, r); + auto z = 1.0 - inc_beta(0.5, 0.5 * nu_dbl, r); zJacobian *= -1; - const T_partials_return Pn = t > 0 ? 1.0 - 0.5 * z : 0.5 * z; + const auto Pn = t > 0 ? 1.0 - 0.5 * z : 0.5 * z; - T_partials_return d_ibeta - = pow(1.0 - r, 0.5 * nu_dbl - 1) * pow(r, -0.5) / betaNuHalf; + auto d_ibeta = pow(1.0 - r, 0.5 * nu_dbl - 1) * pow(r, -0.5) / betaNuHalf; P += log(Pn); @@ -151,12 +126,10 @@ inline return_type_t student_t_lcdf( } if constexpr (is_autodiff_v) { - T_partials_return g1 = 0; - T_partials_return g2 = 0; - - grad_reg_inc_beta(g1, g2, (T_partials_return)0.5, 0.5 * nu_dbl, r, - digammaHalf, digammaNu_vec[n], - digammaNuPlusHalf_vec[n], betaNuHalf); + // nu / 2 is the second shape of inc_beta here and 1 / 2 the first + const auto g2 = grad_reg_inc_beta( + 0.5, 0.5 * nu_dbl, r, 0.0, digamma(0.5 * nu_dbl), + digamma(0.5 + 0.5 * nu_dbl), 0.0); partials<1>(ops_partials)[n] += zJacobian * (-d_ibeta * (r / t) * (r / t) + 0.5 * g2) / Pn; diff --git a/stan/math/rev/fun.hpp b/stan/math/rev/fun.hpp index e5ef938a302..efb00634da1 100644 --- a/stan/math/rev/fun.hpp +++ b/stan/math/rev/fun.hpp @@ -71,6 +71,7 @@ #include #include #include +#include #include #include #include diff --git a/stan/math/rev/fun/grad_reg_inc_beta.hpp b/stan/math/rev/fun/grad_reg_inc_beta.hpp new file mode 100644 index 00000000000..b345e8ffd7d --- /dev/null +++ b/stan/math/rev/fun/grad_reg_inc_beta.hpp @@ -0,0 +1,774 @@ +#ifndef STAN_MATH_REV_FUN_GRAD_REG_INC_BETA_HPP +#define STAN_MATH_REV_FUN_GRAD_REG_INC_BETA_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace stan { +namespace math { + +namespace internal { + +/** + * Hessian-vector product H v of the regularized incomplete beta function + * I_x(a, b), H its Hessian in (a, b), by the power series + * + * I_x(a, b) = N S, N = x^a / B(a, b), S = sum_n P_n / (a + n), + * P_n = prod_{k <= n} (k - b) x / k. + * + * Differentiating I = N S twice, + * + * H v = N (D_v grad S + grad S D_v log N + grad log N D_v S + * + S (grad log N D_v log N + D_v grad log N)), + * + * and a diagonal entry is + * + * H_jj = N (d2S + 2 dlog_N dS + S (dlog_N^2 + d2log_N)), + * + * d and d2 the first and second derivatives in shape j. The a derivatives + * of each term t = P_n / (a + n) are closed form, and its b derivatives need + * those of P_n, carried by recurrence. + * + * Names: for G1 and G2 only the requested shape + * varies, and dX and d2X are the first and second derivatives of X in it. For + * BOTH, X_a and X_b are the partials of X and Dv_X its + * derivative along v. + * + * @tparam grads G1 for H_aa v_a, G2 for H_bb v_b, + * BOTH for H v + * @param a first shape + * @param b second shape + * @param x upper limit, in (0, 1) + * @param v_a a component of v; not used for G2 + * @param v_b b component of v; not used for G1 + * @param digamma_a digamma(a); not used for G2 + * @param digamma_b digamma(b); not used for G1 + * @param digamma_ab digamma(a + b) + * @return H_aa v_a for G1, H_bb v_b for G2, + * ((H v)_a, (H v)_b) for BOTH + */ +template +inline auto inc_beta_hvp_power_series(double a, double b, double x, double v_a, + double v_b, double digamma_a, + double digamma_b, double digamma_ab) { + constexpr int max_iterations = 600; + const double tolerance = std::numeric_limits::epsilon() / a; + const double log_x = std::log(x); + const double N = std::exp(a * log_x - lbeta(a, b)); + // boost's trigamma: stan::math::trigamma(double) is good to about 1e-9 below + // 5, and psi1(a + b) enters H_ab alone, where that error does not cancel + const double trigamma_ab = boost::math::trigamma(a + b, boost_policy_t<>()); + // n = 0: P_0 = 1 and the term is 1 / a, with a derivatives -1 / a^2 and + // 2 / a^3, and it does not depend on b + double product = 1.0; + double S = 1.0 / a; + if constexpr (grads == GRAD_RIB::G1) { + double dS = -S / a; + double d2S = 2.0 * S / (a * a); + for (int n = 1; n <= max_iterations; ++n) { + product *= (n - b) * x / n; + // dt/da = -t / (a + n), d2t/da2 = 2 t / (a + n)^2 + const double inv_apn = 1.0 / (a + n); + const double term = product * inv_apn; + S += term; + dS -= term * inv_apn; + d2S += 2.0 * term * inv_apn * inv_apn; + // for integer b the term and its a derivatives are 0 from n = b on + if (std::fabs(term) <= tolerance) { + break; + } + } + const double dlog_N = log_x - digamma_a + digamma_ab; + const double d2log_N + = trigamma_ab - boost::math::trigamma(a, boost_policy_t<>()); + return v_a * N + * (d2S + 2.0 * dlog_N * dS + S * (dlog_N * dlog_N + d2log_N)); + } else if constexpr (grads == GRAD_RIB::G2) { + double dproduct = 0.0; + double d2product = 0.0; + double dS = 0.0; + double d2S = 0.0; + for (int n = 1; n <= max_iterations; ++n) { + // P_n = f P_{n-1} with f = (n - b) x / n and df/db = -x / n + const double x_div_n = x / n; + const double factor = (n - b) * x_div_n; + d2product = factor * d2product - 2.0 * x_div_n * dproduct; + dproduct = factor * dproduct - x_div_n * product; + product *= factor; + const double inv_apn = 1.0 / (a + n); + const double term = product * inv_apn; + const double dterm = dproduct * inv_apn; + const double d2term = d2product * inv_apn; + S += term; + dS += dterm; + d2S += d2term; + // for integer b the term is 0 from n = b on while its b derivatives are + // not, so every carried series has to be small before stopping + if (std::fabs(term) <= tolerance && std::fabs(dterm) <= tolerance + && std::fabs(d2term) <= tolerance) { + break; + } + } + const double dlog_N = digamma_ab - digamma_b; + const double d2log_N + = trigamma_ab - boost::math::trigamma(b, boost_policy_t<>()); + return v_b * N + * (d2S + 2.0 * dlog_N * dS + S * (dlog_N * dlog_N + d2log_N)); + } else { + double product_b = 0.0; + double product_bb = 0.0; + double S_a = -S / a; + double S_b = 0.0; + double Dv_S_a = 2.0 * v_a * S / (a * a); + double Dv_S_b = 0.0; + for (int n = 1; n <= max_iterations; ++n) { + // P_n = f P_{n-1} with f = (n - b) x / n and df/db = -x / n + const double x_div_n = x / n; + const double factor = (n - b) * x_div_n; + product_bb = factor * product_bb - 2.0 * x_div_n * product_b; + product_b = factor * product_b - x_div_n * product; + product *= factor; + const double inv_apn = 1.0 / (a + n); + const double term = product * inv_apn; + const double term_b = product_b * inv_apn; + // d2t/da2 = 2 t / (a + n)^2, d2t/dadb = -t_b / (a + n), + // d2t/db2 = P''_n / (a + n) + const double term_aa = 2.0 * term * inv_apn * inv_apn; + const double term_ab = -term_b * inv_apn; + const double term_bb = product_bb * inv_apn; + S += term; + S_a -= term * inv_apn; + S_b += term_b; + Dv_S_a += v_a * term_aa + v_b * term_ab; + Dv_S_b += v_a * term_ab + v_b * term_bb; + // for integer b the term is 0 from n = b on while its b derivatives are + // not, so every carried series has to be small before stopping + if (std::fabs(term) <= tolerance && std::fabs(term_b) <= tolerance + && std::fabs(term_bb) <= tolerance) { + break; + } + } + const double log_N_a = log_x - digamma_a + digamma_ab; + const double log_N_b = digamma_ab - digamma_b; + const double Dv_log_N = log_N_a * v_a + log_N_b * v_b; + const double Dv_S = S_a * v_a + S_b * v_b; + // the Hessian of log N is + // [[psi1(a+b) - psi1(a), psi1(a+b)], [psi1(a+b), psi1(a+b) - psi1(b)]] + const double Dv_log_N_a + = (trigamma_ab - boost::math::trigamma(a, boost_policy_t<>())) * v_a + + trigamma_ab * v_b; + const double Dv_log_N_b + = trigamma_ab * v_a + + (trigamma_ab - boost::math::trigamma(b, boost_policy_t<>())) * v_b; + return std::pair{N + * (Dv_S_a + S_a * Dv_log_N + log_N_a * Dv_S + + S * (log_N_a * Dv_log_N + Dv_log_N_a)), + N + * (Dv_S_b + S_b * Dv_log_N + log_N_b * Dv_S + + S * (log_N_b * Dv_log_N + Dv_log_N_b))}; + } +} + +/** + * Hessian-vector product H v of the regularized incomplete beta function + * I_x(a, b), H its Hessian in (a, b), by the continued fraction + * https://dlmf.nist.gov/8.17.E22, I_x(a, b) = Q h with + * Q = x^a (1 - x)^b / (a B(a, b)) and h the continued fraction, evaluated + * with the modified Lentz method. Differentiating I = Q h twice, + * + * H = I (grad log I grad log I^T + Hess log Q + Hess log h). + * + * log h is the sum over the Lentz steps of log C and log D, so its + * derivatives are sums of theirs. With d the partial numerator, one step + * takes C' = 1 + d / C and D' = 1 / (1 + d D), and + * + * grad log C' = (grad d - d grad log C) / (C + d), + * grad log D' = -(grad d + d grad log D) D D'. + * + * Differentiating those once more along a direction v, + * + * D_v grad log C' = (Hess d v - D_v d grad log C - d D_v grad log C) + * / (C + d) - grad log C' D_v log(C C'), + * D_v log(C C') = (C D_v log C + D_v d) / (C + d), + * D_v grad log D' = -(Hess d v + D_v d grad log D + d D_v grad log D) D D' + * + grad log D' D_v log(D D'), + * D_v log(D D') = D_v log D - D D' (D_v d + d D_v log D). + * + * G1 and G2 take v = (1, 0) and (0, 1), so their + * D_v are second derivatives in one shape. + * + * Names: for G1 and G2 only the requested shape + * varies, and dX and d2X are the first and second derivatives of X in it. For + * BOTH, X_a and X_b are the partials of X and Dv_X its + * derivative along v. + * + * @tparam grads G1 for H_aa v_a, G2 for H_bb v_b, + * BOTH for H v + * @param a first shape + * @param b second shape + * @param x upper limit, in (0, 1) + * @param v_a a component of v; not used for G2 + * @param v_b b component of v; not used for G1 + * @param digamma_a digamma(a); not used for G2 + * @param digamma_b digamma(b); not used for G1 + * @param digamma_ab digamma(a + b) + * @return H_aa v_a for G1, H_bb v_b for G2, + * ((H v)_a, (H v)_b) for BOTH + */ +template +inline auto inc_beta_hvp_continued_fraction(double a, double b, double x, + double v_a, double v_b, + double digamma_a, double digamma_b, + double digamma_ab) { + constexpr int max_iterations = 300; + const double tolerance = 3.0 * std::numeric_limits::epsilon(); + const double small = std::sqrt(std::numeric_limits::min()); + const double log_x = std::log(x); + const double log1m_x = log1p(-x); + const double Q + = std::exp(a * log_x + b * log1m_x - std::log(a) - lbeta(a, b)); + // boost's trigamma, as in the power series + const double trigamma_ab = boost::math::trigamma(a + b, boost_policy_t<>()); + + // Partial numerator d_n = num / den over den = (a + n - 1)(a + n), + // https://dlmf.nist.gov/8.17.E23, with its a and b derivatives: + // n = 2m: num = m (b - m) x, num_b = m x + // n = 2m + 1: num = -(a + m)(a + b + m) x, num_a = -(2 (a + m) + b) x, + // num_b = -(a + m) x, num_aa = -2 x, num_ab = -x + // n = 1 is the odd case with m = 0, and d_bb = 0 for every n. + + // first step, n = 1, with C held at 1: D = 1 / (1 + d_1) + const double ap1 = a + 1.0; + const double d_1 = -(a + b) * x / ap1; + double D = 1.0 + d_1; + if (std::fabs(D) < small) { + D = small; + } + D = 1.0 / D; + double C = 1.0; + double h = D; + if constexpr (grads == GRAD_RIB::G1) { + const double dd_1 = -(x + d_1) / ap1; + const double d2d_1 = 2.0 * (x + d_1) / (ap1 * ap1); + double dlog_C = 0.0; + double d2log_C = 0.0; + double dlog_D = -dd_1 * D; + double d2log_D = -(d2d_1 + dlog_D * dd_1) * D; + double dlog_h = dlog_D; + double d2log_h = d2log_D; + double d2log_h_prev = d2log_h; + for (int n = 2; n <= 2 * max_iterations + 1; ++n) { + const int m = n / 2; + const bool odd = n % 2 == 1; + const double a_plus_n = a + n; + const double inv_den = 1.0 / (a_plus_n * (a_plus_n - 1.0)); + const double dden = 2.0 * a_plus_n - 1.0; + double num; + double dnum = 0.0; + double d2num = 0.0; + if (odd) { + num = -(a + m) * (a + b + m) * x; + dnum = -(2.0 * (a + m) + b) * x; + d2num = -2.0 * x; + } else { + num = m * (b - m) * x; + } + const double d_n = num * inv_den; + const double dd_n = (dnum - d_n * dden) * inv_den; + const double d2d_n = (d2num - 2.0 * dd_n * dden - 2.0 * d_n) * inv_den; + + // one step of the modified Lentz method. In each half, the second + // derivative takes its old-step part, the first derivative updates in + // place, then scales dlog(C C_new) or dlog(D D_new). + // C_new = 1 + d_n / C, with C C_new = C + d_n + double C_new = 1.0 + d_n / C; + double C_C_new = C + d_n; + if (std::fabs(C_new) < small) { + C_new = small; + C_C_new = C * small; + } + const double inv_C_C_new = 1.0 / C_C_new; + const double dlog_C_C_new = (C * dlog_C + dd_n) * inv_C_C_new; + d2log_C = (d2d_n - dd_n * dlog_C - d_n * d2log_C) * inv_C_C_new; + dlog_C = (dd_n - d_n * dlog_C) * inv_C_C_new; + d2log_C -= dlog_C * dlog_C_C_new; + + // D_new = 1 / (1 + d_n D) + double D_new = 1.0 + d_n * D; + if (std::fabs(D_new) < small) { + D_new = small; + } + D_new = 1.0 / D_new; + const double D_D_new = D * D_new; + const double dlog_D_D_new = dlog_D - D_D_new * (dd_n + d_n * dlog_D); + d2log_D = -(d2d_n + dd_n * dlog_D + d_n * d2log_D) * D_D_new; + dlog_D = -(dd_n + d_n * dlog_D) * D_D_new; + d2log_D += dlog_D * dlog_D_D_new; + + dlog_h += dlog_C + dlog_D; + d2log_h += d2log_C + d2log_D; + C = C_new; + D = D_new; + const double delta = C_new * D_new; + h *= delta; + if (odd) { + if (std::fabs(delta - 1.0) < tolerance + && std::fabs(d2log_h - d2log_h_prev) + <= tolerance * std::fabs(d2log_h)) { + break; + } + d2log_h_prev = d2log_h; + } + } + const double dlog_I = log_x - 1.0 / a + digamma_ab - digamma_a + dlog_h; + const double d2log_Q = trigamma_ab + - boost::math::trigamma(a, boost_policy_t<>()) + + 1.0 / (a * a); + return v_a * h * Q * (dlog_I * dlog_I + d2log_Q + d2log_h); + } else if constexpr (grads == GRAD_RIB::G2) { + // d_n is linear in b, so d2d_n = 0 + const double dd_1 = -x / ap1; + double dlog_C = 0.0; + double d2log_C = 0.0; + double dlog_D = -dd_1 * D; + double d2log_D = -dlog_D * dd_1 * D; + double dlog_h = dlog_D; + double d2log_h = d2log_D; + double d2log_h_prev = d2log_h; + for (int n = 2; n <= 2 * max_iterations + 1; ++n) { + const int m = n / 2; + const bool odd = n % 2 == 1; + const double a_plus_n = a + n; + const double inv_den = 1.0 / (a_plus_n * (a_plus_n - 1.0)); + double dnum; + double num; + if (odd) { + dnum = -(a + m) * x; + num = dnum * (a + b + m); + } else { + dnum = m * x; + num = dnum * (b - m); + } + const double d_n = num * inv_den; + const double dd_n = dnum * inv_den; + + // one step of the modified Lentz method, as for G1 with d2d_n = 0. + // C_new = 1 + d_n / C, with C C_new = C + d_n + double C_new = 1.0 + d_n / C; + double C_C_new = C + d_n; + if (std::fabs(C_new) < small) { + C_new = small; + C_C_new = C * small; + } + const double inv_C_C_new = 1.0 / C_C_new; + const double dlog_C_C_new = (C * dlog_C + dd_n) * inv_C_C_new; + d2log_C = -(dd_n * dlog_C + d_n * d2log_C) * inv_C_C_new; + dlog_C = (dd_n - d_n * dlog_C) * inv_C_C_new; + d2log_C -= dlog_C * dlog_C_C_new; + + // D_new = 1 / (1 + d_n D) + double D_new = 1.0 + d_n * D; + if (std::fabs(D_new) < small) { + D_new = small; + } + D_new = 1.0 / D_new; + const double D_D_new = D * D_new; + const double dlog_D_D_new = dlog_D - D_D_new * (dd_n + d_n * dlog_D); + d2log_D = -(dd_n * dlog_D + d_n * d2log_D) * D_D_new; + dlog_D = -(dd_n + d_n * dlog_D) * D_D_new; + d2log_D += dlog_D * dlog_D_D_new; + + dlog_h += dlog_C + dlog_D; + d2log_h += d2log_C + d2log_D; + C = C_new; + D = D_new; + const double delta = C_new * D_new; + h *= delta; + if (odd) { + if (std::fabs(delta - 1.0) < tolerance + && std::fabs(d2log_h - d2log_h_prev) + <= tolerance * std::fabs(d2log_h)) { + break; + } + d2log_h_prev = d2log_h; + } + } + const double dlog_I = log1m_x + digamma_ab - digamma_b + dlog_h; + const double d2log_Q + = trigamma_ab - boost::math::trigamma(b, boost_policy_t<>()); + return v_b * h * Q * (dlog_I * dlog_I + d2log_Q + d2log_h); + } else { + double log_C_a = 0.0; + double log_C_b = 0.0; + double Dv_log_C_a = 0.0; + double Dv_log_C_b = 0.0; + const double d_1_a = -(x + d_1) / ap1; + const double d_1_b = -x / ap1; + double log_D_a = -d_1_a * D; + double log_D_b = -d_1_b * D; + // D_v log D = -D D_v d_1 + const double d_1_aa = 2.0 * (x + d_1) / (ap1 * ap1); + const double d_1_ab = x / (ap1 * ap1); + const double Dv_d_1 = d_1_a * v_a + d_1_b * v_b; + double Dv_log_D_a + = -(d_1_aa * v_a + d_1_ab * v_b) * D - log_D_a * D * Dv_d_1; + double Dv_log_D_b = -d_1_ab * v_a * D - log_D_b * D * Dv_d_1; + double log_h_a = log_D_a; + double log_h_b = log_D_b; + double Dv_log_h_a = Dv_log_D_a; + double Dv_log_h_b = Dv_log_D_b; + double Dv_log_h_a_prev = Dv_log_h_a; + double Dv_log_h_b_prev = Dv_log_h_b; + for (int n = 2; n <= 2 * max_iterations + 1; ++n) { + const int m = n / 2; + const bool odd = n % 2 == 1; + const double a_plus_n = a + n; + const double inv_den = 1.0 / (a_plus_n * (a_plus_n - 1.0)); + const double den_a = 2.0 * a_plus_n - 1.0; + double num_b; + double num; + double num_a = 0.0; + double num_aa = 0.0; + double num_ab = 0.0; + if (odd) { + num_b = -(a + m) * x; + num = num_b * (a + b + m); + num_a = -(2.0 * (a + m) + b) * x; + num_aa = -2.0 * x; + num_ab = -x; + } else { + num_b = m * x; + num = num_b * (b - m); + } + const double d_n = num * inv_den; + const double d_n_a = (num_a - d_n * den_a) * inv_den; + const double d_n_b = num_b * inv_den; + const double d_n_aa + = (num_aa - 2.0 * d_n_a * den_a - 2.0 * d_n) * inv_den; + const double d_n_ab = (num_ab - d_n_b * den_a) * inv_den; + const double Dv_d_n = d_n_a * v_a + d_n_b * v_b; + const double Dv_d_n_a = d_n_aa * v_a + d_n_ab * v_b; + const double Dv_d_n_b = d_n_ab * v_a; + + // one step of the modified Lentz method. In each half, the Dv_ + // derivatives take their old-step part, the partials update in place, + // then scale D_v log(C C_new) or D_v log(D D_new). + // C_new = 1 + d_n / C, with C C_new = C + d_n + double C_new = 1.0 + d_n / C; + double C_C_new = C + d_n; + if (std::fabs(C_new) < small) { + C_new = small; + C_C_new = C * small; + } + const double inv_C_C_new = 1.0 / C_C_new; + const double Dv_log_C = log_C_a * v_a + log_C_b * v_b; + const double Dv_log_C_C_new = (C * Dv_log_C + Dv_d_n) * inv_C_C_new; + Dv_log_C_a + = (Dv_d_n_a - Dv_d_n * log_C_a - d_n * Dv_log_C_a) * inv_C_C_new; + log_C_a = (d_n_a - d_n * log_C_a) * inv_C_C_new; + Dv_log_C_a -= log_C_a * Dv_log_C_C_new; + Dv_log_C_b + = (Dv_d_n_b - Dv_d_n * log_C_b - d_n * Dv_log_C_b) * inv_C_C_new; + log_C_b = (d_n_b - d_n * log_C_b) * inv_C_C_new; + Dv_log_C_b -= log_C_b * Dv_log_C_C_new; + + // D_new = 1 / (1 + d_n D) + double D_new = 1.0 + d_n * D; + if (std::fabs(D_new) < small) { + D_new = small; + } + D_new = 1.0 / D_new; + const double D_D_new = D * D_new; + const double Dv_log_D = log_D_a * v_a + log_D_b * v_b; + const double Dv_log_D_D_new + = Dv_log_D - D_D_new * (Dv_d_n + d_n * Dv_log_D); + Dv_log_D_a = -(Dv_d_n_a + Dv_d_n * log_D_a + d_n * Dv_log_D_a) * D_D_new; + log_D_a = -(d_n_a + d_n * log_D_a) * D_D_new; + Dv_log_D_a += log_D_a * Dv_log_D_D_new; + Dv_log_D_b = -(Dv_d_n_b + Dv_d_n * log_D_b + d_n * Dv_log_D_b) * D_D_new; + log_D_b = -(d_n_b + d_n * log_D_b) * D_D_new; + Dv_log_D_b += log_D_b * Dv_log_D_D_new; + + log_h_a += log_C_a + log_D_a; + log_h_b += log_C_b + log_D_b; + Dv_log_h_a += Dv_log_C_a + Dv_log_D_a; + Dv_log_h_b += Dv_log_C_b + Dv_log_D_b; + C = C_new; + D = D_new; + const double delta = C_new * D_new; + h *= delta; + if (odd) { + if (std::fabs(delta - 1.0) < tolerance + && std::fabs(Dv_log_h_a - Dv_log_h_a_prev) + <= tolerance * std::fabs(Dv_log_h_a) + && std::fabs(Dv_log_h_b - Dv_log_h_b_prev) + <= tolerance * std::fabs(Dv_log_h_b)) { + break; + } + Dv_log_h_a_prev = Dv_log_h_a; + Dv_log_h_b_prev = Dv_log_h_b; + } + } + const double I = h * Q; + const double log_I_a = log_x - 1.0 / a + digamma_ab - digamma_a + log_h_a; + const double log_I_b = log1m_x + digamma_ab - digamma_b + log_h_b; + const double Dv_log_I = log_I_a * v_a + log_I_b * v_b; + // the Hessian of log Q is that of -log B(a, b) plus diag(1 / a^2, 0) + const double Dv_log_Q_a + = (trigamma_ab - boost::math::trigamma(a, boost_policy_t<>()) + + 1.0 / (a * a)) + * v_a + + trigamma_ab * v_b; + const double Dv_log_Q_b + = trigamma_ab * v_a + + (trigamma_ab - boost::math::trigamma(b, boost_policy_t<>())) * v_b; + return std::pair{I * (log_I_a * Dv_log_I + Dv_log_Q_a + Dv_log_h_a), + I * (log_I_b * Dv_log_I + Dv_log_Q_b + Dv_log_h_b)}; + } +} + +/** + * Hessian-vector product H v of the regularized incomplete beta function + * I_z(a, b), H its Hessian in (a, b): the reverse pass of + * (d/da I_z(a, b), d/db I_z(a, b)) with output adjoints v. + * + * The branch choice follows grad_reg_inc_beta. Under the symmetry relation + * I_z(a, b) = 1 - I_{1-z}(b, a), H v = -swap(H' swap(v)), H' the Hessian of + * I_{1-z}(b, a) in its own arguments, so H_aa = -H'_bb and H_bb = -H'_aa. + * + * @tparam grads G1 for H_aa v_a, G2 for H_bb v_b, + * BOTH for H v + * @param a first shape + * @param b second shape + * @param z upper limit, in [0, 1] + * @param v_a a component of v; not used for G2 + * @param v_b b component of v; not used for G1 + * @param digamma_a digamma(a); not used for G2 + * @param digamma_b digamma(b); not used for G1 + * @param digamma_ab digamma(a + b) + * @return H_aa v_a for G1, H_bb v_b for G2, + * ((H v)_a, (H v)_b) for BOTH; 0 at z = 0 and z = 1 + */ +template +inline std::conditional_t, + double> +inc_beta_hvp(double a, double b, double z, double v_a, double v_b, + double digamma_a, double digamma_b, double digamma_ab) { + if (z == 0.0 || z == 1.0) { + return {}; + } + // the symmetry relation swaps which shape is requested + constexpr GRAD_RIB swapped_grads = swap_grad_rib(grads); + const double mean = a / (a + b); + const bool use_power_series + = (z < mean && b * z <= 1.0 && z <= 0.95) + || (z >= mean && a * (1.0 - z) <= 1.0 && z >= 0.05); + if (use_power_series) { + if (z >= mean) { + const auto hv = inc_beta_hvp_power_series( + b, a, 1.0 - z, v_b, v_a, digamma_b, digamma_a, digamma_ab); + if constexpr (grads != GRAD_RIB::BOTH) { + return -hv; + } else { + return {-hv.second, -hv.first}; + } + } else { + return inc_beta_hvp_power_series(a, b, z, v_a, v_b, digamma_a, + digamma_b, digamma_ab); + } + } else { + if (z >= (a - 1.0) / (a + b - 2.0)) { + const auto hv = inc_beta_hvp_continued_fraction( + b, a, 1.0 - z, v_b, v_a, digamma_b, digamma_a, digamma_ab); + if constexpr (grads != GRAD_RIB::BOTH) { + return -hv; + } else { + return {-hv.second, -hv.first}; + } + } else { + return inc_beta_hvp_continued_fraction( + a, b, z, v_a, v_b, digamma_a, digamma_b, digamma_ab); + } + } +} + +} // namespace internal + +/** + * The requested gradients of the regularized incomplete beta function + * inc_beta(a, b, z) with respect to a and + * b, for var arguments. + * + * The values are computed in double by the prim overload. The reverse pass + * does not differentiate through its series or continued fraction: with + * output adjoints v, the a and b adjoints are H v for H the Hessian of + * I_z(a, b) in (a, b), computed directly by internal::inc_beta_hvp, and the + * z adjoint is the closed form + * + * d/dz d/da I_z(a, b) = inc_beta_ddz(a, b, z) (log z - psi(a) + psi(a+b)), + * d/dz d/db I_z(a, b) = inc_beta_ddz(a, b, z) (log1p(-z) - psi(b) + + * psi(a+b)). + * + * The digamma arguments only make the values cheap and must equal + * digamma(a), digamma(b) and digamma(a + b); their adjoints are 0 and the + * a and b adjoints above are total derivatives. The cross derivative needs + * the digamma value a one-sided call does not take, so the reverse pass of + * G1 computes digamma(b) when b is a var, and that of + * G2 digamma(a) when a is a var. betaAB is not used. + * + * Each argument is a var or arithmetic; mixing with + * fvar goes to the prim overload. + * + * @tparam grads_requested G1 for the derivative with respect to + * a, G2 for the derivative with respect to b, + * BOTH for both + * @tparam T1 type of a + * @tparam T2 type of b + * @tparam T3 type of z + * @tparam T4 type of digammaA + * @tparam T5 type of digammaB + * @tparam T6 type of digammaSum + * @tparam T7 type of betaAB + * @param a a + * @param b b + * @param z z + * @param digammaA the value of digamma(a); not used for + * G2 + * @param digammaB the value of digamma(b); not used for + * G1 + * @param digammaSum the value of digamma(a + b) + * @param betaAB not used + * @return d/da I_z(a, b) for G1, d/db for G2, the + * pair (d/da, d/db) for BOTH + */ +template >>* = nullptr> +inline auto grad_reg_inc_beta(const T1& a, const T2& b, const T3& z, + const T4& digammaA, const T5& digammaB, + const T6& digammaSum, const T7& /* betaAB */) { + const double digamma_a = value_of(digammaA); + const double digamma_b = value_of(digammaB); + const double digamma_ab = value_of(digammaSum); + if constexpr (grads_requested == GRAD_RIB::G1) { + return make_callback_var( + grad_reg_inc_beta(value_of(a), value_of(b), + value_of(z), digamma_a, digamma_b, + digamma_ab, 0.0), + [a, b, z, digamma_a, digamma_ab](auto& vi) { + const double a_val = value_of(a); + const double b_val = value_of(b); + const double z_val = value_of(z); + if constexpr (is_var_v && is_var_v) { + // the b adjoint is the cross term H_ab v_a, which needs both + // sides and digamma(b), which the G1 call does not take + const auto [hv_a, hv_b] = internal::inc_beta_hvp( + a_val, b_val, z_val, vi.adj(), 0.0, digamma_a, digamma(b_val), + digamma_ab); + a.adj() += hv_a; + b.adj() += hv_b; + } else if constexpr (is_var_v) { + // H_aa v_a does not use digamma(b) + a.adj() += internal::inc_beta_hvp( + a_val, b_val, z_val, vi.adj(), 0.0, digamma_a, 0.0, digamma_ab); + } else if constexpr (is_var_v) { + b.adj() += internal::inc_beta_hvp( + a_val, b_val, z_val, vi.adj(), 0.0, digamma_a, + digamma(b_val), digamma_ab) + .second; + } + if constexpr (is_var_v) { + if (z_val > 0.0 && z_val < 1.0) { + const double log_term = std::log(z_val) - digamma_a + digamma_ab; + z.adj() + += vi.adj() * inc_beta_ddz(a_val, b_val, z_val) * log_term; + } + } + }); + } else if constexpr (grads_requested == GRAD_RIB::G2) { + return make_callback_var( + grad_reg_inc_beta(value_of(a), value_of(b), + value_of(z), digamma_a, digamma_b, + digamma_ab, 0.0), + [a, b, z, digamma_b, digamma_ab](auto& vi) { + const double a_val = value_of(a); + const double b_val = value_of(b); + const double z_val = value_of(z); + if constexpr (is_var_v && is_var_v) { + // the a adjoint is the cross term H_ab v_b, which needs both + // sides and digamma(a), which the G2 call does not take + const auto [hv_a, hv_b] = internal::inc_beta_hvp( + a_val, b_val, z_val, 0.0, vi.adj(), digamma(a_val), digamma_b, + digamma_ab); + a.adj() += hv_a; + b.adj() += hv_b; + } else if constexpr (is_var_v) { + a.adj() += internal::inc_beta_hvp( + a_val, b_val, z_val, 0.0, vi.adj(), digamma(a_val), + digamma_b, digamma_ab) + .first; + } else if constexpr (is_var_v) { + // H_bb v_b does not use digamma(a) + b.adj() += internal::inc_beta_hvp( + a_val, b_val, z_val, 0.0, vi.adj(), 0.0, digamma_b, digamma_ab); + } + if constexpr (is_var_v) { + if (z_val > 0.0 && z_val < 1.0) { + const double log_term = log1p(-z_val) - digamma_b + digamma_ab; + z.adj() + += vi.adj() * inc_beta_ddz(a_val, b_val, z_val) * log_term; + } + } + }); + } else { // both + const auto [g1_val, g2_val] = grad_reg_inc_beta( + value_of(a), value_of(b), value_of(z), digamma_a, digamma_b, digamma_ab, + 0.0); + // Both outputs feed one reverse pass, which runs after everything that + // uses them has added into their adjoints. + var g1 = g1_val; + var g2 = g2_val; + reverse_pass_callback([a, b, z, g1, g2, digamma_a, digamma_b, + digamma_ab]() { + const double a_val = value_of(a); + const double b_val = value_of(b); + const double z_val = value_of(z); + const double v_a = g1.adj(); + const double v_b = g2.adj(); + if constexpr (is_var_v || is_var_v) { + const auto [hv_a, hv_b] = internal::inc_beta_hvp( + a_val, b_val, z_val, v_a, v_b, digamma_a, digamma_b, digamma_ab); + if constexpr (is_var_v) { + a.adj() += hv_a; + } + if constexpr (is_var_v) { + b.adj() += hv_b; + } + } + if constexpr (is_var_v) { + if (z_val > 0.0 && z_val < 1.0) { + z.adj() += inc_beta_ddz(a_val, b_val, z_val) + * (v_a * (std::log(z_val) - digamma_a + digamma_ab) + + v_b * (log1p(-z_val) - digamma_b + digamma_ab)); + } + } + }); + return std::pair{g1, g2}; + } +} + +} // namespace math +} // namespace stan +#endif diff --git a/stan/math/rev/fun/inc_beta.hpp b/stan/math/rev/fun/inc_beta.hpp index be8a65b7b6a..3c4481b81ec 100644 --- a/stan/math/rev/fun/inc_beta.hpp +++ b/stan/math/rev/fun/inc_beta.hpp @@ -3,41 +3,60 @@ #include #include -#include -#include -#include +#include #include #include -#include +#include +#include namespace stan { namespace math { -namespace internal { - -class inc_beta_vvv_vari : public op_vvv_vari { - public: - inc_beta_vvv_vari(vari* avi, vari* bvi, vari* cvi) - : op_vvv_vari(inc_beta(avi->val_, bvi->val_, cvi->val_), avi, bvi, cvi) {} - void chain() { - double d_a; - double d_b; - const double beta_ab = beta(avi_->val_, bvi_->val_); - grad_reg_inc_beta(d_a, d_b, avi_->val_, bvi_->val_, cvi_->val_, - digamma(avi_->val_), digamma(bvi_->val_), - digamma(avi_->val_ + bvi_->val_), beta_ab); - - avi_->adj_ += adj_ * d_a; - bvi_->adj_ += adj_ * d_b; - cvi_->adj_ += adj_ * std::pow(1 - cvi_->val_, bvi_->val_ - 1) - * std::pow(cvi_->val_, avi_->val_ - 1) / beta_ab; - } -}; - -} // namespace internal - -inline var inc_beta(const var& a, const var& b, const var& c) { - return var(new internal::inc_beta_vvv_vari(a.vi_, b.vi_, c.vi_)); +/** + * The regularized incomplete beta function I_z(a, b) for var + * arguments. Only the partials of the var arguments are + * computed in the reverse pass. Each argument is a var or + * arithmetic; mixing with fvar is not supported. + * + * @tparam T1 type of a + * @tparam T2 type of b + * @tparam T3 type of z + * @param a first shape, positive + * @param b second shape, positive + * @param z upper limit, in [0, 1] + * @return I_z(a, b) + */ +template * = nullptr, + require_all_var_or_arithmetic_t* = nullptr, + require_any_var_t* = nullptr> +inline var inc_beta(const T1& a, const T2& b, const T3& z) { + return make_callback_var( + inc_beta(value_of(a), value_of(b), value_of(z)), [a, b, z](auto& vi) { + const double a_val = value_of(a); + const double b_val = value_of(b); + const double z_val = value_of(z); + if constexpr (is_autodiff_v && is_autodiff_v) { + const auto [d_a, d_b] = grad_reg_inc_beta( + a_val, b_val, z_val, digamma(a_val), digamma(b_val), + digamma(a_val + b_val), 0.0); + a.adj() += vi.adj() * d_a; + b.adj() += vi.adj() * d_b; + } else if constexpr (is_autodiff_v) { + a.adj() += vi.adj() + * grad_reg_inc_beta( + a_val, b_val, z_val, digamma(a_val), 0.0, + digamma(a_val + b_val), 0.0); + } else if constexpr (is_autodiff_v) { + b.adj() += vi.adj() + * grad_reg_inc_beta( + a_val, b_val, z_val, 0.0, digamma(b_val), + digamma(a_val + b_val), 0.0); + } + if constexpr (is_autodiff_v) { + z.adj() += vi.adj() * inc_beta_ddz(a_val, b_val, z_val); + } + }); } } // namespace math diff --git a/test/unit/math/mix/fun/grad_reg_inc_beta_test.cpp b/test/unit/math/mix/fun/grad_reg_inc_beta_test.cpp index e6613d43ea1..7fef4350b67 100644 --- a/test/unit/math/mix/fun/grad_reg_inc_beta_test.cpp +++ b/test/unit/math/mix/fun/grad_reg_inc_beta_test.cpp @@ -1,145 +1,87 @@ -#include -#include -#include +#include +#include -TEST(ProbInternalMath, grad_reg_inc_beta_fv) { - using stan::math::digamma; - using stan::math::exp; - using stan::math::fvar; - using stan::math::lbeta; - using stan::math::var; +namespace { - fvar a = 1.0; - fvar b = 1.0; - fvar g = 0.4; - a.d_ = 1.0; - b.d_ = 1.0; - g.d_ = 1.0; - fvar dig_a = digamma(a); - fvar dig_b = digamma(b); - fvar dig_sum = digamma(a + b); - fvar beta_ab = exp(lbeta(a, b)); - fvar g_a; - fvar g_b; - - stan::math::grad_reg_inc_beta(g_a, g_b, a, b, g, dig_a, dig_b, dig_sum, - beta_ab); - EXPECT_FLOAT_EQ(-0.36651629442883944183907601651838247842001142107486495485, - g_a.val_.val()); - EXPECT_NEAR(0.306495375042422864944011633197968575202046200428315551199, - g_b.val_.val(), 1e-6); +// Points in each region of grad_reg_inc_beta: +// (1, 1, 0.4), (2, 3, 0.25), (0.5, 0.5, 0.3): power series +// (10, 1, 0.5): power series, integer b +// (2, 3, 0.6): power series through the symmetry relation +// (20, 20, 0.3), (5, 30, 0.1): continued fraction +// (20, 20, 0.7), (5, 30, 0.3): continued fraction through the symmetry +// relation +template +void expect_ad_grad_reg_inc_beta(const F& f) { + stan::test::expect_ad(f, 1.0, 1.0, 0.4); + stan::test::expect_ad(f, 2.0, 3.0, 0.25); + stan::test::expect_ad(f, 0.5, 0.5, 0.3); + stan::test::expect_ad(f, 10.0, 1.0, 0.5); + stan::test::expect_ad(f, 2.0, 3.0, 0.6); + stan::test::expect_ad(f, 20.0, 20.0, 0.3); + stan::test::expect_ad(f, 5.0, 30.0, 0.1); + stan::test::expect_ad(f, 20.0, 20.0, 0.7); + stan::test::expect_ad(f, 5.0, 30.0, 0.3); } -TEST(ProbInternalMath, grad_reg_inc_beta_fv_1stDeriv1) { - using stan::math::digamma; - using stan::math::exp; - using stan::math::fvar; - using stan::math::lbeta; - using stan::math::var; - - fvar a = 1.0; - fvar b = 1.0; - fvar g = 0.4; - a.d_ = 1.0; - b.d_ = 0.0; - g.d_ = 0.0; - fvar dig_a = digamma(a); - fvar dig_b = digamma(b); - fvar dig_sum = digamma(a + b); - fvar beta_ab = exp(lbeta(a, b)); - fvar g_a; - fvar g_b; - stan::math::grad_reg_inc_beta(g_a, g_b, a, b, g, dig_a, dig_b, dig_sum, - beta_ab); +} // namespace - std::vector y1{a.val_}; - std::vector grad1; - g_a.val_.grad(y1, grad1); - EXPECT_FLOAT_EQ(0.33583548212738989400284958327902414335945097838423129, - grad1[0]); +TEST(mathMixScalFun, grad_reg_inc_beta_g1) { + auto f = [](const auto& a, const auto& b, const auto& z) { + using stan::math::digamma; + return stan::math::grad_reg_inc_beta( + a, b, z, digamma(a), 0.0, digamma(a + b), 0.0); + }; + expect_ad_grad_reg_inc_beta(f); } -TEST(ProbInternalMath, grad_reg_inc_beta_fv_1stDeriv2) { - using stan::math::digamma; - using stan::math::exp; - using stan::math::fvar; - using stan::math::lbeta; - using stan::math::var; - - fvar a = 1.0; - fvar b = 1.0; - fvar g = 0.4; - a.d_ = 0.0; - b.d_ = 1.0; - g.d_ = 0.0; - fvar dig_a = digamma(a); - fvar dig_b = digamma(b); - fvar dig_sum = digamma(a + b); - fvar beta_ab = exp(lbeta(a, b)); - fvar g_a; - fvar g_b; - - stan::math::grad_reg_inc_beta(g_a, g_b, a, b, g, dig_a, dig_b, dig_sum, - beta_ab); - std::vector y1{b.val_}; - std::vector grad1; - g_b.val_.grad(y1, grad1); - EXPECT_NEAR(-0.156565690737548079304827886, grad1[0], 1e-6); +TEST(mathMixScalFun, grad_reg_inc_beta_g2) { + auto f = [](const auto& a, const auto& b, const auto& z) { + using stan::math::digamma; + return stan::math::grad_reg_inc_beta( + a, b, z, 0.0, digamma(b), digamma(a + b), 0.0); + }; + expect_ad_grad_reg_inc_beta(f); } -TEST(ProbInternalMath, grad_reg_inc_beta_fv_2ndDeriv1) { - using stan::math::digamma; - using stan::math::exp; - using stan::math::fvar; - using stan::math::lbeta; - using stan::math::var; - fvar a = 1.0; - fvar b = 1.0; - fvar g = 0.4; - a.d_ = 1.0; - b.d_ = 0.0; - g.d_ = 0.0; - fvar dig_a = digamma(a); - fvar dig_b = digamma(b); - fvar dig_sum = digamma(a + b); - fvar beta_ab = exp(lbeta(a, b)); - fvar g_a; - fvar g_b; - - stan::math::grad_reg_inc_beta(g_a, g_b, a, b, g, dig_a, dig_b, dig_sum, - beta_ab); - - std::vector y1{a.val_}; - std::vector grad1; - g_a.d_.grad(y1, grad1); - EXPECT_FLOAT_EQ(-0.30772293970781581317390510390046098438962772318921, - grad1[0]); +TEST(mathMixScalFun, grad_reg_inc_beta_both) { + auto f = [](const auto& a, const auto& b, const auto& z) { + using stan::math::digamma; + auto grads = stan::math::grad_reg_inc_beta( + a, b, z, digamma(a), digamma(b), digamma(a + b), 0.0); + Eigen::Matrix ret(2); + ret << grads.first, grads.second; + return ret; + }; + expect_ad_grad_reg_inc_beta(f); } -TEST(ProbInternalMath, grad_reg_inc_beta_fv_2ndDeriv2) { - using stan::math::digamma; - using stan::math::exp; - using stan::math::fvar; - using stan::math::lbeta; - using stan::math::var; - fvar a = 1.0; - fvar b = 1.0; - fvar g = 0.4; - a.d_ = 0.0; - b.d_ = 1.0; - g.d_ = 0.0; - fvar dig_a = digamma(a); - fvar dig_b = digamma(b); - fvar dig_sum = digamma(a + b); - fvar beta_ab = exp(lbeta(a, b)); - fvar g_a; - fvar g_b; - - stan::math::grad_reg_inc_beta(g_a, g_b, a, b, g, dig_a, dig_b, dig_sum, - beta_ab); +// the student_t_cdf call: a = nu / 2, b = 1 / 2 fixed, z = 1 - r +TEST(mathMixScalFun, grad_reg_inc_beta_student) { + auto f = [](const auto& a, const auto& /* b */, const auto& z) { + using stan::math::digamma; + return stan::math::grad_reg_inc_beta( + 0.5 * a, 0.5, 1.0 - z, digamma(0.5 * a), 0.0, digamma(0.5 + 0.5 * a), + 0.0); + }; + expect_ad_grad_reg_inc_beta(f); +} - std::vector y1{b.val_}; - std::vector grad1; - g_b.d_.grad(y1, grad1); - EXPECT_NEAR(0.079977766631361187517939795, grad1[0], 1e-4); +TEST(mathMixScalFun, grad_reg_inc_beta_out_params) { + // the overload writing both partials through out-parameters takes every + // argument as the same type + auto f = [](const auto& a, const auto& b, const auto& z) { + using stan::math::digamma; + using T = stan::return_type_t; + const T a_t = a; + const T b_t = b; + const T z_t = z; + T g1 = 0; + T g2 = 0; + stan::math::grad_reg_inc_beta(g1, g2, a_t, b_t, z_t, digamma(a_t), + digamma(b_t), digamma(a_t + b_t), T(0)); + Eigen::Matrix ret(2); + ret << g1, g2; + return ret; + }; + expect_ad_grad_reg_inc_beta(f); } diff --git a/test/unit/math/prim/fun/grad_reg_inc_beta_test.cpp b/test/unit/math/prim/fun/grad_reg_inc_beta_test.cpp index 881057a5cf6..c84a259b796 100644 --- a/test/unit/math/prim/fun/grad_reg_inc_beta_test.cpp +++ b/test/unit/math/prim/fun/grad_reg_inc_beta_test.cpp @@ -22,6 +22,22 @@ void expect_grads(double a, double b, double z, double dda, double ddb, << "a=" << a << " b=" << b << " z=" << z; EXPECT_NEAR(ddb, g2, rtol * std::fabs(ddb)) << "a=" << a << " b=" << b << " z=" << z; + + // each one-sided request carries only its own derivative through the + // expansion, with the other digamma unused + using stan::math::GRAD_RIB; + const double only_g1 = stan::math::grad_reg_inc_beta( + a, b, z, digamma(a), 0.0, digamma(a + b), 0.0); + EXPECT_NEAR(dda, only_g1, rtol * std::fabs(dda)) + << "G1 a=" << a << " b=" << b << " z=" << z; + const double only_g2 = stan::math::grad_reg_inc_beta( + a, b, z, 0.0, digamma(b), digamma(a + b), 0.0); + EXPECT_NEAR(ddb, only_g2, rtol * std::fabs(ddb)) + << "G2 a=" << a << " b=" << b << " z=" << z; + const auto both = stan::math::grad_reg_inc_beta( + a, b, z, digamma(a), digamma(b), digamma(a + b), 0.0); + EXPECT_EQ(g1, both.first) << "BOTH a=" << a << " b=" << b << " z=" << z; + EXPECT_EQ(g2, both.second) << "BOTH a=" << a << " b=" << b << " z=" << z; } } // namespace diff --git a/test/unit/math/prim/fun/inc_beta_dda_test.cpp b/test/unit/math/prim/fun/inc_beta_dda_test.cpp index 10c0f250c08..953ee47b661 100644 --- a/test/unit/math/prim/fun/inc_beta_dda_test.cpp +++ b/test/unit/math/prim/fun/inc_beta_dda_test.cpp @@ -71,3 +71,47 @@ TEST(MathFunctions, inc_beta_dda_regressions) { expect_dda(0.5, 0.5, 0.5, -0.58312180806163756, 1e-10); expect_dda(20, 20, 0.5, -0.063740468830439859, 1e-10); } + +TEST(MathFunctions, inc_beta_dda_vectorized) { + using stan::math::digamma; + using stan::math::inc_beta_dda; + + Eigen::VectorXd a(4); + a << 1.145, 446, 2, 600; + Eigen::VectorXd b(4); + b << 6786, 5, 3, 600; + Eigen::VectorXd z(4); + z << 1.8e-4, 0.76, 0.25, 0.5; + Eigen::VectorXd digamma_a(4); + Eigen::VectorXd digamma_ab(4); + for (int i = 0; i < 4; ++i) { + digamma_a[i] = digamma(a[i]); + digamma_ab[i] = digamma(a[i] + b[i]); + } + + Eigen::ArrayXd all_eigen = inc_beta_dda(a, b, z, digamma_a, digamma_ab); + ASSERT_EQ(4, all_eigen.size()); + for (int i = 0; i < 4; ++i) { + EXPECT_EQ(inc_beta_dda(a[i], b[i], z[i], digamma_a[i], digamma_ab[i]), + all_eigen[i]); + } + + std::vector z_std{1.8e-4, 0.76, 0.25, 0.5}; + Eigen::ArrayXd std_z = inc_beta_dda(a, b, z_std, digamma_a, digamma_ab); + for (int i = 0; i < 4; ++i) { + EXPECT_EQ(all_eigen[i], std_z[i]); + } + + // scalar shapes broadcast against a vector of upper limits + Eigen::ArrayXd only_z = inc_beta_dda(2.0, 3.0, z, digamma(2.0), digamma(5.0)); + ASSERT_EQ(4, only_z.size()); + for (int i = 0; i < 4; ++i) { + EXPECT_EQ(inc_beta_dda(2.0, 3.0, z[i], digamma(2.0), digamma(5.0)), + only_z[i]); + } + + Eigen::VectorXd too_short(3); + too_short << 0.1, 0.2, 0.3; + EXPECT_THROW(inc_beta_dda(a, b, too_short, digamma_a, digamma_ab), + std::invalid_argument); +} diff --git a/test/unit/math/prim/fun/inc_beta_ddz_test.cpp b/test/unit/math/prim/fun/inc_beta_ddz_test.cpp index f8aa0579778..e94808a62bf 100644 --- a/test/unit/math/prim/fun/inc_beta_ddz_test.cpp +++ b/test/unit/math/prim/fun/inc_beta_ddz_test.cpp @@ -42,3 +42,40 @@ TEST(MathFunctions, inc_beta_ddz) { EXPECT_FLOAT_EQ(0.0, inc_beta_ddz(large_a, large_b, large_z)) << "reasonable values for a, b, x"; } + +TEST(MathFunctions, inc_beta_ddz_vectorized) { + using stan::math::inc_beta_ddz; + + Eigen::VectorXd a(3); + a << 1.5, 15.0, 0.75; + Eigen::VectorXd b(3); + b << 1.25, 4.0, 30.0; + Eigen::VectorXd z(3); + z << 0.001, 0.5, 0.999; + + Eigen::VectorXd all_eigen = inc_beta_ddz(a, b, z); + ASSERT_EQ(3, all_eigen.size()); + for (int i = 0; i < 3; ++i) { + EXPECT_EQ(inc_beta_ddz(a[i], b[i], z[i]), all_eigen[i]); + } + + std::vector a_std{1.5, 15.0, 0.75}; + std::vector b_std{1.25, 4.0, 30.0}; + std::vector z_std{0.001, 0.5, 0.999}; + std::vector all_std = inc_beta_ddz(a_std, b_std, z_std); + ASSERT_EQ(3, all_std.size()); + for (int i = 0; i < 3; ++i) { + EXPECT_EQ(inc_beta_ddz(a_std[i], b_std[i], z_std[i]), all_std[i]); + } + + Eigen::VectorXd scalar_a = inc_beta_ddz(2.5, b, z); + Eigen::VectorXd scalar_b = inc_beta_ddz(a, 3.5, z); + Eigen::VectorXd scalar_z = inc_beta_ddz(a, b, 0.3); + Eigen::VectorXd only_z = inc_beta_ddz(2.5, 3.5, z); + for (int i = 0; i < 3; ++i) { + EXPECT_EQ(inc_beta_ddz(2.5, b[i], z[i]), scalar_a[i]); + EXPECT_EQ(inc_beta_ddz(a[i], 3.5, z[i]), scalar_b[i]); + EXPECT_EQ(inc_beta_ddz(a[i], b[i], 0.3), scalar_z[i]); + EXPECT_EQ(inc_beta_ddz(2.5, 3.5, z[i]), only_z[i]); + } +} From 2c96499277b35b963f17ac8cab35bc49ac5401bf Mon Sep 17 00:00:00 2001 From: Stan Jenkins Date: Wed, 21 Oct 2026 18:02:28 -0400 Subject: [PATCH 06/10] [Jenkins] auto-formatting by clang-format version 10.0.0-4ubuntu1 --- stan/math/prim/prob/neg_binomial_2_cdf.hpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/stan/math/prim/prob/neg_binomial_2_cdf.hpp b/stan/math/prim/prob/neg_binomial_2_cdf.hpp index 3890f5149b4..6a96393e9cb 100644 --- a/stan/math/prim/prob/neg_binomial_2_cdf.hpp +++ b/stan/math/prim/prob/neg_binomial_2_cdf.hpp @@ -94,7 +94,8 @@ inline return_type_t neg_binomial_2_cdf( P *= P_i; if constexpr (is_any_autodiff_v) { const auto d_dbl = square(inv_mu_plus_phi); - const auto inc_beta_ddz_i = inc_beta_ddz(phi_dbl, n_dbl_p1, p_dbl) * d_dbl / P_i; + const auto inc_beta_ddz_i + = inc_beta_ddz(phi_dbl, n_dbl_p1, p_dbl) * d_dbl / P_i; if constexpr (is_autodiff_v) { partials<0>(ops_partials)[i] -= inc_beta_ddz_i * phi_dbl; } @@ -102,8 +103,8 @@ inline return_type_t neg_binomial_2_cdf( partials<1>(ops_partials)[i] += inc_beta_dda(phi_dbl, n_dbl_p1, p_dbl, digamma_phi_vec[i], digamma_sum_vec[i]) - / P_i - + inc_beta_ddz_i * mu_dbl; + / P_i + + inc_beta_ddz_i * mu_dbl; } } } From a34e6f7499b864dcc995518d99df5fe3e99c4175 Mon Sep 17 00:00:00 2001 From: Aki Vehtari Date: Thu, 22 Oct 2026 10:59:42 +0300 Subject: [PATCH 07/10] Return infinity from inc_beta_ddz where the density is infinite --- stan/math/prim/fun/inc_beta_ddz.hpp | 9 +++++++-- test/unit/math/mix/fun/inc_beta_test.cpp | 19 +++++++++++++++++++ test/unit/math/prim/fun/inc_beta_ddz_test.cpp | 11 +++++++++++ test/unit/math/rev/fun/inc_beta_test.cpp | 15 +++++++++++++++ 4 files changed, 52 insertions(+), 2 deletions(-) diff --git a/stan/math/prim/fun/inc_beta_ddz.hpp b/stan/math/prim/fun/inc_beta_ddz.hpp index 2700b0d34b5..762e46bce9d 100644 --- a/stan/math/prim/fun/inc_beta_ddz.hpp +++ b/stan/math/prim/fun/inc_beta_ddz.hpp @@ -38,7 +38,9 @@ inline auto inc_beta_ddz(T1 a, T2 b, T3 z) { * Returns the partial derivative of the regularized * incomplete beta function, I_{z}(a, b) with respect to z. * - * Overload for double arguments, evaluated by Boost. + * Overload for double arguments, evaluated by Boost. Where the density is + * infinite (z = 0 with a < 1, or z = 1 with b < 1) it returns infinity + * instead of throwing. * * @param a first argument * @param b second argument @@ -47,7 +49,10 @@ inline auto inc_beta_ddz(T1 a, T2 b, T3 z) { */ inline double inc_beta_ddz(double a, double b, double z) { using boost::math::ibeta_derivative; - return ibeta_derivative(a, b, z); + using overflow_ignore + = boost::math::policies::policy>; + return ibeta_derivative(a, b, z, overflow_ignore()); } /** diff --git a/test/unit/math/mix/fun/inc_beta_test.cpp b/test/unit/math/mix/fun/inc_beta_test.cpp index 4f00deab729..37829e37254 100644 --- a/test/unit/math/mix/fun/inc_beta_test.cpp +++ b/test/unit/math/mix/fun/inc_beta_test.cpp @@ -28,6 +28,25 @@ TEST(mathMixScalFun, inc_beta) { stan::test::expect_ad(f, nan, nan, nan); } +TEST(mathMixScalFun, inc_beta_upper_bound_infinite_density) { + // At z = 1 with b < 1 the density is infinite; the forward and reverse + // sweeps of fvar must not throw there. + using stan::math::fvar; + using stan::math::var; + var a_v = 1.05; + var b_v = 0.35; + fvar a(a_v, 1.0); + fvar b(b_v, 0.0); + fvar z(1.0, 0.0); + fvar f; + EXPECT_NO_THROW(f = stan::math::inc_beta(a, b, z)); + EXPECT_EQ(1.0, f.val_.val()); + EXPECT_NO_THROW(f.val_.grad()); + stan::math::set_zero_all_adjoints(); + EXPECT_NO_THROW(f.d_.grad()); + stan::math::recover_memory(); +} + TEST(mathMixScalFun, inc_beta_vec) { auto f = [](const auto& x1, const auto& x2, const auto& x3) { return stan::math::inc_beta(x1, x2, x3); diff --git a/test/unit/math/prim/fun/inc_beta_ddz_test.cpp b/test/unit/math/prim/fun/inc_beta_ddz_test.cpp index e94808a62bf..d44b613c015 100644 --- a/test/unit/math/prim/fun/inc_beta_ddz_test.cpp +++ b/test/unit/math/prim/fun/inc_beta_ddz_test.cpp @@ -43,6 +43,17 @@ TEST(MathFunctions, inc_beta_ddz) { << "reasonable values for a, b, x"; } +TEST(MathFunctions, inc_beta_ddz_infinite_density) { + using stan::math::inc_beta_ddz; + // the density is infinite at z = 1 with b < 1 and at z = 0 with a < 1 + double at_one = 0; + double at_zero = 0; + EXPECT_NO_THROW(at_one = inc_beta_ddz(1.05, 0.35, 1.0)); + EXPECT_NO_THROW(at_zero = inc_beta_ddz(0.5, 2.0, 0.0)); + EXPECT_TRUE(std::isinf(at_one) && at_one > 0); + EXPECT_TRUE(std::isinf(at_zero) && at_zero > 0); +} + TEST(MathFunctions, inc_beta_ddz_vectorized) { using stan::math::inc_beta_ddz; diff --git a/test/unit/math/rev/fun/inc_beta_test.cpp b/test/unit/math/rev/fun/inc_beta_test.cpp index 9b2aefb1272..9d83916ebaf 100644 --- a/test/unit/math/rev/fun/inc_beta_test.cpp +++ b/test/unit/math/rev/fun/inc_beta_test.cpp @@ -41,3 +41,18 @@ TEST_F(AgradRev, inc_beta_shape_gradients) { expect_inc_beta_grads(10, 1, 0.5, -0.00067690154351557159, 0.00226574342615196, 1e-12); } + +TEST_F(AgradRev, inc_beta_upper_bound_infinite_density) { + // At z = 1, I is 1 for all shapes, so the shape gradients are 0; with + // b < 1 the density, the z gradient, is infinite. + using stan::math::var; + var a = 1.05; + var b = 0.35; + var z = 1.0; + var f = stan::math::inc_beta(a, b, z); + EXPECT_NO_THROW(f.grad()); + EXPECT_EQ(1.0, f.val()); + EXPECT_EQ(0.0, a.adj()); + EXPECT_EQ(0.0, b.adj()); + EXPECT_TRUE(std::isinf(z.adj()) && z.adj() > 0); +} From 3fd3fe2f6a3d68d9aa14af6fd982773b1311a2ec Mon Sep 17 00:00:00 2001 From: Aki Vehtari Date: Thu, 22 Oct 2026 11:20:42 +0300 Subject: [PATCH 08/10] Start the inc_beta upper-bound mix test from an empty autodiff stack --- test/unit/math/mix/fun/inc_beta_test.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/test/unit/math/mix/fun/inc_beta_test.cpp b/test/unit/math/mix/fun/inc_beta_test.cpp index 37829e37254..b90ea81b74c 100644 --- a/test/unit/math/mix/fun/inc_beta_test.cpp +++ b/test/unit/math/mix/fun/inc_beta_test.cpp @@ -33,6 +33,8 @@ TEST(mathMixScalFun, inc_beta_upper_bound_infinite_density) { // sweeps of fvar must not throw there. using stan::math::fvar; using stan::math::var; + // grad() runs every node on the stack; start from an empty one + stan::math::recover_memory(); var a_v = 1.05; var b_v = 0.35; fvar a(a_v, 1.0); From 207f47cdd182f16852252185b11b7dface769994 Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Thu, 22 Oct 2026 15:04:26 -0400 Subject: [PATCH 09/10] adds fixes for grad_reg_inc_beta --- stan/math/prim/fun/grad_reg_inc_beta.hpp | 32 +++++++++++++------ stan/math/rev/fun/grad_reg_inc_beta.hpp | 8 +++-- .../math/mix/fun/grad_reg_inc_beta_test.cpp | 4 +++ 3 files changed, 31 insertions(+), 13 deletions(-) diff --git a/stan/math/prim/fun/grad_reg_inc_beta.hpp b/stan/math/prim/fun/grad_reg_inc_beta.hpp index 37c3c4426e8..d83e9df5996 100644 --- a/stan/math/prim/fun/grad_reg_inc_beta.hpp +++ b/stan/math/prim/fun/grad_reg_inc_beta.hpp @@ -196,6 +196,12 @@ inline auto inc_beta_continued_fraction_kernel(const T1& a, const T2& b, using ret_t = return_type_t; constexpr bool calc_a = grads_requested != GRAD_RIB::G2; constexpr bool calc_b = grads_requested != GRAD_RIB::G1; + // For integer b the partial numerator d_{2b} is 0, so the value and the a + // derivatives stop changing from that step on while the b derivative of + // d_{2b} is not 0. When b is an autodiff type its derivatives are carried + // through this loop, so the b derivative has to converge before stopping + // even when it is not requested, as track_b does in the power series. + constexpr bool track_b = calc_b || is_autodiff_v; constexpr int max_iterations = 300; const double tolerance = 3.0 * std::numeric_limits::epsilon(); const double small = std::sqrt(std::numeric_limits::min()); @@ -229,7 +235,7 @@ inline auto inc_beta_continued_fraction_kernel(const T1& a, const T2& b, e_a = (x + d_1) / ap1 * D; r_a = e_a; } - if constexpr (calc_b) { + if constexpr (track_b) { // d d_1 / db = -x / (a + 1) e_b = x / ap1 * D; r_b = e_b; @@ -282,7 +288,7 @@ inline auto inc_beta_continued_fraction_kernel(const T1& a, const T2& b, e_a = -(dd_a + d * e_a) * D_D_new; r_a += c_a + e_a; } - if constexpr (calc_b) { + if constexpr (track_b) { c_b = (dd_b - d * c_b) / C_C_new; e_b = -(dd_b + d * e_b) * D_D_new; r_b += c_b + e_b; @@ -298,7 +304,7 @@ inline auto inc_beta_continued_fraction_kernel(const T1& a, const T2& b, if (fabs(value_of_rec(delta) - 1.0) < tolerance && (!calc_a || fabs(dh_a_now - dh_a_prev) <= tolerance * fabs(dh_a_now)) - && (!calc_b + && (!track_b || fabs(dh_b_now - dh_b_prev) <= tolerance * fabs(dh_b_now))) { break; } @@ -395,9 +401,13 @@ inline auto grad_reg_inc_beta(const T1& a, const T2& b, const T3& z, const double a_d = value_of_rec(a); const double b_d = value_of_rec(b); const double mean = a_d / (a_d + b_d); + // The series terms shrink by about z per step (1 - z for the flipped + // series), so it is limited to z <= 0.5 (z >= 0.5). For b < 1, b z <= 1 + // alone does not bound z, and at z = 0.95 the series takes most of its 600 + // terms where the continued fraction needs tens of steps. const bool use_power_series - = (z_d < mean && b_d * z_d <= 1.0 && z_d <= 0.95) - || (z_d >= mean && a_d * (1.0 - z_d) <= 1.0 && z_d >= 0.05); + = (z_d < mean && b_d * z_d <= 1.0 && z_d <= 0.5) + || (z_d >= mean && a_d * (1.0 - z_d) <= 1.0 && z_d >= 0.5); if (use_power_series) { // For z >= a / (a + b) the series is evaluated for I_{1-z}(b, a) through // the symmetry relation; its (d/da, d/db) are the negated (d/db, d/da) @@ -410,11 +420,13 @@ inline auto grad_reg_inc_beta(const T1& a, const T2& b, const T3& z, a, b, z, digammaA, digammaB, digammaSum); } } else { - // The continued fraction converges rapidly for z < (a - 1) / (a + b - 2); - // otherwise it is evaluated for I_{1-z}(b, a) through the symmetry - // relation. IEEE division: (a - 1) / 0 is +-inf and the comparison is - // well defined. - if (z_d >= (a_d - 1.0) / (a_d + b_d - 2.0)) { + // The continued fraction converges rapidly for z < (a + 1) / (a + b + 2) + // (Numerical Recipes, 3rd ed., section 6.4); otherwise it is evaluated + // for I_{1-z}(b, a) through the symmetry relation, where 1 - z is below + // the matching (b + 1) / (a + b + 2). The mode (a - 1) / (a + b - 2) is + // not a usable threshold: it exceeds 1 for b < 1, so the fraction would + // run at z near 1, where it needs hundreds of steps. + if (z_d >= (a_d + 1.0) / (a_d + b_d + 2.0)) { return internal::inc_beta_continued_fraction_kernel( b, a, 1.0 - z, digammaB, digammaA, digammaSum); } else { diff --git a/stan/math/rev/fun/grad_reg_inc_beta.hpp b/stan/math/rev/fun/grad_reg_inc_beta.hpp index b345e8ffd7d..90be0e2a247 100644 --- a/stan/math/rev/fun/grad_reg_inc_beta.hpp +++ b/stan/math/rev/fun/grad_reg_inc_beta.hpp @@ -571,9 +571,10 @@ inc_beta_hvp(double a, double b, double z, double v_a, double v_b, // the symmetry relation swaps which shape is requested constexpr GRAD_RIB swapped_grads = swap_grad_rib(grads); const double mean = a / (a + b); + // the same branch rule as grad_reg_inc_beta const bool use_power_series - = (z < mean && b * z <= 1.0 && z <= 0.95) - || (z >= mean && a * (1.0 - z) <= 1.0 && z >= 0.05); + = (z < mean && b * z <= 1.0 && z <= 0.5) + || (z >= mean && a * (1.0 - z) <= 1.0 && z >= 0.5); if (use_power_series) { if (z >= mean) { const auto hv = inc_beta_hvp_power_series( @@ -588,7 +589,8 @@ inc_beta_hvp(double a, double b, double z, double v_a, double v_b, digamma_b, digamma_ab); } } else { - if (z >= (a - 1.0) / (a + b - 2.0)) { + // the same threshold as grad_reg_inc_beta, so both take the same branch + if (z >= (a + 1.0) / (a + b + 2.0)) { const auto hv = inc_beta_hvp_continued_fraction( b, a, 1.0 - z, v_b, v_a, digamma_b, digamma_a, digamma_ab); if constexpr (grads != GRAD_RIB::BOTH) { diff --git a/test/unit/math/mix/fun/grad_reg_inc_beta_test.cpp b/test/unit/math/mix/fun/grad_reg_inc_beta_test.cpp index 7fef4350b67..eeadf6281b6 100644 --- a/test/unit/math/mix/fun/grad_reg_inc_beta_test.cpp +++ b/test/unit/math/mix/fun/grad_reg_inc_beta_test.cpp @@ -10,6 +10,8 @@ namespace { // (20, 20, 0.3), (5, 30, 0.1): continued fraction // (20, 20, 0.7), (5, 30, 0.3): continued fraction through the symmetry // relation +// (10, 1, 0.7), (4, 4, 0.5): continued fraction, integer b, where the +// partial numerator d_{2b} is 0 but its b derivative is not template void expect_ad_grad_reg_inc_beta(const F& f) { stan::test::expect_ad(f, 1.0, 1.0, 0.4); @@ -21,6 +23,8 @@ void expect_ad_grad_reg_inc_beta(const F& f) { stan::test::expect_ad(f, 5.0, 30.0, 0.1); stan::test::expect_ad(f, 20.0, 20.0, 0.7); stan::test::expect_ad(f, 5.0, 30.0, 0.3); + stan::test::expect_ad(f, 10.0, 1.0, 0.7); + stan::test::expect_ad(f, 4.0, 4.0, 0.5); } } // namespace From 6e1938c385b4f9f3fdb6ed5a019139888311c93d Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Thu, 22 Oct 2026 17:32:07 -0400 Subject: [PATCH 10/10] Move grad_reg_inc_beta notes into the function docs Multi-line comments in the prim and rev grad_reg_inc_beta kernels and in the new tests are now part of the doc comments, and the docs are rewritten to read in order: the expansion, its derivatives, how the code evaluates them, and when the loop stops. The swap of requested partials and the track_b flag are now written inline, the overloads dispatch with require_all_not_var_t / require_any_var_t, and the continued-fraction stop test is fixed to compile after inlining. --- stan/math/fwd/fun/inc_beta.hpp | 2 +- stan/math/prim/fun/grad_reg_inc_beta.hpp | 313 ++++++++++-------- stan/math/prim/fun/inc_beta_dda.hpp | 8 +- stan/math/prim/fun/inc_beta_ddb.hpp | 8 +- stan/math/prim/fun/inc_beta_ddz.hpp | 4 +- stan/math/rev/fun/grad_reg_inc_beta.hpp | 269 +++++++++------ stan/math/rev/fun/inc_beta.hpp | 2 +- .../math/mix/fun/grad_reg_inc_beta_test.cpp | 26 +- test/unit/math/mix/fun/inc_beta_test.cpp | 6 +- .../math/prim/fun/grad_reg_inc_beta_test.cpp | 18 +- test/unit/math/prim/fun/inc_beta_dda_test.cpp | 6 +- test/unit/math/prim/fun/inc_beta_ddb_test.cpp | 8 +- test/unit/math/rev/fun/inc_beta_test.cpp | 12 +- test/unit/math/rev/prob/beta_cdf_log_test.cpp | 6 +- test/unit/math/rev/prob/beta_cdf_test.cpp | 7 +- 15 files changed, 402 insertions(+), 293 deletions(-) diff --git a/stan/math/fwd/fun/inc_beta.hpp b/stan/math/fwd/fun/inc_beta.hpp index 1c928880ac0..d0c75a711aa 100644 --- a/stan/math/fwd/fun/inc_beta.hpp +++ b/stan/math/fwd/fun/inc_beta.hpp @@ -15,7 +15,7 @@ namespace math { /** * The regularized incomplete beta function I_x(a, b) for fvar * arguments. Only the partials of the fvar arguments are - * computed. Each argument is an fvar or arithmetic; mixing with + * computed. Each argument is an fvar or arithmetic. Mixing with * var is not supported. * * @tparam T1 type of a diff --git a/stan/math/prim/fun/grad_reg_inc_beta.hpp b/stan/math/prim/fun/grad_reg_inc_beta.hpp index d83e9df5996..c34e463f968 100644 --- a/stan/math/prim/fun/grad_reg_inc_beta.hpp +++ b/stan/math/prim/fun/grad_reg_inc_beta.hpp @@ -30,42 +30,39 @@ enum class GRAD_RIB : uint8_t { G1 = 0, G2 = 1, BOTH = 2 }; namespace internal { /** - * The partials requested after the symmetry relation - * I_z(a, b) = 1 - I_{1-z}(b, a) swaps the roles of a and b. - */ -inline constexpr GRAD_RIB swap_grad_rib(GRAD_RIB grads_requested) { - return grads_requested == GRAD_RIB::G1 - ? GRAD_RIB::G2 - : grads_requested == GRAD_RIB::G2 ? GRAD_RIB::G1 : GRAD_RIB::BOTH; -} - -/** - * True for the arguments of grad_reg_inc_beta handled by the var overload in - * stan/math/rev/fun/grad_reg_inc_beta.hpp: every argument is a var or - * arithmetic and one of a, b, z or the digamma values is a var. The prim - * overload takes every other combination. - */ -template -inline constexpr bool is_grad_reg_inc_beta_rev_v - = is_var_or_arithmetic::value - && (is_var::value || is_var::value || is_var::value - || is_var::value || is_var::value || is_var::value); - -/** - * Partial derivatives of the regularized incomplete beta function - * I_x(a, b) with respect to a and b, by the power series + * Partial derivatives of the regularized incomplete beta function I_x(a, b) + * with respect to a and b, computed from its power series. + * + * The series is + * + * I_x(a, b) = N S, N = x^a / B(a, b), S = sum_{n >= 0} t_n, + * t_n = P_n / (a + n), P_0 = 1, P_n = P_{n-1} (n - b) x / n, + * + * where S is 2F1(a, 1 - b; a + 1; x) / a. The derivatives of each term are * - * I_x(a, b) = x^a 2F1(a, 1 - b; a + 1; x) / (a B(a, b)), + * dt_n/da = -t_n / (a + n), dt_n/db = (dP_n/db) / (a + n), * - * differentiated term by term. Stops when a term is below eps / a in - * magnitude, or after 600 terms. Only the series for the requested partials - * are accumulated. + * and dP_n/db follows from the recurrence for P_n by the product rule. With + * log N = a log x - log B(a, b) and psi the digamma function, + * + * dI/da = N (dS/da + S (log x - psi(a) + psi(a + b))), + * dI/db = N (dS/db + S (psi(a + b) - psi(b))). + * + * Only the sums needed for the requested partials are accumulated. + * + * The loop stops once a term is below eps / a in magnitude, or after 600 + * terms. For a positive integer b, P_n is exactly 0 for n >= b but dP_n/db is + * not, so the sum for dS/db keeps changing after S has converged. The stop + * test therefore also requires dt_n/db to be small whenever that sum + * matters. That is the case when d/db is requested, and also when b is an + * autodiff type, because autodiff then differentiates the d/da result in b + * and needs the same terms. * * @tparam grads_requested which partials to compute - * @tparam flip_sign true when the caller passes (b, a, 1 - x) to evaluate - * its I_x(a, b) through the symmetry relation I_x(a, b) = 1 - I_{1-x}(b, a), - * so the results are negated and returned in the caller's order of a and b + * @tparam flip_sign true when the caller evaluates its I_z(a, b) through the + * symmetry relation I_z(a, b) = 1 - I_{1-z}(b, a) and so passes + * (b, a, 1 - z). The results are then negated, and for BOTH + * they are returned in the caller's order of a and b. * @tparam T1 type of a * @tparam T2 type of b * @tparam T3 type of x @@ -78,9 +75,10 @@ inline constexpr bool is_grad_reg_inc_beta_rev_v * @param digamma_a digamma(a) * @param digamma_b digamma(b) * @param digamma_ab digamma(a + b) - * @return d/da I_x(a, b) for G1, d/db for G2, the - * pair (d/da, d/db) for BOTH; with flip_sign - * these are negated and the pair is returned as (-d/db, -d/da) + * @return d/da I_x(a, b) for G1, d/db for G2, and + * the pair (d/da, d/db) for BOTH. With flip_sign + * these are negated and the pair is returned as (-d/db, -d/da), which are + * the caller's (d/da, d/db). */ template @@ -89,20 +87,11 @@ inline auto inc_beta_power_series_kernel(const T1& a, const T2& b, const T3& x, const T5& digamma_b, const T6& digamma_ab) { using ret_t = return_type_t; - constexpr bool calc_a = grads_requested != GRAD_RIB::G2; - constexpr bool calc_b = grads_requested != GRAD_RIB::G1; - // db_increment is d(term)/db. For a positive integer b the term is exactly - // 0 at n = b while d(term)/db is not, so the series must also run until - // d(term)/db is small: when d/db is requested, and when b is an autodiff - // type, because autodiff of the d/da series in b needs the same tail. - constexpr bool track_b = calc_b || is_autodiff_v; constexpr int max_iterations = 600; const double tolerance = std::numeric_limits::epsilon() / value_of_rec(a); - // series for 2F1(a, 1 - b; a + 1; x) / a and its a and b derivatives. The - // n-th term is product / (a + n), so its a derivative is -term / (a + n) - // and its b derivative is product_grad_b / (a + n). + // series for 2F1(a, 1 - b; a + 1; x) / a and its a and b derivatives ret_t product = 1.0; ret_t series_sum = inv(a); ret_t da = -series_sum / a; @@ -115,15 +104,15 @@ inline auto inc_beta_power_series_kernel(const T1& a, const T2& b, const T3& x, const auto new_product = product * factor; const auto term = new_product / apn; bool db_converged = true; - if constexpr (track_b) { + if constexpr (grads_requested != GRAD_RIB::G1 || is_autodiff_v) { product_grad_b = factor * product_grad_b - product * x_div_n; const auto db_increment = product_grad_b / apn; - if constexpr (calc_b) { + if constexpr (grads_requested != GRAD_RIB::G1) { db += db_increment; } db_converged = fabs(value_of_rec(db_increment)) <= tolerance; } - if constexpr (calc_a) { + if constexpr (grads_requested != GRAD_RIB::G2) { da -= term / apn; } product = new_product; @@ -136,41 +125,87 @@ inline auto inc_beta_power_series_kernel(const T1& a, const T2& b, const T3& x, // the symmetry relation negates the partials constexpr double sign = flip_sign ? -1.0 : 1.0; const auto log_x = log(x); - const ret_t normalization = sign * exp(a * log_x - lbeta(a, b)); + const auto normalize = sign * exp(a * log_x - lbeta(a, b)); if constexpr (grads_requested == GRAD_RIB::G1) { - return ret_t(normalization - * (da + series_sum * (digamma_ab - digamma_a + log_x))); + return normalize * (da + series_sum * (digamma_ab - digamma_a + log_x)); } else if constexpr (grads_requested == GRAD_RIB::G2) { - return ret_t(normalization * (db + series_sum * (digamma_ab - digamma_b))); + return normalize * (db + series_sum * (digamma_ab - digamma_b)); } else { - const ret_t grad_a - = normalization * (da + series_sum * (digamma_ab - digamma_a + log_x)); - const ret_t grad_b - = normalization * (db + series_sum * (digamma_ab - digamma_b)); + const auto grad_a + = normalize * (da + series_sum * (digamma_ab - digamma_a + log_x)); + const auto grad_b + = normalize * (db + series_sum * (digamma_ab - digamma_b)); if constexpr (flip_sign) { // the caller swapped a and b, so its (d/da, d/db) are (grad_b, grad_a) - return std::pair{grad_b, grad_a}; + return std::pair{grad_b, grad_a}; } else { - return std::pair{grad_a, grad_b}; + return std::pair{grad_a, grad_b}; } } } /** - * Partial derivatives of the regularized incomplete beta function - * I_x(a, b) with respect to a and b, by the continued fraction - * https://dlmf.nist.gov/8.17.E22 evaluated with the modified Lentz method, - * with the partial numerators https://dlmf.nist.gov/8.17.E23 differentiated - * in a and b and the derivatives carried through the recurrence (Thompson - * and Barnett (1986), appendix). Only the requested derivatives are carried. - * Stops when - * the multiplicative update is within 3 eps of 1 and the requested - * derivatives have converged, or after 300 double steps. + * Partial derivatives of the regularized incomplete beta function I_x(a, b) + * with respect to a and b, computed from its continued fraction. + * + * The continued fraction https://dlmf.nist.gov/8.17.E22 is + * + * I_x(a, b) = Q h, Q = x^a (1 - x)^b / (a B(a, b)), + * h = 1 / (1 + d_1 / (1 + d_2 / (1 + ...))), + * + * with the partial numerators https://dlmf.nist.gov/8.17.E23 + * + * d_1 = -(a + b) x / (a + 1), + * d_{2m} = m (b - m) x / ((a + 2m - 1)(a + 2m)), + * d_{2m+1} = -(a + m)(a + b + m) x / ((a + 2m)(a + 2m + 1)). + * + * Each d_n is num / den with den = (a + n - 1)(a + n), so its a derivative is + * (num_a - d_n den_a) / den, with den_a = 2 (a + n) - 1. The even numerators + * do not depend on a, and the odd ones have num_a = -(2 (a + m) + b) x. Every + * numerator is linear in b, so the b derivative is num_b / den. + * + * h is evaluated with the modified Lentz method. With h_n = A_n / B_n the + * n-th convergent, it keeps the ratios C = A_n / A_{n-1} and + * D = B_{n-1} / B_n, and step n updates them as + * + * C <- 1 + d_n / C, D <- 1 / (1 + d_n D), h <- h C D. + * + * The first step holds C at 1, so D = 1 / (1 + d_1) and h = D. When + * 1 + d_n / C or 1 + d_n D is smaller in magnitude than sqrt(DBL_MIN), it is + * replaced by that value, so that nothing divides by 0. + * + * The derivatives are carried through the same steps, following Thompson and + * Barnett (1986), appendix. They are kept as logarithmic derivatives + * c = (dC/da) / C, e = (dD/da) / D, and r = (dh/da) / h, and likewise for b. + * Differentiating the updates gives + * + * c <- (dd_n/da - d_n c) / (C C_new), + * e <- -(dd_n/da + d_n e) D D_new, + * r <- r + c + e, + * + * where C, D, c, and e on the right are the values before the step and + * C_new, D_new are the values after it. At the end, with psi the digamma + * function, + * + * dI/da = I (r_a + log x - 1 / a + psi(a + b) - psi(a)), + * dI/db = I (r_b + log(1 - x) + psi(a + b) - psi(b)). + * + * Only the requested derivatives are carried. + * + * The test for stopping runs after each odd n. The loop stops once C D is + * within 3 eps of 1 and each carried derivative dh = h r has changed by at + * most a relative 3 eps since the previous test, or after n = 601. For + * integer b the partial numerator d_{2b} is 0, so from that step on h and its + * a derivative no longer change. The b derivative of d_{2b} is b x / den, + * which is not 0, so dh/db is still changing. The b derivative is therefore + * tested whenever it matters, which is when d/db is requested and also when b + * is an autodiff type, as in the power series. * * @tparam grads_requested which partials to compute - * @tparam flip_sign true when the caller passes (b, a, 1 - x) to evaluate - * its I_x(a, b) through the symmetry relation I_x(a, b) = 1 - I_{1-x}(b, a), - * so the results are negated and returned in the caller's order of a and b + * @tparam flip_sign true when the caller evaluates its I_z(a, b) through the + * symmetry relation I_z(a, b) = 1 - I_{1-z}(b, a) and so passes + * (b, a, 1 - z). The results are then negated, and for BOTH + * they are returned in the caller's order of a and b. * @tparam T1 type of a * @tparam T2 type of b * @tparam T3 type of x @@ -183,9 +218,10 @@ inline auto inc_beta_power_series_kernel(const T1& a, const T2& b, const T3& x, * @param digamma_a digamma(a) * @param digamma_b digamma(b) * @param digamma_ab digamma(a + b) - * @return d/da I_x(a, b) for G1, d/db for G2, the - * pair (d/da, d/db) for BOTH; with flip_sign - * these are negated and the pair is returned as (-d/db, -d/da) + * @return d/da I_x(a, b) for G1, d/db for G2, and + * the pair (d/da, d/db) for BOTH. With flip_sign + * these are negated and the pair is returned as (-d/db, -d/da), which are + * the caller's (d/da, d/db). */ template @@ -194,27 +230,11 @@ inline auto inc_beta_continued_fraction_kernel(const T1& a, const T2& b, const T5& digamma_b, const T6& digamma_ab) { using ret_t = return_type_t; - constexpr bool calc_a = grads_requested != GRAD_RIB::G2; - constexpr bool calc_b = grads_requested != GRAD_RIB::G1; - // For integer b the partial numerator d_{2b} is 0, so the value and the a - // derivatives stop changing from that step on while the b derivative of - // d_{2b} is not 0. When b is an autodiff type its derivatives are carried - // through this loop, so the b derivative has to converge before stopping - // even when it is not requested, as track_b does in the power series. - constexpr bool track_b = calc_b || is_autodiff_v; constexpr int max_iterations = 300; const double tolerance = 3.0 * std::numeric_limits::epsilon(); const double small = std::sqrt(std::numeric_limits::min()); - // Modified Lentz state: C = A_n / A_{n-1} is the ratio of successive - // numerators, D = B_{n-1} / B_n the ratio of successive denominators and - // h = A_n / B_n the convergent, the product of every step's C * D. The - // derivatives are carried as logarithmic derivatives, c = (dC/da) / C, - // e = (dD/da) / D and r = (dh/da) / h for a and likewise for b, so each step - // adds c + e to r. - - // first step, with partial numerator d_1 = -(a + b) x / (a + 1) and C held - // at 1: D = 1 / (1 + d_1) and e = -(d d_1) D + // first step, with C held at 1 const auto ap1 = a + 1.0; const auto d_1 = -(a + b) * x / ap1; ret_t C = 1.0; @@ -230,12 +250,12 @@ inline auto inc_beta_continued_fraction_kernel(const T1& a, const T2& b, ret_t c_b = 0.0; ret_t e_b = 0.0; ret_t r_b = 0.0; - if constexpr (calc_a) { + if constexpr (grads_requested != GRAD_RIB::G2) { // d d_1 / da = -(x + d_1) / (a + 1) e_a = (x + d_1) / ap1 * D; r_a = e_a; } - if constexpr (track_b) { + if constexpr (grads_requested != GRAD_RIB::G1 || is_autodiff_v) { // d d_1 / db = -x / (a + 1) e_b = x / ap1 * D; r_b = e_b; @@ -245,12 +265,7 @@ inline auto inc_beta_continued_fraction_kernel(const T1& a, const T2& b, double dh_a_prev = value_of_rec(h) * value_of_rec(r_a); double dh_b_prev = value_of_rec(h) * value_of_rec(r_b); for (int n = 2; n <= 2 * max_iterations + 1; ++n) { - // partial numerators https://dlmf.nist.gov/8.17.E23 over the shared - // denominator (a + n - 1)(a + n): - // d_{2m} = m (b - m) x / ((a + 2m - 1)(a + 2m)) - // d_{2m+1} = -(a + m)(a + b + m) x / ((a + 2m)(a + 2m + 1)) - // d log(denominator) / da = (2 (a + n) - 1) / denominator, and for odd n - // the numerator's a derivative is -(2 (a + m) + b) x + // the partial numerator d_n and its b derivative dd_b const int m = n / 2; const bool odd = n % 2 == 1; const auto a_plus_n = a + n; @@ -279,7 +294,7 @@ inline auto inc_beta_continued_fraction_kernel(const T1& a, const T2& b, D_new = inv(D_new); // c_new = (dd - d c) / (C C_new) and e_new = -(dd + d e) D D_new const auto D_D_new = D * D_new; - if constexpr (calc_a) { + if constexpr (grads_requested != GRAD_RIB::G2) { auto dd_a = -d * (2.0 * a_plus_n - 1.0) / denominator; if (odd) { dd_a -= x * (2.0 * (a + m) + b) / denominator; @@ -288,7 +303,7 @@ inline auto inc_beta_continued_fraction_kernel(const T1& a, const T2& b, e_a = -(dd_a + d * e_a) * D_D_new; r_a += c_a + e_a; } - if constexpr (track_b) { + if constexpr (grads_requested != GRAD_RIB::G1 || is_autodiff_v) { c_b = (dd_b - d * c_b) / C_C_new; e_b = -(dd_b + d * e_b) * D_D_new; r_b += c_b + e_b; @@ -302,9 +317,9 @@ inline auto inc_beta_continued_fraction_kernel(const T1& a, const T2& b, const double dh_a_now = h_d * value_of_rec(r_a); const double dh_b_now = h_d * value_of_rec(r_b); if (fabs(value_of_rec(delta) - 1.0) < tolerance - && (!calc_a + && (grads_requested == GRAD_RIB::G2 || fabs(dh_a_now - dh_a_prev) <= tolerance * fabs(dh_a_now)) - && (!track_b + && ((grads_requested == GRAD_RIB::G1 && !is_autodiff_v) || fabs(dh_b_now - dh_b_prev) <= tolerance * fabs(dh_b_now))) { break; } @@ -340,20 +355,46 @@ inline auto inc_beta_continued_fraction_kernel(const T1& a, const T2& b, } // namespace internal /** - * Computes the requested gradients of the regularized incomplete beta - * function inc_beta(a, b, z) with respect to the arguments + * Computes the requested partial derivatives of the regularized incomplete + * beta function I_z(a, b) = inc_beta(a, b, z) with respect to * a and b. * - * The algorithm is the one of Boik and Robinson-Cox (1998), "Derivatives - * of the Incomplete Beta Function", Journal of Statistical Software 3(1), - * in the form implemented by TensorFlow Probability: a power series near - * the endpoints and a continued fraction elsewhere, both differentiated - * term by term, with the prefactor evaluated in log space and the - * symmetry relation I_z(a, b) = 1 - I_{1-z}(b, a) applied where the - * expansion converges faster on the other side. Only the requested - * derivatives are carried through the expansion. + * The method is that of Boik and Robinson-Cox (1998), "Derivatives of the + * Incomplete Beta Function", Journal of Statistical Software 3(1), as + * implemented in TensorFlow Probability. I_z(a, b) is expanded either as a + * power series or as a continued fraction, and the expansion is + * differentiated term by term. Only the requested derivatives are carried + * through it. The normalizing factor is evaluated in log space, so it stays + * finite where B(a, b) underflows. + * + * Either expansion can instead be applied to the right side of the symmetry + * relation I_z(a, b) = 1 - I_{1-z}(b, a), which converges faster for some + * arguments. The expansion is then evaluated for I_{1-z}(b, a), and the + * partials of I_z(a, b) are its partials negated, with the roles of a and b + * swapped. + * + * With mean = a / (a + b), the branch is chosen as follows. * - * The gradients are 0 at z = 0 and z = 1. + * - The power series for I_z(a, b) if z < mean, b z <= 1, and z <= 0.5. + * - The power series for I_{1-z}(b, a) if z >= mean, a (1 - z) <= 1, and + * z >= 0.5. + * - Otherwise, the continued fraction for I_z(a, b) if + * z < (a + 1) / (a + b + 2), and for I_{1-z}(b, a) if not. + * + * The terms of the power series shrink by a factor of about z per term + * (1 - z for the series of I_{1-z}(b, a)), so the series is limited to + * z <= 0.5. The condition b z <= 1 does not bound z by itself when b < 1. + * Without the limit, the series would take most of its 600 terms at z = 0.95, + * where the continued fraction needs tens of steps. + * + * The continued fraction converges rapidly for z < (a + 1) / (a + b + 2) + * (Numerical Recipes, 3rd ed., section 6.4). When z is at or above that + * point, 1 - z is at or below (b + 1) / (a + b + 2), the matching point for + * the fraction of I_{1-z}(b, a). The mode (a - 1) / (a + b - 2) is not used + * as the threshold because it exceeds 1 for b < 1. The fraction of I_z(a, b) + * would then be used for z near 1, where it needs hundreds of steps. + * + * The partials are 0 at z = 0 and z = 1. * * @tparam grads_requested G1 for the derivative with respect to * a, G2 for the derivative with respect to b, @@ -368,27 +409,30 @@ inline auto inc_beta_continued_fraction_kernel(const T1& a, const T2& b, * @param a a * @param b b * @param z z - * @param digammaA the value of digamma(a); not used for + * @param digammaA the value of digamma(a), not used for * G2 - * @param digammaB the value of digamma(b); not used for + * @param digammaB the value of digamma(b), not used for * G1 * @param digammaSum the value of digamma(a + b) - * @param betaAB accepted for interface compatibility and not used, the - * normalization is evaluated in log space - * @return d/da I_z(a, b) for G1, d/db for G2, the - * pair (d/da, d/db) for BOTH + * @param betaAB accepted for interface compatibility and not used. The + * normalization is evaluated in log space. + * @return d/da I_z(a, b) for G1, d/db for G2, and + * the pair (d/da, d/db) for BOTH */ template >>* = nullptr> + require_all_not_var_t* = nullptr> inline auto grad_reg_inc_beta(const T1& a, const T2& b, const T3& z, const T4& digammaA, const T5& digammaB, const T6& digammaSum, const T7& /* betaAB */) { // T7 is left out so that every branch returns the same type using ret_t = return_type_t; // the partials requested after the symmetry relation swaps a and b - constexpr GRAD_RIB swapped_grads = internal::swap_grad_rib(grads_requested); + constexpr GRAD_RIB swapped_grads + = grads_requested == GRAD_RIB::G1 + ? GRAD_RIB::G2 + : grads_requested == GRAD_RIB::G2 ? GRAD_RIB::G1 : GRAD_RIB::BOTH; + const double z_d = value_of_rec(z); if (z_d == 0.0 || z_d == 1.0) { if constexpr (grads_requested == GRAD_RIB::BOTH) { @@ -401,17 +445,10 @@ inline auto grad_reg_inc_beta(const T1& a, const T2& b, const T3& z, const double a_d = value_of_rec(a); const double b_d = value_of_rec(b); const double mean = a_d / (a_d + b_d); - // The series terms shrink by about z per step (1 - z for the flipped - // series), so it is limited to z <= 0.5 (z >= 0.5). For b < 1, b z <= 1 - // alone does not bound z, and at z = 0.95 the series takes most of its 600 - // terms where the continued fraction needs tens of steps. const bool use_power_series = (z_d < mean && b_d * z_d <= 1.0 && z_d <= 0.5) || (z_d >= mean && a_d * (1.0 - z_d) <= 1.0 && z_d >= 0.5); if (use_power_series) { - // For z >= a / (a + b) the series is evaluated for I_{1-z}(b, a) through - // the symmetry relation; its (d/da, d/db) are the negated (d/db, d/da) - // of I_z(a, b). if (z_d >= mean) { return internal::inc_beta_power_series_kernel( b, a, 1.0 - z, digammaB, digammaA, digammaSum); @@ -420,12 +457,6 @@ inline auto grad_reg_inc_beta(const T1& a, const T2& b, const T3& z, a, b, z, digammaA, digammaB, digammaSum); } } else { - // The continued fraction converges rapidly for z < (a + 1) / (a + b + 2) - // (Numerical Recipes, 3rd ed., section 6.4); otherwise it is evaluated - // for I_{1-z}(b, a) through the symmetry relation, where 1 - z is below - // the matching (b + 1) / (a + b + 2). The mode (a - 1) / (a + b - 2) is - // not a usable threshold: it exceeds 1 for b < 1, so the fraction would - // run at z near 1, where it needs hundreds of steps. if (z_d >= (a_d + 1.0) / (a_d + b_d + 2.0)) { return internal::inc_beta_continued_fraction_kernel( b, a, 1.0 - z, digammaB, digammaA, digammaSum); @@ -440,11 +471,11 @@ inline auto grad_reg_inc_beta(const T1& a, const T2& b, const T3& z, /** * Computes the gradients of the regularized incomplete beta * function. Specifically, this function computes gradients of - * inc_beta(a, b, z), with respect to the arguments + * inc_beta(a, b, z) with respect to the arguments * a and b. * - * Writes both partials; see the overload templated on - * GRAD_RIB to compute only one. + * Writes both partials. To compute only one, use the overload templated on + * GRAD_RIB. * * @tparam T type of arguments * @param[out] g1 partial derivative of inc_beta(a, b, z) @@ -457,9 +488,9 @@ inline auto grad_reg_inc_beta(const T1& a, const T2& b, const T3& z, * @param[in] digammaA the value of digamma(a) * @param[in] digammaB the value of digamma(b) * @param[in] digammaSum the value of digamma(a + b) - * @param[in] betaAB the value of beta(a, b); accepted for - * interface compatibility and not used, the normalization is evaluated in - * log space + * @param[in] betaAB the value of beta(a, b), accepted for + * interface compatibility and not used. The normalization is evaluated in + * log space. */ template inline void grad_reg_inc_beta(T& g1, T& g2, const T& a, const T& b, const T& z, diff --git a/stan/math/prim/fun/inc_beta_dda.hpp b/stan/math/prim/fun/inc_beta_dda.hpp index f5c7d4efad1..bf5d08258db 100644 --- a/stan/math/prim/fun/inc_beta_dda.hpp +++ b/stan/math/prim/fun/inc_beta_dda.hpp @@ -12,11 +12,11 @@ namespace stan { namespace math { /** - * Returns the partial derivative of the regularized - * incomplete beta function, I_{z}(a, b) with respect to a. + * Returns the partial derivative of the regularized incomplete beta function + * I_z(a, b) with respect to a. * - * Evaluated by grad_reg_inc_beta, computing only the derivative - * with respect to a. + * This calls grad_reg_inc_beta, which computes + * only the derivative with respect to a. * * @tparam T1 type of a * @tparam T2 type of b diff --git a/stan/math/prim/fun/inc_beta_ddb.hpp b/stan/math/prim/fun/inc_beta_ddb.hpp index ee42baa91e3..1e60ec1df01 100644 --- a/stan/math/prim/fun/inc_beta_ddb.hpp +++ b/stan/math/prim/fun/inc_beta_ddb.hpp @@ -8,11 +8,11 @@ namespace stan { namespace math { /** - * Returns the partial derivative of the regularized - * incomplete beta function, I_{z}(a, b) with respect to b. + * Returns the partial derivative of the regularized incomplete beta function + * I_z(a, b) with respect to b. * - * Evaluated by grad_reg_inc_beta, computing only the derivative - * with respect to b. + * This calls grad_reg_inc_beta, which computes + * only the derivative with respect to b. * * @tparam T1 type of a * @tparam T2 type of b diff --git a/stan/math/prim/fun/inc_beta_ddz.hpp b/stan/math/prim/fun/inc_beta_ddz.hpp index 762e46bce9d..4b955b2a189 100644 --- a/stan/math/prim/fun/inc_beta_ddz.hpp +++ b/stan/math/prim/fun/inc_beta_ddz.hpp @@ -38,8 +38,8 @@ inline auto inc_beta_ddz(T1 a, T2 b, T3 z) { * Returns the partial derivative of the regularized * incomplete beta function, I_{z}(a, b) with respect to z. * - * Overload for double arguments, evaluated by Boost. Where the density is - * infinite (z = 0 with a < 1, or z = 1 with b < 1) it returns infinity + * This overload for double arguments is evaluated by Boost. Where the density + * is infinite (z = 0 with a < 1, or z = 1 with b < 1) it returns infinity * instead of throwing. * * @param a first argument diff --git a/stan/math/rev/fun/grad_reg_inc_beta.hpp b/stan/math/rev/fun/grad_reg_inc_beta.hpp index 90be0e2a247..daf151d46f0 100644 --- a/stan/math/rev/fun/grad_reg_inc_beta.hpp +++ b/stan/math/rev/fun/grad_reg_inc_beta.hpp @@ -22,38 +22,68 @@ namespace internal { /** * Hessian-vector product H v of the regularized incomplete beta function - * I_x(a, b), H its Hessian in (a, b), by the power series + * I_x(a, b), where H is the Hessian of I_x(a, b) in (a, b), computed from its + * power series. * - * I_x(a, b) = N S, N = x^a / B(a, b), S = sum_n P_n / (a + n), - * P_n = prod_{k <= n} (k - b) x / k. + * As in internal::inc_beta_power_series_kernel, the series is * - * Differentiating I = N S twice, + * I_x(a, b) = N S, N = x^a / B(a, b), S = sum_{n >= 0} t_n, + * t_n = P_n / (a + n), P_0 = 1, P_n = P_{n-1} (n - b) x / n. + * + * Write grad for the gradient in (a, b) and D_v for the derivative along v. + * Differentiating I = N S twice and applying the result to v gives * * H v = N (D_v grad S + grad S D_v log N + grad log N D_v S - * + S (grad log N D_v log N + D_v grad log N)), + * + S (grad log N D_v log N + D_v grad log N)). * - * and a diagonal entry is + * When only one shape j is requested, the product needs only the diagonal + * entry * * H_jj = N (d2S + 2 dlog_N dS + S (dlog_N^2 + d2log_N)), * - * d and d2 the first and second derivatives in shape j. The a derivatives - * of each term t = P_n / (a + n) are closed form, and its b derivatives need - * those of P_n, carried by recurrence. + * where d and d2 are the first and second derivatives in that shape. + * + * The derivatives of S are the sums of those of its terms. The a derivatives + * of a term are closed form. Its b derivatives come from those of P_n, which + * follow from the recurrence for P_n: + * + * dt/da = -t / (a + n), d2t/da2 = 2 t / (a + n)^2, + * dt/db = P'_n / (a + n), d2t/dadb = -(dt/db) / (a + n), + * d2t/db2 = P''_n / (a + n). + * + * For n = 0 the term is 1 / a, with a derivatives -1 / a^2 and 2 / a^3, and + * it does not depend on b. + * + * With psi the digamma function and psi1 the trigamma function, the + * derivatives of log N = a log x - log B(a, b) are + * + * grad log N = (log x - psi(a) + psi(a + b), psi(a + b) - psi(b)), + * Hess log N = [[psi1(a + b) - psi1(a), psi1(a + b)], + * [psi1(a + b), psi1(a + b) - psi1(b)]]. * - * Names: for G1 and G2 only the requested shape - * varies, and dX and d2X are the first and second derivatives of X in it. For - * BOTH, X_a and X_b are the partials of X and Dv_X its - * derivative along v. + * psi1 is taken from Boost. stan::math::trigamma(double) is accurate only to + * about 1e-9 for arguments below 5, and psi1(a + b) appears alone in H_ab, + * where that error does not cancel. + * + * The loop stops once the terms are below eps / a in magnitude, or after 600 + * terms. For integer b, t_n and its a derivatives are 0 for n >= b, but its b + * derivatives are not. G2 and BOTH therefore also + * wait for the b derivatives of the terms to be small. + * + * The variable names follow two conventions. For G1 and + * G2 only the requested shape varies, and dX and d2X are the + * first and second derivatives of X in that shape. For BOTH, + * X_a and X_b are the partials of X, and Dv_X is the derivative of X along v. * * @tparam grads G1 for H_aa v_a, G2 for H_bb v_b, * BOTH for H v * @param a first shape * @param b second shape * @param x upper limit, in (0, 1) - * @param v_a a component of v; not used for G2 - * @param v_b b component of v; not used for G1 - * @param digamma_a digamma(a); not used for G2 - * @param digamma_b digamma(b); not used for G1 + * @param v_a a component of v, not used for G2 + * @param v_b b component of v, not used for G1 + * @param digamma_a digamma(a), not used for G2 + * @param digamma_b digamma(b), not used for G1 * @param digamma_ab digamma(a + b) * @return H_aa v_a for G1, H_bb v_b for G2, * ((H v)_a, (H v)_b) for BOTH @@ -66,11 +96,8 @@ inline auto inc_beta_hvp_power_series(double a, double b, double x, double v_a, const double tolerance = std::numeric_limits::epsilon() / a; const double log_x = std::log(x); const double N = std::exp(a * log_x - lbeta(a, b)); - // boost's trigamma: stan::math::trigamma(double) is good to about 1e-9 below - // 5, and psi1(a + b) enters H_ab alone, where that error does not cancel const double trigamma_ab = boost::math::trigamma(a + b, boost_policy_t<>()); - // n = 0: P_0 = 1 and the term is 1 / a, with a derivatives -1 / a^2 and - // 2 / a^3, and it does not depend on b + // n = 0: P_0 = 1 and the term is 1 / a double product = 1.0; double S = 1.0 / a; if constexpr (grads == GRAD_RIB::G1) { @@ -78,7 +105,6 @@ inline auto inc_beta_hvp_power_series(double a, double b, double x, double v_a, double d2S = 2.0 * S / (a * a); for (int n = 1; n <= max_iterations; ++n) { product *= (n - b) * x / n; - // dt/da = -t / (a + n), d2t/da2 = 2 t / (a + n)^2 const double inv_apn = 1.0 / (a + n); const double term = product * inv_apn; S += term; @@ -113,8 +139,6 @@ inline auto inc_beta_hvp_power_series(double a, double b, double x, double v_a, S += term; dS += dterm; d2S += d2term; - // for integer b the term is 0 from n = b on while its b derivatives are - // not, so every carried series has to be small before stopping if (std::fabs(term) <= tolerance && std::fabs(dterm) <= tolerance && std::fabs(d2term) <= tolerance) { break; @@ -142,8 +166,6 @@ inline auto inc_beta_hvp_power_series(double a, double b, double x, double v_a, const double inv_apn = 1.0 / (a + n); const double term = product * inv_apn; const double term_b = product_b * inv_apn; - // d2t/da2 = 2 t / (a + n)^2, d2t/dadb = -t_b / (a + n), - // d2t/db2 = P''_n / (a + n) const double term_aa = 2.0 * term * inv_apn * inv_apn; const double term_ab = -term_b * inv_apn; const double term_bb = product_bb * inv_apn; @@ -152,8 +174,6 @@ inline auto inc_beta_hvp_power_series(double a, double b, double x, double v_a, S_b += term_b; Dv_S_a += v_a * term_aa + v_b * term_ab; Dv_S_b += v_a * term_ab + v_b * term_bb; - // for integer b the term is 0 from n = b on while its b derivatives are - // not, so every carried series has to be small before stopping if (std::fabs(term) <= tolerance && std::fabs(term_b) <= tolerance && std::fabs(term_bb) <= tolerance) { break; @@ -163,8 +183,6 @@ inline auto inc_beta_hvp_power_series(double a, double b, double x, double v_a, const double log_N_b = digamma_ab - digamma_b; const double Dv_log_N = log_N_a * v_a + log_N_b * v_b; const double Dv_S = S_a * v_a + S_b * v_b; - // the Hessian of log N is - // [[psi1(a+b) - psi1(a), psi1(a+b)], [psi1(a+b), psi1(a+b) - psi1(b)]] const double Dv_log_N_a = (trigamma_ab - boost::math::trigamma(a, boost_policy_t<>())) * v_a + trigamma_ab * v_b; @@ -182,21 +200,36 @@ inline auto inc_beta_hvp_power_series(double a, double b, double x, double v_a, /** * Hessian-vector product H v of the regularized incomplete beta function - * I_x(a, b), H its Hessian in (a, b), by the continued fraction - * https://dlmf.nist.gov/8.17.E22, I_x(a, b) = Q h with - * Q = x^a (1 - x)^b / (a B(a, b)) and h the continued fraction, evaluated - * with the modified Lentz method. Differentiating I = Q h twice, + * I_x(a, b), where H is the Hessian of I_x(a, b) in (a, b), computed from its + * continued fraction. + * + * As in internal::inc_beta_continued_fraction_kernel, I_x(a, b) = Q h with + * Q = x^a (1 - x)^b / (a B(a, b)) and h the continued fraction + * https://dlmf.nist.gov/8.17.E22, which is evaluated with the modified Lentz + * method. Write grad for the gradient in (a, b), Hess for the Hessian, and + * D_v for the derivative along v. Differentiating I = Q h twice gives + * + * H = I (grad log I grad log I^T + Hess log Q + Hess log h), + * + * with grad log I = grad log Q + grad log h. With psi the digamma function + * and psi1 the trigamma function, * - * H = I (grad log I grad log I^T + Hess log Q + Hess log h). + * grad log Q = (log x - 1 / a - psi(a) + psi(a + b), + * log(1 - x) - psi(b) + psi(a + b)), + * Hess log Q = [[psi1(a + b) - psi1(a) + 1 / a^2, psi1(a + b)], + * [psi1(a + b), psi1(a + b) - psi1(b)]], * - * log h is the sum over the Lentz steps of log C and log D, so its - * derivatives are sums of theirs. With d the partial numerator, one step - * takes C' = 1 + d / C and D' = 1 / (1 + d D), and + * with psi1 taken from Boost, as in the power series. + * + * Each Lentz step multiplies h by C D, so log h is the sum of log C and + * log D over the steps, and its derivatives are the sums of theirs. With d + * the partial numerator of the step, the step sets C' = 1 + d / C and + * D' = 1 / (1 + d D). Their log derivatives are * * grad log C' = (grad d - d grad log C) / (C + d), - * grad log D' = -(grad d + d grad log D) D D'. + * grad log D' = -(grad d + d grad log D) D D', * - * Differentiating those once more along a direction v, + * using C C' = C + d. Differentiating these once more along v gives * * D_v grad log C' = (Hess d v - D_v d grad log C - d D_v grad log C) * / (C + d) - grad log C' D_v log(C C'), @@ -205,23 +238,45 @@ inline auto inc_beta_hvp_power_series(double a, double b, double x, double v_a, * + grad log D' D_v log(D D'), * D_v log(D D') = D_v log D - D D' (D_v d + d D_v log D). * - * G1 and G2 take v = (1, 0) and (0, 1), so their - * D_v are second derivatives in one shape. + * The code applies each of the two second-derivative formulas in three parts. + * It first computes the part that uses only values from before the step. It + * then overwrites grad log C (or grad log D) with its new value, and last it + * adds the final term, which uses that new value. + * + * G1 and G2 take v = (1, 0) and v = (0, 1), so D_v + * is the derivative in the requested shape, and the result is scaled by v_a + * or v_b at the end. + * + * The partial numerators https://dlmf.nist.gov/8.17.E23 are d_n = num / den, + * with den = (a + n - 1)(a + n) and + * + * n = 2m: num = m (b - m) x, num_b = m x, + * n = 2m + 1: num = -(a + m)(a + b + m) x, num_a = -(2 (a + m) + b) x, + * num_b = -(a + m) x, num_aa = -2 x, num_ab = -x. + * + * Every other derivative of num is 0, so d_bb = 0. The derivatives of d_n + * follow from the quotient rule, with den_a = 2 (a + n) - 1 and den_aa = 2. + * n = 1 is the odd case with m = 0. The first step holds C at 1, so + * D = 1 / (1 + d_1). + * + * The test for stopping runs after each odd n. The loop stops once C D is + * within 3 eps of 1 and the second derivatives of log h have changed by at + * most a relative 3 eps since the previous test, or after n = 601. * - * Names: for G1 and G2 only the requested shape - * varies, and dX and d2X are the first and second derivatives of X in it. For - * BOTH, X_a and X_b are the partials of X and Dv_X its - * derivative along v. + * The variable names follow the same two conventions as in the power series. + * For G1 and G2, dX and d2X are the first and + * second derivatives of X in the requested shape. For BOTH, + * X_a and X_b are the partials of X, and Dv_X is the derivative of X along v. * * @tparam grads G1 for H_aa v_a, G2 for H_bb v_b, * BOTH for H v * @param a first shape * @param b second shape * @param x upper limit, in (0, 1) - * @param v_a a component of v; not used for G2 - * @param v_b b component of v; not used for G1 - * @param digamma_a digamma(a); not used for G2 - * @param digamma_b digamma(b); not used for G1 + * @param v_a a component of v, not used for G2 + * @param v_b b component of v, not used for G1 + * @param digamma_a digamma(a), not used for G2 + * @param digamma_b digamma(b), not used for G1 * @param digamma_ab digamma(a + b) * @return H_aa v_a for G1, H_bb v_b for G2, * ((H v)_a, (H v)_b) for BOTH @@ -240,15 +295,7 @@ inline auto inc_beta_hvp_continued_fraction(double a, double b, double x, = std::exp(a * log_x + b * log1m_x - std::log(a) - lbeta(a, b)); // boost's trigamma, as in the power series const double trigamma_ab = boost::math::trigamma(a + b, boost_policy_t<>()); - - // Partial numerator d_n = num / den over den = (a + n - 1)(a + n), - // https://dlmf.nist.gov/8.17.E23, with its a and b derivatives: - // n = 2m: num = m (b - m) x, num_b = m x - // n = 2m + 1: num = -(a + m)(a + b + m) x, num_a = -(2 (a + m) + b) x, - // num_b = -(a + m) x, num_aa = -2 x, num_ab = -x - // n = 1 is the odd case with m = 0, and d_bb = 0 for every n. - - // first step, n = 1, with C held at 1: D = 1 / (1 + d_1) + // first step, n = 1, with C held at 1 const double ap1 = a + 1.0; const double d_1 = -(a + b) * x / ap1; double D = 1.0 + d_1; @@ -288,9 +335,7 @@ inline auto inc_beta_hvp_continued_fraction(double a, double b, double x, const double dd_n = (dnum - d_n * dden) * inv_den; const double d2d_n = (d2num - 2.0 * dd_n * dden - 2.0 * d_n) * inv_den; - // one step of the modified Lentz method. In each half, the second - // derivative takes its old-step part, the first derivative updates in - // place, then scales dlog(C C_new) or dlog(D D_new). + // one step of the modified Lentz method // C_new = 1 + d_n / C, with C C_new = C + d_n double C_new = 1.0 + d_n / C; double C_C_new = C + d_n; @@ -461,9 +506,7 @@ inline auto inc_beta_hvp_continued_fraction(double a, double b, double x, const double Dv_d_n_a = d_n_aa * v_a + d_n_ab * v_b; const double Dv_d_n_b = d_n_ab * v_a; - // one step of the modified Lentz method. In each half, the Dv_ - // derivatives take their old-step part, the partials update in place, - // then scale D_v log(C C_new) or D_v log(D D_new). + // one step of the modified Lentz method // C_new = 1 + d_n / C, with C C_new = C + d_n double C_new = 1.0 + d_n / C; double C_C_new = C + d_n; @@ -540,25 +583,30 @@ inline auto inc_beta_hvp_continued_fraction(double a, double b, double x, /** * Hessian-vector product H v of the regularized incomplete beta function - * I_z(a, b), H its Hessian in (a, b): the reverse pass of - * (d/da I_z(a, b), d/db I_z(a, b)) with output adjoints v. + * I_z(a, b), where H is the Hessian of I_z(a, b) in (a, b). When v holds the + * adjoints of (d/da I_z(a, b), d/db I_z(a, b)), H v is what the reverse pass + * adds to the adjoints of a and b. * - * The branch choice follows grad_reg_inc_beta. Under the symmetry relation - * I_z(a, b) = 1 - I_{1-z}(b, a), H v = -swap(H' swap(v)), H' the Hessian of - * I_{1-z}(b, a) in its own arguments, so H_aa = -H'_bb and H_bb = -H'_aa. + * The power series or the continued fraction, and whether to apply the + * symmetry relation, are chosen by the same rule as in grad_reg_inc_beta. + * Under the symmetry relation I_z(a, b) = 1 - I_{1-z}(b, a), let H' be the + * Hessian of I_{1-z}(b, a) in its own arguments (b, a), and let P swap the + * two components of a vector. Then H = -P H' P, so H v = -P H' (P v). In + * particular H_aa = -H'_bb and H_bb = -H'_aa. * * @tparam grads G1 for H_aa v_a, G2 for H_bb v_b, * BOTH for H v * @param a first shape * @param b second shape * @param z upper limit, in [0, 1] - * @param v_a a component of v; not used for G2 - * @param v_b b component of v; not used for G1 - * @param digamma_a digamma(a); not used for G2 - * @param digamma_b digamma(b); not used for G1 + * @param v_a a component of v, not used for G2 + * @param v_b b component of v, not used for G1 + * @param digamma_a digamma(a), not used for G2 + * @param digamma_b digamma(b), not used for G1 * @param digamma_ab digamma(a + b) * @return H_aa v_a for G1, H_bb v_b for G2, - * ((H v)_a, (H v)_b) for BOTH; 0 at z = 0 and z = 1 + * ((H v)_a, (H v)_b) for BOTH. The result is 0 at z = 0 and + * z = 1. */ template inline std::conditional_t, @@ -569,7 +617,10 @@ inc_beta_hvp(double a, double b, double z, double v_a, double v_b, return {}; } // the symmetry relation swaps which shape is requested - constexpr GRAD_RIB swapped_grads = swap_grad_rib(grads); + constexpr GRAD_RIB swapped_grads + = grads == GRAD_RIB::G1 + ? GRAD_RIB::G2 + : grads == GRAD_RIB::G2 ? GRAD_RIB::G1 : GRAD_RIB::BOTH; const double mean = a / (a + b); // the same branch rule as grad_reg_inc_beta const bool use_power_series @@ -608,29 +659,42 @@ inc_beta_hvp(double a, double b, double z, double v_a, double v_b, } // namespace internal /** - * The requested gradients of the regularized incomplete beta function - * inc_beta(a, b, z) with respect to a and - * b, for var arguments. + * Computes the requested partial derivatives of the regularized incomplete + * beta function I_z(a, b) = inc_beta(a, b, z) with respect to + * a and b, for var arguments. * - * The values are computed in double by the prim overload. The reverse pass - * does not differentiate through its series or continued fraction: with - * output adjoints v, the a and b adjoints are H v for H the Hessian of - * I_z(a, b) in (a, b), computed directly by internal::inc_beta_hvp, and the - * z adjoint is the closed form + * The values come from the prim overload, evaluated in double. The reverse + * pass does not differentiate through that series or continued fraction. + * Instead, with v the adjoints of the outputs, it adds H v to the adjoints of + * a and b, where H is the Hessian of I_z(a, b) in (a, b), computed directly + * by internal::inc_beta_hvp. To the adjoint of z it adds v times the + * closed-form mixed derivatives * * d/dz d/da I_z(a, b) = inc_beta_ddz(a, b, z) (log z - psi(a) + psi(a+b)), * d/dz d/db I_z(a, b) = inc_beta_ddz(a, b, z) (log1p(-z) - psi(b) + - * psi(a+b)). + * psi(a+b)), + * + * where psi is the digamma function. + * + * The digamma arguments are passed only so that the values are cheap to + * compute, and they must equal digamma(a), digamma(b), and digamma(a + b). + * They get no adjoint. The a and b adjoints above are total derivatives, so + * they already include any dependence on a and b that runs through the + * digamma values. + * + * A one-sided request still needs the cross term of H when both shapes are + * vars. For G1 the adjoint of b is H_ab v_a, which the + * BOTH product gives with v_b = 0. That product needs + * digamma(b), which a G1 call does not receive, so the reverse + * pass computes it when b is a var. When only a is a var, H_aa v_a does not + * need digamma(b). G2 is the same with a and b swapped. * - * The digamma arguments only make the values cheap and must equal - * digamma(a), digamma(b) and digamma(a + b); their adjoints are 0 and the - * a and b adjoints above are total derivatives. The cross derivative needs - * the digamma value a one-sided call does not take, so the reverse pass of - * G1 computes digamma(b) when b is a var, and that of - * G2 digamma(a) when a is a var. betaAB is not used. + * For BOTH there are two outputs, and one reverse pass callback + * reads both of their adjoints. It runs after every use of the outputs has + * added into those adjoints. * - * Each argument is a var or arithmetic; mixing with - * fvar goes to the prim overload. + * Each argument is a var or arithmetic. Calls that mix in + * fvar go to the prim overload. * * @tparam grads_requested G1 for the derivative with respect to * a, G2 for the derivative with respect to b, @@ -645,19 +709,18 @@ inc_beta_hvp(double a, double b, double z, double v_a, double v_b, * @param a a * @param b b * @param z z - * @param digammaA the value of digamma(a); not used for + * @param digammaA the value of digamma(a), not used for * G2 - * @param digammaB the value of digamma(b); not used for + * @param digammaB the value of digamma(b), not used for * G1 * @param digammaSum the value of digamma(a + b) - * @param betaAB not used - * @return d/da I_z(a, b) for G1, d/db for G2, the - * pair (d/da, d/db) for BOTH + * @param betaAB accepted for interface compatibility and not used + * @return d/da I_z(a, b) for G1, d/db for G2, and + * the pair (d/da, d/db) for BOTH */ template >>* = nullptr> + require_any_var_t* = nullptr> inline auto grad_reg_inc_beta(const T1& a, const T2& b, const T3& z, const T4& digammaA, const T5& digammaB, const T6& digammaSum, const T7& /* betaAB */) { @@ -674,8 +737,6 @@ inline auto grad_reg_inc_beta(const T1& a, const T2& b, const T3& z, const double b_val = value_of(b); const double z_val = value_of(z); if constexpr (is_var_v && is_var_v) { - // the b adjoint is the cross term H_ab v_a, which needs both - // sides and digamma(b), which the G1 call does not take const auto [hv_a, hv_b] = internal::inc_beta_hvp( a_val, b_val, z_val, vi.adj(), 0.0, digamma_a, digamma(b_val), digamma_ab); @@ -709,8 +770,6 @@ inline auto grad_reg_inc_beta(const T1& a, const T2& b, const T3& z, const double b_val = value_of(b); const double z_val = value_of(z); if constexpr (is_var_v && is_var_v) { - // the a adjoint is the cross term H_ab v_b, which needs both - // sides and digamma(a), which the G2 call does not take const auto [hv_a, hv_b] = internal::inc_beta_hvp( a_val, b_val, z_val, 0.0, vi.adj(), digamma(a_val), digamma_b, digamma_ab); @@ -738,8 +797,6 @@ inline auto grad_reg_inc_beta(const T1& a, const T2& b, const T3& z, const auto [g1_val, g2_val] = grad_reg_inc_beta( value_of(a), value_of(b), value_of(z), digamma_a, digamma_b, digamma_ab, 0.0); - // Both outputs feed one reverse pass, which runs after everything that - // uses them has added into their adjoints. var g1 = g1_val; var g2 = g2_val; reverse_pass_callback([a, b, z, g1, g2, digamma_a, digamma_b, diff --git a/stan/math/rev/fun/inc_beta.hpp b/stan/math/rev/fun/inc_beta.hpp index 3c4481b81ec..037f3f07433 100644 --- a/stan/math/rev/fun/inc_beta.hpp +++ b/stan/math/rev/fun/inc_beta.hpp @@ -16,7 +16,7 @@ namespace math { * The regularized incomplete beta function I_z(a, b) for var * arguments. Only the partials of the var arguments are * computed in the reverse pass. Each argument is a var or - * arithmetic; mixing with fvar is not supported. + * arithmetic. Mixing with fvar is not supported. * * @tparam T1 type of a * @tparam T2 type of b diff --git a/test/unit/math/mix/fun/grad_reg_inc_beta_test.cpp b/test/unit/math/mix/fun/grad_reg_inc_beta_test.cpp index eeadf6281b6..7831ab106ee 100644 --- a/test/unit/math/mix/fun/grad_reg_inc_beta_test.cpp +++ b/test/unit/math/mix/fun/grad_reg_inc_beta_test.cpp @@ -3,15 +3,17 @@ namespace { -// Points in each region of grad_reg_inc_beta: -// (1, 1, 0.4), (2, 3, 0.25), (0.5, 0.5, 0.3): power series -// (10, 1, 0.5): power series, integer b -// (2, 3, 0.6): power series through the symmetry relation -// (20, 20, 0.3), (5, 30, 0.1): continued fraction -// (20, 20, 0.7), (5, 30, 0.3): continued fraction through the symmetry -// relation -// (10, 1, 0.7), (4, 4, 0.5): continued fraction, integer b, where the -// partial numerator d_{2b} is 0 but its b derivative is not +/** + * Runs expect_ad at points in each region of grad_reg_inc_beta: + * (1, 1, 0.4), (2, 3, 0.25), (0.5, 0.5, 0.3): power series + * (10, 1, 0.5): power series, integer b + * (2, 3, 0.6): power series through the symmetry relation + * (20, 20, 0.3), (5, 30, 0.1): continued fraction + * (20, 20, 0.7), (5, 30, 0.3): continued fraction through the symmetry + * relation + * (10, 1, 0.7), (4, 4, 0.5): continued fraction, integer b, where the + * partial numerator d_{2b} is 0 but its b derivative is not + */ template void expect_ad_grad_reg_inc_beta(const F& f) { stan::test::expect_ad(f, 1.0, 1.0, 0.4); @@ -70,9 +72,11 @@ TEST(mathMixScalFun, grad_reg_inc_beta_student) { expect_ad_grad_reg_inc_beta(f); } +/** + * The overload writing both partials through out-parameters takes every + * argument as the same type. + */ TEST(mathMixScalFun, grad_reg_inc_beta_out_params) { - // the overload writing both partials through out-parameters takes every - // argument as the same type auto f = [](const auto& a, const auto& b, const auto& z) { using stan::math::digamma; using T = stan::return_type_t; diff --git a/test/unit/math/mix/fun/inc_beta_test.cpp b/test/unit/math/mix/fun/inc_beta_test.cpp index b90ea81b74c..8805333dddd 100644 --- a/test/unit/math/mix/fun/inc_beta_test.cpp +++ b/test/unit/math/mix/fun/inc_beta_test.cpp @@ -28,9 +28,11 @@ TEST(mathMixScalFun, inc_beta) { stan::test::expect_ad(f, nan, nan, nan); } +/** + * At z = 1 with b < 1 the density is infinite. The forward and reverse + * sweeps of fvar must not throw there. + */ TEST(mathMixScalFun, inc_beta_upper_bound_infinite_density) { - // At z = 1 with b < 1 the density is infinite; the forward and reverse - // sweeps of fvar must not throw there. using stan::math::fvar; using stan::math::var; // grad() runs every node on the stack; start from an empty one diff --git a/test/unit/math/prim/fun/grad_reg_inc_beta_test.cpp b/test/unit/math/prim/fun/grad_reg_inc_beta_test.cpp index c84a259b796..19b8e5c95dd 100644 --- a/test/unit/math/prim/fun/grad_reg_inc_beta_test.cpp +++ b/test/unit/math/prim/fun/grad_reg_inc_beta_test.cpp @@ -5,9 +5,13 @@ namespace { -// Both shape derivatives of I_z(a, b) against a high-precision reference -// (mpmath, 80 digits, density-form quadrature checked against numerical -// differentiation at 110 digits). +/** + * Both shape derivatives of I_z(a, b) against a high-precision reference + * (mpmath, 80 digits, density-form quadrature checked against numerical + * differentiation at 110 digits). Each one-sided request is also checked. + * It carries only its own derivative through the expansion, with the other + * digamma unused. + */ void expect_grads(double a, double b, double z, double dda, double ddb, double rtol) { using stan::math::digamma; @@ -23,8 +27,6 @@ void expect_grads(double a, double b, double z, double dda, double ddb, EXPECT_NEAR(ddb, g2, rtol * std::fabs(ddb)) << "a=" << a << " b=" << b << " z=" << z; - // each one-sided request carries only its own derivative through the - // expansion, with the other digamma unused using stan::math::GRAD_RIB; const double only_g1 = stan::math::grad_reg_inc_beta( a, b, z, digamma(a), 0.0, digamma(a + b), 0.0); @@ -83,9 +85,11 @@ TEST(grad_reg_inc_beta, prefactor_underflow) { 4.5390666074606649e-138, 1e-12); } +/** + * For integer b the power series of the value terminates but the derivative + * series does not. The stop test must check both. + */ TEST(grad_reg_inc_beta, integer_b) { - // for integer b the power series of the value terminates but the - // derivative series does not; the stop test must watch both expect_grads(10, 1, 0.5, -0.00067690154351557159, 0.00226574342615196, 1e-12); expect_grads(1, 100, 0.01, -0.42882413065860014, 0.0036787479630394138, 1e-12); diff --git a/test/unit/math/prim/fun/inc_beta_dda_test.cpp b/test/unit/math/prim/fun/inc_beta_dda_test.cpp index 953ee47b661..934a73c17f0 100644 --- a/test/unit/math/prim/fun/inc_beta_dda_test.cpp +++ b/test/unit/math/prim/fun/inc_beta_dda_test.cpp @@ -4,8 +4,10 @@ namespace { -// d/da of I_z(a, b) against a high-precision reference (mpmath, 80 digits, -// density-form quadrature checked against numerical differentiation). +/** + * d/da of I_z(a, b) against a high-precision reference (mpmath, 80 digits, + * density-form quadrature checked against numerical differentiation). + */ void expect_dda(double a, double b, double z, double expected, double rtol) { using stan::math::digamma; using stan::math::inc_beta_dda; diff --git a/test/unit/math/prim/fun/inc_beta_ddb_test.cpp b/test/unit/math/prim/fun/inc_beta_ddb_test.cpp index dfc8d50cfb0..ec6c28928cc 100644 --- a/test/unit/math/prim/fun/inc_beta_ddb_test.cpp +++ b/test/unit/math/prim/fun/inc_beta_ddb_test.cpp @@ -4,9 +4,11 @@ namespace { -// d/db of I_z(a, b) against a high-precision reference (mpmath, 80 digits, -// density-form quadrature checked against numerical differentiation). -// The fourth argument of inc_beta_ddb is digamma(b). +/** + * d/db of I_z(a, b) against a high-precision reference (mpmath, 80 digits, + * density-form quadrature checked against numerical differentiation). + * The fourth argument of inc_beta_ddb is digamma(b). + */ void expect_ddb(double a, double b, double z, double expected, double rtol) { using stan::math::digamma; using stan::math::inc_beta_ddb; diff --git a/test/unit/math/rev/fun/inc_beta_test.cpp b/test/unit/math/rev/fun/inc_beta_test.cpp index 9d83916ebaf..a66e629482a 100644 --- a/test/unit/math/rev/fun/inc_beta_test.cpp +++ b/test/unit/math/rev/fun/inc_beta_test.cpp @@ -6,8 +6,10 @@ namespace { -// Shape gradients of inc_beta against a high-precision reference (mpmath, -// 80 digits), at points where a finite-difference check is not informative. +/** + * Shape gradients of inc_beta against a high-precision reference (mpmath, + * 80 digits), at points where a finite-difference check is not informative. + */ void expect_inc_beta_grads(double a, double b, double z, double dda, double ddb, double rtol) { using stan::math::var; @@ -42,9 +44,11 @@ TEST_F(AgradRev, inc_beta_shape_gradients) { 0.00226574342615196, 1e-12); } +/** + * At z = 1, I is 1 for all shapes, so the shape gradients are 0. With b < 1 + * the density, which is the z gradient, is infinite. + */ TEST_F(AgradRev, inc_beta_upper_bound_infinite_density) { - // At z = 1, I is 1 for all shapes, so the shape gradients are 0; with - // b < 1 the density, the z gradient, is infinite. using stan::math::var; var a = 1.05; var b = 0.35; diff --git a/test/unit/math/rev/prob/beta_cdf_log_test.cpp b/test/unit/math/rev/prob/beta_cdf_log_test.cpp index f4ab7ed85c3..ec46553b456 100644 --- a/test/unit/math/rev/prob/beta_cdf_log_test.cpp +++ b/test/unit/math/rev/prob/beta_cdf_log_test.cpp @@ -6,8 +6,10 @@ namespace { -// Shape gradients of beta_lcdf against a high-precision reference: the -// reference is (d/da I) / I and (d/db I) / I from mpmath at 80 digits. +/** + * Shape gradients of beta_lcdf against a high-precision reference. The + * reference is (d/da I) / I and (d/db I) / I from mpmath at 80 digits. + */ void expect_beta_lcdf_grads(double y, double alpha, double beta, double dalpha, double dbeta, double rtol) { using stan::math::var; diff --git a/test/unit/math/rev/prob/beta_cdf_test.cpp b/test/unit/math/rev/prob/beta_cdf_test.cpp index ad01be2862f..3f842144099 100644 --- a/test/unit/math/rev/prob/beta_cdf_test.cpp +++ b/test/unit/math/rev/prob/beta_cdf_test.cpp @@ -4,10 +4,11 @@ #include #include +/** + * For one observation the shape partials of beta_cdf are d/da I and d/db I. + * Reference: mpmath, 80 digits. The point has b >> a with a small y. + */ TEST_F(AgradRev, ProbDistributionsBeta_cdf_shape_gradients) { - // For one observation the shape partials of beta_cdf are d/da I and - // d/db I. Reference: mpmath, 80 digits. The point has b >> a with a - // small y. using stan::math::var; var a = 1.145; var b = 6786;