causarray 0.0.2__tar.gz → 0.0.4__tar.gz

{causarray-0.0.2 → causarray-0.0.4}/.gitignore

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

{causarray-0.0.2 → causarray-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: causarray
-Version: 0.0.2
+Version: 0.0.4
 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.2 → causarray-0.0.4}/causarray/DR_estimation.py

@@ -33,7 +33,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 +55,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.
 
@@ -69,7 +75,7 @@ def cross_fitting(
     pi_hat : array
         Estimated propensity score.
     '''
-    func_ps, params_ps = _get_func_ps(ps_model, verbose=verbose, **kwargs)
+    func_ps, params_ps = _get_func_ps(ps_model, verbose=False, **kwargs)
     params_glm = _filter_params(fit_glm, {**kwargs, 'verbose': verbose})
 
     if verbose:
@@ -86,7 +92,8 @@ def cross_fitting(
     # Initialize lists to store results
     fit_pi = True if pi_hat is None else False
     pi_hat = np.zeros_like(A, dtype=float) if fit_pi else pi_hat
-    Y_hat = np.zeros((Y.shape[0],Y.shape[1],A.shape[1],2), dtype=float)
+    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 cross-fitting
     for train_index, test_index in folds:
@@ -103,7 +110,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)
@@ -111,17 +123,15 @@ def cross_fitting(
                     pi[A_train[:,j] == 0., j] = 1 - prob
                 else:
                     pi[:,j] = func_ps(XA_train[i_cells], A_train[i_cells][:,j], XA_test)
-
-        if verbose: pprint.pprint('Fit outcome models...')
-        # 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)
-        
-        # Store results
-        if fit_pi: pi_hat[test_index] = pi
-
-        Y_hat[test_index,:,:,0] = res[1][0]
-        Y_hat[test_index,:,:,1] = res[1][1]
+            pi_hat[test_index] = pi
+
+        if fit_Y:
+            if verbose: pprint.pprint('Fit outcome models...')
+            # 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)
+            Y_hat[test_index,:,:,0] = res[1][0]
+            Y_hat[test_index,:,:,1] = res[1][1]
 
     pi_hat = np.clip(pi_hat, 0.01, 0.99)
     Y_hat = np.clip(Y_hat, None, 1e5)

{causarray-0.0.2 → causarray-0.0.4}/causarray/DR_inference.py

@@ -5,16 +5,14 @@ from scipy.stats import norm
 from statsmodels.stats.multitest import multipletests, fdrcorrection
 
 
-def multiplier_bootstrap(resid, var_est, B):
+def multiplier_bootstrap(eta, B):
     '''
     Multiplier bootstrap for inference.
 
     Parameters
     ----------
-    resid : array-like
-        [n, p] Residuals.
-    var_est : array-like
-        [p,] Variance of the parameter estimates.
+    eta : array-like
+        [n, p] Influence function values scaled by inversion of standard deviation.
     B : int
         Number of bootstrap samples.
 
@@ -23,17 +21,13 @@ def multiplier_bootstrap(resid, var_est, B):
     z_init : array-like
         [B, p] Bootstrap statistics for each hypothesis.
     '''
-    n, p = resid.shape
+    n, p = eta.shape
     B = int(B)
     z_init = np.zeros((B,p))
-    eta = resid / np.sqrt(n * var_est[None,:])
     for b in range(B):
         g = np.random.normal(size=n)
         z_init[b, :] = np.sum(eta * g[:, None], axis=0)
 
-    # g = np.random.normal(size=(B,n,1))
-    # z_init = np.sum(resid[None, :,:] * g, axis=1) / np.sqrt(n * var_est[None,:])
-
     return z_init
     
 
@@ -106,9 +100,9 @@ def augmentation(V, tvalues, c):
 
 
 def fdx_control(
-    tau_est, var_est, tvalues_init, eta_est,
+    tau_est, tvalues_init, eta_est, std_est,
     fdx, B, alpha, c, 
-    min_var=1e-4, min_diff=0.1
+    min_var=1e-8, min_diff=0.1
     ):
     '''
     Perform FDX control.
@@ -117,12 +111,12 @@ def fdx_control(
     ----------
     tau_est : array-like
         The estimated causal effect.    
-    var_est : array-like
-        The estimated variance.
     tvalues_init : array-like
         The estimated t-values.
     eta_est : array-like
-        The estimated influence function values.
+        The estimated influence function values, scaled by inversion of standard deviation.
+    std_est : array-like
+        The estimated standard deviation.
     fdx : bool
         If True, perform FDX control.
     B : int
@@ -141,8 +135,8 @@ def fdx_control(
     n = eta_est.shape[0]
 
     if fdx:
-        id_test = var_est >= min_var
-        z_init = multiplier_bootstrap(eta_est, var_est, B)
+        id_test = std_est >= min_var
+        z_init = multiplier_bootstrap(eta_est / std_est[None,:], B)
         z_init[:, ~id_test] = 0.
         tvalues = tvalues_init.copy()
         tvalues[~id_test] = 0.

{causarray-0.0.2 → causarray-0.0.4}/causarray/DR_learner.py

@@ -11,7 +11,7 @@ from causarray.utils import reset_random_seeds, pprint, tqdm, comp_size_factor,
 def compute_causal_estimand(
     estimand,
     Y, W, A, W_A=None, family='nb', offset=False,    
-    Y_hat=None, pi_hat=None, eps_var=1e-3,
+    Y_hat=None, pi_hat=None, mask=None,
     fdx=False, fdx_B=1000, fdx_alpha=0.05, fdx_c=0.1,     
     verbose=False, random_state=0, **kwargs):
     '''
@@ -39,6 +39,10 @@ def compute_causal_estimand(
         Predicted outcomes under treatment of shape (n, p, a, 2).
     pi_hat : array, optional
         Predicted propensity scores of shape (n, a).
+    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.
 
     fdx : bool
         Whether to use FDX control, P(FDP > c) < alpha.
@@ -61,19 +65,16 @@ 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:
-        A = A.reshape(-1,1)
+    if A.ndim == 1: A = A[:, None]
     if isinstance(A, pd.DataFrame):
         trt_names = A.columns
         A = A.values
@@ -87,7 +88,16 @@ def compute_causal_estimand(
         W_A = W
     elif isinstance(W_A, pd.DataFrame):
         W_A = W_A.values
-    
+
+    if mask is not None:
+        mask = np.array(mask).astype(bool)
+        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]
@@ -103,10 +113,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...')
@@ -117,25 +126,25 @@ def compute_causal_estimand(
     # normalize the influence function values
     etas /= size_factors[:,None,None,None]
 
-
-    i_ctrl = (np.sum(A, axis=1) == 0.)
-
     res = []
     iters = range(A.shape[1]) if A.shape[1]==1 else tqdm(range(A.shape[1]))
     for j in iters:
-        i_case = (A[:,j] == 1.)
-        i_cells = i_ctrl | i_case
-        n_cells = np.sum(i_cells)        
-        eta_est, tau_est, var_est = estimand(etas[i_cells,:,j], **kwargs)
+        if mask is not None:
+            i_cells = mask[:, j]
+        else:
+            i_ctrl = (np.sum(A, axis=1) == 0.)
+            i_case = (A[:,j] == 1.)
+            i_cells = i_ctrl | i_case
+        eta_est, tau_est, var_est = estimand(etas[i_cells,:,j], A[i_cells,j], **kwargs)
 
-        std_est = np.sqrt((var_est + eps_var)/ n_cells)
+        std_est = np.sqrt(var_est)
         tvalues_init = tau_est / std_est
 
         # Multiple testing procedure
-        V = fdx_control(tau_est, var_est, tvalues_init, eta_est, fdx, fdx_B, fdx_alpha, fdx_c)
+        V = fdx_control(tau_est, tvalues_init, eta_est, std_est, fdx, fdx_B, fdx_alpha, fdx_c)
 
         # BH correction
-        tvalues_init[np.isinf(var_est)] = np.nan
+        tvalues_init[np.isinf(std_est)] = np.nan
         pvals, qvals, pvals_adj, qvals_adj = bh_correction(tvalues_init)
         
         df_res = pd.DataFrame({
@@ -159,7 +168,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, 
+    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,
     fdx=False, fdx_alpha=0.05, fdx_c=0.1,     
     verbose=False, **kwargs):
@@ -187,12 +196,21 @@ def LFC(
         Predicted propensity scores of shape (n, a).
     cross_est : bool
         Whether to use cross-estimation.
+    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.
+    usevar : str
+        The method to use for estimating the variance of treatment effects. 
+        Options are 'pooled' (default) or 'unequal'.
     
     thres_min : float
         The minimum threshold for the treatment effect.
     thres_diff : float
         The minimum threshold for the difference in treatment effect.
-    
+    eps_var : float
+        The minimum threshold for the variance of treatment.
+
     fdx : bool
         Whether to use FDX control, P(FDP > c) < alpha.
     fdx_alpha : float
@@ -211,7 +229,7 @@ def LFC(
         Dataframe of test results.
     '''
 
-    def estimand(etas, **kwargs):
+    def estimand(etas, A, **kwargs):
         eta_0, eta_1 = etas[..., 0], etas[..., 1]
         tau_0, tau_1 = np.mean(eta_0, axis=0), np.mean(eta_1, axis=0)
 
@@ -219,7 +237,18 @@ def LFC(
         tau_0 = np.clip(tau_0, thres_diff, None)
         tau_est = np.log(tau_1/tau_0)
         eta_est = eta_1 / tau_1[None,:] -  eta_0 / tau_0[None,:]
-        var_est = np.var(eta_est, axis=0, ddof=1)
+
+        if usevar == 'pooled':
+            var_est = (np.var(eta_est, axis=0, ddof=1) + eps_var) / eta_est.shape[0]
+        elif usevar == 'unequal':
+            # Estimate the variance using Welch's t-test
+            var_0 = np.var(eta_est[A==0], axis=0, ddof=1)
+            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
+        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)
@@ -229,7 +258,7 @@ def LFC(
 
     return compute_causal_estimand(
         estimand, Y, W, A, W_A, family, offset,    
-        Y_hat=Y_hat, pi_hat=pi_hat,
+        Y_hat=Y_hat, pi_hat=pi_hat, mask=mask,
         fdx=fdx, fdx_alpha=fdx_alpha, fdx_c=fdx_c, verbose=verbose, **kwargs)
 
 

causarray-0.0.4/causarray/__about__.py

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

{causarray-0.0.2 → causarray-0.0.4}/causarray/gcate.py

@@ -82,7 +82,8 @@ def fit_gcate(Y, X, A, r, family='nb', disp_glm=None, disp_family=None, offset=T
     kwargs : dict
         Additional keyword arguments.
     '''
-
+    if X.ndim == 1: X = X[:, None]
+    if A.ndim == 1: A = A[:, None]
     X = np.hstack((X, A))
     a = A.shape[1]
     Y, kwargs_glm, lam1 = _check_input(Y, X, family, disp_glm, disp_family, offset, c1, **kwargs)    
@@ -195,8 +196,10 @@ def estimate_r(Y, X, A, r_max, c=1.,
     df_r : DataFrame
         Results of the number of latent factors.
     '''
-    X = np.hstack((X, A))
+    if X.ndim == 1: X = X[:, None]
+    if A.ndim == 1: A = A[:, None]
     a, d = A.shape[1], X.shape[1]
+    X = np.hstack((X, A))    
     n, p = Y.shape
 
     Y, kwargs_glm, _ = _check_input(Y, X, family, disp_glm, disp_family, offset, None, **kwargs)
@@ -210,22 +213,31 @@ def estimate_r(Y, X, A, r_max, c=1.,
         r_list = np.arange(1, int(r_max)+1)
     else:
         r_list = np.array(r_max, dtype=int)
-        
-    for r in r_list:
-        res_1, res_2 = estimate(Y, X, r, a,
-            0, kwargs_glm, kwargs_ls_1, kwargs_es_1, kwargs_ls_2, kwargs_es_2, **kwargs)
-        A01, A02, A1, A2 = res_1['X_U'], res_1['B_Gamma'], res_2['X_U'], res_2['B_Gamma']
-
-        logh = log_h(Y, family, nuisance)
-
-        if r==1:
-            ll = 2 * ( 
-                nll(Y, A01, A02, family, nuisance, size_factor) / p 
-                - np.sum(logh) / (n*p) ) 
-            nu = (d+a) * np.maximum(n,p) * np.log(n * p / np.maximum(n,p)) / (n*p)
-            jic = ll + c * nu
-            res.append([0, ll, nu, jic])
-        
+    r_max = np.max(r_list)
+
+    # Estimate the residual deviance
+    res_glm = fit_glm(Y, X, offset=np.log(size_factor[:,0]), family=family, disp_glm=nuisance[0], maxiter=100, verbose=False)
+    u, s, vt = svds(res_glm[-1], k=r_max)
+    if u.shape[1]<r_max:
+        raise ValueError(f'The number of latent factors is larger than the rank of deviance residuals ({u.shape[1]}). Try to decrease the value of r.')
+    Q, _ = sp.linalg.qr(X, mode='economic')
+    P1 = np.identity(n) - Q @ Q.T
+    P1 = P1.astype(type_f)
+    A1 = np.c_[X, P1 @ u]
+
+    logh = log_h(Y, family, nuisance)
+    ll = 2 * ( 
+        nll(Y, X, res_glm[0], family, nuisance, size_factor) / p 
+        - np.sum(logh) / (n*p) ) 
+    nu = (d+a) * np.maximum(n,p) * np.log(n * p / np.maximum(n,p)) / (n*p)
+    jic = ll + c * nu
+    res.append([0, ll, nu, jic])
+
+    for r in r_list[::-1]:
+        _, res_2 = estimate(Y, X, r, a,
+            0, kwargs_glm, kwargs_ls_1, kwargs_es_1, kwargs_ls_2, kwargs_es_2, A=A1[:,:d+a+r], **kwargs)
+        A1, A2 = res_2['X_U'], res_2['B_Gamma']
+
         ll = 2 * ( 
             nll(Y, A1, A2, family, nuisance, size_factor) / p 
             - np.sum(logh) / (n*p) ) 
@@ -233,5 +245,5 @@ def estimate_r(Y, X, A, r_max, c=1.,
         jic = ll + c * nu
         res.append([r, ll, nu, jic])
 
-    df_r = pd.DataFrame(res, columns=['r', 'deviance', 'nu', 'JIC'])
+    df_r = pd.DataFrame(res, columns=['r', 'deviance', 'nu', 'JIC']).sort_values(by='r')
     return df_r 

{causarray-0.0.2 → causarray-0.0.4}/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
     -------
@@ -156,8 +170,7 @@ def fit_glm(Y, X, A=None, family='gaussian', disp_family='poisson',
 
 
     results = Parallel(n_jobs=n_jobs)(delayed(fit_model)(
-        j, Y, X, offsets, family, disp_glm, impute, alpha) for j in tqdm(range(Y.shape[1])))
-    pprint.pprint('Fitting GLM done.')
+        j, Y, X, offsets, family, disp_glm, impute, alpha) for j in tqdm(range(Y.shape[1]), disable=not verbose))
     if verbose: pprint.pprint('Fitting GLM done.')
 
     B, Yhat_0, Yhat_1, resid_deviance = zip(*results)

{causarray-0.0.2 → causarray-0.0.4}/causarray/gcate_opt.py

@@ -208,19 +208,24 @@ def alter_min(
         a = d
 
     # initialization for Theta = A @ B^T    
-    if A is None or B is None:
+    if A is None:
         if verbose:
             pprint.pprint('Estimating initial latent variables with GLMs...')
-        res_glm = fit_glm(Y, X, offset=np.log(size_factor[:,0]), family=family, disp_glm=nuisance[0], maxiter=100)
+        res_glm = fit_glm(Y, X, offset=np.log(size_factor[:,0]), family=family, disp_glm=nuisance[0], maxiter=100, verbose=verbose)
         u, s, vt = svds(res_glm[-1], k=r)
 
         if u.shape[1]<r:
             raise ValueError(f'The number of latent factors is larger than the rank of deviance residuals ({u.shape[1]}). Try to decrease the value of r.')
+        
+        A = np.c_[X, P1 @ u]
+    else:
+        assert A.shape[1] == d+r
 
+    if B is None:
         if verbose:
             pprint.pprint('Estimating initial coefficients with GLMs...')
-        A = np.c_[X, P1 @ u]
-        B = fit_glm(Y, A, offset=np.log(size_factor[:,0]), family=family, disp_glm=nuisance[0], maxiter=100)[0]
+        
+        B = fit_glm(Y, A, offset=np.log(size_factor[:,0]), family=family, disp_glm=nuisance[0], maxiter=100, verbose=verbose)[0]
         
         E = A[:, -r:] @ B[:, -r:].T
         u, s, vh = sp.sparse.linalg.svds(E, k=r)        
@@ -228,19 +233,10 @@ def alter_min(
         B[:, d:] = vh.T * s[None,:]**(1/2)
         del E, u, s, vh
 
-        # if offset==1:
-        #     scale = np.sqrt(np.median(np.abs(X[:,0])))
-        #     B[:, :offset] = scale
-        #     A[:, :offset] /= scale 
 
     if P2 is not None:
         P2 = P2.astype(type_f)
-        # E = A[:,d-a:] @ B[:,d-a:].T @ (np.identity(p) - P2)
-        # u, s, vh = sp.sparse.linalg.svds(E, k=r)
         B[:, d-a:d] = P2 @ B[:, d-a:d]
-        # A[:, d:] = u * s[None,:]**(1/2)
-        # B[:, d:] = vh.T * s[None,:]**(1/2)
-        # del E, u, s, vh
     
 
     Y = Y.astype(type_f)
@@ -265,10 +261,10 @@ 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})...')
     hist = [func_val_pre]
     es = Early_Stopping(**kwargs_es)
-    with tqdm(np.arange(kwargs_es['max_iters'])) as pbar:
+    with tqdm(np.arange(kwargs_es['max_iters']), disable=not verbose) as pbar:
         for t in pbar:
             func_val, A, B = update(
                 Y, A, B, d, weights, P1, P2,
@@ -281,7 +277,7 @@ def alter_min(
                 pprint.pprint('Encountered large or infinity values. Try to decrease the value of C for the norm constraints.')
                 break
             elif es(func_val):
-                pbar.set_postfix_str('Early stopped.' + es.info)
+                pbar.set_postfix_str('Early stopped. ' + es.info)
                 pbar.close()
                 break
             else:

{causarray-0.0.2 → causarray-0.0.4}/causarray/utils.py

@@ -49,16 +49,19 @@ 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)    
-    intercept_col = np.ones((X.shape[0], 1)) if intercept else np.empty((X.shape[0], 0))
-    X = np.hstack((intercept_col, X))
-
+    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)
     loglibsize = np.log2(np.sum(np.asarray(Y), axis=1))
     loglibsize = (loglibsize - np.mean(loglibsize)) / np.std(loglibsize, ddof=1)
     X_A = np.hstack((X_A, loglibsize[:, None]))
 
+    intercept_col = np.ones((X.shape[0], 1)) if intercept else np.empty((X.shape[0], 0))
+    X = np.hstack((intercept_col, X))
+    X_A = np.hstack((intercept_col, X_A))
+
     return Y, A, X, X_A
 
 

causarray-0.0.2/causarray/__about__.py

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