ferritin_plms/utils/
residue_constants.rs

1//! Residue chemistry constants for structure prediction.
2//!
3//! Centralises the biochemical lookup tables needed by the ESMFold2 forward
4//! pass and ESMC structure utilities:
5//!
6//! - [`ATOM37_NAMES`] / [`atom_order`] — canonical 37-slot heavy-atom ordering
7//! - [`chi_angles_atoms`] — 4-atom tuples defining each sidechain dihedral
8//! - [`atom14_to_atom37_for_residue`] — per-residue atom14 → atom37 slot mapping
9//! - [`vdw_radius`] — van der Waals radii for common protein elements
10//!
11//! The atom37 ordering matches the [`crate::featurize::utilities::AAAtom`] enum.
12
13use std::collections::HashMap;
14
15// ── Atom37 ordering ────────────────────────────────────────────────────────
16
17pub const NUM_ATOM37: usize = 37;
18pub const NUM_RESIDUES: usize = 21; // 20 standard AAs + UNK
19
20/// Canonical atom37 names in slot order.
21///
22/// The index of each name is its atom37 slot number, matching the
23/// `AAAtom` enum in `featurize::utilities`.
24#[rustfmt::skip]
25pub const ATOM37_NAMES: [&str; NUM_ATOM37] = [
26    "N",   "CA",  "C",   "CB",  "O",
27    "CG",  "CG1", "CG2", "OG",  "OG1",
28    "SG",  "CD",  "CD1", "CD2", "ND1",
29    "ND2", "OD1", "OD2", "SD",  "CE",
30    "CE1", "CE2", "CE3", "NE",  "NE1",
31    "NE2", "OE1", "OE2", "CH2", "NH1",
32    "NH2", "OH",  "CZ",  "CZ2", "CZ3",
33    "NZ",  "OXT",
34];
35
36/// Map from atom name (e.g. `"CA"`) to its atom37 slot index.
37pub fn atom_order() -> HashMap<String, usize> {
38    ATOM37_NAMES
39        .iter()
40        .enumerate()
41        .map(|(i, &name)| (name.to_string(), i))
42        .collect()
43}
44
45/// Returns the atom37 slot index for a named atom, or `None` if unknown.
46pub fn atom37_index(name: &str) -> Option<usize> {
47    ATOM37_NAMES.iter().position(|&n| n == name)
48}
49
50// ── Chi angle definitions ──────────────────────────────────────────────────
51
52/// Sidechain chi-angle atom tuples for a residue given by 3-letter code.
53///
54/// Returns a slice of `(a1, a2, a3, a4)` tuples, one per chi angle (chi1
55/// first). An empty slice means the residue has no rotatable sidechains.
56///
57/// Source: OpenFold / AlphaFold2 `residue_constants.py`.
58#[rustfmt::skip]
59pub fn chi_angles_atoms(res3: &str) -> &'static [(&'static str, &'static str, &'static str, &'static str)] {
60    match res3 {
61        "ALA" | "GLY" | "UNK" => &[],
62        "ARG" => &[
63            ("N","CA","CB","CG"), ("CA","CB","CG","CD"),
64            ("CB","CG","CD","NE"), ("CG","CD","NE","CZ"),
65        ],
66        "ASN" => &[
67            ("N","CA","CB","CG"), ("CA","CB","CG","OD1"),
68        ],
69        "ASP" => &[
70            ("N","CA","CB","CG"), ("CA","CB","CG","OD1"),
71        ],
72        "CYS" => &[
73            ("N","CA","CB","SG"),
74        ],
75        "GLN" => &[
76            ("N","CA","CB","CG"), ("CA","CB","CG","CD"), ("CB","CG","CD","OE1"),
77        ],
78        "GLU" => &[
79            ("N","CA","CB","CG"), ("CA","CB","CG","CD"), ("CB","CG","CD","OE1"),
80        ],
81        "HIS" => &[
82            ("N","CA","CB","CG"), ("CA","CB","CG","ND1"),
83        ],
84        "ILE" => &[
85            ("N","CA","CB","CG1"), ("CA","CB","CG1","CD1"),
86        ],
87        "LEU" => &[
88            ("N","CA","CB","CG"), ("CA","CB","CG","CD1"),
89        ],
90        "LYS" => &[
91            ("N","CA","CB","CG"), ("CA","CB","CG","CD"),
92            ("CB","CG","CD","CE"), ("CG","CD","CE","NZ"),
93        ],
94        "MET" => &[
95            ("N","CA","CB","CG"), ("CA","CB","CG","SD"), ("CB","CG","SD","CE"),
96        ],
97        "PHE" => &[
98            ("N","CA","CB","CG"), ("CA","CB","CG","CD1"),
99        ],
100        "PRO" => &[
101            ("N","CA","CB","CG"), ("CA","CB","CG","CD"),
102        ],
103        "SER" => &[
104            ("N","CA","CB","OG"),
105        ],
106        "THR" => &[
107            ("N","CA","CB","OG1"),
108        ],
109        "TRP" => &[
110            ("N","CA","CB","CG"), ("CA","CB","CG","CD1"),
111        ],
112        "TYR" => &[
113            ("N","CA","CB","CG"), ("CA","CB","CG","CD1"),
114        ],
115        "VAL" => &[
116            ("N","CA","CB","CG1"),
117        ],
118        _ => &[],
119    }
120}
121
122// ── atom14 → atom37 mapping ────────────────────────────────────────────────
123
124/// For a residue given by its 3-letter code, return the atom37 slot index for
125/// each of the 14 atom14 positions, or `None` if that slot is unoccupied.
126///
127/// The atom14 slot ordering matches the `Residue::atoms14()` method in
128/// `featurize::utilities`: [N, CA, C, O, CB, sidechain…].
129#[rustfmt::skip]
130pub fn atom14_to_atom37_for_residue(res3: &str) -> [Option<usize>; 14] {
131    let names: [Option<&str>; 14] = match res3 {
132        "ALA" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),None,None,None,None,None,None,None,None,None],
133        "CYS" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("SG"),None,None,None,None,None,None,None,None],
134        "ASP" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG"),Some("OD1"),Some("OD2"),None,None,None,None,None,None],
135        "GLU" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG"),Some("CD"),Some("OE1"),Some("OE2"),None,None,None,None,None],
136        "PHE" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG"),Some("CD1"),Some("CD2"),Some("CE1"),Some("CE2"),Some("CZ"),None,None,None],
137        "GLY" => [Some("N"),Some("CA"),Some("C"),Some("O"),None,None,None,None,None,None,None,None,None,None],
138        "HIS" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG"),Some("ND1"),Some("CD2"),Some("CE1"),Some("NE2"),None,None,None,None],
139        "ILE" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG1"),Some("CG2"),Some("CD1"),None,None,None,None,None,None],
140        "LYS" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG"),Some("CD"),Some("CE"),Some("NZ"),None,None,None,None,None],
141        "LEU" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG"),Some("CD1"),Some("CD2"),None,None,None,None,None,None],
142        "MET" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG"),Some("SD"),Some("CE"),None,None,None,None,None,None],
143        "ASN" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG"),Some("OD1"),Some("ND2"),None,None,None,None,None,None],
144        "PRO" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG"),Some("CD"),None,None,None,None,None,None,None],
145        "GLN" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG"),Some("CD"),Some("OE1"),Some("NE2"),None,None,None,None,None],
146        "ARG" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG"),Some("CD"),Some("NE"),Some("CZ"),Some("NH1"),Some("NH2"),None,None,None],
147        "SER" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("OG"),None,None,None,None,None,None,None,None],
148        "THR" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("OG1"),Some("CG2"),None,None,None,None,None,None,None],
149        "VAL" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG1"),Some("CG2"),None,None,None,None,None,None,None],
150        "TRP" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG"),Some("CD1"),Some("CD2"),Some("CE2"),Some("CE3"),Some("NE1"),Some("CZ2"),Some("CZ3"),Some("CH2")],
151        "TYR" => [Some("N"),Some("CA"),Some("C"),Some("O"),Some("CB"),Some("CG"),Some("CD1"),Some("CD2"),Some("CE1"),Some("CE2"),Some("CZ"),Some("OH"),None,None],
152        _     => [None; 14],
153    };
154
155    let mut out = [None; 14];
156    for (i, name_opt) in names.iter().enumerate() {
157        out[i] = name_opt.and_then(atom37_index);
158    }
159    out
160}
161
162// ── Van der Waals radii ────────────────────────────────────────────────────
163
164/// Van der Waals radius (Å) for elements commonly found in proteins.
165///
166/// Source: Bondi (1964) / standard crystallographic values.
167pub fn vdw_radius(element: &str) -> f32 {
168    match element {
169        "H" => 1.20,
170        "C" => 1.70,
171        "N" => 1.55,
172        "O" => 1.52,
173        "S" => 1.80,
174        "P" => 1.80,
175        "F" => 1.47,
176        "CL" | "Cl" => 1.75,
177        "BR" | "Br" => 1.85,
178        "I" => 1.98,
179        "SE" | "Se" => 1.90,
180        _ => 1.70,
181    }
182}
183
184// ── Tests ──────────────────────────────────────────────────────────────────
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189
190    #[test]
191    fn test_atom_order_completeness() {
192        let order = atom_order();
193        assert_eq!(order.len(), NUM_ATOM37);
194        // Backbone atoms must be present at well-known indices
195        assert_eq!(order["N"], 0);
196        assert_eq!(order["CA"], 1);
197        assert_eq!(order["C"], 2);
198        assert_eq!(order["CB"], 3);
199        assert_eq!(order["O"], 4);
200    }
201
202    #[test]
203    fn test_atom37_index_roundtrip() {
204        for (i, &name) in ATOM37_NAMES.iter().enumerate() {
205            assert_eq!(atom37_index(name), Some(i));
206        }
207        assert_eq!(atom37_index("ZZZ"), None);
208    }
209
210    #[test]
211    fn test_chi_gly_ala_empty() {
212        assert!(chi_angles_atoms("GLY").is_empty());
213        assert!(chi_angles_atoms("ALA").is_empty());
214    }
215
216    #[test]
217    fn test_chi_arg_four_angles() {
218        assert_eq!(chi_angles_atoms("ARG").len(), 4);
219        assert_eq!(chi_angles_atoms("ARG")[0], ("N", "CA", "CB", "CG"));
220    }
221
222    #[test]
223    fn test_chi_val_one_angle() {
224        let chis = chi_angles_atoms("VAL");
225        assert_eq!(chis.len(), 1);
226        assert_eq!(chis[0], ("N", "CA", "CB", "CG1"));
227    }
228
229    #[test]
230    fn test_atom14_to_atom37_gly() {
231        let mapping = atom14_to_atom37_for_residue("GLY");
232        // GLY: N=0, CA=1, C=2, O=4, rest None
233        assert_eq!(mapping[0], Some(0)); // N
234        assert_eq!(mapping[1], Some(1)); // CA
235        assert_eq!(mapping[2], Some(2)); // C
236        assert_eq!(mapping[3], Some(4)); // O
237        assert_eq!(mapping[4], None); // no CB for GLY
238        for i in 4..14 {
239            assert_eq!(mapping[i], None, "GLY slot {i} should be None");
240        }
241    }
242
243    #[test]
244    fn test_atom14_to_atom37_ala() {
245        let mapping = atom14_to_atom37_for_residue("ALA");
246        assert_eq!(mapping[4], Some(3)); // CB → atom37 slot 3
247        assert_eq!(mapping[5], None);
248    }
249
250    #[test]
251    fn test_atom14_to_atom37_trp_full() {
252        // TRP fills all 14 slots
253        let mapping = atom14_to_atom37_for_residue("TRP");
254        assert!(
255            mapping.iter().all(|m| m.is_some()),
256            "TRP should fill all 14 slots"
257        );
258    }
259
260    #[test]
261    fn test_atom14_to_atom37_unk() {
262        let mapping = atom14_to_atom37_for_residue("UNK");
263        assert!(mapping.iter().all(|m| m.is_none()));
264    }
265
266    #[test]
267    fn test_vdw_radius_known_elements() {
268        assert!((vdw_radius("C") - 1.70).abs() < 1e-6);
269        assert!((vdw_radius("N") - 1.55).abs() < 1e-6);
270        assert!((vdw_radius("O") - 1.52).abs() < 1e-6);
271        assert!((vdw_radius("S") - 1.80).abs() < 1e-6);
272    }
273
274    #[test]
275    fn test_vdw_radius_unknown_defaults_to_carbon() {
276        assert!((vdw_radius("X") - 1.70).abs() < 1e-6);
277    }
278}
ferritin_plms/utils/residue_constants.rs

ferritin_plms/utils/
residue_constants.rs
