diff --git a/src/baskerville/dna.py b/src/baskerville/dna.py index a341bfb..24ae99a 100644 --- a/src/baskerville/dna.py +++ b/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. diff --git a/src/baskerville/snps.py b/src/baskerville/snps.py index 635f021..44f2d32 100644 --- a/src/baskerville/snps.py +++ b/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() diff --git a/src/tests/test_dna.py b/src/tests/test_dna.py index 34210a2..6af89d6 100644 --- a/src/tests/test_dna.py +++ b/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)