The generic framework for solving approximate inference problem with BP

README.md

@@ -186,6 +186,42 @@ BPDecoderPlus/
     └── belief_propagation_qec_plan.tex
 ```
 
+## PyTorch BP Module (UAI)
+
+This repository also includes a PyTorch implementation of belief propagation for
+UAI factor graphs under `src/bpdecoderplus/pytorch_bp`.
+
+### Python Setup
+
+```bash
+pip install -e .
+```
+
+### Quick Example
+
+```python
+from bpdecoderplus.pytorch_bp import (
+    read_model_file,
+    BeliefPropagation,
+    belief_propagate,
+    compute_marginals,
+)
+
+model = read_model_file("examples/simple_model.uai")
+bp = BeliefPropagation(model)
+state, info = belief_propagate(bp)
+print(info)
+print(compute_marginals(state, bp))
+```
+
+### Examples and Tests
+
+```bash
+python examples/simple_example.py
+python examples/evidence_example.py
+pytest tests/test_bp_basic.py tests/test_uai_parser.py tests/test_integration.py tests/testcase.py
+```
+
 ## Available Decoders
 
 | Decoder | Symbol | Description |

docs/api_reference.md

@@ -0,0 +1,45 @@
+## PyTorch BP API Reference
+
+This reference documents the public API exported from `bpdecoderplus.pytorch_bp`.
+
+### UAI Parsing
+
+- `read_model_file(path, factor_eltype=torch.float64) -> UAIModel`
+  Parse a UAI `.uai` model file.
+
+- `read_model_from_string(content, factor_eltype=torch.float64) -> UAIModel`
+  Parse a UAI model from an in-memory string.
+
+- `read_evidence_file(path) -> Dict[int, int]`
+  Parse a UAI `.evid` file and return evidence as 1-based indices.
+
+### Data Structures
+
+- `Factor(vars: List[int], values: torch.Tensor)`
+  Container for a factor scope and its tensor.
+
+- `UAIModel(nvars: int, cards: List[int], factors: List[Factor])`
+  Holds all model metadata for BP.
+
+### Belief Propagation
+
+- `BeliefPropagation(uai_model: UAIModel)`
+  Builds factor graph adjacency for BP.
+
+- `initial_state(bp: BeliefPropagation) -> BPState`
+  Initialize messages to uniform vectors.
+
+- `collect_message(bp, state, normalize=True)`
+  Update factor-to-variable messages in place.
+
+- `process_message(bp, state, normalize=True, damping=0.2)`
+  Update variable-to-factor messages in place.
+
+- `belief_propagate(bp, max_iter=100, tol=1e-6, damping=0.2, normalize=True)`
+  Run the full BP loop and return `(BPState, BPInfo)`.
+
+- `compute_marginals(state, bp) -> Dict[int, torch.Tensor]`
+  Compute marginal distributions after convergence.
+
+- `apply_evidence(bp, evidence: Dict[int, int]) -> BeliefPropagation`
+  Return a new BP object with evidence applied to factor tensors.

docs/mathematical_description.md

@@ -0,0 +1,41 @@
+## Belief Propagation (BP) Overview
+
+This document summarizes the BP message-passing rules implemented in
+`src/bpdecoderplus/pytorch_bp/belief_propagation.py` for discrete factor graphs. The approach
+mirrors the tensor-contraction perspective used in TensorInference.jl.
+See https://github.com/TensorBFS/TensorInference.jl for the Julia reference.
+
+### Factor Graph Notation
+
+- Variables are indexed by x_i with domain size d_i.
+- Factors are indexed by f and connect a subset of variables.
+- Each factor has a tensor (potential) phi_f defined over its variables.
+
+### Messages
+
+Factor to variable message:
+
+mu_{f->x}(x) = sum_{all y in ne(f), y != x} phi_f(x, y, ...) * product_{y != x} mu_{y->f}(y)
+
+Variable to factor message:
+
+mu_{x->f}(x) = product_{g in ne(x), g != f} mu_{g->x}(x)
+
+### Damping
+
+To improve stability on loopy graphs, a damping update is applied:
+
+mu_new = damping * mu_old + (1 - damping) * mu_candidate
+
+### Convergence
+
+We use an L1 difference threshold between consecutive factor->variable
+messages to determine convergence.
+
+### Marginals
+
+After convergence, variable marginals are computed as:
+
+b(x) = (1 / Z) * product_{f in ne(x)} mu_{f->x}(x)
+
+The normalization constant Z is obtained by summing the unnormalized vector.

docs/usage_guide.md

@@ -0,0 +1,43 @@
+## PyTorch Belief Propagation Usage
+
+This guide shows how to parse a UAI file, run BP, and apply evidence.
+The implementation follows the tensor-contraction viewpoint in
+TensorInference.jl: https://github.com/TensorBFS/TensorInference.jl
+
+### Quick Start
+
+```python
+from bpdecoderplus.pytorch_bp import (
+    read_model_file,
+    BeliefPropagation,
+    belief_propagate,
+    compute_marginals,
+)
+
+model = read_model_file("examples/simple_model.uai")
+bp = BeliefPropagation(model)
+state, info = belief_propagate(bp, max_iter=50, tol=1e-8, damping=0.1)
+print(info)
+
+marginals = compute_marginals(state, bp)
+print(marginals[1])
+```
+
+### Evidence
+
+```python
+from bpdecoderplus.pytorch_bp import read_model_file, read_evidence_file, apply_evidence
+from bpdecoderplus.pytorch_bp import BeliefPropagation, belief_propagate, compute_marginals
+
+model = read_model_file("examples/simple_model.uai")
+evidence = read_evidence_file("examples/simple_model.evid")
+bp = apply_evidence(BeliefPropagation(model), evidence)
+state, info = belief_propagate(bp)
+marginals = compute_marginals(state, bp)
+```
+
+### Tips
+
+- For loopy graphs, use damping between 0.1 and 0.5.
+- Normalize messages to avoid numerical underflow.
+- Use float64 for consistent comparisons in tests.

examples/evidence_example.py

@@ -0,0 +1,24 @@
+from bpdecoderplus.pytorch_bp import (
+    read_model_file,
+    read_evidence_file,
+    BeliefPropagation,
+    belief_propagate,
+    compute_marginals,
+    apply_evidence,
+)
+
+
+def main():
+    model = read_model_file("examples/simple_model.uai")
+    evidence = read_evidence_file("examples/simple_model.evid")
+    bp = apply_evidence(BeliefPropagation(model), evidence)
+    state, info = belief_propagate(bp, max_iter=50, tol=1e-8, damping=0.1)
+    print(info)
+
+    marginals = compute_marginals(state, bp)
+    for var_idx, marginal in marginals.items():
+        print(f"Variable {var_idx} marginal: {marginal}")
+
+
+if __name__ == "__main__":
+    main()

examples/simple_example.py

@@ -0,0 +1,21 @@
+from bpdecoderplus.pytorch_bp import (
+    read_model_file,
+    BeliefPropagation,
+    belief_propagate,
+    compute_marginals,
+)
+
+
+def main():
+    model = read_model_file("examples/simple_model.uai")
+    bp = BeliefPropagation(model)
+    state, info = belief_propagate(bp, max_iter=50, tol=1e-8, damping=0.1)
+    print(info)
+
+    marginals = compute_marginals(state, bp)
+    for var_idx, marginal in marginals.items():
+        print(f"Variable {var_idx} marginal: {marginal}")
+
+
+if __name__ == "__main__":
+    main()

examples/simple_model.evid

@@ -0,0 +1 @@
+1 0 1

examples/simple_model.uai

@@ -0,0 +1,10 @@
+MARKOV
+2
+2 2
+2
+1 0
+2 0 1
+2
+0.6 0.4
+4
+0.9 0.1 0.2 0.8

pyproject.toml

@@ -23,6 +23,7 @@ classifiers = [
 dependencies = [
     "stim>=1.12.0",
     "numpy>=1.24.0",
+    "torch>=2.0.0",
 ]
 
 [project.optional-dependencies]

requirements.txt

@@ -0,0 +1 @@
+torch

src/bpdecoderplus/pytorch_bp/__init__.py

@@ -0,0 +1,42 @@
+"""
+PyTorch Belief Propagation (BP) submodule for approximate inference.
+"""
+
+from .uai_parser import (
+    read_model_file,
+    read_model_from_string,
+    read_evidence_file,
+    UAIModel,
+    Factor
+)
+
+from .belief_propagation import (
+    BeliefPropagation,
+    BPState,
+    BPInfo,
+    initial_state,
+    collect_message,
+    process_message,
+    belief_propagate,
+    compute_marginals,
+    apply_evidence
+)
+
+__all__ = [
+    # UAI parsing
+    'read_model_file',
+    'read_model_from_string',
+    'read_evidence_file',
+    'UAIModel',
+    'Factor',
+    # Belief Propagation
+    'BeliefPropagation',
+    'BPState',
+    'BPInfo',
+    'initial_state',
+    'collect_message',
+    'process_message',
+    'belief_propagate',
+    'compute_marginals',
+    'apply_evidence',
+]