Feature/dbn

causalnex/network/__init__.py

@@ -33,3 +33,4 @@
 __all__ = ["BayesianNetwork"]
 
 from .network import BayesianNetwork
+from .network import DynamicBayesianNetwork

causalnex/network/network.py

@@ -43,7 +43,7 @@
 from pgmpy.models import BayesianModel
 
 from causalnex.estimator.em import EMSingleLatentVariable
-from causalnex.structure import StructureModel
+from causalnex.structure import StructureModel, DynamicStructureModel
 from causalnex.utils.pgmpy_utils import pd_to_tabular_cpd
 
 
@@ -736,3 +736,84 @@ def _predict_probability_from_incomplete_data(
         probability = probability[cols]
         probability.columns = cols
         return probability
+
+
+class DynamicBayesianNetwork(BayesianNetwork):
+    """
+    Base class for Dynamic Bayesian Network (DBN), a probabilistic weighted DAG where nodes represent variables,
+    edges represent the causal relationships between variables.
+
+    ``DynamicBayesianNetwork`` stores nodes with their possible states, edges and
+    conditional probability distributions (CPDs) of each node.
+
+    ``DynamicBayesianNetwork`` is built on top of the ``StructureModel``, which is an extension of ``networkx.DiGraph``
+    (see :func:`causalnex.structure.structuremodel.StructureModel`).
+
+    In order to define the ``DynamicBayesianNetwork``, users should provide a relevant ``StructureModel``.
+    Once ``DynamicBayesianNetwork`` is initialised, no changes to the ``StructureModel`` can be made
+    and CPDs can be learned from the data.
+
+    The learned CPDs can be then used for likelihood estimation and predictions.
+
+    Example:
+    ::
+        >>> # Create a Dynamic Bayesian Network with a manually defined DAG.
+        >>> from causalnex.structure import StructureModel
+        >>> from causalnex.network import DynamicBayesianNetwork
+        >>>
+        >>> sm = StructureModel()
+        >>> sm.add_edges_from([
+        >>>                    ('rush_hour', 'traffic'),
+        >>>                    ('weather', 'traffic')
+        >>>                    ])
+        >>> dbn = DynamicBayesianNetwork(sm)
+        >>> # A created ``DynamicBayesianNetwork`` stores nodes and edges defined by the ``StructureModel``
+        >>> dbn.nodes
+        ['rush_hour', 'traffic', 'weather']
+        >>>
+        >>> dbn.edges
+        [('rush_hour', 'traffic'), ('weather', 'traffic')]
+        >>> # A ``DynamicBayesianNetwork`` doesn't store any CPDs yet
+        >>> dbn.cpds
+        >>> {}
+        >>>
+        >>> # Learn the nodes' states from the data
+        >>> import pandas as pd
+        >>> data = pd.DataFrame({
+        >>>                      'rush_hour': [True, False, False, False, True, False, True],
+        >>>                      'weather': ['Terrible', 'Good', 'Bad', 'Good', 'Bad', 'Bad', 'Good'],
+        >>>                      'traffic': ['heavy', 'light', 'heavy', 'light', 'heavy', 'heavy', 'heavy']
+        >>>                      })
+        >>> dbn = dbn.fit_node_states(data)
+        >>> dbn.node_states
+        {'rush_hour': {False, True}, 'weather': {'Bad', 'Good', 'Terrible'}, 'traffic': {'heavy', 'light'}}
+        >>> # Learn the CPDs from the data
+        >>> dbn = dbn.fit_cpds(data)
+        >>> # Use the learned CPDs to make predictions on the unseen data
+        >>> test_data = pd.DataFrame({
+        >>>                           'rush_hour': [False, False, True, True],
+        >>>                           'weather': ['Good', 'Bad', 'Good', 'Bad']
+        >>>                           })
+        >>> dbn.predict(test_data, "traffic").to_dict()
+        >>> {'traffic_prediction': {0: 'light', 1: 'heavy', 2: 'heavy', 3: 'heavy'}}
+        >>> dbn.predict_probability(test_data, "traffic").to_dict()
+        {'traffic_prediction': {0: 'light', 1: 'heavy', 2: 'heavy', 3: 'heavy'}}
+        {'traffic_light': {0: 0.75, 1: 0.25, 2: 0.3333333333333333, 3: 0.3333333333333333},
+         'traffic_heavy': {0: 0.25, 1: 0.75, 2: 0.6666666666666666, 3: 0.6666666666666666}}
+    """
+
+    def __init__(self, structure: DynamicStructureModel):
+        """
+        Create a ``DynamicBayesianNetwork`` with a DAG defined by ``DynamicStructureModel``.
+
+        Args:
+            structure: a graph representing a causal relationship between variables.
+                       In the structure
+                           - cycles are not allowed;
+                           - multiple (parallel) edges are not allowed;
+                           - isolated nodes and multiple components are not allowed.
+
+        Raises:
+            ValueError: If the structure is not a connected DAG.
+        """
+        super().__init__(structure)

causalnex/structure/__init__.py

@@ -35,9 +35,14 @@
     "notears",
     "dynotears",
     "data_generators",
+    "node",
     "DAGRegressor",
     "DAGClassifier",
+    "DynamicStructureModel",
+    "DynamicStructureNode"
 ]
 
 from .pytorch.sklearn import DAGClassifier, DAGRegressor
 from .structuremodel import StructureModel
