`wiener_lpdf`: faster term counts in the 4-parameter density and less repeated work in the 5-parameter gradients

--- a/stan/math/prim/prob/wiener_lpdf.hpp
+++ b/stan/math/prim/prob/wiener_lpdf.hpp
@@ -43,6 +43,7 @@
 #include <stan/math/prim/fun/square.hpp>
 #include <stan/math/prim/fun/sqrt.hpp>
 #include <stan/math/prim/fun/value_of.hpp>
+#include <stan/math/prim/fun/value_of_rec.hpp>
 #include <algorithm>
 #include <cmath>
 #include <string>
@@ -150,17 +151,20 @@
     typename scalar_type<T_beta>::type one_minus_beta = 1.0 - beta_vec[i];
     typename scalar_type<T_alpha>::type alpha2 = square(alpha_vec[i]);
     T_return_type x = (y_vec[i] - tau_vec[i]) / alpha2;
-    T_return_type kl, ks, tmp = 0;
-    T_return_type k, K;
-    T_return_type sqrt_x = sqrt(x);
-    T_return_type log_x = log(x);
-    T_return_type one_over_pi_times_sqrt_x = 1.0 / pi() * sqrt_x;
+    T_return_type tmp = 0;
+    // The number of terms has zero derivative, so compute it on values and
+    // keep it, and the loop counter, off the autodiff stack
+    const double x_val = value_of_rec(x);
+    const double sqrt_x = std::sqrt(x_val);
+    const double log_x = std::log(x_val);
+    const double one_over_pi_times_sqrt_x = 1.0 / pi() * sqrt_x;
+    double kl, ks;
 
     // calculate number of terms needed for large t:
     // if error threshold is set low enough
