packed_seq/
packed_n_seq.rs

1use wide::i8x32;
2
3use crate::packed_seq::read_slice_32;
4
5use super::*;
6
7#[derive(Clone, Copy, Debug, MemSize, MemDbg)]
8// #[cfg_attr(feature = "pyo3", pyo3::pyclass)] // Not supported: pyclass cannot have lifetime parameters
9pub struct PackedNSeq<'s> {
10    pub seq: PackedSeq<'s>,
11    pub ambiguous: BitSeq<'s>,
12}
13
14#[derive(Clone, Debug, MemSize, MemDbg, Default)]
15#[cfg_attr(feature = "pyo3", pyo3::pyclass)]
16#[cfg_attr(feature = "epserde", derive(epserde::Epserde))]
17pub struct PackedNSeqVec {
18    pub seq: PackedSeqVec,
19    pub ambiguous: BitSeqVec,
20}
21
22/// Implement a subset of `Seq` for `PackedNSeq`.
23impl<'s> PackedNSeq<'s> {
24    pub fn to_revcomp(&self) -> PackedNSeqVec {
25        PackedNSeqVec {
26            seq: self.seq.to_revcomp(),
27            ambiguous: self.ambiguous.to_revcomp(),
28        }
29    }
30
31    pub fn slice(&self, range: Range<usize>) -> PackedNSeq<'s> {
32        PackedNSeq {
33            seq: self.seq.slice(range.clone()),
34            ambiguous: self.ambiguous.slice(range),
35        }
36    }
37}
38
39/// Implement a subset of `SeqVec` for `PackedNSeqVec`.
40impl PackedNSeqVec {
41    /// Given a bunch of `PackedNSeqVec` chunks, assemble them into a single `PackedNSeqVec` in parallel.
42    #[cfg(feature = "file_io")]
43    fn assemble_parallel_chunks(chunks: Vec<Self>, num_threads: usize) -> Self {
44        let len = chunks
45            .iter()
46            .map(|chunk| chunk.seq.len().next_multiple_of(8))
47            .sum();
48        let mut output = Self {
49            seq: PackedSeqVec::from_raw_parts(vec![0; len / 4], len),
50            ambiguous: BitSeqVec::with_len(len),
51        };
52        let mut seq_out = &mut output.seq.seq[..len / 4];
53        let mut mask_out = &mut output.ambiguous.seq[..len / 8];
54        let mut jobs = Vec::with_capacity(chunks.len());
55        for chunk in chunks {
56            let aligned_len = chunk.seq.len().next_multiple_of(8);
57            let (seq_dest, rest) = seq_out.split_at_mut(aligned_len / 4);
58            seq_out = rest;
59            let (mask_dest, rest) = mask_out.split_at_mut(aligned_len / 8);
60            mask_out = rest;
61            jobs.push((chunk, seq_dest, mask_dest));
62        }
63        let jobs = std::sync::Mutex::new(jobs.into_iter());
64        std::thread::scope(|scope| {
65            for _ in 0..num_threads.max(1) {
66                let jobs = &jobs;
67                scope.spawn(move || {
68                    loop {
69                        let job = jobs.lock().unwrap().next();
70                        let Some((chunk, seq_dest, mask_dest)) = job else {
71                            break;
72                        };
73                        let chunk_len = chunk.seq.len();
74                        let seq_bytes = chunk_len.div_ceil(4);
75                        let mask_bytes = chunk_len.div_ceil(8);
76                        seq_dest[..seq_bytes].copy_from_slice(&chunk.seq.seq[..seq_bytes]);
77                        mask_dest[..mask_bytes].copy_from_slice(&chunk.ambiguous.seq[..mask_bytes]);
78                        // Fill alignment padding with N in both representations.
79                        seq_dest[seq_bytes..].fill(u8::MAX);
80                        if chunk_len % 4 != 0 {
81                            seq_dest[seq_bytes - 1] |= u8::MAX << (2 * (chunk_len % 4));
82                        }
83                        if chunk_len % 8 != 0 {
84                            mask_dest[mask_bytes - 1] |= u8::MAX << (chunk_len % 8);
85                        }
86                    }
87                });
88            }
89        });
90        output
91    }
92
93    pub fn clear(&mut self) {
94        self.seq.clear();
95        self.ambiguous.clear();
96    }
97
98    #[cfg(feature = "rand")]
99    pub fn random(len: usize, n_frac: f32) -> Self {
100        let seq = PackedSeqVec::random(len);
101        let ambiguous = BitSeqVec::random(len, n_frac);
102        Self { seq, ambiguous }
103    }
104
105    pub fn as_slice(&self) -> PackedNSeq<'_> {
106        PackedNSeq {
107            seq: self.seq.as_slice(),
108            ambiguous: self.ambiguous.as_slice(),
109        }
110    }
111
112    pub fn slice(&self, range: Range<usize>) -> PackedNSeq<'_> {
113        self.as_slice().slice(range).to_owned()
114    }
115
116    pub fn push_ascii(&mut self, seq: &[u8]) -> Range<usize> {
117        let r1 = self.seq.push_ascii(seq);
118        let r2 = self.ambiguous.push_ascii(seq);
119        assert_eq!(r1, r2);
120        r1
121    }
122
123    pub fn from_ascii(seq: &[u8]) -> Self {
124        Self {
125            seq: PackedSeqVec::from_ascii(seq),
126            ambiguous: BitSeqVec::from_ascii(seq),
127        }
128    }
129
130    /// Create a mask that is 1 for all non-ACGT bases and for all low-quality bases with quality `<threshold`.
131    pub fn from_ascii_and_quality(seq: &[u8], quality: &[u8], threshold: usize) -> Self {
132        assert_eq!(seq.len(), quality.len());
133
134        let mut ambiguous = BitSeqVec::from_ascii(seq);
135
136        // Low-quality bases are also ambiguous.
137        {
138            let t = (b'!' + threshold as u8) as i8;
139            let t_simd = i8x32::splat(t);
140            let ambiguous = ambiguous.seq.as_chunks_mut::<4>().0;
141            for i in (0..quality.len()).step_by(32) {
142                let chunk = i8x32::from(unsafe {
143                    std::mem::transmute::<_, i8x32>(read_slice_32(quality, i))
144                });
145
146                let mut mask = chunk.simd_lt(t_simd).to_bitmask() as u32;
147                let remaining = quality.len() - i;
148                if remaining < 32 {
149                    mask &= (1u32 << remaining) - 1;
150                }
151                let ambi = &mut ambiguous[i / 32];
152                *ambi = (u32::from_ne_bytes(*ambi) | mask).to_ne_bytes();
153            }
154        }
155
156        Self {
157            seq: PackedSeqVec::from_ascii(seq),
158            ambiguous,
159        }
160    }
161
162    /// Create a mask that is 1 for all non-ACGT bases and for all low-quality bases with quality `<threshold`.
163    pub fn push_from_ascii_and_quality(&mut self, seq: &[u8], quality: &[u8], min_qual: u8) {
164        let r = self.seq.push_ascii(seq);
165        let r2 = self.ambiguous.push_ascii(seq);
166        assert_eq!(r, r2);
167
168        assert_eq!(seq.len(), quality.len());
169
170        // Low-quality bases are also ambiguous.
171        let t = b'!' + min_qual;
172        let t_simd = i8x32::splat(t as i8);
173
174        let shift = r2.start % 8;
175        let ambiguous = &mut self.ambiguous.seq[r2.start / 8..];
176        for i in (0..quality.len()).step_by(32) {
177            let chunk =
178                i8x32::from(unsafe { std::mem::transmute::<_, i8x32>(read_slice_32(quality, i)) });
179
180            let mut mask = chunk.simd_lt(t_simd).to_bitmask() as u32;
181            let remaining = quality.len() - i;
182            if remaining < 32 {
183                mask &= (1u32 << remaining) - 1;
184            }
185            let mask = (mask as u64) << shift;
186            let offset = i / 8;
187            let word = u32::from_le_bytes(ambiguous[offset..offset + 4].try_into().unwrap());
188            ambiguous[offset..offset + 4].copy_from_slice(&(word | mask as u32).to_le_bytes());
189            ambiguous[offset + 4] |= (mask >> 32) as u8;
190        }
191    }
192
193    /// Read a fasta/fastq file at the given path into a `PackedNSeq`, separating records with `N`.
194    #[cfg(feature = "file_io")]
195    pub fn from_fastx(path: &std::path::Path) -> Self {
196        use helicase::{HelicaseParser, input::FromFile};
197
198        const CONFIG: helicase::Config = helicase::ParserOptions::default()
199            .ignore_headers()
200            .dna_string()
201            .config();
202        let mut ascii = Vec::with_capacity(16000);
203
204        let mut seq = Self::default();
205
206        let mut reader = helicase::FastxParser::<CONFIG>::from_file(path).unwrap();
207        while let Some(event) = reader.next() {
208            event.expect("Invalid FASTA/Q record");
209            ascii.extend_from_slice(reader.get_dna_string());
210            ascii.push(b'N');
211
212            // Convert from ascii dna to packed dna+mask in batches of 16kbp.
213            if ascii.len() > 16000 {
214                seq.push_ascii(&ascii);
215                ascii.clear();
216            }
217        }
218
219        if ascii.len() > 0 {
220            seq.push_ascii(&ascii);
221        }
222
223        seq
224    }
225
226    /// Read FASTA/FASTQ using Helicase's parallel reader, separating records with `N`.
227    ///
228    /// Does *not* preserve the order or records!
229    ///
230    /// Chunk boundaries may contain extra `N`s for byte alignment.
231    #[cfg(feature = "file_io")]
232    pub fn from_fastx_parallel(path: &std::path::Path, num_threads: usize) -> Self {
233        use helicase::HelicaseParser;
234
235        const CONFIG: helicase::Config = helicase::ParserOptions::default()
236            .ignore_headers()
237            .ignore_quality()
238            .dna_string()
239            .config();
240        let chunks = std::sync::Mutex::new(Vec::new());
241        helicase::ParallelFastxParser::<CONFIG>::from_file(path)
242            .unwrap()
243            .threads(num_threads)
244            .chunks(num_threads.max(1))
245            .for_each(|mut parser| {
246                let mut seq = Self::default();
247                let mut ascii = Vec::with_capacity(16000);
248                while let Some(event) = parser.next() {
249                    event?;
250                    ascii.extend_from_slice(parser.get_dna_string());
251                    ascii.push(b'N');
252                    if ascii.len() >= 16000 {
253                        seq.push_ascii(&ascii);
254                        ascii.clear();
255                    }
256                }
257                if !ascii.is_empty() {
258                    seq.push_ascii(&ascii);
259                }
260                chunks.lock().unwrap().push(seq);
261                Ok(())
262            })
263            .unwrap();
264        let chunks = chunks.into_inner().unwrap();
265        Self::assemble_parallel_chunks(chunks, num_threads)
266    }
267
268    /// Read a fasta/fastq file at the given path into a `PackedNSeq`, separating records with `N`.
269    /// Low-quality bases are masked out.
270    #[cfg(feature = "file_io")]
271    pub fn from_fastq_with_quality(path: &std::path::Path, min_qual: u8) -> Self {
272        use helicase::{HelicaseParser, input::FromFile};
273
274        const CONFIG: helicase::Config = helicase::ParserOptions::default()
275            .ignore_headers()
276            .dna_string()
277            .compute_quality()
278            .config();
279        let mut ascii = vec![];
280        let mut qual = vec![];
281
282        let mut seq = Self::default();
283
284        let mut reader = helicase::FastxParser::<CONFIG>::from_file(path).unwrap();
285        while let Some(event) = reader.next() {
286            event.expect("Invalid FASTA/Q record");
287            ascii.extend_from_slice(reader.get_dna_string());
288            qual.extend_from_slice(reader.get_quality().expect("Expected FASTQ quality scores"));
289            ascii.push(b'N');
290            qual.push(0);
291
292            // Convert from ascii dna+qual to packed dna+mask in batches of 16kbp.
293            if ascii.len() > 16000 {
294                seq.push_from_ascii_and_quality(&ascii, &qual, min_qual);
295                ascii.clear();
296                qual.clear();
297            }
298        }
299
300        if ascii.len() > 0 {
301            seq.push_from_ascii_and_quality(&ascii, &qual, min_qual);
302        }
303
304        seq
305    }
306
307    /// Read FASTQ using Helicase's parallel reader, separating records with `N`.
308    ///
309    /// Masks non-ACGT bases and qualities below `min_qual`.
310    /// Does *not* preserve the order or records!
311    /// Chunk boundaries may contain extra `N`s for byte alignment.
312    #[cfg(feature = "file_io")]
313    pub fn from_fastq_with_quality_parallel(
314        path: &std::path::Path,
315        min_qual: u8,
316        num_threads: usize,
317    ) -> Self {
318        use helicase::HelicaseParser;
319
320        const CONFIG: helicase::Config = helicase::ParserOptions::default()
321            .ignore_headers()
322            .dna_string()
323            .compute_quality()
324            .config();
325        let chunks = std::sync::Mutex::new(Vec::new());
326        helicase::ParallelFastxParser::<CONFIG>::from_file(path)
327            .unwrap()
328            .threads(num_threads)
329            .chunks(num_threads.max(1))
330            .for_each(|mut parser| {
331                let mut seq = Self::default();
332                let mut ascii = Vec::with_capacity(16000);
333                let mut quality = Vec::with_capacity(16000);
334                while let Some(event) = parser.next() {
335                    event?;
336                    ascii.extend_from_slice(parser.get_dna_string());
337                    ascii.push(b'N');
338                    quality.extend_from_slice(
339                        parser.get_quality().expect("Expected FASTQ quality scores"),
340                    );
341                    quality.push(0);
342                    if ascii.len() >= 16000 {
343                        seq.push_from_ascii_and_quality(&ascii, &quality, min_qual);
344                        ascii.clear();
345                        quality.clear();
346                    }
347                }
348                if !ascii.is_empty() {
349                    seq.push_from_ascii_and_quality(&ascii, &quality, min_qual);
350                }
351                chunks.lock().unwrap().push(seq);
352                Ok(())
353            })
354            .unwrap();
355        let chunks = chunks.into_inner().unwrap();
356        Self::assemble_parallel_chunks(chunks, num_threads)
357    }
358}
packed_seq/packed_n_seq.rs

packed_seq/
packed_n_seq.rs
