Module gabenet.calibration

Functions

def plot_calibration_histogram(histogram: numpy.ndarray[typing.Any, numpy.dtype[+ScalarType]], n_replicates: Optional[int] = None, alpha: float = 0.01)

Analyse the calibration rank histogram for uniformity.

Args

histogram

Rank statistics generated by simulation-based calibration.

n_replicates

Number of replicates used to compute rank histogram, used to compute expected variation.

alpha

Plot expected variation corresponding to quantile range [alpha/2, 1 - alpha/2], i.e., a false positive rate of alpha.

Expand source code

def plot_calibration_histogram(
    histogram: NDArray, n_replicates: Optional[int] = None, alpha: float = 0.01
):
    """Analyse the calibration rank histogram for uniformity.

    Args:
        histogram: Rank statistics generated by simulation-based calibration.
        n_replicates: Number of replicates used to compute rank histogram, used to
            compute expected variation.
        alpha: Plot expected variation corresponding to quantile range
            `[alpha/2, 1 - alpha/2]`, i.e., a false positive rate of `alpha`.
    """
    x_coords = np.arange(len(histogram))

    if n_replicates is not None:
        q_lower = alpha / 2
        q_upper = 1 - alpha / 2
        kwargs = {"n": n_replicates, "p": 1 / len(histogram)}
        y_lower = binom.ppf(q=q_lower, **kwargs)
        y_median = binom.ppf(q=0.5, **kwargs)
        y_upper = binom.ppf(q=q_upper, **kwargs)
        plt.fill_between(
            x_coords,
            np.full_like(x_coords, y_upper),
            np.full_like(x_coords, y_lower),
            alpha=0.5,
        )
        plt.plot(x_coords, np.full_like(x_coords, y_median), "-")

    return plt.bar(x=x_coords, height=histogram, label="posterior", alpha=0.3)

def simulation_based_calibration(forward_fn, posterior_fn: Callable[[Any, numpy.ndarray[Any, numpy.dtype[+ScalarType]], int], numpy.ndarray[Any, numpy.dtype[+ScalarType]]], key, state_fn: Optional[Callable[[numpy.ndarray[Any, numpy.dtype[+ScalarType]]], numpy.ndarray[Any, numpy.dtype[+ScalarType]]]] = None, n_replicates: int = 200, n_posterior_samples: int = 10) ‑> numpy.ndarray[typing.Any, numpy.dtype[+ScalarType]]

Simulation based calibration to validate model implementation.

Verify posterior sampler by comparing prior with posterior distribution. Relies on the identity: p(w) = integral dy dw' p(w | y) p(y| w') p(w'),

where p(w') is the prior and p(w|y) the posterior distribution.

Args

forward_fn

Haiku function that generates an ancestral sample.

posterior_fn

Function that samples from the posterior. Takes as arguments a PRNG key, the observed data, and number of posterior samples. Returns a pytree where the two leading axes refer to the number of chains and corresponding samples, respectively.

state_fn

Arbitrary function that maps a posterior state sample (first argument) to a pytree. When None, directly compute rank statistics based on state.

n_replicates

Number of replicates of the data. Controls the histogram's bin variance.

n_posterior_samples

Number of posterior samples per replicate. Controls the number of histogram bins.

Returns

A pytree histogram of rank statistics (of size n_posterior_samples + 1) per element returned by state_fn. Is uniformly distributed for a correctly implemented prior and posterior sampler.

Expand source code

def simulation_based_calibration(
    forward_fn,
    posterior_fn: Callable[[Any, NDArray, int], NDArray],
    key,
    state_fn: Optional[Callable[[NDArray], NDArray]] = None,
    n_replicates: int = 200,
    n_posterior_samples: int = 10,
) -> NDArray:
    """Simulation based calibration to validate model implementation.

    Verify posterior sampler by comparing prior with posterior distribution.
    Relies on the identity: p(w) = integral dy dw' p(w | y) p(y| w') p(w'),

    where p(w') is the prior and p(w|y) the posterior distribution.

    Args:
        forward_fn: Haiku function that generates an ancestral sample.
        posterior_fn: Function that samples from the posterior. Takes as arguments
            a PRNG key, the observed data, and number of posterior samples. Returns a
            pytree where the two leading axes refer to the number of chains and
            corresponding samples, respectively.
        state_fn: Arbitrary function that maps a posterior `state` sample
            (first argument) to a pytree. When `None`, directly compute rank statistics
            based on `state`.
        n_replicates: Number of replicates of the data. Controls the histogram's bin
            variance.
        n_posterior_samples: Number of posterior samples per replicate. Controls the
            number of histogram bins.

    Returns:
        A pytree histogram of rank statistics (of size `n_posterior_samples + 1`) per
        element returned by `state_fn`. Is uniformly distributed for a correctly
        implemented prior and posterior sampler.

    """
    if state_fn is None:
        state_fn = lambda x: x  # No-op

    key, rng_key = random.split(key)
    param, state_init = forward_fn.init(rng_key)
    # Jit and peel off the `param` parameter.
    _forward_fn = jax.jit(lambda s, k: forward_fn.apply(param, s, k))  # type: ignore

    rank_statistics = []
    for _ in range(n_replicates):
        # 1) Generate a forward sample.
        key, rng_key = random.split(key)
        X_generative, state = _forward_fn(state_init, rng_key)
        # Convert unobserved state to a real number.
        f_generative = state_fn(state)

        # 2) Sample from posterior using observations.
        key, rng_key = random.split(key)
        posterior_states = posterior_fn(rng_key, X_generative, n_posterior_samples)

        # 3) Pool samples from different chains and apply `state_fn`.
        posterior_states = tree_map(
            lambda x: x.reshape(-1, *x.shape[2:]), posterior_states
        )
        f_posterior = state_fn(posterior_states)

        # 4) Compute rank statistic equation in Sec. 4.1 Ref. [1].
        rank_stat = tree_map(
            lambda a, b: (a < b).sum(axis=0), f_posterior, f_generative
        )
        rank_statistics.append(rank_stat)

    # 4) Build histogram of rank statistics.
    # For each ranking element r(f(theta_1, ..., f(theta_L)|theta')) returned by
    # `state_fn`, increment its histogram.
    # (Count the ranks using a flattened representation of the `state_fn` tensor.)
    tree_histogram = tree_map(
        lambda leave: np.zeros([n_posterior_samples + 1, leave.size]), rank_stat
    )
    hist_leaves = tree_leaves(tree_histogram)
    for rank_stat in rank_statistics:
        rank_leaves = tree_leaves(rank_stat)
        for r_leaf, h_leaf in zip(rank_leaves, hist_leaves):
            # In the flat representation, the elements of the tensor are indexed as 0, 1, ...
            h_leaf[r_leaf.flatten(), np.arange(r_leaf.size)] += 1

    # Reshape histogram consistent with the tensor shape returned by `state_fn`.
    histogram = tree_map(
        lambda h, r: h.reshape((-1,) + r.shape), tree_histogram, rank_stat
    )
    return histogram