Convergence statistic in get_rstat() may not match Gelman–Rubin R̂

[fcotizelati]

In mind_the_gaps/gpmodelling.py, the get_rstat() method is presented as a Gelman–Rubin–style convergence diagnostic. However, the current calculation appears to use a ratio of within-chain variance to the variance of all pooled samples. This differs from the standard Gelman–Rubin R̂ definition, which is based on a combination of within-chain variance and between-chain variance computed from the chain means.

Because of this, the values returned by get_rstat() may not reliably reflect the usual interpretation of R̂ (i.e. "close to 1 indicates convergence"). Updating the implementation to follow the standard R̂ formula would make the convergence assessment clearer.

[andresgur]

Thanks for the report @fcotizelati ! Now I see I misunderstood the GR diagnostic. Tell me if you are happy with this implementation, note I have modified B as we divide later by nsteps again.

    def get_rstat(self, burnin: int = None) -> ArrayLike:
        """Calculate convergence criterion from Gelman & Rubin 1992; Gelman et al. 2004.
        Values close to 1 indicate convergence.

        Parameters
        ----------
        burnin
            Number of samples to burn prior to the estimation. By default
            uses the autocorrelation time to estimate a suitable number

        Returns
        -------
        ArrayLike
            Returns an array of the value of the statistic per chain
        """
        if self._sampler is None:
            raise ValueError(
                "Posteriors have not been derived. Please run derive_posteriors prior to populate the attributes."
            )
        if burnin is None:
            # discard 10 times the mean autocorrelation time
            burnin = int(np.mean(self.tau)) * 10
        samples = self._sampler.get_chain(discard=burnin)

        chain_means = np.mean(
            samples, axis=0
        )  # this has nwalkers, ndim (one per chain and param)

        chain_variances = np.var(
            samples, axis=0, ddof=1
        )  # this has nsteps, ndim (one per step and param)

        nsteps = samples.shape[0]

        B = np.var(
            chain_means, axis=0, ddof=1
        )  # this has ndim (one per param) / this is B but without the nsteps, because we divide later by nsteps

        W = np.mean(chain_variances, axis=0)  # this has ndim (one per param)

        V_hat_W = (
            nsteps - 1
        ) / nsteps + B / W  # no need to divide by nsteps here because we are not multiplying by nsteps in the B term, so it is already included in the formula
        return (V_hat_W) ** 0.5

If happy I'll push the changes

[fcotizelati]

Thank you @andresgur ! Everything seems correct to me.