ferritin_plms/ligandmpnn/
configs.rs

1//! PMPNN Core Config and Builder API
2//!
3//! This module provides configuration structs and builders for the PMPNN protein design system.
4//!
5//! # Core Configuration Types
6//!
7//! - `ModelTypes` - Enum of supported model architectures
8//! - `ProteinMPNNConfig` - Core model parameters
9//! - `AABiasConfig` - Amino acid biasing controls
10//! - `LigandMPNNConfig` - LigandMPNN specific settings
11//! - `MembraneMPNNConfig` - MembraneMPNN specific settings
12//! - `MultiPDBConfig` - Multi-PDB mode configuration
13//! - `ResidueControl` - Residue-level design controls
14//! - `RunConfig` - Runtime execution parameters// Core Configs for handling CLI ARGs and Model Params
15
16use super::model::ProteinMPNN;
17use super::proteinfeatures::ProteinFeatures;
18use crate::StructureFeatures;
19use anyhow::Error;
20use candle_core::pickle::PthTensors;
21use candle_core::{DType, Device, Tensor};
22use candle_nn::VarBuilder;
23use clap::ValueEnum;
24use ferritin_core::load_structure;
25use ferritin_test_data::TestFile;
26
27/// Responsible for taking CLI args and returning the Features and Model
28///
29#[allow(dead_code)]
30pub struct MPNNExecConfig {
31    pub(crate) protein_inputs: String, // Todo: make this optionally plural
32    pub(crate) run_config: RunConfig,
33    pub(crate) aabias_config: Option<AABiasConfig>,
34    pub(crate) ligand_mpnn_config: Option<LigandMPNNConfig>,
35    pub(crate) membrane_mpnn_config: Option<MembraneMPNNConfig>,
36    pub(crate) multi_pdb_config: Option<MultiPDBConfig>,
37    pub(crate) residue_control_config: Option<ResidueControl>,
38    pub(crate) device: Device,
39}
40
41impl MPNNExecConfig {
42    pub fn new(
43        device: Device,
44        pdb_path: String,
45        run_config: RunConfig,
46        residue_config: Option<ResidueControl>,
47        aa_bias: Option<AABiasConfig>,
48        lig_mpnn_specific: Option<LigandMPNNConfig>,
49        membrane_mpnn_specific: Option<MembraneMPNNConfig>,
50        multi_pdb_specific: Option<MultiPDBConfig>,
51    ) -> Result<Self, Error> {
52        Ok(MPNNExecConfig {
53            protein_inputs: pdb_path,
54            run_config,
55            aabias_config: aa_bias,
56            ligand_mpnn_config: lig_mpnn_specific,
57            membrane_mpnn_config: membrane_mpnn_specific,
58            residue_control_config: residue_config,
59            multi_pdb_config: multi_pdb_specific,
60            device,
61        })
62    }
63    // Todo: refactor this to use loader.
64    pub fn load_model(&self, model_type: ModelTypes) -> Result<ProteinMPNN, Error> {
65        let default_dtype = DType::F32;
66        match model_type {
67            ModelTypes::ProteinMPNN => {
68                // this is a hidden dep....
69                // todo: use hf_hub
70                let (mpnn_file, _handle) = TestFile::ligmpnn_pmpnn_01().create_temp()?;
71                let pth = PthTensors::new(mpnn_file, Some("model_state_dict"))?;
72                let vb =
73                    VarBuilder::from_backend(Box::new(pth), default_dtype, self.device.clone());
74                let pconf = ProteinMPNNConfig::proteinmpnn();
75                Ok(ProteinMPNN::load(vb, &pconf).expect("Unable to load the PMPNN Model"))
76            }
77            _ => panic!("not implented!"),
78        }
79    }
80    pub fn generate_protein_features(&self) -> Result<ProteinFeatures, Error> {
81        let device = self.device.clone();
82        let base_dtype = DType::F32;
83
84        // init the Protein Features
85        let ac = load_structure(self.protein_inputs.clone())?;
86
87        let s = ac
88            .encode_amino_acids(&device)
89            .expect("A complete convertion to locations");
90        let x_37 = ac.to_numeric_atom37(&device)?;
91        let x_37_mask = Tensor::ones((x_37.dim(0)?, x_37.dim(1)?), base_dtype, &device)?;
92        let (y, y_t, y_m) = ac.to_numeric_ligand_atoms(&device)?;
93        let res_idx = ac.get_res_index();
94        let res_idx_len = res_idx.len();
95        let res_idx_tensor = Tensor::from_vec(res_idx, (1, res_idx_len), &device)?;
96
97        // chain residues
98        let chain_letters: Vec<String> = ac
99            .iter_residues_aminoacid()
100            .map(|res| res.chain_id().to_string())
101            .collect();
102
103        // unique Chains
104        let chain_list: Vec<String> = chain_letters
105            .clone()
106            .into_iter()
107            .collect::<std::collections::HashSet<_>>()
108            .into_iter()
109            .collect();
110
111        // assert_eq!(true, false);
112
113        // update residue info
114        // residue_config: Option<ResidueControl>,
115        // handle these:
116        // pub fixed_residues: Option<String>,
117        // pub redesigned_residues: Option<String>,
118        // pub symmetry_residues: Option<String>,
119        // pub symmetry_weights: Option<String>,
120        // pub chains_to_design: Option<String>,
121        // pub parse_these_chains_only: Option<String>,
122
123        // update AA bias
124        // handle these:
125        // aa_bias: Option<AABiasConfig>,
126        // pub bias_aa: Option<String>,
127        // pub bias_aa_per_residue: Option<String>,
128        // pub omit_aa: Option<String>,
129        // pub omit_aa_per_residue: Option<String>,
130
131        // update LigmpnnConfif
132        // lig_mpnn_specific: Option<LigandMPNNConfig>,
133        // handle these:
134        // pub checkpoint_ligand_mpnn: Option<String>,
135        // pub ligand_mpnn_use_atom_context: Option<i32>,
136        // pub ligand_mpnn_use_side_chain_context: Option<i32>,
137        // pub ligand_mpnn_cutoff_for_score: Option<String>,
138
139        // update Membrane MPNN Config
140        // membrane_mpnn_specific: Option<MembraneMPNNConfig>,
141        // handle these:
142        // pub global_transmembrane_label: Option<i32>,
143        // pub transmembrane_buried: Option<String>,
144        // pub transmembrane_interface: Option<String>,
145
146        // update multipdb
147        // multi_pdb_specific: Option<MultiPDBConfig>,
148        // pub pdb_path_multi: Option<String>,
149        // pub fixed_residues_multi: Option<String>,
150        // pub redesigned_residues_multi: Option<String>,
151        // pub omit_aa_per_residue_multi: Option<String>,
152        // pub bias_aa_per_residue_multi: Option<String>,
153
154        // println!("Returning Protein Features....");
155        // return ligand MPNN.
156        Ok(ProteinFeatures {
157            s,                       // protein amino acids sequences as 1D Tensor of u32
158            x: x_37,                 // protein co-oords by residue [1, 37, 4]
159            x_mask: Some(x_37_mask), // protein mask by residue
160            y,                       // ligand coords
161            y_t,                     // encoded ligand atom names
162            y_m: Some(y_m),          // ligand mask
163            r_idx: res_idx_tensor,   // protein residue indices shape=[length]
164            chain_labels: None,      //  # protein chain letters shape=[length]
165            chain_letters,           // chain_letters: shape=[length]
166            mask_c: None,            // mask_c:  shape=[length]
167            chain_list,
168        })
169    }
170}
171
172#[derive(Debug, Clone, ValueEnum, Copy)]
173pub enum ModelTypes {
174    #[value(name = "protein_mpnn")]
175    ProteinMPNN,
176    #[value(name = "ligand_mpnn")]
177    LigandMPNN,
178}
179
180#[derive(Debug)]
181/// Amino Acid Biasing
182pub struct AABiasConfig {
183    pub bias_aa: Option<String>,
184    pub bias_aa_per_residue: Option<String>,
185    pub omit_aa: Option<String>,
186    pub omit_aa_per_residue: Option<String>,
187}
188
189/// LigandMPNN Specific
190pub struct LigandMPNNConfig {
191    pub checkpoint_ligand_mpnn: Option<String>,
192    pub ligand_mpnn_use_atom_context: Option<i32>,
193    pub ligand_mpnn_use_side_chain_context: Option<i32>,
194    pub ligand_mpnn_cutoff_for_score: Option<String>,
195}
196
197/// Membrane MPNN Specific
198pub struct MembraneMPNNConfig {
199    pub global_transmembrane_label: Option<i32>,
200    pub transmembrane_buried: Option<String>,
201    pub transmembrane_interface: Option<String>,
202}
203
204/// Multi-PDB Related
205pub struct MultiPDBConfig {
206    pub pdb_path_multi: Option<String>,
207    pub fixed_residues_multi: Option<String>,
208    pub redesigned_residues_multi: Option<String>,
209    pub omit_aa_per_residue_multi: Option<String>,
210    pub bias_aa_per_residue_multi: Option<String>,
211}
212#[derive(Clone, Debug)]
213pub struct ProteinMPNNConfig {
214    pub atom_context_num: usize,
215    pub augment_eps: f32,
216    pub dropout_ratio: f32,
217    pub edge_features: i64,
218    pub hidden_dim: i64,
219    pub k_neighbors: i64,
220    pub ligand_mpnn_use_side_chain_context: bool,
221    pub model_type: ModelTypes,
222    pub node_features: i64,
223    pub num_decoder_layers: i64,
224    pub num_encoder_layers: i64,
225    pub num_letters: i64,
226    pub num_rbf: i64,
227    pub scale_factor: f64,
228    pub vocab: i64,
229}
230
231impl ProteinMPNNConfig {
232    pub fn proteinmpnn() -> Self {
233        Self {
234            atom_context_num: 0,
235            augment_eps: 0.0,
236            dropout_ratio: 0.1,
237            edge_features: 128,
238            hidden_dim: 128,
239            k_neighbors: 24,
240            ligand_mpnn_use_side_chain_context: false,
241            model_type: ModelTypes::ProteinMPNN,
242            node_features: 128,
243            num_decoder_layers: 3,
244            num_encoder_layers: 3,
245            num_letters: 21,
246            num_rbf: 16,
247            scale_factor: 1.0,
248            vocab: 21,
249        }
250    }
251}
252
253#[derive(Debug)]
254pub struct ResidueControl {
255    pub fixed_residues: Option<String>,
256    pub redesigned_residues: Option<String>,
257    pub symmetry_residues: Option<String>,
258    pub symmetry_weights: Option<String>,
259    pub chains_to_design: Option<String>,
260    pub parse_these_chains_only: Option<String>,
261}
262
263#[derive(Debug)]
264pub struct RunConfig {
265    pub model_type: Option<ModelTypes>,
266    pub seed: Option<i32>,
267    pub temperature: Option<f32>,
268    pub verbose: Option<i32>,
269    pub save_stats: Option<bool>,
270    pub batch_size: Option<i32>,
271    pub number_of_batches: Option<i32>,
272    pub file_ending: Option<String>,
273    pub zero_indexed: Option<i32>,
274    pub homo_oligomer: Option<i32>,
275    pub fasta_seq_separation: Option<String>,
276}
ferritin_plms/ligandmpnn/configs.rs

ferritin_plms/ligandmpnn/
configs.rs
