1 files changed, 15 insertions, 61 deletions

diff --git a/beliefs/models/BayesianModel.py b/beliefs/models/BayesianModel.py
index 6257a57..b57f968 100644
--- a/beliefs/models/BayesianModel.py
+++ b/beliefs/models/BayesianModel.py
@@ -1,9 +1,7 @@
 import copy
-import numpy as np
 import networkx as nx
 
 from beliefs.models.DirectedGraph import DirectedGraph
-from beliefs.utils.edges_helper import EdgesHelper
 from beliefs.utils.math_helper import is_kronecker_delta
 
 
@@ -12,74 +10,30 @@ class BayesianModel(DirectedGraph):
     Bayesian model stores nodes and edges described by conditional probability
     distributions.
     """
-    def __init__(self, edges, nodes_dict=None):
+    def __init__(self, edges=[], variables=[], cpds=[]):
         """
-        Input:
-          edges: list of edge tuples of form ('parent', 'child')
-          nodes: (optional) dict
-            a dict key, value pair as {label_id: instance_of_node_class_or_subclass}
-        """
-        if nodes_dict is not None:
-            super().__init__(edges, nodes_dict.keys())
-        else:
-            super().__init__(edges)
-        self.nodes_dict = nodes_dict
-
-    @classmethod
-    def from_node_class(cls, edges, node_class):
-        """Automatically create all nodes from the same node class
+        Base class for Bayesian model.
 
         Input:
-           edges: list of edge tuples of form ('parent', 'child')
-           node_class: (optional) the Node class or subclass from which to
-                 create all the nodes from edges.
-        """
-        nodes = cls.create_nodes(edges, node_class)
-        return cls.__init__(edges=edges, nodes=nodes)
-
-    @staticmethod
-    def create_nodes(edges, node_class):
-        """Returns list of Node instances created from edges using
-        the default node_class"""
-        edges_helper = EdgesHelper(edges)
-        nodes = edges_helper.create_nodes_from_edges(node_class=node_class)
-        label_to_node = dict()
-        for node in nodes:
-            label_to_node[node.label_id] = node
-        return label_to_node
-
-    def set_boundary_conditions(self):
-        """
-        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].pi_agg = self.nodes_dict[root].cpd.values
-
-        for leaf in self.get_leaves():
-            self.nodes_dict[leaf].lambda_agg = np.ones([self.nodes_dict[leaf].cardinality])
-
-    @property
-    def all_nodes_are_fully_initialized(self):
-        """
-        Returns True if, for all nodes in the model, all lambda and pi
-        messages and lambda_agg and pi_agg are not None, else False.
+          edges: (optional) list of edges,
+                tuples of form ('parent', 'child')
+          variables: (optional) list of str or int
+                labels for variables
+          cpds: (optional) list of CPDs
+                TabularCPD class or subclass
         """
-        for node in self.nodes_dict.values():
-            if not node.is_fully_initialized:
-                return False
-        return True
+        super().__init__()
+        super().add_edges_from(edges)
+        super().add_nodes_from(variables)
+        self.cpds = cpds
 
     def copy(self):
         """
         Returns a copy of the model.
         """
-        copy_edges = list(self.edges()).copy()
-        copy_nodes = copy.deepcopy(self.nodes_dict)
-        copy_model = self.__class__(edges=copy_edges, nodes=copy_nodes)
+        copy_model = self.__class__(edges=list(self.edges()).copy(),
+                                    variables=list(self.nodes()).copy(),
+                                    cpds=[cpd.copy() for cpd in self.cpds])
         return copy_model
 
     def get_variables_in_definite_state(self):