From f316497545978e819113c55f0b702ebe9dafb2c5 Mon Sep 17 00:00:00 2001
From: Victor Li <wentaol4@illinois.edu>
Date: Thu, 10 Sep 2026 22:14:16 -0500
Subject: [PATCH] Implement parallel network discovery

---
 causationentropy/core/discovery.py       | 264 ++++++++++++++++-------
 causationentropy/tests/test_discovery.py | 111 +++++++++-
 pyproject.toml                           |   1 +
 3 files changed, 292 insertions(+), 84 deletions(-)

diff --git a/causationentropy/core/discovery.py b/causationentropy/core/discovery.py
index ae5d442..b9402cd 100644
--- a/causationentropy/core/discovery.py
+++ b/causationentropy/core/discovery.py
@@ -3,18 +3,164 @@
 Email: kslote@clarkson.edu
 version = 1.1.0
 """
+
 import copy
 from typing import Dict, Tuple, Union
 
 import networkx as nx
 import numpy as np
 import pandas as pd
+from joblib import Parallel, delayed, effective_n_jobs
 from sklearn.linear_model import Lasso, LassoLarsIC
 
 from causationentropy.core.information.conditional_mutual_information import (
     conditional_mutual_information,
 )
 
+# Default seed for permutation tests. Independent streams are spawned per
+# target so serial and parallel execution consume the same randomness.
+_DISCOVER_NETWORK_SEED = 42
+
+
+def _validate_n_jobs(n_jobs):
+    """Normalize ``n_jobs`` to joblib's integer convention.
+
+    ``None`` means serial (1 job). ``-1`` means all processors. ``0`` is
+    invalid. Booleans are rejected because ``bool`` is a subclass of ``int``.
+    """
+    if n_jobs is None:
+        return 1
+    if isinstance(n_jobs, bool) or not isinstance(n_jobs, (int, np.integer)):
+        raise ValueError(f"n_jobs must be an integer, got {type(n_jobs).__name__}")
+    n_jobs = int(n_jobs)
+    if n_jobs == 0:
+        raise ValueError(
+            "n_jobs=0 is invalid; use n_jobs=1 for serial execution "
+            "or n_jobs=-1 for all processors"
+        )
+    return n_jobs
+
+
+def _resolve_n_workers(n_jobs, n_targets):
+    """Worker count: requested jobs, but never more than the number of targets."""
+    n_jobs = _validate_n_jobs(n_jobs)
+    return max(1, min(effective_n_jobs(n_jobs), int(n_targets)))
+
+
+def _run_target_jobs(job_args, n_workers):
+    """Run per-target inference serially or with joblib.
+
+    Both paths call the same ``_discover_edges_for_target`` helper so the
+    oCSE mathematics is not duplicated.
+    """
+    if n_workers == 1:
+        return [_discover_edges_for_target(*args) for args in job_args]
+    # loky isolates BLAS via inner_max_num_threads so process workers do not
+    # oversubscribe CPU against NumPy's Accelerate/OpenMP threads.
+    return Parallel(n_jobs=n_workers, backend="loky", inner_max_num_threads=1)(
+        delayed(_discover_edges_for_target)(*args) for args in job_args
+    )
+
+
+def _discover_edges_for_target(
+    i,
+    rng_seed,
+    series,
+    X_lagged,
+    Y_all,
+    feature_names,
+    method,
+    max_lag,
+    T,
+    alpha_forward,
+    alpha_backward,
+    n_shuffles,
+    information,
+    metric,
+    k_means,
+    bandwidth,
+    var_name,
+):
+    """Infer parents of one target variable and return edge tuples.
+
+    Workers never touch the shared NetworkX graph. Each target gets its own
+    RNG from ``rng_seed`` so jobs do not share or duplicate a random stream.
+    """
+    rng = np.random.default_rng(rng_seed)
+    print(f"Estimating edges for node {i} ({var_name})")
+
+    Y = Y_all[:, [i]]  # shape: (T - max_lag, 1)
+    if method == "standard":
+        Z_init = []
+        for tau in range(1, max_lag + 1):
+            Z_init.append(series[max_lag - tau : T - tau, i])  # lagged Y_i
+        Z_init = np.column_stack(Z_init)  # shape: (T - max_lag, max_lag)
+        S = standard_optimal_causation_entropy(
+            X_lagged,
+            Y,
+            Z_init,
+            rng,
+            alpha_forward,
+            alpha_backward,
+            n_shuffles,
+            information,
+            metric,
+            k_means,
+            bandwidth,
+        )
+    if method == "alternative":
+        S = alternative_optimal_causation_entropy(
+            X_lagged,
+            Y,
+            rng,
+            alpha_forward,
+            alpha_backward,
+            n_shuffles,
+            information,
+            metric,
+            k_means,
+            bandwidth,
+        )
+    if method == "information_lasso":
+        S = information_lasso_optimal_causation_entropy(X_lagged, Y, rng)
+    if method == "lasso":
+        S = lasso_optimal_causation_entropy(X_lagged, Y, rng)
+
+    edges = []
+    for s in S:
+        src_var, src_lag = feature_names[s]
+
+        X_predictor = X_lagged[:, [s]]
+        Y_target = Y
+        other_selected = [idx for idx in S if idx != s]
+        Z_cond = X_lagged[:, other_selected] if other_selected else None
+
+        cmi = conditional_mutual_information(
+            X_predictor,
+            Y_target,
+            Z_cond,
+            method=information,
+            metric=metric,
+            k=k_means,
+            bandwidth=bandwidth,
+        )
+
+        test_result = shuffle_test(
+            X_predictor,
+            Y_target,
+            Z_cond,
+            cmi,
+            alpha=alpha_backward,
+            rng=rng,
+            n_shuffles=n_shuffles,
+            information=information,
+            metric=metric,
+            k_means=k_means,
+            bandwidth=bandwidth,
+        )
+        edges.append((src_var, src_lag, cmi, test_result["P_value"]))
+    return i, edges
+
 
 def discover_network(
     data: Union[np.ndarray, pd.DataFrame],
@@ -97,7 +243,10 @@ def discover_network(
         Number of permutations for statistical significance testing. Higher values
         provide more accurate p-value estimates but increase computational cost.
     n_jobs : int, default=-1
-        Number of parallel jobs for computation. -1 uses all available processors.
+        Number of parallel jobs. Each target variable is an independent job.
+        ``1`` runs serially in the calling process. ``-1`` uses all available
+        processors. The worker count is never larger than the number of
+        target variables. ``0`` is invalid.
 
     Returns
     -------
@@ -115,7 +264,8 @@ def discover_network(
     NotImplementedError
         If an unsupported method or information type is specified.
     ValueError
-        If the time series is too short for the chosen max_lag.
+        If the time series is too short for the chosen max_lag, or if
+        ``n_jobs`` is not a valid integer.
 
     Notes
     -----
@@ -148,8 +298,6 @@ def discover_network(
 
     .. [2] Schreiber, T. Measuring information transfer. Physical Review Letters 85, 461 (2000).
     """
