causarray 0.0.3__tar.gz → 0.0.5__tar.gz

{causarray-0.0.3 → causarray-0.0.5}/.gitignore

@@ -160,4 +160,5 @@ cython_debug/
 #.idea/
 
 
-test*
+test*
+causarray/___*.py

{causarray-0.0.3 → causarray-0.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: causarray
-Version: 0.0.3
+Version: 0.0.5
 Summary: causarray is a Python module for simultaneous causal inference with an array of outcomes.
 Author-email: Jin-Hong Du <jinhongd@andrew.cmu.com>, Maya Shen <myshen@andrew.cmu.edu>, Hansruedi Mathys <mathysh@pitt.edu>, Kathryn Roeder <jinhongd@andrew.cmu.com>
 Maintainer-email: Jin-Hong Du <jinhongd@andrew.cmu.com>

{causarray-0.0.3 → causarray-0.0.5}/causarray/DR_estimation.py

@@ -1,9 +1,12 @@
 import numpy as np
 from sklearn.linear_model import LogisticRegression
-from sklearn.ensemble import RandomForestClassifier
+from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
+from sklearn_ensemble_cv import reset_random_seeds, Ensemble, ECV
 from causarray.gcate_glm import fit_glm
 from causarray.utils import *
 from causarray.utils import _filter_params
+from joblib import Parallel, delayed
+from tqdm import tqdm
 import pprint
 
 from sklearn.model_selection import KFold, ShuffleSplit
@@ -33,7 +36,7 @@ def _get_func_ps(ps_model, **kwargs):
 def cross_fitting(
     Y, A, X, X_A, family='poisson', K=1, glm_alpha=1e-4,
     ps_model='logistic', 
-    pi_hat=None, Y_hat=None, verbose=False, **kwargs):
+    Y_hat=None, pi_hat=None, mask=None, verbose=False, **kwargs):
     '''
     Cross-fitting for causal estimands.
 
@@ -55,10 +58,16 @@ def cross_fitting(
         The regularization parameter for the generalized linear model. The default is 1e-4.
     ps_model : str, optional
         The propensity score model. The default is 'logistic'.
-    pi_hat : array, optional
-        Propensity score of shape (n, a). The default is None.
+    
     Y_hat : array, optional
         Estimated potential outcome of shape (n, p, a, 2). The default is None.
+    pi_hat : array, optional
+        Propensity score of shape (n, a). The default is None.
+    mask : array, optional
+        Boolean mask of shape (n, a) for the treatment, indicating which samples are used for 
+        the estimation of the estimand. This does not affect the estimation of pseudo-outcomes
+        and propensity scores.
+
     **kwargs : dict
         Additional arguments to pass to the model.
 
@@ -76,10 +85,15 @@ def cross_fitting(
         pprint.pprint(params_ps)
         pprint.pprint(params_glm)
     
-    if K>1:
-        # Initialize KFold cross-validator
-        kf = KFold(n_splits=K, random_state=0, shuffle=True)
-        folds = kf.split(X)
+    if K > 1:
+        n_samples = X.shape[0]
+        if K >= n_samples:
+            # Use Leave-One-Out Cross-Validation
+            folds = [([i for i in range(n_samples) if i != j], [j]) for j in range(n_samples)]
+        else:
+            # Initialize KFold cross-validator
+            kf = KFold(n_splits=int(K), random_state=0, shuffle=True)
+            folds = kf.split(X)
     else:
         folds = [(np.arange(X.shape[0]), np.arange(X.shape[0]))]
 
@@ -89,6 +103,16 @@ def cross_fitting(
     fit_Y = True if Y_hat is None else False
     Y_hat = np.zeros((Y.shape[0],Y.shape[1],A.shape[1],2), dtype=float) if fit_Y else Y_hat
 
+    # perform ECV at once
+    if fit_pi and ps_model == 'random_forest_cv':
+        info_ecv = run_ecv(X_A, A, **params_ps)
+        func_ps, params_ps = _get_func_ps(ps_model, verbose=False, ecv=False, 
+                kwargs_ensemble=info_ecv['best_params_ensemble'], kwargs_regr=info_ecv['best_params_regr'])
+        pprint.pprint('Best parameters for the regression model:')
+        pprint.pprint(info_ecv['best_params_regr'])
+        pprint.pprint('Best parameters for the ensemble model:')
+        pprint.pprint(info_ecv['best_params_ensemble'])
+
     # Perform cross-fitting
     for train_index, test_index in folds:
         # Split data
@@ -104,7 +128,12 @@ def cross_fitting(
             pi = np.zeros_like(A_test, dtype=float)
             for j in range(A.shape[1]):
                 i_case = (A_train[:,j] == 1.)
-                i_cells = i_ctrl | i_case
+
+                if mask is not None:
+                    i_cells = mask[:, j]
+                else:
+                    i_ctrl = (np.sum(A_train, axis=1) == 0.)
+                    i_cells = i_ctrl | i_case
 
                 if ps_model=='logistic' and XA_train.shape[1]==1 and np.all(XA_train==1):
                     prob = np.sum(i_case)/np.sum(i_cells)
@@ -167,8 +196,6 @@ def AIPW_mean(Y, A, mu, pi, positive=False):
     tau = np.mean(pseudo_y, axis=0)
 
     return tau, pseudo_y
-    
-
 
 
 
@@ -177,51 +204,89 @@ def AIPW_mean(Y, A, mu, pi, positive=False):
 
 
 
-from joblib import Parallel, delayed
-from tqdm import tqdm
-from sklearn_ensemble_cv import reset_random_seeds, Ensemble, ECV
-from sklearn.tree import DecisionTreeRegressor
-
-def fit_rf(X, y, X_test=None, sample_weight=None, M=100, M_max=1000,
+def run_ecv(
+    X, y, M=200, M_max=1000,
     # fixed parameters for bagging regressor
-    kwargs_ensemble={'verbose':1},
+    kwargs_ensemble={},
     # fixed parameters for decision tree
-    kwargs_regr={'min_samples_leaf': 3}, # 'min_samples_split': 10, 'max_features':'sqrt'
+    kwargs_regr={},
     # grid search parameters
-    grid_regr = {'max_depth': [11]},
-    grid_ensemble = {'random_state': 0}, #'max_samples':np.linspace(0.25, 1., 4)
-    ):
+    grid_regr={},
+    grid_ensemble={}
+):
+    """
+    Runs Ensemble Cross-Validation (ECV) to find the best hyperparameters.
+    """
+    kwargs_ensemble = {**{'verbose': 1, 'bootstrap': True}, **kwargs_ensemble}
+    kwargs_regr = {**{'min_samples_split': 20, 'min_samples_leaf': 10, 'max_features': 'sqrt', 'ccp_alpha': 0.02, 'class_weight': 'balanced'}, **kwargs_regr}
+    grid_regr = {**{'max_depth': [3, 5, 7]}, **grid_regr}
+    grid_ensemble = {**{'random_state': 0, 'max_samples': [0.4, 0.6, 0.8, 1.]}, **grid_ensemble}
 
     # Validate integer parameters
     M = int(M)
     M_max = int(M_max)
-    # for kwargs in [kwargs_regr, kwargs_ensemble, grid_regr, grid_ensemble]:
-    #     for param in kwargs:
-    #         if param in ['max_depth', 'random_state', 'max_leaf_nodes'] and isinstance(kwargs[param], float):
-    #             kwargs[param] = int(kwargs[param])
 
     # Make sure y is 2D
     y = y.reshape(-1, 1) if y.ndim == 1 else y
 
     # Run ECV
-    res_ecv, info_ecv = ECV(
-        X, y, DecisionTreeRegressor, grid_regr, grid_ensemble, 
-        kwargs_regr, kwargs_ensemble, 
+    _, info_ecv = ECV(
+        X, y, DecisionTreeClassifier, grid_regr, grid_ensemble,
+        kwargs_regr, kwargs_ensemble,
         M=M, M0=M, M_max=M_max, return_df=True
     )
 
     # Replace the in-sample best parameter for 'n_estimators' with extrapolated best parameter
     info_ecv['best_params_ensemble']['n_estimators'] = info_ecv['best_n_estimators_extrapolate']
 
+    return info_ecv
+
+
+def fit_rf(
+    X, y, X_test=None, M=100, M_max=1000, ecv=True,
+    # fixed parameters for bagging regressor
+    kwargs_ensemble={},
+    # fixed parameters for decision tree
+    kwargs_regr={},
+    # grid search parameters
+    grid_regr={},
+    grid_ensemble={}
+):
+    """
+    Fits a Random Forest model using parameters found by ECV.
+    """
+
+    kwargs_ensemble = {**{'verbose': 1, 'bootstrap': True}, **kwargs_ensemble}
+    kwargs_regr = {**{'min_samples_split': 20, 'min_samples_leaf': 10, 'max_features': 'sqrt', 'ccp_alpha': 0.02, 'class_weight': 'balanced'}, **kwargs_regr}
+    grid_regr = {**{'max_depth': [3, 5, 7]}, **grid_regr}
+    grid_ensemble = {**{'random_state': 0, 'max_samples': [0.4, 0.6, 0.8, 1.]}, **grid_ensemble}
+
+    # Make sure y is 2D
+    y_2d = y.reshape(-1, 1) if y.ndim == 1 else y
+
+    if ecv:
+        # Get best parameters from ECV
+        info_ecv = run_ecv(
+            X, y_2d, M=M, M_max=M_max,
+            kwargs_ensemble=kwargs_ensemble,
+            kwargs_regr=kwargs_regr,
+            grid_regr=grid_regr,
+            grid_ensemble=grid_ensemble
+        )
+        params_regr = info_ecv['best_params_regr']
+        params_ensemble = info_ecv['best_params_ensemble']
+    else:
+        params_regr = kwargs_regr
+        params_ensemble = kwargs_ensemble
+        
     # Fit the ensemble with the best CV parameters
     regr = Ensemble(
-        estimator=DecisionTreeRegressor(**info_ecv['best_params_regr']),
-        **info_ecv['best_params_ensemble']).fit(X, y, sample_weight=sample_weight)
-        
+        estimator=DecisionTreeClassifier(**params_regr), **params_ensemble).fit(X, y_2d)
+
     # Predict
     if X_test is None:
         X_test = X
-    return regr.predict(X_test).reshape(-1, y.shape[1])
+    return regr.predict(X_test).reshape(-1, y_2d.shape[1])
 
 
 
@@ -241,11 +306,7 @@ def fit_rf_ind_ps(X, Y, *args, **kwargs):
     def _fit(X, y, i_ctrl, *args, **kwargs):        
         i_case = (y == 1.)
         i_cells = i_ctrl | i_case
-        sample_weight = np.ones(y.shape[0])
-        class_weight =  len(y) / (2 * np.bincount(y.astype(int)))  
-        for a in range(2):
-            sample_weight[y == a] = class_weight[a]     
-        return fit_rf(X[i_cells], y[i_cells], sample_weight=sample_weight[i_cells], *args, **kwargs)
+        return fit_rf(X[i_cells], y[i_cells], *args, **kwargs)
 
     Y_hat = Parallel(n_jobs=-1)(delayed(_fit)(X, Y[:,j], i_ctrl, *args, **kwargs)
         for j in tqdm(range(Y.shape[1])))

{causarray-0.0.3 → causarray-0.0.5}/causarray/DR_learner.py

@@ -65,15 +65,13 @@ def compute_causal_estimand(
     '''
     reset_random_seeds(random_state)
 
