diff --git a/neat/bacterial_wrapper/__init__.py b/neat/bacterial_wrapper/__init__.py new file mode 100644 index 00000000..936ca573 --- /dev/null +++ b/neat/bacterial_wrapper/__init__.py @@ -0,0 +1,4 @@ +""" +Load modules needed for other parts of the program +""" +from .runner import * \ No newline at end of file diff --git a/neat/bacterial_wrapper/runner.py b/neat/bacterial_wrapper/runner.py new file mode 100644 index 00000000..333803e5 --- /dev/null +++ b/neat/bacterial_wrapper/runner.py @@ -0,0 +1,192 @@ +import subprocess +import gzip +import shutil +import yaml +import pysam +import unittest +import os + +from pathlib import Path +from typing import List +from Bio import bgzf +from Bio.bgzf import BgzfWriter, BgzfReader + + +# Rearranges the bacterial chromosome by wrapping it around + +def wrapper(seq): + length = len(seq) + + if (length % 2 == 0): + half_index = length // 2 + else: + half_index = (length // 2) + 1 + + first_half = seq[:half_index] + second_half = seq[half_index:] + + new_seq = second_half + first_half + + return new_seq + + +# Writes the newly rearranged chromosome's sequence to a new fasta file + +def write_fasta_file(new_seq, bacteria_name, fasta_header, output_dir_path): + fasta_file_name = f"wrapped_{bacteria_name}.fna" + fasta_file_path = output_dir_path / fasta_file_name + fasta_file = open(fasta_file_path, "w") + + fasta_file.write(fasta_header + "\n" + new_seq) + + fasta_file.close() + + return fasta_file_path + + +# Writes a yml configuration file for the newly rearranged chromosome's fasta sequence +# Splits the coverage in half for the reference and new config files +# These use default values for all other parameters for NEAT + +def write_config_file(ref_config_file, rearranged_seq_file, bacteria_name, output_dir_path): + new_config_file_name = f"new_{bacteria_name}_config_test.yml" + old_config_file_name = f"{bacteria_name}_config_test.yml" + + new_config_file_path = output_dir_path / new_config_file_name + old_config_file_path = output_dir_path / old_config_file_name + + with open(ref_config_file, 'r') as ref_file, open(new_config_file_path, 'w') as new_file, open(old_config_file_path, 'w') as old_file: + + for line in ref_file: + if line.find("reference:") != -1: + new_file.write(f"reference: {rearranged_seq_file}") + old_file.write(line) + elif line.find("coverage:") != -1: + if line.strip() == "coverage: .": + new_coverage = 5.0 + else: + new_coverage = float((line.split(" "))[1].strip()) // 2 + + new_file.write(f"coverage: {new_coverage}") + old_file.write(f"coverage: {new_coverage}") + else: + new_file.write(line) + old_file.write(line) + + + ref_file.close() + new_file.close() + old_file.close() + + return old_config_file_path, new_config_file_path + + +# Runs the NEAT read simulator using the given config file + +def run_neat(config_file, output_dir, prefix): + subprocess.run(["neat", "read-simulator", "-c", config_file, "-o", output_dir + "/" + prefix]) + + +# General function for bacterial wrapper that calls all of the functions defined above + +def bacterial_wrapper(reference_file, bacteria_name, ref_config_file, output_dir): + + orig_seq = "" + + f = open(reference_file) + fasta_header = f.readline().strip() + + for line in f: + if line[0] != ">": + orig_seq += line.strip() + elif line.find("plasmid") != -1: # exclude plasmids from the sequence to be rearranged + break + + f.close() + + output_dir_path = Path(output_dir) + + rearranged_seq = wrapper(orig_seq) + rearranged_seq_file = write_fasta_file(rearranged_seq, bacteria_name, fasta_header, output_dir_path) + + config_files = write_config_file(ref_config_file, rearranged_seq_file, bacteria_name, output_dir_path) + old_config_file = config_files[0] + new_config_file = config_files[1] + + run_neat(old_config_file, output_dir, "Regular") + run_neat(new_config_file, output_dir, "Wrapped") + + +# Stitching all outputs together - Keshav's script + +def concat_fq(input_files: List[Path], dest: BgzfWriter) -> None: + + if not input_files: + return + + for input_file in input_files: + with bgzf.BgzfReader(input_file) as in_f: + shutil.copyfileobj(in_f, dest) + +def merge_bam(bams: List[Path], dest: Path, threads: int) -> None: + + if not bams: + return + + unsorted = dest.with_suffix(".unsorted.bam") + pysam.merge("--no-PG", "-@", str(threads), "-f", str(unsorted), *map(str, bams)) + pysam.sort("-@", str(threads), "-o", str(dest), str(unsorted)) + unsorted.unlink(missing_ok=True) + +def merge_vcf(vcfs: List[Path], dest: Path) -> None: + if not vcfs: + return + + first, *rest = vcfs + shutil.copy(first, dest) + + with dest.open("ab") as out_f: + for vcf in rest: + with vcf.open("rb") as fh: + for line in fh: + if not line.startswith(b"#"): + out_f.write(line) + +def stitch_all_outputs(files: List[Path], output_dir) -> None: + fq1_list = [] + fq2_list = [] + vcf_list = [] + bam_list = [] + + for file in files: + file_name = file.stem # use stem to differentiate fq1 and fq2 + suffixes = file.suffixes # use suffixes to catch vcf and bam files + + if "r2.fastq" in file_name: + fq2_list.append(file) + elif "r1.fastq" in file_name or ".fastq" in suffixes: + fq1_list.append(file) + elif ".vcf" in suffixes: + vcf_list.append(file) + elif ".bam" in suffixes: + bam_list.append(file) + + dest_fq1 = bgzf.BgzfWriter(f"{output_dir}/stitched_fq1.bgzf") + dest_fq2 = bgzf.BgzfWriter(f"{output_dir}/stitched_fq2.bgzf") + dest_bam = Path(f"{output_dir}/stitched.bam") + dest_vcf = Path(f"{output_dir}/stitched.vcf") + + concat_fq(fq1_list, dest_fq1) + concat_fq(fq2_list, dest_fq2) + merge_bam(bam_list, dest_bam, 2) + merge_vcf(vcf_list, dest_vcf) + + +# Testing functions + +class TestWrapper(unittest.TestCase): + def test_even(self): + self.assertEqual(wrapper("ABBCBB"), "CBBABB") + + def test_odd(self): + self.assertEqual(wrapper("ABBCBBC"), "BBCABBC") \ No newline at end of file diff --git a/neat/cli/commands/bacterial_wrapper.py b/neat/cli/commands/bacterial_wrapper.py new file mode 100644 index 00000000..b8dfc818 --- /dev/null +++ b/neat/cli/commands/bacterial_wrapper.py @@ -0,0 +1,78 @@ +""" +Command line interface for NEAT's bacterial wrapper function +""" + +import argparse +import subprocess +import os +from pathlib import Path + +from ...bacterial_wrapper import bacterial_wrapper +from ...bacterial_wrapper import stitch_all_outputs +from .base import BaseCommand +from .options import output_group + + +class Command(BaseCommand): + """ + Class that generates wrapped bacterial models + """ + name = "bacterial-wrapper" + description = "Generate wrapped bacterial model reads" + + def add_arguments(self, parser: argparse.ArgumentParser): + """ + Add the command's arguments to its parser + + :param parser: The parser to add arguments to + """ + + parser.add_argument('reference', + type=str, + metavar='reference.fa', + help="Reference file for organism in fasta format.") + + parser.add_argument('bacteria_name', + type=str, + metavar='bacteria_name', + help="Name of the bacteria.") + + parser.add_argument( + "-c", "--config", + metavar="config", + type=str, + required=True, + help="Path (including filename) to the configuration file for the reference run." + ) + + output_group.add_to_parser(parser) + + def execute(self, arguments: argparse.Namespace): + """ + Execute the command + + :param arguments: The namespace with arguments and their values. + """ + bacterial_wrapper(arguments.reference, arguments.bacteria_name, arguments.config, arguments.output_dir) + + output_path = Path(arguments.output_dir) + file_list = os.listdir(output_path / "Regular") # same file names for both Regular and Wrapped folders + + output_files = [] + + for file in file_list: + reg_file_path = output_path / "Regular" / file + wrap_file_path = output_path / "Wrapped" / file + + if ("vcf" in file): + subprocess.run(["gzip", "-d", reg_file_path]) + subprocess.run(["gzip", "-d", wrap_file_path]) + + file = file[:-3] + reg_file_path = output_path / "Regular" / file + wrap_file_path = output_path / "Wrapped" / file + + output_files.append(reg_file_path) + output_files.append(wrap_file_path) + + stitch_all_outputs(output_files, arguments.output_dir) \ No newline at end of file diff --git a/neat/read_simulator/utils/options.py b/neat/read_simulator/utils/options.py index d2cdf391..f0153170 100644 --- a/neat/read_simulator/utils/options.py +++ b/neat/read_simulator/utils/options.py @@ -498,4 +498,4 @@ def log_configuration(self): _LOG.info(f'Custom average mutation rate for the run: {self.mutation_rate}') if self.mutation_bed: _LOG.info(f'BED of mutation rates of different regions: {self.mutation_bed}') - _LOG.info(f'RNG seed value for run: {self.rng_seed}') + _LOG.info(f'RNG seed value for run: {self.rng_seed}') \ No newline at end of file