Skip to main content

barkit_extract/
pattern.rs

1#![allow(clippy::result_large_err)]
2
3use std::{fmt, mem::size_of};
4
5use fancy_regex::Regex as FancyRegex;
6use regex::bytes::{Captures, Regex};
7
8use crate::error::Error;
9
10const FUZZY_CHARACTER: &str = ".";
11const ADAPTER_PATTERN_REGEX: &str = r"(?<!\[)\b[atgcryswkmbdhvn]+\b(?!\])";
12
13pub struct BarcodePattern {
14    adapter_pattern: FancyRegex,
15    barcode_pattern: String,
16    max_error: usize,
17}
18
19impl BarcodePattern {
20    pub fn new(pattern: &str, max_error: &usize) -> Result<Self, Error> {
21        Ok(Self {
22            adapter_pattern: FancyRegex::new(ADAPTER_PATTERN_REGEX)?,
23            barcode_pattern: pattern.to_owned(),
24            max_error: *max_error,
25        })
26    }
27
28    /// Generates sequences with errors that may occur during amplification.
29    ///
30    /// # Example
31    ///
32    /// ```
33    /// use barkit_extract::pattern::BarcodePattern;
34    ///
35    /// let barcode_pattern = BarcodePattern::new("^atgc(?<UMI>[ATGCN]{12})", &1).unwrap();
36    ///
37    /// let sequences_with_errors = barcode_pattern.get_sequence_with_errors("ATGC").unwrap();
38    /// assert_eq!(vec!["ATG.", "AT.C", "A.GC", ".TGC"], sequences_with_errors);
39    /// ```
40    pub fn get_sequence_with_errors(&self, sequence: &str) -> Result<Vec<String>, Error> {
41        if self.max_error == 0 {
42            return Ok(vec![sequence.to_string().to_ascii_uppercase()]);
43        }
44
45        if sequence.is_empty() {
46            return Ok(Vec::new());
47        }
48
49        if self.max_error >= sequence.len() {
50            return Ok(vec![FUZZY_CHARACTER.repeat(sequence.len())]);
51        }
52
53        let num_chars = sequence.chars().count();
54        assert!(num_chars <= usize::BITS as usize * 8, "too many characters");
55
56        let max_permutation_mask = usize::MAX
57            .checked_shr(size_of::<usize>() as u32 * 8 - num_chars as u32)
58            .ok_or(Error::PermutationMaskSize)?;
59
60        let mut cases = Vec::new();
61
62        let upper: Vec<char> = sequence.chars().map(|c| c.to_ascii_uppercase()).collect();
63
64        for permutation_mask in 0..=max_permutation_mask {
65            if permutation_mask.count_ones() as usize != num_chars - self.max_error {
66                continue;
67            }
68            let mut s = String::new();
69            for (idx, _) in upper.iter().enumerate().take(num_chars) {
70                if (permutation_mask & (1 << idx)) == 0 {
71                    s.push_str(FUZZY_CHARACTER)
72                } else {
73                    s.push(upper[idx])
74                }
75            }
76            cases.push(s);
77        }
78        Ok(cases)
79    }
80
81    /// Returns regex pattern with PCR errors.
82    ///
83    /// # Example
84    ///
85    /// ```
86    /// use barkit_extract::pattern::BarcodePattern;
87    ///
88    /// let barcode_pattern = BarcodePattern::new("^atgc(?<UMI>[ATGCN]{12})", &1).unwrap();
89    ///
90    /// let pattern_with_pcr_errors = barcode_pattern.get_pattern_with_errors().unwrap();
91    /// assert_eq!("^(ATG.|AT.C|A.GC|.TGC)(?<UMI>[ATGCN]{12})", pattern_with_pcr_errors);
92    /// ```
93    pub fn get_pattern_with_errors(&self) -> Result<String, Error> {
94        let mut result = String::new();
95        let mut last_end = 0;
96
97        for mat in self.adapter_pattern.find_iter(&self.barcode_pattern) {
98            let mat = mat?;
99            result.push_str(&self.barcode_pattern[last_end..mat.start()]);
100
101            let fuzzy_patterns = self.get_sequence_with_errors(mat.as_str());
102            result.push_str(&format!("({})", fuzzy_patterns?.join("|")));
103
104            last_end = mat.end();
105        }
106
107        result.push_str(&self.barcode_pattern[last_end..]);
108        Ok(result)
109    }
110}
111
112#[derive(Clone, Debug, PartialEq, Eq)]
113pub enum BarcodeType {
114    /// Moleculare barcode (UMI)
115    Umi,
116
117    /// Sample barcode
118    Sample,
119
120    /// Cell barcode
121    Cell,
122}
123
124impl BarcodeType {
125    /// Parses type of barcode
126    ///
127    /// # Example
128    ///
129    /// ```
130    /// use barkit_extract::pattern::BarcodeType;
131    /// use barkit_extract::error::Error::UnexpectedCaptureGroupName;
132    ///
133    /// assert_eq!(BarcodeType::Umi, BarcodeType::parse_type("UMI").unwrap());
134    /// assert_eq!(BarcodeType::Sample, BarcodeType::parse_type("SB").unwrap());
135    /// assert_eq!(BarcodeType::Cell, BarcodeType::parse_type("CB").unwrap());
136    /// ```
137    pub fn parse_type(name: &str) -> Result<Self, Error> {
138        match name {
139            "UMI" => Ok(BarcodeType::Umi),
140            "SB" => Ok(BarcodeType::Sample),
141            "CB" => Ok(BarcodeType::Cell),
142            _ => Err(Error::UnexpectedCaptureGroupName(name.to_owned())),
143        }
144    }
145}
146
147impl fmt::Display for BarcodeType {
148    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
149        write!(
150            f,
151            "{}",
152            match self {
153                BarcodeType::Umi => "UMI",
154                BarcodeType::Sample => "SB",
155                BarcodeType::Cell => "CB",
156            }
157        )
158    }
159}
160
161#[derive(Clone)]
162pub struct BarcodeRegex {
163    /// Regex pattern to parse barcode(s) from read sequence
164    regex: Regex,
165
166    /// List of barcode types parsed from provided pattern
167    barcode_types: Vec<BarcodeType>,
168}
169
170impl BarcodeRegex {
171    /// Creates `BarcodeRegex` instance
172    ///
173    /// Example
174    /// ```
175    /// use barkit_extract::pattern::BarcodeRegex;
176    ///
177    /// let barcode_regex = BarcodeRegex::new("^atgc(?<UMI>[ATGCN]{6})", 1);
178    /// ```
179    pub fn new(pattern: &str, max_error: usize) -> Result<Self, Error> {
180        let barcode_pattern = BarcodePattern::new(pattern, &max_error)?;
181        let fuzzy_pattern = barcode_pattern.get_pattern_with_errors()?;
182        let regex = Regex::new(&fuzzy_pattern)?;
183        let barcode_types = Self::parse_capture_groups(&regex)?;
184        Ok(Self {
185            regex,
186            barcode_types,
187        })
188    }
189
190    /// Parses capture groups from regex pattern
191    fn parse_capture_groups(regex: &Regex) -> Result<Vec<BarcodeType>, Error> {
192        let mut capture_groups = Vec::<BarcodeType>::new();
193        for capture_group in regex
194            .capture_names()
195            .collect::<Vec<_>>()
196            .into_iter()
197            .flatten()
198        {
199            capture_groups.push(BarcodeType::parse_type(capture_group)?)
200        }
201        if capture_groups.is_empty() {
202            return Err(Error::BarcodeCaptureGroupNotFound(regex.to_string()));
203        }
204        Ok(capture_groups)
205    }
206
207    /// Captures barcodes in read sequence
208    ///
209    /// Example
210    /// ```
211    /// use barkit_extract::pattern::BarcodeRegex;
212    ///
213    /// let barcode_regex = BarcodeRegex::new("^atgc(?<UMI>[ATGCN]{6})", 1).unwrap();
214    ///
215    /// assert_eq!(
216    ///     b"NNNNNN",
217    ///     barcode_regex
218    ///         .get_captures(b"ATGCNNNNNNCCC")
219    ///         .unwrap()
220    ///         .name("UMI")
221    ///         .unwrap()
222    ///         .as_bytes()
223    /// );
224    /// ```
225    pub fn get_captures<'a>(&self, read_seq: &'a [u8]) -> Result<Captures<'a>, Error> {
226        match self.regex.captures(read_seq) {
227            Some(capture) => Ok(capture),
228            None => Err(Error::PatternNotMatched),
229        }
230    }
231
232    pub fn get_barcode_types(&self) -> Vec<BarcodeType> {
233        self.barcode_types.clone()
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use rstest::rstest;
240
241    use crate::pattern;
242
243    #[rstest]
244    #[case(vec!["."], "a", 1)]
245    #[case(vec!["A"], "a", 0)]
246    #[case(vec![], "", 1)]
247    #[case(vec!["AAA.", "AA.A", "A.AA", ".AAA"], "AAAA", 1)]
248    #[case(vec!["..."], "AAA", 3)]
249    #[case(vec!["..."], "AAA", 4)]
250    fn test_generate_sequences_with_pcr_errors(
251        #[case] expected: Vec<&str>,
252        #[case] text: &str,
253        #[case] max_error: usize,
254    ) {
255        let barcode_pattern = pattern::BarcodePattern::new("", &max_error).unwrap();
256        assert_eq!(
257            expected,
258            barcode_pattern.get_sequence_with_errors(text).unwrap()
259        );
260    }
261
262    #[rstest]
263    #[case("^(AA.|A.A|.AA)(?P<UMI>[ATGCN]{3})", "^aaa(?P<UMI>[ATGCN]{3})", 1)]
264    #[case("^(...)(?P<UMI>[ATGCN]{3})", "^aaa(?P<UMI>[ATGCN]{3})", 3)]
265    #[case("^(...)(?P<UMI>[ATGCN]{3})", "^aaa(?P<UMI>[ATGCN]{3})", 4)]
266    #[case("^((...))(?P<UMI>[ATGCN]{3})", "^(aaa)(?P<UMI>[ATGCN]{3})", 4)]
267    #[case(
268        "^(AA.|A.A|.AA)(?P<UMI>[ATGCN]{3})CCC",
269        "^aaa(?P<UMI>[ATGCN]{3})CCC",
270        1
271    )]
272    #[case("^(?P<UMI>[ATGCN]{3})", "^(?P<UMI>[ATGCN]{3})", 1)]
273    fn test_create_fuzzy(#[case] expected: &str, #[case] pattern: &str, #[case] max_error: usize) {
274        let barcode_pattern = pattern::BarcodePattern::new(pattern, &max_error).unwrap();
275        assert_eq!(expected, barcode_pattern.get_pattern_with_errors().unwrap())
276    }
277}