Inference in Graphical Models

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

import numpy as np
from typing import Dict, List, Set, Tuple
from dataclasses import dataclass
 
@dataclass
class Factor:
    """Represents a factor (potential function) over discrete variables.
    
    Each variable is indexed by position, values stored in a multi-dimensional array.
    """
    variables: Tuple[str, ...]  # Variable names in scope
    cardinalities: Tuple[int, ...]  # Number of values for each variable
    values: np.ndarray  # Multi-dimensional array of factor values
    
    @property
    def scope(self) -> Set[str]:
        return set(self.variables)
    
    def get_value(self, assignment: Dict[str, int]) -> float:
        """Get factor value for a specific variable assignment."""
        indices = tuple(assignment[var] for var in self.variables)
        return self.values[indices]
 
 
def factor_product(phi1: Factor, phi2: Factor) -> Factor:
    """Compute the product of two factors.
    
    Result scope is union of input scopes. Each entry is product of
    corresponding entries from input factors.
    
    Time complexity: O(prod of cardinalities in combined scope)
    Space complexity: O(prod of cardinalities in combined scope)
    """
    # Determine combined scope
    combined_vars = list(phi1.variables)
    phi2_only_vars = [v for v in phi2.variables if v not in phi1.scope]
    combined_vars.extend(phi2_only_vars)
    
    # Build cardinality mapping
    var_to_card = dict(zip(phi1.variables, phi1.cardinalities))
    var_to_card.update(dict(zip(phi2.variables, phi2.cardinalities)))
    combined_cards = tuple(var_to_card[v] for v in combined_vars)
    
    # Compute product values
    result_shape = combined_cards
    result_values = np.zeros(result_shape)
    
    # Iterate over all configurations
    for idx in np.ndindex(result_shape):
        assignment = dict(zip(combined_vars, idx))
        
        # Get corresponding indices in original factors
        idx1 = tuple(assignment[v] for v in phi1.variables)
        idx2 = tuple(assignment[v] for v in phi2.variables)
        
        result_values[idx] = phi1.values[idx1] * phi2.values[idx2]
    
    return Factor(
        variables=tuple(combined_vars),
        cardinalities=combined_cards,
        values=result_values
    )
 
 
def factor_marginalize(phi: Factor, var: str) -> Factor:
    """Marginalize (sum out) a variable from the factor.
    
    Result scope is input scope minus the marginalized variable.
    Each entry is sum over all values of eliminated variable.
    
    Time complexity: O(size of input factor)
    Space complexity: O(size of output factor)
    """
    if var not in phi.scope:
        return phi  # Variable not in scope, return unchanged
    
    # Find position and cardinality of variable to eliminate
    var_idx = phi.variables.index(var)
    
    # New scope and cardinalities
    new_vars = tuple(v for v in phi.variables if v != var)
    new_cards = tuple(c for i, c in enumerate(phi.cardinalities) if i != var_idx)
    
    # Sum over the variable being eliminated
    result_values = np.sum(phi.values, axis=var_idx)
    
    return Factor(
        variables=new_vars,
        cardinalities=new_cards,
        values=result_values
    )
 
 
def factor_reduce(phi: Factor, evidence: Dict[str, int]) -> Factor:
    """Reduce factor by fixing variables to observed values (evidence).
    
    Result scope excludes observed variables.
    Only configurations consistent with evidence are retained.
    
    Time complexity: O(size of output factor)
    Space complexity: O(size of output factor)
    """
    remaining_vars = []
    remaining_cards = []
    slices = []
    
    for i, var in enumerate(phi.variables):
        if var in evidence:
            # This variable is observed: select specific value
            slices.append(evidence[var])
        else:
            # This variable is free: keep all values
            remaining_vars.append(var)
            remaining_cards.append(phi.cardinalities[i])
            slices.append(slice(None))
    
    # Extract sub-array corresponding to evidence
    result_values = phi.values[tuple(slices)]
    
    # Handle case where result is scalar (all vars observed)
    if result_values.ndim == 0:
        result_values = np.array([float(result_values)])
        remaining_vars = []
        remaining_cards = []
    
    return Factor(
        variables=tuple(remaining_vars),
        cardinalities=tuple(remaining_cards),
        values=result_values
    )
 
 
