path: root/beliefs/models/belief_update_node_model.py

import copy
from enum import Enum
import numpy as np
import itertools
from functools import reduce

import networkx as nx

from beliefs.models.base_models import BayesianModel
from beliefs.factors.discrete_factor import DiscreteFactor
from beliefs.factors.bernoulli_or_cpd import BernoulliOrCPD
from beliefs.factors.bernoulli_and_cpd import BernoulliAndCPD


class InvalidLambdaMsgToParent(Exception):
    """Computed invalid lambda msg to send to parent."""
    pass


class MessageType(Enum):
    LAMBDA = 'lambda'
    PI = 'pi'


class BeliefUpdateNodeModel(BayesianModel):
    """
    A Bayesian model storing nodes (e.g. Node or BernoulliOrNode) implementing properties
    and methods for Pearl's belief update algorithm.

    ref: "Fusion, Propagation, and Structuring in Belief Networks"
          Artificial Intelligence 29 (1986) 241-288

    """
    def __init__(self, nodes_dict):
        """
        Args
            nodes_dict: dict
                a dict key, value pair as {label_id: instance_of_node_class_or_subclass}
        """
        super().__init__(edges=self._get_edges_from_nodes(nodes_dict.values()),
                         variables=list(nodes_dict.keys()),
                         cpds=[node.cpd for node in nodes_dict.values()])

        self.nodes_dict = nodes_dict

    @classmethod
    def init_from_edges(cls, edges, node_class):
        """
        Create model from edges where all nodes are a from the same node class.

        Args
            edges: list,
                list of edge tuples of form [('parent', 'child')]
            node_class: Node class or subclass,
                class from which to create all the nodes automatically from edges,
                e.g. BernoulliAndNode or BernoulliOrNode
        """
        nodes = set()
        g = nx.DiGraph(edges)

        for label in set(itertools.chain(*edges)):
            node = node_class(label_id=label,
                              children=list(g.successors(label)),
                              parents=list(g.predecessors(label)))
            nodes.add(node)
        nodes_dict = {node.label_id: node for node in nodes}
        return cls(nodes_dict)

    @staticmethod
    def _get_edges_from_nodes(nodes):
        """
        Return list of all directed edges in nodes.

        Args
            nodes: iterable,
                iterable of objects of the Node class or subclass
        Returns
            edges: list,
                list of edge tuples
        """
        edges = set()
        for node in nodes:
            if node.parents:
                edge_tuples = zip(node.parents, [node.label_id]*len(node.parents))
                edges.update(edge_tuples)
        return list(edges)

    def set_boundary_conditions(self):
        """
        Set boundary conditions for nodes in the model.

          1. Root nodes: if x is a node with no parents, set Pi(x) = prior
             probability of x.

          2. Leaf nodes: if x is a node with no children, set Lambda(x)
             to an (unnormalized) unit vector, of length the cardinality of x.
        """
        for root in self.get_roots():
            self.nodes_dict[root].update_pi_agg(self.nodes_dict[root].cpd.values)

        for leaf in self.get_leaves():
            self.nodes_dict[leaf].update_lambda_agg(np.ones([self.nodes_dict[leaf].cardinality]))

    @property
    def all_nodes_are_fully_initialized(self):
        """
        Check if all nodes in the model are initialized, i.e. lambda and pi messages and
        lambda_agg and pi_agg are not None for every node.

        Returns
            bool, True if all nodes in the model are initialized, else False.
        """
        for node in self.nodes_dict.values():
            if not node.is_fully_initialized:
                return False
        return True

    def copy(self):
        """Return a copy of the model."""
        copy_nodes = copy.deepcopy(self.nodes_dict)
        copy_model = self.__class__(nodes_dict=copy_nodes)
        return copy_model


