Apply BP fixes and update tests

Author: Morningliumengkun

src/bpdecoderplus/pytorch_bp/belief_propagation.py

@@ -2,7 +2,7 @@
 Belief Propagation (BP) algorithm implementation using PyTorch.
 """
 
-from typing import List, Dict, Tuple, Optional
+from typing import List, Dict, Tuple
 import torch
 from copy import deepcopy
 
@@ -88,7 +88,7 @@ def initial_state(bp: BeliefPropagation) -> BPState:
         var_messages_in = []
         var_messages_out = []
 
-        for factor_idx in bp.v2t[var_idx]:
+        for _ in bp.v2t[var_idx]:
             card = bp.cards[var_idx]
             msg = torch.ones(card, dtype=torch.float64)
             var_messages_in.append(msg.clone())
@@ -97,7 +97,7 @@ def initial_state(bp: BeliefPropagation) -> BPState:
         message_in.append(var_messages_in)
         message_out.append(var_messages_out)
 
-    return BPState(deepcopy(message_in), message_out)
+    return BPState(message_in, message_out)
 
 
 def _compute_factor_to_var_message(
@@ -124,7 +124,7 @@ def _compute_factor_to_var_message(
         return factor_tensor.clone()
 
     # Multiply factor tensor by incoming messages (excluding target) and sum out dims.
-    result = factor_tensor
+    result = factor_tensor.clone()
     for dim in range(ndims):
         if dim == target_var_idx:
             continue
@@ -154,11 +154,13 @@ def collect_message(bp: BeliefPropagation, state: BPState, normalize: bool = Tru
     for factor_idx, factor in enumerate(bp.factors):
         # Get incoming messages from variables to this factor
         incoming_messages = []
+        var_factor_positions = []
         for var in factor.vars:
             var_idx_0based = var - 1
             # Find position of this factor in v2t[var_idx_0based]
             factor_pos = bp.v2t[var_idx_0based].index(factor_idx)
             incoming_messages.append(state.message_out[var_idx_0based][factor_pos])
+            var_factor_positions.append(factor_pos)
 
         # Compute outgoing message to each variable
         for var_pos, var in enumerate(factor.vars):
@@ -177,7 +179,7 @@ def collect_message(bp: BeliefPropagation, state: BPState, normalize: bool = Tru
                     outgoing_msg = outgoing_msg / msg_sum
 
             # Update message_in
-            factor_pos = bp.v2t[var_idx_0based].index(factor_idx)
+            factor_pos = var_factor_positions[var_pos]
             state.message_in[var_idx_0based][factor_pos] = outgoing_msg
 
 
@@ -334,19 +336,17 @@ def apply_evidence(bp: BeliefPropagation, evidence: Dict[int, int]) -> BeliefPro
         for var_pos, var in enumerate(factor.vars):
             if var in evidence:
                 evid_value = evidence[var]
-                # Create slice that zeros out non-evidence values
-                slices = [slice(None)] * len(factor.vars)
-                slices[var_pos] = evid_value
-                
-                # Zero out all non-evidence assignments
-                mask = torch.ones_like(factor_tensor)
-                for i in range(factor_tensor.shape[var_pos]):
-                    if i != evid_value:
-                        slices_mask = slices.copy()
-                        slices_mask[var_pos] = i
-                        mask[tuple(slices_mask)] = 0
-                
-                factor_tensor = factor_tensor * mask
+                dim_size = factor_tensor.shape[var_pos]
+                if 0 <= evid_value < dim_size:
+                    all_indices = torch.arange(dim_size, device=factor_tensor.device)
+                    zero_indices = all_indices[all_indices != evid_value]
+                    if zero_indices.numel() > 0:
+                        factor_tensor = factor_tensor.index_fill(
+                            var_pos, zero_indices, 0
+                        )
+                else:
+                    factor_tensor = torch.zeros_like(factor_tensor)
+                    break
 
         new_factors.append(Factor(factor.vars, factor_tensor))
 

src/bpdecoderplus/pytorch_bp/uai_parser.py

@@ -2,7 +2,7 @@
 UAI file format parser for Belief Propagation.
 """
 
-from typing import List, Dict, Tuple
+from typing import List, Dict
 import torch
 
 
@@ -71,6 +71,10 @@ def read_model_from_string(content: str, factor_eltype=torch.float64) -> UAIMode
 
     # Parse header
     network_type = lines[0]  # MARKOV or BAYES
+    if network_type not in ("MARKOV", "BAYES"):
+        raise ValueError(
+            f"Unsupported UAI network type: {network_type!r}. Expected 'MARKOV' or 'BAYES'."
+        )
     nvars = int(lines[1])
     cards = [int(x) for x in lines[2].split()]
     ntables = int(lines[3])
@@ -80,6 +84,11 @@ def read_model_from_string(content: str, factor_eltype=torch.float64) -> UAIMode
     for i in range(ntables):
         parts = lines[4 + i].split()
         scope_size = int(parts[0])
+        if len(parts) - 1 != scope_size:
+            raise ValueError(
+                f"Scope size mismatch on line {4 + i}: "
+                f"declared {scope_size}, found {len(parts) - 1} variables."
+            )
         scope = [int(x) + 1 for x in parts[1:]]  # Convert to 1-based
         scopes.append(scope)
 