+from .structuremodel import DynamicStructureModel
+from .structuremodel import DynamicStructureNode

causalnex/structure/dynotears.py

@@ -38,7 +38,8 @@
 import scipy.linalg as slin
 import scipy.optimize as sopt
 
-from causalnex.structure import StructureModel
+from causalnex.structure import DynamicStructureModel
+from causalnex.structure import DynamicStructureNode
 from causalnex.structure.transformers import DynamicDataTransformer
 
 
@@ -53,7 +54,7 @@ def from_pandas_dynamic(  # pylint: disable=too-many-arguments
     tabu_edges: List[Tuple[int, int, int]] = None,
     tabu_parent_nodes: List[int] = None,
     tabu_child_nodes: List[int] = None,
-) -> StructureModel:
+) -> DynamicStructureModel:
     """
     Learn the graph structure of a Dynamic Bayesian Network describing conditional dependencies between variables in
     data. The input data is a time series or a list of realisations of a same time series.
@@ -122,9 +123,9 @@ def from_pandas_dynamic(  # pylint: disable=too-many-arguments
         tabu_child_nodes,
     )
 
-    sm = StructureModel()
-    sm.add_nodes_from(
-        [f"{var}_lag{l_val}" for var in col_idx.keys() for l_val in range(p + 1)]
+    sm = DynamicStructureModel()
+    sm.add_nodes(
+        [DynamicStructureNode(var, l_val) for var in col_idx.keys() for l_val in range(p + 1)]
     )
     sm.add_weighted_edges_from(
         [
@@ -166,7 +167,7 @@ def from_numpy_dynamic(  # pylint: disable=too-many-arguments
     tabu_edges: List[Tuple[int, int, int]] = None,
     tabu_parent_nodes: List[int] = None,
     tabu_child_nodes: List[int] = None,
-) -> StructureModel:
+) -> DynamicStructureModel:
     """
     Learn the graph structure of a Dynamic Bayesian Network describing conditional dependencies between variables in
     data. The input data is time series data present in numpy arrays X and Xlags.
@@ -254,7 +255,7 @@ def from_numpy_dynamic(  # pylint: disable=too-many-arguments
 
 def _matrices_to_structure_model(
     w_est: np.ndarray, a_est: np.ndarray
-) -> StructureModel:
+) -> DynamicStructureModel:
     """
     Converts the matrices output by dynotears (W and A) into a StructureModel
     We use the following convention:
@@ -268,13 +269,13 @@ def _matrices_to_structure_model(
         StructureModel representing the structure learnt
 
     """
-    sm = StructureModel()
+    sm = DynamicStructureModel()
     lag_cols = [
-        f"{var}_lag{l_val}"
+        DynamicStructureNode(var, l_val)
         for l_val in range(1 + (a_est.shape[0] // a_est.shape[1]))
         for var in range(a_est.shape[1])
     ]
-    sm.add_nodes_from(lag_cols)
+    sm.add_nodes(lag_cols)
     sm.add_edges_from(
         [
             (lag_cols[i], lag_cols[j], dict(weight=w_est[i, j]))