causarray 0.0.1__py3-none-any.whl

causarray/gcate_glm.py

@@ -0,0 +1,256 @@
+from causarray.utils import *
+from causarray.utils import _filter_params
+import numpy as np
+import statsmodels as stats
+import statsmodels.api as sm
+from sklearn.linear_model import LinearRegression
+from joblib import Parallel, delayed
+from tqdm import tqdm
+
+warnings.filterwarnings('ignore')
+
+
+def init_inv_link(Y, family, disp):
+    if family=='gaussian':
+        val = Y/disp
+    elif family=='poisson':
+        val = np.log1p(Y)
+    elif family=='nb':
+        val = np.log1p(Y)
+    elif family=='binomial':
+        eps = (np.mean(Y, axis=0) + np.mean(Y, axis=1)) / 2 
+        val = np.log((Y + eps)/(disp - Y + eps))
+    else:
+        raise ValueError('Family not recognized')
+    return val
+
+
+
+def fit_glm(Y, X, A=None, family='gaussian', disp_family='poisson',
+    disp_glm=None, impute=False, offset=None, shrinkage=False,
+    alpha=1e-4, maxiter=1000, thres_disp=100., n_jobs=-3, random_state=0, verbose=False, **kwargs):
+    '''
+    Fit GLM to each column of Y, with covariate X and treatment A.
+
+    Parameters
+    ----------
+    Y : array
+        n x p matrix of outcomes
+    X : array
+        n x d matrix of covariates
+    A : array
+        n x 1 vector of treatments or None
+    family : str
+        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
+    offset : bool
+        whether to use log of sum of Y as offset    
+    '''
+    np.random.seed(random_state)
+    
+    if family not in ['gaussian', 'poisson', 'nb']:
+        raise ValueError('Family not recognized')
+   
+    d = X.shape[1]
+
+    if A is None:
+        a = 1 # dummy treatment
+        assert impute is False
+    else:
+        if A.ndim==1:
+            A = A[:,None]
+        if impute is not False and isinstance(impute, np.ndarray):
+            X_test = impute
+        else:
+            X_test = X
+        X_test = np.c_[X,np.zeros_like(A)]
+        X = np.c_[X,A]
+        # X_test = X.copy()
+
+        a = A.shape[1]
+
+    if offset is not None and offset is not False:
+        if type(offset)==bool and offset is True:
+            offsets = np.log(comp_size_factor(Y, **_filter_params(comp_size_factor, kwargs)))
+        else:
+            offsets = offset
+    else:
+        offsets = None
+
+    # estimate dispersion parameter for negative binomial GLM if not provided
+    if family=='nb' and disp_glm is None:
+        disp_glm = estimate_disp(Y, X, offset=offsets, disp_family=disp_family, maxiter=1000, verbose=verbose, **kwargs)
+    
+    alpha = np.full(X.shape[1], alpha)
+    pprint.pprint('Fitting {} GLM{}...'.format(family, '' if offsets is None else ' with offset'))
+    is_constant = np.all(X == X[0, :], axis=0)
+    alpha[is_constant] = 0
+    # alpha[:-a] = 0
+
+    families = {
+        'gaussian': lambda disp: sm.families.Gaussian(),
+        'poisson': lambda disp: sm.families.Poisson(),
+        'nb': lambda disp: sm.families.NegativeBinomial(alpha=1/disp)
+    }
+
+    def fit_model(j, Y, X, offsets, family, disp, impute, alpha):
+        if family=='nb' and disp[j]>thres_disp:
+            family = 'poisson'
+        try:            
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore")
+                glm_family = families.get(family, lambda: ValueError('family must be one of: "gaussian", "poisson", "nb"'))(disp_glm[j] if family == 'nb' else None)
+
+                try:
+                    if shrinkage:
+                        raise ValueError('fit regularized GLM')
+                    mod = sm.GLM(Y[:,j], X, family=glm_family, offset=offsets).fit(maxiter=maxiter)
+                    if not np.all(np.isfinite(mod.params)) or np.any(np.abs(mod.params[:d])>50) or np.any(np.abs(mod.params[d:])>10):
+                        raise ValueError('GLM did not converge')
+                    resid_deviance = mod.resid_deviance
+                except:
+                    mod = sm.GLM(Y[:,j], X, family=glm_family, offset=offsets).fit_regularized(alpha=alpha, cnvrg_tol=1e-5)
+                    resid_deviance = np.full(Y.shape[0], 0.)
+
+            B = mod.params
+
+            Yhat_0 = np.zeros((Y.shape[0], a))
+            Yhat_1 = np.zeros((Y.shape[0], a))
+            if impute is not False:
+                for j in range(a):
+                    X_test_copy = X_test.copy()
+                    Yhat_0[:,j] = mod.predict(X_test_copy, offset=offsets)                    
+                    X_test_copy[:, d+j] = 1  # Update the j-th column with all ones
+                    Yhat_1[:,j] = mod.predict(X_test_copy, offset=offsets)
+            else:
+                Yhat_0[:,:] = Yhat_1[:,:] = mod.predict(X, offset=offsets).reshape(-1, a)
+            
+        except:
+            pprint.pprint('Fitting GLM for column {} does not converge.'.format(j))
+            B = np.full(X.shape[1], 0.)
+            
+            Yhat_0 = np.full((Y.shape[0], a), 0)
+            Yhat_1 = np.full((Y.shape[0], a), 0)
+            if impute is not False:
+                for k in range(a):
+                    Yhat_0[:, k] = np.mean(Y[A[:, k] == 1, j])
+                    Yhat_1[:, k] = np.mean(Y[A[:, k] == 0, j])
+            resid_deviance = np.full(Y.shape[0], 0.)
+        return B, Yhat_0, Yhat_1, resid_deviance
+
+
+    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])))
+
+    B, Yhat_0, Yhat_1, resid_deviance = zip(*results)
+    B = np.array(B)    
+    Yhat_0 = np.array(Yhat_0).transpose(1, 0, 2)
+    Yhat_1 = np.array(Yhat_1).transpose(1, 0, 2)
+    resid_deviance = np.array(resid_deviance).T
+
+    if impute is not False:        
+        Yhat = (Yhat_0, Yhat_1)
+    else:
+        Yhat = np.array(Yhat_0)[:,:,0]
+    
+    return B, Yhat, disp_glm, offsets, resid_deviance
+
+
+def estimate_disp(Y, X=None, A=None, Y_hat=None, disp_family='gaussian', offset=None, verbose=False, **kwargs):
+    if offset is not None:
+        if type(offset)==bool and offset is True:
+            offsets = np.log(comp_size_factor(Y, **_filter_params(comp_size_factor, kwargs)))
+        else:
+            offsets = offset
+        sf = np.exp(offsets)[:,None]
+    else:
+        offsets = None
+        sf = 1.
+    
+
+    if Y_hat is None:        
+        if verbose:
+            pprint.pprint('Estimating dispersion parameter...')
+
+        if A is not None:
+            X = np.c_[X,A]
+
+        if disp_family=='gaussian':
+            Y_norm = Y/sf
+            reg = LinearRegression(fit_intercept=False).fit(X, Y_norm)
+            Y_hat = reg.predict(X)
+            # Y_hat = np.clip(Y_hat, 0, 1) * sf            
+        elif disp_family=='poisson':
+            Y_hat = fit_glm(Y, X, None, offset=offsets, family='poisson', impute=False, **kwargs)[1]      
+            Y_hat /= sf
+
+    # Clip Y_hat based on the range of Y per column
+    Y_hat = np.clip(Y_hat, 0., np.max(Y/sf, axis=0))
+
+    disp_glm = np.mean((Y/sf - Y_hat)**2 - Y_hat, axis=0) / np.mean(Y_hat**2, axis=0)
+    disp_glm = 1./np.clip(disp_glm, 0.01, 100.)
+    disp_glm[np.isnan(disp_glm)] = 1.
+
+    return disp_glm
+
+
+
+
+def loess_fit(Y, X, n_jobs=-3, **kwargs):
+    
+    def _loess_fit(y, x, **kwargs):
+        try:
+            from skmisc.loess import loess
+            l = loess(x, y, **kwargs)
+            l.fit()
+            pred = l.predict(x, stderror=True)
+            conf = pred.confidence()
+            pred, lower, upper = pred.values, conf.lower, conf.upper
+        except:
+            pred, lower, upper = np.full(y.shape[0], np.nan), np.full(y.shape[0], np.nan), np.full(y.shape[0], np.nan)
+
+        return pred, lower, upper
+
+    results = Parallel(n_jobs=n_jobs)(delayed(_loess_fit)(Y[:,j], X, **kwargs) for j in range(Y.shape[1]))
+
+    CATE, CATE_lower, CATE_upper = zip(*results)
+    CATE = np.array(CATE).T
+    CATE_lower = np.array(CATE_lower).T
+    CATE_upper = np.array(CATE_upper).T
+    return CATE, CATE_lower, CATE_upper
+
+
+
+def ls_fit(Y, X, n_jobs=-3, **kwargs):
+    
+    def _ls_fit(y, x, **kwargs):
+        # try:
+        model = sm.OLS(y, x)
+        result = model.fit(disp=False)
+
+        # Get the predicted values
+        pred = result.predict(x)
+
+        # Get the confidence intervals
+        conf = result.conf_int()    
+        pred, lower, upper = pred, np.full(y.shape[0], conf[0][0]), np.full(y.shape[0], conf[0][1])
+        # except:
+        #     pred, lower, upper = np.full(y.shape[0], np.nan), np.full(y.shape[0], np.nan), np.full(y.shape[0], np.nan)
+
+        return pred, lower, upper
+
+    results = Parallel(n_jobs=n_jobs)(delayed(_ls_fit)(Y[:,j], X, **kwargs) for j in range(Y.shape[1]))
+
+    CATE, CATE_lower, CATE_upper = zip(*results)
+    CATE = np.array(CATE).T
+    CATE_lower = np.array(CATE_lower).T
+    CATE_upper = np.array(CATE_upper).T
+    return CATE, CATE_lower, CATE_upper
+
+    
+    

