Add causal_selection/discovery/algorithms.py

causal_selection/discovery/algorithms.py

@@ -0,0 +1,395 @@
+"""
+Algorithm adapters: run each causal discovery algorithm with timeout handling.
+All algorithms take a pandas DataFrame (integer-encoded discrete data) and return
+an adjacency matrix (np.ndarray) representing the learned graph.
+"""
+import numpy as np
+import pandas as pd
+import signal
+import time
+import traceback
+import logging
+from functools import wraps
+
+logger = logging.getLogger(__name__)
+
+
+class TimeoutError(Exception):
+    pass
+
+
+def timeout_handler(signum, frame):
+    raise TimeoutError("Algorithm timed out")
+
+
+def run_with_timeout(func, timeout_sec, *args, **kwargs):
+    """Run a function with a timeout (Unix only, uses SIGALRM)."""
+    old_handler = signal.signal(signal.SIGALRM, timeout_handler)
+    signal.alarm(timeout_sec)
+    try:
+        result = func(*args, **kwargs)
+        signal.alarm(0)
+        return result
+    except TimeoutError:
+        logger.warning(f"Timeout after {timeout_sec}s")
+        raise
+    finally:
+        signal.signal(signal.SIGALRM, old_handler)
+        signal.alarm(0)
+
+
+def safe_run(algo_name, func, timeout_sec, *args, **kwargs):
+    """Run algorithm with timeout and exception handling.
+    
+    Returns:
+        (adjmat, runtime, status): adjacency matrix, time in seconds, status string
+    """
+    start = time.time()
+    try:
+        adjmat = run_with_timeout(func, timeout_sec, *args, **kwargs)
+        runtime = time.time() - start
+        return adjmat, runtime, 'success'
+    except TimeoutError:
+        runtime = time.time() - start
+        logger.warning(f"{algo_name}: TIMEOUT after {runtime:.1f}s")
+        return None, runtime, 'timeout'
+    except Exception as e:
+        runtime = time.time() - start
+        logger.error(f"{algo_name}: ERROR after {runtime:.1f}s: {e}")
+        logger.debug(traceback.format_exc())
+        return None, runtime, f'error: {str(e)[:100]}'
+
+
+# ============================================================
+# CONSTRAINT-BASED ALGORITHMS (causal-learn)
+# ============================================================
+
+def run_pc_discrete(df, alpha=0.01, stable=True):
+    """PC algorithm for discrete data using G-squared test."""
+    from causallearn.search.ConstraintBased.PC import pc
+    from causallearn.utils.cit import gsq
+    
+    data = df.values.astype(int)
+    cg = pc(data, alpha=alpha, indep_test=gsq, stable=stable, 
+            show_progress=False)
+    
+    # Extract adjacency matrix from GeneralGraph
+    adj = cg.G.graph  # numpy array
+    return _causallearn_graph_to_adjmat(adj)
+
+
+def run_fci(df, alpha=0.05, depth=4):
+    """FCI algorithm - outputs PAG. We extract the directed edges."""
+    from causallearn.search.ConstraintBased.FCI import fci
+    from causallearn.utils.cit import gsq
+    
+    data = df.values.astype(int)
+    g, edges = fci(data, independence_test_method=gsq, alpha=alpha,
+                   depth=depth, show_progress=False)
+    
+    adj = g.graph
+    return _causallearn_pag_to_adjmat(adj)
+
+
+# ============================================================
+# SCORE-BASED ALGORITHMS (causal-learn)
+# ============================================================
+
+def run_ges_causallearn(df, score_func='local_score_BDeu'):
+    """GES algorithm from causal-learn."""
+    from causallearn.search.ScoreBased.GES import ges
+    
+    data = df.values.astype(int)
+    record = ges(data, score_func=score_func, maxP=None,
+                 parameters=None)
+    
+    adj = record['G'].graph
+    return _causallearn_graph_to_adjmat(adj)
+
+
+def run_boss(df, score_func='local_score_BDeu'):
+    """BOSS (Best Order Score Search) algorithm."""
+    from causallearn.search.PermutationBased.BOSS import boss
+    
+    data = df.values.astype(int)
+    cg = boss(data, score_func=score_func, 
+              parameters=None)
+    
+    adj = cg.graph
+    return _causallearn_graph_to_adjmat(adj)
+
+
+def run_grasp(df, score_func='local_score_BDeu', depth=3):
+    """GRaSP (Greedy relaxation of Sparsest Permutation) algorithm."""
+    from causallearn.search.PermutationBased.GRaSP import grasp
+    
+    data = df.values.astype(int)
+    cg = grasp(data, score_func=score_func, depth=depth,
+               parameters=None)
+    
+    adj = cg.graph
+    return _causallearn_graph_to_adjmat(adj)
+
+
+# ============================================================
+# SCORE-BASED ALGORITHMS (pgmpy)
+# ============================================================
+
+def run_ges_pgmpy(df, scoring_method='bicscore'):
+    """GES algorithm from pgmpy."""
+    from pgmpy.estimators import ExhaustiveSearch, HillClimbSearch, BicScore, BDeuScore
+    
+    # Use HillClimbSearch as proxy since pgmpy doesn't have native GES
+    # We'll use the tabu search for better exploration
+    scoring = BicScore(df) if scoring_method == 'bicscore' else BDeuScore(df)
+    hc = HillClimbSearch(df)
+    best_model = hc.estimate(scoring_method=scoring, max_indegree=4,
+                              max_iter=100000, epsilon=0.0001)
+    
+    return _pgmpy_model_to_adjmat(best_model, sorted(df.columns))
+
+
+def run_hc(df, scoring_method='bicscore', max_indegree=3, max_iter=100000):
+    """Hill-Climbing algorithm."""
+    from pgmpy.estimators import HillClimbSearch, BicScore, BDeuScore
+    
+    scoring = BicScore(df) if scoring_method == 'bicscore' else BDeuScore(df)
+    hc = HillClimbSearch(df)
+    best_model = hc.estimate(scoring_method=scoring, max_indegree=max_indegree,
+                              max_iter=max_iter, epsilon=0.0001)
+    
+    return _pgmpy_model_to_adjmat(best_model, sorted(df.columns))
+
+
+def run_tabu(df, scoring_method='bicscore', tabu_length=100, max_indegree=3, max_iter=100000):
+    """Tabu Search algorithm."""
+    from pgmpy.estimators import HillClimbSearch, BicScore, BDeuScore
+    
+    scoring = BicScore(df) if scoring_method == 'bicscore' else BDeuScore(df)
+    hc = HillClimbSearch(df)
+    best_model = hc.estimate(scoring_method=scoring, max_indegree=max_indegree,
+                              max_iter=max_iter, epsilon=0.0001, 
+                              tabu_length=tabu_length)
+    
+    return _pgmpy_model_to_adjmat(best_model, sorted(df.columns))
+
+
+def run_mmhc(df, scoring_method='bdeuscore', significance_level=0.01):
+    """Max-Min Hill-Climbing (MMHC) hybrid algorithm."""
+    from pgmpy.estimators import MmhcEstimator, BDeuScore, BicScore
+    
+    mmhc = MmhcEstimator(df)
+    scoring = BDeuScore(df) if scoring_method == 'bdeuscore' else BicScore(df)
+    best_model = mmhc.estimate(scoring_method=scoring, 
+                                significance_level=significance_level)
+    
+    return _pgmpy_model_to_adjmat(best_model, sorted(df.columns))
+
+
+def run_k2(df, max_parents=3):
+    """K2 algorithm (requires node ordering - we use alphabetical as default).
+    We run with multiple random orderings and take the best-scoring result.
+    """
+    from pgmpy.estimators import HillClimbSearch, K2Score
+    
+    scoring = K2Score(df)
+    hc = HillClimbSearch(df)
+    best_model = hc.estimate(scoring_method=scoring, max_indegree=max_parents,
+                              max_iter=100000, epsilon=0.0001)
+    
+    return _pgmpy_model_to_adjmat(best_model, sorted(df.columns))
+
+
+# ============================================================
+# CONVERSION UTILITIES
+# ============================================================
+
+def _causallearn_graph_to_adjmat(graph_matrix):
+    """Convert causal-learn's graph representation to standard adjacency matrix.
+    
+    causal-learn encoding:
+        graph[i,j] = -1 and graph[j,i] = 1  means i -> j
+        graph[i,j] = -1 and graph[j,i] = -1 means i -- j (undirected)
+        graph[i,j] = 1 and graph[j,i] = 1   means i <-> j (bidirected)
+    
+    Our encoding:
+        adj[i,j] = 1 and adj[j,i] = 0 means i -> j
+        adj[i,j] = 1 and adj[j,i] = 1 means i -- j (undirected)
+    """
+    n = graph_matrix.shape[0]
+    adj = np.zeros((n, n), dtype=int)
+    
+    for i in range(n):
+        for j in range(n):
+            if i == j:
+                continue
+            if graph_matrix[i, j] == -1 and graph_matrix[j, i] == 1:
+                # i -> j (tail at i, arrowhead at j)
+                adj[i, j] = 1
+            elif graph_matrix[i, j] == -1 and graph_matrix[j, i] == -1:
+                # i -- j (undirected)
+                adj[i, j] = 1
+                adj[j, i] = 1
+            elif graph_matrix[i, j] == 1 and graph_matrix[j, i] == 1:
+                # i <-> j (bidirected) - treat as undirected for CPDAG comparison
+                adj[i, j] = 1
+                adj[j, i] = 1
+    
+    return adj
+
+
+def _causallearn_pag_to_adjmat(pag_matrix):
+    """Convert PAG (from FCI) to adjacency matrix.
+    
+    PAG encoding in causal-learn:
+        1 = arrowhead (>), -1 = tail (-), 2 = circle (o)
+    
+    We extract: definite directed edges and definite adjacencies.
+    """
+    n = pag_matrix.shape[0]
+    adj = np.zeros((n, n), dtype=int)
+    
+    for i in range(n):
+        for j in range(n):
+            if i == j:
+                continue
+            # i -> j: tail at i (-1), arrowhead at j (1)
+            if pag_matrix[i, j] == -1 and pag_matrix[j, i] == 1:
+                adj[i, j] = 1
+            # i -- j or i o-o j or i o-> j: treat as undirected edge
+            elif pag_matrix[i, j] != 0 and pag_matrix[j, i] != 0:
+                adj[i, j] = 1
+                adj[j, i] = 1
+    
+    return adj
+
+
+def _pgmpy_model_to_adjmat(model, node_names):
+    """Convert pgmpy DAGModel to adjacency matrix."""
+    n = len(node_names)
+    node_idx = {name: i for i, name in enumerate(node_names)}
+    adj = np.zeros((n, n), dtype=int)
+    
+    for parent, child in model.edges():
+        if parent in node_idx and child in node_idx:
+            adj[node_idx[parent], node_idx[child]] = 1
+    
+    return adj
+
+
+# ============================================================
+# ALGORITHM REGISTRY
+# ============================================================
+
+ALGORITHM_POOL = {
+    'PC_discrete': {
+        'func': run_pc_discrete,
+        'kwargs': {'alpha': 0.01, 'stable': True},
+        'library': 'causal_learn',
+        'output_type': 'cpdag',
+        'family': 'constraint',
+    },
+    'FCI': {
+        'func': run_fci,
+        'kwargs': {'alpha': 0.05, 'depth': 4},
+        'library': 'causal_learn',
+        'output_type': 'pag',
+        'family': 'constraint',
+    },
+    'GES': {
+        'func': run_ges_causallearn,
+        'kwargs': {'score_func': 'local_score_BDeu'},
+        'library': 'causal_learn',
+        'output_type': 'cpdag',
+        'family': 'score',
+    },
+    'BOSS': {
+        'func': run_boss,
+        'kwargs': {'score_func': 'local_score_BDeu'},
+        'library': 'causal_learn',
+        'output_type': 'cpdag',
+        'family': 'permutation',
+    },
+    'GRaSP': {
+        'func': run_grasp,
+        'kwargs': {'score_func': 'local_score_BDeu', 'depth': 3},
+        'library': 'causal_learn',
+        'output_type': 'cpdag',
+        'family': 'permutation',
+    },
+    'HC': {
+        'func': run_hc,
+        'kwargs': {'scoring_method': 'bicscore', 'max_indegree': 3, 'max_iter': 100000},
+        'library': 'pgmpy',
+        'output_type': 'dag',
+        'family': 'score',
+    },
+    'Tabu': {
+        'func': run_tabu,
+        'kwargs': {'scoring_method': 'bicscore', 'tabu_length': 100, 'max_indegree': 3, 'max_iter': 100000},
+        'library': 'pgmpy',
+        'output_type': 'dag',
+        'family': 'score',
+    },
+    'MMHC': {
+        'func': run_mmhc,
+        'kwargs': {'scoring_method': 'bdeuscore', 'significance_level': 0.01},
+        'library': 'pgmpy',
+        'output_type': 'dag',
+        'family': 'hybrid',
+    },
+    'K2': {
+        'func': run_k2,
+        'kwargs': {'max_parents': 3},
+        'library': 'pgmpy',
+        'output_type': 'dag',
+        'family': 'score',
+    },
+}
+
+
+def run_algorithm(algo_name, df, timeout_sec=600):
+    """Run a single algorithm on a dataset.
+    
+    Returns:
+        dict with keys: adjmat, runtime, status, output_type
+    """
+    if algo_name not in ALGORITHM_POOL:
+        raise ValueError(f"Unknown algorithm: {algo_name}")
+    
+    algo = ALGORITHM_POOL[algo_name]
+    func = algo['func']
+    kwargs = algo['kwargs'].copy()
+    
+    adjmat, runtime, status = safe_run(
+        algo_name, func, timeout_sec, df, **kwargs
+    )
+    
+    return {
+        'adjmat': adjmat,
+        'runtime': runtime,
+        'status': status,
+        'output_type': algo['output_type'],
+        'family': algo['family'],
+    }
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+    
+    # Quick test on Asia
+    from causal_selection.data.generator import load_bn_model, sample_dataset
+    
+    model = load_bn_model('asia')
+    df = sample_dataset(model, 1000, seed=0)
+    
+    print(f"Testing on ASIA (N=1000)...")
+    for algo_name in ALGORITHM_POOL:
+        result = run_algorithm(algo_name, df, timeout_sec=60)
+        status = result['status']
+        runtime = result['runtime']
+        if result['adjmat'] is not None:
+            n_edges = result['adjmat'].sum()
+            print(f"  {algo_name:15s}: {status:10s} {runtime:6.2f}s  edges={n_edges}")
+        else:
+            print(f"  {algo_name:15s}: {status:20s} {runtime:6.2f}s")