-    if (PI_TIMES_WIENER_ERR * x < 1) {
+    if (PI_TIMES_WIENER_ERR * x_val < 1) {
       // compute bound
-      kl = sqrt(-2.0 * SQRT_PI * (LOG_PI_LOG_WIENER_ERR + log_x)) / sqrt_x;
+      kl = std::sqrt(-2.0 * SQRT_PI * (LOG_PI_LOG_WIENER_ERR + log_x)) / sqrt_x;
       // ensure boundary conditions met
       kl = (kl > one_over_pi_times_sqrt_x) ? kl : one_over_pi_times_sqrt_x;
     } else {
@@ -168,28 +172,28 @@
     }
     // calculate number of terms needed for small t:
     // if error threshold is set low enough
-    T_return_type tmp_expr0 = TWO_TIMES_SQRT_TWO_PI_TIMES_WIENER_ERR * sqrt_x;
+    const double tmp_expr0 = TWO_TIMES_SQRT_TWO_PI_TIMES_WIENER_ERR * sqrt_x;
     if (tmp_expr0 < 1) {
       // compute bound
-      ks = 2.0 + sqrt_x * sqrt(-2 * log(tmp_expr0));
+      ks = 2.0 + sqrt_x * std::sqrt(-2 * std::log(tmp_expr0));
       // ensure boundary conditions are met
-      T_return_type sqrt_x_plus_one = sqrt_x + 1.0;
+      const double sqrt_x_plus_one = sqrt_x + 1.0;
       ks = (ks > sqrt_x_plus_one) ? ks : sqrt_x_plus_one;
     } else {     // if error threshold was set too high
       ks = 2.0;  // minimal kappa for that case
     }
-    if (ks < kl) {   // small t
-      K = ceil(ks);  // round to smallest integer meeting error
-      T_return_type tmp_expr1 = (K - 1.0) / 2.0;
-      T_return_type tmp_expr2 = ceil(tmp_expr1);
-      for (k = -floor(tmp_expr1); k <= tmp_expr2; k++) {
-        tmp += (one_minus_beta + 2.0 * k)
-               * exp(-(square(one_minus_beta + 2.0 * k)) * 0.5 / x);
+    if (ks < kl) {                     // small t
+      const double K = std::ceil(ks);  // round to smallest integer meeting error
+      const double tmp_expr1 = (K - 1.0) / 2.0;
+      const double tmp_expr2 = std::ceil(tmp_expr1);
+      for (double k = -std::floor(tmp_expr1); k <= tmp_expr2; k++) {
+        const auto w_plus_2_k = one_minus_beta + 2.0 * k;
+        tmp += w_plus_2_k * exp(-(square(w_plus_2_k)) * 0.5 / x);
       }
-      tmp = log(tmp) - LOG_TWO_OVER_TWO_PLUS_LOG_SQRT_PI - 1.5 * log_x;
-    } else {         // if large t is better...
-      K = ceil(kl);  // round to smallest integer meeting error
-      for (k = 1; k <= K; ++k) {
+      tmp = log(tmp) - LOG_TWO_OVER_TWO_PLUS_LOG_SQRT_PI - 1.5 * log(x);
+    } else {                           // if large t is better...
+      const double K = std::ceil(kl);  // round to smallest integer meeting error
+      for (double k = 1; k <= K; ++k) {
         tmp += k * exp(-(square(k)) * (SQUARE_PI_OVER_TWO * x))
                * sin(k * pi() * one_minus_beta);
       }

--- a/stan/math/prim/prob/wiener_lpdf.hpp
+++ b/stan/math/prim/prob/wiener_lpdf.hpp
@@ -157,14 +157,14 @@
     const double x_val = value_of_rec(x);
     const double sqrt_x = std::sqrt(x_val);
     const double log_x = std::log(x_val);
-    const double one_over_pi_times_sqrt_x = 1.0 / pi() * sqrt_x;
+    const double one_over_pi_times_sqrt_x = 1.0 / (pi() * sqrt_x);
     double kl, ks;
 
     // calculate number of terms needed for large t:
     // if error threshold is set low enough
     if (PI_TIMES_WIENER_ERR * x_val < 1) {
       // compute bound
-      kl = std::sqrt(-2.0 * SQRT_PI * (LOG_PI_LOG_WIENER_ERR + log_x)) / sqrt_x;
+      kl = std::sqrt(-2.0 * (LOG_PI_LOG_WIENER_ERR + log_x)) / (pi() * sqrt_x);
       // ensure boundary conditions met
       kl = (kl > one_over_pi_times_sqrt_x) ? kl : one_over_pi_times_sqrt_x;
     } else {

wiener_lpdf_b.hpp and wiener_lpdf_c.hpp are copies of wiener_lpdf.hpp with patch 1, and with patches 1 + 2. In each copy, the include guard and the function are renamed (wiener_lpdf_b, wiener_lpdf_c).

#include <stan/math/rev.hpp>
#include "wiener_lpdf_b.hpp"  // patch 1
#include "wiener_lpdf_c.hpp"  // patch 1 + 2
#include <chrono>
#include <cstdio>
#include <random>
#include <vector>
using namespace stan::math;
using internal::GradientCalc;

template <typename F>
double best_ns(F&& f, int reps, int n) {
  double best = 1e300;
  for (int r = 0; r < reps; ++r) {
    auto t0 = std::chrono::steady_clock::now();
    f();
    auto t1 = std::chrono::steady_clock::now();
    best = std::min(best, std::chrono::duration<double, std::nano>(t1 - t0).count() / n);
  }
  return best;
}

// d/da log p from d/dy log p: the series depends on y and a only through y/a^2
double grad_a_from_t(double y, double a, double v, double w, double sv, double gt) {
  const double q = 1.0 - w, sv2 = sv * sv, opsy = 1.0 + sv2 * y;
  const double dlet_da = (v * q + sv2 * q * q * a) / opsy - 2.0 / a;
  const double dlet_dy = -0.5 * (sv2 * sv2 * (y + a * a * q * q) + sv2 * (1.0 + 2.0 * a * v * q) + v * v)
                         / (opsy * opsy);
  return dlet_da - (2.0 * y / a) * (gt - dlet_dy);
}

// wiener5_density with the series summed in linear space after factoring out
// the largest exponent (as wiener5_grad_w already does)
double density_linear(double y, double a, double v, double w, double sv, double log_err) {
  const double let = internal::wiener5_compute_log_error_term(y, a, v, w, sv);
  const double n_small
      = internal::wiener5_n_terms_small_t<true, GradientCalc::OFF>(y, a, w, log_err - let);
  const double n_large = internal::wiener5_density_large_reaction_time_terms(y, a, w, log_err - let);
  const double u = y / (a * a), q = 1.0 - w;
  if (2 * n_small <= n_large) {
    const double m = -q * q / (2.0 * u);
    double s = q;
    for (int k = 1; k <= static_cast<int>(n_small); ++k) {
      const double zp = q + 2.0 * k, zm = q - 2.0 * k;
      s += zp * std::exp(-zp * zp / (2.0 * u) - m) + zm * std::exp(-zm * zm / (2.0 * u) - m);
    }
    return let - 0.5 * LOG_TWO - LOG_SQRT_PI - 1.5 * std::log(u) + std::log(std::fabs(s)) + m;
  }
  const double c = 0.5 * pi() * pi() * u;
  double s = 0.0;
  for (int k = 1; k <= static_cast<int>(n_large); ++k)
    s += k * std::sin(k * pi() * q) * std::exp(-(k * k - 1.0) * c);
  return let + std::log(std::fabs(s)) - c + LOG_PI;
}

int main() {
  // Part 1: public functions, value + gradient, reverse mode
  const int N = 2000, reps = 200;
  std::mt19937 rng(1);
  std::uniform_real_distribution<double> ua(0.8, 2.0), ut0(0.15, 0.35), uw(0.3, 0.7), uv(-1.5, 2.5),
      udt(0.15, 1.5);
  std::vector<double> y(N), a0(N), t00(N), w0(N), v0(N);
  for (int i = 0; i < N; ++i) {
    a0[i] = ua(rng); t00[i] = ut0(rng); w0[i] = uw(rng); v0[i] = uv(rng); y[i] = t00[i] + udt(rng);
  }
  std::vector<double> g_ref;
  double lp_ref = 0, t_ref = 0;
  auto run = [&](const char* name, auto&& f) {
    double lp_val = 0;
    std::vector<double> g;
    double t = best_ns([&] {
      std::vector<var> a(a0.begin(), a0.end()), t0(t00.begin(), t00.end()), w(w0.begin(), w0.end()),
          v(v0.begin(), v0.end());
      var lp = f(a, t0, w, v);
      lp.grad();
      lp_val = lp.val();
      g.clear();
      for (int i = 0; i < N; ++i) {
        g.push_back(a[i].adj()); g.push_back(t0[i].adj()); g.push_back(w[i].adj()); g.push_back(v[i].adj());
      }
      recover_memory();
    }, reps, N);
    if (g_ref.empty()) { g_ref = g; lp_ref = lp_val; t_ref = t; }
    double md = 0;
    for (size_t i = 0; i < g.size(); ++i) md = std::max(md, std::fabs(g[i] - g_ref[i]));
    std::printf("| %-46s | %7.0f | %5.2fx | %.1e | %.1e |\n", name, t, t_ref / t, std::fabs(lp_val - lp_ref), md);
  };
  run("4-param, develop", [&](auto& a, auto& t0, auto& w, auto& v) { return wiener_lpdf(y, a, t0, w, v); });
  run("4-param, patch 1", [&](auto& a, auto& t0, auto& w, auto& v) { return wiener_lpdf_b(y, a, t0, w, v); });
  run("4-param, patch 1 + 2", [&](auto& a, auto& t0, auto& w, auto& v) { return wiener_lpdf_c(y, a, t0, w, v); });
  run("5-param with sv = 0", [&](auto& a, auto& t0, auto& w, auto& v) { return wiener_lpdf(y, a, t0, w, v, 0.0); });
  volatile double sink = 0;
  double t_dbl = best_ns([&] { sink = sink + wiener_lpdf(y, a0, t00, w0, v0, 0.0); }, reps, N);
  std::printf("5-param, all arguments double (value only): %.0f ns/obs\n\n", t_dbl);

  // Part 2: 5-parameter internals on 100k random points
  const double le = std::log(1e-12);
  std::mt19937 rng2(42);
  std::uniform_real_distribution<double> uy(0.05, 2.5), ua2(0.5, 3.0), uv2(-4, 4), uw2(0.1, 0.9), usv(0.0, 1.0);
  struct Pt { double y, a, v, w, sv; };
  std::vector<Pt> pts(100000);
  for (auto& p : pts) p = {uy(rng2), ua2(rng2), uv2(rng2), uw2(rng2), (rng2() % 2) ? usv(rng2) : 0.0};
  auto per_point = [&](const char* name, auto&& f) {
    double t = best_ns([&] {
      double acc = 0;
      for (auto& p : pts) acc += f(p);
      sink = sink + acc;
    }, 5, pts.size());
    std::printf("| %-40s | %7.0f |\n", name, t);
  };
  per_point("wiener5_density", [&](const Pt& p) { return internal::wiener5_density<false>(p.y, p.a, p.v, p.w, p.sv, le); });
  per_point("density, linear-space sum", [&](const Pt& p) { return density_linear(p.y, p.a, p.v, p.w, p.sv, le); });
  per_point("wiener5_grad_t", [&](const Pt& p) { return internal::wiener5_grad_t<false>(p.y, p.a, p.v, p.w, p.sv, le); });
  per_point("wiener5_grad_a", [&](const Pt& p) { return internal::wiener5_grad_a<false>(p.y, p.a, p.v, p.w, p.sv, le); });
  per_point("grad_a from grad_t (identity)", [&](const Pt& p) { return grad_a_from_t(p.y, p.a, p.v, p.w, p.sv, 0.3); });
  per_point("wiener5_grad_w", [&](const Pt& p) { return internal::wiener5_grad_w<false>(p.y, p.a, p.v, p.w, p.sv, le); });

  // Part 3: accuracy of the identity and of the linear-space sum
  const double lt = std::log(1e-20);
  double id_tight = 0, id_fd = 0, id_def = 0, lin = 0;
  for (size_t i = 0; i < 20000; ++i) {
    const auto& p = pts[i];
    const double ga_t = internal::wiener5_grad_a<false>(p.y, p.a, p.v, p.w, p.sv, lt);
    const double gi_t = grad_a_from_t(p.y, p.a, p.v, p.w, p.sv,
                                      internal::wiener5_grad_t<false>(p.y, p.a, p.v, p.w, p.sv, lt));
    const double h = 1e-5 * p.a;
    const double fd = (internal::wiener5_density<false>(p.y, p.a + h, p.v, p.w, p.sv, lt)
                       - internal::wiener5_density<false>(p.y, p.a - h, p.v, p.w, p.sv, lt))
                      / (2 * h);
    id_tight = std::max(id_tight, std::fabs(ga_t - gi_t) / std::max(1.0, std::fabs(ga_t)));
    const double lp = internal::wiener5_density<false>(p.y, p.a, p.v, p.w, p.sv, le);
    if (lp > std::log(1e-6)) {
      id_fd = std::max(id_fd, std::fabs(gi_t - fd) / std::max(1.0, std::fabs(fd)));
      const double ga = internal::wiener5_grad_a<false>(p.y, p.a, p.v, p.w, p.sv, le);
      const double gi = grad_a_from_t(p.y, p.a, p.v, p.w, p.sv,
                                      internal::wiener5_grad_t<false>(p.y, p.a, p.v, p.w, p.sv, le));
      id_def = std::max(id_def, std::fabs(ga - gi) / std::max(1.0, std::fabs(ga)));
    }
    lin = std::max(lin, std::fabs(lp - density_linear(p.y, p.a, p.v, p.w, p.sv, le)));
  }
  std::printf("identity vs wiener5_grad_a, 1e-20 tolerance: %.1e\n", id_tight);
  std::printf("identity vs wiener5_grad_a, default tolerance, density > 1e-6: %.1e\n", id_def);
  std::printf("identity vs central difference, density > 1e-6: %.1e\n", id_fd);
  std::printf("linear-space vs log-space log density: %.1e\n", lin);
}

// Accuracy of the 4-parameter wiener_lpdf (develop, patch 1 + 2) against the
// 5-parameter internals with sv = 0 and a 1e-20 series tolerance
#include <stan/math/rev.hpp>
#include "wiener_lpdf_c.hpp"
#include <cstdio>
#include <random>
using namespace stan::math;

struct Res { double lp, ga, gt0, gw, gv; };

template <typename F>
Res eval(F&& f, double y, double a, double t0, double w, double v) {
  var av = a, t0v = t0, wv = w, vv = v;
  var lp = f(y, av, t0v, wv, vv);
  lp.grad();
  Res r{lp.val(), av.adj(), t0v.adj(), wv.adj(), vv.adj()};
  recover_memory();
  return r;
}

int main() {
  const double lt = std::log(1e-20);
  std::mt19937 rng(7);
  std::uniform_real_distribution<double> ua(0.5, 3.0), uw(0.2, 0.8), uv(-4, 4), udt(0.05, 3.0);
  double e_dev[5] = {0}, e_new[5] = {0};
  for (int i = 0; i < 20000; ++i) {
    const double a = ua(rng), w = uw(rng), v = uv(rng), dt = udt(rng), t0 = 0.2, y = t0 + dt;
    const double ref[5] = {internal::wiener5_density<false>(dt, a, v, w, 0.0, lt),
                           internal::wiener5_grad_a<false>(dt, a, v, w, 0.0, lt),
                           -internal::wiener5_grad_t<false>(dt, a, v, w, 0.0, lt),
                           internal::wiener5_grad_w<false>(dt, a, v, w, 0.0, lt),
                           internal::wiener5_grad_v<false>(dt, a, v, w, 0.0, lt)};
    const Res d = eval([](auto&&... x) { return wiener_lpdf(x...); }, y, a, t0, w, v);
    const Res c = eval([](auto&&... x) { return wiener_lpdf_c(x...); }, y, a, t0, w, v);
    const double dv[5] = {d.lp, d.ga, d.gt0, d.gw, d.gv}, cv[5] = {c.lp, c.ga, c.gt0, c.gw, c.gv};
    for (int k = 0; k < 5; ++k) {
      const double scale = k == 0 ? 1.0 : std::max(1.0, std::fabs(ref[k]));
      e_dev[k] = std::max(e_dev[k], std::fabs(dv[k] - ref[k]) / scale);
      e_new[k] = std::max(e_new[k], std::fabs(cv[k] - ref[k]) / scale);
    }
  }
  const char* names[5] = {"log density (abs)", "d/da", "d/dt0", "d/dw", "d/dv"};
  for (int k = 0; k < 5; ++k) std::printf("| %-18s | %9.1e | %13.1e |\n", names[k], e_dev[k], e_new[k]);
}