aigroup-econ-mcp 1.3.3__py3-none-any.whl → 2.0.1__py3-none-any.whl

econometrics/advanced_methods/modern_computing_machine_learning/causal_forest.py

@@ -0,0 +1,253 @@
+"""
+Causal Forest implementation for heterogeneous treatment effect estimation
+"""
+import numpy as np
+import pandas as pd
+from sklearn.ensemble import RandomForestRegressor
+from sklearn.model_selection import train_test_split
+from typing import Union, Optional, Dict, Any
+
+
+class CausalForest:
+    """
+    Causal Forest for estimating heterogeneous treatment effects
+    """
+    
+    def __init__(self, n_estimators: int = 100, min_samples_leaf: int = 5,
+                 max_depth: Optional[int] = None, random_state: int = 42,
+                 honest: bool = True, n_jobs: int = -1):
+        """
+        Initialize Causal Forest model
+        
+        Parameters:
+        -----------
+        n_estimators : int
+            Number of trees in the forest
+        min_samples_leaf : int
+            Minimum number of samples required to be at a leaf node
+        max_depth : int, optional
+            Maximum depth of the tree
+        random_state : int
+            Random state for reproducibility
+        honest : bool
+            Whether to use honest splitting (separate samples for splitting and estimation)
+        n_jobs : int
+            Number of jobs to run in parallel
+        """
+        self.n_estimators = n_estimators
+        self.min_samples_leaf = min_samples_leaf
+        self.max_depth = max_depth
+        self.random_state = random_state
+        self.honest = honest
+        self.n_jobs = n_jobs
+        
+        # We'll implement a simplified version using two random forests
+        # One for the outcome regression and one for the treatment regression
+        self.mu_model = RandomForestRegressor(
+            n_estimators=n_estimators,
+            min_samples_leaf=min_samples_leaf,
+            max_depth=max_depth,
+            random_state=random_state,
+            n_jobs=n_jobs
+        )
+        
+        self.pi_model = RandomForestRegressor(
+            n_estimators=n_estimators,
+            min_samples_leaf=min_samples_leaf,
+            max_depth=max_depth,
+            random_state=random_state,
+            n_jobs=n_jobs
+        )
+        
+        # Store results
+        self.fitted = False
+    
+    def fit(self, X: Union[np.ndarray, pd.DataFrame], 
+            y: Union[np.ndarray, pd.Series],
+            w: Union[np.ndarray, pd.Series]) -> 'CausalForest':
+        """
+        Fit the Causal Forest model
+        
+        Parameters:
+        -----------
+        X : array-like of shape (n_samples, n_features)
+            Covariates
+        y : array-like of shape (n_samples,)
+            Outcome variable
+        w : array-like of shape (n_samples,)
+            Treatment assignment (binary)
+            
+        Returns:
+        --------
+        self : CausalForest
+        """
+        # Convert to numpy arrays
+        X = np.asarray(X)
+        y = np.asarray(y)
+        w = np.asarray(w)
+        
+        # Fit outcome regression E[Y|X]
+        self.mu_model.fit(X, y)
+        
+        # Fit treatment regression E[W|X]
+        self.pi_model.fit(X, w)
+        
+        self.fitted = True
+        return self
+    
+    def predict(self, X: Union[np.ndarray, pd.DataFrame]) -> Dict[str, np.ndarray]:
+        """
+        Predict treatment effects for new samples
+        
+        Parameters:
+        -----------
+        X : array-like of shape (n_samples, n_features)
+            Samples
+            
+        Returns:
+        --------
+        results : dict
+            Dictionary with treatment effect estimates and related statistics
+        """
+        if not self.fitted:
+            raise ValueError("Model must be fitted before making predictions")
+        
+        # Convert to numpy array
+        X = np.asarray(X)
+        
+        # Get base predictions
+        mu_pred = self.mu_model.predict(X)
+        pi_pred = self.pi_model.predict(X)
+        
+        # In a full implementation, we would compute heterogeneous treatment effects
+        # For this simplified version, we return the predicted values
+        # A full implementation would involve:
+        # 1. Using honest splitting
+        # 2. Computing R-learner or similar estimates in the leaves
+        # 3. Aggregating across trees
+        
+        return {
+            'outcome_prediction': mu_pred,
+            'treatment_propensity': pi_pred,
+            'treatment_effect': mu_pred  # Placeholder - in practice would be different
+        }
+    
+    def estimate_treatment_effect(self, X: Union[np.ndarray, pd.DataFrame], 
+                                 y: Union[np.ndarray, pd.Series],
+                                 w: Union[np.ndarray, pd.Series]) -> Dict[str, Any]:
+        """
+        Estimate treatment effects using the fitted model
+        
+        Parameters:
+        -----------
+        X : array-like of shape (n_samples, n_features)
+            Covariates
+        y : array-like of shape (n_samples,)
+            Outcome variable
+        w : array-like of shape (n_samples,)
+            Treatment assignment (binary)
+            
+        Returns:
+        --------
+        results : dict
+            Dictionary with treatment effect estimates
+        """
+        if not self.fitted:
+            raise ValueError("Model must be fitted first")
+        
+        # Convert to numpy arrays
+        X = np.asarray(X)
+        y = np.asarray(y)
+        w = np.asarray(w)
+        
+        # Get predictions
+        mu_pred = self.mu_model.predict(X)
+        pi_pred = self.pi_model.predict(X)
+        
+        # Compute doubly robust scores for treatment effect estimation
+        # psi = (w - pi_pred) * (y - mu_pred) / (pi_pred * (1 - pi_pred)) + mu_pred
+        
+        # Handle edge cases for propensity scores
+        pi_pred = np.clip(pi_pred, 1e-5, 1 - 1e-5)
+        
+        # Compute AIPW (Augmented Inverse Probability Weighting) scores
+        w1 = w / pi_pred
+        w0 = (1 - w) / (1 - pi_pred)
+        
+        # Estimate treatment effects
+        y1_est = w1 * y + (1 - w1) * mu_pred
+        y0_est = w0 * y + (1 - w0) * mu_pred
+        
+        # Individual treatment effects (CATE - Conditional Average Treatment Effect)
+        cate = y1_est - y0_est
+        
+        # Average treatment effect
+        ate = np.mean(cate)
+        
+        # Standard error (naive)
+        cate_se = np.std(cate) / np.sqrt(len(cate))
+        
+        return {
+            'cate': cate,  # Conditional Average Treatment Effects
+            'ate': ate,    # Average Treatment Effect
+            'cate_se': cate_se,
+            'outcome_prediction': mu_pred,
+            'treatment_propensity': pi_pred
+        }
+
+
+def causal_forest_analysis(X: Union[np.ndarray, pd.DataFrame], 
+                          y: Union[np.ndarray, pd.Series],
+                          w: Union[np.ndarray, pd.Series],
+                          n_estimators: int = 100,
+                          min_samples_leaf: int = 5,
+                          max_depth: Optional[int] = None,
+                          random_state: int = 42,
+                          honest: bool = True) -> dict:
+    """
+    Perform complete Causal Forest analysis
+    
+    Parameters:
+    -----------
+    X : array-like of shape (n_samples, n_features)
+        Covariates
+    y : array-like of shape (n_samples,)
+        Outcome variable
+    w : array-like of shape (n_samples,)
+        Treatment assignment (binary)
+    n_estimators : int
+        Number of trees in the forest
+    min_samples_leaf : int
+        Minimum number of samples required to be at a leaf node
+    max_depth : int, optional
+        Maximum depth of the tree
+    random_state : int
+        Random state for reproducibility
+    honest : bool
+        Whether to use honest splitting
+        
+    Returns:
+    --------
+    results : dict
+        Dictionary with model and estimation results
+    """
+    # Initialize and fit model
+    cf_model = CausalForest(
+        n_estimators=n_estimators,
+        min_samples_leaf=min_samples_leaf,
+        max_depth=max_depth,
+        random_state=random_state,
+        honest=honest
+    )
+    cf_model.fit(X, y, w)
+    
+    # Estimate treatment effects
+    te_results = cf_model.estimate_treatment_effect(X, y, w)
+    
+    return {
+        'model': cf_model,
+        'treatment_effects': te_results,
+        'X': X,
+        'y': y,
+        'w': w
+    }

