misc

funsor/sum_product.py

@@ -186,29 +186,20 @@ def partial_sum_product(sum_op, prod_op, factors, eliminate=frozenset(), plates=
     for f in factors:
         ordinal = plates.intersection(f.inputs)
         ordinal_to_factors[ordinal].append(f)
-        # for var in sum_vars.intersection(f.inputs):
-        for var in frozenset(f.inputs) - plates:
+        for var in sum_vars.intersection(f.inputs):
             var_to_ordinal[var] = var_to_ordinal.get(var, ordinal) & ordinal
 
     ordinal_to_vars = defaultdict(set)
-    ordinal_to_not_summed = defaultdict(set)
     for var, ordinal in var_to_ordinal.items():
-        if var in sum_vars:
-            ordinal_to_vars[ordinal].add(var)
-        else:
-            ordinal_to_not_summed[ordinal].add(var)
+        ordinal_to_vars[ordinal].add(var)
 
     results = []
     while ordinal_to_factors:
         leaf = max(ordinal_to_factors, key=len)
         leaf_factors = ordinal_to_factors.pop(leaf)
         leaf_reduce_vars = ordinal_to_vars[leaf]
         for (group_factors, group_vars) in _partition(leaf_factors, leaf_reduce_vars):
-            f = reduce(prod_op, group_factors).reduce(sum_op, group_vars & sum_vars)
-            #  not_summed_vars = frozenset(f.inputs) & ordinal_to_not_summed[leaf]
-            #  parent_vars = sum_vars & frozenset(f.inputs)
-            #  if not_summed_vars and parent_vars:
-            #      raise ValueError("intractable!")
+            f = reduce(prod_op, group_factors).reduce(sum_op, group_vars)
             remaining_sum_vars = sum_vars.intersection(f.inputs)
             if not remaining_sum_vars:
                 results.append(f.reduce(prod_op, leaf & eliminate))
@@ -284,30 +275,22 @@ def modified_partial_sum_product(sum_op, prod_op, factors,
     for f in factors:
         ordinal = plates.intersection(f.inputs)
         ordinal_to_factors[ordinal].append(f)
-        # for var in sum_vars.intersection(f.inputs):
-        for var in frozenset(f.inputs) - plates:
+        for var in sum_vars.intersection(f.inputs):
             var_to_ordinal[var] = var_to_ordinal.get(var, ordinal) & ordinal
 
     ordinal_to_vars = defaultdict(set)
-    ordinal_to_not_summed = defaultdict(set)
     for var, ordinal in var_to_ordinal.items():
-        if var in sum_vars:
             ordinal_to_vars[ordinal].add(var)
-        else:
-            ordinal_to_not_summed[ordinal].add(var)
 
     results = []
     while ordinal_to_factors:
         leaf = max(ordinal_to_factors, key=len)
         leaf_factors = ordinal_to_factors.pop(leaf)
         leaf_reduce_vars = ordinal_to_vars[leaf]
         for (group_factors, group_vars) in _partition(leaf_factors, leaf_reduce_vars | markov_prod_vars):
+            # eliminate markov vars
             nonmarkov_vars = group_vars - markov_sum_vars - markov_prod_vars
             f = reduce(prod_op, group_factors).reduce(sum_op, nonmarkov_vars)
-            #  not_summed_vars = frozenset(f.inputs) & ordinal_to_not_summed[leaf]
-            #  if not_summed_vars and (sum_vars & frozenset(f.inputs)):
-            #      raise ValueError("intractable!")
-            # eliminate markov vars
             markov_vars = group_vars.intersection(markov_sum_vars)
             #  cond_vars |= frozenset(f.inputs) - plates
             if markov_vars:

test/test_sum_product.py

@@ -81,8 +81,6 @@ def test_partition(inputs, dims, expected_num_components):
 @pytest.mark.parametrize('sum_op,prod_op', [(ops.add, ops.mul), (ops.logaddexp, ops.add)])
 @pytest.mark.parametrize('inputs,plates', [('a,abi,bcij', 'ij')])
 @pytest.mark.parametrize('vars1,vars2', [
-    #  ('acj', 'bi'),
-    #  ('a', 'bcij'),
     ('', 'abcij'),
     ('c', 'abij'),
     ('cj', 'abi'),
@@ -163,8 +161,6 @@ def test_modified_partial_sum_product_0(sum_op, prod_op, vars1, vars2,
 
 
 @pytest.mark.parametrize('vars1,vars2', [
-    (frozenset({"time", "x_0", "x_prev", "x_curr"}),
-     frozenset({"y_0", "y_curr"})),
     (frozenset(),
      frozenset({"time", "x_0", "x_prev", "x_curr", "y_0", "y_curr"})),
     (frozenset({"y_0", "y_curr"}),