Fix the underflow and gradient errors of the Wiener functions

[avehtari]

Fixes #3329.

This PR improves wiener_lpdf (5- and 7-parameter), wiener_lcdf_unnorm and wiener_lccdf_unnorm (4- and 7-parameter).

This was assisted by Claude

Summary

This PR fixes the underflow of #3329 and other accuracy defects of the Wiener functions that I found while testing a fix for it. It is an alternative to #3435 (comparison: #3435 (comment)).

Functions and files. The Stan functions, the C++ files that compute them (in stan/math/prim/prob), and the internal functions that they call. Each function has an optional last argument for the precision.

Stan function	C++ file	computes with
wiener_lpdf(y | alpha, tau, beta, delta), 4-parameter	wiener_lpdf.hpp	its own series; accurate and does not need fixing (at 300 points against mpmath: value error at most 3.6e-9, gradient error at most 5.2e-8), so this PR does not change it
wiener_lpdf(y | a, t0, w, v, sv), 5-parameter	wiener5_lpdf.hpp	the series of the 5-parameter density and its gradients (internal::wiener5_density, internal::wiener5_grad_*)
wiener_lpdf(y | a, t0, w, v, sv, sw, st0), 7-parameter	wiener_full_lpdf.hpp	hcubature integrals of the 5-parameter density and gradients over sw and st0 (internal::wiener7_integrate); with sw = st0 = 0 it calls the 5-parameter wiener_lpdf
wiener_lcdf_unnorm(y, a, t0, w, v), 4-parameter	wiener4_lcdf_unnorm.hpp	the series of the 4-parameter cdf and its gradients (internal::wiener4_distribution, internal::wiener4_cdf_grad_*, logMill)
wiener_lccdf_unnorm(y, a, t0, w, v), 4-parameter	wiener4_lccdf_unnorm.hpp	P - F with the 4-parameter cdf, or the large-time series of P - F (new); P is the probability of the upper boundary
wiener_lcdf_unnorm(y, a, t0, w, v, sv, sw, st0), 7-parameter	wiener_full_lcdf_unnorm.hpp	hcubature integrals of the 4-parameter cdf and gradients over sv, sw and st0 (internal::wiener7_integrate_cdf); with sv = sw = st0 = 0 it calls the 4-parameter wiener_lcdf_unnorm
wiener_lccdf_unnorm(y, a, t0, w, v, sv, sw, st0), 7-parameter	wiener_full_lccdf_unnorm.hpp	the same integrals of the 4-parameter ccdf and gradients (internal::wiener4_ccdf, internal::wiener4_ccdf_grad_*); with sv = sw = st0 = 0 it calls the 4-parameter wiener_lccdf_unnorm

In the rest of this text, "4-parameter lcdf" means wiener_lcdf_unnorm with 4 parameters, "7-parameter lpdf" means wiener_lpdf with 7 parameters, and so on. Because of these calls, a fix in the 5-parameter functions also changes the 7-parameter lpdf, and a fix in the 4-parameter cdf functions also changes the 4-parameter lccdf and the 7-parameter lcdf and lccdf (see the table "Several fixes are in shared code" below).

7-parameter lpdf, #3329. The integrands of the 7-parameter density are natural-scale 5-parameter densities, which underflow to 0 when the log density is below about -708. Per observation, the integrands are now scaled by exp(-c), where c is the largest 5-parameter log density on a 3 x 3 grid of the integration box (a 17 x 17 grid if the scaled integral is still outside [1e-100, 1e100]). The log density is c + log(integral); the partials are ratios of integrals, so c cancels. The series tolerances are relative to exp(c). wiener5_density and wiener5_grad_* get an optional log_shift argument (default 0); existing callers are unchanged. The existing hcubature is used; all autodiff types use the same code path.

hcubature. hcubature is used only by the 7-parameter wiener_lpdf, wiener_lcdf_unnorm and wiener_lccdf_unnorm, so this fix changes only these three functions. Two bugs: (1) the loop condition used the error and value of the first box, so once the first box failed the test, the loop ran to the evaluation limit; (2) a split box stayed in the list with error 0 and could be split and counted again. With the fix the integration stops at the requested tolerance. This makes the 7-parameter lpdf 11 to 14 times faster when st0 > 0.

d/dsw and d/dst0 of the 7-parameter lpdf.