causarray/gcate_likelihood.py

@@ -0,0 +1,143 @@
+import numpy as np
+# from scipy.special import expit, xlogy, xlog1py, logsumexp, factorial
+# from scipy.stats import binom, poisson, norm, nbinom
+
+import numba as nb
+from numba import njit, prange
+
+type_f = np.float64
+
+
+from scipy.special import xlogy, gammaln
+from numba.extending import get_cython_function_address
+import ctypes
+
+_PTR = ctypes.POINTER
+_dble = ctypes.c_double
+_ptr_dble = _PTR(_dble)
+
+addr = get_cython_function_address("scipy.special.cython_special", "gammaln")
+functype = ctypes.CFUNCTYPE(_dble, _dble)
+gammaln_float64 = functype(addr)
+
+@nb.vectorize
+def gammaln_nb(x):
+  return gammaln_float64(x)
+
+
+def log_h(y, family, nuisance):
+    if family=='nb':
+        return gammaln(y + nuisance) - gammaln(nuisance) - gammaln(y+1)
+    elif family=='poisson':
+        return - gammaln(y+1)
+
+
+@nb.vectorize
+def log1mexp(a):
+    '''
+    A numeral stable function to compute log(1-exp(a)) for a in [-inf,0].
+    '''
+    if(a >= -np.log(type_f(2.))):
+        return np.log(-np.expm1(a)) 
+    else:
+        return np.log1p(-np.exp(a))
+    
+
+@njit
+def nll(Y, A, B, family, nuisance=np.ones((1,1)), Tys=np.zeros((1,1)), thres_disp=10.
+    #  size_factor=np.ones((1,1))
+    ):
+    """
+    Compute the negative log likelihood for generalized linear models with optional nuisance parameters.
+    
+    Parameters:
+    Y : array-like of shape (n_samples, n_features)
+        The response variable.
+    A : array-like of shape (n_samples, n_factors)
+        The input data matrix.
+    B : array-like of shape (n_features, n_factors)
+        The input data matrix.
+    family : str, optional (default='gaussian')
+        The family of the generalized linear model. Options include 'poisson', and 'nb'.
+    nuisance : float or array-like of shape (n_samples,), optional (default=1)
+        The nuisance parameter for the family. For the Gaussian family, this is the variance; for the Poisson
+        family, this is the scaling factor; and for the negative binomial family, this is the overdispersion
+        parameter.
+    size_factor : float or array-like of shape (n_samples,), optional (default=1)
+        The size factor for the response variable.
+    
+    Returns:
+    nll : float
+        The negative log likelihood.
+    """
+    
+    Theta = A @ B.T
+    Ty = Y.copy()
+    n = Y.shape[0]
+    
+    if family == 'poisson':
+        Theta = np.clip(Theta, -np.inf, type_f(1e2))
+        b = np.exp(Theta)
+    elif family == 'nb':
+        Xi = np.clip(Theta, -np.inf, type_f(1e2))
+        exp_Xi = np.exp(Xi)
+        tmp = np.clip(1 / (type_f(1.) + exp_Xi / nuisance), 1e-6, 1-1e-6)
+
+        Theta = np.where(nuisance > thres_disp, Xi, np.log1p(-tmp))
+        b = np.where(nuisance > thres_disp, exp_Xi, - nuisance * np.log(tmp) #+ gammaln_nb(nuisance+Y) - gammaln_nb(nuisance)
+        ) # ignoring a common factor - gammaln_nb(Y)
+    else:
+        raise ValueError('Family not recognized')
+
+    nll = - np.sum(Ty * Theta - b + Tys) / type_f(n) # * size_factor to get back the likelihood
+
+    return nll
+
+
+
+@njit
+def grad(Y, A, B, family, nuisance=np.ones((1,1)), thres_disp=10.
+    #size_factor=np.ones((1,1)),
+        ):
+    """
+    Compute the gradient of log likelihood with respect to B
+    for generalized linear models with optional nuisance parameters.
+    
+    The natural parameter of Y is Theta = A @ B^T.
+    
+    Parameters:
+    Y : array-like of shape (n_samples, n_features)
+        The response variable.
+    A : array-like of shape (n_samples, n_factors)
+        The input data matrix.
+    B : array-like of shape (n_features, n_factors)
+        The input data matrix.
+    family : str, optional (default='gaussian')
+        The family of the generalized linear model. Options include 'poisson', and 'nb'.
+    nuisance : float or array-like of shape (n_samples,), optional (default=1)
+        The nuisance parameter for the family. For the Gaussian family, this is the variance; for the Poisson
+        family, this is the scaling factor; and for the negative binomial family, this is the overdispersion
+        parameter.
+    
+    Returns:
+    grad : array-like of shape (n_features, n_factors)
+        The gradient of log likelihood.
+    """
+    Theta = A @ B.T
+    Ty = Y.copy()
+    n = Y.shape[0]
+    
+    if family == 'nb':
+        Xi = np.clip(Theta, -np.inf, type_f(1e2))
+        b_p = np.exp(Xi)
+        tmp = np.clip(1 / (type_f(1.) + b_p / nuisance), 1e-6, 1-1e-6)
+        grad = - (Ty - b_p) * np.where(nuisance > thres_disp, 1., tmp) # * size_factor to get back the likelihood
+    elif family == 'poisson':
+        Theta = np.clip(Theta, -np.inf, type_f(1e2))
+        b_p = np.exp(Theta)
+        grad = - (Ty - b_p) # * size_factor to get back the likelihood
+    else:
+        raise ValueError('Family not recognized')
+    grad = grad.T @ A / type_f(n)
+    return grad
+

