Skip to main content

sioul_learn/
diagnose.rs

1// SPDX-License-Identifier: GPL-3.0-or-later
2// Copyright © 2026 Aurélien Pierre
3
4//! What the header features say of each account's mail (`sioul spam
5//! features`): each feature's mean over the corpus, by account and label,
6//! and how often it is known. A feature whose mean differs from one account
7//! to another more than from ham to spam tells the classifier where mail was
8//! collected rather than what it is: it learns each account's share of spam
9//! instead of spam. Optionally, only the mail of one sender's domain: what
10//! the features carry of a bank's or a provider's real messages. Numbers
11//! only, never a message.
12
13use crate::labels::{self, Copy, Label};
14use crate::spamcore::{N, NAMES};
15use crate::{Dirs, LearnError, corpus, train};
16use serde::Serialize;
17use std::collections::{BTreeMap, HashMap};
18
19/// One account's mail of one label: how many, each feature's mean over
20/// those where it is known, and how many those are.
21#[derive(Debug, Clone, PartialEq, Serialize)]
22pub struct Group {
23    pub account: String,
24    pub label: Label,
25    pub messages: u64,
26    /// By feature, in `NAMES` order; none when no message knows it.
27    pub means: Vec<Option<f64>>,
28    pub known: Vec<u64>,
29}
30
31/// The feature names, in the order of `Group::means`.
32pub fn names() -> &'static [&'static str] {
33    &NAMES
34}
35
36/// Each account's and label's feature means, over the copies a training
37/// learns from; only the mail of `domain` (the sender's registrable domain,
38/// or one under it) when given.
39pub fn feature_means(dirs: &Dirs, trusted: &BTreeMap<String, Vec<String>>, domain: Option<&str>) -> Result<Vec<Group>, LearnError> {
40    let mut copies = Vec::new();
41    corpus::read_all(dirs, |record| copies.push(Copy::of(&record)))?;
42    let (labeled, _) = labels::decide(copies, &labels::read_log(dirs), &labels::read_moved(dirs), &labels::read_flagged(dirs));
43    let chosen: HashMap<corpus::Place, Label> = labeled.into_iter().map(|l| (l.place, l.label)).collect();
44    let domain = domain.map(|d| d.trim().trim_start_matches('@').to_ascii_lowercase());
45    // By account and label (spam: true), in order.
46    let mut sums: BTreeMap<(String, bool), (u64, [f64; N], [u64; N])> = BTreeMap::new();
47    corpus::read_all(dirs, |record| {
48        let Some(&label) = chosen.get(&record.place()) else { return };
49        if let Some(domain) = &domain {
50            let from = labels::from_address(&record.header_bytes()).and_then(|a| a.rsplit_once('@').map(|(_, d)| d.to_string())).unwrap_or_default();
51            if !(from == *domain || from.ends_with(&format!(".{domain}"))) {
52                return;
53            }
54        }
55        // The header features only: the words read do not count.
56        let (_, features) = train::read_message(&record, trusted, &crate::spamcore::Lexicon::builtin());
57        let group = sums.entry((record.account.clone(), label == Label::Spam)).or_insert((0, [0.0; N], [0; N]));
58        group.0 += 1;
59        for (h, x) in features.iter().enumerate() {
60            if !x.is_nan() {
61                group.1[h] += f64::from(*x);
62                group.2[h] += 1;
63            }
64        }
65    })?;
66    Ok(sums
67        .into_iter()
68        .map(|((account, spam), (messages, sum, known))| Group {
69            account,
70            label: if spam { Label::Spam } else { Label::Ham },
71            messages,
72            means: sum.iter().zip(&known).map(|(s, &k)| (k > 0).then(|| s / k as f64)).collect(),
73            known: known.to_vec(),
74        })
75        .collect())
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81    use crate::synthetic;
82
83    #[test]
84    fn means_by_account_and_label() {
85        let root = std::env::temp_dir().join(format!("sioul-learn-diagnose-{}", std::process::id()));
86        let _ = std::fs::remove_dir_all(&root);
87        let dirs = Dirs::under(&root);
88        let mail = synthetic::mailbox(2, 60, 40);
89        let records: Vec<_> = mail.iter().enumerate().map(|(i, m)| synthetic::record_of(m, 1, i as u32 + 1)).collect();
90        corpus::store(&dirs, &records).unwrap();
91        let groups = feature_means(&dirs, &BTreeMap::new(), None).unwrap();
92        assert_eq!(groups.iter().map(|g| g.messages).sum::<u64>(), 100);
93        assert!(groups.iter().all(|g| g.means.len() == N && g.known.len() == N));
94        let list = NAMES.iter().position(|n| *n == "list_unsubscribe").unwrap();
95        assert!(groups.iter().all(|g| g.means[list].is_some_and(|m| (0.0..=1.0).contains(&m))), "{groups:?}");
96        // One sender's domain alone: the invented spam's.
97        let lottery = feature_means(&dirs, &BTreeMap::new(), Some("lottery.test")).unwrap();
98        assert!(!lottery.is_empty() && lottery.iter().all(|g| g.label == Label::Spam), "{lottery:?}");
99        let _ = std::fs::remove_dir_all(root);
100    }
101}