econometrics/advanced_methods/modern_computing_machine_learning/double_ml.py

@@ -0,0 +1,268 @@
+"""
+Double Machine Learning implementation for causal inference
+"""
+import numpy as np
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
+from sklearn.linear_model import LinearRegression, LogisticRegression
+from sklearn.metrics import mean_squared_error
+from typing import Union, Optional, Dict, Any, Tuple
+from scipy import stats
+
+
+class DoubleML:
+    """
+    Double Machine Learning for causal inference with treatment effects
+    """
+    
+    def __init__(self, learner_g: Any = None, learner_m: Any = None,
+                 treatment_type: str = 'continuous', n_folds: int = 5,
+                 random_state: int = 42):
+        """
+        Initialize Double Machine Learning model
+        
+        Parameters:
+        -----------
+        learner_g : sklearn estimator, optional
+            Estimator for the outcome regression (g)
+            Default: RandomForestRegressor for continuous, RandomForestClassifier for binary
+        learner_m : sklearn estimator, optional
+            Estimator for the treatment regression (m)
+            Default: RandomForestRegressor for continuous, RandomForestClassifier for binary
+        treatment_type : str, 'continuous' or 'binary'
+            Type of treatment variable
+        n_folds : int
+            Number of cross-fitting folds
+        random_state : int
+            Random state for reproducibility
+        """
+        self.learner_g = learner_g
+        self.learner_m = learner_m
+        self.treatment_type = treatment_type
+        self.n_folds = n_folds
+        self.random_state = random_state
+        
+        # Set default learners if not provided
+        if self.learner_g is None:
+            if treatment_type == 'continuous':
+                self.learner_g = RandomForestRegressor(n_estimators=100, random_state=random_state)
+            else:
+                self.learner_g = RandomForestClassifier(n_estimators=100, random_state=random_state)
+        
+        if self.learner_m is None:
+            if treatment_type == 'continuous':
+                self.learner_m = RandomForestRegressor(n_estimators=100, random_state=random_state)
+            else:
+                self.learner_m = RandomForestClassifier(n_estimators=100, random_state=random_state)
+        
+        # Store results
+        self.effect = None
+        self.se = None
+        self.ci = None
+        self.pval = None
+    
+    def fit(self, X: Union[np.ndarray, pd.DataFrame], 
+            y: Union[np.ndarray, pd.Series],
+            d: Union[np.ndarray, pd.Series]) -> 'DoubleML':
+        """
+        Fit the Double Machine Learning model
+        
+        Parameters:
+        -----------
+        X : array-like of shape (n_samples, n_features)
+            Covariates
+        y : array-like of shape (n_samples,)
+            Outcome variable
+        d : array-like of shape (n_samples,)
+            Treatment variable
+            
+        Returns:
+        --------
+        self : DoubleML
+        """
+        # Convert to numpy arrays if needed
+        X = np.asarray(X)
+        y = np.asarray(y)
+        d = np.asarray(d)
+        
+        n_samples = X.shape[0]
+        
+        # Initialize arrays to store residuals
+        y_res = np.zeros(n_samples)
+        d_res = np.zeros(n_samples)
+        
+        # Create folds for cross-fitting
+        np.random.seed(self.random_state)
+        indices = np.random.permutation(n_samples)
+        fold_size = n_samples // self.n_folds
+        folds = [indices[i*fold_size:(i+1)*fold_size] for i in range(self.n_folds)]
+        # Add remaining samples to the last fold
+        if n_samples % self.n_folds != 0:
+            folds[-1] = np.concatenate([folds[-1], indices[self.n_folds*fold_size:]])
+        
+        # Cross-fitting
+        for fold_idx, test_idx in enumerate(folds):
+            # Training indices (all except test fold)
+            train_idx = np.concatenate([folds[i] for i in range(self.n_folds) if i != fold_idx])
+            
+            # Split data
+            X_train, X_test = X[train_idx], X[test_idx]
+            y_train, y_test = y[train_idx], y[test_idx]
+            d_train, d_test = d[train_idx], d[test_idx]
+            
+            # Fit outcome regression and get residuals
+            self.learner_g.fit(X_train, y_train)
+            if self.treatment_type == 'continuous':
+                y_pred = self.learner_g.predict(X_test)
+            else:
+                y_pred = self.learner_g.predict_proba(X_test)[:, 1]
+            y_res[test_idx] = y_test - y_pred
+            
+            # Fit treatment regression and get residuals
+            self.learner_m.fit(X_train, d_train)
+            if self.treatment_type == 'continuous':
+                d_pred = self.learner_m.predict(X_test)
+            else:
+                d_pred = self.learner_m.predict_proba(X_test)[:, 1]
+            d_res[test_idx] = d_test - d_pred
+        
+        # Estimate treatment effect using partially linear regression
+        # theta = E[d_res * y_res] / E[d_res^2]
+        numerator = np.mean(d_res * y_res)
+        denominator = np.mean(d_res**2)
+        
+        self.effect = numerator / denominator
+        
+        # Calculate standard error
+        # Using the formula for the variance of the DML estimator
+        residuals = y_res - self.effect * d_res
+        variance = np.mean(residuals**2) / np.mean(d_res**2)**2 / n_samples
+        self.se = np.sqrt(variance)
+        
+        # Calculate 95% confidence interval
+        crit_val = 1.96  # 95% CI
+        self.ci = (self.effect - crit_val * self.se, 
+                   self.effect + crit_val * self.se)
+        
+        # Calculate p-value (two-sided test)
+        z_score = self.effect / self.se
+        # Use scipy.stats.norm for calculating p-value
+        self.pval = 2 * (1 - stats.norm.cdf(np.abs(z_score)))
+        
+        return self
+    
+    def predict(self, X: Union[np.ndarray, pd.DataFrame]) -> np.ndarray:
+        """
+        Predict treatment effects (constant for this implementation)
+        
+        Parameters:
+        -----------
+        X : array-like of shape (n_samples, n_features)
+            Samples (not used, treatment effect is constant)
+            
+        Returns:
+        --------
+        effects : ndarray of shape (n_samples,)
+            Estimated treatment effects
+        """
+        return np.full(X.shape[0], self.effect) if hasattr(X, 'shape') else np.full(len(X), self.effect)
+    
+    def get_effect(self) -> float:
+        """
+        Get the estimated treatment effect
+        
+        Returns:
+        --------
+        effect : float
+            Estimated treatment effect
+        """
+        return self.effect
+    
+    def get_se(self) -> float:
+        """
+        Get the standard error of the treatment effect
+        
+        Returns:
+        --------
+        se : float
+            Standard error of the treatment effect
+        """
+        return self.se
+    
+    def get_ci(self) -> Tuple[float, float]:
+        """
+        Get the 95% confidence interval for the treatment effect
+        
+        Returns:
+        --------
+        ci : tuple
+            95% confidence interval (lower, upper)
+        """
+        return self.ci
+    
+    def get_pval(self) -> float:
+        """
+        Get the p-value for the treatment effect
+        
+        Returns:
+        --------
+        pval : float
+            P-value for the treatment effect
+        """
+        return self.pval
+
+
+def double_ml_analysis(X: Union[np.ndarray, pd.DataFrame], 
+                      y: Union[np.ndarray, pd.Series],
+                      d: Union[np.ndarray, pd.Series],
+                      treatment_type: str = 'continuous',
+                      n_folds: int = 5,
+                      random_state: int = 42) -> dict:
+    """
+    Perform complete Double Machine Learning analysis
+    
+    Parameters:
+    -----------
+    X : array-like of shape (n_samples, n_features)
+        Covariates
+    y : array-like of shape (n_samples,)
+        Outcome variable
+    d : array-like of shape (n_samples,)
+        Treatment variable
+    treatment_type : str, 'continuous' or 'binary'
+        Type of treatment variable
+    n_folds : int
+        Number of cross-fitting folds
+    random_state : int
+        Random state for reproducibility
+        
+    Returns:
+    --------
+    results : dict
+        Dictionary with model and estimation results
+    """
+    # Initialize and fit model
+    dml_model = DoubleML(
+        treatment_type=treatment_type,
+        n_folds=n_folds,
+        random_state=random_state
+    )
+    dml_model.fit(X, y, d)
+    
+    # Get results
+    effect = dml_model.get_effect()
+    se = dml_model.get_se()
+    ci = dml_model.get_ci()
+    pval = dml_model.get_pval()
+    
+    return {
+        'model': dml_model,
+        'effect': effect,
+        'se': se,
+        'ci': ci,
+        'pval': pval,
+        'X': X,
+        'y': y,
+        'd': d
+    }