Understanding reduce_sum and ll_chunk in TEXAS-PSM#
A visual guide for paleogeologists#
This note explains two important pieces of computational machinery used in the TEXAS
inverse-temperature Stan models β reduce_sum and ll_chunk β using a small toy
example you can trace by hand.
Background in one sentence#
When we reconstruct past temperatures from Ring Index (RI) values, the Stan model must
evaluate β for every proposed temperature β how well it fits the data under every one of
the M calibration curves drawn from the forward posterior. For large N (many downcore
samples) and large M (many calibration draws), this gets expensive fast. reduce_sum
and ll_chunk are how we split that work across CPU cores.
1. The Problem: N Γ M evaluations per sampling step#
Suppose you have:
N = 100 downcore sediment samples (each with a measured RI value)
M = 200 plausible calibration curves drawn from the forward posterior
For every step the sampler takes, it must evaluate the likelihood:
For each sample n (1 to 100):
For each calibration draw m (1 to 200):
Compute: expected RI = f(T_n ; curve_m)
Ask: how likely is the observed RI given this expected RI?
Average the M likelihoods for sample n
Total evaluations per step: 100 Γ 200 = 20,000
Total evaluations (4 chains Γ 1,500 steps): 120,000,000
Thatβs a lot of repeated computation. If your computer has 8 CPU cores, why not use all of them at the same time?
2. The Solution: Split N samples across CPU cores#
reduce_sum divides the N samples into chunks and assigns each chunk to a
separate CPU core. All cores work simultaneously; their results are summed at the end.
N = 8 samples to process
βββββ¬ββββ¬ββββ¬ββββ¬ββββ¬ββββ¬ββββ¬ββββ
Observations (n): β 1 β 2 β 3 β 4 β 5 β 6 β 7 β 8 β
βββββ΄ββββ΄ββββ΄ββββ΄ββββ΄ββββ΄ββββ΄ββββ
β β
ββββββββββββββββββ€ ββββββββββββββββββ
βΌ βΌ βΌ βΌ
βββββββββββ βββββββββββ βββββββββββ βββββββββββ
β Core 1 β β Core 2 β β Core 3 β β Core 4 β
β n=1,2 β β n=3,4 β β n=5,6 β β n=7,8 β
β β β β β β β β
βll_chunk β βll_chunk β βll_chunk β βll_chunk β
β(1 β 2) β β(3 β 4) β β(5 β 6) β β(7 β 8) β
ββββββ¬βββββ ββββββ¬βββββ ββββββ¬βββββ ββββββ¬βββββ
β β β β
ββββββββββββββββββ΄βββββββββββββββ΄ββββββββββββββββββ
β
βΌ
lp_1 + lp_2 + lp_3 + lp_4
= total log-probability
The grainsize parameter controls the chunk size. grainsize = 1 makes the
smallest possible chunks (maximum parallelism); larger values reduce overhead.
3. What ll_chunk does inside each core#
ll_chunk is the function that processes one chunk of samples. For each sample n
in the chunk, it:
Applies the prior on temperature: βhow plausible is this T given our prior knowledge?β
Loops over all M calibration draws and computes how well each curve predicts the observed RI
Averages those M likelihoods (the log-sum-exp step)
Adds the result to the chunkβs running total
Here is the flow for a single sample inside ll_chunk:
For sample n=2 (RI_obs = 0.68, T_proposed = 22.5Β°C):
ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
β Step 1 β Prior β
β Is T = 22.5Β°C consistent with our prior guess? β
β prior_mu_t[2] = 20Β°C, prior_sigma_t = 10Β°C β
β β log N(22.5 | 20, 10) = β2.33 β
ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
β Step 2 β Likelihood under each calibration draw β
β β
β Draw m=1: Tβ=28, k=0.18, b=0.15, Ξ½=0.9, Ο=0.05 β
β β expected RI = 0.65 β log N(0.68 | 0.65, 0.05) β
β β lp[1] = β1.58 β
β β
β Draw m=2: Tβ=27, k=0.20, b=0.12, Ξ½=1.1, Ο=0.06 β
β β expected RI = 0.67 β log N(0.68 | 0.67, 0.06) β
β β lp[2] = β1.39 β
β β
β Draw m=3: Tβ=30, k=0.15, b=0.18, Ξ½=0.8, Ο=0.04 β
β β expected RI = 0.60 β log N(0.68 | 0.60, 0.04) β
β β lp[3] = β3.13 β
ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
β Step 3 β Average over M draws (log-sum-exp trick) β
β β
β We want: log[ (1/3) Γ (e^lp[1] + e^lp[2] + e^lp[3]) ]β
β = log_sum_exp(lp) β log(3) β
β = log_sum_exp(β1.58, β1.39, β3.13) β 1.10 β
β β β1.23 β
β β
β This is the marginalized likelihood for sample n=2. β
ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
Contribution of sample n=2 to log-probability:
prior + likelihood = β2.33 + (β1.23) = β3.56
4. Why the same draw index m matters#
Each calibration draw m represents one self-consistent set of curve parameters
sampled together from the forward posterior. Think of it like a single geochemistβs
βbest guessβ for the entire calibration curve β all parameters in that guess were
estimated together and are correlated.
Forward calibration posterior: 3 draws
βββββββββββββββββββββββββββββββββββββ
Draw β Tβ β k β b β Ο
ββββββΌββββββββΌβββββββΌβββββββΌββββββ
1 β 28.0 β 0.18 β 0.15 β 0.05 β always use ALL from row 1
2 β 27.0 β 0.20 β 0.12 β 0.06 β always use ALL from row 2
3 β 30.0 β 0.15 β 0.18 β 0.04 β always use ALL from row 3
What we do (correct β
):
For draw m=1, use Tβ[1]=28.0, k[1]=0.18, b[1]=0.15, Ο[1]=0.05 together.
What we must NOT do (wrong β):
Mix parameters from different rows β e.g., Tβ[1] + k[3] + b[2].
This would break the correlations among parameters and overestimate uncertainty.
In the Stan code, this constraint is enforced because every parameter uses the
same index [m] inside the inner loop:
for (m in 1:M) {
real mu = b[m] + (1 - b[m])
/ pow(1 + exp(-k[m] * (t_est[n] - t0[m])), 1.0 / v[m]);
// every parameter uses the SAME draw index [m]:
// b[m], k[m], t0[m], v[m] all come from calibration draw m
}
5. The log-sum-exp trick explained#
When averaging probabilities across M calibration draws, we work in log-space to avoid numerical underflow (very small probabilities becoming exactly zero in a computer).
The naive (unstable) way:
average_likelihood = (1/M) Γ (exp(lp[1]) + exp(lp[2]) + exp(lp[3]))
If lp values are very negative (e.g., β300), exp(β300) = 0 in floating-point β lost!
The log-sum-exp way (stable):
log_average_likelihood = log_sum_exp(lp[1], lp[2], lp[3]) β log(M)
Stanβs log_sum_exp shifts all values by the maximum before exponentiating, keeping
numbers in a safe range. The result is mathematically identical but numerically stable.
Visually:
lp = [β1.58, β1.39, β3.13]
β
βΌ
shift by max = β1.39: [β0.19, 0.00, β1.74]
β
βΌ
exp: [0.83, 1.00, 0.18] β safe numbers, no underflow
β
βΌ
sum: 2.01 β log(2.01) = 0.70
β
βΌ
shift back: 0.70 + (β1.39) β log(3) = β1.79
6. Putting it all together: the full picture#
TEXAS Inverse-T Model (marginal + reduce_sum)
INPUT: N=4 RI observations, M=3 calibration draws, 4 CPU cores
Forward posterior (M=3 draws):
ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
β m=1: Tβ=28 k=0.18 b=0.15 Ξ½=0.9 Ο=0.05 β
β m=2: Tβ=27 k=0.20 b=0.12 Ξ½=1.1 Ο=0.06 β
β m=3: Tβ=30 k=0.15 b=0.18 Ξ½=0.8 Ο=0.04 β
ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
Downcore samples: n=1 n=2 n=3 n=4
RI_obs: 0.72 0.68 0.55 0.41
prior_mu_t: 25 20 15 10
β reduce_sum splits across cores β
ββββββββββββββββββββββββ ββββββββββββββββββββββββ
β Core A: n=1, n=2 β β Core B: n=3, n=4 β
β β β β
β For n=1: β β For n=3: β
β prior(T_1) β β prior(T_3) β
β lp[1,2,3] β β lp[1,2,3] β
β log_sum_exp β logMβ β log_sum_exp β logMβ
β β β β
β For n=2: β β For n=4: β
β prior(T_2) β β prior(T_4) β
β lp[1,2,3] β β lp[1,2,3] β
β log_sum_exp β logMβ β log_sum_exp β logMβ
β β β β
β return lp_A β β return lp_B β
ββββββββββββ¬ββββββββββββ ββββββββββββ¬ββββββββββββ
β β
ββββββββββββββ¬βββββββββββββββ
βΌ
target += lp_A + lp_B
(total log-posterior for this proposed T vector)
Stanβs HMC sampler repeats this for thousands of proposed temperature vectors until it has drawn enough samples from the posterior distribution of temperatures.
7. Summary table#
Concept |
What it means in plain language |
|---|---|
|
A function that computes the βscoreβ (log-probability) for one batch of sediment samples |
|
Runs |
|
How many samples go into each batch (1 = smallest, N = no parallelism) |
|
A numerically safe way to average M likelihoods in log-space |
Same index |
Ensures all calibration parameters in each draw came from the same original MCMC sample, preserving their correlations |
|
Extracts the portion of vector |
8. When does parallelism help?#
reduce_sum only speeds things up when compiled with STAN_THREADS=True and
threads_per_chain > 1. Otherwise it runs sequentially (same result, no speedup).
Rule of thumb: parallelism helps most when N is large (hundreds of samples) and the calibration integral (M loop) is expensive. For small reconstructions (N < 20), the overhead of parallelism may outweigh the benefit.
See also: Why marginalization improves inverse sampling and the Stan models overview.