Use omeco ordering in tropical MPE.

Update docs/imports and improve validation coverage.

Author: Morningliumengkun

tropical_in_new/README.md

@@ -1,8 +1,12 @@
 # Tropical Tensor Network for MPE
 
 This folder contains an independent implementation of tropical tensor network
-contraction for Most Probable Explanation (MPE). It does not depend on the
-`bpdecoderplus` package; all code lives under `tropical_in_new/src`.
+contraction for Most Probable Explanation (MPE). It uses `omeco` for contraction
+order optimization and does not depend on the `bpdecoderplus` package; all code
+lives under `tropical_in_new/src`.
+
+`omeco` provides high-quality contraction order heuristics (greedy and
+simulated annealing). Install it alongside Torch to run the examples and tests.
 
 ## Structure
 
@@ -20,7 +24,8 @@ tropical_in_new/
 ├── tests/
 │   ├── test_primitives.py
 │   ├── test_contraction.py
-│   └── test_mpe.py
+│   ├── test_mpe.py
+│   └── test_utils.py
 ├── examples/
 │   └── asia_network/
 │       ├── main.py
@@ -34,5 +39,13 @@ tropical_in_new/
 ## Quick Start
 
 ```bash
+pip install -r tropical_in_new/requirements.txt
 python tropical_in_new/examples/asia_network/main.py
 ```
+
+## Notes on omeco
+
+`omeco` is a Rust-backed Python package. If a prebuilt wheel is not available
+for your Python version, you will need a Rust toolchain with `cargo` on PATH to
+build it from source. See the omeco repository for details:
+https://github.com/GiggleLiu/omeco

tropical_in_new/docs/api_reference.md

@@ -34,8 +34,8 @@ Public APIs exported from `tropical_in_new/src`.
 - `build_network(factors: Iterable[Factor]) -> list[TensorNode]`
   Convert factors into log-domain tensors with scopes.
 
-- `choose_order(nodes, heuristic="min_fill") -> list[int]`
-  Select variable elimination order (min-fill / min-degree).
+- `choose_order(nodes, heuristic="omeco") -> list[int]`
+  Select variable elimination order using `omeco`.
 
 - `build_contraction_tree(order, nodes) -> ContractionTree`
   Construct a contraction plan from order and nodes.

tropical_in_new/docs/usage_guide.md

@@ -3,10 +3,19 @@
 This guide shows how to parse a UAI model and compute MPE using the
 tropical tensor network implementation.
 
+### Install Dependencies
+
+```bash
+pip install -r tropical_in_new/requirements.txt
+```
+
+`omeco` provides contraction order optimization. If a prebuilt wheel is not
+available for your Python version, you may need a Rust toolchain installed.
+
 ### Quick Start
 
 ```python
-from src import mpe_tropical, read_model_file
+from tropical_in_new.src import mpe_tropical, read_model_file
 
 model = read_model_file("tropical_in_new/examples/asia_network/model.uai")
 assignment, score, info = mpe_tropical(model)
@@ -16,7 +25,7 @@ print(assignment, score, info)
 ### Evidence
 
 ```python
-from src import mpe_tropical, read_model_file
+from tropical_in_new.src import mpe_tropical, read_model_file
 
 model = read_model_file("tropical_in_new/examples/asia_network/model.uai")
 evidence = {1: 0}  # variable index is 1-based

tropical_in_new/examples/asia_network/main.py

@@ -1,12 +1,6 @@
 """Run tropical MPE on a small UAI model."""
 
-from pathlib import Path
-import sys
-
-ROOT = Path(__file__).resolve().parents[2]
-sys.path.insert(0, str(ROOT))
-
-from src import mpe_tropical, read_model_file  # noqa: E402
+from tropical_in_new.src import mpe_tropical, read_model_file
 
 
 def main() -> None:

tropical_in_new/requirements.txt

@@ -1,2 +1,2 @@
 torch>=2.0.0
-numpy>=1.24.0
+omeco

tropical_in_new/src/__init__.py

@@ -4,7 +4,14 @@
 from .mpe import mpe_tropical, recover_mpe_assignment
 from .network import TensorNode, build_network
 from .primitives import argmax_trace, safe_log, tropical_einsum
