Source code for dit.bgm

"""
Joint distributions from Bayesian graphical models, aka Bayesian networks.

Currently, this does not support factorings which include more than one
random variable, such as P(X) P(Y,Z|X). Effectively, we assume each random
variable can be factored so that it is "alone", as in P(X) P(Y|X,Z) P(Z|X)
or P(X) P(Z|X,Y) P(Y|X). While one can always do this in probability, it may
not be the case that this corresponds to how it is causally generated.

An example of this is an edge-emitting hidden Markov model: P(X_0, S_1 | S_0).
There is no causal relationship between X_0 and S_1...they are jointly
generated from S_0.
"""

import numpy as np

import dit

__all__ = ("distribution_from_bayesnet",)


def sanitize_inputs(digraph, nodes, attr):
    """
    Quick sanity checks on the input.

    """
    all_nodes = set(digraph.nodes())
    ops = dit.math.get_ops("linear")
    is_callable = []
    for rv in digraph:
        # Make sure we have dists for each node.

        try:
            val = digraph.nodes[rv][attr]
        except KeyError as err:
            msg = f"Node {rv} is missing its distributions."
            raise ValueError(msg) from err

        if callable(val):
            is_callable.append(1)
            continue
        else:
            is_callable.append(0)

        # Make sure the rv names are appropriate.
        if digraph.in_degree(rv) == 0:
            # No parents
            dists = [val]
        else:
            # A distribution for each value of the parents.

            # This helps find mistakes more easily!
            # We need [parents, dists] rather than a single distribution.
            if isinstance(val, dit.Distribution):
                msg = "Node {} has an invalid dist specification."
                raise Exception(msg.format(rv))

            if isinstance(val, dict):
                dists = val.values()
            else:
                outcomes, dists = val
                if len(outcomes) != len(dists):
                    msg = "Node {} has an invalid dist specification."
                    raise Exception(msg.format(rv))

        # No worries if this gets overwritten with each rv, as it had better
        # be the same for each rv.
        ops = next(iter(dists)).ops

    # Get a good set of random variable names.
    if nodes is None:
        rv_names = sorted(digraph.nodes())
    else:
        if len(nodes) != len(all_nodes):
            msg = f"`nodes` is missing required nodes: {nodes}"
            raise ValueError(msg)
        if set(nodes) != all_nodes:
            msg = "`set(nodes)` does not contain all required nodes."
            raise ValueError(msg)
        rv_names = nodes

    n_callable = sum(is_callable)
    if 0 < n_callable < len(rv_names):
        # Then some distributions were callable while others were not.
        msg = "All distributions must be callable if any are."
        raise Exception(msg)

    all_callable = bool(n_callable)

    return rv_names, ops, all_callable


def build_samplespace(digraph, rv_names, attr):
    """
    Builds the sample space for the joint distribution.

    """
    sample_spaces = []
    product = None
    for rv in rv_names:
        if digraph.in_degree(rv) == 0:
            # No parents
            dist = digraph.nodes[rv][attr]
        else:
            # Grab the first distribution.
            val = digraph.nodes[rv][attr]
            try:
                dist = next(iter(val.values()))
            except AttributeError:
                print(val)
                dist = val[1][0]

        # Since we are assuming each rv is completely alone and factored, we
        # can just take the alphabet. This will need to change eventually.
        try:
            alphabet = dist._sample_space.alphabets[0]
        except AttributeError:
            alphabet = list(dist._sample_space)
        sample_spaces.append(alphabet)
    else:
        # Use the last dist to get a product function.
        # We'll assume they are all the same.
        product = dist._product

    ss = dit.CartesianProduct(sample_spaces, product=product)
    return ss


def build_pfuncs(digraph, rv_names, attr, outcome_ctor):
    """
    Build probability functions for each rv.

    The function takes a random variable and the joint outcome and returns
    the probability contribution to the joint probability.


    """
    rv_index = dict(zip(rv_names, range(len(rv_names)), strict=True))
    pfuncs = {}
    parents_index = {}

    for rv in rv_names:
        parents = list(digraph.predecessors(rv))
        parents.sort(key=rv_index.__getitem__)
        parents_index[rv] = parents

        val = digraph.nodes[rv][attr]
        if callable(val):
            pfuncs[rv] = val
            continue

        if not parents:
            # Immediately bind variables since we are in a for loop.
            # http://docs.python-guide.org/en/latest/writing/gotchas/#late-binding-closures
            def prob(outcome, dist=val, rv=rv):
                rv_outcome = outcome_ctor([outcome[rv_index[rv]]])
                return dist[rv_outcome]

        else:
            if isinstance(val, dict):
                val.values()
                dists = val
            else:
                outcomes, dists = val
                dists = dict(zip(outcomes, dists, strict=True))

            def prob(outcome, dists=dists, parents=parents, rv=rv):
                node_outcome = outcome_ctor([outcome[rv_index[rv]]])
                parent_vals = [outcome[rv_index[parent]] for parent in parents]
                parent_outcome = outcome_ctor(parent_vals)
                dist = dists[parent_outcome]
                return dist[node_outcome]

        pfuncs[rv] = prob

    # Create a function for callable dists that returns the node value and
    # the parents via a dict.
    def get_values(rv, outcome):
        node_val = outcome[rv_index[rv]]
        parents = parents_index[rv]
        parent_vals = [outcome[rv_index[parent]] for parent in parents]
        parents = dict(zip(parents, parent_vals, strict=True))
        return node_val, parents

    return pfuncs, get_values


