File size: 8,432 Bytes
3290550 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 |
import torch
import click
import cooler
import warnings
import numpy as np
from torch import nn
from tqdm import tqdm
from torch.cuda.amp import autocast
from importlib_resources import files
from polaris.utils.util_loop import bedpewriter
from polaris.model.polarisnet import polarisnet
from scipy.sparse import coo_matrix
from scipy.sparse import SparseEfficiencyWarning
warnings.filterwarnings("ignore", category=SparseEfficiencyWarning)
def getLocal(mat, i, jj, w, N):
if i >= 0 and jj >= 0 and i+w <= N and jj+w <= N:
mat = mat[i:i+w,jj:jj+w].toarray()
# print(f"global: {mat.shape}")
return mat[None,...]
# pad_width = ((up, down), (left, right))
slice_pos = [[i, i+w], [jj, jj+w]]
pad_width = [[0, 0], [0, 0]]
if i < 0:
pad_width[0][0] = -i
slice_pos[0][0] = 0
if jj < 0:
pad_width[1][0] = -jj
slice_pos[1][0] = 0
if i+w > N:
pad_width[0][1] = i+w-N
slice_pos[0][1] = N
if jj+w > N:
pad_width[1][1] = jj+w-N
slice_pos[1][1] = N
_mat = mat[slice_pos[0][0]:slice_pos[0][1],slice_pos[1][0]:slice_pos[1][1]].toarray()
padded_mat = np.pad(_mat, pad_width, mode='constant', constant_values=0)
# print(f"global: {padded_mat.shape}",slice_pos, pad_width)
return padded_mat[None,...]
def upperCoo2symm(row,col,data,N=None):
# print(np.max(row),np.max(col),N)
if N:
shape=(N,N)
else:
shape=(row.max() + 1,col.max() + 1)
sparse_matrix = coo_matrix((data, (row, col)), shape=shape)
symm = sparse_matrix + sparse_matrix.T
diagVal = symm.diagonal(0)/2
symm = symm.tocsr()
symm.setdiag(diagVal)
return symm
def processCoolFile(coolfile, cchrom):
extent = coolfile.extent(cchrom)
N = extent[1] - extent[0]
ccdata = coolfile.matrix(balance=True, sparse=True, as_pixels=True).fetch(cchrom)
ccdata['balanced'] = ccdata['balanced'].fillna(0)
ccdata['bin1_id'] -= extent[0]
ccdata['bin2_id'] -= extent[0]
ccdata['distance'] = ccdata['bin2_id'] - ccdata['bin1_id']
d_means = ccdata.groupby('distance')['balanced'].transform('mean')
ccdata['oe'] = ccdata['balanced'] / d_means
ccdata['oe'] = ccdata['oe'].fillna(0)
ccdata['oe'] = ccdata['oe'] / ccdata['oe'].max()
oeMat = upperCoo2symm(ccdata['bin1_id'].ravel(), ccdata['bin2_id'].ravel(), ccdata['oe'].ravel(), N)
return oeMat, N
@click.command()
@click.option('--batchsize', type=int, default=16, help='Batch size [16]')
@click.option('--cpu', type=bool, default=False, help='Use CPU [False]')
@click.option('--gpu', type=str, default=None, help='Comma-separated GPU indices [auto select]')
@click.option('--chrom', type=str, default=None, help='Comma separated chroms')
@click.option('--max_distance', type=int, default=3000000, help='Max distance (bp) between contact pairs')
@click.option('--resol',type=int,default=500,help ='Resolution')
@click.option('--image',type=int,default=1024,help ='Resolution')
@click.option('--center_size',type=int,default=224,help ='Resolution')
@click.option('-i','--input', type=str,required=True,help='Hi-C contact map path')
@click.option('-o','--output', type=str,required=True,help='.bedpe file path to save loop candidates')
def dev(batchsize, cpu, gpu, chrom, max_distance, resol, input, output, image, center_size):
""" *development function* Coming soon...
"""
print('polaris loop dev START :) ')
# center_size = 224
# center_size = image // 2
start_idx = (image - center_size) // 2
end_idx = (image + center_size) // 2
slice_obj_pred = (slice(None), slice(None), slice(start_idx, end_idx), slice(start_idx, end_idx))
slice_obj_coord = (slice(None), slice(start_idx, end_idx), slice(start_idx, end_idx))
max_distance_bin=max_distance//resol
loopwriter = bedpewriter(output,resol,max_distance)
if cpu:
assert gpu is None, "\033[91m QAQ The CPU and GPU modes cannot be used simultaneously. Please check the command. \033[0m\n"
gpu = ['None']
device = torch.device("cpu")
print('Using CPU mode... (This may take significantly longer than using GPU mode.)')
else:
if torch.cuda.is_available():
if gpu is not None:
print("Using the specified GPU: " + gpu)
gpu=[int(i) for i in gpu.split(',')]
device = torch.device(f"cuda:{gpu[0]}")
else:
gpuIdx = torch.cuda.current_device()
device = torch.device(gpuIdx)
print("Automatically selected GPU: " + str(gpuIdx))
gpu=[gpu]
else:
device = torch.device("cpu")
gpu = ['None']
cpu = True
print('GPU is not available!')
print('Using CPU mode... (This may take significantly longer than using GPU mode.)')
coolfile = cooler.Cooler(input + '::/resolutions/' + str(resol))
modelstate = str(files('polaris').joinpath('model/sft_loop.pt'))
_modelstate = torch.load(modelstate, map_location=device.type)
parameters = _modelstate['parameters']
if chrom is None:
chrom =coolfile.chromnames
else:
chrom = chrom.split(',')
for rmchr in ['chrMT','MT','chrM','M','Y','chrY','X','chrX']: # 'Y','chrY','X','chrX'
if rmchr in chrom:
chrom.remove(rmchr)
print(f"\nAnalysing chroms: {chrom}")
model = polarisnet(
image_size=parameters['image_size'],
in_channels=parameters['in_channels'],
out_channels=parameters['out_channels'],
embed_dim=parameters['embed_dim'],
depths=parameters['depths'],
channels=parameters['channels'],
num_heads=parameters['num_heads'],
drop=parameters['drop'],
drop_path=parameters['drop_path'],
pos_embed=parameters['pos_embed']
).to(device)
model.load_state_dict(_modelstate['model_state_dict'])
if not cpu and len(gpu) > 1:
model = nn.DataParallel(model, device_ids=gpu)
model.eval()
chrom = tqdm(chrom, dynamic_ncols=True)
for _chrom in chrom:
chrom.desc = f"[analyzing {_chrom}]"
oeMat, N = processCoolFile(coolfile, _chrom)
start_point = -(image - center_size) // 2
joffset = np.repeat(np.linspace(0, image, image, endpoint=False, dtype=int)[np.newaxis, :], image, axis=0)
ioffset = np.repeat(np.linspace(0, image, image, endpoint=False, dtype=int)[:, np.newaxis], image, axis=1)
data, i_list, j_list = [], [], []
for i in range(start_point, N - image - start_point, center_size):
for j in range(0, max_distance_bin, center_size):
jj = j + i
# if jj + w <= N and i + w <= N:
_oeMat = getLocal(oeMat, i, jj, image, N)
if np.sum(_oeMat == 0) <= (image*image*0.9):
data.append(_oeMat)
i_list.append(i + ioffset)
j_list.append(jj + joffset)
while len(data) >= batchsize or (i + center_size > N - image - start_point and len(data) > 0):
bin_i = torch.tensor(np.stack(i_list[:batchsize], axis=0)).to(device)
bin_j = torch.tensor(np.stack(j_list[:batchsize], axis=0)).to(device)
targetX = torch.tensor(np.stack(data[:batchsize], axis=0)).to(device)
bin_i = bin_i*resol
bin_j = bin_j*resol
data = data[batchsize:]
i_list = i_list[batchsize:]
j_list = j_list[batchsize:]
print(targetX.shape)
print(bin_i.shape)
print(bin_j.shape)
with torch.no_grad():
with autocast():
pred = torch.sigmoid(model(targetX.float().to(device)))[slice_obj_pred].flatten()
loop = torch.nonzero(pred>0.5).flatten().cpu()
prob = pred[loop].cpu().numpy().flatten().tolist()
frag1 = bin_i[slice_obj_coord].flatten().cpu().numpy()[loop].flatten().tolist()
frag2 = bin_j[slice_obj_coord].flatten().cpu().numpy()[loop].flatten().tolist()
loopwriter.write(_chrom,frag1,frag2,prob)
if __name__ == '__main__':
dev() |