diff --git a/stan/math/rev/functor/integrate_1d_adjoint.hpp b/stan/math/rev/functor/integrate_1d_adjoint.hpp index 064e9ba9acc..5026c39a4b2 100644 --- a/stan/math/rev/functor/integrate_1d_adjoint.hpp +++ b/stan/math/rev/functor/integrate_1d_adjoint.hpp @@ -38,6 +38,12 @@ namespace internal { * zero. Any other NaN propagates into the component integral and is reported as * a `domain_error` naming the (flattened) parameter index. * + * `shift_gradient_integrand`: when true (and the value integral `I` is finite + * and non-zero) each component adjoint is computed as + * `integrator(d f / d arg + c f) - c I` instead of `integrator(d f / d arg)`, + * which can make the computation better behaving. + * + * @tparam shift_gradient_integrand see above * @tparam F Type of f * @tparam T_a type of first limit * @tparam T_b type of second limit @@ -52,8 +58,8 @@ namespace internal { * @param args additional arguments to pass to f * @return numeric integral of function f */ -template +template inline return_type_t integrate_1d_adjoint( const char* function, const F& f, const T_a& a, const T_b& b, Integrator&& integrator, std::ostream* msgs, const Args&... args) { @@ -94,13 +100,18 @@ inline return_type_t integrate_1d_adjoint( // Argument adjoints. if constexpr (is_any_var_scalar_v) { + const bool shift = shift_gradient_integrand && integral != 0.0 + && !is_inf(integral) && !is_nan(integral); + // Shift constant. The inverse golden ratio is unlikely to be a + // saturated gradient. + constexpr double shift_c = 0.6180339887498949; auto args_adj = make_zeroed_arena(std::forward_as_tuple(args...)); { nested_rev_autodiff argument_nest; auto args_copy = deep_copy_vargs(std::forward_as_tuple(args...)); auto args_copy_filter = filter_var_scalar_types(args_copy); auto integrate_grad = [&](auto&& target) -> double { - return integrator([&](const auto& x, const auto& xc) { + const double result = integrator([&](const auto& x, const auto& xc) { argument_nest.set_zero_all_adjoints(); nested_rev_autodiff gradient_nest; var fx = stan::math::apply( @@ -113,8 +124,9 @@ inline return_type_t integrate_1d_adjoint( if (is_nan(gradient) && fx.val() == 0) { gradient = 0.0; } - return gradient; + return shift ? gradient + shift_c * fx.val() : gradient; }); + return shift ? result - shift_c * integral : result; }; std::size_t param_index = 0; auto assign_grad = [&](auto&& adj, auto&& target) { diff --git a/stan/math/rev/functor/integrate_1d_gauss_kronrod.hpp b/stan/math/rev/functor/integrate_1d_gauss_kronrod.hpp index 135c4437d40..7a2d1ed063a 100644 --- a/stan/math/rev/functor/integrate_1d_gauss_kronrod.hpp +++ b/stan/math/rev/functor/integrate_1d_gauss_kronrod.hpp @@ -45,7 +45,8 @@ inline return_type_t integrate_1d_gauss_kronrod_tol( check_less_or_equal(function, "lower limit", a, b); check_nonnegative(function, "max_depth", max_depth); check_nonnegative(function, "absolute_tolerance", absolute_tolerance); - return internal::integrate_1d_adjoint( + // `true`: gradient integrands are shifted by f + return internal::integrate_1d_adjoint( function, f, a, b, [&](auto &&integrand) { return integrate_gk(std::forward(integrand), diff --git a/test/unit/math/rev/functor/integrate_1d_gauss_kronrod_test.cpp b/test/unit/math/rev/functor/integrate_1d_gauss_kronrod_test.cpp index 7dfc53065d6..d9e8554d396 100644 --- a/test/unit/math/rev/functor/integrate_1d_gauss_kronrod_test.cpp +++ b/test/unit/math/rev/functor/integrate_1d_gauss_kronrod_test.cpp @@ -481,4 +481,95 @@ TEST_F(AgradRev, StanMath_integrate_1d_gk_rev_TestUniform) { EXPECT_FLOAT_EQ(1, 1 + g[1]); } +// ── gradient-integrand shift (integrate_1d_adjoint) ──────────── +// +// The wrapper integrates (d f / d theta_i + f) and subtracts the +// value integral. These tests pin (a) the evaluation count on such a +// component, which the accuracy tests above cannot see, and (b) the +// guard paths: an integral that is exactly zero (shift disabled) and +// a negative integral (shift active with I < 0). + +long n_evals = 0; + +// d f / d theta is analytically zero (cos^2 + sin^2 = 1) but autodiff +// evaluates it as round-off noise of order 1e-16 * x * f. +struct f_noisy_gradient_counted { + template + inline stan::return_type_t operator()( + const T1 &x, const T2 &xc, std::ostream *msgs, + const std::vector &theta, const std::vector &x_r, + const std::vector &x_i) const { + ++n_evals; + auto tx = theta[0] * x; + return exp(-x * x) * (cos(tx) * cos(tx) + sin(tx) * sin(tx)); + } +}; + +TEST_F(AgradRev, StanMath_integrate_1d_gk_rev_GradientShift_noisy_gradient) { + using stan::math::var; + const double I_ref = std::sqrt(stan::math::pi()) * std::erf(1.0); + std::vector theta = {2.5}; + n_evals = 0; + var I = stan::math::integrate_1d_gauss_kronrod_tol( + f_noisy_gradient_counted{}, -1.0, 1.0, 1e-6, 0.0, 15, msgs, theta, + std::vector{}, std::vector{}); + std::vector g; + I.grad(theta, g); + EXPECT_NEAR(I_ref, I.val(), 1e-6); + EXPECT_NEAR(0.0, g[0], 1e-6); + // Value + one gradient component. Without the shift the gradient + // integral hits max_depth (2^15 * 21 ≈ 6.9e5 evaluations). + EXPECT_LT(n_evals, 5000L); +} + +struct f_odd { + template + inline stan::return_type_t operator()( + const T1 &x, const T2 &xc, std::ostream *msgs, + const std::vector &theta, const std::vector &x_r, + const std::vector &x_i) const { + return theta[0] * x * exp(-x * x); // odd: integral over [-1, 1] is 0 + } +}; + +struct f_negative { + template + inline stan::return_type_t operator()( + const T1 &x, const T2 &xc, std::ostream *msgs, + const std::vector &theta, const std::vector &x_r, + const std::vector &x_i) const { + return -theta[0] * exp(-x * x); // negative everywhere, I < 0 + } +}; + +TEST_F(AgradRev, StanMath_integrate_1d_gk_rev_GradientShift_guards) { + using stan::math::var; + // (a) exactly-zero integral: the shift is disabled and the gradient is + // the (zero) integral of x exp(-x^2). + { + std::vector theta = {2.5}; + var I; + EXPECT_NO_THROW(I = stan::math::integrate_1d_gauss_kronrod_tol( + f_odd{}, -1.0, 1.0, 1e-6, 0.0, 15, msgs, theta, + std::vector{}, std::vector{})); + std::vector g; + I.grad(theta, g); + EXPECT_NEAR(0.0, I.val(), 1e-8); + EXPECT_NEAR(0.0, g[0], 1e-8); + } + // (b) negative integral: shift active with I < 0; d I / d theta = I / theta. + { + const double th = 2.5; + const double I_ref = -th * std::sqrt(stan::math::pi()) * std::erf(1.0); + std::vector theta = {th}; + var I = stan::math::integrate_1d_gauss_kronrod_tol( + f_negative{}, -1.0, 1.0, 1e-8, 0.0, 15, msgs, theta, + std::vector{}, std::vector{}); + std::vector g; + I.grad(theta, g); + EXPECT_NEAR(I_ref, I.val(), 1e-7); + EXPECT_NEAR(I_ref / th, g[0], 1e-7); + } +} + } // namespace integrate_1d_gk_test