-from .utils import Factor, UAIModel, build_tropical_factors, read_model_file, read_model_from_string
+from .utils import (
+    Factor,
+    UAIModel,
+    build_tropical_factors,
+    read_evidence_file,
+    read_model_file,
+    read_model_from_string,
+)
 
 __all__ = [
     "Factor",
@@ -17,6 +24,7 @@
     "choose_order",
     "contract_tree",
     "mpe_tropical",
+    "read_evidence_file",
     "read_model_file",
     "read_model_from_string",
     "recover_mpe_assignment",

tropical_in_new/src/contraction.py

@@ -7,10 +7,7 @@
 
 import torch
 
-try:  # Optional heuristic provider
-    import omeco  # type: ignore
-except Exception:  # pragma: no cover - optional dependency
-    omeco = None
+import omeco
 
 from .network import TensorNode
 from .primitives import Backpointer, tropical_reduce_max
@@ -45,59 +42,86 @@ class ContractionTree:
     nodes: Tuple[TensorNode, ...]
 
 
-def _build_var_graph(nodes: Iterable[TensorNode]) -> dict[int, set[int]]:
-    graph: dict[int, set[int]] = {}
+def _infer_var_sizes(nodes: Iterable[TensorNode]) -> dict[int, int]:
+    sizes: dict[int, int] = {}
     for node in nodes:
-        vars = list(node.vars)
+        for var, dim in zip(node.vars, node.values.shape):
+            if var in sizes and sizes[var] != dim:
+                raise ValueError(
+                    f"Variable {var} has inconsistent sizes: {sizes[var]} vs {dim}."
+                )
+            sizes[var] = int(dim)
+    return sizes
+
+
+def _extract_leaf_index(node_dict: dict) -> int | None:
+    for key in ("leaf", "leaf_index", "index", "tensor"):
+        if key in node_dict:
+            value = node_dict[key]
+            if isinstance(value, int):
+                return value
+    return None
+
+
+def _elim_order_from_tree_dict(tree_dict: dict, ixs: list[list[int]]) -> list[int]:
+    total_counts: dict[int, int] = {}
+    for vars in ixs:
         for var in vars:
-            graph.setdefault(var, set()).update(v for v in vars if v != var)
-    return graph
-
-
-def _min_fill_order(graph: dict[int, set[int]]) -> list[int]:
-    order: list[int] = []
-    graph = {k: set(v) for k, v in graph.items()}
-    while graph:
-        best_var = None
-        best_fill = None
-        best_degree = None
-        for var, neighbors in graph.items():
-            fill = 0
-            neighbor_list = list(neighbors)
-            for i in range(len(neighbor_list)):
-                for j in range(i + 1, len(neighbor_list)):
-                    if neighbor_list[j] not in graph[neighbor_list[i]]:
-                        fill += 1
-            degree = len(neighbors)
-            if best_fill is None or (fill, degree) < (best_fill, best_degree):
-                best_var = var
-                best_fill = fill
-                best_degree = degree
-        if best_var is None:
-            break
-        neighbors = list(graph[best_var])
-        for i in range(len(neighbors)):
-            for j in range(i + 1, len(neighbors)):
-                graph[neighbors[i]].add(neighbors[j])
-                graph[neighbors[j]].add(neighbors[i])
-        for neighbor in neighbors:
-            graph[neighbor].discard(best_var)
-        graph.pop(best_var, None)
-        order.append(best_var)
-    return order
-
-
-def choose_order(nodes: list[TensorNode], heuristic: str = "min_fill") -> list[int]:
-    """Select elimination order over variable indices."""
-    if heuristic == "omeco" and omeco is not None:
-        if hasattr(omeco, "min_fill_order"):
-            return list(omeco.min_fill_order([node.vars for node in nodes]))
-    graph = _build_var_graph(nodes)
-    if heuristic in ("min_fill", "omeco"):
-        return _min_fill_order(graph)
-    if heuristic == "min_degree":
-        return sorted(graph, key=lambda v: len(graph[v]))
-    raise ValueError(f"Unknown heuristic: {heuristic!r}")
+            total_counts[var] = total_counts.get(var, 0) + 1
+
+    eliminated: set[int] = set()
+
+    def visit(node: dict) -> tuple[dict[int, int], list[int]]:
+        leaf_index = _extract_leaf_index(node)
+        if leaf_index is not None:
+            counts: dict[int, int] = {}
+            for var in ixs[leaf_index]:
+                counts[var] = counts.get(var, 0) + 1
+            return counts, []
+
+        children = node.get("children", [])
+        if not isinstance(children, list) or not children:
+            return {}, []
+
+        counts: dict[int, int] = {}
+        order: list[int] = []
+        for child in children:
+            child_counts, child_order = visit(child)
+            order.extend(child_order)
+            for var, count in child_counts.items():
+                counts[var] = counts.get(var, 0) + count
+
+        newly_eliminated = [
+            var
+            for var, count in counts.items()
+            if count == total_counts.get(var, 0) and var not in eliminated
+        ]
+        for var in sorted(newly_eliminated):
+            eliminated.add(var)
+            order.append(var)
+        return counts, order
+
+    _, order = visit(tree_dict)
+    remaining = sorted([var for var in total_counts if var not in eliminated])
+    return order + remaining
+
+
+def choose_order(nodes: list[TensorNode], heuristic: str = "omeco") -> list[int]:
+    """Select elimination order over variable indices using omeco."""
+    if heuristic != "omeco":
+        raise ValueError("Only the 'omeco' heuristic is supported.")
+    ixs = [list(node.vars) for node in nodes]
+    sizes = _infer_var_sizes(nodes)
+    method = omeco.GreedyMethod() if hasattr(omeco, "GreedyMethod") else None
+    tree = (
+        omeco.optimize_code(ixs, [], sizes, method)
+        if method is not None
+        else omeco.optimize_code(ixs, [], sizes)
+    )
+    tree_dict = tree.to_dict() if hasattr(tree, "to_dict") else tree
+    if not isinstance(tree_dict, dict):
+        raise ValueError("omeco.optimize_code did not return a usable tree.")
+    return _elim_order_from_tree_dict(tree_dict, ixs)
 
 
 def build_contraction_tree(order: Iterable[int], nodes: list[TensorNode]) -> ContractionTree:

tropical_in_new/src/mpe.py

@@ -29,6 +29,14 @@ def recover_mpe_assignment(root) -> Dict[int, int]:
     """Recover MPE assignment from a contraction tree with backpointers."""
     assignment: Dict[int, int] = {}
 
+    def require_vars(required: Iterable[int], available: Dict[int, int]) -> None:
+        missing = [v for v in required if v not in available]
+        if missing:
+            raise KeyError(
+                "Missing assignment values for variables: "
+                f"{missing}. Provided assignment keys: {sorted(available.keys())}"
+            )
+
     def traverse(node, out_assignment: Dict[int, int]) -> None:
         assignment.update(out_assignment)
         if isinstance(node, TensorNode):
@@ -38,6 +46,7 @@ def traverse(node, out_assignment: Dict[int, int]) -> None:
                 argmax_trace(node.backpointer, out_assignment) if node.backpointer else {}
             )
             combined = {**out_assignment, **elim_assignment}
