api_download

.gitignore

@@ -1 +1,2 @@
-dpacman/data_files
+dpacman/data_files
+dpacman/data/tfclust/*.log

dpacman/data/tfclust/api_download.py

@@ -0,0 +1,448 @@
+import requests
+from time import sleep
+import json
+import logging
+import multiprocessing
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import os
+import pandas as pd
+
+def get_all_tfs(genome: str = "hg38"):
+    """
+    Get all the transcription factors from the appropriate encRegTfbsClusteredWithCells.genome.bed file.
+    Available in data_files/raw/tfclust for genomes hg38 and hg19
+    """
+    # Read raw file
+    raw_data = pd.read_csv(
+        "../../data_files/encode3TfbsClusteredWithCells.bed", sep="\t", header=None
+    )
+    raw_data.columns = ["chrom", "start", "end", "tf_name", "score", "cell_line"]
+
+    # Extract all unique TF names
+    all_tfs = encode_raw["tf_name"].unique().tolist()
+    logging.info(f"Found {len(all_tfs)} transcription factors in genome {genome}.")
+
+    return all_tfs
+
+
+def get_all_chroms(genome: str = "hg38", exclude: list=None, include: list=None, logger: logging.Logger=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:
+        raise ValueError(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 = [chrom for chrom in r.json()["chromosomes"]]
+    if include:
+        logger.info(f"Including only the following chromosomes: {include}")
+        all_chroms = [chrom for chrom in all_chroms if chrom in include]
+    if exclude:
+        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 fetch_tfbs_track(chrom: str, genome: str = "hg38", logger:logging.Logger=None):
+    """
+    Fetch raw data from the track encRegTfbsClustered.
+    Returns json data for the specified chromosome, where key information appears as follows:
+    "encRegTfbsClustered": [
+        {
+            "bin": 585,
+            "chrom": "chr1",
+            "chromStart": 9917,
+            "chromEnd": 10247,
+            "name": "NUFIP1",
+            "score": 680,
+            "sourceCount": 1,
+            "sourceIds": "1063",
+            "sourceScores": "680"
+        },...
+        ]
+
+    """
+    if logger is None:
+        logger = logging.getLogger(__name__)
+        
+    params = {"genome": genome, "track": "encRegTfbsClustered", "chrom": chrom}
+    url = f"https://api.genome.ucsc.edu/getData/track?genome={params['genome']};track={params['track']};chrom={params['chrom']}"
+    try:
+        r = requests.get(url)
+        r.raise_for_status()
+    except:
+        raise ValueError(
+            f"Failed to fetch encRegTfbsClustered for {chrom} in genome {genome}"
+        )
+
+    # Extract the output and save it
+    json_out_dir = f"../../data_files/raw/tfclust/encRegTfbsClustered_data/{genome}"
+    os.makedirs(json_out_dir, exist_ok=True)
+
+    # Save it
+    json_output = r.json()
+    with open(
+        f"{json_out_dir}/{params['genome']}_{params['track']}_{params['chrom']}.json",
+        "w",
+    ) as f:
+        json.dump(json_output, f, indent=4)
+
+    logger.info(
+        f"Saved to {json_out_dir}/{params['genome']}_{params['track']}_{params['chrom']}.json"
+    )
+    return json_output
+
+
+def get_sequence(
+    chrom: str,
+    start: int,
+    end: int,
+    flank5: int = 0,
+    flank3: int = 0,
+    genome: str = "hg38",
+    logger: logging.Logger=None
+):
+    """
+    Given genome, start position, end position, chromosome, and desired flank size, extract the raw DNA sequence
+    """
+    if logger is None:
+        logger = logging.getLogger(__name__)
+        
+    new_start = max(0, start - flank5)
+    new_end = end + flank3
+    region = f"{chrom}:{new_start}-{new_end}"
+    url = f"https://api.genome.ucsc.edu/getData/sequence?genome={genome};chrom={chrom};start={new_start};end={new_end}"
+    
+    try:
+        r = requests.get(url)
+        r.raise_for_status()
+    except:
+        raise ValueError(f"Failed to fetch sequence for {region} in genome {genome}")
+
+    results_dict = {
+        "chromStart": new_start,
+        "chromEnd": new_end,
+        "seq": r.json()["dna"],
+    }
+    return results_dict
+
+
+def extract_tfbs_with_context(
+    genome: str = "hg38",
+    flank5: int = 500,
+    flank3: int = 500,
+    control_run: bool = True,  # if there's a flank, whether to also run without flank
+    out_dir: str = "../../data_files/processed/tfclust",
+    allowed_tfs: list = None,  # e.g., ['CTCF', 'MAX']
+    chroms: list = None,
+    logger: logging.Logger = None 
+):
+    """
+    Loop through raw downloads and extract TF binding sites (bs) with flanks
+    Builds a DataFrame with all the available data for each TF. Columns = ["bin", "chrom", "chromStart", "chromEnd", "name", "score", "scoreCount", "sourceIds", "sourceScores", "seq", "seq_flanked", "chromStart_flanked", "chromEnd_flanked"]
+    """
+    # Prepare logger
+    if logger is None:
+        logger = logging.getLogger(__name__)
+    # Prepare to save output
+    os.makedirs(out_dir, exist_ok=True)
+
+    # Get chromosomes
+    if chroms is None:
+        logger.info(
+            "No chromosomes provided, fetching all chromosomes for the given genome..."
+        )
+        chroms = get_all_chroms(genome, logger = logger)
+    count = 0
+    # Initialize the final DF
+    results_cols = [
+        "bin",
+        "chrom",
+        "chromStart",
+        "chromEnd",
+        "name",
+        "score",
+        "scoreCount",
+        "sourceIds",
+        "sourceScores",
+        "seq",
+        "seq_flanked",
+        "chromStart_flanked",
+        "chromEnd_flanked",
+        "flank5",
+        "flank3",
+    ]
+    results_init = pd.DataFrame(columns=results_cols)
+
+    # Make a list of the types of runs we need
+    queries = [{"flank5": flank5, "flank3": flank3}]
+    if not ((flank5 == 0) and (flank3 == 0) and control_run):
+        queries.append({"type": "control", "flank5": 0, "flank3": 0})
+        queries[0]["type"] = "flank"
+    elif (flank5 == 0) and (flank3 == 0):
+        queries[0]["type"] = "control"
+
+    # For each chromosome, download the encRegTfbsClustered track, extract the features, and fetch the sequences
+    # Loop through chroms
+    for chrom in chroms:
+        chrom_write_count = 0
+        chrom_output_fname = f"{out_dir}/encRegTfbsClustered_{genome}_{chrom}.csv"
+        results_init.to_csv(
+            chrom_output_fname, index=False
+        )
+        logger.info(f"Fetching {chrom}...")
+        # Fetch the data json (has start and end positions in the chrom, but not the sequence)
+        try:
+            data = fetch_tfbs_track(chrom, genome=genome, logger=logger)
+            logger.info(f"  → Fetched {chrom} successfully")
+            features = data.get("encRegTfbsClustered", {})
+            logger.info(f"  → Found {len(features)} features")
+        except Exception as e:
+            logger.info(f"  Failed to fetch {chrom}: {e}")
+            continue
+
+        # Get the sequences of the DNA binding sites
+        for feature_no, feature in enumerate(features):
+            # Initialize new results row
+            new_row = {}
+
+            # Check if tf is valid
+            tf_name = feature.get("name", "UnknownTF")
+            if allowed_tfs and tf_name not in allowed_tfs:
+                logger.warning(f"TF name {tf_name} not in allowed_tfs. Skipping.")
+                continue
+            # Make sure the chromosomes match and we have the right sequence!
+            assert (
+                feature["chrom"] == chrom
+            ), f"Chromosome mismatch: {feature['chrom']} != {chrom}"
+
+            # Add all the cols already in the json, add
+            for c in results_cols:
+                if c in feature:
+                    new_row[c] = feature[c]
+
+            ### Extract sequence
+            start = feature["chromStart"]
+            end = feature["chromEnd"]
+
+            for query in queries:
+                try:
+                    results_dict = get_sequence(
+                        chrom,
+                        start,
+                        end,
+                        flank5=query["flank5"],
+                        flank3=query["flank3"],
+                        genome=genome,
+                        logger = logger
+                    )
+                    # Add the returned info
+                    if query["type"] == "control":
+                        new_row["seq"] = results_dict["seq"]
+                    elif query["type"] == "flank":
+                        new_row["seq_flanked"] = results_dict["seq"]
+                        new_row["chromStart_flanked"] = results_dict["chromStart"]
+                        new_row["chromEnd_flanked"] = results_dict["chromEnd"]
+                        new_row["flank5"] = flank5
+                        new_row["flank3"] = flank3
+                    logger.info(
+                        f" Success on feat. {feature_no} {chrom}:{start}-{end}, type {query['type']}"
+                    )
+                except Exception as e:
+                    logger.info(
+                        f"  Skipped feat. {feature_no} {chrom}:{start}-{end} due to error: {e}"
+                    )
+                    continue
+
+                sleep(0.05)  # Stay within UCSC's 20 req/sec rate limit
+
+            # Fill out any blank columns
+            try:
+                for c in results_cols:
+                    if c not in new_row:
+                        new_row[c] = None
+
+                new_row_df = pd.DataFrame(data=new_row, index=[0])
+                if new_row_df["seq"] is not None:
+                    new_row_df.to_csv(
+                        chrom_output_fname,
+                        mode="a",
+                        index=False,
+                        header=False,
+                    )
+                    logger.info(
+                        f"Wrote new row to {out_dir}/encRegTfbsClustered_{chrom}.csv"
+                    )
+                    chrom_write_count += 1
+                else:
+                    logger.info(f"Did not write new row. {new_row}")
+            except Exception as e:
+                logger.error(F"Failed to write new row to {out_dir}/encRegTfbsClustered_{chrom}.csv: error {e}")
+
+        logger.info(f"Done. Wrote {chrom_write_count} sequences to {out_dir}/encRegTfbsClustered_{chrom}.csv")
+        count += chrom_write_count
+    
+    logger.info(f"Done with all chroms. Wrote {count} sequences to {out_dir}.")
+    
+def setup_chrom_logger(chrom: str, genome: str, out_dir: str) -> logging.Logger:
+    """Set up a dedicated logger for a given chromosome."""
+    logger = logging.getLogger(f"{genome}_{chrom}")
+    logger.setLevel(logging.DEBUG)
+    logger.propagate = False
+
+    # Avoid duplicate handlers if reused
+    if not logger.handlers:
+        os.makedirs(out_dir, exist_ok=True)
+        log_path = os.path.join(out_dir, f"log_{genome}_{chrom}.log")
+        handler = logging.FileHandler(log_path, mode='w', encoding='utf-8')
+        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)
+
+    return logger
+
+# Thread function for one chromosome
+def process_chrom(
+    chrom: str = "chr1",
+    genome: str = "hg38",
+    flank5: int = 500,
+    flank3: int = 500,
+    control_run: bool = True,
+    out_dir: str = "../../data_files/processed/tfclust",
+    allowed_tfs: list = None,
+):
+    """
+    Called within parallel method to strat a thread
+    """
+    chrom_logger = setup_chrom_logger(chrom, genome, f"{out_dir}/logs")
+    chrom_logger.info(f"Starting thread for {chrom}")
+    
+    logging.info(f"Starting thread for {chrom}")
+    try:
+        extract_tfbs_with_context(
+            genome=genome,
+            flank5=flank5,
+            flank3=flank3,
+            control_run=control_run,
+            out_dir=out_dir,
+            allowed_tfs=allowed_tfs,
+            chroms=[chrom],  # important: wrap in list
+            logger=chrom_logger
+        )
+        chrom_logger.info(f"Finished {chrom}")
+    except Exception as e:
+        chrom_logger.error(f"Error processing {chrom}: {e}")
+
+
+import multiprocessing
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+def parallel_extract_tfbs_for_genome(
+    genome: str,
+    flank5: int,
+    flank3: int,
+    control_run: bool,
+    out_dir: str,
+    allowed_tfs: list,
+    chroms: list,
+    max_workers: int,
+):
+    logger = logging.getLogger(f"{genome}")
+    logger.info(f"Using {max_workers} threads for {genome}...")
+
+    if chroms is None:
+        chroms = get_all_chroms(genome=genome)
+
+    futures = {}
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        for chrom in chroms:
+            future = executor.submit(
+                process_chrom,
+                chrom=chrom,
+                genome=genome,
+                flank5=flank5,
+                flank3=flank3,
+                control_run=control_run,
+                out_dir=f"{out_dir}/{genome}",
+                allowed_tfs=allowed_tfs,
+            )
+            futures[future] = f"{genome}:{chrom}"
+
+        for future in as_completed(futures):
+            label = futures[future]
+            try:
+                future.result()
+            except Exception as e:
+                logger.error(f"{label} raised an exception: {e}")
+
+
+def parallel_extract_tfbs_with_context(
+    genomes=["hg38", "hg19"],
+    flank5=500,
+    flank3=500,
+    control_run=True,
+    out_dir="../../data_files/processed/tfclust",
+    allowed_tfs=None,
+    chroms=None,
+):
+    total_cpus = multiprocessing.cpu_count()
+    cpu_per_genome = total_cpus // len(genomes)
+
+    logging.info(f"Total CPUs: {total_cpus}")
+    logging.info(f"Launching {len(genomes)} genome pipelines with {cpu_per_genome} threads each")
+
+    processes = []
+    for genome in genomes:
+        p = multiprocessing.Process(
+            target=parallel_extract_tfbs_for_genome,
+            args=(
+                genome,
+                flank5,
+                flank3,
+                control_run,
+                out_dir,
+                allowed_tfs,
+                chroms,
+                cpu_per_genome
+            )
+        )
+        p.start()
+        processes.append(p)
+
+    for p in processes:
+        p.join()
+
+def main():
+    genomes = ["hg38", "hg19"]
+
+    parallel_extract_tfbs_with_context(
+        genomes=genomes,
+        flank5=500,
+        flank3=500,
+        control_run=True,  # if there's a flank, whether to also run without flank
+        out_dir=f"../../data_files/processed/tfclust",
+        allowed_tfs=None,  # e.g., ['CTCF', 'MAX']
+        chroms=None,
+    )
+
+if __name__ == "__main__":
+    logger = logging.getLogger(__name__)
+    logging.basicConfig(
+        filename="download.log",
+        encoding="utf-8",
+        level=logging.DEBUG,
+        filemode="w",
+    )
+    main()