fix integrate_1d_gauss_kronrod gradient speed

Author: avehtari

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 + f) - 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 <typename F, typename T_a, typename T_b, typename Integrator,
-          typename... Args>
+template <bool shift_gradient_integrand = false, typename F, typename T_a,
+          typename T_b, typename Integrator, typename... Args>
 inline return_type_t<T_a, T_b, Args...> 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,15 @@ inline return_type_t<T_a, T_b, Args...> integrate_1d_adjoint(
 
   // Argument adjoints.
   if constexpr (is_any_var_scalar_v<Args...>) {
+    const bool shift = shift_gradient_integrand && integral != 0.0
+                       && !is_inf(integral) && !is_nan(integral);
     auto args_adj = make_zeroed_arena(std::forward_as_tuple(args...));
     {
       nested_rev_autodiff argument_nest;
       auto args_copy = deep_copy_vargs<var>(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 +121,9 @@ inline return_type_t<T_a, T_b, Args...> integrate_1d_adjoint(
           if (is_nan(gradient) && fx.val() == 0) {
             gradient = 0.0;
           }
-          return gradient;
+          return shift ? gradient + fx.val() : gradient;
         });
+        return shift ? result - integral : result;
       };
       std::size_t param_index = 0;
       auto assign_grad = [&](auto&& adj, auto&& target) {

stan/math/rev/functor/integrate_1d_gauss_kronrod.hpp

@@ -45,7 +45,8 @@ inline return_type_t<T_a, T_b, Args...> 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<true>(
       function, f, a, b,
       [&](auto &&integrand) {
         return integrate_gk(std::forward<decltype(integrand)>(integrand),

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]);
 }
 
-}  // namespace integrate_1d_gk_test
+// ── gradient-integrand shift (integrate_1d_adjoint<true>) ────────────
+//
+// 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 <typename T1, typename T2, typename T3>
+  inline stan::return_type_t<T1, T2, T3> operator()(
+      const T1 &x, const T2 &xc, std::ostream *msgs,
+      const std::vector<T3> &theta, const std::vector<double> &x_r,
+      const std::vector<int> &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<var> 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<double>{}, std::vector<int>{});
+  std::vector<double> 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 <typename T1, typename T2, typename T3>
+  inline stan::return_type_t<T1, T2, T3> operator()(
+      const T1 &x, const T2 &xc, std::ostream *msgs,
+      const std::vector<T3> &theta, const std::vector<double> &x_r,
+      const std::vector<int> &x_i) const {
+    return theta[0] * x * exp(-x * x);   // odd: integral over [-1, 1] is 0
+  }
+};
+
+struct f_negative {
+  template <typename T1, typename T2, typename T3>
+  inline stan::return_type_t<T1, T2, T3> operator()(
+      const T1 &x, const T2 &xc, std::ostream *msgs,
+      const std::vector<T3> &theta, const std::vector<double> &x_r,
+      const std::vector<int> &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<var> 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<double>{}, std::vector<int>{}));
+    std::vector<double> 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<var> 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<double>{}, std::vector<int>{});
+    std::vector<double> g;
+    I.grad(theta, g);
+    EXPECT_NEAR(I_ref, I.val(), 1e-7);
+    EXPECT_NEAR(I_ref / th, g[0], 1e-7);
+  }
+}
+
+}