""" 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")