tigramite-fast 5.2.10.1__py3-none-any.whl

tigramite/causal_effects.py

@@ -0,0 +1,1525 @@
+"""Tigramite causal inference for time series."""
+
+# Author: Jakob Runge <jakob@jakob-runge.com>
+#
+# License: GNU General Public License v3.0
+
+import numpy as np
+import math
+import itertools
+from copy import deepcopy
+from collections import defaultdict
+from tigramite.models import Models
+from tigramite.graphs import Graphs
+import struct
+
+class CausalEffects(Graphs):
+    r"""Causal effect estimation.
+
+    Methods for the estimation of linear or non-parametric causal effects 
+    between (potentially multivariate) X and Y (potentially conditional 
+    on S) by (generalized) backdoor adjustment. Various graph types are 
+    supported, also including hidden variables.
+    
+    Linear and non-parametric estimators are based on sklearn. For the 
+    linear case without hidden variables also an efficient estimation 
+    based on Wright's path coefficients is available. This estimator 
+    also allows to estimate mediation effects.
+
+    See the corresponding paper [6]_ and tigramite tutorial for an 
+    in-depth introduction. 
+
+    References
+    ----------
+
+    .. [6] J. Runge, Necessary and sufficient graphical conditions for
+           optimal adjustment sets in causal graphical models with 
+           hidden variables, Advances in Neural Information Processing
+           Systems, 2021, 34 
+           https://proceedings.neurips.cc/paper/2021/hash/8485ae387a981d783f8764e508151cd9-Abstract.html
+
+
+    Parameters
+    ----------
+    graph : array of either shape [N, N], [N, N, tau_max+1], or [N, N, tau_max+1, tau_max+1]
+        Different graph types are supported, see tutorial.
+    X : list of tuples
+        List of tuples [(i, -tau), ...] containing cause variables.
+    Y : list of tuples
+        List of tuples [(j, 0), ...] containing effect variables.
+    S : list of tuples
+        List of tuples [(i, -tau), ...] containing conditioned variables.  
+    graph_type : str
+        Type of graph.
+    hidden_variables : list of tuples
+        Hidden variables in format [(i, -tau), ...]. The internal graph is 
+        constructed by a latent projection.
+    check_SM_overlap : bool
+        Whether to check whether S overlaps with M.
+    verbosity : int, optional (default: 0)
+        Level of verbosity.
+    """
+
+    def __init__(self,
+                 graph,
+                 graph_type,
+                 X,
+                 Y,
+                 S=None,
+                 hidden_variables=None,
+                 check_SM_overlap=True,
+                 verbosity=0):
+        
+        self.verbosity = verbosity
+        self.N = graph.shape[0]
+
+        if S is None:
+            S = []
+
+        self.listX = list(X)
+        self.listY = list(Y)
+        self.listS = list(S)
+
+        self.X = set(X)
+        self.Y = set(Y)
+        self.S = set(S)    
+
+        # TODO?: check that masking aligns with hidden samples in variables
+        if hidden_variables is None:
+            hidden_variables = []
+
+        # 
+        # Checks regarding graph type
+        #
+        supported_graphs = ['dag', 
+                            'admg',
+                            'tsg_dag',
+                            'tsg_admg',
+                            'stationary_dag',
+                            'stationary_admg',
+
+                            # 'mag',
+                            # 'tsg_mag',
+                            # 'stationary_mag',
+                            # 'pag',
+                            # 'tsg_pag',
+                            # 'stationary_pag',
+                            ]
+        if graph_type not in supported_graphs:
+            raise ValueError("Only graph types %s supported!" %supported_graphs)
+
+        # Determine tau_max
+        if graph_type in ['dag', 'admg']: 
+            self.tau_max = 0
+
+        elif graph_type in ['tsg_dag', 'tsg_admg']:
+            # tau_max is implicitely derived from
+            # the dimensions 
+            self.tau_max = graph.shape[2] - 1
+
+        elif graph_type in ['stationary_dag', 'stationary_admg']:
+            # For a stationary DAG without hidden variables it's sufficient to consider
+            # a tau_max that includes the parents of X, Y, M, and S. A conservative
+            # estimate thereof is simply the lag-dimension of the stationary DAG plus
+            # the maximum lag of XYS.
+            statgraph_tau_max = graph.shape[2] - 1
+            maxlag_XYS = 0
+            for varlag in self.X.union(self.Y).union(self.S):
+                maxlag_XYS = max(maxlag_XYS, abs(varlag[1]))
+            self.tau_max = maxlag_XYS + statgraph_tau_max
+        else:
+            raise ValueError("graph_type invalid.")
+
+        self.hidden_variables = set(hidden_variables)
+        if len(self.hidden_variables.intersection(self.X.union(self.Y).union(self.S))) > 0:
+            raise ValueError("XYS overlaps with hidden_variables!")
+
+        # self.tau_max is needed in the Graphs class
+        Graphs.__init__(self, 
+                        graph=graph,
+                        graph_type=graph_type,
+                        tau_max=self.tau_max,
+                        hidden_variables=self.hidden_variables,
+                        verbosity=verbosity)
+
+        self._check_XYS()
+
+        self.ancX = self._get_ancestors(X)
+        self.ancY = self._get_ancestors(Y)
+        self.ancS = self._get_ancestors(S)
+
+        # If X is not in anc(Y), then no causal link exists
+        if self.ancY.intersection(set(X)) == set():
+            self.no_causal_path = True
+            if self.verbosity > 0:
+                print("No causal path from X to Y exists.")
+        else:
+            self.no_causal_path = False
+
+        # Get mediators
+        mediators = self.get_mediators(start=self.X, end=self.Y) 
+
+        M = set(mediators)
+        self.M = M
+
+        self.listM = list(self.M)
+
+        for varlag in self.X.union(self.Y).union(self.S):
+            if abs(varlag[1]) > self.tau_max:
+                raise ValueError("X, Y, S must have time lags inside graph.")
+
+        if len(self.X.intersection(self.Y)) > 0:
+            raise ValueError("Overlap between X and Y.")
+
+        if len(self.S.intersection(self.Y.union(self.X))) > 0:
+            raise ValueError("Conditions S overlap with X or Y.")
+
+        # # TODO: need to prove that this is sufficient for non-identifiability!
+        # if len(self.X.intersection(self._get_descendants(self.M))) > 0:
+        #     raise ValueError("Not identifiable: Overlap between X and des(M)")
+
+        if check_SM_overlap and len(self.S.intersection(self.M)) > 0:
+            raise ValueError("Conditions S overlap with mediators M.")
+
+        self.desX = self._get_descendants(self.X)
+        self.desY = self._get_descendants(self.Y)
+        self.desM = self._get_descendants(self.M)
+        self.descendants = self.desY.union(self.desM)
+
+        # Define forb as X and descendants of YM
+        self.forbidden_nodes = self.descendants.union(self.X)  #.union(S)
+
+        # Define valid ancestors
+        self.vancs = self.ancX.union(self.ancY).union(self.ancS) - self.forbidden_nodes
+
+        if self.verbosity > 0:
+            if len(self.S.intersection(self.desX)) > 0:
+                print("Warning: Potentially outside assumptions: Conditions S overlap with des(X)")
+
+        # Here only check if S overlaps with des(Y), leave the option that S
+        # contains variables in des(M) to the user
+        if len(self.S.intersection(self.desY)) > 0:
+            raise ValueError("Not identifiable: Conditions S overlap with des(Y).")
+
+        if self.verbosity > 0:
+            print("\n##\n## Initializing CausalEffects class\n##"
+                  "\n\nInput:")
+            print("\ngraph_type = %s" % graph_type
+                  + "\nX = %s" % self.listX
+                  + "\nY = %s" % self.listY
+                  + "\nS = %s" % self.listS
+                  + "\nM = %s" % self.listM
+                  )
+            if len(self.hidden_variables) > 0:
+                print("\nhidden_variables = %s" % self.hidden_variables
+                      ) 
+            print("\n\n")
+            if self.no_causal_path:
+                print("No causal path from X to Y exists!")
+
+
+    def _check_XYS(self):
+        """Check whether XYS are sober.
+        """
+
+        XYS = self.X.union(self.Y).union(self.S)
+        for xys in XYS:
+            var, lag = xys 
+            if var < 0 or var >= self.N:
+                raise ValueError("XYS vars must be in [0...N]")
+            if lag < -self.tau_max or lag > 0:
+                raise ValueError("XYS lags must be in [-taumax...0]")
+
+
+    def check_XYS_paths(self):
+        """Check whether one can remove nodes from X and Y with no proper causal paths.
+
+        Returns
+        -------
+        X, Y : cleaned lists of X and Y with irrelevant nodes removed.
+        """
+
+        # TODO: Also check S...
+        oldX = self.X.copy()
+        oldY = self.Y.copy()
+
+        # anc_Y = self._get_ancestors(self.Y)
+        # anc_S = self._get_ancestors(self.S)
+
+        # Remove first from X those nodes with no causal path to Y or S
+        X = set([x for x in self.X if x in self.ancY.union(self.ancS)])
+        
+        # Remove from Y those nodes with no causal path from X
+        # des_X = self._get_descendants(X)
+
+        Y = set([y for y in self.Y if y in self.desX])
+
+        # Also require that all x in X have proper path to Y or S,
+        # that is, the first link goes out of x 
+        # and into path nodes
+        mediators_S = self.get_mediators(start=self.X, end=self.S)
+        path_nodes = list(self.M.union(Y).union(mediators_S)) 
+        X = X.intersection(self._get_all_parents(path_nodes))
+
+        if set(oldX) != set(X) and self.verbosity > 0:
+            print("Consider pruning X = %s to X = %s " %(oldX, X) +
+                  "since only these have causal path to Y")
+
+        if set(oldY) != set(Y) and self.verbosity > 0:
+            print("Consider pruning Y = %s to Y = %s " %(oldY, Y) +
+                  "since only these have causal path from X")
+
+        return (list(X), list(Y))
+
+
+    def get_optimal_set(self, 
+        alternative_conditions=None,
+        minimize=False,
+        return_separate_sets=False,
+        ):
+        """Returns optimal adjustment set.
+        
+        See Runge NeurIPS 2021.
+
+        Parameters
+        ----------
+        alternative_conditions : set of tuples
+            Used only internally in optimality theorem. If None, self.S is used.
+        minimize : {False, True, 'colliders_only'} 
+            Minimize optimal set. If True, minimize such that no subset 
+            can be removed without making it invalid. If 'colliders_only',
+            only colliders are minimized.
+        return_separate_sets : bool
+            Whether to return tuple of parents, colliders, collider_parents, and S.
+        
+        Returns
+        -------
+        Oset_S : False or list or tuple of lists
+            Returns optimal adjustment set if a valid set exists, otherwise False.
+        """
+
+
+        # Needed for optimality theorem where Osets for alternative S are tested
+        if alternative_conditions is None:
+            S = self.S.copy()
+            vancs = self.vancs.copy()
+        else:
+            S = alternative_conditions
+            newancS = self._get_ancestors(S)
+            vancs = self.ancX.union(self.ancY).union(newancS) - self.forbidden_nodes
+
+            # vancs = self._get_ancestors(list(self.X.union(self.Y).union(S))) - self.forbidden_nodes
+
+        # descendants = self._get_descendants(self.Y.union(self.M))
+
+        # Sufficient condition for non-identifiability
+        if len(self.X.intersection(self.descendants)) > 0:
+            return False  # raise ValueError("Not identifiable: Overlap between X and des(M)")
+
+        ##
+        ## Construct O-set
+        ##
+
+        # Start with parents 
+        parents = self._get_all_parents(self.Y.union(self.M)) # set()
+
+        # Remove forbidden nodes
+        parents = parents - self.forbidden_nodes
+
+        # Construct valid collider path nodes
+        colliders = set()
+        for w in self.Y.union(self.M):
+            j, tau = w 
+            this_level = [w]
+            non_suitable_nodes = []
+            while len(this_level) > 0:
+                next_level = []
+                for varlag in this_level:
+                    suitable_spouses = set(self._get_spouses(varlag)) - set(non_suitable_nodes)
+                    for spouse in suitable_spouses:
+                        i, tau = spouse
+                        if spouse in self.X:
+                            return False
+
+                        if (# Node not already in set
+                            spouse not in colliders  #.union(parents)
+                            # not forbidden
+                            and spouse not in self.forbidden_nodes 
+                            # in time bounds
+                            and (-self.tau_max <= tau <= 0) # or self.ignore_time_bounds)
+                            and (spouse in vancs
+                                or not self._check_path(#graph=self.graph, 
+                                    start=self.X, end=[spouse], 
+                                                    conditions=list(parents.union(vancs)) + list(S),
+                                                    ))
+                                ):
+                                colliders = colliders | {spouse}
+                                next_level.append(spouse)
+                        else:
+                            if spouse not in colliders:
+                                non_suitable_nodes.append(spouse)
+
+
+                this_level = set(next_level) - set(non_suitable_nodes)  
+
+        # Add parents and raise Error if not identifiable
+        collider_parents = self._get_all_parents(colliders)
+        if len(self.X.intersection(collider_parents)) > 0:
+            return False
+
+        colliders_and_their_parents = colliders.union(collider_parents)
+
+        # Add valid collider path nodes and their parents
+        Oset = parents.union(colliders_and_their_parents)
+
+
+        if minimize: 
+            removable = []
+            # First remove all those that have no path from X
+            sorted_Oset =  Oset
+            if minimize == 'colliders_only':
+                sorted_Oset = [node for node in sorted_Oset if node not in parents]
+
+            for node in sorted_Oset:
+                if (not self._check_path(#graph=self.graph, 
+                    start=self.X, end=[node], 
+                                conditions=list(Oset - {node}) + list(S))):
+                    removable.append(node) 
+
+            Oset = Oset - set(removable)
+            if minimize == 'colliders_only':
+                sorted_Oset = [node for node in Oset if node not in parents]
+
+            removable = []
+            # Next remove all those with no direct connection to Y
+            for node in sorted_Oset:
+                if (not self._check_path(#graph=self.graph, 
+                    start=[node], end=self.Y, 
+                            conditions=list(Oset - {node}) + list(S) + list(self.X),
+                            ends_with=['**>', '**+'])): 
+                    removable.append(node) 
+
+            Oset = Oset - set(removable)
+
+        Oset_S = Oset.union(S)
+
+        if return_separate_sets:
+            return parents, colliders, collider_parents, S
+        else:
+            return list(Oset_S)
+
+
+    def _get_collider_paths_optimality(self, source_nodes, target_nodes,
+        condition, 
+        inside_set=None, 
+        start_with_tail_or_head=False, 
+        ):
+        """Returns relevant collider paths to check optimality.
+
+        Iterates over collider paths within O-set via depth-first search
+
+        """
+
+        for w in source_nodes:
+            # Only used to return *all* collider paths 
+            # (needed in optimality theorem)
+            
+            coll_path = []
+
+            queue = [(w, coll_path)]
+
+            non_valid_subsets = []
+
+            while queue:
+
+                varlag, coll_path = queue.pop()
+
+                coll_path = coll_path + [varlag]
+                coll_path_set = set(coll_path)
+
+                suitable_nodes = set(self._get_spouses(varlag))
+
+                if start_with_tail_or_head and coll_path == [w]:
+                    children = set(self._get_children(varlag))
+                    suitable_nodes = suitable_nodes.union(children)
+
+                for node in suitable_nodes:
+                    i, tau = node
+                    if ((-self.tau_max <= tau <= 0) # or self.ignore_time_bounds)
+                        and node not in coll_path_set):
+
+                        if condition == 'II' and node not in target_nodes and node not in self.vancs:
+                            continue
+
+                        if node in inside_set:
+                            if condition == 'I':
+                                non_valid = False
+                                extended_set = coll_path_set | {node}
+                                for pathset in non_valid_subsets[::-1]:
+                                    if pathset.issubset(extended_set):
+                                        non_valid = True
+                                        break
+                                if non_valid is False:
+                                    queue.append((node, coll_path))
+                                else:
+                                    continue
+                            elif condition == 'II':
+                                queue.append((node, coll_path))
+
+                        if node in target_nodes:
+                            # yield coll_path
+                            # collider_paths[node].append(coll_path)
+                            if condition == 'I':
+                                # Construct OπiN
+                                Sprime = self.S.union(coll_path_set)
+                                OpiN = self.get_optimal_set(alternative_conditions=Sprime)
+                                if OpiN is False:
+                                    queue = [(q_node, q_path) for (q_node, q_path) in queue if not coll_path_set.issubset(set(q_path + [q_node]))]
+                                    non_valid_subsets.append(coll_path_set)
+                                else:
+                                    return False
+
+                            elif condition == 'II':
+                                return True
+                                # yield coll_path
+ 
+        if condition == 'I':
+            return True
+        elif condition == 'II':
+            return False
+        # return collider_paths
+
+
+    def check_optimality(self):
+        """Check whether optimal adjustment set exists according to Thm. 3 in Runge NeurIPS 2021.
+
+        Returns
+        -------
+        optimality : bool
+            Returns True if an optimal adjustment set exists, otherwise False.
+        """
+
+        # Cond. 0: Exactly one valid adjustment set exists
+        cond_0 = (self._get_all_valid_adjustment_sets(check_one_set_exists=True))
+
+        #
+        # Cond. I
+        #
+        parents, colliders, collider_parents, _ = self.get_optimal_set(return_separate_sets=True)
+        Oset = parents.union(colliders).union(collider_parents)
+        n_nodes = self._get_all_spouses(self.Y.union(self.M).union(colliders)) - self.forbidden_nodes - Oset - self.S - self.Y - self.M - colliders
+
+        if (len(n_nodes) == 0):
+            # # (1) There are no spouses N ∈ sp(YMC) \ (forbOS)
+            cond_I = True
+        else:
+            
+            # (2) For all N ∈ N and all its collider paths i it holds that 
+            # OπiN does not block all non-causal paths from X to Y
+            # cond_I = True
+            cond_I = self._get_collider_paths_optimality(
+                source_nodes=list(n_nodes), target_nodes=list(self.Y.union(self.M)),
+                condition='I', 
+                inside_set=Oset.union(self.S), start_with_tail_or_head=False,
+                )
+
+        #
+        # Cond. II
+        #
+        e_nodes = Oset.difference(parents)
+        cond_II = True
+        for E in e_nodes:
+            Oset_minusE = Oset - {E}
+            if self._check_path(#graph=self.graph, 
+                start=list(self.X), end=[E], 
+                                conditions=list(self.S) + list(Oset_minusE)):
+                   
+                cond_II = self._get_collider_paths_optimality(
+                    target_nodes=self.Y.union(self.M), 
+                    source_nodes=[E],
+                    condition='II', 
+                    inside_set=list(Oset.union(self.S)),
+                    start_with_tail_or_head = True)
+               
+                if cond_II is False:
+                    if self.verbosity > 1:
+                        print("Non-optimal due to E = ", E)
+                    break
+
+        optimality = (cond_0 or (cond_I and cond_II))
+        if self.verbosity > 0:
+            print("Optimality = %s with cond_0 = %s, cond_I = %s, cond_II = %s"
+                    %  (optimality, cond_0, cond_I, cond_II))
+        return optimality
+
+    def _check_validity(self, Z):
+        """Checks whether Z is a valid adjustment set."""
+
+        # causal_children = list(self.M.union(self.Y))
+        backdoor_path = self._check_path(#graph=self.graph, 
+            start=list(self.X), end=list(self.Y), 
+                            conditions=list(Z), 
+                            # causal_children=causal_children,
+                            path_type = 'non_causal')
+
+        if backdoor_path:
+            return False
+        else:
+            return True
+    
+    def _get_adjust_set(self, 
+        minimize=False,
+        ):
+        """Returns Adjust-set.
+        
+        See van der Zander, B.; Liśkiewicz, M. & Textor, J.
+        Separators and adjustment sets in causal graphs: Complete 
+        criteria and an algorithmic framework 
+        Artificial Intelligence, Elsevier, 2019, 270, 1-40
+
+        """
+
+        vancs = self.vancs.copy()
+
+        if minimize:
+            # Get removable nodes by computing minimal valid set from Z
+            if minimize == 'keep_parentsYM':
+                minimize_nodes = vancs - self._get_all_parents(list(self.Y.union(self.M)))
+
+            else:
+                minimize_nodes = vancs
+
+            # Zprime2 = Zprime
+            # First remove all nodes that have no unique path to X given Oset
+            for node in minimize_nodes:
+                # path = self.oracle.check_shortest_path(X=X, Y=[node], 
+                #     Z=list(vancs - {node}), 
+                #     max_lag=None, 
+                #     starts_with=None, #'arrowhead', 
+                #     forbidden_nodes=None, #list(Zprime - {node}), 
+                #     return_path=False)
+                path = self._check_path(#graph=self.graph, 
+                    start=self.X, end=[node], 
+                    conditions=list(vancs - {node}), 
+                     )
+  
+                if path is False:
+                    vancs = vancs - {node}
+
+            if minimize == 'keep_parentsYM':
+                minimize_nodes = vancs - self._get_all_parents(list(self.Y.union(self.M)))
+            else:
+                minimize_nodes = vancs
+
+            # print(Zprime2) 
+            # Next remove all nodes that have no unique path to Y given Oset_min
+            # Z = Zprime2
+            for node in minimize_nodes:
+
+                path = self._check_path(#graph=self.graph, 
+                    start=[node], end=self.Y, 
+                    conditions=list(vancs - {node}) + list(self.X),
+                    )
+
+                if path is False:
+                   vancs = vancs - {node}  
+
+        if self._check_validity(list(vancs)) is False:
+            return False
+        else:
+            return list(vancs)
+
+
+    def _get_all_valid_adjustment_sets(self, 
+        check_one_set_exists=False, yield_index=None):
+        """Constructs all valid adjustment sets or just checks whether one exists.
+        
+        See van der Zander, B.; Liśkiewicz, M. & Textor, J.
+        Separators and adjustment sets in causal graphs: Complete 
+        criteria and an algorithmic framework 
+        Artificial Intelligence, Elsevier, 2019, 270, 1-40
+
+        """
+
+        cond_set = set(self.S)
+        all_vars = [(i, -tau) for i in range(self.N)
+                    for tau in range(0, self.tau_max + 1)]
+
+        all_vars_set = set(all_vars) - self.forbidden_nodes
+
+
+        def find_sep(I, R):
+            Rprime = R - self.X - self.Y
+            # TODO: anteriors and NOT ancestors where
+            # anteriors include --- links in causal paths
+            # print(I)
+            XYI = list(self.X.union(self.Y).union(I))
+            # print(XYI)
+            ancs = self._get_ancestors(list(XYI))
+            Z = ancs.intersection(Rprime)
+            if self._check_validity(Z) is False:
+                return False
+            else:
+                return Z
+
+
+        def list_sep(I, R):
+            # print(find_sep(X, Y, I, R))
+            if find_sep(I, R) is not False:
+                # print(I,R)
+                if I == R: 
+                    # print('--->', I)
+                    yield I
+                else:
+                    # Pick arbitrary node from R-I
+                    RminusI = list(R - I)
+                    # print(R, I, RminusI)
+                    v = RminusI[0]
+                    # print("here ", X, Y, I.union(set([v])), R)
+                    yield from list_sep(I | {v}, R)
+                    yield from list_sep(I, R - {v})
+
+        # print("all ", X, Y, cond_set, all_vars_set)
+        all_sets = []
+        I = cond_set
+        R = all_vars_set
+        for index, valid_set in enumerate(list_sep(I, R)):
+            # print(valid_set)
+            all_sets.append(list(valid_set))
+            if check_one_set_exists and index > 0:
+                break
+
+            if yield_index is not None and index == yield_index:
+                return valid_set
+
+        if yield_index is not None:
+            return None
+
+        if check_one_set_exists:
+            if len(all_sets) == 1:
+                return True
+            else:
+                return False
+
+        return all_sets
+
+
+    def fit_total_effect(self,
+        dataframe, 
+        estimator,
+        adjustment_set='optimal',
+        conditional_estimator=None,  
+        data_transform=None,
+        mask_type=None,
+        ignore_identifiability=False,
+        ):
+        """Returns a fitted model for the total causal effect of X on Y 
+           conditional on S.
+
+        Parameters
+        ----------
+        dataframe : data object
+            Tigramite dataframe object. It must have the attributes dataframe.values
+            yielding a numpy array of shape (observations T, variables N) and
+            optionally a mask of the same shape and a missing values flag.
+        estimator : sklearn model object
+            For example, sklearn.linear_model.LinearRegression() for a linear
+            regression model.
+        adjustment_set : str or list of tuples
+            If 'optimal' the Oset is used, if 'minimized_optimal' the minimized Oset,
+            and if 'colliders_minimized_optimal', the colliders-minimized Oset.
+            If a list of tuples is passed, this set is used.
+        conditional_estimator : sklearn model object, optional (default: None)
+            Used to fit conditional causal effects in nested regression. 
+            If None, the same model as for estimator is used.
+        data_transform : sklearn preprocessing object, optional (default: None)
+            Used to transform data prior to fitting. For example,
+            sklearn.preprocessing.StandardScaler for simple standardization. The
+            fitted parameters are stored.
+        mask_type : {None, 'y','x','z','xy','xz','yz','xyz'}
+            Masking mode: Indicators for which variables in the dependence
+            measure I(X; Y | Z) the samples should be masked. If None, the mask
+            is not used. Explained in tutorial on masking and missing values.
+        ignore_identifiability : bool
+            Only applies to adjustment sets supplied by user. Ignores if that 
+            set leads to a non-identifiable effect.
+        """
+
+        if self.no_causal_path:
+            if self.verbosity > 0:
+                print("No causal path from X to Y exists.")
+            return self
+
+        self.dataframe = dataframe
+        self.conditional_estimator = conditional_estimator
+
+        # if self.dataframe.has_vector_data:
+        #     raise ValueError("vector_vars in DataFrame cannot be used together with CausalEffects!"
+        #                      " You can estimate vector-valued effects by using multivariate X, Y, S."
+        #                      " Note, however, that this requires assuming a graph at the level "
+        #                      "of the components of X, Y, S, ...")
+
+        if self.N != self.dataframe.N:
+            raise ValueError("Dataset dimensions inconsistent with number of variables in graph.")
+
+        if adjustment_set == 'optimal':
+            # Check optimality and use either optimal or colliders_only set
+            adjustment_set = self.get_optimal_set()
+        elif adjustment_set == 'colliders_minimized_optimal':
+            adjustment_set = self.get_optimal_set(minimize='colliders_only')
+        elif adjustment_set == 'minimized_optimal':
+            adjustment_set = self.get_optimal_set(minimize=True)
+        else:
+            if ignore_identifiability is False and self._check_validity(adjustment_set) is False:
+                raise ValueError("Chosen adjustment_set is not valid.")
+
+        if adjustment_set is False:
+            raise ValueError("Causal effect not identifiable via adjustment.")
+
+        self.adjustment_set = adjustment_set
+
+        # Fit model of Y on X and Z (and conditions)
+        # Build the model
+        self.model = Models(
+                        dataframe=dataframe,
+                        model=estimator,
+                        conditional_model=conditional_estimator,
+                        data_transform=data_transform,
+                        mask_type=mask_type,
+                        verbosity=self.verbosity)      
+
+        self.model.get_general_fitted_model(
+                Y=self.listY, X=self.listX, Z=list(self.adjustment_set),
+                conditions=self.listS,
+                tau_max=self.tau_max,
+                cut_off='tau_max',
+                return_data=False)
+
+        return self
+
+    def predict_total_effect(self, 
+        intervention_data, 
+        conditions_data=None,
+        pred_params=None,
+        return_further_pred_results=False,
+        aggregation_func=np.mean,
+        transform_interventions_and_prediction=False,
+        intervention_type='hard',
+        ):
+        """Predict effect of intervention with fitted model.
+
+        Uses the model.predict() function of the sklearn model.
+
+        Parameters
+        ----------
+        intervention_data : numpy array
+            Numpy array of shape (n_interventions, len(X)) that contains the do(X) values.
+        conditions_data : data object, optional
+            Numpy array of shape (n_interventions, len(S)) that contains the S=s values.
+        pred_params : dict, optional
+            Optional parameters passed on to sklearn prediction function.
+        return_further_pred_results : bool, optional (default: False)
+            In case the predictor class returns more than just the expected value,
+            the entire results can be returned.
+        aggregation_func : callable
+            Callable applied to output of 'predict'. Default is 'np.mean'.
+        transform_interventions_and_prediction : bool (default: False)
+            Whether to perform the inverse data_transform on prediction results.
+        intervention_type : {'hard', 'soft'}
+            Specify whether intervention is 'hard' (set value) or 'soft' 
+            (add value to observed data).
+        
+        Returns
+        -------
+        Results from prediction: an array of shape  (n_interventions, len(Y)).
+        If estimate_confidence = True, then a tuple is returned.
+        """
+
+        def get_vectorized_length(W):
+            return sum([len(self.dataframe.vector_vars[w[0]]) for w in W])
+
+        # lenX = len(self.listX)
+        # lenS = len(self.listS)
+
+        lenX = get_vectorized_length(self.listX)
+        lenS = get_vectorized_length(self.listS)
+
+        if intervention_data.shape[1] != lenX:
+            raise ValueError("intervention_data.shape[1] must be len(X).")
+
+        if intervention_type not in {'hard', 'soft'}:
+            raise ValueError("intervention_type must be 'hard' or 'soft'.")
+
+        if conditions_data is not None and lenS > 0:
+            if conditions_data.shape[1] != lenS:
+                raise ValueError("conditions_data.shape[1] must be len(S).")
+            if conditions_data.shape[0] != intervention_data.shape[0]:
+                raise ValueError("conditions_data.shape[0] must match intervention_data.shape[0].")
+        elif conditions_data is not None and lenS == 0:
+            raise ValueError("conditions_data specified, but S=None or empty.")
+        elif conditions_data is None and lenS > 0:
+            raise ValueError("S specified, but conditions_data is None.")
+
+
+        if self.no_causal_path:
+            if self.verbosity > 0:
+                print("No causal path from X to Y exists.")
+            return np.zeros((len(intervention_data), len(self.listY)))
+
+        effect = self.model.get_general_prediction(
+            intervention_data=intervention_data,
+            conditions_data=conditions_data,
+            pred_params=pred_params,
+            return_further_pred_results=return_further_pred_results,
+            transform_interventions_and_prediction=transform_interventions_and_prediction,
+            aggregation_func=aggregation_func,
+            intervention_type=intervention_type,) 
+
+        return effect
+
+    def fit_wright_effect(self,
+        dataframe, 
+        mediation=None,
+        method='parents',
+        links_coeffs=None,  
+        data_transform=None,
+        mask_type=None,
+        ):
+        """Returns a fitted model for the total or mediated causal effect of X on Y 
+           potentially through mediator variables.
+
+        Parameters
+        ----------
+        dataframe : data object
+            Tigramite dataframe object. It must have the attributes dataframe.values
+            yielding a numpy array of shape (observations T, variables N) and
+            optionally a mask of the same shape and a missing values flag.
+        mediation : None, 'direct', or list of tuples
+            If None, total effect is estimated, if 'direct' then only the direct effect is estimated,
+            else only those causal paths are considerd that pass at least through one of these mediator nodes.
+        method : {'parents', 'links_coeffs', 'optimal'}
+            Method to use for estimating Wright's path coefficients. If 'optimal', 
+            the Oset is used, if 'links_coeffs', the coefficients in links_coeffs are used,
+            if 'parents', the parents are used (only valid for DAGs).
+        links_coeffs : dict
+            Only used if method = 'links_coeffs'.
+            Dictionary of format: {0:[((i, -tau), coeff),...], 1:[...],
+            ...} for all variables where i must be in [0..N-1] and tau >= 0 with
+            number of variables N. coeff must be a float.
+        data_transform : None
+            Not implemented for Wright estimator. Complicated for missing samples.
+        mask_type : {None, 'y','x','z','xy','xz','yz','xyz'}
+            Masking mode: Indicators for which variables in the dependence
+            measure I(X; Y | Z) the samples should be masked. If None, the mask
+            is not used. Explained in tutorial on masking and missing values.
+        """
+
+        if self.no_causal_path:
+            if self.verbosity > 0:
+                print("No causal path from X to Y exists.")
+            return self
+
+        if data_transform is not None:
+            raise ValueError("data_transform not implemented for Wright estimator."
+                             " You can preprocess data yourself beforehand.")
+
+        import sklearn.linear_model
+
+        self.dataframe = dataframe
+        if self.dataframe.has_vector_data:
+            raise ValueError("vector_vars in DataFrame cannot be used together with Wright method!"
+                             " You can either 1) estimate vector-valued effects by using multivariate (X, Y, S)"
+                             " together with assuming a graph at the level of the components of (X, Y, S), "
+                             " or 2) use vector_vars together with fit_total_effect and an estimator"
+                             " that supports multiple outputs.")
+
+        estimator = sklearn.linear_model.LinearRegression()
+
+        # Fit model of Y on X and Z (and conditions)
+        # Build the model
+        self.model = Models(
+                        dataframe=dataframe,
+                        model=estimator,
+                        data_transform=None, #data_transform,
+                        mask_type=mask_type,
+                        verbosity=self.verbosity)
+
+        mediators = self.M  # self.get_mediators(start=self.X, end=self.Y)
+
+        if mediation == 'direct':
+            causal_paths = {}         
+            for w in self.X:
+                causal_paths[w] = {}
+                for z in self.Y:
+                    if w in self._get_parents(z):
+                        causal_paths[w][z] = [[w, z]]
+                    else:
+                        causal_paths[w][z] = []
+        else:
+            causal_paths = self._get_causal_paths(source_nodes=self.X, 
+                target_nodes=self.Y, mediators=mediators, 
+                mediated_through=mediation, proper_paths=True)
+
+        if method == 'links_coeffs':
+            coeffs = {}
+            max_lag = 0
+            for medy in [med for med in mediators] + [y for y in self.listY]:
+                coeffs[medy] = {}
+                j, tauj = medy
+                for ipar, par_coeff in enumerate(links_coeffs[medy[0]]):
+                    par, coeff, _ = par_coeff
+                    i, taui = par
+                    taui_shifted = taui + tauj
+                    max_lag = max(abs(par[1]), max_lag)
+                    coeffs[medy][(i, taui_shifted)] = coeff #self.fit_results[j][(j, 0)]['model'].coef_[ipar]
+
+            self.model.tau_max = max_lag
+            # print(coeffs)
+
+        elif method == 'optimal':
+            # all_parents = {}
+            coeffs = {}
+            for medy in [med for med in mediators] + [y for y in self.listY]:
+                coeffs[medy] = {}
+                mediator_parents = self._get_all_parents([medy]).intersection(mediators.union(self.X).union(self.Y)) - {medy}
+                all_parents = self._get_all_parents([medy]) - {medy}
+                for par in mediator_parents:
+                    Sprime = set(all_parents) - {par, medy}
+                    causal_effects = CausalEffects(graph=self.graph, 
+                                        X=[par], Y=[medy], S=Sprime,
+                                        graph_type=self.graph_type,
+                                        check_SM_overlap=False,
+                                        )
+                    oset = causal_effects.get_optimal_set()
+                    # print(medy, par, list(set(all_parents)), oset)
+                    if oset is False:
+                        raise ValueError("Not identifiable via Wright's method.")
+                    fit_res = self.model.get_general_fitted_model(
+                        Y=[medy], X=[par], Z=oset,
+                        tau_max=self.tau_max,
+                        cut_off='tau_max',
+                        return_data=False)
+                    coeffs[medy][par] = fit_res['model'].coef_[0]
+
+        elif method == 'parents':
+            coeffs = {}
+            for medy in [med for med in mediators] + [y for y in self.listY]:
+                coeffs[medy] = {}
+                # mediator_parents = self._get_all_parents([medy]).intersection(mediators.union(self.X)) - {medy}
+                all_parents = self._get_all_parents([medy]) - {medy}
+                if 'dag' not in self.graph_type:
+                    spouses = self._get_all_spouses([medy]) - {medy}
+                    if len(spouses) != 0:
+                        raise ValueError("method == 'parents' only possible for "
+                                         "causal paths without adjacent bi-directed links!")
+
+                # print(j, all_parents[j])
+                # if len(all_parents[j]) > 0:
+                # print(medy, list(all_parents))
+                fit_res = self.model.get_general_fitted_model(
+                    Y=[medy], X=list(all_parents), Z=[],
+                    conditions=None,
+                    tau_max=self.tau_max,
+                    cut_off='tau_max',
+                    return_data=False)
+
+                for ipar, par in enumerate(list(all_parents)):
+                    # print(par, fit_res['model'].coef_)
+                    coeffs[medy][par] = fit_res['model'].coef_[0][ipar]
+
+        else:
+            raise ValueError("method must be 'optimal', 'links_coeffs', or 'parents'.")
+        
+        # Effect is sum over products over all path coefficients
+        # from x in X to y in Y
+        effect = {}
+        for (x, y) in itertools.product(self.listX, self.listY):
+            effect[(x, y)] = 0.
+            for causal_path in causal_paths[x][y]:
+                effect_here = 1.
+                # print(x, y, causal_path)
+                for index, node in enumerate(causal_path[:-1]):
+                    i, taui = node
+                    j, tauj = causal_path[index + 1]
+                    # tau_ij = abs(tauj - taui)
+                    # print((j, tauj), (i, taui))
+                    effect_here *= coeffs[(j, tauj)][(i, taui)]
+
+                effect[(x, y)] += effect_here
+               
+        # Make fitted coefficients available as attribute
+        self.coeffs = coeffs
+
+        # Modify and overwrite variables in self.model
+        self.model.Y = self.listY
+        self.model.X = self.listX  
+        self.model.Z = []
+        self.model.conditions = [] 
+        self.model.cut_off = 'tau_max' # 'max_lag_or_tau_max'
+
+        class dummy_fit_class():
+            def __init__(self, y_here, listX_here, effect_here):
+                dim = len(listX_here)
+                self.coeff_array = np.array([effect_here[(x, y_here)] for x in listX_here]).reshape(dim, 1)
+            def predict(self, X):
+                return np.dot(X, self.coeff_array).squeeze()
+
+        fit_results = {}
+        for y in self.listY:
+            fit_results[y] = {}
+            fit_results[y]['model'] = dummy_fit_class(y, self.listX, effect)
+            fit_results[y]['data_transform'] = deepcopy(data_transform)
+
+        # self.effect = effect
+        self.model.fit_results = fit_results
+        return self
+ 
+    def predict_wright_effect(self, 
+        intervention_data, 
+        pred_params=None,
+        ):
+        """Predict linear effect of intervention with fitted Wright-model.
+
+        Parameters
+        ----------
+        intervention_data : numpy array
+            Numpy array of shape (n_interventions, len(X)) that contains the do(X) values.
+        pred_params : dict, optional
+            Optional parameters passed on to sklearn prediction function.
+
+        Returns
+        -------
+        Results from prediction: an array of shape  (n_interventions, len(Y)).
+        """
+
+        lenX = len(self.listX)
+        lenY = len(self.listY)
+
+        if intervention_data.shape[1] != lenX:
+            raise ValueError("intervention_data.shape[1] must be len(X).")
+
+        if self.no_causal_path:
+            if self.verbosity > 0:
+                print("No causal path from X to Y exists.")
+            return np.zeros((len(intervention_data), len(self.Y)))
+
+        n_interventions, _ = intervention_data.shape
+
+
+        predicted_array = np.zeros((n_interventions, lenY))
+        pred_dict = {}
+        for iy, y in enumerate(self.listY):
+            # Print message
+            if self.verbosity > 1:
+                print("\n## Predicting target %s" % str(y))
+                if pred_params is not None:
+                    for key in list(pred_params):
+                        print("%s = %s" % (key, pred_params[key]))
+            # Default value for pred_params
+            if pred_params is None:
+                pred_params = {}
+            # Check this is a valid target
+            if y not in self.model.fit_results:
+                raise ValueError("y = %s not yet fitted" % str(y))
+
+            # data_transform is too complicated for Wright estimator
+            # Transform the data if needed
+            # fitted_data_transform = self.model.fit_results[y]['fitted_data_transform']
+            # if fitted_data_transform is not None:
+            #     intervention_data = fitted_data_transform['X'].transform(X=intervention_data)
+
+            # Now iterate through interventions (and potentially S)
+            for index, dox_vals in enumerate(intervention_data):
+                # Construct XZS-array
+                intervention_array = dox_vals.reshape(1, lenX) 
+                predictor_array = intervention_array
+
+                predicted_vals = self.model.fit_results[y]['model'].predict(
+                X=predictor_array, **pred_params)
+                predicted_array[index, iy] = predicted_vals.mean()
+
+                # data_transform is too complicated for Wright estimator
+                # if fitted_data_transform is not None:
+                #     rescaled = fitted_data_transform['Y'].inverse_transform(X=predicted_array[index, iy].reshape(-1, 1))
+                #     predicted_array[index, iy] = rescaled.squeeze()
+
+        return predicted_array
+
+
+    def fit_bootstrap_of(self, method, method_args, 
+                        boot_samples=100,
+                        boot_blocklength=1,
+                        seed=None):
+        """Runs chosen method on bootstrap samples drawn from DataFrame.
+        
+        Bootstraps for tau=0 are drawn from [max_lag, ..., T] and all lagged
+        variables constructed in DataFrame.construct_array are consistently
+        shifted with respect to this bootsrap sample to ensure that lagged
+        relations in the bootstrap sample are preserved.
+
+        This function fits the models, predict_bootstrap_of can then be used
+        to get confidence intervals for the effect of interventions.
+
+        Parameters
+        ----------
+        method : str
+            Chosen method among valid functions in this class.
+        method_args : dict
+            Arguments passed to method.
+        boot_samples : int
+            Number of bootstrap samples to draw.
+        boot_blocklength : int, optional (default: 1)
+            Block length for block-bootstrap.
+        seed : int, optional(default = None)
+            Seed for RandomState (default_rng)
+        """
+
+        # if dataframe.analysis_mode != 'single':
+        #     raise ValueError("CausalEffects class currently only supports single "
+        #                      "datasets.")
+
+        valid_methods = ['fit_total_effect',
+                         'fit_wright_effect',
+                          ]
+
+        if method not in valid_methods:
+            raise ValueError("method must be one of %s" % str(valid_methods))
+
+        # First call the method on the original dataframe 
+        # to make available adjustment set etc
+        getattr(self, method)(**method_args)
+
+        self.original_model = deepcopy(self.model)
+
+        if self.verbosity > 0:
+            print("\n##\n## Running Bootstrap of %s " % method +
+                  "\n##\n" +
+                  "\nboot_samples = %s \n" % boot_samples +
+                  "\nboot_blocklength = %s \n" % boot_blocklength
+                  )
+
+        method_args_bootstrap = deepcopy(method_args)
+        self.bootstrap_results = {}
+
+        for b in range(boot_samples):
+            # # Replace dataframe in method args by bootstrapped dataframe
+            # method_args_bootstrap['dataframe'].bootstrap = boot_draw
+            if seed is None:
+                random_state = np.random.default_rng(None)
+            else:
+                random_state = np.random.default_rng(seed*boot_samples + b)
+
+            method_args_bootstrap['dataframe'].bootstrap = {'boot_blocklength':boot_blocklength,
+                                                            'random_state':random_state}
+
+            # Call method and save fitted model
+            getattr(self, method)(**method_args_bootstrap)
+            self.bootstrap_results[b] = deepcopy(self.model)
+
+        # Reset model
+        self.model = self.original_model
+
+        return self
+
+
+    def predict_bootstrap_of(self, method, method_args, 
+                        conf_lev=0.9,
+                        return_individual_bootstrap_results=False):
+        """Predicts with fitted bootstraps.
+
+        To be used after fitting with fit_bootstrap_of. Only uses the 
+        expected values of the predict function, not potential other output.
+
+        Parameters
+        ----------
+        method : str
+            Chosen method among valid functions in this class.
+        method_args : dict
+            Arguments passed to method.
+        conf_lev : float, optional (default: 0.9)
+            Two-sided confidence interval.
+        return_individual_bootstrap_results : bool
+            Returns the individual bootstrap predictions.
+
+        Returns
+        -------
+        confidence_intervals : numpy array
+        """
+
+        valid_methods = ['predict_total_effect',
+                         'predict_wright_effect',
+                          ]
+
+        if method not in valid_methods:
+            raise ValueError("method must be one of %s" % str(valid_methods))
+
+        # def get_vectorized_length(W):
+        #     return sum([len(self.dataframe.vector_vars[w[0]]) for w in W])
+
+        lenX = len(self.listX)
+        lenS = len(self.listS)
+        lenY = len(self.listY)
+
+        n_interventions, _ = method_args['intervention_data'].shape
+
+        boot_samples = len(self.bootstrap_results)
+        # bootstrap_predicted_array = np.zeros((boot_samples, n_interventions, lenY))
+        
+        for b in range(boot_samples): #self.bootstrap_results.keys():
+            self.model = self.bootstrap_results[b]
+            boot_effect = getattr(self, method)(**method_args)
+
+            if isinstance(boot_effect, tuple):
+                boot_effect = boot_effect[0]
+            
+            if b == 0:
+                bootstrap_predicted_array = np.zeros((boot_samples, ) + boot_effect.shape, 
+                                            dtype=boot_effect.dtype)
+            bootstrap_predicted_array[b] = boot_effect
+
+        # Reset model
+        self.model = self.original_model
+
+        # Confidence intervals for val_matrix; interval is two-sided
+        c_int = (1. - (1. - conf_lev)/2.)
+        confidence_interval = np.percentile(
+                bootstrap_predicted_array, axis=0,
+                q = [100*(1. - c_int), 100*c_int])   #[:,:,0]
+
+        if return_individual_bootstrap_results:
+            return bootstrap_predicted_array, confidence_interval
+
+        return confidence_interval
+
+
+if __name__ == '__main__':
+    
+    # Consider some toy data
+    import tigramite
+    import tigramite.toymodels.structural_causal_processes as toys
+    import tigramite.data_processing as pp
+    import tigramite.plotting as tp
+    from matplotlib import pyplot as plt
+    import sys
+
+    import sklearn
+    from sklearn.linear_model import LinearRegression, LogisticRegression
+    from sklearn.preprocessing import StandardScaler
+    from sklearn.neural_network import MLPRegressor
+
+
+    # def lin_f(x): return x
+    # coeff = .5
+ 
+    # links_coeffs = {0: [((0, -1), 0.5, lin_f)],
+    #          1: [((1, -1), 0.5, lin_f), ((0, -1), 0.5, lin_f)],
+    #          2: [((2, -1), 0.5, lin_f), ((1, 0), 0.5, lin_f)]
+    #          }
+    # T = 1000
+    # data, nonstat = toys.structural_causal_process(
+    #     links_coeffs, T=T, noises=None, seed=7)
+    # dataframe = pp.DataFrame(data)
+
+    # graph = CausalEffects.get_graph_from_dict(links_coeffs)
+
+    # original_graph = np.array([[['', ''],
+    #     ['-->', ''],
+    #     ['-->', ''],
+    #     ['', '']],
+
+    #    [['<--', ''],
+    #     ['', '-->'],
+    #     ['-->', ''],
+    #     ['-->', '']],
+
+    #    [['<--', ''],
+    #     ['<--', ''],
+    #     ['', '-->'],
+    #     ['-->', '']],
+
+    #    [['', ''],
+    #     ['<--', ''],
+    #     ['<--', ''],
+    #     ['', '-->']]], dtype='<U3')
+    # graph = np.copy(original_graph)
+
+    # # Add T <-> Reco and T 
+    # graph[2,3,0] = '+->' ; graph[3,2,0] = '<-+'
+    # graph[1,3,1] = '<->' #; graph[2,1,0] = '<--'
+
+    # added = np.zeros((4, 4, 1), dtype='<U3')
+    # added[:] = ""
+    # graph = np.append(graph, added , axis=2)
+
+
+    # X = [(1, 0)]
+    # Y = [(3, 0)]
+
+    # # # Initialize class as `stationary_dag`
+    # causal_effects = CausalEffects(graph, graph_type='stationary_admg', 
+    #                             X=X, Y=Y, S=None, 
+    #                             hidden_variables=None, 
+    #                             verbosity=0)
+
+    # print(causal_effects.get_optimal_set())
+
+    # tp.plot_time_series_graph(
+    #     graph = graph,
+    #     save_name='Example_graph_in.pdf',
+    #     # special_nodes=special_nodes,
+    #     # var_names=var_names,
+    #     figsize=(6, 4),
+    #     )
+
+    # tp.plot_time_series_graph(
+    #     graph = causal_effects.graph,
+    #     save_name='Example_graph_out.pdf',
+    #     # special_nodes=special_nodes,
+    #     # var_names=var_names,
+    #     figsize=(6, 4),
+    #     )
+
+    # causal_effects.fit_wright_effect(dataframe=dataframe, 
+    #                         # links_coeffs = links_coeffs,
+    #                         # mediation = [(1, 0), (1, -1), (1, -2)]
+    #                         )
+
+    # intervention_data = 1.*np.ones((1, 1))
+    # y1 = causal_effects.predict_wright_effect( 
+    #         intervention_data=intervention_data,
+    #         )
+
+    # intervention_data = 0.*np.ones((1, 1))
+    # y2 = causal_effects.predict_wright_effect( 
+    #         intervention_data=intervention_data,
+    #         )
+
+    # beta = (y1 - y2)
+    # print("Causal effect is %.5f" %(beta))
+
+    # tp.plot_time_series_graph(
+    #     graph = causal_effects.graph,
+    #     save_name='Example_graph.pdf',
+    #     # special_nodes=special_nodes,
+    #     var_names=var_names,
+    #     figsize=(8, 4),
+    #     )
+
+    T = 10000
+    def lin_f(x): return x
+
+    auto_coeff = 0.
+    coeff = 2.
+
+    links = {
+            0: [((0, -1), auto_coeff, lin_f)], 
+            1: [((1, -1), auto_coeff, lin_f)], 
+            2: [((2, -1), auto_coeff, lin_f), ((0, 0), coeff, lin_f)],
+            3: [((3, -1), auto_coeff, lin_f)], 
+            }
+    data, nonstat = toys.structural_causal_process(links, T=T, 
+                                noises=None, seed=7)
+
+
+    # # Create some missing values
+    # data[-10:,:] = 999.
+    # var_names = range(2)
+
+    dataframe = pp.DataFrame(data,
+                    vector_vars={0:[(0,0), (1,0)], 
+                                 1:[(2,0), (3,0)]}
+                    )
+
+    # # Construct expert knowledge graph from links here 
+    aux_links = {0: [(0, -1)],
+                 1: [(1, -1), (0, 0)],
+              }
+    # # Use staticmethod to get graph
+    graph = CausalEffects.get_graph_from_dict(aux_links, tau_max=2)
+    # graph = np.array([['', '-->'],
+    #                   ['<--', '']], dtype='<U3')
+    
+    # # We are interested in lagged total effect of X on Y
+    X = [(0, 0), (0, -1)]
+    Y = [(1, 0), (1, -1)]
+
+    # # Initialize class as `stationary_dag`
+    causal_effects = CausalEffects(graph, graph_type='stationary_dag', 
+                                X=X, Y=Y, S=None, 
+                                hidden_variables=None, 
+                                verbosity=1)
+
+    # print(data)
+    # # Optimal adjustment set (is used by default)
+    # # print(causal_effects.get_optimal_set())
+
+    # # # Fit causal effect model from observational data
+    causal_effects.fit_total_effect(
+        dataframe=dataframe, 
+        # mask_type='y',
+        estimator=LinearRegression(),
+        )
+
+    # # Fit causal effect model from observational data
+    # causal_effects.fit_bootstrap_of(
+    #     method='fit_total_effect',
+    #     method_args={'dataframe':dataframe,  
+    #     # mask_type='y',
+    #     'estimator':LinearRegression()
+    #     },
+    #     boot_samples=3,
+    #     boot_blocklength=1,
+    #     seed=5
+    #     )
+
+
+    # Predict effect of interventions do(X=0.), ..., do(X=1.) in one go
+    lenX = 4 # len(dataframe.vector_vars[X[0][0]])
+    dox_vals = np.linspace(0., 1., 3)
+    intervention_data = np.tile(dox_vals.reshape(len(dox_vals), 1), lenX)
+
+    intervention_data = np.array([[1., 0., 0., 0.]])
+
+    print(intervention_data)
+
+    pred_Y = causal_effects.predict_total_effect( 
+            intervention_data=intervention_data)
+    print(pred_Y, pred_Y.shape)
+
+
+
+
+
+    # # Predict effect of interventions do(X=0.), ..., do(X=1.) in one go
+    # # dox_vals = np.array([1.]) #np.linspace(0., 1., 1)
+    # intervention_data = np.tile(dox_vals.reshape(len(dox_vals), 1), len(X))
+    # conf = causal_effects.predict_bootstrap_of(
+    #     method='predict_total_effect',
+    #     method_args={'intervention_data':intervention_data})
+    # print(conf, conf.shape)
+
+
+
+    # # # Predict effect of interventions do(X=0.), ..., do(X=1.) in one go
+    # # dox_vals = np.array([1.]) #np.linspace(0., 1., 1)
+    # # intervention_data = dox_vals.reshape(len(dox_vals), len(X))
+    # # pred_Y = causal_effects.predict_total_effect( 
+    # #         intervention_data=intervention_data)
+    # # print(pred_Y)
+
+
+
+    # # Fit causal effect model from observational data
+    # causal_effects.fit_wright_effect(
+    #     dataframe=dataframe, 
+    #     # mask_type='y',
+    #     # estimator=LinearRegression(),
+    #     # data_transform=StandardScaler(),
+    #     )
+
+    # # # Predict effect of interventions do(X=0.), ..., do(X=1.) in one go
+    # dox_vals = np.linspace(0., 1., 5)
+    # intervention_data = dox_vals.reshape(len(dox_vals), len(X))
+    # pred_Y = causal_effects.predict_wright_effect( 
+    #         intervention_data=intervention_data)
+    # print(pred_Y)