• The d/dsw integral over the st0 variable used the limits of the sw variable; for (y - t0)/st0 < 0.004, d/dsw was -1/sw.
• With sw > 0 and st0 = 0, d/dsw used an absolute series tolerance (relative errors up to 2e3 where the density is small).
• The end-point forms $\partial \log p / \partial\,\mathrm{sw} = I_\mathrm{sw} / I - 1/\mathrm{sw}$ and $\partial \log p / \partial\,\mathrm{st0} = f / (\mathrm{st0}\, I) - 1/\mathrm{st0}$ multiply the errors of two separate integrations by 1/sw (1/st0). Here $I$ is the density integral, $I_\mathrm{sw}$ the integral over the st0 variable of the density at $w \pm \mathrm{sw}/2$, and $f$ the integral over the sw variable of the density at $t = y - t_0 - \mathrm{st0}$. When both sw and st0 are non-zero and sw (st0) is below 1e-3, the partial is now computed by differentiation under the integral (the idea of the fallback in Fix full Wiener density underflow #3435). With the integration variables $x_w, x_t \in [0, 1]$ and the 5-parameter density $p_5$: $\partial \log p / \partial\,\mathrm{sw} = \int (x_w - \tfrac{1}{2})\, \partial_w p_5 \, dx \,/ \int p_5 \, dx$, with the integrand folded at $x_w = \tfrac{1}{2}$ so that its odd part cancels at each node; $\partial \log p / \partial\,\mathrm{st0} = -\int x_t\, \partial_t p_5 \, dx \,/ \int p_5 \, dx$. For larger sw and st0 the end-point forms are kept: their error is about 3e-10/sw at the default precision, and the score integrals cost one more 2-D integral. In develop, the d/dsw error at sw = 1e-8 is up to 1.1 at the default precision.
• The density integral uses a relative tolerance of at most 0.9e-6 (log density error at the default precision 1.7e-6 -> 1.0e-7).

5-parameter lpdf. Three faults on the path for small densities: assign_err took its argument by value, so the re-estimate in estimate_with_err_check used the old tolerance; with LogResult == false the check compared |log density| with the log tolerance and never fired; the gradient series of t, a and w used an absolute tolerance. Value errors were up to 1e-2 at ordinary points.

4-parameter lcdf. The small-time series of the derivatives with respect to a and w had sign errors (gradient errors up to 1.6e-4; through these functions also the 7-parameter lcdf and lccdf gradients, up to 3e-3). logMill uses erfcx (recently merged) for x >= 0.

4-parameter lccdf. The ccdf P - F loses all digits where F is close to P (-inf at 8 of 300 ordinary points). Where F >= P/2 it now uses the large-time series of P - F, with its derivatives.

7-parameter lcdf and lccdf. The integration tolerance is at most 0.9e-6. The lccdf now includes the trials whose non-decision time ends after y (they contribute P; errors in log ccdf up to 1.4 where st0 > y - t0). The lccdf integrates the 4-parameter ccdf above, so it is finite down to a ccdf of about 1e-308; develop formed P - F and returned -inf below about 1e-16 of P (for example at a log ccdf of -39, -120 and -299). The d/dt series tolerance is relative to the cdf or ccdf. Where the value integral underflows to 0, the value is -inf and the gradient integrals are skipped.

Several fixes are in shared code, so they reach more than one function:

fix	4-parameter lcdf	4-parameter lccdf	5-parameter lpdf	7-parameter lpdf	7-parameter lcdf	7-parameter lccdf
hcubature: loop condition, split box counted twice				x	x	x
assign_err by reference	x	x	x	x	x	x
estimate_with_err_check with LogResult == false	x		x	(x)	(x)
tolerance of the 5-parameter gradient series (t, a, w)			x	(x)
logMill with erfcx	x	x			x	x
sign errors in the 4-parameter cdf gradients (a, w)	x	x			x	x
4-parameter ccdf from the large-time series of P - F		x				x
log shift of the integrand (#3329)				x
d/dsw: integration limits, series tolerance				x
d/dsw, d/dst0 by differentiation under the integral (small sw, st0)				x
integration tolerance at most 0.9e-6				x	x	x
d/dt series tolerance relative to the value					x	x
guards: no NaN where the integrand weight is 0; -inf without gradient integrals where the value underflows					x	x
lccdf: mass of the trials with t0 + τ > y						x

(x): only through the call to the 5-parameter lpdf (7-parameter lpdf with sw = st0 = 0) or to the 4-parameter lcdf (7-parameter lcdf with sv = sw = st0 = 0). The 4-parameter lccdf uses the 4-parameter cdf and its gradients where F < P/2; the 7-parameter lcdf and lccdf use them in their integrands. The older 4-parameter wiener_lpdf (5 arguments) is not changed.

Accuracy, develop -> this PR (largest error over the test points; log scale for values, relative to max(|ref|, 1) for gradients; default precision):

function	value	gradients
4-parameter lcdf	1.1e-8 -> 1.1e-8	1.6e-4 -> 9.5e-9
4-parameter lccdf	-inf -> 1.4e-14	-inf -> 2.9e-9
5-parameter lpdf	1.1e-2 -> 4.6e-12	1.8e-2 -> 2.1e-14
7-parameter lpdf, ordinary	1.7e-6 -> 1.0e-7	2.3e-5 -> 2.1e-5
7-parameter lpdf, #3329 point and scan	-inf -> 4.8e-11	d/dsw -1e8 -> 5.3e-10
7-parameter lpdf, sw or st0 = 1e-8		1.1 -> 3.3e-9
7-parameter lcdf, 1-2 integration dims	3.1e-7 -> 3.1e-7	2.9e-3 -> 3.9e-7
7-parameter lccdf, 1-2 integration dims	-inf, 1.3 -> 8.4e-9	6.3e2 -> 2.7e-8

References: mpmath (50 digits) for the 4- and 5-parameter functions; for the 7-parameter functions an independent adaptive Gauss-Kronrod integration of the 5-parameter density (own code, not hcubature), with gradients by differentiation under the integral or by finite differences of the reference. Ordinary points: sets of 300 random points (a 0.5-3, v -5..5 or -3..3, w 0.2-0.8, sv 0-1.5, sw up to 0.8 of its range, st0 up to 0.3, y - t0 0.05-3); the #3329 scan has 16 points with y - t0 from 3e-6 to 3e-5; the small sw and st0 set has 240 points with sizes from 1e-8 to 0.1.

Time, this PR / develop (gradients): 4-parameter lcdf 0.89, lccdf 0.62; 5-parameter lpdf 1.09-1.13; 7-parameter lpdf 0.07 (sw and st0), 0.09 (st0 only), 1.05 (sw only); 7-parameter lcdf 0.04-0.67, lccdf 0.06-0.87.

Not fixed (follow-up PR): the 7-parameter lcdf and lccdf integrate natural-scale values, so they return -inf where the cdf or ccdf is below about 1e-308 (the #3329 mechanism; for the lccdf this PR moves the limit from about 1e-16 of P to 1e-308); with sv, sw and st0 all non-zero, the 3-dimensional lcdf/lccdf integrals reach the evaluation limit of 6000 (errors up to 4.3e-5, as in develop).

Diff size against develop (43bc568), code:

file	+	−
stan/math/prim/functor/hcubature.hpp	8	14
stan/math/prim/prob/wiener4_lcdf_unnorm.hpp	22	5
stan/math/prim/prob/wiener4_lccdf_unnorm.hpp	227	49
stan/math/prim/prob/wiener5_lpdf.hpp	45	20
stan/math/prim/prob/wiener_full_lpdf.hpp	162	35
stan/math/prim/prob/wiener_full_lcdf_unnorm.hpp	37	12
stan/math/prim/prob/wiener_full_lccdf_unnorm.hpp	41	14
code, 7 files	542	149

Tests combined: 6 files, +454/−10. Four of these files are new (198 lines); the other two are hcubature_test.cpp (+47) and wiener_full_test.cpp (+209/−10).

Tests

New tests (each fails with the develop versions of the 7 headers): hcubature_test (a step integrand, for which develop splits finished boxes again and counts them twice: 99.5 instead of 0.5; sqrt(x) and a 2-D Gaussian, which must stop at the tolerance, not at the evaluation limit); prim wiener4_lcdf_unnorm_test (logMill); rev wiener4_lcdf_unnorm_test (gradients, mpmath); rev wiener4_lccdf_unnorm_test (small ccdf, mpmath); rev wiener5_lpdf_test (small densities, mpmath); rev wiener_full_test: the #3329 point and scan (values, d/dy, d/dst0, d/dsw), d/dsw with st0 > y - t0, 7-parameter lccdf with st0 > y - t0 and with a small ccdf, 7-parameter lccdf at log ccdf -120 and -299 (independent references; develop -inf), and the 7-parameter lcdf and lccdf below the smallest double (-inf, without the gradient integrals).

Changed test: in mathRevCorrectValues_wiener_lccdf_unnorm the reference of the third point (WienR) computes cdf(inf) - cdf and loses digits (the ccdf is 1.6e-14 of P): WienR -33.8123, mpmath -33.7956003. The row now has the mpmath values.

All Wiener prim, rev and mix suites and the hcubature suite pass.

Side Effects

None.

Release notes

Improved the accuracy and speed of wiener_lpdf (5- and 7-parameter), wiener_lcdf_unnorm and wiener_lccdf_unnorm (4- and 7-parameter).

Checklist

• Copyright holder: Aalto University

  The copyright holder is typically you or your assignee, such as a university or company. By submitting this pull request, the copyright holder is agreeing to the license the submitted work under the following licenses:
  - Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
  - Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)

• the basic tests are passing

  • unit tests pass (to run, use: ./runTests.py test/unit)
  • header checks pass, (make test-headers)
  • dependencies checks pass, (make test-math-dependencies)
  • docs build, (make doxygen)
  • code passes the built in C++ standards checks (make cpplint)

• the code is written in idiomatic C++ and changes are documented in the doxygen

• the new changes are tested

[stan-buildbot]

Name	Old Result	New Result	Ratio	Performance change( 1 - new / old )
stat_comp_benchmarks/benchmarks/gp_regr/gen_gp_data.stan	0.04	0.04	0.96	-4.1% slower
stat_comp_benchmarks/benchmarks/gp_regr/gp_regr.stan	0.16	0.17	0.99	-0.73% slower
stat_comp_benchmarks/benchmarks/garch/garch.stan	0.85	0.89	0.96	-4.27% slower
stat_comp_benchmarks/benchmarks/low_dim_corr_gauss/low_dim_corr_gauss.stan	0.01	0.02	0.98	-2.13% slower
stat_comp_benchmarks/benchmarks/eight_schools/eight_schools.stan	0.11	0.11	1.0	-0.1% slower
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix/low_dim_gauss_mix.stan	8.03	8.02	1.0	0.18% faster
stat_comp_benchmarks/benchmarks/gp_pois_regr/gp_pois_regr.stan	3.8	3.75	1.01	1.37% faster
stat_comp_benchmarks/benchmarks/pkpd/one_comp_mm_elim_abs.stan	33.01	32.53	1.01	1.44% faster
stat_comp_benchmarks/benchmarks/pkpd/sim_one_comp_mm_elim_abs.stan	0.38	0.37	1.02	2.17% faster
stat_comp_benchmarks/benchmarks/irt_2pl/irt_2pl.stan	6.92	6.93	1.0	-0.02% slower
stat_comp_benchmarks/benchmarks/sir/sir.stan	136.43	138.35	0.99	-1.41% slower
stat_comp_benchmarks/benchmarks/arK/arK.stan	3.56	3.41	1.05	4.36% faster
stat_comp_benchmarks/benchmarks/arma/arma.stan	0.61	0.61	1.0	0.17% faster
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix_collapse/low_dim_gauss_mix_collapse.stan	25.89	25.83	1.0	0.23% faster
performance.compilation	404.47	401.84	1.01	0.65% faster
Mean result: 0.9990045408676703

Jenkins Console Log
Jenkins Build Stages
Commit hash: 44a50912ff3f2e71023ea29a6f91adf7c6b1cf87

Machine information

Distributor ID:	Ubuntu
Description:	Ubuntu 20.04.3 LTS
Release:	20.04
Codename:	focal

CPU:

Architecture:                            x86_64
CPU op-mode(s):                          32-bit, 64-bit
Byte Order:                              Little Endian
Address sizes:                           52 bits physical, 57 bits virtual
CPU(s):                                  192
On-line CPU(s) list:                     0-191
Thread(s) per core:                      2
Core(s) per socket:                      48
Socket(s):                               2
NUMA node(s):                            2
Vendor ID:                               AuthenticAMD
CPU family:                              25
Model:                                   17
Model name:                              AMD EPYC 9474F 48-Core Processor
Stepping:                                1
Frequency boost:                         enabled
CPU MHz:                                 1496.938
CPU max MHz:                             4114.4229
CPU min MHz:                             1500.0000
BogoMIPS:                                7189.39
Virtualization:                          AMD-V
L1d cache:                               3 MiB
L1i cache:                               3 MiB
L2 cache:                                96 MiB
L3 cache:                                512 MiB
NUMA node0 CPU(s):                       0-47,96-143
NUMA node1 CPU(s):                       48-95,144-191
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Old microcode:             Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Mitigation; Safe RET
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; Enhanced / Automatic IBRS; IBPB conditional; STIBP always-on; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Mitigation; Clear CPU buffers
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Mitigation; IBPB before exit to userspace
Flags:                                   fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good amd_lbr_v2 nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba perfmon_v2 ibrs ibpb stibp ibrs_enhanced vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local user_shstk avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin cppc arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif x2avic v_spec_ctrl vnmi avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid overflow_recov succor smca fsrm flush_l1d debug_swap

G++:

g++ (Ubuntu 9.4.0-1ubuntu1~20.04) 9.4.0
Copyright (C) 2019 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

Clang:

clang version 10.0.0-4ubuntu1 
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin