e_csc = e_coo.tocsc()
e_csc_chunk = e_csc[:, (rank * p_chunk_size):((rank + 1) * p_chunk_size)]
e_coo_chunk = e_csc_chunk.tocoo()
e_values = TType(e_coo_chunk.data)
e_rows = torch.LongTensor(e_coo_chunk.row)
e_cols = torch.LongTensor(e_coo_chunk.col)
if e_values.is_cuda:
e_rows = e_rows.cuda()
e_cols = e_cols.cuda()
e_indices = torch.stack([e_rows, e_cols], dim=1).t()
e_shape = e_coo_chunk.shape
e_size = torch.Size([int(e_shape[0]), int(e_shape[1])])
e_chunk = TType_sp(e_indices, e_values, e_size).t()
if not args.dense:
e_dist = THDistMat.from_chunks(e_chunk).t()
else:
e_chunk = e_chunk.to_dense().t()
e_dist = THDistMat.from_chunks(e_chunk, force_bycol=True)
G_coo = sparse.coo_matrix(
(datafile['G_values'],
(datafile['G_indices'][0, :], datafile['G_indices'][1, :])),
shape=datafile['G_shape'])
G_csr = G_coo.tocsr()
G_csr_chunk = G_csr[(rank * p_chunk_size):((rank + 1) * p_chunk_size), :]
G_coo_chunk = G_csr_chunk.tocoo()
G_values = TType(G_coo_chunk.data).type(TType)
G_rows = torch.LongTensor(G_coo_chunk.row)
G_cols = torch.LongTensor(G_coo_chunk.col)
if G_values.is_cuda:
X_chunk = G[:, start_ind:end_ind].data.compute()
else:
X_chunk = da.hstack([G[:,start_ind:].data, da.zeros((n, 6))]).compute()
X_chunk[:, -11:] = pheno[:, 1:p_pheno + 1]
from utils import impute_na
X_chunk = impute_na(X_chunk)
# normalize
if args.normalize:
X_chunk -= X_chunk.mean(0)
X_chunk /= X_chunk.std(0)
X_chunk = torch.tensor(X_chunk)
print(X_chunk.shape)
X = THDistMat.from_chunks(X_chunk, force_bycol=True)
time = torch.tensor(pheno[:, 12]).view(-1, 1).type(TType)
delta = torch.tensor(pheno[:, 13]).view(-1, 1).type(TType)
print(args.lambd)
#X = distmat.distgen_normal(p, n, TType=TType, set_from_master=True)
#torch.manual_seed(seed+100)
#delta = torch.multinomial(torch.tensor([1., 1.]), n, replacement=True).float().view(-1, 1).type(TType)
if args.datnormest:
cox_driver = cox.COX(X, delta, float(args.lambd), nonsnps=11, time=time, TType=TType, sigma=1.0/(2*float(args.datnormest)**2))
elif args.quicknorm:
cox_driver = cox.COX(X, delta, float(args.lambd), nonsnps=11, time=time, TType=TType, sigma='quicknorm')
else:
cox_driver = cox.COX(X, delta, float(args.lambd), nonsnps=11, time=time, TType=TType, sigma='power')
cox_driver.run(int(args.iter), tol=float(args.tol),check_interval=int(args.step), check_obj=True)