Merge pull request #169 from openvax/codex/fix-allele-counts-cli-output

[codex] Restore allele-count CLI output

Author: iskandr

isovar/__init__.py

@@ -10,7 +10,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "1.4.18"
+__version__ = "1.4.19"
 
 
 from .allele_read import AlleleRead

isovar/cli/isovar_allele_counts.py

@@ -18,7 +18,7 @@
 
 from ..logging import get_logger
 
-from .rna_args import make_rna_reads_arg_parser, read_evidence_dataframe_from_args
+from .rna_args import make_rna_reads_arg_parser, allele_counts_dataframe_from_args
 from .output_args import add_output_args, write_dataframe
 
 
@@ -28,15 +28,15 @@
 parser = add_output_args(
     parser,
     filename="isovar-allele-counts-result.csv",
-    description="Name of CSV file which contains read sequences")
+    description="Name of CSV file which contains read and fragment counts")
 
 
 def run(args=None):
     if args is None:
         args = sys.argv[1:]
     args = parser.parse_args(args)
     logger.info(args)
-    df = read_evidence_dataframe_from_args(args)
+    df = allele_counts_dataframe_from_args(args)
     logger.info(df)
     write_dataframe(
         df=df,

isovar/cli/rna_args.py

@@ -21,7 +21,11 @@
 from ..default_parameters import MIN_READ_MAPPING_QUALITY
 
 from ..read_collector import ReadCollector
-from ..dataframe_helpers import allele_reads_to_dataframe, read_evidence_generator_to_dataframe
+from ..dataframe_helpers import (
+    allele_counts_dataframe,
+    allele_reads_to_dataframe,
+    read_evidence_generator_to_dataframe,
+)
 
 
 def add_rna_args(
@@ -160,6 +164,15 @@ def read_evidence_dataframe_from_args(args):
         read_evidence_generator_from_args(args))
 
 
+def allele_counts_dataframe_from_args(args):
+    """
+    Collect read and fragment counts for each variant and turn them into a
+    DataFrame.
+    """
+    return allele_counts_dataframe(
+        read_evidence_generator_from_args(args))
+
+
 def allele_reads_dataframe_from_args(args):
     """
     Collect all allele reads for each variant and turn them into a DataFrame.

tests/test_allele_counts.py

@@ -25,18 +25,26 @@ def test_allele_count_dataframe():
             trimmed_alt="G",
             ref_reads=[
                 AlleleRead(prefix="AAA", allele="C", suffix="TTT", name="C1"),
-                AlleleRead(prefix="AAC", allele="C", suffix="TTA", name="C2"),
+                AlleleRead(prefix="AAC", allele="C", suffix="TTA", name="C1"),
             ],
             alt_reads=[
-                AlleleRead(prefix="AAA", allele="G", suffix="TTT", name="G1")
+                AlleleRead(prefix="AAA", allele="G", suffix="TTT", name="G1"),
+                AlleleRead(prefix="AAT", allele="G", suffix="TTC", name="G1"),
             ],
-            other_reads=[])
+            other_reads=[
+                AlleleRead(prefix="CCA", allele="T", suffix="GGA", name="T1"),
+                AlleleRead(prefix="CCG", allele="T", suffix="GGT", name="T2"),
+                AlleleRead(prefix="CCT", allele="T", suffix="GGC", name="T2"),
+            ])
     df = allele_counts_dataframe([(variant, read_evidence)])
     assert len(df) == 1, "Wrong number of rows in DataFrame: %s" % (df,)
     row = df.iloc[0]
     eq_(row.num_ref_reads, 2)
-    eq_(row.num_alt_reads, 1)
-    eq_(row.num_other_reads, 0)
+    eq_(row.num_alt_reads, 2)
+    eq_(row.num_other_reads, 3)
+    eq_(row.num_ref_fragments, 1)
+    eq_(row.num_alt_fragments, 1)
+    eq_(row.num_other_fragments, 2)
 
 
 if __name__ == "__main__":

tests/test_cli.py

@@ -26,6 +26,7 @@
 from isovar.cli.isovar_variant_sequences import run as isovar_variant_sequences
 from isovar.cli.isovar_main import run as isovar_main
 from isovar.cli.rna_args import (
+    allele_counts_dataframe_from_args,
     make_rna_reads_arg_parser,
     variants_reads_dataframe_from_args,
 )
@@ -60,7 +61,16 @@ def run_cli_fn(fn, include_bam_in_args=True, return_dataframe=False):
 
 
 def test_cli_allele_counts():
-    run_cli_fn(isovar_allele_counts)
+    df = run_cli_fn(isovar_allele_counts, return_dataframe=True)
+    args = make_rna_reads_arg_parser().parse_args(args_with_bam)
+    expected_df = allele_counts_dataframe_from_args(args)
+    assert list(df.columns) == list(expected_df.columns)
+    assert "ref_reads" not in df.columns
+    assert len(df) == 4
+    assert df[["num_ref_reads", "num_alt_reads", "num_other_reads"]].to_dict(orient="records") == \
+        expected_df[["num_ref_reads", "num_alt_reads", "num_other_reads"]].to_dict(orient="records")
+    assert df[["num_ref_fragments", "num_alt_fragments", "num_other_fragments"]].to_dict(orient="records") == \
+        expected_df[["num_ref_fragments", "num_alt_fragments", "num_other_fragments"]].to_dict(orient="records")
 
 
 def test_cli_allele_reads():