//  [python]
def train(self, max_iters = 50, n_intervals = 1, lamb = 2.1e-4, alpha = 2, rho = 1.7, beta = 0, theta_init_offset = 0.1, njobs = 1):
        n_batches = int(np.ceil(self.ntimes/self.batchsize))
        for batch in range(n_batches):
            print("start running batch " + str(batch))
            start_time = time.time()
            # select a minibatch, and load to cuda
            start_idx = batch * self.batchsize
            if batch < n_batches - 1:
                end_idx = (batch + 1) * self.batchsize
                w_empir_cov = self.w_empir_cov[start_idx:end_idx, :, :].to(self.device)
                if self.mask.shape[0] == self.ntimes:
                    mask = self.mask[start_idx:end_idx, :, :].to(self.device)
                else:
                    mask = self.mask.to(self.device)
            else:
                w_empir_cov = self.w_empir_cov[start_idx:, :, :].to(self.device)
                if self.mask.shape[0] == self.ntimes:
                    mask = self.mask[start_idx:, :, :].to(self.device)
                else:
                    mask = self.mask.to(self.device)
            # initialize mini-batch, Z of the shape (batch_size, ngenes, ngenes)
            Z = torch.diag_embed(1/(torch.diagonal(w_empir_cov, offset=0, dim1=-2, dim2=-1) + theta_init_offset))
            # make Z positive definite matrix
            ll = torch.linalg.cholesky(Z)
            Z = torch.matmul(ll, ll.transpose(-1, -2))
            U = torch.zeros(Z.shape).to(self.device)
            I = torch.eye(self.ngenes).expand(Z.shape).to(self.device)

            it = 0
            its = []
            losses = []
            losses1 = []
            losses2 = []
            losses3 = []
            # hyper-parameter for batches
            if rho is None:
                updating_rho = True
                # rho of the shape (ntimes, 1, 1)
                b_rho = torch.ones((Z.shape[0], 1, 1)).to(self.device) * 1.7
            else:
                b_rho = torch.FloatTensor([rho] * Z.shape[0])[:, None, None].to(self.device)
                updating_rho = False
            b_alpha = alpha 
            b_beta = beta
            b_lamb = lamb
            while(it < max_iters): 
                # Primal 
                Y = U - Z + w_empir_cov/b_rho    # (ntimes, ngenes, ngenes)
                thetas = - 0.5 * Y + torch.stack([torch_sqrtm(mat) for mat in (torch.transpose(Y,1,2) @ Y * 0.25 + I/b_rho)])
                Z_pre = Z.detach().clone()
                # over-relaxation
                thetas = b_alpha * thetas + (1 - b_alpha) * Z_pre            
                Z = torch.sign(thetas + U) * torch.max((b_rho * (thetas + U).abs() - b_lamb)/(b_rho + b_beta * mask), torch.Tensor([0]).to(self.device))

                # Dual
                U = U + thetas - Z

                # calculate residual
                # primal_residual and dual_residual of the shape (ntimes, 1, 1)
                primal_residual = torch.sqrt((thetas - Z).pow(2).sum(1).sum(1))
                dual_residual = b_rho.squeeze() * torch.sqrt((Z - Z_pre).pow(2).sum(1).sum(1))

                # updating rho, rho should be of shape (ntimes, 1, 1)
                if updating_rho:
                    mask_inc = (primal_residual > 10 * dual_residual)
                    b_rho[mask_inc, :, :] = b_rho[mask_inc, :, :] * 2
                    mask_dec = (dual_residual > 10 * primal_residual)
                    b_rho[mask_dec, :, :] = b_rho[mask_dec, :, :] / 2
                
                # print(rho.squeeze())
                # free-up memory
                del Z_pre
                
                # Stopping criteria
                if (it + 1) % n_intervals == 0:
                    # calculate sum of all duality gap
                    # loss = self.neg_lkl_loss(thetas, w_empir_cov).sum() + b_lamb * Z.abs().sum() + b_beta * (self.mask * Z).pow(2).sum()
                    # primal_val = loss  + rho/2 * (thetas - Z).pow(2).sum()
                    # dual_val = loss + rho/2 * (thetas - Z + U).pow(2).sum() - rho/2 * U.pow(2).sum()
                    # duality_gap = primal_val - dual_val

                    # simplify min of all duality gap
                    duality_gap = b_rho.squeeze() * torch.stack([torch.trace(mat) for mat in torch.bmm(U.permute(0,2,1), Z - thetas)])
                    duality_gap = duality_gap.abs()

                    # loss function
                    loss1 = self.neg_lkl_loss(Z, w_empir_cov).sum()
                    loss2 = Z.abs().sum()
                    loss3 = (mask * Z).pow(2).sum()
                    loss = loss1 + b_lamb * loss2 + b_beta * loss3
                    losses.append(loss.item())
                    losses1.append(loss1.item())
                    losses2.append(loss2.item())
                    losses3.append(loss3.item())
                    its.append(it)
                    print("n_iter: {}, duality gap: {:.4e}, primal residual: {:.4e}, dual residual: {:4e}, loss1: {:4e}, loss2: {:4e}, loss3: {:.4e}".format(it+1, duality_gap.max().item(), primal_residual.max().item(), dual_residual.max().item(), loss1.item(), loss2.item(), loss3.item()))
                    
                    # if duality_gap < 1e-8:
                    #     break
                    primal_eps = 1e-3 # 1e-6
                    dual_eps = 1e-3 # 1e-6
                    if (primal_residual.max() < primal_eps) and (dual_residual.max() < dual_eps):
                        break                
                it += 1
            
            loss1 = self.neg_lkl_loss(Z, w_empir_cov).sum()
            loss2 = Z.abs().sum()
            loss3 = (mask * Z).pow(2).sum()
            print("Batch loss: loss1: {:.5f}, loss2: {:.5f}, loss3: {:.5f}".format(loss1.item(), loss2.item(), loss3.item()))

            # # plot loss curves
            # fig = plt.figure(figsize = (15,5))  
            # ax = fig.subplots(nrows = 1, ncols = 4)
            # ax[0].plot(its, losses)
            # ax[1].plot(its, losses1)
            # ax[2].plot(its, losses2)
            # ax[3].plot(its, losses3)
            # plt.tight_layout()

            # store values
            if batch < n_batches - 1:
                self.thetas[start_idx:end_idx] = Z.detach().cpu().numpy()
            else:
                self.thetas[start_idx:] = Z.detach().cpu().numpy()
            del thetas, U, I, Y, ll, Z

            print("Finished running batch {:d}, running time: {:.2f} sec".format(batch, time.time() - start_time))

            # calculate the partial correlation matrix
            print("calculate the partial correlation matrix")
            self.Gs = construct_weighted_G(self.thetas, njobs = njobs)

            # remove target-target if beta > 0, after the partial correlation matrix is calculated
            if beta > 0:
                for t, G in enumerate(self.Gs):
                    self.Gs[t] = G * (1 - self.mask[0].detach().cpu().numpy())
                for t, G in enumerate(self.Gs):
                    assert np.sum(G * self.mask[0].detach().cpu().numpy()) == 0            

        return self.Gs