src/bpdecoderplus/pytorch_bp/utils.py

@@ -4,5 +4,7 @@
 
 def add_project_root_to_path():
     project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))
-    if project_root not in sys.path:
-        sys.path.insert(0, project_root)
+    src_root = os.path.join(project_root, "src")
+    for path in (src_root, project_root):
+        if path not in sys.path:
+            sys.path.insert(0, path)

tests/_path.py

@@ -1,7 +1,13 @@
 import unittest
-import itertools
 import torch
 
+try:
+    from ._path import add_project_root_to_path
+except ImportError:
+    from _path import add_project_root_to_path
+
+add_project_root_to_path()
+
 from bpdecoderplus.pytorch_bp import (
     read_model_from_string,
     BeliefPropagation,
@@ -10,32 +16,7 @@
     apply_evidence,
 )
 
-
-def exact_marginals(model, evidence=None):
-    evidence = evidence or {}
-    assignments = list(itertools.product(*[range(c) for c in model.cards]))
-    weights = []
-    for assignment in assignments:
-        if any(assignment[var_idx - 1] != val for var_idx, val in evidence.items()):
-            weights.append(0.0)
-            continue
-        weight = 1.0
-        for factor in model.factors:
-            idx = tuple(assignment[v - 1] for v in factor.vars)
-            weight *= float(factor.values[idx])
-        weights.append(weight)
-    total = sum(weights)
-    marginals = {}
-    for var_idx, card in enumerate(model.cards):
-        values = []
-        for value in range(card):
-            mass = 0.0
-            for assignment, weight in zip(assignments, weights):
-                if assignment[var_idx] == value:
-                    mass += weight
-            values.append(mass / total if total > 0 else 0.0)
-        marginals[var_idx + 1] = torch.tensor(values, dtype=torch.float64)
-    return marginals
+from tests.test_utils import exact_marginals
 
 
 class TestBeliefPropagationBasic(unittest.TestCase):

tests/test_bp_basic.py

@@ -1,5 +1,12 @@
 import unittest
 
+try:
+    from ._path import add_project_root_to_path
+except ImportError:
+    from _path import add_project_root_to_path
+
+add_project_root_to_path()
+
 from bpdecoderplus.pytorch_bp import (
     read_model_file,
     read_evidence_file,
@@ -15,7 +22,7 @@ def test_example_file_runs(self):
         model = read_model_file("examples/simple_model.uai")
         bp = BeliefPropagation(model)
         state, info = belief_propagate(bp, max_iter=30, tol=1e-8)
-        self.assertTrue(info.iterations > 0)
+        self.assertGreater(info.iterations, 0)
         marginals = compute_marginals(state, bp)
         self.assertEqual(set(marginals.keys()), {1, 2})
 
@@ -24,7 +31,7 @@ def test_example_with_evidence(self):
         evidence = read_evidence_file("examples/simple_model.evid")
         bp = apply_evidence(BeliefPropagation(model), evidence)
         state, info = belief_propagate(bp, max_iter=30, tol=1e-8)
-        self.assertTrue(info.iterations > 0)
+        self.assertGreater(info.iterations, 0)
         marginals = compute_marginals(state, bp)
         self.assertEqual(set(marginals.keys()), {1, 2})
 

tests/test_integration.py

@@ -1,6 +1,13 @@
 import unittest
 import torch
 
+try:
+    from ._path import add_project_root_to_path
+except ImportError:
+    from _path import add_project_root_to_path
+
+add_project_root_to_path()
+
 from bpdecoderplus.pytorch_bp import read_model_from_string, read_evidence_file
 
 

tests/test_uai_parser.py

@@ -1,7 +1,13 @@
 import unittest
-import itertools
 import torch
 
+try:
+    from ._path import add_project_root_to_path
+except ImportError:
+    from _path import add_project_root_to_path
+
+add_project_root_to_path()
+
 from bpdecoderplus.pytorch_bp import (
     read_model_from_string,
     BeliefPropagation,
@@ -13,31 +19,7 @@
 )
 
 
-def exact_marginals(model, evidence=None):
-    evidence = evidence or {}
-    assignments = list(itertools.product(*[range(c) for c in model.cards]))
-    weights = []
-    for assignment in assignments:
-        if any(assignment[var_idx - 1] != val for var_idx, val in evidence.items()):
-            weights.append(0.0)
-            continue
-        weight = 1.0
-        for factor in model.factors:
-            idx = tuple(assignment[v - 1] for v in factor.vars)
-            weight *= float(factor.values[idx])
-        weights.append(weight)
-    total = sum(weights)
-    marginals = {}
-    for var_idx, card in enumerate(model.cards):
-        values = []
-        for value in range(card):
-            mass = 0.0
-            for assignment, weight in zip(assignments, weights):
-                if assignment[var_idx] == value:
-                    mass += weight
-            values.append(mass / total if total > 0 else 0.0)
-        marginals[var_idx + 1] = torch.tensor(values, dtype=torch.float64)
-    return marginals
+from tests.test_utils import exact_marginals
 
 
 class TestBPAdditionalCases(unittest.TestCase):