causarray/gcate_opt.py

@@ -0,0 +1,298 @@
+from causarray.gcate_likelihood import *
+from causarray.utils import *
+from causarray.gcate_glm import *
+import scipy as sp
+from scipy.sparse.linalg import svds
+from joblib import Parallel, delayed
+from collections import defaultdict
+
+
+@njit
+def project_norm_ball(X, radius):
+    """
+    Projects a vector X to the norm ball of a given radius.
+
+    Parameters:
+        X (ndarray): The vector to be projected.
+        radius (float): The radius of the norm ball.
+
+    Returns:
+        ndarray: The projected matrix.
+    """
+    norms = np.linalg.norm(X, ord=np.inf)
+    if norms > radius:
+        X *= radius / norms
+
+    return X
+
+
+
+@njit
+def line_search(Y, A, x0, g, d, family, nuisance, Ys, thres_disp, start,
+                lam, alpha=1., beta=0.5, max_iters=100, tol=1e-3, 
+                ):
+    """
+    Performs line search to find the step size that minimizes a given function.
+    
+    Parameters:
+        f (callable): A scalar-valued function of a vector argument.
+        x0 (ndarray): The starting point for the line search.
+        g (ndarray): The search direction vector.
+        alpha (float, optional): The initial step size. Default is 10.
+        beta (float, optional): The shrinkage factor used to reduce the step size. Default is 0.5.
+        max_iters (int, optional): The maximum number of iterations. Default is 100.
+        tol (float, optional): The tolerance for the step size. Default is 1e-3.
+    
+    Returns:
+        float: The step size that minimizes the function along the search direction.
+    """
+
+    # Evaluate the function at the starting point.
+    f0 = nll(Y, A, x0, family, nuisance, Ys, thres_disp)  + np.mean(lam * np.abs(x0[start:d]))
+    
+    # Initialize the step size.    
+    alpha = type_f(alpha)
+    beta = type_f(beta)
+    t = alpha
+    norm_g = np.linalg.norm(g)
+
+    # Iterate until the maximum number of iterations is reached or the step size is small enough.
+    for i in range(max_iters):
+        
+        # Compute the new point.
+        x1 = x0 - t*g
+        x1[start:d] = np.sign(x1[start:d]) * np.maximum(np.abs(x1[start:d]) - lam, type_f(0.))
+        
+        # Evaluate the function at the new point.
+        f1 = nll(Y, A, x1, family, nuisance, Ys, thres_disp) + np.mean(lam * np.abs(x1[start:d]))
+        if i==0:
+            f01 = f1
+
+        # Check if the function has decreased sufficiently.
+        if f1 < f0 - tol*t*norm_g:
+            return x1
+        if t<1e-4:
+            break
+        t *= beta
+
+    # Return the maximum step size.
+    if f01<1.1*f1:
+        x1 = x0 - alpha*g
+        x1[start:d] = np.sign(x1[start:d]) * np.maximum(np.abs(x1[start:d]) - lam, type_f(0.))
+    return x1
+
+
+
+
+
+
+@njit(parallel=True)
+def update(Y, A, B, d, lam, P1, P2,
+          family, nuisance, Ys, thres_disp, a, C,
+          alpha, beta, max_iters, tol):
+    n, p = Y.shape
+    
+    g = grad(Y.T, B, A, family, nuisance.T, thres_disp)
+    g[:, :d] = 0.
+    for i in prange(n):
+        g[i, d:] = project_norm_ball(g[i, d:], 2*C)
+
+        A[i, :] = line_search(Y[i, :], B, A[i, :], g[i, :], d, 
+                              family, nuisance[0], Ys[i, :], thres_disp, 0, 
+                              type_f(0.), alpha, beta, max_iters, tol)
+    if P1 is not None:
+        A[:, d:] = P1 @ A[:, d:]
+    
+    g = grad(Y, A, B, family, nuisance, thres_disp)
+    if P1 is not None:
+        g[:, :d] = 0.
+    elif P2 is not None:
+        g[:, d-a:d] = P2 @ g[:, d-a:d]
+        g[:, d:] = 0.
+        g[:, :d-a] = 0.
+
+    for j in prange(p):
+        if P2 is None:
+            g[j, d:] = project_norm_ball(g[j, d:], 2*C)
+        else:
+            g[j, :d] = project_norm_ball(g[j, :d], 2*C)
+
+        B[j, :] = line_search(Y[:, j], A, B[j, :], g[j, :], d, 
+                              family, nuisance[:,j], Ys[:, j], thres_disp, d-a,
+                              lam[j,:], alpha, beta, max_iters, tol
+                             )
+
+    if P2 is None:
+        B[:, d:] = np.clip(B[:, d:], -10., 10.)
+    func_val = nll(Y, A, B, family, nuisance, Ys, thres_disp)
+
+    return func_val, A, B
+
+
+def alter_min(
+    Y, r, X=None, P1=None, P2=None, A=None, B=None,
+    kwargs_glm={}, kwargs_ls={}, kwargs_es={}, 
+    lam=0., a=None, verbose=False, thres_disp=100.):
+    '''
+    Alternative minimization of latent factorization for generalized linear models.
+
+    Parameters
+    ----------
+    Y : array-like, shape (n, p)
+        Response matrix.
+    r : int
+        The number of latent factors.
+    X : array-like, shape (n, d)
+        Observed covariate matrix.
+    P1 : array-like, shape (n, n)
+        The projection matrix for the rows.
+    P2 : array-like, shape (p, p)
+        The projection matrix onto the orthogonal column space of Gamma.
+    A : array-like, shape (n, d+r)
+        The initial matrix for the covariate and latent factors. If None, initialize internally.
+    B : array-like, shape (p, d+r)
+        The initial matrix for the covariate and latent coefficients. If None, initialize internally.
+    kwargs_glm : dict
+        Keyword arguments for the generalized linear model.
+    kwargs_ls : dict
+        Keyword arguments for the line search.
+    kwargs_es : dict
+        Keyword arguments for the early stopping.
+    lam : float
+        The regularization parameter for the l1 norm of the coefficients.
+    a : int
+        The number of columns to be regularized. Assume the last 'a' columns of the covariates are the regularized coefficients. If 'a' is None, it is set to be 'd-offset' by default.
+    verbose : bool
+        The indicator of whether to print the progress.
+    thres_disp : float
+        The threshold for the dispersion parameter.
+
+    Returns
+    -------
+    res : dict
+        A dictionary containing the information of the optimization, including 
+            'A': the matrix (n, d+r) for the covariate and updated latent factors,
+            'B': the matrix (p, d+r) for the updated covariate and latent coefficients,
+            'U': the matrix (n, r) for the estimated confouders.
+    '''
+
+    n, p = Y.shape
+    d = X.shape[1]
+    
+    assert d>0
+    if verbose:
+        pprint.pprint({'n':n,'p':p,'d':d,'r':r})
+
+    kwargs_glm = {**{'family':'poisson', 'disp_glm':np.ones(p), 'size_factor':np.ones(n)
+        }, **kwargs_glm}
+    kwargs_ls = {**{'alpha':0.1, 'beta':0.5, 'max_iters':20, 'tol':1e-4, 'C':None}, **kwargs_ls}
+    kwargs_es = {**{'max_iters':500, 'warmup':0, 'patience':5, 'tolerance':1e-3}, **kwargs_es}
+
+    if kwargs_es['max_iters'] == 0:
+        return A, B, {'n_iter':0, 'func_val':0., 'resid':0., 'hist':[0.], 'kwargs_glm':kwargs_glm, 'kwargs_ls':kwargs_ls, 'kwargs_es':kwargs_es}
+    
+    family, nuisance, size_factor = kwargs_glm['family'], kwargs_glm['disp_glm'].astype(type_f), kwargs_glm['size_factor'].astype(type_f)
+    nuisance = nuisance.reshape(1,-1)
+    size_factor = size_factor.reshape(-1,1)
+
+    if kwargs_ls['C'] is None:
+        kwargs_ls['C'] = 1e3 if family=='nb' else 1e5
+    C = kwargs_ls['C']
+
+    if P1 is True:
+        Q, _ = sp.linalg.qr(X, mode='economic')
+        P1 = np.identity(n) - Q @ Q.T
+        P1 = P1.astype(type_f)
+            
+    if a is None:
+        a = d
+
+    # initialization for Theta = A @ B^T    
+    if A is None or B 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)
+        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.')
+
+        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]
+        
+        E = A[:, -r:] @ B[:, -r:].T
+        u, s, vh = sp.sparse.linalg.svds(E, k=r)        
+        A[:, d:] = u * s[None,:]**(1/2)
+        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)
+    Y /= size_factor
+    Ys = - gammaln_nb(Y+1)
+    if family=='nb':
+        Ys += np.where(nuisance > thres_disp, 0., gammaln_nb(nuisance+Y) - gammaln_nb(nuisance))    
+    
+    A = A.astype(type_f)
+    B = B.astype(type_f)
+    
+    lam = type_f(lam)
+    weights = lam / (np.abs(B[:, d-a:d]) + 1e-6)
+    
+    assert ~np.any(np.isnan(A))
+    assert ~np.any(np.isnan(B))
+    
+    t = 0
+    func_val_pre = (nll(Y, A, B, family, nuisance, Ys, thres_disp) + np.sum(np.abs(B[:,d-a:d]) * weights)) / p
+    func_val = func_val_pre
+
+    kwargs_ls['alpha'] = kwargs_ls['alpha']
+    if verbose:
+        pprint.pprint({'kwargs_glm':kwargs_glm,'kwargs_ls':kwargs_ls,'kwargs_es':kwargs_es}, compact=True)
+
+    hist = [func_val_pre]
+    es = Early_Stopping(**kwargs_es)
+    with tqdm(np.arange(kwargs_es['max_iters'])) as pbar:
+        for t in pbar:
+            func_val, A, B = update(
+                Y, A, B, d, weights, P1, P2,
+                family, nuisance, Ys, thres_disp, a, C,
+                kwargs_ls['alpha'], kwargs_ls['beta'], kwargs_ls['max_iters'], kwargs_ls['tol'], 
+            )
+            func_val = (func_val + np.sum(np.abs(B[:,d-a:d]) * weights)) / p
+            hist.append(func_val)
+            if not np.isfinite(func_val) or func_val>np.maximum(1e3*np.abs(func_val_pre),1e3):
+                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.close()
+                break
+            else:
+                func_val_pre = func_val
+            pbar.set_postfix(nll='{:.02f}'.format(func_val))
+
+    kwargs_glm['disp_glm'] = nuisance[0]
+    res = {'n_iter':t, 'func_val':func_val, 'resid':func_val_pre - func_val,
+           'hist':hist, 'kwargs_glm':kwargs_glm, 'kwargs_ls':kwargs_ls, 'kwargs_es':kwargs_es}
+    res['X_U'] = A; res['B_Gamma'] = B; res['U'] = A[:,d:]
+    return res
+
+
+