Skip to main content

ferritin_core/info/
atom37.rs

1//! Canonical 37-slot heavy-atom representation and atom-name lookups.
2//!
3//! The atom37 scheme is the standard representation used by ESMFold, AlphaFold2,
4//! and related models: every amino acid is represented with up to 37 possible
5//! heavy-atom positions, each at a fixed slot index defined by [`ATOM37_NAMES`].
6//!
7//! # Key items
8//! - [`NUM_ATOM37`] / [`NUM_RESIDUES`] — dimension constants
9//! - [`ATOM37_NAMES`] — canonical name→slot mapping (array)
10//! - [`atom_order`] — `HashMap` version of the same mapping
11//! - [`atom37_index`] — slot lookup by atom name
12//! - [`AAAtom`] — typed enum mirroring the slot indices
13
14use std::collections::HashMap;
15use std::sync::OnceLock;
16use strum::{Display, EnumIter, EnumString};
17
18// ── Dimension constants ───────────────────────────────────────────────────────
19
20pub const NUM_ATOM37: usize = 37;
21pub const NUM_RESIDUES: usize = 21; // 20 standard AAs + UNK
22
23// ── Atom37 slot names ─────────────────────────────────────────────────────────
24
25/// Canonical atom37 names in slot order.
26///
27/// The index of each name is its atom37 slot number, matching [`AAAtom`].
28#[rustfmt::skip]
29pub const ATOM37_NAMES: [&str; NUM_ATOM37] = [
30    "N",   "CA",  "C",   "CB",  "O",
31    "CG",  "CG1", "CG2", "OG",  "OG1",
32    "SG",  "CD",  "CD1", "CD2", "ND1",
33    "ND2", "OD1", "OD2", "SD",  "CE",
34    "CE1", "CE2", "CE3", "NE",  "NE1",
35    "NE2", "OE1", "OE2", "CH2", "NH1",
36    "NH2", "OH",  "CZ",  "CZ2", "CZ3",
37    "NZ",  "OXT",
38];
39
40static ATOM_ORDER: OnceLock<HashMap<&'static str, usize>> = OnceLock::new();
41
42/// Map from atom name (e.g. `"CA"`) to its atom37 slot index.
43pub fn atom_order() -> &'static HashMap<&'static str, usize> {
44    ATOM_ORDER.get_or_init(|| {
45        ATOM37_NAMES
46            .iter()
47            .enumerate()
48            .map(|(i, &n)| (n, i))
49            .collect()
50    })
51}
52
53/// Returns the atom37 slot index for a named atom, or `None` if unknown.
54pub fn atom37_index(name: &str) -> Option<usize> {
55    ATOM37_NAMES.iter().position(|&n| n == name)
56}
57
58// ── AAAtom enum ───────────────────────────────────────────────────────────────
59
60/// Typed atom37 slots; discriminants match [`ATOM37_NAMES`] indices.
61///
62/// `Unknown = -1` is used as a sentinel for "atom not present in this residue"
63/// in atom14 tables. Callers working with `usize` slot indices should use
64/// [`atom37_index`] instead.
65#[rustfmt::skip]
66#[derive(Debug, Clone, Copy, PartialEq, Display, EnumString, EnumIter)]
67pub enum AAAtom {
68    N = 0,    CA = 1,   C = 2,    CB = 3,   O = 4,
69    CG = 5,   CG1 = 6,  CG2 = 7,  OG = 8,   OG1 = 9,
70    SG = 10,  CD = 11,  CD1 = 12, CD2 = 13, ND1 = 14,
71    ND2 = 15, OD1 = 16, OD2 = 17, SD = 18,  CE = 19,
72    CE1 = 20, CE2 = 21, CE3 = 22, NE = 23,  NE1 = 24,
73    NE2 = 25, OE1 = 26, OE2 = 27, CH2 = 28, NH1 = 29,
74    NH2 = 30, OH = 31,  CZ = 32,  CZ2 = 33, CZ3 = 34,
75    NZ = 35,  OXT = 36,
76    Unknown = -1,
77}
78
79impl AAAtom {
80    pub fn to_index(&self) -> usize {
81        *self as usize
82    }
83}
84
85// ── Tests ─────────────────────────────────────────────────────────────────────
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn test_atom_order_completeness() {
93        let order = atom_order();
94        assert_eq!(order.len(), NUM_ATOM37);
95        assert_eq!(order["N"], 0);
96        assert_eq!(order["CA"], 1);
97        assert_eq!(order["C"], 2);
98        assert_eq!(order["CB"], 3);
99        assert_eq!(order["O"], 4);
100    }
101
102    #[test]
103    fn test_atom37_index_roundtrip() {
104        for (i, &name) in ATOM37_NAMES.iter().enumerate() {
105            assert_eq!(atom37_index(name), Some(i));
106        }
107        assert_eq!(atom37_index("ZZZ"), None);
108    }
109
110    #[test]
111    fn test_aaatom_discriminants_match_names() {
112        assert_eq!(AAAtom::N as i32, 0);
113        assert_eq!(AAAtom::CA as i32, 1);
114        assert_eq!(AAAtom::OXT as i32, 36);
115        assert_eq!(AAAtom::Unknown as i32, -1);
116    }
117}