bascet 0.4.0

Bascet is a tool to preprocess single-cell data, handling barcode detection, trimming, QC, and managing the execution of custom tools for each cell
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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243

use std::{cmp::Ordering, hash::Hasher};

use fasthash::FastHasher;

////////////// Lookup table for N where N is any of ATCG. Maps to 0..3
const NT1_LOOKUP: [u8; (b'T' - b'A' + 1) as usize] = {
    let mut table = [0u8; (b'T' - b'A' + 1) as usize];
    table[(b'A' - b'A') as usize] = 0b00;
    table[(b'T' - b'A') as usize] = 0b01;
    table[(b'G' - b'A') as usize] = 0b10;
    table[(b'C' - b'A') as usize] = 0b11;
    table
};


const fn generate_nt4_value(a: u8, b: u8, c: u8, d: u8) -> u8 {
    (NT1_LOOKUP[(a - b'A') as usize] << 6)
        | (NT1_LOOKUP[(b - b'A') as usize] << 4)
        | (NT1_LOOKUP[(c - b'A') as usize] << 2)
        | NT1_LOOKUP[(d - b'A') as usize]
}


const fn calculate_index(a: u8, b: u8, c: u8, d: u8) -> usize {
    const DIM: usize = (b'T' - b'A' + 1) as usize;
    ((a - b'A') as usize)
        + ((b - b'A') as usize * DIM)
        + ((c - b'A') as usize * DIM * DIM)
        + ((d - b'A') as usize * DIM * DIM * DIM)
}


////////////// Lookup table for NNNN where N is any of ATCG.
////////////// Maps compressed ATCG (usize) to 0..255 ie compresses it to a single byte
const fn generate_nt4_table() -> [u8; NT4_DIMSIZE * NT4_DIMSIZE * NT4_DIMSIZE * NT4_DIMSIZE] {
    const NUCLEOTIDES: [u8; 4] = [b'A', b'T', b'G', b'C'];
    let mut table = [0u8; NT4_DIMSIZE * NT4_DIMSIZE * NT4_DIMSIZE * NT4_DIMSIZE];

    let mut i = 0;
    while i < 256 {
        // 4^4 combinations
        let n1 = NUCLEOTIDES[(i >> 6) & 0b11];
        let n2 = NUCLEOTIDES[(i >> 4) & 0b11];
        let n3 = NUCLEOTIDES[(i >> 2) & 0b11];
        let n4 = NUCLEOTIDES[i & 0b11];

        let idx = calculate_index(n1, n2, n3, n4);
        table[idx] = generate_nt4_value(n1, n2, n3, n4);

        i += 1;
    }
    table
}
const NT4_DIMSIZE: usize = (b'T' - b'A' + 1) as usize;
const NT4_LOOKUP: [u8; NT4_DIMSIZE * NT4_DIMSIZE * NT4_DIMSIZE * NT4_DIMSIZE] =  /////// Map compressed ATCG => single byte
    generate_nt4_table();

const NT_REVERSE: [u8; 4] = [b'A', b'T', b'G', b'C'];






//NOTE: all of this can probably make use of SIMD operations but I do not know how that'd work

//////////// KMER encoder, for a given KMER-size
#[derive(Clone, Copy)]
pub struct KMERCodec {
    pub kmer_size: usize,
}
impl KMERCodec {
    pub const fn new(kmer_size: usize) -> Self {
        Self {
            kmer_size: kmer_size,
        }
    }

    //////////// Encode a kmer + count + random value
    #[inline(always)]
    pub unsafe fn encode(&self, bytes: &[u8], count: u32) -> KMERandCount {
        let chunk_size: usize = 4;
        let kmer_size = self.kmer_size as usize;
        let full_chunks = kmer_size / chunk_size;
        let remainder = kmer_size % chunk_size;

        let mut encoded = 0;
        let ptr = bytes.as_ptr();

        // Compress chunks of 4 nucleotides to 1-byte encoding
        for i in 0..full_chunks {
            let chunk_ptr = ptr.add(i * chunk_size);
            let idx = unsafe {
                (*chunk_ptr.offset(0) - b'A') as usize
                    + ((*chunk_ptr.offset(1) - b'A') as usize * NT4_DIMSIZE)
                    + ((*chunk_ptr.offset(2) - b'A') as usize * NT4_DIMSIZE * NT4_DIMSIZE)
                    + ((*chunk_ptr.offset(3) - b'A') as usize * NT4_DIMSIZE * NT4_DIMSIZE * NT4_DIMSIZE)
            };
            encoded = (encoded << 8) | u64::from(NT4_LOOKUP[idx]);
        }

        // Compress remaining nucleotides
        let start = full_chunks * chunk_size;
        for i in 0..remainder {
            encoded = (encoded << 2) | u64::from(NT1_LOOKUP[(bytes[start + i] - b'A') as usize]);
        }

        KMERandCount::new(encoded, count)
    }



    //////////// Encode a kmer + count + random value ............... cannot just use as_bytes + above?
    #[inline(always)]
    pub unsafe fn encode_str(&self, kmer: &str, count: u32) -> KMERandCount {
        let chunk_size: usize = 4;
        let kmer_size = self.kmer_size as usize;
        let full_chunks = kmer_size / chunk_size;
        let remainder = kmer_size % chunk_size;

        let bytes = kmer.as_bytes();
        let mut encoded = 0;
        let ptr = bytes.as_ptr();

        // Process chunks of 4 nucleotides
        for i in 0..full_chunks {
            let chunk_ptr = ptr.add(i * chunk_size);
            let idx = unsafe {
                (*chunk_ptr.offset(0) - b'A') as usize
                    + ((*chunk_ptr.offset(1) - b'A') as usize * NT4_DIMSIZE)
                    + ((*chunk_ptr.offset(2) - b'A') as usize * NT4_DIMSIZE * NT4_DIMSIZE)
                    + ((*chunk_ptr.offset(3) - b'A') as usize * NT4_DIMSIZE * NT4_DIMSIZE * NT4_DIMSIZE)
            };
            encoded = (encoded << 8) | u64::from(NT4_LOOKUP[idx]);
        }

        // Handle remaining nucleotides
        let start = full_chunks * chunk_size;
        for i in 0..remainder {
            encoded = (encoded << 2) | u64::from(NT1_LOOKUP[(bytes[start + i] - b'A') as usize]);
        }

        KMERandCount::new(encoded, count)
    }

    #[inline(always)]
    pub unsafe fn decode(&self, encoded: &KMERandCount) -> String {
        let mut sequence = Vec::with_capacity(self.kmer_size);
        let mut temp = encoded.kmer; 
        for _ in 0..self.kmer_size {
            let nuc = (temp & 0b11) as usize;
            sequence.push(NT_REVERSE[nuc]);
            temp >>= 2;
        }
        sequence.reverse();
        String::from_utf8_unchecked(sequence)
    }


    #[inline(always)]
    pub fn hash_for_kmer(kmer: u64) -> u32 {
        //Use a fast hash function https://docs.rs/fasthash/latest/fasthash/sea/index.html
        let mut hasher=fasthash::sea::Hasher64::new();
        hasher.write_u64(kmer);
        let f= hasher.finish();

        let hash = f ^ (f>>32); //fit hash in u32
        hash as u32
    }


    #[inline(always)]
    pub fn g_hash_for_kmer(kmer: u64) -> u32 {

        //https://docs.rs/fasthash/latest/fasthash/murmur3/index.html    note!!! will get different results depending on architecture. this can be a pain!
        let mut hasher = fasthash::murmur3::Hasher32::new();
        hasher.write_u64(kmer);
        let f= hasher.finish();

        let hash = f ^ (f>>32); //fit hash in u32
        hash as u32
    }


    /* 
    // Return +1 for even numbers, and -1 for odd numbers
    #[inline(always)]
    pub fn plusmin_one_hash_for_kmer(kmer: u32) -> i32 {
        1 - ((kmer & 1) << 1) as i32
    }
    */


    pub fn plusmin_one_hash_for_kmer(kmer: u32) -> i32 {
        1 - ((kmer & 1) << 1) as i32
    }


}






#[derive(Clone, Copy)]
pub struct KMERandCount {
    pub kmer: u64,
    pub hash: u32,
    pub count: u32,
}
impl KMERandCount {
    pub fn new(
        kmer: u64,
        count: u32
   ) -> KMERandCount {
       Self {
           kmer: kmer,
           hash: KMERCodec::hash_for_kmer(kmer), 
           count: count
       }
   }
}

impl PartialOrd for KMERandCount {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.hash.cmp(&other.hash))
    }
}

impl Ord for KMERandCount {
    fn cmp(&self, other: &Self) -> Ordering {
        self.hash.cmp(&other.hash)
    }
}

impl PartialEq for KMERandCount {
    fn eq(&self, other: &Self) -> bool {
        self.hash == other.hash
    }
}

impl Eq for KMERandCount { }