class Node:
    """
    A node in a DAG with methods to compute the belief (marginal probability of
    the node given evidence) and compute pi/lambda messages to/from its neighbors
    to update its belief.

    Implemented from Pearl's belief propagation algorithm for polytrees.
    """
    def __init__(self, children, cpd):
        """
        Args
            children: list,
                list of strings
            cpd: an instance of TabularCPD or one of its subclasses,
                e.g. BernoulliOrCPD or BernoulliAndCPD
        """
        self.label_id = cpd.variable
        self.children = children
        self.parents = cpd.parents
        self.cardinality = cpd.cardinality[0]
        self.cpd = cpd

        self.pi_agg = self._init_factors_for_variables([self.label_id])[self.label_id]
        self.lambda_agg = self._init_factors_for_variables([self.label_id])[self.label_id]

        self.pi_received_msgs = self._init_factors_for_variables(self.parents)
        self.lambda_received_msgs = \
                {child: self._init_factors_for_variables([self.label_id])[self.label_id]
                 for child in children}


    @property
    def belief(self):
        """
        Calculate the marginal probability of the variable from its aggregate values.

        Returns
            belief, an np.array of ndim 1 and shape (self.cardinality,)
        """
        if any(self.pi_agg.values) and any(self.lambda_agg.values):
            belief = (self.lambda_agg * self.pi_agg).values
            return self._normalize(belief)
        else:
            return None

    def _normalize(self, value):
        return value/value.sum()

    def _init_factors_for_variables(self, variables):
        """
        Args
            variables: list,
                 list of ints/strings, e.g. the single node variable or list
                 of parent ids of the node
        Returns
            factors: dict,
                where the dict has key, value pair as {variable_id: instance of a DiscreteFactor},
                where DiscreteFactor.values is an np.array of ndim 1 and
                shape (cardinality of variable_id,)
        """
        variables = list(variables)
        factors = {}

        for var in variables:
            if self.cpd.state_names is not None:
                state_names = {var: self.cpd.state_names[var]}
            else:
                state_names = None

            cardinality = self.cpd.cardinality[self.cpd.variables.index(var)]
            factors[var] = DiscreteFactor(variables=[var],
                                          cardinality=[cardinality],
                                          values=None,
                                          state_names=state_names)
        return factors

    def _return_msgs_received_for_msg_type(self, message_type):
        """
        Args
            message_type: MessageType enum
        Returns
            msg_values: list,
                list of DiscreteFactors with property `values` containing
                the values of the messages (np.arrays)
        """
        if message_type == MessageType.LAMBDA:
            msgs = [msg for msg in self.lambda_received_msgs.values()]
        elif message_type == MessageType.PI:
            msgs = [msg for msg in self.pi_received_msgs.values()]
        return msgs

    def validate_and_return_msgs_received_for_msg_type(self, message_type):
        """
        Check that all messages have been received from children (parents).
        Raise error if all messages have not been received.  Called
        before calculating lambda_agg (pi_agg).

        Args
            message_type: MessageType enum
        Returns
            msgs: list,
                list of DiscreteFactors with property `values` containing
                the values of the messages (np.arrays)
        """
        msgs = self._return_msgs_received_for_msg_type(message_type)

        if any(msg.values is None for msg in msgs):
            raise ValueError(
                "Missing value for {msg_type} msg from child: can't compute {msg_type}_agg."
                .format(msg_type=message_type.value)
            )
        else:
            return msgs

    def compute_pi_agg(self):
        """
        Compute and update pi_agg, the prior probability, given the current state
        of messages received from parents.
        """
        if len(self.parents) == 0:
            self.update_pi_agg(self.cpd.values)
        else:
            factors_to_multiply = [self.cpd]
            pi_msgs = self.validate_and_return_msgs_received_for_msg_type(MessageType.PI)
            factors_to_multiply.extend(pi_msgs)

            factor_product = reduce(lambda phi1, phi2: phi1*phi2, factors_to_multiply)
            self.update_pi_agg(factor_product.marginalize(self.parents).values)
            pi_msgs = self.validate_and_return_msgs_received_for_msg_type(MessageType.PI)

    def compute_lambda_agg(self):
        """
        Compute and update lambda_agg, the likelihood, given the current state
        of messages received from children.
        """
        if len(self.children) != 0:
            lambda_msg_values = [
                msg.values for msg in
                self.validate_and_return_msgs_received_for_msg_type(MessageType.LAMBDA)
            ]
            self.update_lambda_agg(reduce(np.multiply, lambda_msg_values))

    def update_pi_agg(self, new_value):
        self.pi_agg.update_values(new_value)

    def update_lambda_agg(self, new_value):
        self.lambda_agg.update_values(new_value)

    def _update_received_msg_by_key(self, received_msg_dict, key, new_value, message_type):
        if key not in received_msg_dict.keys():
            raise ValueError("Label id '{}' to update message isn't in allowed set of keys: {}"
                             .format(key, received_msg_dict.keys()))

        if not isinstance(new_value, np.ndarray):
            raise TypeError("Expected a new value of type numpy.ndarray, but got type {}"
                            .format(type(new_value)))

        if message_type == MessageType.LAMBDA:
            expected_shape = (self.cardinality,)
        elif message_type == MessageType.PI:
            expected_shape = (self.cpd.cardinality[self.cpd.variables.index(key)],)

        if new_value.shape != expected_shape:
            raise ValueError("Expected new value to be of dimensions ({},) but got {} instead"
                             .format(expected_shape, new_value.shape))
        received_msg_dict[key].update_values(new_value)

    def update_pi_msg_from_parent(self, parent, new_value):
        self._update_received_msg_by_key(received_msg_dict=self.pi_received_msgs,
                                         key=parent,
                                         new_value=new_value,
                                         message_type=MessageType.PI)

    def update_lambda_msg_from_child(self, child, new_value):
        self._update_received_msg_by_key(received_msg_dict=self.lambda_received_msgs,
                                         key=child,
                                         new_value=new_value,
                                         message_type=MessageType.LAMBDA)

    def compute_pi_msg_to_child(self, child_k):
        """
        Compute pi_msg to child.

        Args
            child_k: string or int,
                the label_id of the child receiving the pi_msg
        Returns
            np.array of ndim 1 and shape (self.cardinality,)
        """
        lambda_msg_from_child = self.lambda_received_msgs[child_k].values
        if lambda_msg_from_child is not None:
            with np.errstate(divide='ignore', invalid='ignore'):
                # 0/0 := 0
                return self._normalize(
                    np.nan_to_num(np.divide(self.belief, lambda_msg_from_child)))
        else:
            raise ValueError("Can't compute pi message to child_{} without having received a lambda message from that child.")

    def compute_lambda_msg_to_parent(self, parent_k):
        """
        Compute lambda_msg to parent.

        Args
            parent_k: string or int,
                the label_id of the parent receiving the lambda_msg
        Returns
            np.array of ndim 1 and shape (cardinality of parent_k,)
        """
        if np.array_equal(self.lambda_agg.values, np.ones([self.cardinality])):
            return np.ones([self.cardinality])
        else:
            factors_to_multiply = [self.cpd]
            pi_msgs_excl_k = [msg for par_id, msg in self.pi_received_msgs.items()
                              if par_id != parent_k]
            factors_to_multiply.extend(pi_msgs_excl_k)
            factor_product = reduce(lambda phi1, phi2: phi1*phi2, factors_to_multiply)
            new_factor = factor_product.marginalize(list(set(self.parents) - set([parent_k])))
            lambda_msg_to_k = (self.lambda_agg * new_factor).marginalize([self.lambda_agg.variables[0]])
            return self._normalize(lambda_msg_to_k.values)

    @property
    def is_fully_initialized(self):
        """
        Returns True if all lambda and pi messages and lambda_agg and
        pi_agg are not None, else False.
        """
        lambda_msgs = self._return_msgs_received_for_msg_type(MessageType.LAMBDA)
        if any(msg is None for msg in lambda_msgs):
            return False

        pi_msgs = self._return_msgs_received_for_msg_type(MessageType.PI)
        if any(msg is None for msg in pi_msgs):
            return False

        if (self.pi_agg.values is None) or (self.lambda_agg.values is None):
            return False

        return True


