kmers_rs/
frequency.rs

1use crate::Kmer;
2
3/// An adaptor that coverts an iterator over k-mers into a k-mer frequency iterator
4pub fn unique<Src>(itr: Src) -> Unique<Src>
5where
6    Src: Iterator<Item = Kmer>,
7{
8    Unique::new(itr)
9}
10
11/// An iterator producing k-mer frequencies.
12pub struct Unique<Src>
13where
14    Src: Iterator<Item = Kmer>,
15{
16    src: Src,
17}
18
19impl<Src> Unique<Src>
20where
21    Src: Iterator<Item = Kmer>,
22{
23    fn new(src: Src) -> Unique<Src> {
24        Unique { src }
25    }
26}
27
28impl<Src> Iterator for Unique<Src>
29where
30    Src: Iterator<Item = Kmer>,
31{
32    type Item = (Kmer, usize);
33
34    fn next(&mut self) -> Option<Self::Item> {
35        self.src.next().map(|x| (x, 1))
36    }
37}
38
39/// Construct a k-mer frequency iterator from an iterator over frequencies.
40pub fn frequency_vector_iter<'a, Src>(itr: Src) -> KmerFrequencyIterator<'a, Src>
41where
42    Src: Iterator<Item = &'a usize>,
43{
44    KmerFrequencyIterator::new(itr)
45}
46
47/// An iterator producing k-mer frequencies.
48pub struct KmerFrequencyIterator<'a, Src>
49where
50    Src: Iterator<Item = &'a usize>,
51{
52    x: usize,
53    src: Src,
54}
55
56impl<'a, Src> KmerFrequencyIterator<'a, Src>
57where
58    Src: Iterator<Item = &'a usize>,
59{
60    fn new(src: Src) -> KmerFrequencyIterator<'a, Src> {
61        KmerFrequencyIterator { x: 0, src }
62    }
63}
64
65impl<'a, Src> Iterator for KmerFrequencyIterator<'a, Src>
66where
67    Src: Iterator<Item = &'a usize>,
68{
69    type Item = (Kmer, usize);
70
71    fn next(&mut self) -> Option<Self::Item> {
72        while let Some(freq) = self.src.next() {
73            if *freq == 0 {
74                self.x += 1;
75                continue;
76            }
77            let x = Kmer::from_u64(self.x as u64);
78            self.x += 1;
79            return Some((x, *freq));
80        }
81        None
82    }
83}
84
85/// Take a sorted iterator and yield k-mer frequencies.
86pub fn counting_kmer_frequency_iterator<'a, Src>(src: Src) -> CountingIterator<'a, Src>
87where
88    Src: Iterator<Item = &'a Kmer>,
89{
90    CountingIterator::new(src)
91}
92
93/// Take a sorted iterator and yield k-mer frequencies.
94pub struct CountingIterator<'a, Src>
95where
96    Src: Iterator<Item = &'a Kmer>,
97{
98    src: Src,
99    src_next: Option<&'a Kmer>
100}
101
102impl<'a, Src> CountingIterator<'a, Src>
103where
104    Src: Iterator<Item = &'a Kmer>,
105{
106    fn new(mut src: Src) -> CountingIterator<'a, Src> {
107        let src_next = src.next();
108        CountingIterator { src, src_next }
109    }
110}
111
112impl<'a, Src> Iterator for CountingIterator<'a, Src>
113where
114    Src: Iterator<Item = &'a Kmer>,
115{
116    type Item = (Kmer, usize);
117
118    fn next(&mut self) -> Option<Self::Item> {
119        if let Some(x) = self.src_next {
120            let x = x.clone();
121            let mut f = 1;
122            self.src_next = self.src.next();
123            while let Some(y) = self.src_next {
124                if *y != x {
125                    break;
126                }
127                f += 1;
128                self.src_next = self.src.next();
129            }
130            return Some((x, f));
131        }
132        None
133    }
134}
135
136/// Merge two iterators
137pub fn merge<Lhs, Rhs>(lhs: Lhs, rhs: Rhs) -> MergeIterator<Lhs, Rhs>
138where
139    Lhs: Iterator<Item = (Kmer, usize)>,
140    Rhs: Iterator<Item = (Kmer, usize)>,
141{
142    MergeIterator::new(lhs, rhs)
143}
144
145/// Merge two k-mer frequency iterators.
146pub struct MergeIterator<Lhs, Rhs>
147where
148    Lhs: Iterator<Item = (Kmer, usize)>,
149    Rhs: Iterator<Item = (Kmer, usize)>,
150{
151    lhs: Lhs,
152    lhs_next: Option<(Kmer, usize)>,
153    rhs: Rhs,
154    rhs_next: Option<(Kmer, usize)>,
155}
156
157impl<Lhs, Rhs> MergeIterator<Lhs, Rhs>
158where
159    Lhs: Iterator<Item = (Kmer, usize)>,
160    Rhs: Iterator<Item = (Kmer, usize)>,
161{
162    fn new(mut lhs: Lhs, mut rhs: Rhs) -> MergeIterator<Lhs, Rhs> {
163        let lhs_next = lhs.next();
164        let rhs_next = rhs.next();
165        MergeIterator {
166            lhs,
167            lhs_next,
168            rhs,
169            rhs_next,
170        }
171    }
172}
173
174impl<Lhs, Rhs> Iterator for MergeIterator<Lhs, Rhs>
175where
176    Lhs: Iterator<Item = (Kmer, usize)>,
177    Rhs: Iterator<Item = (Kmer, usize)>,
178{
179    type Item = (Kmer, usize);
180
181    fn next(&mut self) -> Option<Self::Item> {
182        while let (Some((x, x_f)), Some((y, y_f))) = (&self.lhs_next, &self.rhs_next) {
183            if x < y {
184                let res = (x.clone(), *x_f);
185                self.lhs_next = self.lhs.next();
186                return Some(res);
187            }
188            if y < x {
189                let res = (y.clone(), *y_f);
190                self.rhs_next = self.rhs.next();
191                return Some(res);
192            }
193            let res = (x.clone(), *x_f + *y_f);
194            self.lhs_next = self.lhs.next();
195            self.rhs_next = self.rhs.next();
196            return Some(res);
197        }
198        while let Some((x, x_f)) = &self.lhs_next {
199            let res = (x.clone(), *x_f);
200            self.lhs_next = self.lhs.next();
201            return Some(res);
202        }
203        while let Some((y, y_f)) = &self.rhs_next {
204            let res = (y.clone(), *y_f);
205            self.rhs_next = self.rhs.next();
206            return Some(res);
207        }
208        None
209    }
210}
kmers_rs/frequency.rs

kmers_rs/
frequency.rs