-    kwargs = {k:v for k,v in kwargs.items() if k not in 
-        ['kwargs_ls_1', 'kwargs_ls_2', 'kwargs_es_1', 'kwargs_es_2', 'c1', 'num_d']
-    }
-
+    # check the input data
     if isinstance(Y, pd.DataFrame):
         gene_names = Y.columns
         Y = Y.values
     else:
         gene_names = range(Y.shape[1])
+    Y = Y.astype('float')
     n, p = Y.shape
 
     if len(A.shape) == 1:
@@ -97,7 +95,10 @@ def compute_causal_estimand(
         if len(mask.shape) == 1: mask = mask.reshape(-1,1)
         if mask.shape != A.shape:
             raise ValueError('Mask must have the same shape as the treatment matrix')
-    
+
+    kwargs = {k:v for k,v in kwargs.items() if k not in 
+        ['kwargs_ls_1', 'kwargs_ls_2', 'kwargs_es_1', 'kwargs_es_2', 'c1', 'num_d']
+    }
 
     if verbose:
         d_A = W_A.shape[1]
@@ -113,10 +114,9 @@ def compute_causal_estimand(
     else:
         offset = None
         size_factors = np.ones(n)
-
-    Y = Y.astype('float')
+    
     Y_hat, pi_hat = cross_fitting(Y, A, W, W_A, family=family, offset=offset, 
-        Y_hat=Y_hat, pi_hat=pi_hat, random_state=random_state, verbose=verbose, **kwargs)
+        Y_hat=Y_hat, pi_hat=pi_hat, mask=mask, random_state=random_state, verbose=verbose, **kwargs)
     pi_hat = pi_hat.reshape(*A.shape)
 
     if verbose: pprint.pprint('Estimating AIPW mean...')
@@ -170,7 +170,7 @@ def compute_causal_estimand(
 def LFC(
     Y, W, A, W_A=None, family='nb', offset=False,    
     Y_hat=None, pi_hat=None, cross_est=False,  mask=None, usevar='pooled',
-    thres_min=1e-4, thres_diff=1e-6, eps_var=1e-3,
+    thres_min=1e-2, thres_diff=1e-2, eps_var=1e-4,
     fdx=False, fdx_alpha=0.05, fdx_c=0.1,     
     verbose=False, **kwargs):
     '''
@@ -244,12 +244,12 @@ def LFC(
             var_1 = np.var(eta_est[A==1], axis=0, ddof=1)
             n_0 = np.sum(A==0)
             n_1 = np.sum(A==1)
-            var_est = (var_0 + eps_var) / n_0 + (var_1 + eps_var) / n_1
+            var_est = ((var_0 + eps_var) / n_0 + (var_1 + eps_var) / n_1) / 2
         else:
             raise ValueError('usevar must be either "pooled" or "unequal"')
 
         # filter out low-expressed genes
-        idx = (np.maximum(tau_0,tau_1)<thres_min) & ((tau_1-tau_0)<thres_diff)
+        idx = (np.maximum(np.abs(tau_0),np.abs(tau_1))<thres_min) | (np.abs(tau_1-tau_0)<thres_diff)
         tau_est[idx] = 0.; eta_est[:,idx] = 0.; var_est[idx] = np.inf
 
         return eta_est, tau_est, var_est

causarray-0.0.5/causarray/__about__.py

@@ -0,0 +1 @@
+__version__ = "0.0.5"

{causarray-0.0.3 → causarray-0.0.5}/causarray/gcate.py

@@ -196,7 +196,7 @@ def estimate_r(Y, X, A, r_max, c=1.,
         Results of the number of latent factors.
     '''
     a, d = A.shape[1], X.shape[1]
-    X = np.hstack((X, A))    
+    X = np.hstack((X, A))
     n, p = Y.shape
 
     Y, kwargs_glm, _ = _check_input(Y, X, family, disp_glm, disp_family, offset, None, **kwargs)

{causarray-0.0.3 → causarray-0.0.5}/causarray/gcate_glm.py

@@ -41,15 +41,29 @@ def fit_glm(Y, X, A=None, family='gaussian', disp_family='poisson',
     A : array
         n x 1 vector of treatments or None
     family : str
-        family of GLM to fit, can be one of: 'gaussian', 'poisson', 'nb'
+        Family of GLM to fit, can be one of: 'gaussian', 'poisson', 'nb'
     disp_glm : array or None
-        dispersion parameter for negative binomial GLM
-    return_df : bool
-        whether to return results as DataFrame
-    impute : bool
-        whether to impute potential outcomes and get predicted values
+        Dispersion parameter for negative binomial GLM.
+    impute : bool or None
+        Whether to impute missing values in Y.        
     offset : bool
-        whether to use log of sum of Y as offset
+        Whether to use log of sum of Y as offset.
+    shrinkage : bool
+        Whether to use regularized GLM.
+    alpha : float
+        Regularization parameter for regularized GLM.
+    maxiter : int
+        Maximum number of iterations for GLM fitting.
+    thres_disp : float
+        Threshold for dispersion parameter for negative binomial GLM.
+    n_jobs : int
+        Number of jobs to run in parallel.
+    random_state : int
+        Random seed for reproducibility.
+    verbose : bool
+        Whether to print progress messages.
+    kwargs : dict
+        Additional arguments to pass to GLM fitting.
 
     Returns
     -------

{causarray-0.0.3 → causarray-0.0.5}/causarray/gcate_opt.py

@@ -261,7 +261,7 @@ def alter_min(
     kwargs_ls['alpha'] = kwargs_ls['alpha']
     if verbose:
         pprint.pprint({'kwargs_glm':kwargs_glm,'kwargs_ls':kwargs_ls,'kwargs_es':kwargs_es}, compact=True)
-    pprint.pprint(f'Fitting GCATE (step {1 if P1 is None else 2})...')
+    pprint.pprint(f'Fitting GCATE (step {2 if P1 is None else 1})...')
     hist = [func_val_pre]
     es = Early_Stopping(**kwargs_es)
     with tqdm(np.arange(kwargs_es['max_iters']), disable=not verbose) as pbar:

{causarray-0.0.3 → causarray-0.0.5}/causarray/utils.py

@@ -49,6 +49,8 @@ def prep_causarray_data(Y, A, X=None, X_A=None, intercept=True):
     Y = np.minimum(Y, np.round(np.quantile(np.max(Y, 0), 0.999)))
     if not isinstance(A, pd.DataFrame):
         A = np.asarray(A)
+    if A.ndim == 1:
+        A = A[:, None]
 
     X = np.zeros((Y.shape[0], 0)) if X is None else np.asarray(X)        
     X_A = X if X_A is None else np.asarray(X_A)

causarray-0.0.3/causarray/__about__.py

@@ -1 +0,0 @@
-__version__ = "0.0.3"