bamojax.marginal_likelihoods.utility

iid_likelihood(L)

We typically have multiple observations and assume the likelihood factorizes as:

\[ \log p\left(Y \mid \theta\right) = \sum_{i=1}^N \log p\left(y_i \mid \theta\right) \enspace. \]

Source code in bamojax/marginal_likelihoods/utility.py

def iid_likelihood(L: Callable):
    r"""

    We typically have multiple observations and assume the likelihood factorizes 
    as: 

    $$    
        \log p\left(Y \mid \theta\right) = \sum_{i=1}^N \log p\left(y_i \mid \theta\right) \enspace.
    $$

    """
    return lambda x: jnp.sum(L()(x))

flatten_dict_to_array(samples)

Bamojax states are dictionaries, with entries per model variable. Here we flatten them so the proposal distribution can be one single multivariate distribution.

Source code in bamojax/marginal_likelihoods/utility.py

def flatten_dict_to_array(samples: dict):
    """ Bamojax states are dictionaries, with entries per model variable. Here we flatten them so the 
    proposal distribution can be one single multivariate distribution.

    """

    leaves, treedef = jax.tree_util.tree_flatten(samples)

    Ns = [leaf.shape[0] for leaf in leaves]
    N = Ns[0]

    flat_leaves = [leaf.reshape(N, -1) for leaf in leaves]

    sizes = [f.shape[1] for f in flat_leaves]
    cumulative_sizes = jnp.cumsum(jnp.array(sizes))
    samples_flattened = jnp.concatenate(flat_leaves, axis=-1)  # (N, D)

    leaf_shapes = [leaf.shape[1:] for leaf in leaves]

    def unravel_one_sample(vec: jnp.ndarray):

        splits = jnp.split(vec, cumulative_sizes[:-1])
        reshaped = [v.reshape(s) for v, s in zip(splits, leaf_shapes)]
        return jax.tree_util.tree_unflatten(treedef, reshaped)

    return samples_flattened, unravel_one_sample