Repository: pgmpy/pgmpy

Issue:
Implement Network Pruning for optimizing Inference
### Subject of the issue
In some cases, some of the variables in the model do not effect the results of inference. The network can, therefore, be pruned before any inference is run on it. And since pruning would reduce the size of the network, it should make inference more efficient.

Example Implementation: https://github.com/ankurankan/pgmpy/blob/feature/optimize_ve/pgmpy/inference/ExactInference.py#L17
Reference: https://arxiv.org/pdf/1304.1112.pdf

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/pgmpy/pgmpy/pull/1398?src=pr&el=h1) Report
> Merging [#1398](https://codecov.io/gh/pgmpy/pgmpy/pull/1398?src=pr&el=desc) (7d87060) into [dev](https://codecov.io/gh/pgmpy/pgmpy/commit/2ae8c904ee59ffffe2ff7bbc794a075cc37646ae?el=desc) (2ae8c90) will **increase** coverage by `0.02%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/pgmpy/pgmpy/pull/1398/graphs/tree.svg?width=650&height=150&src=pr&token=UaJMCdHaEF)](https://codecov.io/gh/pgmpy/pgmpy/pull/1398?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##              dev    #1398      +/-   ##
==========================================
+ Coverage   93.28%   93.30%   +0.02%     
==========================================
  Files         134      134              
  Lines       13476    13521      +45     
==========================================
+ Hits        12571    12616      +45     
  Misses        905      905              
```


| [Impacted Files](https://codecov.io/gh/pgmpy/pgmpy/pull/1398?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pgmpy/base/DAG.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1398/diff?src=pr&el=tree#diff-cGdtcHkvYmFzZS9EQUcucHk=) | `94.11% <ø> (ø)` | |
| [pgmpy/tests/test\_inference/test\_ExactInference.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1398/diff?src=pr&el=tree#diff-cGdtcHkvdGVzdHMvdGVzdF9pbmZlcmVuY2UvdGVzdF9FeGFjdEluZmVyZW5jZS5weQ==) | `100.00% <ø> (ø)` | |
| [pgmpy/inference/ExactInference.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1398/diff?src=pr&el=tree#diff-cGdtcHkvaW5mZXJlbmNlL0V4YWN0SW5mZXJlbmNlLnB5) | `95.22% <100.00%> (+0.57%)` | :arrow_up: |
| [pgmpy/inference/base.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1398/diff?src=pr&el=tree#diff-cGdtcHkvaW5mZXJlbmNlL2Jhc2UucHk=) | `100.00% <100.00%> (ø)` | |
| [pgmpy/inference/bn\_inference.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1398/diff?src=pr&el=tree#diff-cGdtcHkvaW5mZXJlbmNlL2JuX2luZmVyZW5jZS5weQ==) | `96.29% <100.00%> (+0.06%)` | :arrow_up: |
| [pgmpy/inference/dbn\_inference.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1398/diff?src=pr&el=tree#diff-cGdtcHkvaW5mZXJlbmNlL2Ribl9pbmZlcmVuY2UucHk=) | `96.55% <100.00%> (+0.02%)` | :arrow_up: |
| [pgmpy/inference/mplp.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1398/diff?src=pr&el=tree#diff-cGdtcHkvaW5mZXJlbmNlL21wbHAucHk=) | `97.14% <100.00%> (ø)` | |
| [pgmpy/tests/test\_inference/test\_Inference.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1398/diff?src=pr&el=tree#diff-cGdtcHkvdGVzdHMvdGVzdF9pbmZlcmVuY2UvdGVzdF9JbmZlcmVuY2UucHk=) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pgmpy/pgmpy/pull/1398?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pgmpy/pgmpy/pull/1398?src=pr&el=footer). Last update [2ae8c90...7d87060](https://codecov.io/gh/pgmpy/pgmpy/pull/1398?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/pgmpy/base/DAG.py b/pgmpy/base/DAG.py
index 46775525..1868d7fb 100644
--- a/pgmpy/base/DAG.py
+++ b/pgmpy/base/DAG.py
@@ -586,12 +586,15 @@ class DAG(nx.DiGraph):
         """
         Returns a dictionary with the given variables as keys and all the nodes reachable
         from that respective variable as values.
+
         Parameters
         ----------
         variables: str or array like
             variables whose active trails are to be found.
+
         observed : List of nodes (optional)
             If given the active trails would be computed assuming these nodes to be observed.
+
         Examples
         --------
         >>> from pgmpy.base import DAG
diff --git a/pgmpy/inference/ExactInference.py b/pgmpy/inference/ExactInference.py
index 1fb58a95..a9a3612b 100644
--- a/pgmpy/inference/ExactInference.py
+++ b/pgmpy/inference/ExactInference.py
@@ -254,7 +254,13 @@ class VariableElimination(Inference):
                 f"Can't have the same variables in both `variables` and `evidence`. Found in both: {common_vars}"
             )
 
-        return self._variable_elimination(
+        # Make a copy of the original model as it will be replaced during pruning.
+        orig_model = self.model
+        if isinstance(self.model, BayesianModel):
+            self.model, evidence = self._prune_bayesian_model(variables, evidence)
+        self._initialize_structures()
+
+        result = self._variable_elimination(
             variables=variables,
             operation="marginalize",
             evidence=evidence,
@@ -262,6 +268,9 @@ class VariableElimination(Inference):
             joint=joint,
             show_progress=show_progress,
         )
+        self.model = orig_model
+
+        return result
 
     def max_marginal(
         self,
@@ -277,9 +286,11 @@ class VariableElimination(Inference):
         ----------
         variables: list
             list of variables over which we want to compute the max-marginal.
+
         evidence: dict
             a dict key, value pair as {var: state_of_var_observed}
             None if no evidence
+
         elimination_order: list
             order of variable eliminations (if nothing is provided) order is
             computed automatically
@@ -308,6 +319,12 @@ class VariableElimination(Inference):
                 f"Can't have the same variables in both `variables` and `evidence`. Found in both: {common_vars}"
             )
 
+        # Make a copy of the original model and replace self.model with it later.
+        orig_model = self.model
+        if isinstance(self.model, BayesianModel):
+            self.model, evidence = self._prune_bayesian_model(variables, evidence)
+        self._initialize_structures()
+
         final_distribution = self._variable_elimination(
             variables=variables,
             operation="maximize",
@@ -316,6 +333,7 @@ class VariableElimination(Inference):
             show_progress=show_progress,
         )
 
+        self.model = orig_model
         return np.max(final_distribution.values)
 
     def map_query(
@@ -355,14 +373,22 @@ class VariableElimination(Inference):
         >>> inference = VariableElimination(model)
         >>> phi_query = inference.map_query(['A', 'B'])
         """
+        variables = [] if variables is None else variables
         common_vars = set(evidence if evidence is not None else []).intersection(
-            set(variables if variables is not None else [])
+            variables
         )
         if common_vars:
             raise ValueError(
                 f"Can't have the same variables in both `variables` and `evidence`. Found in both: {common_vars}"
             )
 
+        # Make a copy of the original model and replace self.model with it later
+        orig_model = self.model
+
+        if isinstance(self.model, BayesianModel):
+            self.model, evidence = self._prune_bayesian_model(variables, evidence)
+        self._initialize_structures()
+
         # TODO:Check the note in docstring. Change that behavior to return the joint MAP
         final_distribution = self._variable_elimination(
             variables=variables,
@@ -372,6 +398,8 @@ class VariableElimination(Inference):
             joint=True,
             show_progress=show_progress,
         )
+        self.model = orig_model
+
         argmax = np.argmax(final_distribution.values)
         assignment = final_distribution.assignment([argmax])[0]
 
@@ -410,6 +438,8 @@ class VariableElimination(Inference):
         >>> inference = VariableElimination(model)
         >>> inference.induced_graph(['C', 'D', 'A', 'B', 'E'])
         """
+        self._initialize_structures()
+
         # If the elimination order does not contain the same variables as the model
         if set(elimination_order) != set(self.variables):
             raise ValueError(
@@ -881,6 +911,11 @@ class BeliefPropagation(Inference):
                 f"Can't have the same variables in both `variables` and `evidence`. Found in both: {common_vars}"
             )
 
+        orig_model = self.model
+        if isinstance(self.model, BayesianModel):
+            self.model, evidence = self._prune_bayesian_model(variables, evidence)
+        self._initialize_structures()
+
         result = self._query(
             variables=variables,
             operation="marginalize",
@@ -888,6 +923,8 @@ class BeliefPropagation(Inference):
             joint=joint,
             show_progress=show_progress,
         )
+        self.model = orig_model
+
         if joint:
             return result.normalize(inplace=False)
         else:
@@ -935,9 +972,11 @@ class BeliefPropagation(Inference):
         >>> belief_propagation.map_query(variables=['J', 'Q'],
         ...                              evidence={'A': 0, 'R': 0, 'G': 0, 'L': 1})
         """
+        variables = [] if variables is None else variables
         common_vars = set(evidence if evidence is not None else []).intersection(
-            set(variables if variables is not None else [])
+            variables
         )
+
         if common_vars:
             raise ValueError(
                 f"Can't have the same variables in both `variables` and `evidence`. Found in both: {common_vars}"
@@ -945,7 +984,13 @@ class BeliefPropagation(Inference):
 
         # TODO:Check the note in docstring. Change that behavior to return the joint MAP
         if not variables:
-            variables = set(self.variables)
+            variables = list(self.model.nodes())
+
+        # Make a copy of the original model and then replace self.model with it later.
+        orig_model = self.model
+        if isinstance(self.model, BayesianModel):
+            self.model, evidence = self._prune_bayesian_model(variables, evidence)
+        self._initialize_structures()
 
         final_distribution = self._query(
             variables=variables,
@@ -954,6 +999,8 @@ class BeliefPropagation(Inference):
             show_progress=show_progress,
         )
 
+        self.model = orig_model
+
         # To handle the case when no argument is passed then
         # _variable_elimination returns a dict.
         argmax = np.argmax(final_distribution.values)
diff --git a/pgmpy/inference/base.py b/pgmpy/inference/base.py
index c419f9ba..a8ad63c8 100644
--- a/pgmpy/inference/base.py
+++ b/pgmpy/inference/base.py
@@ -62,18 +62,23 @@ class Inference(object):
         self.model = model
         model.check_model()
 
-        if isinstance(model, JunctionTree):
-            self.variables = set(chain(*model.nodes()))
+    def _initialize_structures(self):
+        """
+        Initializes all the data structures which will
+        later be used by the inference algorithms.
+        """
+        if isinstance(self.model, JunctionTree):
+            self.variables = set(chain(*self.model.nodes()))
         else:
-            self.variables = model.nodes()
+            self.variables = self.model.nodes()
 
         self.cardinality = {}
         self.factors = defaultdict(list)
 
-        if isinstance(model, BayesianModel):
+        if isinstance(self.model, BayesianModel):
             self.state_names_map = {}
-            for node in model.nodes():
-                cpd = model.get_cpds(node)
+            for node in self.model.nodes():
+                cpd = self.model.get_cpds(node)
                 if isinstance(cpd, TabularCPD):
                     self.cardinality[node] = cpd.variable_card
                     cpd = cpd.to_factor()
@@ -81,21 +86,71 @@ class Inference(object):
                     self.factors[var].append(cpd)
                 self.state_names_map.update(cpd.no_to_name)
 
-        elif isinstance(model, (MarkovModel, FactorGraph, JunctionTree)):
-            self.cardinality = model.get_cardinality()
+        elif isinstance(self.model, (MarkovModel, FactorGraph, JunctionTree)):
+            self.cardinality = self.model.get_cardinality()
 
-            for factor in model.get_factors():
+            for factor in self.model.get_factors():
                 for var in factor.variables:
                     self.factors[var].append(factor)
 
-        elif isinstance(model, DynamicBayesianNetwork):
-            self.start_bayesian_model = BayesianModel(model.get_intra_edges(0))
-            self.start_bayesian_model.add_cpds(*model.get_cpds(time_slice=0))
-            cpd_inter = [model.get_cpds(node) for node in model.get_interface_nodes(1)]
-            self.interface_nodes = model.get_interface_nodes(0)
+        elif isinstance(self.model, DynamicBayesianNetwork):
+            self.start_bayesian_model = BayesianModel(self.model.get_intra_edges(0))
+            self.start_bayesian_model.add_cpds(*self.model.get_cpds(time_slice=0))
+            cpd_inter = [
+                self.model.get_cpds(node) for node in self.model.get_interface_nodes(1)
+            ]
+            self.interface_nodes = self.model.get_interface_nodes(0)
             self.one_and_half_model = BayesianModel(
-                model.get_inter_edges() + model.get_intra_edges(1)
+                self.model.get_inter_edges() + self.model.get_intra_edges(1)
             )
             self.one_and_half_model.add_cpds(
-                *(model.get_cpds(time_slice=1) + cpd_inter)
+                *(self.model.get_cpds(time_slice=1) + cpd_inter)
             )
+
+    def _prune_bayesian_model(self, variables, evidence):
+        """
+        Prunes unnecessary nodes from the model to optimize the computation.
+
+        Parameters
+        ----------
+        variables: list
+            The variables on which the query is done i.e. the variables whose
+            values we are interested in.
+
+        evidence: dict (default: None)
+            The variables whose values we know. The values can be specified as
+            {variable: state}.
+
+        Returns
+        -------
+        instance of pgmpy.models.BayesianModel: The pruned model.
+
+        Examples
+        --------
+        >>>
+        >>>
+
+        References
+        ----------
+        [1] Baker, M., & Boult, T. E. (2013). Pruning Bayesian networks for efficient computation. arXiv preprint arXiv:1304.1112.
+        """
+        bn = self.model.copy()
+        evidence = {} if evidence is None else evidence
+        variables = list(self.model.nodes()) if len(variables) == 0 else list(variables)
+
+        # Step 1: Remove all the variables that are d-separated from `variables` when conditioned
+        #         on `evidence`
+        d_connected = bn.active_trail_nodes(
+            variables=variables, observed=list(evidence.keys())
+        )
+        d_connected = set.union(*d_connected.values()).union(evidence.keys())
+        bn = bn.subgraph(d_connected)
+        evidence = {var: state for var, state in evidence.items() if var in d_connected}
+
+        # Step 2: Reduce the model to ancestral graph of [`variables` + `evidence`]
+        bn = bn.get_ancestral_graph(list(variables) + list(evidence.keys()))
+
+        # Step 3: Since all the CPDs are lost, add them back.
+        bn.add_cpds(*[self.model.get_cpds(var) for var in bn.nodes()])
+
+        return bn, evidence
diff --git a/pgmpy/inference/bn_inference.py b/pgmpy/inference/bn_inference.py
index 5c0c37e0..ac752057 100644
--- a/pgmpy/inference/bn_inference.py
+++ b/pgmpy/inference/bn_inference.py
@@ -25,6 +25,7 @@ class BayesianModelInference(Inference):
                 "Model expected type: BayesianModel, got type: ", type(model)
             )
         super(BayesianModelInference, self).__init__(model)
+        self._initialize_structures()
 
         self.topological_order = list(nx.topological_sort(model))
 
diff --git a/pgmpy/inference/dbn_inference.py b/pgmpy/inference/dbn_inference.py
index 6a42dc03..72540d1a 100644
--- a/pgmpy/inference/dbn_inference.py
+++ b/pgmpy/inference/dbn_inference.py
@@ -56,6 +56,8 @@ class DBNInference(Inference):
             http://www.cs.ubc.ca/~murphyk/Thesis/thesis.pdf
         """
         super(DBNInference, self).__init__(model)
+        self._initialize_structures()
+
         self.interface_nodes_0 = model.get_interface_nodes(time_slice=0)
         self.interface_nodes_1 = model.get_interface_nodes(time_slice=1)
 
diff --git a/pgmpy/inference/mplp.py b/pgmpy/inference/mplp.py
index 4659fcd4..45f8ed76 100644
--- a/pgmpy/inference/mplp.py
+++ b/pgmpy/inference/mplp.py
@@ -52,7 +52,7 @@ class Mplp(Inference):
             raise TypeError("Only MarkovModel is supported")
 
         super(Mplp, self).__init__(model)
-        self.model = model
+        self._initialize_structures()
 
         # S = \{c \cap c^{'} : c, c^{'} \in C, c \cap c^{'} \neq \emptyset\}
         self.intersection_set_variables = set()


Test patch:
diff --git a/pgmpy/tests/test_inference/test_ExactInference.py b/pgmpy/tests/test_inference/test_ExactInference.py
index 4695d39e..48bfdfbe 100644
--- a/pgmpy/tests/test_inference/test_ExactInference.py
+++ b/pgmpy/tests/test_inference/test_ExactInference.py
@@ -93,7 +93,6 @@ class TestVariableElimination(unittest.TestCase):
                 variables=["J"], cardinality=[2], values=np.array([0.416, 0.584])
             ),
         )
-
         query_result = self.bayesian_inference.query(["Q", "J"])
         query_result = self.bayesian_inference.query(["Q", "J"])
         self.assertEqual(
@@ -143,17 +142,17 @@ class TestVariableElimination(unittest.TestCase):
 
     def test_max_marginal_var(self):
         np_test.assert_almost_equal(
-            self.bayesian_inference.max_marginal(["G"]), 0.5714, decimal=4
+            self.bayesian_inference.max_marginal(["G"]), 0.6, decimal=4
         )
 
     def test_max_marginal_var1(self):
         np_test.assert_almost_equal(
-            self.bayesian_inference.max_marginal(["G", "R"]), 0.3740, decimal=4
+            self.bayesian_inference.max_marginal(["G", "R"]), 0.36, decimal=4
         )
 
     def test_max_marginal_var2(self):
         np_test.assert_almost_equal(
-            self.bayesian_inference.max_marginal(["G", "R", "A"]), 0.3061, decimal=4
+            self.bayesian_inference.max_marginal(["G", "R", "A"]), 0.288, decimal=4
         )
 
     def test_max_marginal_common_var(self):
diff --git a/pgmpy/tests/test_inference/test_Inference.py b/pgmpy/tests/test_inference/test_Inference.py
index 85644b45..ac162ffb 100644
--- a/pgmpy/tests/test_inference/test_Inference.py
+++ b/pgmpy/tests/test_inference/test_Inference.py
@@ -36,6 +36,7 @@ class TestInferenceBase(unittest.TestCase):
 
     def test_bayesian_inference_init(self):
         infer_bayesian = Inference(self.bayesian)
+        infer_bayesian._initialize_structures()
         self.assertEqual(set(infer_bayesian.variables), {"a", "b", "c", "d", "e"})
         self.assertEqual(
             infer_bayesian.cardinality, {"a": 2, "b": 2, "c": 2, "d": 2, "e": 2}
@@ -84,6 +85,7 @@ class TestInferenceBase(unittest.TestCase):
 
     def test_markov_inference_init(self):
         infer_markov = Inference(self.markov)
+        infer_markov._initialize_structures()
         self.assertEqual(set(infer_markov.variables), {"a", "b", "c", "d"})
         self.assertEqual(infer_markov.cardinality, {"a": 2, "b": 2, "c": 2, "d": 2})
         self.assertEqual(