Graph: the hierarchical model

A Graph defines a joint probability law through conditional distributions and deterministic transformations.

From population to observations

A population parameter can govern several latent variables; deterministic operators map those variables into predicted signals; observation distributions connect the signals to measured data. Hierarchy lives in these conditional dependencies. Constants and deterministic nodes participate in the calculation without adding a sampled coordinate of their own.

Normal distributions below use variance as their second mathematical parameter; the Python API uses standard deviation.

Joint density

For this model, with fixed noise scales and observations y, the ordinary site factors give:

The deterministic signal is substituted into the likelihood. The posterior conditions this same joint model on y; it does not reverse the declared arrows. General graphs may have multiple parents and richer dependence than this simple chain.

"""A flowchart whose dependencies define a hierarchical probability model."""

import jax.numpy as jnp
import numpy as np
import numpyro.distributions as dist

import bayesmith as bs


def model(data):
    population_mean = bs.sample("population_mean", lambda: dist.Normal(0.0, 2.0))
    units = bs.plate("units", len(data))
    level = bs.sample(
        "level", lambda mean: dist.Normal(mean, 0.5), population_mean, plate=units
    )
    signal = bs.det("signal", lambda value: 2.0 * value, level, plate=units)
    bs.observe(
        "data", lambda mean: dist.Normal(mean, 0.2), signal, obs=data, plate=units
    )


def normal_log_density(value, mean, scale):
    """Independent scalar Gaussian formula, summed over repeated units."""
    return np.sum(
        -0.5 * np.log(2.0 * np.pi * scale**2) - 0.5 * ((value - mean) / scale) ** 2
    )


def main():
    data = jnp.array([0.5, 1.7, -0.3])
    graph = bs.trace(model, data)
    point = {"population_mean": jnp.asarray(0.3), "level": jnp.array([0.2, 0.8, -0.2])}
    actual = bs.log_joint(graph, point)
    mean, levels = (
        float(point["population_mean"]),
        np.asarray(point["level"], dtype=float),
    )
    expected = (
        normal_log_density(mean, 0.0, 2.0)
        + normal_log_density(levels, mean, 0.5)
        + normal_log_density(np.asarray(data, dtype=float), 2.0 * levels, 0.2)
    )
    np.testing.assert_allclose(actual, expected, rtol=1e-6)
    print("Graph log joint:", float(actual))
    print("Independent hierarchical factorization:", expected)


if __name__ == "__main__":
    main()

Download this example

This executable example compares log_joint with an independent Gaussian calculation. log_joint returns the joint log-density at the supplied latent values and observed data; it is not a normalized posterior density.

Node roles

Declare these inside the function passed to trace. Tracing executes the declaration once to build a static directed acyclic graph; it does not draw a simulated dataset. Declare parents before their children. Parent arguments are node references. Distribution constructors and numerical operations run inside the attached callables. Model changes should produce a newly traced graph and fresh fingerprints.

joint_prior adds a factor to the site contributions; it does not automatically replace their priors. Use flat site measures for covered latents when the graph-level factor supplies their prior, and respect that prior object's contract. Ordinary conditional hierarchies such as the example above need no joint_prior.

Separate batch, event and draw axes

NumPyro distributions distinguish batch shape from event shape. A vector of independent normal observations and one multivariate-normal observation may have the same visible array shape but different density grouping and covariance semantics. Use .to_event(...) only when the model needs that event interpretation.

plate declarations record a named repetition axis; they do not justify treating dependent data as independent. The current Graph supports at most one plate per node, without nested plates. Units may be conditionally independent given a population parameter while remaining marginally correlated through that shared parameter. Posterior and artifact arrays have additional sample axes. Read each NamedArray's dimensions and method contract before flattening chain and draw axes.

Represent missing observations with masks

bs.observe(
    "signal", lambda x: dist.Normal(x, 0.5), level,
    obs=jnp.array([0.8, 0.0, 1.2]),
    mask=jnp.array([True, False, True]),
)

True means the sample was observed. Supply a boolean mask shaped like the observation. The finite placeholder at a masked entry is not data used in the likelihood. Do not replace a missing measurement by infinite noise or introduce NaNs into intermediate calculations.

Changing data or masks changes the conditioning fingerprint. A posterior fitted under an old mask is not interchangeable with a new fit. Held-out checks use the distinction between included and withheld observations.

Structural declarations

Automatic affinity analysis can certify supported affine computations. A user declaration alone must not make a nonlinear model exact. Gaussian prior eligibility and covariance dependence are checked separately from the mean's affinity. Unsupported operations or derivative semantics can leave a claim unresolved.

For correlated noise, declare its probability law rather than hiding its covariance in a scale vector. For complex variables, consult ComplexNormal and its real-coordinate convention. Continue with route eligibility.