# 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)