def factor_normalize(phi: Factor) -> Factor:
    """Normalize factor to sum to 1 (valid probability distribution)."""
    total = np.sum(phi.values)
    if total == 0:
        raise ValueError("Cannot normalize factor with zero sum")
    return Factor(
        variables=phi.variables,
        cardinalities=phi.cardinalities,
        values=phi.values / total
    )

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

from typing import List, Dict, Set, Optional
from functools import reduce
 
class VariableElimination:
    """Complete Variable Elimination implementation for exact inference.
    
    Computes marginal and conditional probabilities in graphical models
    by systematically eliminating variables through factor operations.
    """
    
    def __init__(self, factors: List[Factor], variables: Set[str]):
        """
        Initialize with factors and the set of all variables.
        
        Args:
            factors: List of Factor objects representing the model
            variables: Set of all variable names in the model
        """
        self.original_factors = factors
        self.variables = variables
    
    def query(
        self, 
        query_vars: List[str],
        evidence: Optional[Dict[str, int]] = None,
        elimination_order: Optional[List[str]] = None
    ) -> Factor:
        """
        Compute P(query_vars | evidence) via variable elimination.
        
        Args:
            query_vars: Variables whose marginal distribution we want
            evidence: Observed variable assignments (variable -> value)
            elimination_order: Order to eliminate hidden variables
            
        Returns:
            Normalized factor representing P(query_vars | evidence)
        """
        evidence = evidence or {}
        
        # Step 1: Create working copy of factors
        factors = [Factor(f.variables, f.cardinalities, f.values.copy()) 
                   for f in self.original_factors]
        
        # Step 2: Incorporate evidence by reducing factors
        if evidence:
            factors = [factor_reduce(f, evidence) for f in factors]
        
        # Step 3: Determine variables to eliminate
        # Hidden = all variables - query variables - evidence variables
        hidden_vars = self.variables - set(query_vars) - set(evidence.keys())
        
        # Step 4: Determine elimination order
        if elimination_order is None:
            # Default: arbitrary order (not optimal!)
            elim_order = list(hidden_vars)
        else:
            # Use specified order, filtered to hidden variables
            elim_order = [v for v in elimination_order if v in hidden_vars]
        
        # Step 5: Eliminate variables one by one
        for var in elim_order:
            factors = self._eliminate_variable(factors, var)
        
        # Step 6: Multiply remaining factors
        result = self._multiply_all_factors(factors)
        
        # Step 7: Normalize to get proper probability distribution
        result = factor_normalize(result)
        
        return result
    
    def _eliminate_variable(
        self, 
        factors: List[Factor], 
        var: str
    ) -> List[Factor]:
        """
        Eliminate a single variable from the factor list.
        
        1. Identify all factors containing the variable
        2. Multiply them together
        3. Marginalize the variable from the product
        4. Add the new factor back (replacing the originals)
        
        This is the core operation that makes VE efficient.
        """
        # Partition factors by whether they contain var
        factors_with_var = [f for f in factors if var in f.scope]
        factors_without_var = [f for f in factors if var not in f.scope]
        
        if not factors_with_var:
            # Variable doesn't appear in any factor
            return factors
        
        # Multiply all factors containing var
        product = factors_with_var[0]
        for f in factors_with_var[1:]:
            product = factor_product(product, f)
        
        # Marginalize var from the product
        new_factor = factor_marginalize(product, var)
        
        # Return updated factor list
        return factors_without_var + [new_factor]
    
    def _multiply_all_factors(self, factors: List[Factor]) -> Factor:
        """Multiply all factors together."""
        if not factors:
            raise ValueError("No factors to multiply")
        return reduce(factor_product, factors)
    
    def compute_marginal_likelihood(
        self, 
        evidence: Dict[str, int],
        elimination_order: Optional[List[str]] = None
    ) -> float:
        """
        Compute P(evidence) - the marginal likelihood of observations.
        
        This is done by eliminating ALL variables (none are query variables).
        The result is a scalar representing the probability of evidence.
        """
        # Eliminate all non-evidence variables
        factors = [Factor(f.variables, f.cardinalities, f.values.copy()) 
                   for f in self.original_factors]
        factors = [factor_reduce(f, evidence) for f in factors]
        
        hidden_vars = self.variables - set(evidence.keys())
        elim_order = elimination_order or list(hidden_vars)
        elim_order = [v for v in elim_order if v in hidden_vars]
        
        for var in elim_order:
            factors = self._eliminate_variable(factors, var)
        
        # Result should be a single scalar factor
        result = self._multiply_all_factors(factors)
        return float(np.sum(result.values))

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