+            require_vars(node.child.vars, combined)
             child_assignment = {v: combined[v] for v in node.child.vars}
             traverse(node.child, child_assignment)
             return
@@ -46,7 +55,9 @@ def traverse(node, out_assignment: Dict[int, int]) -> None:
                 argmax_trace(node.backpointer, out_assignment) if node.backpointer else {}
             )
             combined = {**out_assignment, **elim_assignment}
+            require_vars(node.left.vars, combined)
             left_assignment = {v: combined[v] for v in node.left.vars}
+            require_vars(node.right.vars, combined)
             right_assignment = {v: combined[v] for v in node.right.vars}
             traverse(node.left, left_assignment)
             traverse(node.right, right_assignment)
@@ -66,7 +77,7 @@ def mpe_tropical(
     factors = build_tropical_factors(model, evidence)
     nodes = build_network(factors)
     if order is None:
-        order = choose_order(nodes, heuristic="min_fill")
+        order = choose_order(nodes, heuristic="omeco")
     tree = build_contraction_tree(order, nodes)
     root = _contract_tree(tree, einsum_fn=tropical_einsum)
     if root.vars:

tropical_in_new/src/primitives.py

@@ -33,19 +33,6 @@ class Backpointer:
     argmax_flat: torch.Tensor
 
 
-@dataclass(frozen=True)
-class TropicalTensor:
-    """Lightweight wrapper for tropical (max-plus) tensors."""
-
-    vars: Tuple[int, ...]
-    values: torch.Tensor
-
-    def __add__(self, other: "TropicalTensor") -> "TropicalTensor":
-        if self.vars != other.vars:
-            raise ValueError("TropicalTensor.__add__ requires identical variable order.")
-        return TropicalTensor(self.vars, torch.maximum(self.values, other.values))
-
-
 def safe_log(tensor: torch.Tensor) -> torch.Tensor:
     """Convert potentials to log domain; zeros map to -inf."""
     neg_inf = torch.tensor(float("-inf"), dtype=tensor.dtype, device=tensor.device)
@@ -83,6 +70,12 @@ def tropical_reduce_max(
     if not elim_vars:
         return tensor, None
     target_vars = tuple(vars)
+    missing_elim_vars = [v for v in elim_vars if v not in target_vars]
+    if missing_elim_vars:
+        raise ValueError(
+            "tropical_reduce_max: elim_vars "
+            f"{missing_elim_vars} are not present in vars {target_vars}."
+        )
     elim_axes = [target_vars.index(v) for v in elim_vars]
     keep_axes = [i for i in range(len(target_vars)) if i not in elim_axes]
     perm = keep_axes + elim_axes
@@ -138,6 +131,12 @@ def argmax_trace(backpointer: Backpointer, assignment: Dict[int, int]) -> Dict[i
     if not backpointer.elim_vars:
         return {}
     if backpointer.out_vars:
+        missing = [v for v in backpointer.out_vars if v not in assignment]
+        if missing:
+            raise KeyError(
+                "Missing assignment values for output variables: "
+                f"{missing}. Provided assignment keys: {sorted(assignment.keys())}"
+            )
         idx = tuple(assignment[v] for v in backpointer.out_vars)
         flat = int(backpointer.argmax_flat[idx].item())
     else: