seqalign.py

#!/usr/bin/env python3.9

"""
seqalign - 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 multiprocessing
import os
import subprocess
import sys
from pathlib import Path, PurePath

import inquirer

STAR_BIN_LOCATION_DEFAULT = Path("/opt/STAR-2.5.0c/bin/Linux_x86_64/STAR")
WORKING_DIR_DEFAULT = Path("/data/working_directory/projects/")
ORIGINAL_FILES_DIR_DEFAULT = Path("/data/originalFiles/")
GENOMES_LOCATION_DEFAULT = Path("/data/referenceFiles/genome/")

genomes_location = GENOMES_LOCATION_DEFAULT


def parse_args() -> None:
    """Parse arguments provided by the user on runtime."""
    parser = argparse.ArgumentParser(description="Get genomes from Ensembl.")
    parser.add_argument(
        "-g",
        "--genomeDir",
        type=str,
        metavar="/data/referenceFiles/genome/" "boops_boops/star/",
        help="directory with STAR genome index",
    )
    parser.add_argument(
        "-o",
        "--outdir",
        type=str,
        metavar="/data/working_directory/PROJECT/",
        help="directory with the outputs of your project",
    )
    args = parser.parse_args()

    if args.genomeDir is not None:
        global genomeDir
        genomeDir = Path(args.genomeDir)


def check_genome_index(genome_location: Path) -> None:
    """Check if there's a STAR index inside a directory."""
    index_exists = os.path.isdir(genomes_location + "star")
    if not index_exists:
        print("No STAR genome index could be found.")
        print("Refer to SOP DAM002 for index creation instructions.")
        sys.exit()


def align_single(r1) -> None:
    global star_bin_location
    global genomeDir
    global outDir
    outDir_single = f"{outDir}/{r1.stem}/"
    Path(outDir_single).mkdir(parents=True, exist_ok=True)
    p = subprocess.Popen(
        [
            str(PurePath(star_bin_location)),
            "--genomeDir",
            str(PurePath(genomeDir)),
            "--readFilesIn",
            r1.name,
            "--outFilterType",
            "BySJout",
            "--outFilterMultimapNmax",
            "1",
            "--alignSJoverhangMin",
            "8",
            "--alignSJDBoverhangMin",
            "5",
            "--outFilterMismatchNmax",
            "1",
            "--alignIntronMax",
            "100000",
            "--outSAMmultNmax",
            "1",
            "--outSAMattributes",
            "All",
            "--outSAMstrandField",
            "intronMotif",
            "--outSAMtype",
            "BAM",
            "Unsorted",
            "SortedByCoordinate",
            "--quantMode",
            "TranscriptomeSAM",
            "GeneCounts",
            "--outSAMunmapped",
            "Within",
            "--outReadsUnmapped",
            "Fastx",
            "--twopassMode",
            "Basic",
            "--outSJfilterCountTotalMin",
            "8",
            "5",
            "5",
            "5",
            "--outSJfilterCountUniqueMin",
            "-1",
            "-1",
            "-1",
            "-1",
            "--outWigType",
            "wiggle",
            "read1_5p",
            "--limitBAMsortRAM",
            "10000000000",
            "--outFileNamePrefix",
            outDir_single,
        ],
        cwd=sequences_dir,
    )
    p.communicate()


def align_pair(r1, r2):
    global star_bin_location
    global genomeDir
    global outDir
    r1 = PurePath(r1[0]).name
    r2 = PurePath(r2[0]).name
    outDir_pair = f"{outDir}/{r1.stem}/"
    Path(outDir_pair).mkdir(parents=True, exist_ok=True)
    p = subprocess.Popen(
        [
            str(PurePath(star_bin_location)),
            "--genomeDir",
            str(PurePath(genomeDir)),
            "--readFilesIn",
            r1,
            r2,
            "--outFilterType",
            "BySJout",
            "--outFilterMultimapNmax",
            "1",
            "--alignSJoverhangMin",
            "8",
            "--alignSJDBoverhangMin",
            "5",
            "--outFilterMismatchNmax",
            "1",
            "--alignIntronMax",
            "100000",
            "--outSAMmultNmax",
            "1",
            "--outSAMattributes",
            "All",
            "--outSAMstrandField",
            "intronMotif",
            "--outSAMtype",
            "BAM",
            "Unsorted",
            "SortedByCoordinate",
            "--quantMode",
            "TranscriptomeSAM",
            "GeneCounts",
            "--outSAMunmapped",
            "Within",
            "--outReadsUnmapped",
            "Fastx",
            "--twopassMode",
            "Basic",
            "--outSJfilterCountTotalMin",
            "8",
            "5",
            "5",
            "5",
            "--outSJfilterCountUniqueMin",
            "-1",
            "-1",
            "-1",
            "-1",
            "--outWigType",
            "wiggle",
            "read1_5p",
            "--limitBAMsortRAM",
            "10000000000",
            "--outFileNamePrefix",
            outDir_pair,
        ],
        cwd=sequences_dir,
    )
    p.communicate()


def find_genomes() -> dict:
    """Find indexed genomes in the system."""
    available_genomes = genomes_location.iterdir()
    organisms = {}
    for species_path in available_genomes:
        slug = species_path.name
        organisms[slug] = species_path
    return organisms


def query_genomeDir() -> Path:
    """Query indexed genomeDir form filesystem."""
    available_genomes = find_genomes()
    print(
        "Use your up/down arrows to find the species you want, and select "
        "it with ENTER."
    )
    questions = [
        inquirer.List("species", message="Species", choices=available_genomes.keys())
    ]
    answer = inquirer.prompt(questions)["species"]
    chosen_dir = genomes_location / answer / "star"
    return chosen_dir


def query_sequences_dir() -> str:
    """Query the user for the directory which contains the filtered
    sequences"""
    project_paths = working_dir.glob("*/fastfilter")
    print(
        "Use your up/down arrows to find the folder you want, and select "
        "it with ENTER."
    )
    questions = [
        inquirer.List(
            "project_path", message="Sequences folder", choices=list(project_paths)
        )
    ]
    answer = inquirer.prompt(questions)["project_path"]
    return Path(answer)


def query_outDir() -> str:
    """Query the user interactively for the folder to save the reports to."""
    project_paths = working_dir.iterdir()
    print(
        "Use your up/down arrows to find the folder you want, and select "
        "it with ENTER."
    )
    questions = [
        inquirer.List(
            "project_path", message="Output folder", choices=list(project_paths)
        )
    ]
    answer = inquirer.prompt(questions)["project_path"]
    return answer


def get_fasta_pairs() -> list:
    global sequences_dir
    all_r1 = list(sequences_dir.glob("*_R1_*_FILTERED.fastq"))
    all_r2 = list(sequences_dir.glob("*_R2_*_FILTERED.fastq"))
    if len(list(all_r1)) != len(list(all_r2)):
        raise Exception("Unmatching paired ends found! Process aborted.")
    paired_ends = list(zip(all_r1, all_r2))
    return paired_ends


def main() -> None:
    """Main function"""
    global working_dir
    global star_bin_location
    global sequences_dir
    global genomeDir
    global outDir
    working_dir = WORKING_DIR_DEFAULT
    star_bin_location = STAR_BIN_LOCATION_DEFAULT
    parse_args()
    if "sequences_dir" not in globals():
        sequences_dir = Path(query_sequences_dir())
    if "genomeDir" not in globals():
        genomeDir = Path(query_genomeDir())
    if "outDir" not in globals():
        outDir = Path(query_outDir())
    outDir = outDir / "STAR_results"
    paired_ends = get_fasta_pairs()
    pool = multiprocessing.Pool()
    if len(paired_ends) == 2:
        print("Detected single-end reading.")
        align_pair(paired_ends[0], paired_ends[1])
    elif len(paired_ends) == 0:
        print("Detected single-end readings.")
        fastq_files = list(sequences_dir.glob("*_FILTERED.fastq"))
        pool.map(align_single, fastq_files)
    else:
        print("Detected paired-end readings.")
        pool.starmap(align_pair, paired_ends)


if __name__ == "__main__":
    main()