Make FactorGraph aware of factor groups (#59)

* Make factor graph aware of factor groups

* Docstrings

* Fix tests

* Address comments

Author: StannisZhou

.isort.cfg

@@ -1,2 +1,2 @@
 [settings]
-profile=black
+profile=black

pgmax/fg/graph.py

@@ -14,10 +14,7 @@
 
 @dataclass
 class FactorGraph:
-    """Base class to represent a factor graph.
-
-    Concrete factor graphs inherits from this class, and specifies get_evidence to generate
-    the evidence array, and optionally init_msgs (default to initializing all messages to 0)
+    """Class for representing a factor graph
 
     Args:
         variable_groups: A container containing multiple VariableGroups, or a CompositeVariableGroup.
@@ -32,7 +29,7 @@ class FactorGraph:
 
     Attributes:
         _composite_variable_group: CompositeVariableGroup. contains all involved VariableGroups
-        _factors: list. contains all involved factors
+        _factor_groups: List of added factor groups
         num_var_states: int. represents the sum of all variable states of all variables in the
             FactorGraph
         _vars_to_starts: MappingProxyType[nodes.Variable, int]. maps every variable to an int
@@ -84,7 +81,7 @@ def __post_init__(self):
 
         self._vars_to_evidence: Dict[nodes.Variable, np.ndarray] = {}
 
-        self._factors: List[nodes.EnumerationFactor] = []
+        self._factor_groups: List[groups.FactorGroup] = []
 
     def add_factors(
         self,
@@ -113,20 +110,24 @@ def add_factors(
         """
         factor_factory = kwargs.pop("factor_factory", None)
         if factor_factory is not None:
-            factors = factor_factory(
+            factor_group = factor_factory(
                 self._composite_variable_group, *args, **kwargs
-            ).factors
+            )
         else:
             if len(args) > 0:
                 new_args = list(args)
-                new_args[0] = tuple(self._composite_variable_group[args[0]])
-                factors = [nodes.EnumerationFactor(*new_args, **kwargs)]
+                new_args[0] = [args[0]]
+                factor_group = groups.EnumerationFactorGroup(
+                    self._composite_variable_group, *new_args, **kwargs
+                )
             else:
                 keys = kwargs.pop("keys")
-                kwargs["variables"] = self._composite_variable_group[keys]
-                factors = [nodes.EnumerationFactor(**kwargs)]
+                kwargs["connected_var_keys"] = [keys]
+                factor_group = groups.EnumerationFactorGroup(
+                    self._composite_variable_group, **kwargs
+                )
 
-        self._factors.extend(factors)
+        self._factor_groups.append(factor_group)
 
     @property
     def wiring(self) -> nodes.EnumerationWiring:
@@ -138,7 +139,8 @@ def wiring(self) -> nodes.EnumerationWiring:
             compiled wiring from each individual factor
         """
         wirings = [
-            factor.compile_wiring(self._vars_to_starts) for factor in self._factors
+            factor_group.compile_wiring(self._vars_to_starts)
+            for factor_group in self._factor_groups
         ]
         wiring = fg_utils.concatenate_enumeration_wirings(wirings)
         return wiring
@@ -154,7 +156,10 @@ def factor_configs_log_potentials(self) -> np.ndarray:
                 valid configuration
         """
         return np.concatenate(
-            [factor.factor_configs_log_potentials for factor in self._factors]
+            [
+                factor_group.factor_group_log_potentials
+                for factor_group in self._factor_groups
+            ]
         )
 
     @property
@@ -182,6 +187,11 @@ def evidence(self) -> np.ndarray:
 
         return evidence
 
+    @property
+    def factors(self) -> Tuple[nodes.EnumerationFactor, ...]:
+        """List of individual factors in the factor graph"""
+        return sum([factor_group.factors for factor_group in self._factor_groups], ())
+
     def get_init_msgs(self, context: Any = None):
         """Function to initialize messages.
 
@@ -201,7 +211,7 @@ def set_evidence(
         self,
         key: Union[Tuple[Any, ...], Any],
         evidence: Union[Dict[Any, np.ndarray], np.ndarray],
-    ):
+    ) -> None:
         """Function to update the evidence for variables in the FactorGraph.
 
         Args:

pgmax/fg/groups.py

@@ -7,6 +7,7 @@
 import numpy as np
 
 import pgmax.fg.nodes as nodes
+from pgmax.fg import fg_utils
 from pgmax.utils import cached_property
 
 
@@ -375,6 +376,35 @@ def __post_init__(self) -> None:
     def factors(self) -> Tuple[nodes.EnumerationFactor, ...]:
         raise NotImplementedError("Needs to be overriden by subclass")
 
+    def compile_wiring(
+        self, vars_to_starts: Mapping[nodes.Variable, int]
+    ) -> nodes.EnumerationWiring:
+        """Function to compile wiring for the factor group.
+
+        Args:
+            vars_to_starts: A dictionary that maps variables to their global starting indices
+                For an n-state variable, a global start index of m means the global indices
+                of its n variable states are m, m + 1, ..., m + n - 1
+
+        Returns:
+            compiled wiring for the factor group
+        """
+        wirings = [factor.compile_wiring(vars_to_starts) for factor in self.factors]
+        wiring = fg_utils.concatenate_enumeration_wirings(wirings)
+        return wiring
+
+    @cached_property
+    def factor_group_log_potentials(self) -> np.ndarray:
+        """Function to compile potential array for the factor group
+
+        Returns:
+            a jnp array representing the log of the potential function for
+            the factor group
+        """
+        return np.concatenate(
+            [factor.factor_configs_log_potentials for factor in self.factors]
+        )
+
 
 @dataclass(frozen=True, eq=False)
 class EnumerationFactorGroup(FactorGroup):
@@ -387,14 +417,14 @@ class EnumerationFactorGroup(FactorGroup):
     Args:
         factor_configs: Array of shape (num_val_configs, num_variables)
             An array containing explicit enumeration of all valid configurations
+        factor_configs_log_potentials: Optional array of shape (num_val_configs,).
+            If specified, it contains the log of the potential value for every possible configuration.
+            If none, it is assumed the log potential is uniform 0 and such an array is automatically
+            initialized.
 
     Attributes:
         factors: a tuple of all the factors belonging to this group. These are constructed
             internally by invoking the _get_connected_var_keys_for_factors method.
-        factor_configs_log_potentials: Optional ndarray of shape (num_val_configs,).
-            if specified. Must contain the log of the potential value for every possible
-            configuration. If left unspecified, it is assumed the log potential is uniform
-            0 and such an array is automatically initialized.
     """
 
     factor_configs: np.ndarray

tests/test_pgmax.py

@@ -306,7 +306,7 @@ def create_valid_suppression_config_arr(suppression_diameter):
             else:
                 fg.add_factors(
                     keys=curr_keys,
-                    configs=valid_configs_non_supp,
+                    factor_configs=valid_configs_non_supp,
                     factor_configs_log_potentials=np.zeros(
                         valid_configs_non_supp.shape[0], dtype=float
                     ),
@@ -417,4 +417,4 @@ def binary_connected_variables(
 
     assert isinstance(fg.evidence, np.ndarray)
 
-    assert len(fg._factors) == 7056
+    assert len(fg.factors) == 7056