update ESS calculation code to be consistent with bayes-kit

[magland]

There are some unresolved questions here about what should be the exact formulae used. (will need to discuss with @bob-carpenter)

Here is the bayes-kit implemention, which may have some issues:

https://github.com/flatironinstitute/bayes-kit/blob/main/bayes_kit/ess.py

Here's is the current mcmc-monitor implementation, which may need to be adjusted:

https://github.com/flatironinstitute/mcmc-monitor/blob/459e8b6814c745bbf7681f3035b409177a540057/src/MCMCMonitorDataManager/stats/ess.ts

The critical functions in question are

first_neg_pair_start

mcmc-monitor/src/MCMCMonitorDataManager/stats/ess.ts

Lines 72 to 82 in 459e8b6

	export function first_neg_pair_start(chain: number[]): number {
	const N = chain.length
	let n = 0
	while (n + 1 < N) {
	if (chain[n] + chain[n + 1] < 0) {
	return n
	}
	n = n + 1
	}
	return N
	}

and ess_imse
https://github.com/flatironinstitute/bayes-kit/blob/22a3e9ff31f2268f47e13a737dc57c81a26ae917/bayes_kit/ess.py#L99-L135

They way mcmc-monitor does it now, the sigma_sq_hat (I think aka IAT) is never going to be less than 1, which I believe should be a desirable property. But of course we'll want to be consistent with bayes-kit.

tagging: @jsoules @WardBrian

[WardBrian]

I believe you will want to be able to report an ESS greater than the number of draws. This assumes the function people are interested in is the mean, but in that case anticorrelated draws (commonly produced by HMC) are actually better than independent samples

[magland]

I believe you will want to be able to report an ESS greater than the number of draws. This assumes the function people are interested in is the mean, but in that case anticorrelated draws (commonly produced by HMC) are actually better than independent samples

I guess that makes sense. But from what I understand, bayes-kit will essentially compute the area under the main (positve) lobe... and the only time you get IAT<1 is when there is some technical circumstance where a bit of the negative dip after the main lobe is counted as part of the area. The present mcmc-monitor code doesn't ever include that negative piece. I would be surprised if the bayes-kit version is reliably picking up anti-correlated behavior, but perhaps I am misunderstanding something.

[WardBrian]

I can't really speak on the bayes_kit code, just the behavior I expect from tools like stansummary (which I believe is calculated with this code)

[bob-carpenter]

The bayes-kit code just implements the mathematical definition from the paper as described in the Stan reference manual.

The work in Stan's implementation is delegated to the function defined in stan/analyze/mcmc/compute_effective_sample_size.hpp

There's a comment starting at line 98 which says how this is adjusted for antiautocorrleation. The monotonicity condition is implemented differently, though it may work out to the same thing.

P.S. We need to update the doc in our reference manual description of what's being computed to match the actual computation. What happened, I believe, is that we used to use the standard definition from Geyer's paper and then it was updated to deal with the kinds of antiautocorrelated chains we see with Hamiltonian Monte Carlo.

[magland]

I guess we should wait until the bayes_kit implementation is solidified, and then we can just copy it exactly. For now there's a working calculation for ESS that should be pretty close.

[bob-carpenter]

I sat down and did the algebra and I think the confusion is the offset on the pairwise positivity constrain. I originally had the wrong implementation there using pairs (1, 2), (3, 4), .... Instead, what you want to do is take pairs (0, 1), (2, 3), ...

For example, If we have a simple time series model like an AR(1) process with autocorrrelation rho in (-1,1), then we know the overall autocorrelations are

ac[-2] = rho^2
ac[-1] = rho
ac[0] = 1
ac[1] = rho
ac[2] = rho^2

Our estimator for integrated autocorrelation time is

IAT = ... + ac[-2] + ac[-1] + ac[0] + ac[1] + ac[2] + ...

but we have symmetry, ac[n] = ac[-n], so we evaluate with just the positive terms

IAT = -1 + 2 * [ ac[0] + ac[1] + ... ]

where the -1 term is to avoid double counting the lag-zero autocorrelation ac[0] = 1.

Now, if we have rho = -0.9, we'll get something like this:

IAT = -1 + 2 * [ (1 + -0.9) + (.81 + -.72) + ... ]
    < 1

Here's the behavior on the current branch, where sample_ar1 uses (lag 1) autocorrleation rho and generates a sample of size N.

>>> y = sample_ar1(rho = 0, N = 1000)
>>> len(y)
1000
>>> ess(y)
971.3854635513478
>>> 
>>> y = sample_ar1(rho = 0.5, N = 1000)
>>> ess(y)
376.41893467549767
>>> 
>>> y = sample_ar1(rho = -0.5, N = 1000)
>>> ess(y)
2999.0927300430913

So you can see in the last case that estimated ESS > N.

[magland]

Thanks @bob-carpenter. Okay, I will adjust our ess.ts to exactly match what's in your ess.py.

One minor note. Shouldn't the comment on first_neg_pair_start be adjusted?

https://github.com/flatironinstitute/bayes-kit/blob/40a34129a6313056bc4fedc42606578cee6dff77/bayes_kit/ess.py#L65-L67

Right now it reads

Return:
index of first element whose sum with following element is negative, or
the number of elements if there is no such element

But really it should be "Index of the first even-indexed element...

[magland]

@bob-carpenter There is still a typo in ess_imse where prevmin should be min_prev, which I guess will be resolved with PR flatironinstitute/bayes-kit#16

But also, could you double-check the indexing there, because in this function (ess_imse) you are taking pairs (1, 2) (3, 4) etc. which seems inconsistent with first_neg_pair_start. (but maybe that's okay?)

[bob-carpenter]

That's a much more accurate way to doc, so I'll update. I'll also check code again. It should be checking pairs (0, 1), (1, 2), ....

The estimator without the montonic-downward constraint on pair sums will have a bit more variance, but should still be OK.

[jsoules]

Sorry, Bob, I'm still confused. You've said:

That's a much more accurate way to doc, so I'll update.

presumably referring to Jeremy's note that

really it should be "Index of the first even-indexed element...

But then you go on to say:

It should be checking pairs (0, 1), (1, 2),

which looks like it would be every pair, not just the even-indexed ones?

[magland]

@jsoules My understanding is that "That's a much more accurate way to doc, so I'll update." refers to the first_neg_pair_start function, whereas the rest of Bob's paragraph refers to ess_imse.

[bob-carpenter]

Ack, I mean pairs (0, 1), (2, 3), .... Doc is hard, especially before coffee.