Add causal inference: bn.inference.fit(..., do={...}) (closes #79)

bnlearn/inference.py

@@ -27,6 +27,7 @@ def fit(model,
         groupby=None,
         plot=False,
         verbose=3,
+        do=None,
         ):
     """Inference using using Variable Elimination.
 
@@ -47,6 +48,14 @@ def fit(model,
         For exact inference, P(variables | evidence). The default is None.
             * {'Rain':1}
             * {'Rain':1, 'Sprinkler':0, 'Cloudy':1}
+    do : dict, optional
+        Interventions for causal inference, P(variables | do(X=x), evidence).
+        Whereas evidence conditions on passively observed values, do simulates
+        setting the variable by intervention: incoming edges of the intervened
+        variables are cut (Pearl's do-operator, comparable to mutilated() in the
+        R version of bnlearn). The query runs on the mutilated network, so it
+        combines freely with evidence and the other query options. The default is None.
+            * {'Sprinkler':1}
     to_df : Bool, (default is True)
         The output is converted in the dataframe [query.df]. Enabling this function may impact the processing speed.
     elimination_order: str or list (default='greedy')
@@ -87,15 +96,25 @@ def fit(model,
     >>> print(query)
     >>> query.df
     >>>
+    >>> # Causal inference: P(Wet_Grass | do(Sprinkler=1)) differs from the
+    >>> # observational P(Wet_Grass | Sprinkler=1) because the intervention
+    >>> # cuts the Cloudy->Sprinkler edge.
+    >>> query = bn.inference.fit(model, variables=['Wet_Grass'], do={'Sprinkler':1})
+    >>> query.df
+    >>>
 
     """
     if not isinstance(model, dict): raise Exception('[bnlearn] >Error: Input requires a object that contains the key: model.')
     adjmat = model['adjmat']
     if not np.all(np.isin(variables, adjmat.columns)):
         raise Exception('[bnlearn] >Error: [variables] should match names in the model (Case sensitive!)')
-    if not np.all(np.isin([*evidence.keys()], adjmat.columns)):
+    if evidence is not None and not np.all(np.isin([*evidence.keys()], adjmat.columns)):
         raise Exception('[bnlearn] >Error: [evidence] should match names in the model (Case sensitive!)')
-    if verbose>=3: print('[bnlearn] >Variable Elimination.')
+    if do is not None and not np.all(np.isin([*do.keys()], adjmat.columns)):
+        raise Exception('[bnlearn] >Error: [do] should match names in the model (Case sensitive!)')
+    if do is not None and evidence is not None and (set(do.keys()) & set(evidence.keys())):
+        raise Exception('[bnlearn] >Error: A variable can not be in both [do] and [evidence]: %s' %(set(do.keys()) & set(evidence.keys())))
+    if verbose>=3: print('[bnlearn] >Causal inference with do-operator.' if do else '[bnlearn] >Variable Elimination.')
 
     # Extract model
     if isinstance(model, dict):
@@ -111,20 +130,35 @@ def fit(model,
         model = bnlearn.to_bayesiannetwork(adjmat, verbose=verbose)
 
     try:
+        if do:
+            # Query the mutilated network (incoming edges of the intervened nodes
+            # are cut) with the interventions fixed as evidence. This is exact for
+            # any mix of do and evidence, including interventions on causally
+            # related nodes, where pgmpy's adjustment-based CausalInference.query
+            # is not, and it keeps elimination_order/joint applicable.
+            model = model.do(list(do.keys()))
         model_infer = VariableElimination(model)
     except ValueError as e:
         raise Exception(f'[bnlearn] >Error: {e}')
         # Input model does not contain learned CPDs. hint: did you run parameter_learning.fit()?
 
-    # Computing the probability P(class | evidence)
-    query = model_infer.query(variables=variables, evidence=evidence, elimination_order=elimination_order, joint=joint, show_progress=(verbose>=3))
+    # Computing the probability P(class | do, evidence): in the mutilated network,
+    # fixing the intervened variables as evidence equals intervening on them.
+    query_evidence = {**do, **(evidence or {})} if do else evidence
+    query = model_infer.query(variables=variables, evidence=query_evidence, elimination_order=elimination_order, joint=joint, show_progress=(verbose>=3))
 
     # Store dataframe in query
+    if isinstance(query, dict):
+        # joint=False returns a dict of per-variable factors; there is no single
+        # joint table to attach a dataframe or summary to. (Attaching attributes
+        # to the dict raised AttributeError before, so this also unbreaks joint=False.)
+        return query
     if to_df or plot:
         # Convert to Dataframe
         query.df = bnlearn.query2df(query, variables=variables, groupby=groupby, verbose=verbose)
-        # Make readable text
-        query.text = summarize_inference(variables, evidence, query, plot=plot, verbose=verbose)
+        # Make readable text; label interventions as do(X) to keep them apart from observations
+        summary_given = {**{f'do({k})': v for k, v in (do or {}).items()}, **(evidence or {})}
+        query.text = summarize_inference(variables, summary_given, query, plot=plot, verbose=verbose)
         if verbose>=3 and query.text is not None: print(query.text)
     else:
         query.df = None

bnlearn/tests/test_inferences.py

@@ -86,6 +86,69 @@ def test_inference_sprinkler_example():
 
 
 
+def test_inference_do_operator():
+    # Ground truth computed by hand from the sprinkler CPDs:
+    # P(W=1|S=1) weights Cloudy by P(C|S=1) (observation), whereas
+    # P(W=1|do(S=1)) cuts Cloudy->Sprinkler and weights by the prior P(C).
+    model = bn.import_DAG('sprinkler', verbose=0)
+
+    def p1(query):
+        return float(query.df.loc[query.df['Wet_Grass'] == 1, 'p'].iloc[0])
+
+    observational = bn.inference.fit(model, variables=['Wet_Grass'], evidence={'Sprinkler': 1}, verbose=0)
+    assert p1(observational) == pytest.approx(0.927, abs=1e-3)
+
+    interventional = bn.inference.fit(model, variables=['Wet_Grass'], do={'Sprinkler': 1}, verbose=0)
+    assert p1(interventional) == pytest.approx(0.945, abs=1e-3)
+    assert abs(interventional.df['p'].sum() - 1.0) < 1e-6
+
+    # do and evidence combine: conditioning the intervened network on Cloudy=1
+    combined = bn.inference.fit(model, variables=['Wet_Grass'], do={'Sprinkler': 1}, evidence={'Cloudy': 1}, verbose=0)
+    assert p1(combined) == pytest.approx(0.972, abs=1e-3)
+
+    # Evidence on a variable outside the intervention's adjustment set: in the
+    # mutilated network Wet_Grass depends only on its parents, so this is
+    # exactly P(W=1|S=1,R=1)=0.99 from the CPT. (pgmpy's adjustment-based
+    # CausalInference.query returns 0.9612 here.)
+    outside = bn.inference.fit(model, variables=['Wet_Grass'], do={'Sprinkler': 1}, evidence={'Rain': 1}, verbose=0)
+    assert p1(outside) == pytest.approx(0.99, abs=1e-6)
+
+    # The other query options keep working with do
+    marginals = bn.inference.fit(model, variables=['Wet_Grass', 'Rain'], do={'Sprinkler': 1}, joint=False, to_df=False, verbose=0)
+    assert set(marginals.keys()) == {'Wet_Grass', 'Rain'}
+
+    # interventions are labeled as do(X) in the readable summary
+    assert 'do(Sprinkler)=1' in interventional.text
+
+
+def test_inference_do_operator_related_targets():
+    # Intervening on causally related nodes (A is a parent of B): the mutilated
+    # network holds both fixed, so P(C=1|do(A=1),do(B=1)) = P(C=1|A=1,B=1) = 1.
+    # (pgmpy's adjustment-based CausalInference.query returns 0.7 here.)
+    edges = [('A', 'B'), ('A', 'C'), ('B', 'C')]
+    cpt_a = TabularCPD(variable='A', variable_card=2, values=[[0.6], [0.4]])
+    cpt_b = TabularCPD(variable='B', variable_card=2,
+                       values=[[0.7, 0.2],
+                               [0.3, 0.8]],
+                       evidence=['A'], evidence_card=[2])
+    cpt_c = TabularCPD(variable='C', variable_card=2,
+                       values=[[0.5, 0.5, 0.5, 0.0],
+                               [0.5, 0.5, 0.5, 1.0]],
+                       evidence=['A', 'B'], evidence_card=[2, 2])
+    model = bn.make_DAG(edges, CPD=[cpt_a, cpt_b, cpt_c], verbose=0)
+
+    query = bn.inference.fit(model, variables=['C'], do={'A': 1, 'B': 1}, verbose=0)
+    assert float(query.df.loc[query.df['C'] == 1, 'p'].iloc[0]) == pytest.approx(1.0, abs=1e-6)
+
+
+def test_inference_do_operator_validation():
+    model = bn.import_DAG('sprinkler', verbose=0)
+    with pytest.raises(Exception, match='do'):
+        bn.inference.fit(model, variables=['Wet_Grass'], do={'NotANode': 1}, verbose=0)
+    with pytest.raises(Exception, match='both'):
+        bn.inference.fit(model, variables=['Wet_Grass'], do={'Rain': 1}, evidence={'Rain': 1}, verbose=0)
+
+
 def test_make_DAG_naivebayes():
     edges = [('A', 'B'), ('A', 'C'), ('A', 'D')]
     DAG = bn.make_DAG(edges, methodtype='naivebayes')