agentgraph/causal/component_analysis.py

#!/usr/bin/env python3
"""
Causal Component Analysis

This script implements causal inference methods to analyze the causal relationship
between knowledge graph components and perturbation scores.
"""

import os
import sys
import pandas as pd
import numpy as np
import logging
import argparse
from typing import Dict, List, Optional, Tuple, Set
from sklearn.linear_model import LinearRegression

# Import from utils directory  
from .utils.dataframe_builder import create_component_influence_dataframe
# Import shared utilities
from .utils.shared_utils import list_available_components

# Configure logging for this module
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

def calculate_average_treatment_effect(
    df: pd.DataFrame, 
    component_id: str,
    outcome_var: str = "perturbation",
    control_vars: Optional[List[str]] = None
) -> Dict[str, float]:
    """
    Calculates the Average Treatment Effect (ATE) of a component on perturbation score.
    
    Args:
        df: DataFrame with binary component features and perturbation score
        component_id: ID of the component to analyze (including 'entity_' or 'relation_' prefix)
        outcome_var: Name of the outcome variable (default: 'perturbation')
        control_vars: List of control variables to include in the model (other components)
        
    Returns:
        Dictionary with ATE estimates and confidence intervals
    """
    if component_id not in df.columns:
        logger.error(f"Component {component_id} not found in DataFrame")
        return {
            "ate": 0.0,
            "std_error": 0.0,
            "p_value": 1.0,
            "confidence_interval_95": (0.0, 0.0)
        }
    
    # Check if there's enough variation in the treatment variable
    if df[component_id].std() == 0:
        logger.warning(f"No variation in component {component_id}, cannot estimate causal effect")
        return {
            "ate": 0.0,
            "std_error": 0.0,
            "p_value": 1.0,
            "confidence_interval_95": (0.0, 0.0)
        }
    
    # Check if there's enough variation in the outcome variable
    if df[outcome_var].std() == 0:
        logger.warning(f"No variation in outcome {outcome_var}, cannot estimate causal effect")
        return {
            "ate": 0.0,
            "std_error": 0.0,
            "p_value": 1.0,
            "confidence_interval_95": (0.0, 0.0)
        }
        
    # Select control variables (other components that could confound the relationship)
    if control_vars is None:
        # Use all other components as control variables
        control_vars = [col for col in df.columns if (col.startswith("entity_") or col.startswith("relation_")) and col != component_id]
    
    # Create treatment and control groups
    treatment_group = df[df[component_id] == 1]
    control_group = df[df[component_id] == 0]
    
    # Calculate naive ATE (without controlling for confounders)
    naive_ate = treatment_group[outcome_var].mean() - control_group[outcome_var].mean()
    
    # Implement regression adjustment to control for confounders
    X = df[control_vars + [component_id]]
    y = df[outcome_var]
    
    # Use linear regression for adjustment
    model = LinearRegression()
    model.fit(X, y)
    
    # Extract coefficient for the component of interest (the ATE)
    component_idx = control_vars.index(component_id) if component_id in control_vars else -1
    ate = model.coef_[component_idx]
    
    # Use bootstrapping to calculate standard errors and confidence intervals
    # Simplified implementation for demonstration
    n_bootstrap = 1000
    bootstrap_ates = []
    
    for _ in range(n_bootstrap):
        # Sample with replacement
        sample_idx = np.random.choice(len(df), len(df), replace=True)
        sample_df = df.iloc[sample_idx]
        
        # Calculate ATE for this sample
        sample_X = sample_df[control_vars + [component_id]]
        sample_y = sample_df[outcome_var]
        
        try:
            sample_model = LinearRegression()
            sample_model.fit(sample_X, sample_y)
            sample_ate = sample_model.coef_[component_idx]
            bootstrap_ates.append(sample_ate)
        except:
            # Skip problematic samples
            continue
    
    # Calculate standard error and confidence intervals
    std_error = np.std(bootstrap_ates)
    ci_lower = np.percentile(bootstrap_ates, 2.5)
    ci_upper = np.percentile(bootstrap_ates, 97.5)
    
    # Calculate p-value (simplified approach)
    z_score = ate / std_error if std_error > 0 else 0
    p_value = 2 * (1 - abs(z_score)) if z_score != 0 else 1.0
    
    return {
        "ate": ate,
        "naive_ate": naive_ate,
        "std_error": std_error,
        "p_value": p_value,
        "confidence_interval_95": (ci_lower, ci_upper)
    }