-    rng = np.random.default_rng(42)
-
     if method not in ["standard", "alternative", "information_lasso", "lasso"]:
         raise NotImplementedError(f"discover_network: method={method} not supported.")
     supported_information_types = ["gaussian", "knn", "kde", "geometric_knn", "poisson"]
@@ -158,6 +306,7 @@ def discover_network(
             f"discover_network: information={information} not supported. "
             f"Supported types: {supported_information_types}"
         )
+    n_jobs = _validate_n_jobs(n_jobs)
 
     # Convert DataFrame to ndarray while keeping column labels
     if isinstance(data, pd.DataFrame):
@@ -188,89 +337,44 @@ def discover_network(
     G = nx.MultiDiGraph()
     G.add_nodes_from(var_names)
 
-    # Step 3: Loop over each variable and infer parents from lagged predictors
-    for i in range(n):
-        print(f"Estimating edges for node {i} ({var_names[i]})")
-
-        Y = Y_all[:, [i]]  # shape: (T - max_lag, 1)
-        if method == "standard":
-            Z_init = []
-            for tau in range(1, max_lag + 1):
-                Z_init.append(series[max_lag - tau : T - tau, i])  # lagged Y_i
-            Z_init = np.column_stack(Z_init)  # shape: (T - max_lag, max_lag)
-            S = standard_optimal_causation_entropy(
-                X_lagged,
-                Y,
-                Z_init,
-                rng,
-                alpha_forward,
-                alpha_backward,
-                n_shuffles,
-                information,
-                metric,
-                k_means,
-                bandwidth,
-            )
-        if method == "alternative":
-            S = alternative_optimal_causation_entropy(
-                X_lagged,
-                Y,
-                rng,
-                alpha_forward,
-                alpha_backward,
-                n_shuffles,
-                information,
-                metric,
-                k_means,
-                bandwidth,
-            )
-        if method == "information_lasso":
-            S = information_lasso_optimal_causation_entropy(X_lagged, Y, rng)
-        if method == "lasso":
-            S = lasso_optimal_causation_entropy(X_lagged, Y, rng)
-        for s in S:
-            src_var, src_lag = feature_names[s]
-
-            # Compute CMI and p-value for this edge
-            X_predictor = X_lagged[:, [s]]  # predictor at this lag
-            Y_target = Y  # target variable
-
-            # Conditioning set: all other selected predictors for this target
-            other_selected = [idx for idx in S if idx != s]
-            Z_cond = X_lagged[:, other_selected] if other_selected else None
-
-            # Compute conditional mutual information
-            cmi = conditional_mutual_information(
-                X_predictor,
-                Y_target,
-                Z_cond,
-                method=information,
-                metric=metric,
-                k=k_means,
-                bandwidth=bandwidth,
-            )
+    # Independent RNG per target so n_jobs=1 and n_jobs>1 match, and so
+    # workers never share or clone a single Generator stream.
+    target_seeds = np.random.SeedSequence(_DISCOVER_NETWORK_SEED).spawn(n)
+    job_args = [
+        (
+            i,
+            target_seeds[i],
+            series,
+            X_lagged,
+            Y_all,
+            feature_names,
+            method,
+            max_lag,
+            T,
+            alpha_forward,
+            alpha_backward,
+            n_shuffles,
+            information,
+            metric,
+            k_means,
+            bandwidth,
+            var_names[i],
+        )
+        for i in range(n)
+    ]
 
-            # Compute p-value using shuffle test
-            test_result = shuffle_test(
-                X_predictor,
-                Y_target,
-                Z_cond,
-                cmi,
-                alpha=alpha_backward,  # Use backward elimination alpha
-                rng=rng,
-                n_shuffles=n_shuffles,
-                information=information,
-                metric=metric,
-                k_means=k_means,
-                bandwidth=bandwidth,
-            )
+    # Workers return edge lists; only this process mutates G.
+    n_workers = _resolve_n_workers(n_jobs, n)
+    results = _run_target_jobs(job_args, n_workers)
 
+    for i, edges in sorted(results, key=lambda item: item[0]):
+        for src_var, src_lag, cmi, p_value in edges:
             G.add_edge(
                 var_names[src_var],
                 var_names[i],
                 lag=src_lag,
                 cmi=cmi,
-                p_value=test_result["P_value"],
+                p_value=p_value,
             )
 
     return G
diff --git a/causationentropy/tests/test_discovery.py b/causationentropy/tests/test_discovery.py
index d681a99..2005031 100644
--- a/causationentropy/tests/test_discovery.py
+++ b/causationentropy/tests/test_discovery.py
@@ -209,7 +209,7 @@ def test_discover_network_cmi_integration(self, mock_cmi):
         mock_cmi.return_value = 0.5  # Mock CMI value
 
         data = np.random.normal(0, 1, (20, 2))
-        G = discover_network(data, max_lag=1, n_shuffles=10)
+        G = discover_network(data, max_lag=1, n_shuffles=10, n_jobs=1)
 
         # Verify CMI was called
         assert mock_cmi.called
@@ -239,7 +239,12 @@ def test_parameter_passing_metric(self, mock_cmi):
         # Test with different metric values
         for metric in ["euclidean", "cityblock", "chebyshev"]:
             discover_network(
-                data, information="knn", metric=metric, max_lag=1, n_shuffles=5
+                data,
+                information="knn",
+                metric=metric,
+                max_lag=1,
+                n_shuffles=5,
+                n_jobs=1,
             )
 
             # Verify that the metric parameter was passed to conditional_mutual_information
@@ -256,7 +261,12 @@ def test_parameter_passing_bandwidth(self, mock_cmi):
         # Test with different bandwidth values
         for bandwidth in ["silverman", "scott", 0.5]:
             discover_network(
-                data, information="kde", bandwidth=bandwidth, max_lag=1, n_shuffles=5
+                data,
+                information="kde",
+                bandwidth=bandwidth,
+                max_lag=1,
+                n_shuffles=5,
+                n_jobs=1,
             )
 
             # Verify that the bandwidth parameter was passed to conditional_mutual_information
@@ -273,7 +283,12 @@ def test_parameter_passing_k_means(self, mock_cmi):
         # Test with different k_means values
         for k_means in [1, 3, 5, 10]:
             discover_network(
-                data, information="knn", k_means=k_means, max_lag=1, n_shuffles=5
+                data,
+                information="knn",
+                k_means=k_means,
+                max_lag=1,
+                n_shuffles=5,
+                n_jobs=1,
             )
 
             # Verify that the k_means parameter was passed as 'k' to conditional_mutual_information
@@ -300,6 +315,7 @@ def test_parameter_passing_all_three(self, mock_cmi):
             k_means=k_means,
             max_lag=1,
             n_shuffles=5,
+            n_jobs=1,
         )
 
         # Verify all parameters were passed correctly
