dataprep.py

#!/usr/bin/env python3.9

"""
dataprep.py - version 1.0a1
Author: Gil Oliveira <gpo@ciencias.ulisboa.pt>
PI: Margarida Gama-Carvalho <mhcarvalho@ciencias.ulisboa.pt>

RNA Systems Biology Lab
BioISI - Biosystems and Integrative Sciences Institute
Department of Chemistry and Biochemistry
Faculty of Sciences, University of Lisbon

(C) 2022-2023
"""

import argparse
import gzip
import shutil
import subprocess
import sys
from multiprocessing import Pool
from pathlib import Path

import inquirer

ORIGINAL_FILES_DIR = Path("/data/original_files/")
WORKING_DIR = Path("/data/working_directory/projects/")
CUTADAPT_BIN = Path("/opt/anaconda2/bin/cutadapt")
FASTQC_BIN = Path("/opt/FastQC/fastqc")

ADAPTER_DEFAULT = "AGATCGGAAGAGCGTCGTGTA"
ADAPTER_R1_DEFAULT = "GATCGGAAGAGCACACGTCTGA"
ADAPTER_R2_DEFAULT = "AGATCGGAAGAGCGTCGTGTA"
TRIM_ENDS_DEFAULT = 10


def cutadapt_paired_end(r1_path, r2_path):
    global trim_ends
    r1 = r1_path.name
    r2 = r2_path.name
    log_path = (r1_path.parents[1] / "cutadapt") / (
        r1_path.stem + "_cutadapt_stdout.txt"
    )
    with open(log_path, "w") as f:
        instance = subprocess.Popen(
            [
                CUTADAPT_BIN,
                "-a",
                adapter_r1,
                "-A",
                adapter_r2,
                "-u",
                str(trim_ends),
                "-U",
                str(trim_ends),
                "-o",
                r1_path.stem + "_cutadapt.fastq",
                "-p",
                r2_path.stem + "_cutadapt.fastq",
                r1,
                r2,
            ],
            cwd=raw_fastq_dir,
            stdout=f,
        )
    instance.wait()


def cutadapt_single_end(fastq_path):
    global trim_ends
    fastq_name = fastq_path.name
    log_path = (fastq_path.parents[1] / "cutadapt") / (
        fastq_path.stem + "_cutadapt_stdout.txt"
    )
    with open(log_path, "w") as f:
        instance = subprocess.Popen(
            [
                CUTADAPT_BIN,
                "-a",
                adapter,
                "-u",
                str(trim_ends),
                "-o",
                fastq_path.stem + "_cutadapt.fastq",
                fastq_name,
            ],
            cwd=raw_fastq_dir,
            stdout=f,
        )
    instance.wait()


def cutadapt_all_files():
    trim_ends_response = input(
        "How many nucleotides should be cut off each end? (default: 0): "
    )
    adapter_response = input(
        f"Insert adapter sequence (default {ADAPTER_DEFAULT}): "
    )
    global trim_ends
    global adapter

    if trim_ends_response == "":
        trim_ends = 0
    else:
        trim_ends = int(trim_ends_response)

    if adapter_response == "":
        adapter = ADAPTER_DEFAULT
    else:
        adapter = adapter_response

    print("Trimming adapters with cutadapt. This may take a while...")

    global cutadapt_dir
    global raw_fastq_dir

    cutadapt_dir.mkdir(parents=True, exist_ok=True)
    r1_files = raw_fastq_dir.glob("*R1*.fastq")
    r2_files = raw_fastq_dir.glob("*R2*.fastq")
    pool = Pool()
    if len(list(r1_files)) + len(list(r2_files)) > 2:
        # Paired-end files
        paired_ends = zip(r1_files, r2_files)
        pool.starmap(cutadapt_paired_end, paired_ends)
    else:
        fastq_files = raw_fastq_dir.glob("*.fastq")
        pool.map(cutadapt_single_end, fastq_files)
    cutadapt_files = raw_fastq_dir.glob("*_cutadapt.fastq")
    for file in cutadapt_files:
        shutil.move(file, cutadapt_dir)