def granger_causality_test(
    df: pd.DataFrame,
    component_id: str,
    outcome_var: str = "perturbation",
    max_lag: int = 2
) -> Dict[str, float]:
    """
    Implements a simplified Granger causality test to assess if a component
    'Granger-causes' the perturbation score.
    
    Note: This is a simplified implementation and requires time-series data.
    If the data doesn't have a clear time dimension, the results should be
    interpreted with caution.
    
    Args:
        df: DataFrame with binary component features and perturbation score
        component_id: ID of the component to analyze (including 'entity_' or 'relation_' prefix)
        outcome_var: Name of the outcome variable (default: 'perturbation')
        max_lag: Maximum number of lags to include in the model
        
    Returns:
        Dictionary with Granger causality test results
    """
    if component_id not in df.columns:
        logger.error(f"Component {component_id} not found in DataFrame")
        return {"f_statistic": 0.0, "p_value": 1.0, "causal_direction": "none"}
    
    # Check if there's enough data points
    if len(df) <= max_lag + 1:
        logger.warning(f"Not enough data points for Granger causality test with max_lag={max_lag}")
        return {"f_statistic": 0.0, "p_value": 1.0, "causal_direction": "none"}
    
    # Check if there's enough variation in the variables
    if df[component_id].std() == 0 or df[outcome_var].std() == 0:
        logger.warning(f"No variation in component or outcome, cannot test Granger causality")
        return {"f_statistic": 0.0, "p_value": 1.0, "causal_direction": "none"}
    
    # Implement Granger causality test using OLS and F-test
    # This is a simplified approach - in practice, use statsmodels or other libraries
    
    # First, create lagged versions of the data
    lagged_df = df.copy()
    for i in range(1, max_lag + 1):
        lagged_df[f"{component_id}_lag{i}"] = df[component_id].shift(i)
        lagged_df[f"{outcome_var}_lag{i}"] = df[outcome_var].shift(i)
    
    # Drop rows with NaN values (due to lagging)
    lagged_df = lagged_df.dropna()
    
    # Model 1: Outcome ~ Past Outcomes
    X1 = lagged_df[[f"{outcome_var}_lag{i}" for i in range(1, max_lag + 1)]]
    y = lagged_df[outcome_var]
    model1 = LinearRegression()
    model1.fit(X1, y)
    y_pred1 = model1.predict(X1)
    ssr1 = np.sum((y - y_pred1) ** 2)
    
    # Model 2: Outcome ~ Past Outcomes + Past Component
    X2 = lagged_df[[f"{outcome_var}_lag{i}" for i in range(1, max_lag + 1)] + 
                  [f"{component_id}_lag{i}" for i in range(1, max_lag + 1)]]
    model2 = LinearRegression()
    model2.fit(X2, y)
    y_pred2 = model2.predict(X2)
    ssr2 = np.sum((y - y_pred2) ** 2)
    
    # Calculate F-statistic
    n = len(lagged_df)
    df1 = max_lag
    df2 = n - 2 * max_lag - 1
    
    if ssr1 == 0 or df2 <= 0:
        f_statistic = 0
        p_value = 1.0
    else:
        f_statistic = ((ssr1 - ssr2) / df1) / (ssr2 / df2)
        # Simplified p-value calculation (for demonstration)
        p_value = 1 / (1 + f_statistic)
    
    # Test reverse causality
    # Model 3: Component ~ Past Components
    X3 = lagged_df[[f"{component_id}_lag{i}" for i in range(1, max_lag + 1)]]
    y_comp = lagged_df[component_id]
    model3 = LinearRegression()
    model3.fit(X3, y_comp)
    y_pred3 = model3.predict(X3)
    ssr3 = np.sum((y_comp - y_pred3) ** 2)
    
    # Model 4: Component ~ Past Components + Past Outcomes
    X4 = lagged_df[[f"{component_id}_lag{i}" for i in range(1, max_lag + 1)] + 
                  [f"{outcome_var}_lag{i}" for i in range(1, max_lag + 1)]]
    model4 = LinearRegression()
    model4.fit(X4, y_comp)
    y_pred4 = model4.predict(X4)
    ssr4 = np.sum((y_comp - y_pred4) ** 2)
    
    # Calculate F-statistic for reverse causality
    if ssr3 == 0 or df2 <= 0:
        f_statistic_reverse = 0
        p_value_reverse = 1.0
    else:
        f_statistic_reverse = ((ssr3 - ssr4) / df1) / (ssr4 / df2)
        # Simplified p-value calculation
        p_value_reverse = 1 / (1 + f_statistic_reverse)
    
    # Determine causality direction
    causal_direction = "none"
    if p_value < 0.05 and p_value_reverse >= 0.05:
        causal_direction = "component -> outcome"
    elif p_value >= 0.05 and p_value_reverse < 0.05:
        causal_direction = "outcome -> component"
    elif p_value < 0.05 and p_value_reverse < 0.05:
        causal_direction = "bidirectional"
    
    return {
        "f_statistic": f_statistic,
        "p_value": p_value,
        "f_statistic_reverse": f_statistic_reverse,
        "p_value_reverse": p_value_reverse,
        "causal_direction": causal_direction
    }

