dpacman/data_tasks/download/genome.py

"""
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()