diff options
| author | tslil clingman <tslil@posteo.de> | 2022-09-28 00:12:11 +0200 |
|---|---|---|
| committer | tslil clingman <tslil@posteo.de> | 2022-09-28 00:20:29 +0200 |
| commit | 785fa20b4e4a5372857f6f033b9d478685460cea (patch) | |
| tree | df5ef4b2a3ead571a7de84c1bb2395917b5f4fee /srchr/src/corpus.rs | |
| parent | c4e7a97fd7efecdb85f362f04434ea8470a1927f (diff) | |
Throwing in kicad stuff
Diffstat (limited to 'srchr/src/corpus.rs')
| -rw-r--r-- | srchr/src/corpus.rs | 204 |
1 files changed, 204 insertions, 0 deletions
diff --git a/srchr/src/corpus.rs b/srchr/src/corpus.rs new file mode 100644 index 0000000..8e4b437 --- /dev/null +++ b/srchr/src/corpus.rs @@ -0,0 +1,204 @@ +use std::fmt; +use std::fs; + +use crate::layout::*; + +const fn build_lookup_table() -> [usize; 128] { + let mut result = [0; 128]; + let mut i = 0; + + while i < NUM_KEYS { + result[KEY_CHARS[i] as usize] = i; + i += 1; + } + return result; +} +const CHAR_TO_INDEX: [usize; 128] = build_lookup_table(); + +const NUM_BIGRAMS: usize = NUM_KEYS * NUM_KEYS; +pub struct Corpus { + bigram_count: [u32; NUM_BIGRAMS], + character_count: [u32; 256], + total_count: u32, + index_threshold: u32, +} + +fn pair_to_index(x: char, y: char) -> usize { + let xi = CHAR_TO_INDEX[x as usize]; + let yi = CHAR_TO_INDEX[y as usize]; + xi + NUM_KEYS * yi +} + +impl Corpus { + pub fn get_character_count(&self, c: char) -> u32 { + self.character_count[c as usize] + } + + pub fn load(path: &str) -> Result<Corpus, std::io::Error> { + let contents = fs::read_to_string(path)?; + + let mut bigram_count = [0; NUM_BIGRAMS]; + let mut character_count = [0; 256]; + let mut total_count = 0; + + let mut last_char = None; + for c in contents.chars() { + if let Some(c) = canonicalise(c) { + if let Some(lc) = last_char { + // We don't count these anyway + if lc != c { + bigram_count[pair_to_index(c, lc)] += 1; + bigram_count[pair_to_index(lc, c)] += 1; + } + } + last_char = Some(c); + character_count[c as usize] += 1; + total_count += 1; + } else { + last_char = None; + } + } + + let index_threshold = (total_count as f32 * 0.15) as u32; + + return Ok(Corpus { + bigram_count, + character_count, + total_count, + index_threshold, + }); + } + + pub fn prelayout_fitness(&self, layout: &Prelayout) -> u32 { + let mut score: u32 = 0; + + for i in 0..6 { + let keys = layout.get_standard_column(i); + let k1 = keys[0]; + let k2 = keys[1]; + let k3 = keys[2]; + + score += self.bigram_count[pair_to_index(k1, k2)] + + self.bigram_count[pair_to_index(k1, k3)] + + self.bigram_count[pair_to_index(k2, k3)]; + } + + for i in 0..2 { + let keys = layout.get_index_column(i); + let k1 = keys[0]; + let k2 = keys[1]; + let k3 = keys[2]; + let k4 = keys[3]; + let k5 = keys[4]; + let k6 = keys[5]; + + // We want index usage! + let index_count: u32 = keys.iter().map(|&k| self.character_count[k as usize]).sum(); + score += (self.index_threshold as i64 - index_count as i64).abs() as u32 / 64; + + score += self.bigram_count[pair_to_index(k1, k2)] + + self.bigram_count[pair_to_index(k1, k3)] + + self.bigram_count[pair_to_index(k2, k3)] + + self.bigram_count[pair_to_index(k4, k5)] + + self.bigram_count[pair_to_index(k4, k6)] + + self.bigram_count[pair_to_index(k5, k6)] + + self.bigram_count[pair_to_index(k1, k4)] + + self.bigram_count[pair_to_index(k1, k5)] + + self.bigram_count[pair_to_index(k1, k6)] + + self.bigram_count[pair_to_index(k2, k4)] + + self.bigram_count[pair_to_index(k2, k5)] + + self.bigram_count[pair_to_index(k2, k6)] + + self.bigram_count[pair_to_index(k3, k4)] + + self.bigram_count[pair_to_index(k3, k5)] + + self.bigram_count[pair_to_index(k3, k6)]; + } + + score + } + + pub fn evaluate_layout(&self, layout: &Layout) -> Evaluation { + let mut keypress: [u32; NUM_KEYS] = [0; NUM_KEYS]; + let mut sfb: [u32; 8] = [0; 8]; + + for (i, &c) in self.character_count.iter().enumerate() { + if c > 0 { + keypress[layout.get_index((i as u8) as char)] = c; + } + } + + // TODO: We make assumptions about NUM_KEYS here + for i in 0..8 { + let ind = if i < 4 { i } else { i + 2 }; + + let k1 = layout.get_key(ind + 10 * 0); + let k2 = layout.get_key(ind + 10 * 1); + let k3 = layout.get_key(ind + 10 * 2); + + sfb[i] = self.bigram_count[pair_to_index(k1, k2)] + + self.bigram_count[pair_to_index(k1, k3)] + + self.bigram_count[pair_to_index(k2, k3)]; + + if i == 3 { + let k4 = layout.get_key(4 + 10 * 0); + let k5 = layout.get_key(4 + 10 * 1); + let k6 = layout.get_key(4 + 10 * 2); + sfb[i] += self.bigram_count[pair_to_index(k4, k5)] + + self.bigram_count[pair_to_index(k4, k6)] + + self.bigram_count[pair_to_index(k5, k6)] + + self.bigram_count[pair_to_index(k1, k4)] + + self.bigram_count[pair_to_index(k1, k5)] + + self.bigram_count[pair_to_index(k1, k6)] + + self.bigram_count[pair_to_index(k2, k4)] + + self.bigram_count[pair_to_index(k2, k5)] + + self.bigram_count[pair_to_index(k2, k6)] + + self.bigram_count[pair_to_index(k3, k4)] + + self.bigram_count[pair_to_index(k3, k5)] + + self.bigram_count[pair_to_index(k3, k6)]; + } else if i == 4 { + let k4 = layout.get_key(5 + 10 * 0); + let k5 = layout.get_key(5 + 10 * 1); + let k6 = layout.get_key(5 + 10 * 2); + sfb[i] += self.bigram_count[pair_to_index(k4, k5)] + + self.bigram_count[pair_to_index(k5, k6)] + + self.bigram_count[pair_to_index(k4, k6)] + + self.bigram_count[pair_to_index(k1, k4)] + + self.bigram_count[pair_to_index(k1, k5)] + + self.bigram_count[pair_to_index(k1, k6)] + + self.bigram_count[pair_to_index(k2, k4)] + + self.bigram_count[pair_to_index(k2, k5)] + + self.bigram_count[pair_to_index(k2, k6)] + + self.bigram_count[pair_to_index(k3, k4)] + + self.bigram_count[pair_to_index(k3, k5)] + + self.bigram_count[pair_to_index(k3, k6)]; + } + } + + return Evaluation::new(keypress, self.total_count, sfb); + } +} + +fn dump_bigrams(corpus: &Corpus) -> Vec<(String, u32)> { + let mut result: Vec<(String, u32)> = Vec::new(); + + for (i, &x) in KEY_CHARS.iter().enumerate() { + for &y in &KEY_CHARS[i..] { + let mut pair = String::from(x); + pair.push(y); + result.push((pair, corpus.bigram_count[pair_to_index(x, y)])); + } + } + + result.sort_by(|(_, c1), (_, c2)| c1.cmp(c2).reverse()); + + return result; +} + +impl fmt::Display for Corpus { + fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str(&format!( + "Counted {} characters, top 5 bigrams {:?}", + self.total_count, + &dump_bigrams(self)[0..5] + )) + } +} |