def compute_causal_effect_strength(
    df: pd.DataFrame,
    control_group: Optional[List[str]] = None,
    outcome_var: str = "perturbation"
) -> Dict[str, float]:
    """
    Computes the strength of causal effects for all components.
    
    Args:
        df: DataFrame with binary component features and perturbation score
        control_group: List of components to use as control variables
        outcome_var: Name of the outcome variable (default: 'perturbation')
        
    Returns:
        Dictionary mapping component IDs to their causal effect strengths
    """
    # Get all component columns
    component_cols = [col for col in df.columns if col.startswith(("entity_", "relation_"))]
    
    if not component_cols:
        logger.error("No component features found in DataFrame")
        return {}
    
    # Calculate ATE for each component
    effect_strengths = {}
    for component_id in component_cols:
        try:
            ate_results = calculate_average_treatment_effect(
                df, 
                component_id,
                outcome_var=outcome_var,
                control_vars=control_group
            )
            effect_strengths[component_id] = ate_results["ate"]
        except Exception as e:
            logger.warning(f"Error calculating ATE for {component_id}: {e}")
            effect_strengths[component_id] = 0.0
    
    return effect_strengths

# Note: create_mock_perturbation_scores and list_available_components 
# moved to utils.shared_utils to avoid duplication

def main():
    """Main function to run the causal component analysis."""
    parser = argparse.ArgumentParser(description='Analyze causal relationships between components and perturbation scores')
    parser.add_argument('--input', '-i', required=True, help='Path to the knowledge graph JSON file')
    parser.add_argument('--output', '-o', help='Path to save the output analysis (CSV format)')
    args = parser.parse_args()
    
    print(f"Loading knowledge graph")
    
    # Create DataFrame
    df = create_component_influence_dataframe(args.input)
    
    if df is None or df.empty:
        logger.error("Failed to create or empty DataFrame. Cannot proceed with analysis.")
        return
    
    # Print basic DataFrame info
    print(f"\nDataFrame info:")
    print(f"Rows: {len(df)}")
    entity_features = [col for col in df.columns if col.startswith("entity_")]
    relation_features = [col for col in df.columns if col.startswith("relation_")]
    print(f"Entity features: {len(entity_features)}")
    print(f"Relation features: {len(relation_features)}")
    
    # Check if we have any variance in perturbation scores
    if df['perturbation'].std() == 0:
        logger.warning("All perturbation scores are identical. This might lead to uninformative results.")
        print("\nWARNING: All perturbation scores are identical (value: %.2f). Results may not be meaningful." % df['perturbation'].iloc[0])
    else:
        print(f"\nPerturbation score distribution:")
        print(f"Min: {df['perturbation'].min():.2f}, Max: {df['perturbation'].max():.2f}")
        print(f"Mean: {df['perturbation'].mean():.2f}, Std: {df['perturbation'].std():.2f}")
    
    # Compute causal effect strengths
    print("\nComputing causal effect strengths...")
    effect_strengths = compute_causal_effect_strength(df)
    print(f"Found {len(effect_strengths)} components with causal effects")
    
    # Sort components by effect strength
    sorted_components = sorted(effect_strengths.items(), key=lambda x: abs(x[1]), reverse=True)
    
    print("\nTop 10 Components by Causal Effect Strength:")
    print("=" * 50)
    print(f"{'Rank':<5}{'Component':<30}{'Effect Strength':<15}")
    print("-" * 50)
    
    for i, (component, strength) in enumerate(sorted_components[:10], 1):
        print(f"{i:<5}{component:<30}{strength:.6f}")
    
    # Save results
    if args.output:
        # Create results DataFrame
        results_df = pd.DataFrame({
            'Component': [comp for comp, _ in sorted_components],
            'Effect_Strength': [strength for _, strength in sorted_components]
        })
        
        # Save to specified output path
        print(f"\nSaving results to: {args.output}")
        try:
            results_df.to_csv(args.output, index=False)
            print(f"Successfully saved results to: {args.output}")
        except Exception as e:
            print(f"Error saving to {args.output}: {str(e)}")
        
        # Also save to default location in the causal_analysis directory
        default_output = os.path.join(os.path.dirname(__file__), 'causal_component_effects.csv')
        print(f"Also saving results to: {default_output}")
        try:
            results_df.to_csv(default_output, index=False)
            print(f"Successfully saved results to: {default_output}")
        except Exception as e:
            print(f"Error saving to {default_output}: {str(e)}")
    
    print("\nAnalysis complete.")

if __name__ == "__main__":
    main()