ferritin_plms/esm3/tokenization/
sasa.rs

1//! ESM3 SASA (solvent-accessible surface area) tokenization.
2//!
3//! Continuous SASA values (Å²) are discretized into 16 bins using
4//! the 15 boundary values from `esm/utils/constants/esm3.py`.
5//!
6//! Token layout:
7//!   0 = PAD, 1 = MASK, 2 = UNK
8//!   3..18 = bins 0..15   (total 19 tokens)
9
10use crate::esm3::utils::constants::SASA_DISCRETIZATION_BOUNDARIES;
11
12/// Discretize a slice of per-residue SASA values into token IDs.
13///
14/// Each value is placed in the leftmost bin where `value < boundary`.
15/// Values above all boundaries fall into the highest bin.
16pub fn tokenize_sasa(sasa_values: &[f32]) -> Vec<u32> {
17    sasa_values
18        .iter()
19        .map(|&v| {
20            let bin = SASA_DISCRETIZATION_BOUNDARIES
21                .iter()
22                .position(|&b| v < b)
23                .unwrap_or(SASA_DISCRETIZATION_BOUNDARIES.len());
24            bin as u32 + 3 // offset by 3 for PAD, MASK, UNK
25        })
26        .collect()
27}
28
29#[cfg(test)]
30mod tests {
31    use super::*;
32
33    #[test]
34    fn test_tokenize_sasa_lowest_bin() {
35        // 0.0 < 0.8 → bin 0 → token 3
36        assert_eq!(tokenize_sasa(&[0.0]), vec![3]);
37    }
38
39    #[test]
40    fn test_tokenize_sasa_highest_bin() {
41        // 200.0 > all boundaries → bin 15 → token 18
42        assert_eq!(tokenize_sasa(&[200.0]), vec![18]);
43    }
44
45    #[test]
46    fn test_tokenize_sasa_midpoint() {
47        // 5.0 >= 0.8, 5.0 >= 4.0, 5.0 < 9.6 → bin 2 → token 5
48        assert_eq!(tokenize_sasa(&[5.0]), vec![5]);
49    }
50
51    #[test]
52    fn test_tokenize_sasa_batch() {
53        let ids = tokenize_sasa(&[0.0, 200.0]);
54        assert_eq!(ids, vec![3, 18]);
55    }
56}
ferritin_plms/esm3/tokenization/sasa.rs

ferritin_plms/esm3/tokenization/
sasa.rs
