Speed up SNP scoring: vectorize dna_1hot, keep the CUDA cache between clusters

src/baskerville/dna.py

@@ -71,27 +71,27 @@ def dna_1hot(
     else:
         seq_code = np.zeros((seq_len, 4), dtype="bool")
 
-    for i in range(seq_len):
-        if i >= seq_start and i - seq_start < len(seq):
-            nt = seq[i - seq_start]
-            if nt == "A":
-                seq_code[i, 0] = 1
-            elif nt == "C":
-                seq_code[i, 1] = 1
-            elif nt == "G":
-                seq_code[i, 2] = 1
-            elif nt == "T":
-                seq_code[i, 3] = 1
-            else:
-                if n_uniform:
-                    seq_code[i, :] = 0.25
-                elif n_sample:
-                    ni = random.randint(0, 3)
-                    seq_code[i, ni] = 1
+    seq_end = min(seq_len, seq_start + len(seq))
+    rows = np.arange(seq_start, seq_end)
+    cols = _NT_COLUMN[np.frombuffer(seq.encode("ascii", "replace"), np.uint8)]
+    cols = cols[: seq_end - seq_start]
+    known = cols >= 0
+    seq_code[rows[known], cols[known]] = 1
+
+    if n_uniform:
+        seq_code[rows[~known]] = 0.25
+    elif n_sample:
+        for i in rows[~known]:
+            seq_code[i, random.randint(0, 3)] = 1
 
     return seq_code
 
 
+# column of each ASCII byte in a 1-hot encoding; -1 for N and anything else
+_NT_COLUMN = np.full(256, -1, dtype=np.int8)
+_NT_COLUMN[np.frombuffer(b"ACGT", np.uint8)] = np.arange(4)
+
+
 def dna_1hot_index(seq: str, n_sample: bool = False):
     """Convert a DNA sequence to an index encoding.
 

src/baskerville/snps.py

@@ -564,7 +564,6 @@ def hdf5_write_pair(cov_scores, gene_head_scores, score_index):
 
                 # clean reference
                 del rp_cov, ref_cov_preds
-                torch.cuda.empty_cache()
 
     # surface any cov-write exceptions (executor swallows them otherwise)
     for f in cov_write_futures:
@@ -903,7 +902,6 @@ def hdf5_write_pair(cov_scores, gene_head_scores, score_index):
                 # clean up memory
                 del rp_cov, ref_cov_preds
                 del ap_cov, alt_cov_preds
-                torch.cuda.empty_cache()
 
     # close open files
     genome_open.close()

src/tests/test_dna.py

@@ -2,9 +2,12 @@
 Tests for DNA one-hot utilities in baskerville.dna.
 """
 
+import random
+
 import numpy as np
+import pytest
 
-from baskerville.dna import hot1_rc
+from baskerville.dna import dna_1hot, hot1_rc
 
 
 def test_hot1_rc_known_value():
@@ -54,3 +57,46 @@ def test_hot1_rc_batched_matches_singleton():
     assert out.shape == batch.shape
     for i in range(batch.shape[0]):
         np.testing.assert_array_equal(out[i], hot1_rc(batch[i]))
+
+
+def _dna_1hot_per_base(seq, seq_len=None, n_uniform=False, n_sample=False):
+    """dna_1hot's original per-base loop, the reference for the vectorized one."""
+    if seq_len is None:
+        seq_len, seq_start = len(seq), 0
+    elif seq_len <= len(seq):
+        seq_trim = (len(seq) - seq_len) // 2
+        seq, seq_start = seq[seq_trim : seq_trim + seq_len], 0
+    else:
+        seq_start = (seq_len - len(seq)) // 2
+    seq = seq.upper()
+    seq_code = np.zeros((seq_len, 4), dtype="float16" if n_uniform else "bool")
+    for i in range(seq_len):
+        if i >= seq_start and i - seq_start < len(seq):
+            nt = seq[i - seq_start]
+            if nt in "ACGT":
+                seq_code[i, "ACGT".index(nt)] = 1
+            elif n_uniform:
+                seq_code[i, :] = 0.25
+            elif n_sample:
+                seq_code[i, random.randint(0, 3)] = 1
+    return seq_code
+
+
+def test_dna_1hot_known_value():
+    expected = [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1], [0, 0, 0, 0]]
+    np.testing.assert_array_equal(dna_1hot("ACgtN"), np.array(expected, dtype=bool))
+
+
+@pytest.mark.parametrize("seq_len", [None, 37, 64, 101])
+@pytest.mark.parametrize("mode", ["plain", "n_uniform", "n_sample"])
+def test_dna_1hot_matches_per_base_loop(seq_len, mode):
+    """Trimming, padding, N handling and n_sample's random draws all match."""
+    rng = np.random.default_rng(0)
+    seq = "".join(rng.choice(list("ACGTNacgtn"), 64))
+    kwargs = {"n_uniform": mode == "n_uniform", "n_sample": mode == "n_sample"}
+    random.seed(1)
+    expected = _dna_1hot_per_base(seq, seq_len, **kwargs)
+    random.seed(1)
+    got = dna_1hot(seq, seq_len, **kwargs)
+    assert got.dtype == expected.dtype
+    np.testing.assert_array_equal(got, expected)