class BernoulliOrNode(Node):
    """
    A node in a DAG associated with a Bernoulli random variable with state_names ['False', 'True']
    and conditional probability distribution described by 'Or' logic.
    """
    def __init__(self, label_id, children, parents):
        super().__init__(children=children, cpd=BernoulliOrCPD(label_id, parents))

    def compute_pi_agg(self):
        """
        Compute and update pi_agg, the prior probability, given the current state
        of messages received from parents.  Sidestep explicit factor product and
        marginalization.
        """
        if len(self.parents) == 0:
            self.update_pi_agg(self.cpd.values)
        else:
            pi_msgs = self.validate_and_return_msgs_received_for_msg_type(MessageType.PI)
            parents_p0 = [p.get_value_for_state_vector({p.variables[0]: 'False'})
                          for p in pi_msgs]
            p_0 = reduce(lambda x, y: x*y, parents_p0)
            p_1 = 1 - p_0
            self.update_pi_agg(np.array([p_0, p_1]))

    def compute_lambda_msg_to_parent(self, parent_k):
        """
        Compute lambda_msg to parent.  Sidestep explicit factor product and
        marginalization.

        Args
            parent_k: string or int,
                the label_id of the parent receiving the lambda_msg
        Returns
            np.array of ndim 1 and shape (cardinality of parent_k,)
        """
        if np.array_equal(self.lambda_agg.values, np.ones([self.cardinality])):
            return np.ones([self.cardinality])
        else:
            # TODO: cleanup this validation
            _ = self.validate_and_return_msgs_received_for_msg_type(MessageType.PI)
            p0_excluding_k = [p.get_value_for_state_vector({p.variables[0]: 'False'})
                              for par_id, p in self.pi_received_msgs.items() if par_id != parent_k]
            p0_product = reduce(lambda x, y: x*y, p0_excluding_k, 1)

            lambda_agg_0 = self.lambda_agg.get_value_for_state_vector({self.label_id: 'False'})
            lambda_agg_1 = self.lambda_agg.get_value_for_state_vector({self.label_id: 'True'})
            lambda_0 = lambda_agg_1 + (lambda_agg_0 - lambda_agg_1)*p0_product
            lambda_1 = lambda_agg_1
            lambda_msg = np.array([lambda_0, lambda_1])
            if not any(lambda_msg):
                raise InvalidLambdaMsgToParent
            return self._normalize(lambda_msg)


