causarray 0.0.4__py3-none-any.whl → 0.0.6__py3-none-any.whl

causarray/DR_estimation.py

@@ -1,9 +1,13 @@
 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
+import causarray.gcate_glm as _gcate_glm  # module-qualified so _USE_FAST_BACKEND changes take effect at call time
 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
@@ -82,10 +86,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]))]
 
@@ -93,7 +102,29 @@ def cross_fitting(
     fit_pi = True if pi_hat is None else False
     pi_hat = np.zeros_like(A, dtype=float) if fit_pi else pi_hat
     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
+    if fit_Y:
+        _yhat_gb = Y.shape[0] * Y.shape[1] * A.shape[1] * 2 * 8 / 1e9
+        _mem_limit_gb = kwargs.get('mem_limit_gb', None)
+        if _mem_limit_gb is not None and _yhat_gb > _mem_limit_gb:
+            import warnings
+            warnings.warn(
+                f"Y_hat allocation ({_yhat_gb:.1f} GB as float64) exceeds "
+                f"mem_limit_gb={_mem_limit_gb} GB; using float32 to halve peak memory.",
+                ResourceWarning, stacklevel=3,
+            )
+            Y_hat = np.zeros((Y.shape[0], Y.shape[1], A.shape[1], 2), dtype=np.float32)
+        else:
+            Y_hat = np.zeros((Y.shape[0], Y.shape[1], A.shape[1], 2), dtype=float)
+
+    # 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:
@@ -127,9 +158,19 @@ def cross_fitting(
 
         if fit_Y:
             if verbose: pprint.pprint('Fit outcome models...')
+            # Subset offset to training fold (for fitting) and test fold (for
+            # imputation) when it is a pre-computed array, so that
+            # ``fit_glm_auto`` receives arrays with matching leading
+            # dimensions in both stages.
+            params_glm_fold = params_glm
+            offset_test_arr = None
+            if 'offset' in params_glm and isinstance(params_glm['offset'], np.ndarray):
+                params_glm_fold = dict(params_glm)
+                params_glm_fold['offset'] = params_glm['offset'][train_index]
+                offset_test_arr = params_glm['offset'][test_index]
             # Fit GLM on training data and predict on test data
-            res = fit_glm(Y_train, X_train, A_train, family=family, alpha=glm_alpha,
-                impute=X_test, **params_glm)
+            res = _gcate_glm.fit_glm_auto(Y_train, X_train, A_train, family=family, alpha=glm_alpha,
+                impute=X_test, offset_test=offset_test_arr, **params_glm_fold)
             Y_hat[test_index,:,:,0] = res[1][0]
             Y_hat[test_index,:,:,1] = res[1][1]
 
@@ -178,8 +219,6 @@ def AIPW_mean(Y, A, mu, pi, positive=False):
     tau = np.mean(pseudo_y, axis=0)
 
     return tau, pseudo_y
-    
-
 
 
 
@@ -188,51 +227,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])
 
 
 
@@ -252,11 +329,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/DR_inference.py

@@ -150,7 +150,7 @@ def fdx_control(
     return V
 
 
-def bh_correction(tvalues_init):
+def bh_correction(tvalues_init, df=None):
     '''
     Perform BH correction.
 
@@ -179,7 +179,10 @@ def bh_correction(tvalues_init):
     qvals_adj = np.full(tvalues_init.shape, np.nan)
 
     if np.sum(idx) > 0:
-        pvals[idx] = sp.stats.norm.sf(np.abs(tvalues_init[idx]))*2
+        if df is not None:
+            pvals[idx] = sp.stats.t.sf(np.abs(tvalues_init[idx]), df=df[idx]) * 2
+        else:
+            pvals[idx] = sp.stats.norm.sf(np.abs(tvalues_init[idx])) * 2
         qvals[idx] = multipletests(pvals[idx], alpha=0.05, method='fdr_bh')[1]
 
         # BH correction with empirical null adjustment
@@ -189,7 +192,10 @@ def bh_correction(tvalues_init):
         if mad>0:
             tvalues_init_adj = (tvalues_init - med) / mad
 
-            pvals_adj[idx] = sp.stats.norm.sf(np.abs(tvalues_init_adj[idx]))*2
+            if df is not None:
+                pvals_adj[idx] = sp.stats.t.sf(np.abs(tvalues_init_adj[idx]), df=df[idx]) * 2
+            else:
+                pvals_adj[idx] = sp.stats.norm.sf(np.abs(tvalues_init_adj[idx])) * 2
             qvals_adj[idx] = multipletests(pvals_adj[idx], alpha=0.05, method='fdr_bh')[1]
     
     return pvals, qvals, pvals_adj, qvals_adj