barkit_extract/
pattern.rs1#![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 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 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 Umi,
116
117 Sample,
119
120 Cell,
122}
123
124impl BarcodeType {
125 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: Regex,
165
166 barcode_types: Vec<BarcodeType>,
168}
169
170impl BarcodeRegex {
171 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(®ex)?;
184 Ok(Self {
185 regex,
186 barcode_types,
187 })
188 }
189
190 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 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}