ferritin_plms/featurize/
structure_features.rs

1//!  Protein->Tensor utilities useful for Machine Learning
2use super::utilities::{AAAtom, aa1to_int, aa3to1, int_to_aa1, get_nearest_neighbours};
3use crate::ligandmpnn::proteinfeatures::ProteinFeatures;
4use candle_core::{D, DType, Device, IndexOp, Result, Tensor};
5use ferritin_core::{AtomCollection, Model};
6use ferritin_core::info::elements::Element;
7use std::collections::HashSet;
8use strum::IntoEnumIterator;
9
10const LIGAND_CUTOFF_SCORE: f32 = 5.;
11
12// Helper Fns --------------------------------------
13fn is_heavy_atom(element: &Element) -> bool {
14    !matches!(element, Element::H | Element::He)
15}
16
17///. Trait defining Protein->Tensor utilities useful for Machine Learning
18pub trait StructureFeatures {
19    /// Convert amino acid sequence to numeric representation
20    fn decode_amino_acids(&self, device: &Device) -> Result<Tensor>;
21
22    /// Convert amino acid sequence to numeric representation
23    fn encode_amino_acids(&self, device: &Device) -> Result<Tensor>;
24
25    /// Convert amino acid sequence to numeric representation
26    fn create_cb(&self, device: &Device) -> Result<Tensor>;
27
28    /// Prepare for ProteinMPNN
29    fn featurize_lmpnn(&self, device: &Device) -> Result<ProteinFeatures>; // need more control over this featurization process
30
31    /// Get residue indices
32    fn get_res_index(&self) -> Vec<u32>;
33
34    /// Extract backbone atom coordinates (N, CA, C, O)
35    fn to_numeric_backbone_atoms(&self, device: &Device) -> Result<Tensor>;
36
37    /// Extract all atom coordinates in standard ordering
38    fn to_numeric_atom37(&self, device: &Device) -> Result<Tensor>;
39
40    /// Extract ligand atom coordinates and properties
41    fn to_numeric_ligand_atoms(&self, device: &Device) -> Result<(Tensor, Tensor, Tensor)>;
42}
43
44impl StructureFeatures for AtomCollection {
45    /// Decode amino acid integer indices back to one-letter codes as ASCII bytes.
46    ///
47    /// This is the inverse of `encode_amino_acids`. It iterates over the amino acid
48    /// residues in the structure, converts each three-letter residue name to a
49    /// one-letter code, then encodes it as an integer via `aa1to_int`, decodes it
50    /// back via `int_to_aa1`, and returns the ASCII byte values in a tensor of
51    /// shape `[1, n]` where `n` is the number of amino acid residues.
52    ///
53    /// Unknown residues map to the sentinel index 20, which decodes to `'X'` (ASCII 88).
54    fn decode_amino_acids(&self, device: &Device) -> Result<Tensor> {
55        let n = self.iter_residues_aminoacid().count();
56        let s: Vec<u8> = self
57            .iter_residues_aminoacid()
58            .map(|res| res.residue_name().to_string())
59            .map(|res| aa3to1(&res))
60            .map(|ch| aa1to_int(ch))
61            .map(|idx| int_to_aa1(idx) as u8)
62            .collect();
63        Ok(Tensor::from_iter(s.into_iter(), device)?.reshape((1, n))?)
64    }
65
66    /// Convert amino acid sequence to numeric representation
67    fn encode_amino_acids(&self, device: &Device) -> Result<Tensor> {
68        let n = self.iter_residues_aminoacid().count();
69        let s = self
70            .iter_residues_aminoacid()
71            .map(|res| res.residue_name().to_string())
72            .map(|res| aa3to1(&res))
73            .map(|res| aa1to_int(res));
74
75        Ok(Tensor::from_iter(s, device)?.reshape((1, n))?)
76    }
77
78    /// Calculate CB for each residue
79    fn create_cb(&self, device: &Device) -> Result<Tensor> {
80        let backbone = self.to_numeric_backbone_atoms(device)?.squeeze(0)?;
81
82        // Extract N, CA, C coordinates
83        let n = backbone.i((.., 0, ..))?;
84        let ca = backbone.i((.., 1, ..))?;
85        let c = backbone.i((.., 2, ..))?;
86
87        // Constants for CB calculation
88        let a_coeff = -0.58273431_f64;
89        let b_coeff = 0.56802827_f64;
90        let c_coeff = -0.54067466_f64;
91
92        // Calculate vectors
93        let b = (&ca - &n)?;
94        let c = (&c - &ca)?;
95
96        // Manual cross product components
97        // a_x = b_y * c_z - b_z * c_y
98        // a_y = b_z * c_x - b_x * c_z
99        // a_z = b_x * c_y - b_y * c_x
100        let b_x = b.i((.., 0))?;
101        let b_y = b.i((.., 1))?;
102        let b_z = b.i((.., 2))?;
103        let c_x = c.i((.., 0))?;
104        let c_y = c.i((.., 1))?;
105        let c_z = c.i((.., 2))?;
106
107        let a_x = ((&b_y * &c_z)? - (&b_z * &c_y)?)?;
108        let a_y = ((&b_z * &c_x)? - (&b_x * &c_z)?)?;
109        let a_z = ((&b_x * &c_y)? - (&b_y * &c_x)?)?;
110        let a = Tensor::stack(&[&a_x, &a_y, &a_z], D::Minus1)?;
111
112        // Final CB calculation: -0.58273431 * a + 0.56802827 * b - 0.54067466 * c + CA
113        let cb = ((&a * a_coeff)? + (&b * b_coeff)? + (&c * c_coeff)? + &ca)?;
114        let cb = cb.unsqueeze(0)?;
115        Ok(cb)
116    }
117
118    // Convert AtomCollection to ProteinFeatures
119    fn featurize_lmpnn(&self, device: &Device) -> Result<ProteinFeatures> {
120        let x_37 = self.to_numeric_atom37(device)?;
121        let x_37_m = Tensor::ones((x_37.dim(0)?, x_37.dim(1)?), DType::F32, device)?;
122        let (y, y_t, y_m) = self.to_numeric_ligand_atoms(device)?;
123        let _cb = self.create_cb(device);
124        let _chain_labels = self.get_resids(); //  <-- need to double-check shape. I think this is all-atom
125        let residue_ids = self.get_res_index();
126        let residue_length = residue_ids.len();
127        let r_idx = Tensor::from_iter(residue_ids, device)?.reshape((1, residue_length))?;
128        let chain_letters: Vec<String> = self
129            .iter_residues_aminoacid()
130            .map(|res| res.chain_id().to_string())
131            .collect();
132        let chain_list: Vec<String> = self
133            .iter_residues_aminoacid()
134            .map(|res| res.chain_id().to_string())
135            .collect::<HashSet<_>>()
136            .into_iter()
137            .collect();
138        // Numeric chain labels (optional)
139        let chain_labels: Option<Vec<f64>> = None; // Could populate if needed
140        let s = self.encode_amino_acids(device)?;
141        // coordinates of the backbone atoms
142        let indices = Tensor::from_slice(
143            &[0i64, 1i64, 2i64, 4i64], // index of N/CA/C/O as integers
144            (4,),
145            &device,
146        )?;
147        let x = x_37.index_select(&indices, 2)?;
148        Ok(ProteinFeatures {
149            s,
150            x,
151            x_mask: Some(x_37_m),
152            y,
153            y_t,
154            y_m: Some(y_m),
155            r_idx,
156            chain_labels,
157            chain_letters,
158            mask_c: None,
159            chain_list,
160        })
161    }
162    /// Get residue indices
163    fn get_res_index(&self) -> Vec<u32> {
164        self.iter_residues_aminoacid()
165            .map(|res| res.residue_id() as u32)
166            .collect()
167    }
168
169    /// create numeric Tensor of shape [1, <sequence-length>, 4, 3] where the 4 is N/CA/C/O
170    fn to_numeric_backbone_atoms(&self, device: &Device) -> Result<Tensor> {
171        let res_count = self.iter_residues_aminoacid().count();
172        let mut backbone_data = Vec::with_capacity(res_count * 4 * 3);
173
174        for residue in self.iter_residues_aminoacid() {
175            for atom_name in ["N", "CA", "C", "O"] {
176                if let Some(atom) = residue.find_atom_by_name(atom_name) {
177                    let [x, y, z] = atom.coords();
178                    backbone_data.extend_from_slice(&[*x, *y, *z]);
179                } else {
180                    backbone_data.extend_from_slice(&[0.0, 0.0, 0.0]);
181                }
182            }
183        }
184        Tensor::from_vec(backbone_data, (1, res_count, 4, 3), &device)
185    }
186
187    /// create numeric Tensor of shape [1, <sequence-length>, 37, 3]
188    fn to_numeric_atom37(&self, device: &Device) -> Result<Tensor> {
189        let res_count = self.iter_residues_aminoacid().count();
190        let mut atom37_data = vec![0.0; res_count * 37 * 3];
191        for (res_idx, residue) in self.iter_residues_aminoacid().enumerate() {
192            for atom_type in AAAtom::iter().filter(|&a| a != AAAtom::Unknown) {
193                if let Some(atom) = residue.find_atom_by_name(&atom_type.to_string()) {
194                    let [x, y, z] = atom.coords();
195                    let base_idx = (res_idx * 37 + atom_type as usize) * 3;
196                    atom37_data[base_idx..base_idx + 3].copy_from_slice(&[*x, *y, *z]);
197                }
198            }
199        }
200        Tensor::from_vec(atom37_data, (1, res_count, 37, 3), &device)
201    }
202
203    // The purpose of this function it to create 3 output tensors that relate
204    // key information about a protein sequence and ligands it interacts with.
205    //
206    // The outputs are:
207    //  - y: 4D tensor of dimensions (<batch=1>, <num_residues>, <number_of_ligand_atoms>, <coords=3>)
208    //  - y_t: 1D tensor of dimension = <num_residues>
209    //  - y_m: 3D tensor of dimensions: (<batch=1>, <num_residues>, <number_of_ligand_atoms>))
210    //
211    fn to_numeric_ligand_atoms(&self, device: &Device) -> Result<(Tensor, Tensor, Tensor)> {
212        let mut coords = Vec::new();
213        let mut elements = Vec::new();
214        for residue in self.iter_residues() {
215            let res_name = residue.residue_name();
216            if residue.is_amino_acid() || res_name == "HOH" || res_name == "WAT" {
217                continue;
218            }
219            let atoms: Vec<_> = residue
220                .iter_atoms()
221                .filter(|atom| is_heavy_atom(atom.element()))
222                .collect();
223            for atom in atoms {
224                coords.push(*atom.coords());
225                elements.push(*atom.element());
226            }
227        }
228
229        // When there are no ligand atoms, backends like Metal cannot allocate zero-size
230        // buffers. Return a single dummy ligand slot with a zeroed mask so it has no
231        // effect on the model output.
232        if coords.is_empty() {
233            let cb = self.create_cb(device)?;
234            let (batch, res_num, _) = cb.dims3()?;
235            let y = Tensor::zeros((batch, res_num, 1, 3), DType::F32, device)?;
236            let y_t = Tensor::zeros((batch, res_num, 1), DType::I64, device)?;
237            let y_m = Tensor::zeros((batch, res_num, 1), DType::F32, device)?;
238            return Ok((y, y_t, y_m));
239        }
240
241        // raw starting tensors
242        let y = Tensor::from_slice(&coords.concat(), (coords.len(), 3), device)?;
243        let y_m = Tensor::ones_like(&y)?;
244        let y_t = Tensor::from_slice(
245            &elements
246                .iter()
247                .map(|e| e.atomic_number() as f32)
248                .collect::<Vec<_>>(),
249            (elements.len(),),
250            device,
251        )?;
252        let cb = self.create_cb(device)?;
253        let (batch, res_num, _coords) = cb.dims3()?;
254        let (number_of_ligand_atoms, _coords) = y.dims2()?;
255        let mask = Tensor::zeros((batch, res_num), DType::F32, device)?;
256        let (y, y_t, y_m, d_xy) =
257            get_nearest_neighbours(&cb, &mask, &y, &y_t, &y_m, number_of_ligand_atoms as i64)?;
258        let distance_mask = d_xy.lt(LIGAND_CUTOFF_SCORE)?.to_dtype(DType::F32)?;
259        let y_m_first = y_m.i((.., 0))?;
260        let mask = mask.squeeze(0)?;
261        let _mask_xy = distance_mask.mul(&mask)?.mul(&y_m_first)?;
262        let y = y.unsqueeze(0)?;
263        let y_t = y_t.to_dtype(DType::I64)?.unsqueeze(0)?;
264        let y_m = y_m.unsqueeze(0)?;
265        Ok((y, y_t, y_m))
266    }
267}
268
269/// Delegate all `StructureFeatures` methods to an `AtomCollection` adapter.
270///
271/// This lets callers pass a `&Model` directly to ML featurisation routines
272/// without manually calling `AtomCollection::from(&model)` at every call site.
273impl StructureFeatures for Model {
274    fn decode_amino_acids(&self, device: &Device) -> Result<Tensor> {
275        AtomCollection::from(self).decode_amino_acids(device)
276    }
277    fn encode_amino_acids(&self, device: &Device) -> Result<Tensor> {
278        AtomCollection::from(self).encode_amino_acids(device)
279    }
280    fn create_cb(&self, device: &Device) -> Result<Tensor> {
281        AtomCollection::from(self).create_cb(device)
282    }
283    fn featurize_lmpnn(&self, device: &Device) -> Result<ProteinFeatures> {
284        AtomCollection::from(self).featurize_lmpnn(device)
285    }
286    fn get_res_index(&self) -> Vec<u32> {
287        AtomCollection::from(self).get_res_index()
288    }
289    fn to_numeric_backbone_atoms(&self, device: &Device) -> Result<Tensor> {
290        AtomCollection::from(self).to_numeric_backbone_atoms(device)
291    }
292    fn to_numeric_atom37(&self, device: &Device) -> Result<Tensor> {
293        AtomCollection::from(self).to_numeric_atom37(device)
294    }
295    fn to_numeric_ligand_atoms(&self, device: &Device) -> Result<(Tensor, Tensor, Tensor)> {
296        AtomCollection::from(self).to_numeric_ligand_atoms(device)
297    }
298}
299
300#[cfg(test)]
301mod tests {
302    use super::*;
303    use ferritin_core::load_structure;
304    use ferritin_test_data::TestFile;
305
306    /// `decode_amino_acids` must round-trip with `encode_amino_acids`.
307    ///
308    /// `encode_amino_acids` produces integer indices (u32); `decode_amino_acids`
309    /// produces ASCII byte values (u8). For every standard amino acid the cycle
310    ///   residue_name -> aa3to1 -> aa1to_int -> int_to_aa1 -> u8
311    /// must yield the same one-letter code that `aa3to1` returned.
312    #[test]
313    fn test_decode_amino_acids_roundtrip() -> candle_core::Result<()> {
314        let device = Device::Cpu;
315        let (pdb_file, _temp) = TestFile::protein_01().create_temp().map_err(|e| {
316            candle_core::Error::Msg(format!("test file setup failed: {e}"))
317        })?;
318        let ac = load_structure(pdb_file).map_err(|e| {
319            candle_core::Error::Msg(format!("load_structure failed: {e}"))
320        })?;
321
322        let encoded = ac.encode_amino_acids(&device)?;
323        let decoded = ac.decode_amino_acids(&device)?;
324
325        // Both tensors must have shape [1, n].
326        assert_eq!(encoded.dims(), decoded.dims());
327
328        let n = encoded.dim(1)?;
329        let enc_vals: Vec<u32> = encoded.reshape(n)?.to_vec1()?;
330        let dec_bytes: Vec<u8> = decoded.reshape(n)?.to_vec1()?;
331
332        // For each position: int_to_aa1(encode_val) as u8 == decoded byte.
333        for (idx, (&enc, &dec)) in enc_vals.iter().zip(dec_bytes.iter()).enumerate() {
334            use super::super::utilities::int_to_aa1;
335            let expected = int_to_aa1(enc) as u8;
336            assert_eq!(
337                dec, expected,
338                "Mismatch at position {idx}: encoded={enc}, decoded byte={dec}, expected={expected}"
339            );
340        }
341        Ok(())
342    }
343
344    /// Index 20 is the unknown-residue sentinel; it must decode to `'X'` (ASCII 88).
345    #[test]
346    fn test_decode_amino_acids_unknown_sentinel() {
347        use super::super::utilities::int_to_aa1;
348        let ch = int_to_aa1(20);
349        assert_eq!(ch, 'X', "sentinel index 20 must decode to 'X'");
350        // Any out-of-range index must also fall back to 'X'.
351        let ch_oob = int_to_aa1(99);
352        assert_eq!(ch_oob, 'X', "out-of-range index must decode to 'X'");
353    }
354
355    /// `decode_amino_acids` output shape must be [1, sequence_length].
356    #[test]
357    fn test_decode_amino_acids_shape() -> candle_core::Result<()> {
358        let device = Device::Cpu;
359        let (pdb_file, _temp) = TestFile::protein_01().create_temp().map_err(|e| {
360            candle_core::Error::Msg(format!("test file setup failed: {e}"))
361        })?;
362        let ac = load_structure(pdb_file).map_err(|e| {
363            candle_core::Error::Msg(format!("load_structure failed: {e}"))
364        })?;
365
366        let n = ac.iter_residues_aminoacid().count();
367        let decoded = ac.decode_amino_acids(&device)?;
368        assert_eq!(decoded.dims(), &[1, n]);
369        Ok(())
370    }
371
372    /// `encode_amino_acids` produces u32 integer indices in [0, 20].
373    #[test]
374    fn test_encode_amino_acids_shape_and_range() -> candle_core::Result<()> {
375        let device = Device::Cpu;
376        let (pdb_file, _temp) = TestFile::protein_01().create_temp().map_err(|e| {
377            candle_core::Error::Msg(format!("test file setup failed: {e}"))
378        })?;
379        let ac = load_structure(pdb_file).map_err(|e| {
380            candle_core::Error::Msg(format!("load_structure failed: {e}"))
381        })?;
382
383        let n = ac.iter_residues_aminoacid().count();
384        let encoded = ac.encode_amino_acids(&device)?;
385        assert_eq!(encoded.dims(), &[1, n]);
386
387        let vals: Vec<u32> = encoded.reshape(n)?.to_vec1()?;
388        for v in vals {
389            assert!(v <= 20, "encoded index {v} out of range [0, 20]");
390        }
391        Ok(())
392    }
393
394    /// `get_res_index` returns one entry per amino acid residue.
395    #[test]
396    fn test_get_res_index_length() -> candle_core::Result<()> {
397        let (pdb_file, _temp) = TestFile::protein_01().create_temp().map_err(|e| {
398            candle_core::Error::Msg(format!("test file setup failed: {e}"))
399        })?;
400        let ac = load_structure(pdb_file).map_err(|e| {
401            candle_core::Error::Msg(format!("load_structure failed: {e}"))
402        })?;
403
404        let n = ac.iter_residues_aminoacid().count();
405        let res_index = ac.get_res_index();
406        assert_eq!(res_index.len(), n);
407        Ok(())
408    }
409
410    /// `create_cb` output has the right shape [1, n, 3].
411    #[test]
412    fn test_create_cb_shape() -> candle_core::Result<()> {
413        let device = Device::Cpu;
414        let (pdb_file, _temp) = TestFile::protein_01().create_temp().map_err(|e| {
415            candle_core::Error::Msg(format!("test file setup failed: {e}"))
416        })?;
417        let ac = load_structure(pdb_file).map_err(|e| {
418            candle_core::Error::Msg(format!("load_structure failed: {e}"))
419        })?;
420
421        let n = ac.iter_residues_aminoacid().count();
422        let cb = ac.create_cb(&device)?;
423        assert_eq!(cb.dims(), &[1, n, 3]);
424        Ok(())
425    }
426}
ferritin_plms/featurize/structure_features.rs

ferritin_plms/featurize/
structure_features.rs