class BernoulliAndNode(Node):
    """
    A node in a DAG associated with a Bernoulli random variable with state_names ['False', 'True']
    and conditional probability distribution described by 'And' logic.
    """
    def __init__(self, label_id, children, parents):
        super().__init__(children=children, cpd=BernoulliAndCPD(label_id, parents))

    def compute_pi_agg(self):
        """
        Compute and update pi_agg, the prior probability, given the current state
        of messages received from parents.  Sidestep explicit factor product and
        marginalization.
        """
        if len(self.parents) == 0:
            self.update_pi_agg(self.cpd.values)
        else:
            pi_msgs = self.validate_and_return_msgs_received_for_msg_type(MessageType.PI)
            parents_p1 = [p.get_value_for_state_vector({p.variables[0]: 'True'})
                          for p in pi_msgs]
            p_1 = reduce(lambda x, y: x*y, parents_p1)
            p_0 = 1 - p_1
            self.update_pi_agg(np.array([p_0, p_1]))

    def compute_lambda_msg_to_parent(self, parent_k):
        """
        Compute lambda_msg to parent.  Sidestep explicit factor product and
        marginalization.

        Args
            parent_k: string or int,
                the label_id of the parent receiving the lambda_msg
        Returns
            np.array of ndim 1 and shape (cardinality of parent_k,)
        """
        if np.array_equal(self.lambda_agg.values, np.ones([self.cardinality])):
            return np.ones([self.cardinality])
        else:
            # TODO: cleanup this validation
            _ = self.validate_and_return_msgs_received_for_msg_type(MessageType.PI)
            p1_excluding_k = [p.get_value_for_state_vector({p.variables[0]: 'True'})
                              for par_id, p in self.pi_received_msgs.items() if par_id != parent_k]
            p1_product = reduce(lambda x, y: x*y, p1_excluding_k, 1)

            lambda_agg_0 = self.lambda_agg.get_value_for_state_vector({self.label_id: 'False'})
            lambda_agg_1 = self.lambda_agg.get_value_for_state_vector({self.label_id: 'True'})

            lambda_0 = lambda_agg_0
            lambda_1 = lambda_agg_0 + (lambda_agg_1 - lambda_agg_0)*p1_product
            lambda_msg = np.array([lambda_0, lambda_1])
            if not any(lambda_msg):
                raise InvalidLambdaMsgToParent
            return self._normalize(lambda_msg)