File size: 7,574 Bytes
e3814d7 | 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 | # This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# Based on the HiCARN implementation, a script to generate data file for multichannel matrices. Also fix bugs to ensure locuses near edges are covered.
# References:
# HiCARN: https://github.com/OluwadareLab/HiCARN
# --------------------------------------------------------
import sys
from typing import Any
import numpy as np
from dataset_informations import *
import argparse
from utils import *
from tqdm import tqdm
def data_divider_parser():
parser = argparse.ArgumentParser(description='Divide data for train and predict')
req_args = parser.add_argument_group('Required Arguments')
req_args.add_argument('-c', dest='cell_line', help='REQUIRED: Cell line for analysis[example:GM12878]',
required=True)
req_args.add_argument('-hr', dest='high_res', help='REQUIRED: High resolution specified[example:10kb]',
default='10kb', choices=res_map.keys())
req_args.add_argument('-lr', dest='low_res', help='REQUIRED: Low resolution specified[example:40kb]',
default='10kb_d16_seed0')
req_args.add_argument('-s', dest='dataset', help='REQUIRED: Dataset for train/valid/predict(all)',
choices=set_dict.keys(), required = True )
misc_args = parser.add_argument_group('Miscellaneous Arguments')
misc_args.add_argument('-trs', '--transform-names', dest='transform_names', type=str, help='List of transforms used. Group transforms should be in the correct order(i.e. the next of Lp should be Lr)',
nargs='+', default = ['HiC', 'OE', '01TAD', 'Lp', 'Lr'])
misc_args.add_argument('-pref', '--save-prefix', dest='save_prefix', type=str, help='The prefix to save generated data',
default = 'Multi')
misc_args.add_argument('--chunk', dest='chunk', help='REQUIRED: chunk size for dividing[example:40]',
default=40, type=int, )
misc_args.add_argument('--stride', dest='stride', help='REQUIRED: stride for dividing[example:40]',
default=40, type=int, )
misc_args.add_argument('--diagonal-stride', type=int, help='Allow the submatrices moving along the diagonal to obatin more data. Default: 40, equals to the stride, which means no additional move along diagonal',
default = 40 )
misc_args.add_argument('--bound', dest='bound', help='REQUIRED: distance boundary interested[example:201]',
default=200, type=int, )
return parser
def divide_multichannel(mat, chr_num, chunk_size=40, stride=40, bound=200, diagonal_stride = 40, species='hsa', verbose=False):
"""
Dividing method.
"""
chr_str = str(chr_num)
if isinstance(chr_num, str): chr_num = except_chr[species][chr_num]
result = []
index = []
if len(mat.shape) <= 2:
mat = np.expand_dims(mat, axis=0)
channel, size, _ = mat.shape
# if (diagonal_stride < chunk_size and padding):
# pad_len = (chunk_size - diagonal_stride) // 2
# mat = np.pad(mat, ((0,0), (pad_len, pad_len), (pad_len, pad_len)), 'constant')
_, height, width = mat.shape
assert height == width, 'Now, we just assumed matrix is squared!'
assert diagonal_stride <= stride
for i in range(0, height - chunk_size + diagonal_stride, diagonal_stride):
i = i if i + chunk_size <= height else height - chunk_size
for j in range(i, i + bound - chunk + 2*stride, stride):
j = min(j, width - chunk_size)
subImage = mat[..., i:i + chunk_size, j:j + chunk_size]
result.append(subImage)
index.append((chr_num, size, i, j))
if j + chunk_size >= width : break
for j in range(i - stride, i - bound + chunk - 2 * stride, - stride):
j = max(j, 0)
subImage = mat[..., i:i + chunk_size, j:j + chunk_size]
result.append(subImage)
index.append((chr_num, size, i, j))
if j <= 0 : break
result = np.array(result)
if verbose: print(
f'[Chr{chr_str}] Deviding HiC matrix ({channel}x{size}x{size}) into {len(result)} samples with chunk={chunk_size}, '
f'stride={stride}, bound={bound} and diagonal_stride={diagonal_stride}')
index = np.array(index)
return result, index
def carn_divider(n,
high_file,
down_file,
chunk=40,
stride=40,
bound=200,
diagonal_stride = 40):
hic_data = np.load(high_file)
down_data = np.load(down_file)
compact_idx = hic_data['compact']
norm = hic_data['norm']
full_size = hic_data['hic'].shape[-1]
hic = hic_data['hic']
down_hic = down_data['hic']
print(f'[Chr{n}]File loaded.')
# Compacting
hic = compactM(hic, compact_idx)
down_hic = compactM(down_hic, compact_idx)
print(f'[Chr{n}]Compacted.')
# Deviding and Pooling
div_dhic, div_inds = divide_multichannel(down_hic, n, chunk, stride, bound, diagonal_stride)
div_hhic, _ = divide_multichannel(hic, n, chunk, stride, bound, diagonal_stride, verbose=True)
print(f'[Chr{n}]Finished.')
return n, div_dhic, div_hhic, div_inds, compact_idx, norm, full_size
if __name__ == '__main__':
parser = data_divider_parser()
args = parser.parse_args(sys.argv[1:])
cell_line = args.cell_line
high_res = args.high_res
low_res = args.low_res
dataset = args.dataset
chunk = args.chunk
stride = args.stride
bound = args.bound
diagonal_stride = args.diagonal_stride
trs = args.transform_names
prefix = args.save_prefix
chr_list = set_dict[dataset]
abandon_chromosome = abandon_chromosome_dict.get(cell_line, [])
postfix = cell_line.lower() if dataset == 'all' else dataset
print(f'Going to read {high_res} and {low_res} data with {trs}, then deviding matrices')
# pool_num = 23 if multiprocessing.cpu_count() > 23 else multiprocessing.cpu_count()
out_dir = os.path.join(root_dir, data_dir)
mkdir(out_dir)
# start = time.time()
# pool = multiprocessing.Pool(processes=pool_num)
# print(f'Start a multiprocess pool with processes = {pool_num} for generating HiCARN data')
results = []
for n in tqdm(chr_list):
if n in abandon_chromosome:
continue
high_file = os.path.join(data_dir, f'chr{n}_{high_res}.npz')
down_file = os.path.join(data_dir, f'chr{n}_{low_res}.npz')
kwargs = {'chunk':chunk, 'stride':stride, 'bound':bound, 'diagonal_stride' : diagonal_stride}
res = carn_divider(n, high_file, down_file, **kwargs)
results.append(res)
# return: n, div_dhic, div_hhic, div_inds, compact_idx, full_size
data = np.concatenate([r[1] for r in results])
target = np.concatenate([r[2] for r in results])
inds = np.concatenate([r[3] for r in results])
compacts = {r[0]: r[4] for r in results}
norms = {r[0]: r[5] for r in results}
sizes = {r[0]: r[6] for r in results}
filename = f'{prefix}_{high_res}{low_res}_c{chunk}_s{stride}_ds{diagonal_stride}_b{bound}_{cell_line}_{postfix}.npz'
datafile = os.path.join(out_dir, filename)
np.savez_compressed(datafile, data=data, target=target, inds=inds, compacts=compacts, norms=norms, sizes=sizes)
print('Saving file:', datafile)
|