File size: 2,636 Bytes
8081b08 | 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 | import math
import os
import pickle
import crick
import pandas as pd
from loguru import logger
def get_filename() -> str:
current_file_name = os.path.basename(__file__)
log_name = "{}.log"
return log_name.format(current_file_name.split('.')[0])
def set_log() -> None:
filename = get_filename()
logger.add(f'../log/{filename}')
set_log()
gene_info_table_path = '/scDifformer/data/240412_test/material/gene_info_table.csv'
gene_info = pd.read_csv(gene_info_table_path)
func_gene_list = [i for i in gene_info[(gene_info["gene_type"] == "protein_coding") | (gene_info["gene_type"] == "miRNA")]["ensembl_id"]]
home_path = r'/scDifformer/data/240412_test/tdigest/'
gene_median_dict_path = "/scDifformer/data/240412_test/material/gene_median_dictionary.pkl"
def combine_digests(rootdir: str) -> None:
# merge new tdigests into total tdigest dict
def merge_digest(dict_key_ensembl_id, dict_value_tdigest, new_tdigest_dict):
new_gene_tdigest = new_tdigest_dict.get(dict_key_ensembl_id)
if new_gene_tdigest is not None:
dict_value_tdigest.merge(new_gene_tdigest)
return dict_value_tdigest
elif new_gene_tdigest is None:
return dict_value_tdigest
# initiate tdigests
median_digests = [crick.tdigest.TDigest() for _ in range(len(func_gene_list))]
total_tdigest_dict = dict(zip(func_gene_list, median_digests))
for subdir, dirs, files in os.walk(rootdir):
for file in files:
if file.endswith(".gene_median_digest_dict.pickle"):
logger.info(f"{rootdir}{file} is walking")
with open(f"{rootdir}{file}", "rb") as fp:
tdigest_dict = pickle.load(fp)
total_tdigest_dict = {k: merge_digest(k, v, tdigest_dict) for k, v in total_tdigest_dict.items()}
# save dict of merged tdigests
# with open(f"/mnt/nfs/geneformer/data/pkl/total_gene_tdigest_dict.pkl", "wb") as fp:
# pickle.dump(total_tdigest_dict, fp)
# extract medians and save dict
total_median_dict = {k: v.quantile(0.5) for k, v in total_tdigest_dict.items()}
# with open(f"/mnt/nfs/geneformer/data/pkl/total_gene_median_dict.pkl", "wb") as fp:
# pickle.dump(total_median_dict, fp)
# save dict of only detected genes' medians
detected_median_dict = {k: v for k, v in total_median_dict.items() if not math.isnan(v)}
with open(gene_median_dict_path, "wb") as fp:
pickle.dump(detected_median_dict, fp)
del detected_median_dict, total_median_dict, total_tdigest_dict
if __name__ == '__main__':
combine_digests(rootdir=home_path)
|