def min_fill_heuristic(
    variables: Set[str], 
    factors: List[Factor]
) -> List[str]:
    """
    Compute elimination order using min-fill heuristic.
    
    At each step, eliminate the variable that adds the fewest
    fill edges to the induced graph. This greedy approach often
    produces orderings with low width.
    
    Args:
        variables: Set of variables to eliminate
        factors: Original factors defining the graph structure
        
    Returns:
        Ordered list of variables (elimination order)
    """
    # Build initial adjacency from factor scopes
    # Two variables are adjacent if they appear in the same factor
    adjacency = {v: set() for v in variables}
    for f in factors:
        for v1 in f.scope:
            for v2 in f.scope:
                if v1 != v2:
                    adjacency[v1].add(v2)
                    adjacency[v2].add(v1)
    
    remaining = set(variables)
    order = []
    
    while remaining:
        # Find variable with minimum fill-in
        best_var = None
        best_fill = float('inf')
        
        for var in remaining:
            # Compute fill: number of edges added if we eliminate var
            neighbors = adjacency[var] & remaining
            fill = 0
            neighbors_list = list(neighbors)
            for i, n1 in enumerate(neighbors_list):
                for n2 in neighbors_list[i+1:]:
                    if n2 not in adjacency[n1]:
                        fill += 1
            
            if fill < best_fill:
                best_fill = fill
                best_var = var
        
        # Eliminate best_var: add fill edges
        neighbors = adjacency[best_var] & remaining
        neighbors_list = list(neighbors)
        for i, n1 in enumerate(neighbors_list):
            for n2 in neighbors_list[i+1:]:
                adjacency[n1].add(n2)
                adjacency[n2].add(n1)
        
        order.append(best_var)
        remaining.remove(best_var)
    
    return order
 
 
def min_neighbors_heuristic(
    variables: Set[str], 
    factors: List[Factor]
) -> List[str]:
    """
    Simpler heuristic: always eliminate variable with fewest neighbors.
    
    Faster than min-fill but often produces slightly worse orderings.
    """
    adjacency = {v: set() for v in variables}
    for f in factors:
        for v1 in f.scope:
            for v2 in f.scope:
                if v1 != v2:
                    adjacency[v1].add(v2)
    
    remaining = set(variables)
    order = []
    
    while remaining:
        # Find variable with minimum remaining neighbors
        best_var = min(remaining, 
                       key=lambda v: len(adjacency[v] & remaining))
        
        # Add fill edges
        neighbors = list(adjacency[best_var] & remaining)
        for i, n1 in enumerate(neighbors):
            for n2 in neighbors[i+1:]:
                adjacency[n1].add(n2)
                adjacency[n2].add(n1)
        
        order.append(best_var)
        remaining.remove(best_var)
    
    return order
 
 
def compute_elimination_width(
    order: List[str], 
    factors: List[Factor]
) -> int:
    """
    Compute the width (max clique size - 1) of an elimination order.
    
    This metric determines the worst-case factor size created during VE.
    """
    adjacency = {v: set() for v in order}
    for f in factors:
        for v1 in f.scope:
            for v2 in f.scope:
                if v1 != v2:
                    adjacency[v1].add(v2)
    
    remaining = set(order)
    max_width = 0
    
    for var in order:
        # Width is number of remaining neighbors when eliminated
        width = len(adjacency[var] & remaining)
        max_width = max(max_width, width)
        
        # Add fill edges
        neighbors = list(adjacency[var] & remaining)
        for i, n1 in enumerate(neighbors):
            for n2 in neighbors[i+1:]:
                adjacency[n1].add(n2)
                adjacency[n2].add(n1)
        
        remaining.remove(var)
    
    return max_width

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69

import numpy as np
from scipy.special import logsumexp
 
def log_factor_marginalize(log_phi: np.ndarray, axis: int) -> np.ndarray:
    """
    Marginalize a log-space factor along the specified axis.
    
    Uses logsumexp for numerical stability:
    log(sum_x exp(log_phi)) = logsumexp(log_phi)
    
    Args:
        log_phi: Log-space factor values
        axis: Axis to marginalize over
        
    Returns:
        Log-space marginalized factor values
    """
    return logsumexp(log_phi, axis=axis)
 
 
