haystackfm 0.3.0

GPU-accelerated FM-index construction for DNA sequences via WebGPU
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

//! GPU locate parity tests: assert locate_gpu returns the same multiset of
//! (seq_header, pos) pairs as the CPU locate for random corpora and queries.

#[cfg(feature = "gpu")]
mod tests {
    use haystackfm::alphabet::encode_char;
    use haystackfm::{DnaSequence, FmIndex, FmIndexConfig, SeqId};
    use pollster::FutureExt as _;

    fn cpu_config() -> FmIndexConfig {
        FmIndexConfig {
            sa_sample_rate: 1,
            use_gpu: false,
            ..Default::default()
        }
    }

    fn encode(s: &str) -> Vec<u8> {
        s.chars().map(|c| encode_char(c).unwrap()).collect()
    }

    fn sorted_hits(mut v: Vec<(SeqId, u32)>) -> Vec<(SeqId, u32)> {
        v.sort();
        v
    }

    fn build(seqs: &[&str]) -> FmIndex {
        let dna: Vec<DnaSequence> = seqs
            .iter()
            .map(|s| DnaSequence::from_str(s).unwrap())
            .collect();
        FmIndex::build_cpu(&dna, &cpu_config()).unwrap()
    }

    fn locate_gpu_sync(idx: &FmIndex, queries: &[Vec<u8>]) -> Vec<Vec<(SeqId, u32)>> {
        let refs: Vec<&[u8]> = queries.iter().map(|q| q.as_slice()).collect();
        idx.locate_gpu(&refs).block_on().unwrap()
    }

    #[test]
    fn single_seq_single_query() {
        let idx = build(&["ACGTACGT"]);
        let pat = encode("ACGT");
        let cpu = idx.locate(&pat);
        let gpu = locate_gpu_sync(&idx, &[pat])[0].clone();
        assert_eq!(sorted_hits(cpu), sorted_hits(gpu));
    }

    #[test]
    fn no_match() {
        let idx = build(&["ACGTACGT"]);
        let pat = encode("TTTT");
        let cpu = idx.locate(&pat);
        let gpu = locate_gpu_sync(&idx, &[pat])[0].clone();
        assert!(cpu.is_empty());
        assert!(gpu.is_empty());
    }

    #[test]
    fn empty_pattern_matches_all() {
        let idx = build(&["ACGT"]);
        let cpu = idx.locate(&[]);
        let gpu = locate_gpu_sync(&idx, &[vec![]])[0].clone();
        assert_eq!(sorted_hits(cpu), sorted_hits(gpu));
    }

    #[test]
    fn single_char_repeated() {
        let idx = build(&["AAAACCCGGG"]);
        for ch in ["A", "C", "G"] {
            let pat = encode(ch);
            let cpu = idx.locate(&pat);
            let gpu = locate_gpu_sync(&idx, &[pat.clone()])[0].clone();
            assert_eq!(sorted_hits(cpu), sorted_hits(gpu), "char={ch}");
        }
    }

    #[test]
    fn multi_seq_single_query() {
        let idx = build(&["ACGT", "TGCA"]);
        let pat = encode("A");
        let cpu = idx.locate(&pat);
        let gpu = locate_gpu_sync(&idx, &[pat])[0].clone();
        assert_eq!(sorted_hits(cpu), sorted_hits(gpu));
    }

    #[test]
    fn batch_multiple_queries() {
        let idx = build(&["ACGTACGT", "NNNACGTNNN"]);
        let patterns_valid = vec![encode("ACG"), encode("ACGT")];
        let cpu_results: Vec<Vec<(SeqId, u32)>> =
            patterns_valid.iter().map(|p| idx.locate(p)).collect();
        let gpu_results = locate_gpu_sync(&idx, &patterns_valid);
        for (i, (cpu, gpu)) in cpu_results.iter().zip(gpu_results.iter()).enumerate() {
            assert_eq!(
                sorted_hits(cpu.clone()),
                sorted_hits(gpu.clone()),
                "query {i}"
            );
        }
    }

    #[test]
    fn pattern_with_n() {
        let idx = build(&["ACNGT"]);
        let pat = encode("N");
        let cpu = idx.locate(&pat);
        let gpu = locate_gpu_sync(&idx, &[pat])[0].clone();
        assert_eq!(sorted_hits(cpu), sorted_hits(gpu));
    }

    #[test]
    fn many_sequences() {
        let seqs: Vec<&str> = vec!["ACGT", "TTTT", "CCCC", "GGGG", "ACGTACGT"];
        let idx = build(&seqs);
        let pat = encode("ACGT");
        let cpu = idx.locate(&pat);
        let gpu = locate_gpu_sync(&idx, &[pat])[0].clone();
        assert_eq!(sorted_hits(cpu), sorted_hits(gpu));
    }

    #[test]
    fn large_batch() {
        let idx = build(&["ACGTACGTACGTACGT"]);
        let patterns: Vec<Vec<u8>> = vec![
            encode("A"),
            encode("C"),
            encode("G"),
            encode("T"),
            encode("AC"),
            encode("CG"),
            encode("GT"),
            encode("ACG"),
        ];
        let cpu_results: Vec<Vec<(SeqId, u32)>> = patterns.iter().map(|p| idx.locate(p)).collect();
        let gpu_results = locate_gpu_sync(&idx, &patterns);
        for (i, (cpu, gpu)) in cpu_results.iter().zip(gpu_results.iter()).enumerate() {
            assert_eq!(
                sorted_hits(cpu.clone()),
                sorted_hits(gpu.clone()),
                "query {i}"
            );
        }
    }
}