[docs] def distribution_from_bayesnet(digraph, nodes=None, sample_space=None, attr="dist"): """ Returns a distribution built from a Bayesian network. Each node represents a random variable ``X_i``. Every node must store its conditional probability distribution ``P(X_i | Y_i)`` where ``Y_i`` represents the parents of ``X_i``. If a node has no in-degree, then it must store the probability distribution ``P(X_i)``. Parameters ---------- digraph : NetworkX digraph A directed graph, representing the Bayesian graphical model. nodes : list, None The order of the nodes that will determine the random variable order. If `None`, then we use `sorted(digraph.nodes())`, which assumes the nodes are sortable. The reason we assume they are sortable is because the parent values must correspond to the node order and thus, we need an unambiguous ordering that the user could have known ahead of time. sample_space : SampleSpace If provided, this specifies the outcomes of the distribution to be constructed. The distributions stored on the nodes are assumed to be compatible with this space. If functions are stored on the nodes, then this parameter must be provided. attr : str The attribute for each node that holds the conditional distributions. The attribute value can take a variety of forms. It can be a function. The function must take two arguments. The first is the value of random variable for the current node. The second is a dictionary of keyed by parents whose values are the values of the random variables corresponding to the parents. The function should return the probability P(node_val|parent_vals). It can be a list, such as [parents, dists], that holds the parents and the conditional distributions: `dists[i] = P(X | Y_i = parents[i])`. It can also be a dict-like structure so that `dists[y]` is a distribution representing P(X | Y = y)`. If the node has no in-degree, then the attribute value should store the distribution only. When using distributions, each should have random variable names assigned that match the nodes in the graph, or alternatively, all nodes in the graph should be integers and then random variable names are not necessary. The order of elements within these lists (or the dict) does not matter, but for nodes that have parents, the order of the random variables that specify the parents, must match the order of ``nodes``. So for example, if the node order is [2, 1, 0] and node 1 has parents 0 and 2. Then the parents for node 1 will be such that the first element corresponds to node 2 and the second to node 0, since 2 precedes 0 in the node order. All distributions will be assumed to have the same base, and this base will determine the base of the constructed distribution. Returns ------- dist : Distribution The joint distribution. Examples -------- The Xor logic gate has the following structure: >>> g = nx.DiGraph() >>> g.add_edge(0, 2) >>> g.add_edge(1, 2) Let's add distributions to it using functions. >>> uniform = lambda node_val, parents: 0.5 >>> def xor(node_val, parents): ... if '1' == parents[0] == parents[1]: ... desired_output = '1' ... else: ... desired_output = '0' ... return int(node_val == desired_output) ... >>> g.nodes[0]['dist'] = uniform >>> g.nodes[1]['dist'] = uniform >>> g.nodes[2]['dist'] = xor >>> ss = ['000', '001', '010', '011', '100', '101', '110', '111'] >>> d = dit.distribution_from_bayesnet(g, sample_space=ss) Alternatively, we could add distributions using Distribution objects. >>> uniform = dit.uniform_distribution(1, 2) >>> sample_space1 = [(0,), (1,)] >>> one = dit.Distribution(sample_space1, [0, 1]) >>> zero = dit.Distribution(sample_space1, [1, 0]) >>> sample_space2 = [(0, 0), (0, 1), (1, 0), (1, 1)] >>> xor = [ sample_space2, [zero, one, one, zero]] >>> g.nodes[0]['dist'] = uniform >>> g.nodes[1]['dist'] = uniform >>> g.nodes[2]['dist'] = xor >>> d = dit.distribution_from_bayesnet(g) We can add noise whenever the output would normally be 1. >>> noisy = dit.Distribution(sample_space, [.1, .9]) >>> dists = [zero, noisy, noisy, zero] >>> g.nodes[2]['dist'][1] = dists >>> d = dit.distribution_from_bayesnet(g) """ rv_names, ops, callables = sanitize_inputs(digraph, nodes, attr) if callables: if sample_space is None: msg = "sample_space must be specified since the " msg += "distributions were callable." raise ValueError(msg) if not isinstance(sample_space, dit.SampleSpace): sample_space = dit.SampleSpace(sample_space) else: sample_space = build_samplespace(digraph, rv_names, attr) ctor = sample_space._outcome_ctor pfuncs, get_values = build_pfuncs(digraph, rv_names, attr, ctor) outcomes = list(sample_space) mult = ops.mult_reduce if callables: pmf = [mult(np.asarray([pfuncs[rv](*get_values(rv, outcome)) for rv in rv_names])) for outcome in outcomes] else: pmf = [mult(np.asarray([pfuncs[rv](outcome) for rv in rv_names])) for outcome in outcomes] pmf = ops.normalize(np.asarray(pmf)) str_names = [str(n) for n in rv_names] dist = dit.Distribution(outcomes, pmf, rv_names=str_names, base=ops.get_base()) dist.set_rv_names(str_names) return dist