diff --git a/reditools/compiled_position.py b/reditools/compiled_position.py index 1b476d9..d4d0450 100644 --- a/reditools/compiled_position.py +++ b/reditools/compiled_position.py @@ -5,7 +5,12 @@ from typing import Iterator from reditools.constants import bases as base_order -from reditools.constants import comp_map +from reditools.constants import ( + comp_map, + forward_strand_symbol, + reverse_strand_symbol, + undetermined_strand_symbol, +) @dataclass @@ -78,17 +83,17 @@ def calculate_strand(self, threshold: float = 0) -> str: pos_count = 0 neg_count = 0 for strand in self.strands: - if strand == "+": + if strand == forward_strand_symbol: pos_count += 1 - elif strand == "-": + elif strand == reverse_strand_symbol: neg_count += 1 if pos_count == neg_count: - return "*" + return undetermined_strand_symbol if pos_count / (pos_count + neg_count) >= threshold: - return "+" + return forward_strand_symbol if neg_count / (pos_count + neg_count) >= threshold: - return "-" - return "*" + return reverse_strand_symbol + return undetermined_strand_symbol def filter_by_strand(self, strand: str) -> None: """Filter observations to keep only those from a specific strand. @@ -98,7 +103,7 @@ def filter_by_strand(self, strand: str) -> None: strand : str The strand to keep ('+', '-', or '*'). If '*', no filtering is done. """ - if strand == "*": + if strand == undetermined_strand_symbol: return keep = [ idx for idx in range(len(self.bases)) diff --git a/reditools/compiled_reads.py b/reditools/compiled_reads.py index b0f50d6..30f72a5 100644 --- a/reditools/compiled_reads.py +++ b/reditools/compiled_reads.py @@ -4,6 +4,11 @@ from typing import TYPE_CHECKING from reditools.compiled_position import CompiledPosition +from reditools.constants import ( + forward_strand_symbol, + reverse_strand_symbol, + undetermined_strand_symbol, +) from reditools.fasta_file import RTFastaFile if TYPE_CHECKING: @@ -99,7 +104,11 @@ class CompiledReads: and strand. """ - _strands = ("-", "+", "*") + _strands = ( + reverse_strand_symbol, + forward_strand_symbol, + undetermined_strand_symbol, + ) def __init__( self, diff --git a/reditools/constants.py b/reditools/constants.py index 97df71d..138030b 100644 --- a/reditools/constants.py +++ b/reditools/constants.py @@ -1,4 +1,6 @@ -"""Dictionary for base complements.""" +"""Constants.""" + +"""Complementary base dictionary.""" comp_map = { "A": "T", "T": "A", @@ -7,4 +9,16 @@ "N": "N", "-": "-", } + +"""Nucleotides in output order.""" bases = ("A", "C", "G", "T") + +"""Output strand symbols.""" +forward_strand_symbol = "+" +reverse_strand_symbol = "-" +undetermined_strand_symbol = "*" + +"""Output strand numbers.""" +forward_strand_number = "1" +reverse_strand_number = "0" +undetermined_strand_number = "2" diff --git a/reditools/tools/analyze/parse_args/parse_args.py b/reditools/tools/analyze/parse_args/parse_args.py index dd56038..2b7ba08 100644 --- a/reditools/tools/analyze/parse_args/parse_args.py +++ b/reditools/tools/analyze/parse_args/parse_args.py @@ -89,6 +89,14 @@ def build_argument_parser() -> argparse.ArgumentParser: # noqa: WPS213, WPS210 action="store_true", help="Appends results to file (and creates if not existing).", ) + output_group.add_argument( + "--number-strand-output", + action="store_true", + help=( + "Replaces the '-', '+', and '*' values with 0, 1, and 2 in the " + "Strand column of the output." + ), + ) bqf_group = parser.add_argument_group( title="Base/Read Quality Controls", ) diff --git a/reditools/tools/analyze/redi_thread.py b/reditools/tools/analyze/redi_thread.py index 913b3ab..621d69a 100644 --- a/reditools/tools/analyze/redi_thread.py +++ b/reditools/tools/analyze/redi_thread.py @@ -44,6 +44,7 @@ def __init__(self, options: argparse.Namespace) -> None: ) self.rtqc = RTChecks(options) self.temp_dir = options.temp_dir + self.number_strand_output = options.number_strand_output def analyze( self, @@ -65,6 +66,7 @@ def analyze( filename, self.rtqc, self.rtools.log, + self.number_strand_output, ) class REDIThreadManager: diff --git a/reditools/tools/analyze/write_results.py b/reditools/tools/analyze/write_results.py index 77772d8..17ac6a7 100644 --- a/reditools/tools/analyze/write_results.py +++ b/reditools/tools/analyze/write_results.py @@ -2,25 +2,34 @@ import csv from pathlib import Path -from typing import Callable, Iterator +from typing import Callable, Iterable from reditools.compiled_position import RTResult +from reditools.constants import ( + forward_strand_number, + forward_strand_symbol, + reverse_strand_number, + reverse_strand_symbol, + undetermined_strand_number, + undetermined_strand_symbol, +) from reditools.logger import Logger from reditools.tools.analyze.rtchecks import RTChecks _empty = "-" def write_results( - rtresults: Iterator[RTResult], + rtresults: Iterable[RTResult], filename: str, filters: RTChecks, logger: Callable, + strand_numbers: bool = False, ) -> None: """Write analysis results to a file. Parameters ---------- - rtresults : Iterator[RTResult] + rtresults : Iterable[RTResult] The analysis results for each position. filename : str Where to save the results. @@ -28,7 +37,22 @@ def write_results( The quality control checks to apply. logger : Callable The logger function for debug messages. + strand_numbers : bool + If False, uses -, +, and * for the Strand column. + If True, uses 0, 1, and 2 for the Strand column. """ + if strand_numbers: + strand_lookup = { + forward_strand_symbol: forward_strand_number, + reverse_strand_symbol: reverse_strand_number, + undetermined_strand_symbol: undetermined_strand_number, + } + else: + strand_lookup = { + forward_strand_symbol: forward_strand_symbol, + reverse_strand_symbol: reverse_strand_symbol, + undetermined_strand_symbol: undetermined_strand_symbol, + } with Path(filename).open("w") as stream: writer = csv.writer(stream, delimiter="\t", lineterminator="\n") for rt_result in rtresults: @@ -41,7 +65,7 @@ def write_results( rt_result.contig, rt_result.position + 1, rt_result.reference, - rt_result.strand, + strand_lookup[rt_result.strand], len(rt_result), f"{rt_result.mean_quality:.2f}", list(rt_result), diff --git a/test/__main__.py b/test/__main__.py index cffc4d7..fd553c9 100644 --- a/test/__main__.py +++ b/test/__main__.py @@ -8,6 +8,7 @@ from test.analyze.rtchecks import TestRTChecks from test.analyze.setup_alignment_manager import TestSetupAlignmentManager from test.analyze.setup_rtools import TestSetupRTools +from test.analyze.write_results import TestWriteResults from test.compiled_position import TestCompiledPosition from test.compiled_reads import TestCompiledReads from test.fasta_file import TestRTFastaFile diff --git a/test/analyze/write_results.py b/test/analyze/write_results.py new file mode 100644 index 0000000..45bafc4 --- /dev/null +++ b/test/analyze/write_results.py @@ -0,0 +1,124 @@ +"""Test methods for analyze tool's write_results().""" +from __future__ import annotations + +import unittest +from pathlib import Path +from tempfile import NamedTemporaryFile + +from reditools.compiled_position import CompiledPosition, RTResult +from reditools.logger import Logger +from reditools.tools.analyze.parse_args.parse_args import parse_args +from reditools.tools.analyze.rtchecks import RTChecks +from reditools.tools.analyze.write_results import write_results + + +class TestWriteResults(unittest.TestCase): + """Test methods for analyze tool's write_results().""" + + def setUp(self) -> None: + """Pre-flight setup.""" + self.rtresults = [] + + cp = CompiledPosition(ref="A", contig="chr1", position=100) + cp.add_base(40, "+", "A") + cp.add_base(30, "+", "G") + cp.add_base(30, "+", "T") + self.rtresults.append(RTResult(cp, "+")) + self.rtresults.append(RTResult(cp, "-")) + self.rtresults.append(RTResult(cp, "*")) + + with NamedTemporaryFile(mode="w", dir=".", delete=False) as temp_file: + self.fname = temp_file.name + self.log = Logger(Logger.silent_level).log + + def tearDown(self) -> None: + """Post-checks cleanup.""" + Path(self.fname).unlink() + + def rtchecks(self, cli_args: list[str]) -> RTChecks: + """Create RTChecks object from CLI arguments. + + Parameters + ---------- + cli_args : list[str] + Command line arguments. + + Returns + ------- + RTChecks + Checks for output filtering. + """ + return RTChecks(parse_args(cli_args)) + + def test_write_results(self) -> None: + """Check basic functionality.""" + write_results( + self.rtresults, + self.fname, + self.rtchecks(["x.bam"]), + self.log, + ) + with Path(self.fname).open("r") as stream: + self.assertEqual( + next(stream).strip().split("\t") , + [ + "chr1", "101", "A", "+", "3", "33.33", "[1, 0, 1, 1]", + "AG AT", "0.50", "-", "-", "-", "-", "-", + ], + ) + self.assertEqual( + next(stream).strip().split("\t") , + [ + "chr1", "101", "A", "-", "3", "33.33", "[1, 0, 1, 1]", + "AG AT", "0.50", "-", "-", "-", "-", "-", + ], + ) + self.assertEqual( + next(stream).strip().split("\t") , + [ + "chr1", "101", "A", "*", "3", "33.33", "[1, 0, 1, 1]", + "AG AT", "0.50", "-", "-", "-", "-", "-", + ], + ) + + def test_write_results_numbers(self) -> None: + """Check strand-number option.""" + write_results( + self.rtresults, + self.fname, + self.rtchecks(["x.bam"]), + self.log, + True, # noqa: FBT003 + ) + with Path(self.fname).open("r") as stream: + self.assertEqual( + next(stream).strip().split("\t") , + [ + "chr1", "101", "A", "1", "3", "33.33", "[1, 0, 1, 1]", + "AG AT", "0.50", "-", "-", "-", "-", "-", + ], + ) + self.assertEqual( + next(stream).strip().split("\t") , + [ + "chr1", "101", "A", "0", "3", "33.33", "[1, 0, 1, 1]", + "AG AT", "0.50", "-", "-", "-", "-", "-", + ], + ) + self.assertEqual( + next(stream).strip().split("\t") , + [ + "chr1", "101", "A", "2", "3", "33.33", "[1, 0, 1, 1]", + "AG AT", "0.50", "-", "-", "-", "-", "-", + ], + ) + + def test_filter(self) -> None: + """Check that RTChecks filters output.""" + write_results( + self.rtresults, + self.fname, + self.rtchecks(["x.bam", "--min-read-depth", "5"]), + self.log, + ) + self.assertEqual(Path(self.fname).stat().st_size, 0)