uno_core/
rt_cache.rs

1//! On-disk cache of abstract RT predictions.
2//!
3//! The RT model's *abstract* RT is a pure function of `(model, peptide
4//! sequence)` — it does not depend on the run being searched (the run-specific
5//! part is the on-data calibration applied downstream). So predicting ~1.2M
6//! unique peptides (~18s) is repeated in full on every search over the same
7//! FASTA/digest/model. This module memoizes those predictions to a file keyed by
8//! the model's content fingerprint, so repeat runs only predict sequences not
9//! seen before.
10//!
11//! Cache values are the model's exact `f64` output, written with Rust's
12//! round-trippable float formatting, so a cache hit yields a bit-identical
13//! prediction to recomputing — search output is unchanged whether cached or not.
14//!
15//! Location: `$UNO_RT_CACHE_DIR` if set, else `$XDG_CACHE_HOME/uno/rt`, else
16//! `~/.cache/uno/rt`. Set `UNO_RT_CACHE=0` to disable (always recompute, never
17//! read or write).
18
19use std::collections::HashMap;
20use std::io::{BufRead, BufReader, BufWriter, Write};
21use std::path::PathBuf;
22
23use crate::rt::RtModel;
24
25/// Predict abstract RT for `seqs`, reusing an on-disk cache keyed by the model
26/// fingerprint. Only sequences missing from the cache are run through the model;
27/// newly predicted `(seq, rt)` pairs are appended to the cache file. The returned
28/// vec is aligned to `seqs` and bit-identical to `model.predict_many(seqs)`.
29pub fn predict_many_cached(model: &RtModel, seqs: &[&str]) -> Vec<f64> {
30    let cache_path = match cache_file(model.fingerprint()) {
31        Some(p) => p,
32        // Caching disabled or no home dir: fall straight through to the model.
33        None => return model.predict_many(seqs),
34    };
35
36    BATCH_CACHE.with(|state| {
37        let mut state = state.borrow_mut();
38        if let Some(models) = state.as_mut() {
39            let cached = models
40                .entry(model.fingerprint())
41                .or_insert_with(|| load_cache(&cache_path));
42            predict_into(model, seqs, &cache_path, cached)
43        } else {
44            predict_into(model, seqs, &cache_path, &mut load_cache(&cache_path))
45        }
46    })
47}
48
49thread_local! {
50    static BATCH_CACHE: std::cell::RefCell<Option<HashMap<u64,HashMap<String,f64>>>> = const { std::cell::RefCell::new(None) };
51}
52
53/// Retain prediction maps across sequential per-run searches on this thread.
54/// The guard restores the previous scope even on early return or panic.
55pub fn with_memory_cache<T>(work: impl FnOnce() -> T) -> T {
56    struct Restore(Option<HashMap<u64, HashMap<String, f64>>>);
57    impl Drop for Restore {
58        fn drop(&mut self) {
59            BATCH_CACHE.with(|s| *s.borrow_mut() = self.0.take());
60        }
61    }
62    let previous = BATCH_CACHE.with(|s| s.replace(Some(HashMap::new())));
63    let _guard = Restore(previous);
64    work()
65}
66
67fn predict_into(
68    model: &RtModel,
69    seqs: &[&str],
70    path: &std::path::Path,
71    cached: &mut HashMap<String, f64>,
72) -> Vec<f64> {
73    let mut miss_seqs = Vec::new();
74    let mut seen = std::collections::HashSet::new();
75    for &s in seqs {
76        if !cached.contains_key(s) && seen.insert(s) {
77            miss_seqs.push(s);
78        }
79    }
80    log::info!(
81        "RT cache {}: {} unique misses of {} requested ({} rows cached)",
82        path.display(),
83        miss_seqs.len(),
84        seqs.len(),
85        cached.len()
86    );
87    if !miss_seqs.is_empty() {
88        let predictions = model.predict_many(&miss_seqs);
89        if let Err(e) = append_cache(path, &miss_seqs, &predictions) {
90            log::warn!("RT cache write {}: {}", path.display(), e);
91        }
92        for (s, p) in miss_seqs.into_iter().zip(predictions) {
93            cached.insert(s.to_owned(), p);
94        }
95    }
96    seqs.iter().map(|s| cached[*s]).collect()
97}
98
99/// The cache file for a model fingerprint, or `None` if caching is disabled or no
100/// cache directory can be determined.
101fn cache_file(fingerprint: u64) -> Option<PathBuf> {
102    if std::env::var("UNO_RT_CACHE").as_deref() == Ok("0") {
103        return None;
104    }
105    let dir = cache_dir()?;
106    Some(dir.join(format!("{fingerprint:016x}.tsv")))
107}
108
109fn cache_dir() -> Option<PathBuf> {
110    if let Ok(d) = std::env::var("UNO_RT_CACHE_DIR") {
111        return Some(PathBuf::from(d));
112    }
113    if let Ok(d) = std::env::var("XDG_CACHE_HOME") {
114        if !d.is_empty() {
115            return Some(PathBuf::from(d).join("uno").join("rt"));
116        }
117    }
118    std::env::var("HOME")
119        .ok()
120        .map(|h| PathBuf::from(h).join(".cache").join("uno").join("rt"))
121}
122
123/// Load `seq<TAB>rt` rows into a map. A missing/corrupt file yields an empty map
124/// (the run simply recomputes and rewrites).
125fn load_cache(path: &PathBuf) -> HashMap<String, f64> {
126    let mut map = HashMap::new();
127    let f = match std::fs::File::open(path) {
128        Ok(f) => f,
129        Err(_) => return map,
130    };
131    for line in BufReader::new(f).lines().map_while(Result::ok) {
132        let mut it = line.split('\t');
133        if let (Some(seq), Some(rt)) = (it.next(), it.next()) {
134            if let Ok(rt) = rt.parse::<f64>() {
135                map.insert(seq.to_string(), rt);
136            }
137        }
138    }
139    map
140}
141
142/// Append newly predicted `(seq, rt)` pairs to the cache, creating the directory
143/// and file on first use. Append-only, so concurrent searches can both add
144/// entries without clobbering each other's rows.
145fn append_cache(path: &std::path::Path, seqs: &[&str], preds: &[f64]) -> std::io::Result<()> {
146    if let Some(parent) = path.parent() {
147        std::fs::create_dir_all(parent)?;
148    }
149    let f = std::fs::OpenOptions::new()
150        .create(true)
151        .append(true)
152        .open(path)?;
153    let mut w = BufWriter::new(f);
154    for (s, p) in seqs.iter().zip(preds) {
155        // Rust's default f64 formatting is the shortest round-trippable decimal,
156        // so parsing it back yields the identical bit pattern.
157        writeln!(w, "{s}\t{p}")?;
158    }
159    w.flush()
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165    #[test]
166    fn warm_predictions_are_bit_exact_and_append_only_misses() {
167        let dir = std::env::temp_dir().join(format!(
168            "uno-cache-test-{}-{}",
169            std::process::id(),
170            std::time::SystemTime::now()
171                .duration_since(std::time::UNIX_EPOCH)
172                .unwrap()
173                .as_nanos()
174        ));
175        let path = dir.join("model.tsv");
176        let model = RtModel::bundled().unwrap();
177        let seqs = ["PEPTIDEK", "AAAAAAAR", "PEPTIDEK"];
178        let expected = model.predict_many(&seqs);
179        let mut map = HashMap::new();
180        let first = predict_into(&model, &seqs, &path, &mut map);
181        let before = std::fs::read(&path).unwrap();
182        let second = predict_into(&model, &seqs, &path, &mut map);
183        assert_eq!(first, expected);
184        assert_eq!(second, expected);
185        assert_eq!(std::fs::read(&path).unwrap(), before);
186        let disk = load_cache(&path);
187        assert_eq!(disk.len(), 2);
188        assert_eq!(disk["PEPTIDEK"].to_bits(), expected[0].to_bits());
189        std::fs::remove_dir_all(dir).unwrap();
190    }
191    #[test]
192    fn batch_cache_scope_restores_previous_state() {
193        with_memory_cache(|| {
194            BATCH_CACHE.with(|s| {
195                s.borrow_mut().as_mut().unwrap().insert(42, HashMap::new());
196            });
197            with_memory_cache(|| {
198                assert!(BATCH_CACHE.with(|s| s.borrow().as_ref().unwrap().is_empty()))
199            });
200            assert!(BATCH_CACHE.with(|s| s.borrow().as_ref().unwrap().contains_key(&42)));
201        });
202        assert!(BATCH_CACHE.with(|s| s.borrow().is_none()));
203    }
204}
uno_core/rt_cache.rs

uno_core/
rt_cache.rs
