ferritin_plms/ligandmpnn/
proteinfeatures.rs

1//! Protein Featurizer for ProteinMPNN/LignadMPNN
2//!
3//! Extract protein features for ligandmpnn
4//!
5//! Returns a set of features calculated from protein structure
6//! including:
7//! - Residue-level features like amino acid type, secondary structure
8//! - Geometric features like distances, angles
9//! - Chemical features like hydrophobicity, charge
10//! - Evolutionary features from MSA profiles
11use crate::featurize::utilities::aa1to_int;
12use candle_core::{Device, Result, Tensor};
13use std::collections::{HashMap, HashSet};
14
15#[allow(dead_code)]
16pub struct ProteinFeatures {
17    /// protein amino acids sequences as 1D Tensor of u32
18    pub(crate) s: Tensor,
19    /// protein co-oords by residue [batch, seqlength, 37, 3]
20    pub(crate) x: Tensor,
21    /// protein mask by residue
22    pub(crate) x_mask: Option<Tensor>,
23    /// ligand coords
24    pub(crate) y: Tensor,
25    /// encoded ligand atom names
26    pub(crate) y_t: Tensor,
27    /// ligand mask
28    pub(crate) y_m: Option<Tensor>,
29    /// R_idx:         Tensor dimensions: torch.Size([93])          # protein residue indices shape=[length]
30    pub(crate) r_idx: Tensor,
31    /// chain_labels:  Tensor dimensions: torch.Size([93])          # protein chain letters shape=[length]
32    pub(crate) chain_labels: Option<Vec<f64>>,
33    /// chain_letters: NumPy array dimensions: (93,)
34    pub(crate) chain_letters: Vec<String>,
35    /// mask_c:        Tensor dimensions: torch.Size([93])
36    pub(crate) mask_c: Option<Tensor>,
37    pub(crate) chain_list: Vec<String>,
38}
39impl ProteinFeatures {
40    pub fn get_coords(&self) -> &Tensor {
41        &self.x
42    }
43    pub fn get_sequence(&self) -> &Tensor {
44        &self.s
45    }
46    pub fn get_sequence_mask(&self) -> Option<&Tensor> {
47        self.x_mask.as_ref()
48    }
49    pub fn get_residue_index(&self) -> &Tensor {
50        &self.r_idx
51    }
52    #[allow(clippy::type_complexity)] // (vocab, tok→id, id→tok) triple; a named type adds little
53    pub fn get_encoded(
54        &self,
55    ) -> Result<(Vec<String>, HashMap<String, usize>, HashMap<usize, String>)> {
56        let r_idx_list = &self.r_idx.flatten_all()?.to_vec1::<u32>()?;
57        let chain_letters_list = &self.chain_letters;
58        let encoded_residues: Vec<String> = r_idx_list
59            .iter()
60            .enumerate()
61            .map(|(i, r_idx)| format!("{}{}", chain_letters_list[i], r_idx))
62            .collect();
63        let encoded_residue_dict: HashMap<String, usize> = encoded_residues
64            .iter()
65            .enumerate()
66            .map(|(i, s)| (s.clone(), i))
67            .collect();
68        let encoded_residue_dict_rev: HashMap<usize, String> = encoded_residues
69            .iter()
70            .enumerate()
71            .map(|(i, s)| (i, s.clone()))
72            .collect();
73        Ok((
74            encoded_residues,
75            encoded_residue_dict,
76            encoded_residue_dict_rev,
77        ))
78    }
79    // Fixed Residue List --> Tensor of 1/0
80    // Inputs: `"C1 C2 C3 C4 C5 C6 C7 C8 C9 C10`
81    pub fn get_encoded_tensor(&self, fixed_residues: String, device: &Device) -> Result<Tensor> {
82        let res_set: HashSet<String> = fixed_residues.split(' ').map(String::from).collect();
83        let (encoded_res, _, _) = &self.get_encoded()?;
84        Tensor::from_iter(
85            encoded_res
86                .iter()
87                .map(|item| u32::from(!res_set.contains(item))),
88            device,
89        )
90    }
91    pub fn get_chain_mask_tensor(
92        &self,
93        chains_to_design: Vec<String>,
94        device: &Device,
95    ) -> Result<Tensor> {
96        let mask_values: Vec<u32> = self
97            .chain_letters
98            .iter()
99            .map(|chain| u32::from(chains_to_design.contains(chain)))
100            .collect();
101        Tensor::from_iter(mask_values, device)
102    }
103    pub fn update_mask(&mut self, tensor: Tensor) -> Result<()> {
104        self.x_mask = match self.x_mask.as_ref() {
105            Some(mask) => Some(mask.mul(&tensor)?),
106            None => Some(tensor),
107        };
108        Ok(())
109    }
110    // Fixed Residue List --> Tensor of length 21
111    // Inputs: `A:10.0"`
112    pub fn create_bias_tensor(&self, bias_aa: Option<String>) -> Result<Tensor> {
113        let device = self.s.device();
114        let dtype = self.s.dtype();
115        match bias_aa {
116            None => Tensor::zeros(21, dtype, device),
117            Some(bias_aa) => {
118                let mut bias_values = vec![0.0f32; 21];
119                for pair in bias_aa.split(',') {
120                    if let Some((aa, value_str)) = pair.split_once(':')
121                        && let Ok(value) = value_str.parse::<f32>()
122                        && let Some(aa_char) = aa.chars().next()
123                    {
124                        let idx = aa1to_int(aa_char) as usize;
125                        bias_values[idx] = value;
126                    }
127                }
128                Tensor::from_slice(&bias_values, 21, device)
129            }
130        }
131    }
132    pub fn save_to_safetensor(&self, path: &str) -> Result<()> {
133        let mut tensors: HashMap<String, Tensor> = HashMap::new();
134        // this is only one field. need to do the rest of the fields
135        tensors.insert("protein_atom_sequence".to_string(), self.s.clone());
136        tensors.insert("protein_atom_positions".to_string(), self.x.clone());
137        tensors.insert("ligand_atom_positions".to_string(), self.y.clone());
138        tensors.insert("ligand_atom_name".to_string(), self.y_t.clone());
139        candle_core::safetensors::save(&tensors, path)?;
140        Ok(())
141    }
142}
ferritin_plms/ligandmpnn/proteinfeatures.rs

ferritin_plms/ligandmpnn/
proteinfeatures.rs
