Merge branch 'c-python' of github.com:esrice/trio_binning into c-python

esrice · Sep 29, 2022 · 6b3cc73 · 6b3cc73
2 parents 74b33ac + 35a3805
commit 6b3cc73
Show file tree

Hide file tree

Showing 9 changed files with 75 additions and 37 deletions.
diff --git a/c/kmers.c b/c/kmers.c
@@ -58,13 +58,13 @@ uint64_t kmer_to_int(char* kmer, unsigned char k) {
                 // no need to waste time doing 'kmer_int |= (0 << (i*2))'
                 break;
             case 'C':
-                kmer_int |= (1 << (i*2));
+                kmer_int |= (UINT64_C(1) << (i*2));
                 break;
             case 'G':
-                kmer_int |= (2 << (i*2));
+                kmer_int |= (UINT64_C(2) << (i*2));
                 break;
             case 'T':
-                kmer_int |= (3 << (i*2));
+                kmer_int |= (UINT64_C(3) << (i*2));
         }
     }
 
@@ -290,7 +290,7 @@ void count_kmers_in_read(
             kmer[j] = read[i+j];
         if (kmer_in_hash_set(kmer, kmer_revcomp, haplotype_A))
             (*count_A)++;
-        if (kmer_in_hash_set(kmer, kmer_revcomp, haplotype_B))
+        else if (kmer_in_hash_set(kmer, kmer_revcomp, haplotype_B))
             (*count_B)++;
     }
 

diff --git a/src/trio_binning/kmers.py b/src/trio_binning/kmers.py
@@ -72,6 +72,27 @@ class _HashSet(Structure):
     POINTER(c_int),
 ]
 
+kmer_to_int_c = lib.kmer_to_int
+kmer_to_int_c.argtypes = [c_char_p, c_ubyte]
+kmer_to_int_c.restype = c_uint64
+
+
+def kmer_to_int(kmer: str) -> int:
+    """Convert a kmer to integer format"""
+    return kmer_to_int_c(bytes(kmer, "utf-8"), c_ubyte(len(kmer)))
+
+
+reverse_complement_c = lib.reverse_complement
+reverse_complement_c.argtypes = [c_char_p, c_char_p, c_ubyte]
+
+
+def reverse_complement(kmer: str) -> str:
+    """Reverse complement a k-mer"""
+    out_kmer = bytes("x" * len(kmer), "utf-8")
+    reverse_complement_c(bytes(kmer, "utf-8"), out_kmer, len(kmer))
+    return out_kmer.decode("utf-8")
+
+
 # this is ugly as sin, but necessary because mypy is ok with `pointer`
 # as a subscriptable type while python runtime is not
 if TYPE_CHECKING:

diff --git a/tests/data/hapA.fastq b/tests/data/hapA.fastq
diff --git a/tests/data/hapA.txt b/tests/data/hapA.txt
@@ -1,6 +1,4 @@
-CTACCGCCCTCCC
-GAAGATCCTCATC
-GAGGAGATTTAGA
-CTACGGCGCTCAC
-GTGAGCCCCCCGA
-CCTGGGCACACAC
+ACCTCTAAGAAGCTTTGAAAA
+CTTATCATGTCTTTGTTTTCA
+CCAGCTGTGTGATTGGTCCGC
+AAAAATAAGACAGCATGTGAC
diff --git a/tests/data/hapB.txt b/tests/data/hapB.txt
@@ -1,7 +1,3 @@
-ATGCGTAGCAGTA
-CATCGCCATCAGC
-CCCCTCCCTGCCC
-GCTCGACTCACAC
-CACTGCGCATCTC
-ATAGAGATATATA
-CTCCCGCCCTCCC
+AACACCAAAAAAAAAAACCTC
+AACCCTGTTGTTACCCTTGTG
+AAACTGTTTTAAAACTAAAGT
diff --git a/tests/data/test.ccs.fastq.gz b/tests/data/test.ccs.fastq.gz
diff --git a/tests/data/unclassified.fastq b/tests/data/unclassified.fastq
diff --git a/tests/test_classify_by_kmers.py b/tests/test_classify_by_kmers.py
@@ -21,7 +21,7 @@ def test_classify_by_kmers(capsys, tmpdir):
         "sys.argv",
         [
             "classify-by-kmers",
-            join(dirname(__file__), "data", "test.fastq"),
+            join(dirname(__file__), "data", "test.ccs.fastq.gz"),
             join(dirname(__file__), "data", "hapA.txt"),
             join(dirname(__file__), "data", "hapB.txt"),
             "--haplotype-a-out-prefix",

diff --git a/tests/test_kmers.py b/tests/test_kmers.py
@@ -1,12 +1,14 @@
 import os.path
 
-import kmers
+import pytest
+
+from trio_binning import kmers
 
 
 def test_build_kmer_hash():
     kmers_path = os.path.join(os.path.dirname(__file__), "data", "hapA.txt")
     kmer_set = kmers.create_kmer_hash_set(kmers_path)
-    assert kmers.get_number_kmers_in_set(kmer_set) == 6
+    assert kmers.get_number_kmers_in_set(kmer_set) == 4
 
 
 def test_count_kmers_in_read():
@@ -17,5 +19,35 @@ def test_count_kmers_in_read():
     hap_b_set = kmers.create_kmer_hash_set(hap_b_path)
 
     assert kmers.count_kmers_in_read(
-        "GAGGAGATTTAGAGTGTGAGTCGAGCATAGAGATATATA", hap_a_set, hap_b_set
-    ) == (1, 2)
+        "CTTATCATGTCTTTGTTTTCAAAGCTTCTTAGAGGTTTTTTTTTTTGGTGTTAATTGGCATAAATTATGGCT",
+        hap_a_set,
+        hap_b_set,
+    ) == (2, 1)
+
+
+@pytest.mark.parametrize(
+    "kmer_str,kmer_int",
+    [
+        ("ATGCTAGCTAGAGAGAGAGGA", 696357446508),
+        ("TTTTTTTTTTTTTTTTTTTTTTTTTTTT", 72057594037927935),
+        ("TTTTTTTTTTTTTTTAGGCCCACTTTTT", 72006304612220927),
+        ("AAAAAAAAAAAAAAAAAAAAAAAAAA", 0),
+        ("GGGAGGGAGGGAGGGAGGGAGGGAGGG", 11868309606246954),
+    ],
+)
+def test_kmer_to_int(kmer_str, kmer_int):
+    assert kmers.kmer_to_int(kmer_str) == kmer_int
+
+
+@pytest.mark.parametrize(
+    "kmer,revcomp_kmer",
+    [
+        ("ATGCTAGCTAGAGAGAGAGGA", "TCCTCTCTCTCTAGCTAGCAT"),
+        ("TTTTTTTTTTTTTTTTTTTTTTTTTTTT", "AAAAAAAAAAAAAAAAAAAAAAAAAAAA"),
+        ("TTTTTTTTTTTTTTTAGGCCCACTTTTT", "AAAAAGTGGGCCTAAAAAAAAAAAAAAA"),
+        ("AAAAAAAAAAAAAAAAAAAAAAAAAA", "TTTTTTTTTTTTTTTTTTTTTTTTTT"),
+        ("GGGAGGGAGGGAGGGAGGGAGGGAGGG", "CCCTCCCTCCCTCCCTCCCTCCCTCCC"),
+    ],
+)
+def test_reverse_complement(kmer, revcomp_kmer):
+    assert kmers.reverse_complement(kmer) == revcomp_kmer