diff options
Diffstat (limited to 'tests/test_belief_propagation.py')
| -rw-r--r-- | tests/test_belief_propagation.py | 13 |
1 files changed, 6 insertions, 7 deletions
diff --git a/tests/test_belief_propagation.py b/tests/test_belief_propagation.py index 24ee94b..264ddae 100644 --- a/tests/test_belief_propagation.py +++ b/tests/test_belief_propagation.py @@ -3,8 +3,8 @@ import pytest from pytest import approx from beliefs.inference.belief_propagation import BeliefPropagation, ConflictingEvidenceError -from beliefs.models.BernoulliOrModel import BernoulliOrModel -from beliefs.types.BernoulliOrNode import BernoulliOrNode +from beliefs.models.beliefupdate.BeliefUpdateNodeModel import BeliefUpdateNodeModel +from beliefs.models.beliefupdate.BernoulliOrNode import BernoulliOrNode @pytest.fixture(scope='module') @@ -37,24 +37,23 @@ def many_parents_edges(): @pytest.fixture(scope='function') def four_node_model(edges_four_nodes): - return BernoulliOrModel(edges_four_nodes) + return BeliefUpdateNodeModel.from_edges(edges_four_nodes, BernoulliOrNode) @pytest.fixture(scope='function') def simple_model(simple_edges): - return BernoulliOrModel(simple_edges) + return BeliefUpdateNodeModel.from_edges(simple_edges, BernoulliOrNode) @pytest.fixture(scope='function') def many_parents_model(many_parents_edges): - return BernoulliOrModel(many_parents_edges) + return BeliefUpdateNodeModel.from_edges(many_parents_edges, BernoulliOrNode) @pytest.fixture(scope='function') def one_node_model(): a_node = BernoulliOrNode(label_id='x', children=[], parents=[]) - # a_node = BernoulliOrNode(label_id='x', children=set(), parents=set()) - return BernoulliOrModel(edges=None, nodes={'x': a_node}) + return BeliefUpdateNodeModel(nodes_dict={'x': a_node}) def get_label_mapped_to_positive_belief(query_result): |