def fastqc_file(file):
    parent_folder = file.parent
    instance = subprocess.Popen(
        [
            FASTQC_BIN,
            file,
            "--outdir=" + str(parent_folder / "fastqc_reports"),
        ],
        cwd=file.parents[0],
    )
    instance.wait()


def fastqc_all_files(fastq_dir):
    fastqc_dir = Path(fastq_dir / "fastqc_reports")
    Path(fastqc_dir / "html").mkdir(parents=True, exist_ok=True)
    fastq_files = fastq_dir.glob("*.fastq")
    pool = Pool()
    pool.map(fastqc_file, fastq_files)
    for file in fastqc_dir.glob("*.html"):
        shutil.move(file, fastqc_dir / "html")


def parse_arguments():
    """Parse arguments supplied on runtime and create global variables to
    store those values. Any unsupplied parameters will be set to the default
    hardcoded values (variables which end in "DEFAULT")
    """
    parser = argparse.ArgumentParser(
        description="Prep the data for RNA-Seq analysis."
    )
    parser.add_argument(
        "-i",
        "--input",
        type=str,
        metavar="/data/originalFiles/PROJECT/",
        help="Directory where the original files are.",
    )
    args = parser.parse_args()

    global project_original_files
    project_original_files = args.input


def decompress_file(path):
    stem = Path(path).stem
    output = project_working_dir / "raw_data_fastq" / stem
    with gzip.open(path, "rb") as f_gz, open(output, "wb") as f_ungz:
        print("Unzipping", f_gz, "to", f_ungz)
        content = f_gz.read()
        f_ungz.write(content)


def decompress_data(path):
    global project_working_dir

    Path(project_working_dir / "raw_data_fastq").mkdir(
        parents=True, exist_ok=True
    )
    files = path.glob("*.gz")
    print("Decompressing files...")
    pool = Pool()
    pool.map(decompress_file, files)
    print(
        "Done! Decompressed FASTQ files saved in "
        f"{project_working_dir}/raw_data_fastq/"
    )


def select_project():
    projects = [x for x in original_files.iterdir() if x.is_dir()]
    print(
        "Use your up/down arrows to find the folder you want, and select "
        "it with ENTER."
    )
    questions = [
        inquirer.List("project_path", message="Folder:", choices=projects)
    ]
    project_path = Path(inquirer.prompt(questions)["project_path"])
    return project_path


def main() -> None:
    global original_files
    global working_dir
    global adapter
    global adapter_r1
    global adapter_r2
    global trim_ends
    global project_original_files

    original_files = ORIGINAL_FILES_DIR
    working_dir = WORKING_DIR
    adapter = ADAPTER_DEFAULT
    adapter_r1 = ADAPTER_R1_DEFAULT
    adapter_r2 = ADAPTER_R2_DEFAULT
    trim_ends = TRIM_ENDS_DEFAULT

    parse_arguments()
    if project_original_files is None:
        project_original_files = select_project()
        slug = project_original_files.name
        project_working_dir = working_dir / slug
        yn = input(f"You've selected project {slug}. Is that correct? (Y/n): ")
        if yn.lower() not in ("y", "yes", "s", "sim", ""):
            sys.exit()
    else:
        project_original_files = Path(project_original_files)
        project_working_dir = working_dir / project_original_files.name
        print(f"Using original files in {project_original_files}.")

    raw_fastq_dir = project_working_dir / "raw_data_fastq"
    cutadapt_dir = project_working_dir / "cutadapt"

    yn = input(
        "Do you want to uncompress the files into "
        f"{project_working_dir}? (Y/n): "
    )
    if yn.lower() in ("y", "yes", "s", "sim", ""):
        decompress_data(project_original_files)
    yn = input("Run FastQC on raw data? (Y/n): ")
    if yn.lower() in ("y", "yes", "s", "sim", ""):
        fastqc_all_files(raw_fastq_dir)
    yn = input("Run cutadapt? (Y/n): ")
    if yn.lower() in ("y", "yes", "s", "sim", ""):
        cutadapt_all_files()
    yn = input("Run FastQC on trimmed files? (Y/n): ")
    if yn.lower() in ("y", "yes", "s", "sim", ""):
        fastqc_all_files(cutadapt_dir)


if __name__ == "__main__":
    main()