""" Script for downloading the genome, hg38 """ import rootutils root = rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True) import os import logging import requests import json import hydra from omegaconf import DictConfig from pathlib import Path import logging import multiprocessing from hydra.core.hydra_config import HydraConfig base_logger = logging.getLogger(__name__) def get_all_chroms( genome: str = "hg38", exclude: list = None, logger: logging.Logger = None, include: list = None, ): """ Fetch all chromosome names for a genome. Note: some chromosomes are in unexpected formats (e.g. there is 'chr15', but also 'chr15_ML143371v1_fix') """ if logger is None: logger = logging.getLogger(__name__) url = f"https://api.genome.ucsc.edu/list/chromosomes?genome={genome}" try: r = requests.get(url) r.raise_for_status() except: logger.error(f"Failed to fetch all chromosomes for genome {genome}") if include is not None and exclude is not None: raise ValueError(f"Must pass EITHER exclude or include. Cannot pass both.") all_chroms = list(r.json()["chromosomes"].keys()) if include is not None: logger.info(f"Including only the following chromosomes: {include}") all_chroms = [chrom for chrom in all_chroms if chrom in include] if exclude is not None: logger.info(f"Excluding the following chromosomes: {exclude}") all_chroms = [chrom for chrom in all_chroms if not (chrom in exclude)] logger.info(f"Found {len(all_chroms)} chromosomes in genome {genome}.") return all_chroms def get_all_chrom_fasta_files( genome: str = "hg38", exclude: list = None, include: list = None, logger: logging.Logger = None, output_dir="../../data_files/raw/genomes", ): """ Get FASTA files for each chromosome for a current genome """ if logger is None: logger = logging.getLogger(__name__) if include is not None and exclude is not None: raise ValueError(f"Must pass EITHER exclude or include. Cannot pass both.") chroms = get_all_chroms( genome=genome, exclude=exclude, include=include, logger=logger ) logger.info(f"Saving downloaded chromosomes to {output_dir}") os.makedirs(output_dir, exist_ok=True) for chrom in chroms: chrom_save_path = os.path.join(output_dir, f"{genome}_{chrom}.json") if not (os.path.exists(chrom_save_path)): url = f"https://api.genome.ucsc.edu/getData/sequence?genome={genome};chrom={chrom}" try: r = requests.get(url) r.raise_for_status() json_output = r.json() with open(chrom_save_path, "w") as f: json.dump(json_output, f, indent=4) logger.info( f"Downloaded {chrom} in genome {genome}. Saved to: {chrom_save_path}" ) except: logger.error(f"Failed to fetch all {chrom} for genome {genome}") else: logger.info(f"Already downloaded {chrom} in genome {genome}. Skipping.") logger.info(f"Downloaded {len(chroms)} chromosomes in genome {genome}.") return chroms def merge_completed_files(genome: str, logs_dir: Path): """ Merge all completed_worker_*.txt files into a single completed.txt file """ merged_path = os.path.join(logs_dir, "completed.txt") with open(merged_path, "w") as outfile: outfile.write("chrom\trow_count\n") # header for fname in os.listdir(logs_dir): if fname.startswith("completed_worker_") and fname.endswith(".txt"): with open(os.path.join(logs_dir, fname), "r") as infile: for line in infile: if line.startswith("chrom"): # skip header lines continue outfile.write(line) print(f"Merged completed_worker_*.txt into {merged_path}") def worker(args): """ Worker function for parallel processing """ # Extract args chrom_group, idx, genome, logs_dir, output_dir = args os.makedirs(logs_dir, exist_ok=True) # Define logger wlogger = logging.getLogger(f"worker_{idx}") wlogger.setLevel(logging.DEBUG) wlogger.propagate = False log_file = os.path.join(logs_dir, f"worker_{idx}.log") fh = logging.FileHandler(log_file, mode="w", encoding="utf-8") fh.setLevel(logging.DEBUG) formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s") fh.setFormatter(formatter) wlogger.addHandler(fh) wlogger.info(f"Starting worker {idx} for chromosomes: {chrom_group}") all_chroms = get_all_chrom_fasta_files( genome=genome, include=chrom_group, logger=wlogger, output_dir=output_dir ) wlogger.info(f"Finished worker {idx}") def parallel_extract( genome: str, include: list = None, exclude: list = None, output_dir: Path = None, logs_dir: Path = None, ): """ Run extract_tfbs_with_context in parallel for groups of chromosomes in the genome to speed up processing. """ # Get all chromosomes whose sequences we want to download. chroms = get_all_chroms(genome, exclude=exclude, include=include) num_cores = multiprocessing.cpu_count() - 1 # Separate primary vs accessory chromosomes primary_chroms = [c for c in chroms if "_" not in c] accessory_chroms = [c for c in chroms if "_" in c] base_logger.info(f"Total primary chromosomes: {len(primary_chroms)}") for pc in primary_chroms: base_logger.info(pc) base_logger.info(f"Total accessory chromosomes: {len(accessory_chroms)}") for ac in accessory_chroms: base_logger.info(ac) # Distribute primary chromosomes round-robin across workers chunks = [[] for _ in range(num_cores)] for i, chrom in enumerate(primary_chroms): chunks[i % num_cores].append(chrom) # Now add accessory chromosomes to the least-loaded chunk (by count) for chrom in accessory_chroms: min_idx = min(range(num_cores), key=lambda i: len(chunks[i])) chunks[min_idx].append(chrom) # Log how we split it up - want to see which chromosomes are in which chunks. logging.info( f"{num_cores} CPU cores available (leaving 1 empty). Primary chromosomes distributed round-robin." ) for chunk_no, chunk in enumerate(chunks): logging.info(f"Chunk {chunk_no}. Chromosomes = {','.join(chunk)}") args_list = [ (chunk, i, genome, logs_dir, output_dir) for i, chunk in enumerate(chunks) ] with multiprocessing.Pool(processes=num_cores) as pool: pool.map(worker, args_list) merge_completed_files(genome, logs_dir) def main(cfg: DictConfig): include = cfg.get("include", None) exclude = cfg.get("exclude", None) output_dir = Path(root) / cfg.data_task.output_dir os.makedirs(output_dir, exist_ok=True) # Download the sequences of all chromosomes for genome in cfg.data_task.genomes: base_logger.info(f"Downloading all chromsoomes for genome {genome}") # Make a subfolder for this specific genome and its logs genome_output_dir = output_dir / genome genome_logs_dir = Path(HydraConfig.get().run.dir) / genome / "logs" parallel_extract( genome, include=include, exclude=exclude, output_dir=genome_output_dir, logs_dir=genome_logs_dir, ) if __name__ == "__main__": main()