@@ -330,6 +346,7 @@ def test_parameter_passing_different_methods(self, mock_cmi):
                 k_means=k_means,
                 max_lag=1,
                 n_shuffles=5,
+                n_jobs=1,
             )
 
             # Verify parameters were passed through for both methods
@@ -496,6 +513,92 @@ def test_pvalue_calculation_correctness(self):
             pytest.skip("No X0->X1 edges found - test setup may need adjustment")
 
 
+def _comparable_edges(G):
+    """Stable edge signature for serial vs parallel comparison."""
+    return sorted(
+        (u, v, d["lag"], float(d["cmi"]), float(d["p_value"]))
+        for u, v, d in G.edges(data=True)
+    )
+
+
+class TestDiscoverNetworkNJobs:
+    """Tests that n_jobs actually controls parallel execution."""
+
+    @pytest.fixture
+    def causal_data(self):
+        rng = np.random.default_rng(0)
+        T, n = 80, 4
+        data = rng.normal(size=(T, n))
+        for t in range(1, T):
+            data[t, 1] += 0.8 * data[t - 1, 0]
+            data[t, 2] += 0.8 * data[t - 1, 1]
+        return data
+
+    def test_n_jobs_1_runs(self, causal_data):
+        G = discover_network(causal_data, max_lag=1, n_shuffles=20, n_jobs=1)
+        assert isinstance(G, nx.MultiDiGraph)
+        assert len(G.nodes()) == 4
+
+    def test_n_jobs_2_runs(self, causal_data):
+        G = discover_network(causal_data, max_lag=1, n_shuffles=20, n_jobs=2)
+        assert isinstance(G, nx.MultiDiGraph)
+        assert len(G.nodes()) == 4
+
+    def test_n_jobs_4_runs(self, causal_data):
+        G = discover_network(causal_data, max_lag=1, n_shuffles=20, n_jobs=4)
+        assert isinstance(G, nx.MultiDiGraph)
+        assert len(G.nodes()) == 4
+
+    def test_n_jobs_minus_one_runs(self, causal_data):
+        G = discover_network(causal_data, max_lag=1, n_shuffles=20, n_jobs=-1)
+        assert isinstance(G, nx.MultiDiGraph)
+        assert len(G.nodes()) == 4
+
+    def test_serial_and_parallel_equivalent(self, causal_data):
+        kwargs = dict(
+            max_lag=1,
+            n_shuffles=30,
+            information="gaussian",
+            method="standard",
+        )
+        G1 = discover_network(causal_data, n_jobs=1, **kwargs)
+        G2 = discover_network(causal_data, n_jobs=2, **kwargs)
+        G4 = discover_network(causal_data, n_jobs=4, **kwargs)
+        G_all = discover_network(causal_data, n_jobs=-1, **kwargs)
+
+        assert (
+            set(G1.nodes()) == set(G2.nodes()) == set(G4.nodes()) == set(G_all.nodes())
+        )
+        assert _comparable_edges(G1) == _comparable_edges(G2)
+        assert _comparable_edges(G1) == _comparable_edges(G4)
+        assert _comparable_edges(G1) == _comparable_edges(G_all)
+
+    def test_n_jobs_zero_invalid(self):
+        data = np.random.normal(0, 1, (20, 3))
+        with pytest.raises(ValueError, match="n_jobs=0 is invalid"):
+            discover_network(data, n_jobs=0)
+
+    def test_n_jobs_non_integer_invalid(self):
+        data = np.random.normal(0, 1, (20, 3))
+        with pytest.raises(ValueError, match="n_jobs must be an integer"):
+            discover_network(data, n_jobs=1.5)
+        with pytest.raises(ValueError, match="n_jobs must be an integer"):
+            discover_network(data, n_jobs=True)
+
+
+def test_resolve_n_workers_caps_at_target_count():
+    """n_jobs=-1 / oversized n_jobs must not exceed the number of targets."""
+    from joblib import effective_n_jobs
+
+    from causationentropy.core.discovery import _resolve_n_workers
+
+    assert _resolve_n_workers(1, 8) == 1
+    assert _resolve_n_workers(2, 8) == 2
+    assert _resolve_n_workers(100, 3) == 3
+    assert _resolve_n_workers(-1, 2) == 2
+    assert _resolve_n_workers(-1, 100) == effective_n_jobs(-1)
+
+
 class TestLassoOptimalCausationEntropy:
     """Test LASSO-based variable selection for causal discovery."""
 
diff --git a/pyproject.toml b/pyproject.toml
index 3281637..7d339f7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -60,6 +60,7 @@ dependencies = [
     "scipy>=1.7.0",
     "scikit-learn>=1.0.0",
     "networkx>=2.6.0",
+    "joblib>=1.3.0",
     "matplotlib>=3.4.0",
     "tqdm>=4.60.0",
     "tigramite>=5.0.0",