tigramite-fast 5.2.10.1__py3-none-any.whl

tigramite/jpcmciplus.py

@@ -0,0 +1,966 @@
+"""Tigramite causal discovery for time series."""
+
+# Authors: Wiebke Günther <wiebke.guenther@dlr.de>, Urmi Ninad, Jakob Runge <jakob@jakob-runge.com>
+#
+# License: GNU General Public License v3.0
+
+from __future__ import print_function
+import numpy as np
+from tigramite.pcmci import PCMCI
+from copy import deepcopy
+import itertools
+
+from tigramite.toymodels.context_model import _group_links
+
+
+class JPCMCIplus(PCMCI):
+    r"""J-PCMCIplus causal discovery for time series datasets from multiple contexts.
+        
+    This class is based on the PCMCI framework as described in
+    [i]. JPCMCIplus enables causal discovery for time series data from
+    different contexts, i.e. datasets, where some of the variables
+    describing the context might be unobserved. The method is described
+    in detail in [ii]. See the tutorial for guidance in applying the
+    method.
+
+    References
+    ----------
+    .. [i] J. Runge, P. Nowack, M. Kretschmer, S. Flaxman, D. Sejdinovic,
+       Detecting and quantifying causal associations in large nonlinear
+       time series datasets. Sci. Adv. 5, eaau4996
+       (2019) https://advances.sciencemag.org/content/5/11/eaau4996
+    
+    .. [ii] W. Günther, U. Ninad, J. Runge, Causal discovery for time
+       series from multiple datasets with latent contexts. UAI 2023
+    
+    Parameters
+    ----------
+    node_classification : dictionary
+        Classification of nodes into system, context, or dummy nodes.
+        Keys of the dictionary are from {0, ..., N-1} where N is the number of nodes.
+        Options for the values are "system", "time_context", "space_context", "time_dummy", or "space_dummy".
+    
+    Attributes
+    ----------
+    all_parents : dictionary
+        Dictionary of form {0:[(0, -1), (3, -2), ...], 1:[], ...} containing
+        the conditioning-parents estimated with PC algorithm.
+    val_min : dictionary
+        Dictionary of form val_min[j][(i, -tau)] = float
+        containing the minimum test statistic value for each link estimated in
+        the PC algorithm.
+    pval_max : dictionary
+        Dictionary of form pval_max[j][(i, -tau)] = float containing the maximum
+        p-value for each link estimated in the PC algorithm.
+    iterations : dictionary
+        Dictionary containing further information on algorithm steps.
+    N : int
+        Number of variables.
+    T : dict
+        Time series sample length of dataset(s).
+    dummy_parents : dictionary or None
+        Dictionary of form {0:[(0, -1), (3, -2), ...], 1:[], ...} containing
+        the dependence of the system nodes on the dummy nodes.
+    observed_context_parents : dictionary or None
+        Dictionary of form {0:[(0, -1), (3, -2), ...], 1:[], ...} containing
+        the dependence of the system nodes on the observed context nodes.
+    dummy_ci_test : conditional independence test object
+        Conditional independence test used to test dependence between system nodes and dummy nodes.
+        Currently, ParCorr is used with one-hot encoded dummies.
+    mode : "system_search" or "context_search" or "dummy_search" (default: "system_search")
+    time_context_nodes : list
+        List with entries from {0, ..., N-1} where N is the number of nodes.
+        This is the list of the temporal context nodes which are assumed to be constant over the different datasets.
+    space_context_nodes :
+        List with entries from {0, ..., N-1} where N is the number of nodes.
+        This is the list of the spatial context nodes which are assumed to be constant over time.
+    time_dummy : int or None (default: None)
+        Node corresponding to the temporal dummy variable.
+    space_dummy : int or None (default: None)
+        Node corresponding to the spatial dummy variable.
+    system_nodes : list
+        List with entries from {0, ..., N-1} where N is the number of nodes.
+        This is the list of the system nodes.
+    """
+
+    def __init__(self, node_classification, **kwargs):
+        # Init base class
+        PCMCI.__init__(self, **kwargs)
+
+        self.system_nodes = self.group_nodes(node_classification, "system")
+        self.time_context_nodes = self.group_nodes(node_classification, "time_context")
+        self.space_context_nodes = self.group_nodes(node_classification, "space_context")
+        self.time_dummy = self.group_nodes(node_classification, "time_dummy")
+        self.space_dummy = self.group_nodes(node_classification, "space_dummy")
+
+        self.dummy_parents = {i: [] for i in range(self.N)}
+        self.observed_context_parents = {i: [] for i in range(self.N)}
+        self.mode = "system_search"
+
+    def group_nodes(self, node_types, node_type):
+        nodes = range(self.N)
+        return [node for node in nodes if node_types[node] == node_type]
+
+    def run_jpcmciplus(self,
+                       contemp_collider_rule='majority',
+                       link_assumptions=None,
+                       tau_min=0,
+                       tau_max=1,
+                       pc_alpha=0.01,
+                       conflict_resolution=True,
+                       reset_lagged_links=False,
+                       max_conds_dim=None,
+                       max_combinations=1,
+                       max_conds_py=None,
+                       max_conds_px=None,
+                       max_conds_px_lagged=None,
+                       fdr_method='none'):
+        """Runs JPCMCIplus time-lagged and contemporaneous causal discovery for time series from multiple contexts.
+        
+        Method described in: W. Günther, U. Ninad, J. Runge, Causal discovery
+        for time series from multiple datasets with latent contexts. UAI
+        2023
+        
+        Notes
+        -----
+        The JPCMCIplus causal discovery method is described in [ii], where
+        also analytical and numerical results are presented. JPCMCIplus can identify the joint causal graph
+        over multiple datasets containing time series data from different contexts under the standard assumptions
+        of Causal Sufficiency, Faithfulness and the Markov condition, as well as some background knowledge assumptions.
+        JPCMCIplus estimates time-lagged and contemporaneous causal links from context to system
+        variables and in between system variables by a four-step procedure:
+
+        1.  **Discovery of supersets of the lagged parents of the system and observed temporal context nodes** by
+        running the :math:`PC_1` lagged phase on this subset of nodes to obtain :math:`\\hat{\\mathcal{B}}^-_t(X_t^j)`.
+
+        2.  Next, the **MCI test is run on pairs of system and context nodes conditional on subsets of system
+        and context**, i.e. perform MCI tests for pairs :math:`((C^j_{t-\\tau}, X^i_t))_{\\tau > 0}`,
+        :math:`(C_t^j, X_t^i)`, :math:`(X_t^i, C_t^j)` for all :math:`i,j`,
+
+        .. math:: C_{t-\\tau}^i \\perp X_t^j | \\mathbf{S}, \\hat{\\mathcal{B}}^-_t(X_t^j)
+                    \\setminus \\{ C_{t-\\tau}^i \\}, \\hat{\\mathcal{B}}^-_{t-\\tau}(C_{t-\\tau}^i)
+
+        with :math:`\\mathbf{S}` being a subset of the contemporaneous adjacencies :math:`\\mathcal{A}_t(X_t^j)` and
+        :math:`\\hat{\\mathcal{B}}^-_t(X_t^j)` are the lagged adjacencies from step one. If :math:`C` is a
+        spatial context variable, we only have to test the contemporaneous pairs
+        :math:`(C_t^j, X_t^i)`, :math:`(X_t^i, C_t^j)` for all :math:`i,j`.
+        If :math:`C_t^j` and :math:`X_t^i` are conditionally independent, all lagged links between :math:`C_t^j` and
+        :math:`X^j_{t-\\tau}` are also removed for all :math:`\\tau`.
+
+        3.  **Perform MCI tests on all system-dummy pairs conditional on the superset of lagged links, the discovered
+        contemporaneous context adjacencies, as well as on subsets of contemporaneous system links**, i.e. test
+        for :math:`(D, X_t^i)`, :math:`(X_t^i, D)` for all :math:`i`, i.e.
+
+        .. math:: D \\perp X_t^j | \\mathbf{S}, \\hat{\\mathcal{B}}^C_t(X_t^j),
+
+        where :math:`\\mathbf{S} \\subset \\mathcal{A}_t(X_t^i)` and :math:`\\hat{\\mathcal{B}}^C_t(X_t^j)`
+        are the lagged and contextual adjacencies found in the previous step.
+        If :math:`D` and :math:`X_t^j` are found to be conditionally independence, links between :math:`D` and
+        :math:`X^j_{t-\\tau}` are removed for all :math:`\\tau`.
+        By assumption context node is the parent in all system-context links.
+
+        4.  Finally, we **perform  MCI tests on all system pairs conditional on discovered lagged, context and dummy
+        adjacencies, as well as on subsets of contemporaneous system links** and **orientation phase**. In more detail,
+        we perform MCI test for pairs :math:`((X^j_{t-\\tau}, X_t^i))_{\\tau > 0}`, :math:`(X_t^i, X_t^j)` for all
+        :math:`i, j`, i.e.
+
+        .. math:: X^i_{t-\\tau} \\perp X_t^j | \\mathbf{S}, \\hat{\\mathcal{B}}^{CD}_t(X_t^j)
+                    \\setminus \\{ X_{t-\\tau}^i \\},\\hat{\\mathcal{B}}^{CD}_t(X_{t-\\tau}^i)
+
+        where :math:`\\mathbf{S} \\subset \\mathcal{A}_t(X_t^i)` and :math:`\\hat{\\mathcal{B}}^{CD}_t(X_t^j)`
+        are the lagged, contextual, and dummy adjacencies found in the previous steps.
+        Finally, all remaining edges (without expert knowledge) are oriented using the PCMCI+ orientation phase while
+        making use of all triples involving one context or dummy variable and two system variables as in the non-time
+        series case.
+
+        JPCMCIplus can be flexibly combined with any kind of conditional
+        independence test statistic adapted to the kind of data (continuous
+        or discrete) and its assumed dependency types. These are available in
+        ``tigramite.independence_tests``.
+        See PCMCIplus for a description of the parameters of JPCMCIplus. Also, guidance on best practices for
+        setting these parameters is given there.
+
+        Parameters
+        ----------
+        link_assumptions : dict
+            Dictionary of form {j:{(i, -tau): link_type, ...}, ...} specifying
+            assumptions about links. This initializes the graph with entries
+            graph[i,j,tau] = link_type. For example, graph[i,j,0] = '-->'
+            implies that a directed link from i to j at lag 0 must exist.
+            Valid link types are 'o-o', '-->', '<--'. In addition, the middle
+            mark can be '?' instead of '-'. Then '-?>' implies that this link
+            may not exist, but if it exists, its orientation is '-->'. Link
+            assumptions need to be consistent, i.e., graph[i,j,0] = '-->'
+            requires graph[j,i,0] = '<--' and acyclicity must hold. If a link
+            does not appear in the dictionary, it is assumed absent. That is,
+            if link_assumptions is not None, then all links have to be specified
+            or the links are assumed absent.
+        tau_min : int, optional (default: 0)
+            Minimum time lag to test.
+        tau_max : int, optional (default: 1)
+            Maximum time lag. Must be larger or equal to tau_min.
+        pc_alpha : float or list of floats, default: 0.01
+            Significance level in algorithm. If a list or None is passed, the
+            pc_alpha level is optimized for every graph across the given
+            pc_alpha values ([0.001, 0.005, 0.01, 0.025, 0.05] for None) using
+            the score computed in cond_ind_test.get_model_selection_criterion().
+        contemp_collider_rule : {'majority', 'conservative', 'none'}
+            Rule for collider phase to use. See the paper for details. Only
+            'majority' and 'conservative' lead to an order-independent
+            algorithm.
+        conflict_resolution : bool, optional (default: True)
+            Whether to mark conflicts in orientation rules. Only for True
+            this leads to an order-independent algorithm.
+        reset_lagged_links : bool, optional (default: False)
+            Restricts the detection of lagged causal links in Step 2 to the
+            significant adjacencies found in the PC1 algorithm in Step 1. For
+            True, *all* lagged links are considered again, which improves
+            detection power for lagged links, but also leads to larger
+            runtimes.
+        max_conds_dim : int, optional (default: None)
+            Maximum number of conditions to test. If None is passed, this number
+            is unrestricted.
+        max_combinations : int, optional (default: 1)
+            Maximum number of combinations of conditions of current cardinality
+            to test. Defaults to 1 for PC_1 algorithm. For original PC algorithm
+            a larger number, such as 10, can be used.
+        max_conds_py : int, optional (default: None)
+            Maximum number of lagged conditions of Y to use in MCI tests. If
+            None is passed, this number is unrestricted.
+        max_conds_px : int, optional (default: None)
+            Maximum number of lagged conditions of X to use in MCI tests. If
+            None is passed, this number is unrestricted.
+        max_conds_px_lagged : int, optional (default: None)
+            Maximum number of lagged conditions of X when X is lagged in MCI
+            tests. If None is passed, this number is equal to max_conds_px.
+        fdr_method : str, optional (default: 'none')
+            Correction method, default is Benjamini-Hochberg False Discovery
+            Rate method.
+        Returns
+        -------
+        graph : array of shape [N, N, tau_max+1]
+            Resulting causal graph, see description above for interpretation.
+        val_matrix : array of shape [N, N, tau_max+1]
+            Estimated matrix of test statistic values regarding adjacencies.
+        p_matrix : array of shape [N, N, tau_max+1]
+            Estimated matrix of p-values regarding adjacencies.
+        sepset : dictionary
+            Separating sets. See paper for details.
+        ambiguous_triples : list
+            List of ambiguous triples, only relevant for 'majority' and
+            'conservative' rules, see paper for details.
+        """
+        observed_context_nodes = self.time_context_nodes + self.space_context_nodes
+
+        # initialize / clean link_assumptions
+        if link_assumptions is not None:
+            _link_assumptions = deepcopy(link_assumptions)
+        else:
+            _link_assumptions = self._set_link_assumptions(link_assumptions, tau_min, tau_max,
+                       remove_contemp=False)
+            # {j: {(i, -tau): ("-?>" if tau > 0 else "o?o") for i in range(self.N)
+            #                          for tau in range(tau_max + 1)} for j in range(self.N)}
+
+        _link_assumptions = self.assume_exogenous_context(_link_assumptions, observed_context_nodes)
+        _link_assumptions = self.clean_link_assumptions(_link_assumptions, tau_max)
+
+        # for j in _link_assumptions:
+        #     print(j, _link_assumptions[j])
+        # self._set_link_assumptions(_link_assumptions, tau_min, tau_max,
+        #                        remove_contemp=False)
+
+        # Check if pc_alpha is chosen to optimize over a list
+        if pc_alpha is None or isinstance(pc_alpha, (list, tuple, np.ndarray)):
+            # Call optimizer wrapper around run_pcmciplus()
+            return self._optimize_pcmciplus_alpha(
+                link_assumptions=link_assumptions,
+                tau_min=tau_min,
+                tau_max=tau_max,
+                pc_alpha=pc_alpha,
+                contemp_collider_rule=contemp_collider_rule,
+                conflict_resolution=conflict_resolution,
+                reset_lagged_links=reset_lagged_links,
+                max_conds_dim=max_conds_dim,
+                max_combinations=max_combinations,
+                max_conds_py=max_conds_py,
+                max_conds_px=max_conds_px,
+                max_conds_px_lagged=max_conds_px_lagged,
+                fdr_method=fdr_method)
+
+        elif pc_alpha < 0. or pc_alpha > 1:
+            raise ValueError("Choose 0 <= pc_alpha <= 1")
+
+        # Check the limits on tau
+        self._check_tau_limits(tau_min, tau_max)
+
+        # Step 0 and 1:
+        context_results = self.discover_lagged_context_system_links(
+            _link_assumptions,
+            tau_min=tau_min,
+            tau_max=tau_max,
+            pc_alpha=pc_alpha,
+            reset_lagged_links=reset_lagged_links,
+            max_conds_dim=max_conds_dim,
+            max_combinations=max_combinations,
+            max_conds_py=max_conds_py,
+            max_conds_px=max_conds_px,
+            max_conds_px_lagged=max_conds_px_lagged,
+            fdr_method=fdr_method
+        )
+        ctxt_res = deepcopy(context_results)
+        # Store the parents in the pcmci member
+        self.observed_context_parents = deepcopy(context_results['parents'])
+        self.all_lagged_parents = deepcopy(
+            context_results['lagged_parents'])  # remove context nodes from lagged parents
+        self.all_lagged_parents = {i: [el for el in self.all_lagged_parents[i] if el[0] not in self.time_context_nodes]
+                                   for i in
+                                   range(self.N)}
+
+        # if self.verbosity > 0:
+        #     print("\nDiscovered observed context parents: ", context_results['parents'])
+
+        if len(self.time_dummy) > 0 or len(self.space_dummy) > 0:
+            # step 2:
+            dummy_system_results = self.discover_dummy_system_links(
+                _link_assumptions,
+                ctxt_res,
+                self.all_lagged_parents,
+                tau_min=tau_min,
+                tau_max=tau_max,
+                pc_alpha=pc_alpha,
+                reset_lagged_links=reset_lagged_links,
+                max_conds_dim=max_conds_dim,
+                max_conds_py=max_conds_py,
+                max_conds_px=max_conds_px,
+                max_conds_px_lagged=max_conds_px_lagged,
+                fdr_method=fdr_method
+            )
+            # Store the parents in the pcmci member
+            self.dummy_parents = dummy_system_results['parents']
+        else:
+            dummy_system_results = deepcopy(context_results)
+
+        # if self.verbosity > 0:
+        #     print("Discovered dummy parents: ", self.dummy_parents)
+
+        # step 3:
+        self.mode = "system_search"
+
+        lagged_context_dummy_parents = {
+            i: list(
+                dict.fromkeys(self.all_lagged_parents[i] + self.observed_context_parents[i] + self.dummy_parents[i]))
+            for i in self.system_nodes}
+        # we only care about the parents of system nodes
+        lagged_context_dummy_parents.update(
+            {i: [] for i in observed_context_nodes + self.time_dummy + self.space_dummy})
+
+        dummy_system_results_copy = deepcopy(dummy_system_results)
+
+        # step 4:
+        system_skeleton_results = self.discover_system_system_links(link_assumptions=_link_assumptions,
+                                                                    lagged_context_dummy_parents=lagged_context_dummy_parents,
+                                                                    tau_min=tau_min,
+                                                                    tau_max=tau_max,
+                                                                    pc_alpha=pc_alpha,
+                                                                    reset_lagged_links=reset_lagged_links,
+                                                                    max_conds_dim=max_conds_dim,
+                                                                    max_conds_py=max_conds_py,
+                                                                    max_conds_px=max_conds_px,
+                                                                    max_conds_px_lagged=max_conds_px_lagged,
+                                                                    fdr_method=fdr_method)
+
+        # orientation phase
+        colliders_step_results = self._pcmciplus_collider_phase(
+            system_skeleton_results['graph'], system_skeleton_results['sepsets'],
+            lagged_context_dummy_parents, pc_alpha,
+            tau_min, tau_max, max_conds_py, max_conds_px, max_conds_px_lagged,
+            conflict_resolution, contemp_collider_rule)
+
+        final_graph = self._pcmciplus_rule_orientation_phase(colliders_step_results['graph'],
+                                                             colliders_step_results['ambiguous_triples'],
+                                                             conflict_resolution)
+
+        # add context-system and dummy-system values and pvalues back in (lost because of link_assumption)
+        for c in observed_context_nodes + self.time_dummy + self.space_dummy:
+            for j in range(self.N):
+                for lag in range(tau_max + 1):
+                    # add context-system links to results
+                    system_skeleton_results['val_matrix'][c, j, lag] = dummy_system_results_copy['val_matrix'][
+                        c, j, lag]
+                    system_skeleton_results['val_matrix'][j, c, lag] = dummy_system_results_copy['val_matrix'][
+                        j, c, lag]
+
+                    system_skeleton_results['p_matrix'][c, j, lag] = dummy_system_results_copy['p_matrix'][c, j, lag]
+                    system_skeleton_results['p_matrix'][j, c, lag] = dummy_system_results_copy['p_matrix'][j, c, lag]
+
+        # No confidence interval estimation here
+        return_dict = {'graph': final_graph, 'p_matrix': system_skeleton_results['p_matrix'],
+                       'val_matrix': system_skeleton_results['val_matrix'],
+                       'sepsets': colliders_step_results['sepsets'],
+                       'ambiguous_triples': colliders_step_results['ambiguous_triples'], 
+                       'conf_matrix': None}
+
+        # Print the results
+        if self.verbosity > 0:
+            self.print_results(return_dict, alpha_level=pc_alpha)
+
+        # Return the dictionary
+        self.results = return_dict
+
+        # Return the dictionary
+        return return_dict
+
+    def assume_exogenous_context(self, link_assumptions, observed_context_nodes):
+        """Helper function to amend the link_assumptions to ensure that all context-system links are oriented
+        such that the context variable is the parent."""
+        for j in link_assumptions:
+            if j in self.system_nodes:
+                for link in link_assumptions[j]:
+                    i, lag = link
+                    if i in observed_context_nodes + self.time_dummy + self.space_dummy:  # is context var
+                        link_type = link_assumptions[j][link]
+                        link_assumptions[j][link] = '-' + link_type[1] + '>'
+        return link_assumptions
+
+    def clean_link_assumptions(self, link_assumptions, tau_max):
+        """Helper function to amend the link_assumptions in the following ways
+            * remove any links where dummy is the child
+            * remove any lagged links to dummy, and space_context (not to observed time context)
+            * and system - context links where context is the child
+            * and any links between spatial and temporal context
+        """
+        for node in self.time_dummy + self.space_dummy:
+            link_assumptions[node] = {}
+
+        for j in self.system_nodes + self.time_context_nodes + self.space_context_nodes:
+            for lag in range(1, tau_max + 1):
+                for c in self.time_dummy + self.space_dummy + self.space_context_nodes:
+                    if (c, -lag) in link_assumptions[j]: link_assumptions[j].pop((c, -lag), None)
+        for c in self.space_context_nodes + self.time_context_nodes:
+            for j in self.system_nodes:
+                for lag in range(tau_max + 1):
+                    if (j, -lag) in link_assumptions[c]: link_assumptions[c].pop((j, -lag), None)
+            if (c, 0) in link_assumptions[c]: link_assumptions[c].pop((c, 0), None)  # remove self-links
+
+        for c in self.space_context_nodes:
+            for k in self.time_context_nodes:
+                for lag in range(tau_max + 1):
+                    if (k, -lag) in link_assumptions[c]: link_assumptions[c].pop((k, -lag), None)
+        for c in self.time_context_nodes:
+            for k in self.space_context_nodes:
+                if (k, 0) in link_assumptions[c]: link_assumptions[c].pop((k, 0), None)
+
+        return link_assumptions
+
+    def remove_dummy_link_assumptions(self, link_assumptions):
+        """Helper function to remove any links to dummy from link_assumptions."""
+        link_assumptions_wo_dummy = deepcopy(link_assumptions)
+        for j in self.system_nodes + self.time_context_nodes + self.space_context_nodes:
+            for dummy_node in self.time_dummy + self.space_dummy:
+                if (dummy_node, 0) in link_assumptions_wo_dummy[j]:
+                    link_assumptions_wo_dummy[j].pop((dummy_node, 0), None)
+        return link_assumptions_wo_dummy
+
+    def add_found_context_link_assumptions(self, link_assumptions, tau_max):
+        """Helper function to add discovered links between system and observed context nodes to link_assumptions."""
+        link_assumptions_dummy = deepcopy(link_assumptions)
+
+        for c in self.space_context_nodes + self.time_context_nodes:
+            link_assumptions_dummy[c] = {}
+        for j in self.system_nodes + self.time_context_nodes + self.space_context_nodes:
+            for c in self.space_context_nodes + self.time_context_nodes:
+                for lag in range(tau_max + 1):
+                    if (c, -lag) in link_assumptions_dummy[j]:
+                        link_assumptions_dummy[j].pop((c, -lag), None)
+            link_assumptions_dummy[j].update({parent: '-->' for parent in self.observed_context_parents[j]})
+
+        return link_assumptions_dummy
+
+    def clean_system_link_assumptions(self, link_assumptions, tau_max):
+        """Helper function to remove any links to dummy and observed context nodes from link_assumptions.
+            Add discovered links to contextual parents (from steps 1 and 2) to the link_assumptions.
+        """
+        dummy_vars = self.time_dummy + self.space_dummy
+        observed_context_nodes = self.time_context_nodes + self.space_context_nodes
+        system_links = deepcopy(link_assumptions)
+
+        for j in self.system_nodes:
+            for C in dummy_vars + observed_context_nodes:
+                for lag in range(tau_max + 1):
+                    if (C, -lag) in system_links[j]:
+                        system_links[j].pop((C, -lag), None)
+
+        for j in system_links:
+            system_links[j].update(
+                {parent: '-->' for parent in self.observed_context_parents[j] + self.dummy_parents[j]})
+
+        for C in observed_context_nodes + dummy_vars:
+            # we are not interested in links between context variables (thus system_links[C] = {})
+            system_links[C] = {}
+        return system_links
+
+    def discover_lagged_context_system_links(self, link_assumptions,
+                                             tau_min=0,
+                                             tau_max=1, pc_alpha=0.01,
+                                             reset_lagged_links=False,
+                                             max_conds_dim=None,
+                                             max_combinations=1,
+                                             max_conds_py=None,
+                                             max_conds_px=None,
+                                             max_conds_px_lagged=None,
+                                             fdr_method='none'):
+        """
+        Step 1 of JPCMCIplus, i.e. discovery of links between observed context nodes and system nodes through an
+        application of the skeleton phase of PCMCIplus to this subset of nodes (observed context nodes and system
+        nodes).
+        See run_jpcmciplus for a description of the parameters.
+
+        Returns
+        -------
+        graph : array of shape [N, N, tau_max+1]
+            Resulting causal graph, see description above for interpretation.
+        val_matrix : array of shape [N, N, tau_max+1]
+            Estimated matrix of test statistic values regarding adjacencies.
+        p_matrix : array of shape [N, N, tau_max+1]
+            Estimated matrix of p-values regarding adjacencies.
+        parents : dictionary
+            Dictionary of form {0:[(0, -1), (3, -2), ...], 1:[], ...} containing
+            the estimated context parents of the system nodes.
+        lagged_parents : dictionary
+            Dictionary of form {0:[(0, -1), (3, -2), ...], 1:[], ...} containing
+            the conditioning-parents estimated with PC algorithm.
+        """
+
+        # Initializing
+        context_parents = {i: [] for i in range(self.N)}
+
+        # find links between expressive context, and between expressive context and system nodes
+        # here, we exclude any links to dummy
+        _link_assumptions_wo_dummy = self.remove_dummy_link_assumptions(link_assumptions)
+        _int_link_assumptions = self._set_link_assumptions(_link_assumptions_wo_dummy, tau_min, tau_max)
+
+        # Step 1: Get a superset of lagged parents from run_pc_stable
+        if self.verbosity > 0:
+            print("\n##\n## J-PCMCI+ Step 1: Selecting lagged conditioning sets\n##")
+
+        lagged_parents = self.run_pc_stable(link_assumptions=link_assumptions,
+                                            tau_min=tau_min,
+                                            tau_max=tau_max,
+                                            pc_alpha=pc_alpha,
+                                            max_conds_dim=max_conds_dim,
+                                            max_combinations=max_combinations)
+
+        self.mode = "context_search"
+
+        p_matrix = self.p_matrix
+        val_matrix = self.val_matrix
+
+        # run PCMCI+ skeleton phase on subset of links to discover context-system links
+        if self.verbosity > 0:
+            print("\n##\n## J-PCMCI+ Step 2: Discovering context-system links\n##")
+            if link_assumptions is not None:
+                print("\nWith link_assumptions = %s" % str(_int_link_assumptions))
+
+        skeleton_results = self._pcmciplus_mci_skeleton_phase(
+            lagged_parents, _int_link_assumptions, pc_alpha,
+            tau_min, tau_max, max_conds_dim, None,
+            max_conds_py, max_conds_px, max_conds_px_lagged,
+            reset_lagged_links, fdr_method,
+            p_matrix, val_matrix
+        )
+
+        skeleton_graph = skeleton_results['graph']
+
+        for j in self.system_nodes + self.time_context_nodes + self.space_context_nodes:
+            for c in self.space_context_nodes + self.time_context_nodes:
+                for k in range(tau_max + 1):
+                    if skeleton_graph[c, j, k] == 'o?o' or skeleton_graph[c, j, k] == '-?>' or skeleton_graph[
+                        c, j, k] == 'o-o' or skeleton_graph[c, j, k] == '-->':
+                        context_parents[j].append((c, -k))
+
+        return_dict = {'graph': skeleton_results['graph'], 'p_matrix': skeleton_results['p_matrix'],
+                       'val_matrix': skeleton_results['val_matrix'],
+                       'parents': context_parents, 'lagged_parents': lagged_parents}
+
+        # Print the results
+        if self.verbosity > 0:
+            self.print_results(return_dict, alpha_level=pc_alpha)
+
+        return return_dict
+
+    def discover_dummy_system_links(self, link_assumptions,
+                                    context_system_results,
+                                    lagged_parents,
+                                    tau_min=0,
+                                    tau_max=1,
+                                    pc_alpha=0.01,
+                                    reset_lagged_links=False,
+                                    max_conds_dim=None,
+                                    max_conds_py=None,
+                                    max_conds_px=None,
+                                    max_conds_px_lagged=None,
+                                    fdr_method='none'):
+        """
+        Step 2 of JPCMCIplus, i.e. discovery of links between observed (time and space) dummy nodes and system nodes
+        through an application of the skeleton phase of PCMCIplus to this subset of nodes (dummy nodes and
+        system nodes).
+        See run_jpcmciplus for a description of the parameters.
+
+        Parameters
+        ----------
+        context_system_results : dictionary
+            Output of discover_lagged_and_context_system_links, i.e. lagged and context parents together with the
+            corresponding estimated test statistic values regarding adjacencies.
+        lagged_parents : dictionary
+            Dictionary of form {0:[(0, -1), (3, -2), ...], 1:[], ...} containing the conditioning-parents
+            estimated with PC algorithm.
+
+        Returns
+        -------
+        graph : array of shape [N, N, tau_max+1]
+            Resulting causal graph, see description above for interpretation.
+        val_matrix : array of shape [N, N, tau_max+1]
+            Estimated matrix of test statistic values regarding adjacencies.
+        p_matrix : array of shape [N, N, tau_max+1]
+            Estimated matrix of p-values regarding adjacencies.
+        parents : dictionary
+            Dictionary of form {0:[(0, -1), (3, -2), ...], 1:[], ...} containing
+            the estimated dummy parents of the system nodes.
+        """
+        lagged_context_parents = {i: list(dict.fromkeys(context_system_results['parents'][i] + lagged_parents[i])) for
+                                  i in range(self.N)}
+        dummy_parents = {i: [] for i in range(self.N)}
+        p_matrix = context_system_results['p_matrix']
+
+        # setup link assumptions without the observed context nodes
+        _link_assumptions_dummy = self.add_found_context_link_assumptions(link_assumptions, tau_max)
+        _int_link_assumptions = self._set_link_assumptions(_link_assumptions_dummy, tau_min, tau_max)
+
+        self.mode = "dummy_search"
+        if self.verbosity > 0:
+            print("\n##\n## J-PCMCI+ Step 3: Discovering dummy-system links\n##")
+            if _link_assumptions_dummy is not None:
+                print("\nWith link_assumptions = %s" % str(_int_link_assumptions))
+
+        skeleton_results_dummy = self._pcmciplus_mci_skeleton_phase(
+            lagged_context_parents, _int_link_assumptions, pc_alpha,
+            tau_min, tau_max, max_conds_dim, None,
+            max_conds_py, max_conds_px, max_conds_px_lagged,
+            reset_lagged_links, fdr_method,
+            self.p_matrix, self.val_matrix
+        )
+
+        skeleton_graph_dummy = skeleton_results_dummy['graph']
+
+        for j in self.system_nodes:
+            for k in range(tau_max + 1):
+                for dummy_node in self.time_dummy + self.space_dummy:
+                    if skeleton_graph_dummy[dummy_node, j, k] == 'o?o' or \
+                            skeleton_graph_dummy[dummy_node, j, k] == '-?>' or \
+                            skeleton_graph_dummy[dummy_node, j, k] == 'o-o' or \
+                            skeleton_graph_dummy[dummy_node, j, k] == '-->':
+                        dummy_parents[j].append((dummy_node, k))
+                for context_node in self.time_context_nodes + self.space_context_nodes:
+                    skeleton_results_dummy['val_matrix'][context_node, j, k] = context_system_results['val_matrix'][
+                        context_node, j, k]
+                    skeleton_results_dummy['val_matrix'][j, context_node, k] = context_system_results['val_matrix'][
+                        j, context_node, k]
+
+                    skeleton_results_dummy['p_matrix'][context_node, j, k] = p_matrix[context_node, j, k]
+                    skeleton_results_dummy['p_matrix'][j, context_node, k] = p_matrix[j, context_node, k]
+
+        return_dict = {'graph': skeleton_results_dummy['graph'], 'p_matrix': skeleton_results_dummy['p_matrix'],
+                       'val_matrix': skeleton_results_dummy['val_matrix'], 'parents': dummy_parents}
+
+        # Print the results
+        if self.verbosity > 0:
+            self.print_results(return_dict, alpha_level=pc_alpha)
+        return return_dict
+
+    def discover_system_system_links(self, link_assumptions,
+                                     lagged_context_dummy_parents,
+                                     tau_min=0,
+                                     tau_max=1,
+                                     pc_alpha=0.01,
+                                     reset_lagged_links=False,
+                                     max_conds_dim=None,
+                                     max_conds_py=None,
+                                     max_conds_px=None,
+                                     max_conds_px_lagged=None,
+                                     fdr_method='none'
+                                     ):
+        """
+        Step 4 of JPCMCIplus and orientation phase, i.e. discovery of links between system nodes given the knowledge
+        about their context parents through an application of PCMCIplus to this subset of nodes (system nodes).
+        See run_jpcmciplus for a description of the other parameters.
+
+
+        Parameters
+        ----------
+        lagged_context_dummy_parents : dictionary
+            Dictionary containing lagged and (dummy and observed) context parents of the system nodes estimated during
+            step 1 and step 2 of J-PCMCI+.
+
+        Returns
+        -------
+        graph : array of shape [N, N, tau_max+1]
+            Resulting causal graph, see description above for interpretation.
+        val_matrix : array of shape [N, N, tau_max+1]
+            Estimated matrix of test statistic values regarding adjacencies.
+        p_matrix : array of shape [N, N, tau_max+1]
+            Estimated matrix of p-values regarding adjacencies.
+        sepset : dictionary
+            Separating sets. See paper for details.
+        ambiguous_triples : list
+            List of ambiguous triples, only relevant for 'majority' and
+            'conservative' rules, see paper for details.
+        """
+        self.mode = "system_search"
+
+        # Get the parents from run_pc_stable only on the system links
+        system_links = self.clean_system_link_assumptions(link_assumptions, tau_max)
+        # Set the selected links
+        _int_link_assumptions = self._set_link_assumptions(system_links, tau_min, tau_max)
+
+        if self.verbosity > 0:
+            print("\n##\n## J-PCMCI+ Step 4: Discovering system-system links \n##")
+            if system_links is not None:
+                print("\nWith link_assumptions = %s" % str(_int_link_assumptions))
+
+        skeleton_results = self._pcmciplus_mci_skeleton_phase(
+            lagged_context_dummy_parents, _int_link_assumptions, pc_alpha,
+            tau_min, tau_max, max_conds_dim, None,
+            max_conds_py, max_conds_px, max_conds_px_lagged,
+            reset_lagged_links, fdr_method,
+            self.p_matrix, self.val_matrix
+        )
+
+        return skeleton_results
+
+    def _remaining_pairs(self, graph, adjt, tau_min, tau_max, p):
+        """Helper function returning the remaining pairs that still need to be
+        tested depending on the JPCMCIplus step, i.e. discovery of context-system links (step 1),
+        dummy-context links (step 2) or system-system links in which case the function of the parent class is called.
+        """
+        all_context_nodes = self.time_context_nodes + self.space_context_nodes
+        if self.mode == "context_search":
+            # during discovery of context-system links we are only
+            # interested in context-context and context-system pairs
+            N = graph.shape[0]
+            pairs = []
+            for (i, j) in itertools.product(range(N), range(N)):
+                for abstau in range(tau_min, tau_max + 1):
+                    if (graph[i, j, abstau] != ""
+                            and len(
+                                [a for a in adjt[j] if a != (i, -abstau)]) >= p
+                            and i in all_context_nodes):
+                        pairs.append((i, j, abstau))
+            return pairs
+        elif self.mode == "dummy_search":
+            # during discovery of dummy-system links we are only
+            # interested in dummy-system pairs
+            N = graph.shape[0]
+            pairs = []
+            for (i, j) in itertools.product(range(N), range(N)):
+                for abstau in range(tau_min, tau_max + 1):
+                    if (graph[i, j, abstau] != ""
+                            and len(
+                                [a for a in adjt[j] if a != (i, -abstau)]) >= p
+                            and i in self.time_dummy + self.space_dummy):
+                        pairs.append((i, j, abstau))
+            return pairs
+        else:
+            return super()._remaining_pairs(graph, adjt, tau_min, tau_max, p)
+
+    def _run_pcalg_test(self, graph, i, abstau, j, S, lagged_parents, max_conds_py,
+                        max_conds_px, max_conds_px_lagged, tau_max, alpha_or_thres=None):
+        """MCI conditional independence tests within PCMCIplus or PC algorithm. Depending on the JPCMCIplus step
+        the setup is adapted slightly. During the discovery of dummy-system links (step 2) we are using
+        the dummy_ci_test and condition on the parents found during step 1; during the discovery of system-system links
+        (step 3) we are conditioning on the found contextual parents.
+        Parameters
+        ----------
+        graph : array
+            ...
+        i : int
+            Variable index.
+        abstau : int
+            Time lag (absolute value).
+        j : int
+            Variable index.
+        S : list
+            List of contemporaneous conditions.
+        lagged_parents : dictionary of lists
+            Dictionary of lagged parents for each node.
+        max_conds_py : int
+            Max number of lagged parents for node j.
+        max_conds_px : int
+            Max number of lagged parents for lagged node i.
+        max_conds_px_lagged : int
+            Maximum number of lagged conditions of X when X is lagged in MCI
+            tests. If None is passed, this number is equal to max_conds_px.
+        tau_max : int
+            Maximum time lag.
+        Returns
+        -------
+        val : float
+            Test statistic value.
+        pval : float
+            Test statistic p-value.
+        Z : list
+            List of conditions.
+        """
+
+        if self.mode == 'dummy_search':
+            # during discovery of dummy-system links we condition on the found contextual parents from step 1.
+            cond = list(S) + self.observed_context_parents[j]
+            cond = list(dict.fromkeys(cond))  # remove overlapps
+            return super()._run_pcalg_test(graph, i, abstau, j, cond, lagged_parents, max_conds_py,
+                                           max_conds_px, max_conds_px_lagged, tau_max, alpha_or_thres)
+
+        elif self.mode == 'system_search':
+            # during discovery of system-system links we are conditioning on the found contextual parents
+            cond = list(S) + self.dummy_parents[j] + self.observed_context_parents[j]
+            cond = list(dict.fromkeys(cond))  # remove overlapps
+            return super()._run_pcalg_test(graph, i, abstau, j, cond, lagged_parents, max_conds_py,
+                                           max_conds_px, max_conds_px_lagged, tau_max, alpha_or_thres)
+        else:
+            return super()._run_pcalg_test(graph, i, abstau, j, S, lagged_parents, max_conds_py,
+                                           max_conds_px, max_conds_px_lagged, tau_max, alpha_or_thres)
+
+if __name__ == '__main__':
+    # Imports
+    from numpy.random import SeedSequence, default_rng
+
+    import tigramite
+    from tigramite.toymodels import structural_causal_processes as toys
+    from tigramite.toymodels.context_model import ContextModel
+    # from tigramite.jpcmciplus import JPCMCIplus
+    from tigramite.independence_tests.parcorr_mult import ParCorrMult
+    import tigramite.data_processing as pp
+    import tigramite.plotting as tp
+
+
+    # Set seeds for reproducibility
+    ss = SeedSequence(12345)
+    child_seeds = ss.spawn(2)
+
+    model_seed = child_seeds[0]
+    context_seed = child_seeds[1]
+
+    random_state = np.random.default_rng(model_seed)
+
+    # Choose the time series length and number of spatial contexts
+    T = 100
+    nb_domains = 50
+
+    transient_fraction=0.2
+    tau_max = 2
+    frac_observed = 0.5
+
+    # Specify the model
+    def lin(x): return x
+
+    links = {0: [((0, -1), 0.3, lin), ((3, -1), 0.7, lin), ((4, 0), 0.9, lin)],
+             1: [((1, -1), 0.4, lin), ((3, -1), 0.8, lin)],
+             2: [((2, -1), 0.3, lin), ((1, 0), -0.5, lin), ((4, 0), 0.5, lin), ((5, 0), 0.6, lin)] ,
+             3: [], 
+             4: [], 
+             5: []
+                }
+
+    # Specify which node is a context node via node_type (can be "system", "time_context", or "space_context")
+    node_classification = {
+        0: "system",
+        1: "system",
+        2: "system",
+        3: "time_context",
+        4: "time_context",
+        5: "space_context"
+    }
+
+    # Specify dynamical noise term distributions, here unit variance Gaussians
+    #random_state = np.random.RandomState(seed)
+    noises = [random_state.standard_normal for j in range(6)]
+
+    contextmodel = ContextModel(links=links, node_classification=node_classification,
+                                noises=noises, 
+                                seed=context_seed)
+
+    data_ens, nonstationary = contextmodel.generate_data(nb_domains, T)
+
+    assert not nonstationary
+
+    system_indices = [0,1,2]
+    # decide which context variables should be latent, and which are observed
+    observed_indices_time = [4]
+    latent_indices_time = [3]
+
+    observed_indices_space = [5]
+    latent_indices_space = []
+
+    # all system variables are also observed, thus we get the following observed data
+    observed_indices = system_indices + observed_indices_time + observed_indices_space
+    data_observed = {key: data_ens[key][:,observed_indices] for key in data_ens}
+
+
+    # Add one-hot-encoding of time-steps and dataset index to the observational data. 
+    # These are the values of the time and space dummy variables.
+    dummy_data_time = np.identity(T)
+
+    data_dict = {}
+    for i in range(nb_domains):
+        dummy_data_space = np.zeros((T, nb_domains))
+        dummy_data_space[:, i] = 1.
+        data_dict[i] = np.hstack((data_observed[i], dummy_data_time, dummy_data_space))
+
+    # Define vector-valued variables including dummy variables as well as observed (system and context) variables
+    nb_observed_context_nodes = len(observed_indices_time) + len(observed_indices_space)
+    N = len(system_indices)
+    process_vars = system_indices
+    observed_temporal_context_nodes = list(range(N, N + len(observed_indices_time)))
+    observed_spatial_context_nodes = list(range(N + len(observed_indices_time), 
+                                                N + len(observed_indices_time) + len(observed_indices_space)))
+    time_dummy_index = N + nb_observed_context_nodes
+    space_dummy_index = N + nb_observed_context_nodes + 1
+    time_dummy = list(range(time_dummy_index, time_dummy_index + T))
+    space_dummy = list(range(time_dummy_index + T, time_dummy_index + T + nb_domains))
+
+    vector_vars = {i: [(i, 0)] for i in process_vars + observed_temporal_context_nodes + observed_spatial_context_nodes}
+    vector_vars[time_dummy_index] = [(i, 0) for i in time_dummy]
+    vector_vars[space_dummy_index] = [(i, 0) for i in space_dummy]
+
+    # Name all the variables and initialize the dataframe object
+    # Be careful to use analysis_mode = 'multiple'
+    sys_var_names = ['X_' + str(i) for i in process_vars]
+    context_var_names = ['t-C_'+str(i) for i in observed_indices_time] + ['s-C_'+str(i) for i in observed_indices_space]
+    var_names = sys_var_names + context_var_names + ['t-dummy', 's-dummy']
+
+    dataframe = pp.DataFrame(
+        data=data_dict,
+        vector_vars = vector_vars,
+        analysis_mode = 'multiple',
+        var_names = var_names
+        )
+
+
+    # Classify all the nodes into system, context, or dummy
+    node_classification_jpcmci = {i: node_classification[var] for i, var in enumerate(observed_indices)}
+    node_classification_jpcmci.update({time_dummy_index : "time_dummy", space_dummy_index : "space_dummy"})
+
+    # Create a J-PCMCI+ object, passing the dataframe and (conditional)
+    # independence test objects, as well as the observed temporal and spatial context nodes 
+    # and the indices of the dummies.
+    JPCMCIplus = JPCMCIplus(dataframe=dataframe,
+                              cond_ind_test=ParCorrMult(significance='analytic'), 
+                              node_classification=node_classification_jpcmci,
+                              verbosity=1,)
+
+    # Define the analysis parameters.
+    tau_max = 2
+    pc_alpha = 0.01
+
+    # Run J-PCMCI+
+    results = JPCMCIplus.run_jpcmciplus(tau_min=0, 
+                                  tau_max=tau_max, 
+                                  pc_alpha=pc_alpha)