Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 51 additions & 0 deletions tropical_in_new/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
# Tropical Tensor Network for MPE

This folder contains an independent implementation of tropical tensor network
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

```
tropical_in_new/
├── README.md
├── requirements.txt
├── src/
│ ├── __init__.py
│ ├── primitives.py
│ ├── network.py
│ ├── contraction.py
│ ├── mpe.py
│ └── utils.py
├── tests/
│ ├── test_primitives.py
│ ├── test_contraction.py
│ ├── test_mpe.py
│ └── test_utils.py
├── examples/
│ └── asia_network/
│ ├── main.py
│ └── model.uai
└── docs/
├── mathematical_description.md
├── api_reference.md
└── usage_guide.md
```

## 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
52 changes: 52 additions & 0 deletions tropical_in_new/docs/api_reference.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
## Tropical Tensor Network API Reference

Public APIs exported from `tropical_in_new/src`.

### 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.

### Data Structures

- `Factor(vars: Tuple[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 MPE.

### Primitives

- `safe_log(tensor: torch.Tensor) -> torch.Tensor`
Convert potentials to log-domain; maps zeros to `-inf`.

- `tropical_einsum(a, b, index_map, track_argmax=True)`
Binary contraction in max-plus semiring; returns `(values, backpointer)`.

- `argmax_trace(backpointer, assignment) -> Dict[int, int]`
Decode assignments for eliminated variables from backpointer metadata.

### Network + Contraction

- `build_network(factors: Iterable[Factor]) -> list[TensorNode]`
Convert factors into log-domain tensors with scopes.

- `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.

- `contract_tree(tree, einsum_fn, track_argmax=True) -> TreeNode`
Execute contractions and return the root node with backpointers.

### MPE

- `mpe_tropical(model, evidence=None, order=None)`
Return `(assignment_dict, score, info)` where `score` is log-domain.

- `recover_mpe_assignment(root) -> Dict[int, int]`
Recover assignments from a contraction tree root.
32 changes: 32 additions & 0 deletions tropical_in_new/docs/mathematical_description.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
## Tropical Tensor Network for MPE

We compute the Most Probable Explanation (MPE) by treating the factor
graph as a tensor network in the tropical semiring (max-plus).

### Tropical Semiring

For log-potentials, multiplication becomes addition and summation becomes
maximization:

- `a ⊗ b = a + b`
- `a ⊕ b = max(a, b)`

### MPE Objective

Let `x` be the set of variables and `phi_f(x_f)` the factor potentials.
We maximize the log-score:

`x* = argmax_x sum_f log(phi_f(x_f))`

### Contraction

Each factor becomes a tensor over its variable scope. Contraction combines
two tensors by adding their log-values and reducing (max) over eliminated
variables. A greedy elimination order (min-fill or min-degree) controls
the intermediate tensor sizes.

### Backpointers

During each reduction, we store the argmax index for eliminated variables.
Traversing the contraction tree with these backpointers recovers the MPE
assignment.
39 changes: 39 additions & 0 deletions tropical_in_new/docs/usage_guide.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
## Tropical Tensor Network Usage

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 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)
print(assignment, score, info)
```

### Evidence

```python
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
assignment, score, info = mpe_tropical(model, evidence=evidence)
```

### Running the Example

```bash
python tropical_in_new/examples/asia_network/main.py
```
15 changes: 15 additions & 0 deletions tropical_in_new/examples/asia_network/main.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
"""Run tropical MPE on a small UAI model."""

from tropical_in_new.src import mpe_tropical, read_model_file


def main() -> None:
model = read_model_file("tropical_in_new/examples/asia_network/model.uai")
assignment, score, info = mpe_tropical(model)
print("MPE assignment:", assignment)
print("MPE log-score:", score)
print("Info:", info)


if __name__ == "__main__":
main()
19 changes: 19 additions & 0 deletions tropical_in_new/examples/asia_network/model.uai
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
MARKOV
3
2 2 2
5
1 0
1 1
1 2
2 0 1
2 1 2
2
0.6 0.4
2
0.5 0.5
2
0.7 0.3
4
1.2 0.2 0.2 1.2
4
1.1 0.3 0.3 1.1
2 changes: 2 additions & 0 deletions tropical_in_new/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
torch>=2.0.0
omeco
33 changes: 33 additions & 0 deletions tropical_in_new/src/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
"""Tropical tensor network tools for MPE (independent package)."""

from .contraction import build_contraction_tree, choose_order, contract_tree
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_evidence_file,
read_model_file,
read_model_from_string,
)

__all__ = [
"Factor",
"TensorNode",
"UAIModel",
"argmax_trace",
"build_contraction_tree",
"build_network",
"build_tropical_factors",
"choose_order",
"contract_tree",
"mpe_tropical",
"read_evidence_file",
"read_model_file",
"read_model_from_string",
"recover_mpe_assignment",
"safe_log",
"tropical_einsum",
]
Loading