def log_factor_product(log_phi1: np.ndarray, log_phi2: np.ndarray,
                       phi1_axes: List[int], phi2_axes: List[int],
                       result_axes: List[Tuple[int, str]]) -> np.ndarray:
    """
    Log-space factor product using broadcasting.
    
    In log space, multiplication becomes addition.
    The complexity is managing axis alignment.
    """
    # Expand dimensions for broadcasting
    shape1 = [1] * len(result_axes)
    shape2 = [1] * len(result_axes)
    
    for i, (_, source) in enumerate(result_axes):
        if source == 'phi1':
            shape1[i] = log_phi1.shape[phi1_axes.pop(0)]
        elif source == 'phi2':
            shape2[i] = log_phi2.shape[phi2_axes.pop(0)]
    
    # Reshape and add (product in log space)
    return log_phi1.reshape(shape1) + log_phi2.reshape(shape2)
 
 
def normalize_log_distribution(log_probs: np.ndarray) -> np.ndarray:
    """
    Normalize a log-probability distribution to sum to 1.
    
    Returns probabilities in standard (not log) space.
    """
    log_Z = logsumexp(log_probs)  # log of normalizing constant
    log_normalized = log_probs - log_Z
    return np.exp(log_normalized)
 
 
# Example: Avoiding underflow
def demo_log_space_necessity():
    """Demonstrate why log space is essential."""
    # Consider a chain of 100 binary variables
    n = 100
    p = 0.1  # Low probability at each step
    
    # Standard space: underflows to 0
    prob_standard = p ** n
    print(f"Standard space: {prob_standard}")  # 0.0 (underflow)
    
    # Log space: handles correctly
    log_prob = n * np.log(p)
    print(f"Log probability: {log_prob}")  # -230.26 (correct)
    print(f"Actual probability: {np.exp(log_prob)}")  # 1e-100 (reconstituted)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

from typing import Dict, Tuple
 
def factor_max_marginalize(phi: Factor, var: str) -> Tuple[Factor, Dict]:
    """
    Max-marginalize a variable (for MAP inference).
    
    Returns:
        - New factor with max values
        - Traceback dictionary mapping remaining configs to argmax values
    """
    var_idx = phi.variables.index(var)
    new_vars = tuple(v for v in phi.variables if v != var)
    new_cards = tuple(c for i, c in enumerate(phi.cardinalities) if i != var_idx)
    
    # Take max along variable axis
    max_values = np.max(phi.values, axis=var_idx)
    argmax_values = np.argmax(phi.values, axis=var_idx)
    
    new_factor = Factor(
        variables=new_vars,
        cardinalities=new_cards,
        values=max_values
    )
    
    return new_factor, argmax_values
 
 
class MAPVariableElimination:
    """Variable elimination for MAP inference."""
    
    def __init__(self, factors: List[Factor], variables: Set[str]):
        self.factors = factors
        self.variables = variables
        
    def map_query(
        self, 
        evidence: Dict[str, int],
        elimination_order: Optional[List[str]] = None
    ) -> Dict[str, int]:
        """
        Compute argmax_X P(X | evidence).
        
        Returns the most probable assignment to all non-evidence variables.
        """
        # Reduce factors by evidence
        factors = [factor_reduce(f, evidence) for f in self.factors]
        
        hidden_vars = self.variables - set(evidence.keys())
        elim_order = elimination_order or list(hidden_vars)
        
        # Track argmax for backtracking
        traceback = {}  # var -> (prev_factors, argmax_array)
        
        # Eliminate variables with max instead of sum
        for var in elim_order:
            factors_with_var = [f for f in factors if var in f.scope]
            factors_without_var = [f for f in factors if var not in f.scope]
            
            if factors_with_var:
                product = reduce(factor_product, factors_with_var)
                new_factor, argmax = factor_max_marginalize(product, var)
                traceback[var] = (product.variables, argmax)
                factors = factors_without_var + [new_factor]
        
        # Backtrack to recover MAP assignment
        assignment = {}
        
        # Reverse elimination order for backtracking
        for var in reversed(elim_order):
            if var in traceback:
                orig_vars, argmax = traceback[var]
                # Build index from already-assigned variables
                idx = []
                for v in orig_vars:
                    if v != var and v in assignment:
                        idx.append(assignment[v])
                assignment[var] = int(argmax[tuple(idx)]) if idx else int(argmax)
        
        return assignment