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::pmpnn_runner::{ProteinMPNNModels, ProteinMPNNRunner};
18use super::proteinfeatures::ProteinFeatures;
19use crate::StructureFeatures;
20use anyhow::Error;
21use candle_core::{DType, Device, Tensor};
22use clap::ValueEnum;
23use ferritin_core::load_structure;
24
25/// Responsible for taking CLI args and returning the Features and Model
26///
27#[allow(dead_code)]
28pub struct MPNNExecConfig {
29    pub(crate) protein_inputs: String, // Todo: make this optionally plural
30    pub(crate) run_config: RunConfig,
31    pub(crate) aabias_config: Option<AABiasConfig>,
32    pub(crate) ligand_mpnn_config: Option<LigandMPNNConfig>,
33    pub(crate) membrane_mpnn_config: Option<MembraneMPNNConfig>,
34    pub(crate) multi_pdb_config: Option<MultiPDBConfig>,
35    pub(crate) residue_control_config: Option<ResidueControl>,
36    pub(crate) device: Device,
37}
38
39impl MPNNExecConfig {
40    #[allow(clippy::too_many_arguments)] // CLI-driven config; consolidation is ferritin-100.8
41    pub fn new(
42        device: Device,
43        pdb_path: String,
44        run_config: RunConfig,
45        residue_config: Option<ResidueControl>,
46        aa_bias: Option<AABiasConfig>,
47        lig_mpnn_specific: Option<LigandMPNNConfig>,
48        membrane_mpnn_specific: Option<MembraneMPNNConfig>,
49        multi_pdb_specific: Option<MultiPDBConfig>,
50    ) -> Result<Self, Error> {
51        Ok(MPNNExecConfig {
52            protein_inputs: pdb_path,
53            run_config,
54            aabias_config: aa_bias,
55            ligand_mpnn_config: lig_mpnn_specific,
56            membrane_mpnn_config: membrane_mpnn_specific,
57            residue_control_config: residue_config,
58            multi_pdb_config: multi_pdb_specific,
59            device,
60        })
61    }
62    /// Load the weights for `model_type`.
63    ///
64    /// Delegates to [`ProteinMPNNRunner::from_pretrained`], which downloads from
65    /// `zcpbx/ligandmpnn-weights`. This previously extracted an embedded test
66    /// fixture (`TestFile::ligmpnn_pmpnn_01`) to a temp file, which made the
67    /// 37 MB `ferritin-test-data` crate a runtime dependency of every
68    /// downstream consumer (ferritin-100.10).
69    pub fn load_model(&self, model_type: ModelTypes) -> Result<ProteinMPNN, Error> {
70        let variant = match model_type {
71            ModelTypes::ProteinMPNN => ProteinMPNNModels::V48_020,
72            ModelTypes::LigandMPNN => ProteinMPNNModels::LigandV32_020_25,
73        };
74        Ok(ProteinMPNNRunner::from_pretrained(variant, self.device.clone())?.into_model())
75    }
76    pub fn generate_protein_features(&self) -> Result<ProteinFeatures, Error> {
77        let device = self.device.clone();
78        let base_dtype = DType::F32;
79
80        // init the Protein Features
81        let ac = load_structure(self.protein_inputs.clone())?;
82
83        let s = ac
84            .encode_amino_acids(&device)
85            .expect("A complete convertion to locations");
86        let x_37 = ac.to_numeric_atom37(&device)?;
87        let x_37_mask = Tensor::ones((x_37.dim(0)?, x_37.dim(1)?), base_dtype, &device)?;
88        let (y, y_t, y_m) = ac.to_numeric_ligand_atoms(&device)?;
89        let res_idx = ac.get_res_index();
90        let res_idx_len = res_idx.len();
91        let res_idx_tensor = Tensor::from_vec(res_idx, (1, res_idx_len), &device)?;
92
93        // chain residues
94        let chain_letters: Vec<String> = ac
95            .iter_residues_aminoacid()
96            .map(|res| res.chain_id().to_string())
97            .collect();
98
99        // unique Chains
100        let chain_list: Vec<String> = chain_letters
101            .clone()
102            .into_iter()
103            .collect::<std::collections::HashSet<_>>()
104            .into_iter()
105            .collect();
106
107        // assert_eq!(true, false);
108
109        // update residue info
110        // residue_config: Option<ResidueControl>,
111        // handle these:
112        // pub fixed_residues: Option<String>,
113        // pub redesigned_residues: Option<String>,
114        // pub symmetry_residues: Option<String>,
115        // pub symmetry_weights: Option<String>,
116        // pub chains_to_design: Option<String>,
117        // pub parse_these_chains_only: Option<String>,
118
119        // update AA bias
120        // handle these:
121        // aa_bias: Option<AABiasConfig>,
122        // pub bias_aa: Option<String>,
123        // pub bias_aa_per_residue: Option<String>,
124        // pub omit_aa: Option<String>,
125        // pub omit_aa_per_residue: Option<String>,
126
127        // update LigmpnnConfif
128        // lig_mpnn_specific: Option<LigandMPNNConfig>,
129        // handle these:
130        // pub checkpoint_ligand_mpnn: Option<String>,
131        // pub ligand_mpnn_use_atom_context: Option<i32>,
132        // pub ligand_mpnn_use_side_chain_context: Option<i32>,
133        // pub ligand_mpnn_cutoff_for_score: Option<String>,
134
135        // update Membrane MPNN Config
136        // membrane_mpnn_specific: Option<MembraneMPNNConfig>,
137        // handle these:
138        // pub global_transmembrane_label: Option<i32>,
139        // pub transmembrane_buried: Option<String>,
140        // pub transmembrane_interface: Option<String>,
141
142        // update multipdb
143        // multi_pdb_specific: Option<MultiPDBConfig>,
144        // pub pdb_path_multi: Option<String>,
145        // pub fixed_residues_multi: Option<String>,
146        // pub redesigned_residues_multi: Option<String>,
147        // pub omit_aa_per_residue_multi: Option<String>,
148        // pub bias_aa_per_residue_multi: Option<String>,
149
150        // println!("Returning Protein Features....");
151        // return ligand MPNN.
152        Ok(ProteinFeatures {
153            s,                       // protein amino acids sequences as 1D Tensor of u32
154            x: x_37,                 // protein co-oords by residue [1, 37, 4]
155            x_mask: Some(x_37_mask), // protein mask by residue
156            y,                       // ligand coords
157            y_t,                     // encoded ligand atom names
158            y_m: Some(y_m),          // ligand mask
159            r_idx: res_idx_tensor,   // protein residue indices shape=[length]
160            chain_labels: None,      //  # protein chain letters shape=[length]
161            chain_letters,           // chain_letters: shape=[length]
162            mask_c: None,            // mask_c:  shape=[length]
163            chain_list,
164        })
165    }
166}
167
168#[derive(Debug, Clone, ValueEnum, Copy)]
169pub enum ModelTypes {
170    #[value(name = "protein_mpnn")]
171    ProteinMPNN,
172    #[value(name = "ligand_mpnn")]
173    LigandMPNN,
174}
175
176#[derive(Debug)]
177/// Amino Acid Biasing
178pub struct AABiasConfig {
179    pub bias_aa: Option<String>,
180    pub bias_aa_per_residue: Option<String>,
181    pub omit_aa: Option<String>,
182    pub omit_aa_per_residue: Option<String>,
183}
184
185/// LigandMPNN Specific
186pub struct LigandMPNNConfig {
187    pub checkpoint_ligand_mpnn: Option<String>,
188    pub ligand_mpnn_use_atom_context: Option<i32>,
189    pub ligand_mpnn_use_side_chain_context: Option<i32>,
190    pub ligand_mpnn_cutoff_for_score: Option<String>,
191}
192
193/// Membrane MPNN Specific
194pub struct MembraneMPNNConfig {
195    pub global_transmembrane_label: Option<i32>,
196    pub transmembrane_buried: Option<String>,
197    pub transmembrane_interface: Option<String>,
198}
199
200/// Multi-PDB Related
201pub struct MultiPDBConfig {
202    pub pdb_path_multi: Option<String>,
203    pub fixed_residues_multi: Option<String>,
204    pub redesigned_residues_multi: Option<String>,
205    pub omit_aa_per_residue_multi: Option<String>,
206    pub bias_aa_per_residue_multi: Option<String>,
207}
208#[derive(Clone, Debug)]
209pub struct ProteinMPNNConfig {
210    pub atom_context_num: usize,
211    pub augment_eps: f32,
212    pub dropout_ratio: f32,
213    pub edge_features: i64,
214    pub hidden_dim: i64,
215    pub k_neighbors: i64,
216    pub ligand_mpnn_use_side_chain_context: bool,
217    pub model_type: ModelTypes,
218    pub node_features: i64,
219    pub num_decoder_layers: i64,
220    pub num_encoder_layers: i64,
221    pub num_letters: i64,
222    pub num_rbf: i64,
223    pub scale_factor: f64,
224    pub vocab: i64,
225}
226
227impl ProteinMPNNConfig {
228    pub fn proteinmpnn() -> Self {
229        Self {
230            atom_context_num: 0,
231            augment_eps: 0.0,
232            dropout_ratio: 0.1,
233            edge_features: 128,
234            hidden_dim: 128,
235            k_neighbors: 48,
236            ligand_mpnn_use_side_chain_context: false,
237            model_type: ModelTypes::ProteinMPNN,
238            node_features: 128,
239            num_decoder_layers: 3,
240            num_encoder_layers: 3,
241            num_letters: 21,
242            num_rbf: 16,
243            scale_factor: 30.0,
244            vocab: 21,
245        }
246    }
247
248    /// `ligandmpnn_v_32_020_25`: 32 neighbours, 25 ligand context atoms.
249    ///
250    /// `k_neighbors` and `atom_context_num` are checkpoint facts — the `.pt`
251    /// carries them as `num_edges` and `atom_context_num` — but they are plain
252    /// Python ints rather than tensors, so candle's pickle reader does not
253    /// surface them and they are declared here instead. Verified against the
254    /// four checkpoints in `ferritin-test-data` (ferritin-100.11).
255    pub fn ligandmpnn() -> Self {
256        Self {
257            atom_context_num: 25,
258            k_neighbors: 32,
259            model_type: ModelTypes::LigandMPNN,
260            ..Self::proteinmpnn()
261        }
262    }
263}
264
265#[derive(Debug)]
266pub struct ResidueControl {
267    pub fixed_residues: Option<String>,
268    pub redesigned_residues: Option<String>,
269    pub symmetry_residues: Option<String>,
270    pub symmetry_weights: Option<String>,
271    pub chains_to_design: Option<String>,
272    pub parse_these_chains_only: Option<String>,
273}
274
275#[derive(Debug)]
276pub struct RunConfig {
277    pub model_type: Option<ModelTypes>,
278    pub seed: Option<i32>,
279    pub temperature: Option<f32>,
280    pub verbose: Option<i32>,
281    pub save_stats: Option<bool>,
282    pub batch_size: Option<i32>,
283    pub number_of_batches: Option<i32>,
284    pub file_ending: Option<String>,
285    pub zero_indexed: Option<i32>,
286    pub homo_oligomer: Option<i32>,
287    pub fasta_seq_separation: Option<String>,
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293
294    fn exec_config() -> MPNNExecConfig {
295        MPNNExecConfig::new(
296            Device::Cpu,
297            // load_model never reads the structure, so a placeholder is fine.
298            "unused.pdb".to_string(),
299            RunConfig {
300                model_type: None,
301                seed: None,
302                temperature: None,
303                verbose: None,
304                save_stats: None,
305                batch_size: None,
306                number_of_batches: None,
307                file_ending: None,
308                zero_indexed: None,
309                homo_oligomer: None,
310                fasta_seq_separation: None,
311            },
312            None,
313            None,
314            None,
315            None,
316            None,
317        )
318        .expect("MPNNExecConfig::new should not fail")
319    }
320
321    /// Both CLI model types map to a real checkpoint.
322    ///
323    /// `--model-type ligand_mpnn` used to `panic!("not implented!")`, then
324    /// returned an "unimplemented" error, and now loads
325    /// `ligandmpnn_v_32_020_25`. Asserted on the mapping rather than on a
326    /// download so the test stays offline; the weights themselves are covered
327    /// by `test_ligandmpnn_parity_vs_python_reference` (ferritin-100.11).
328    #[test]
329    fn test_every_cli_model_type_maps_to_a_registry_row() {
330        for (model_type, expected) in [
331            (ModelTypes::ProteinMPNN, "proteinmpnn-v48-020"),
332            (ModelTypes::LigandMPNN, "ligandmpnn-v32-020-25"),
333        ] {
334            let variant = match model_type {
335                ModelTypes::ProteinMPNN => ProteinMPNNModels::V48_020,
336                ModelTypes::LigandMPNN => ProteinMPNNModels::LigandV32_020_25,
337            };
338            assert_eq!(variant.registry_id(), expected, "{model_type:?}");
339            assert!(
340                crate::registry::lookup(expected).is_some(),
341                "{expected} must be a registry row"
342            );
343        }
344    }
345
346    /// The two checkpoints need different architecture configs, and nothing in
347    /// the tensor shapes would catch a mix-up.
348    #[test]
349    fn test_model_types_carry_distinct_configs() {
350        let protein = ProteinMPNNModels::V48_020.config();
351        let ligand = ProteinMPNNModels::LigandV32_020_25.config();
352        assert_eq!(protein.k_neighbors, 48);
353        assert_eq!(protein.atom_context_num, 0);
354        assert_eq!(ligand.k_neighbors, 32);
355        assert_eq!(ligand.atom_context_num, 25);
356    }
357
358    /// The ProteinMPNN branch loads from HuggingFace rather than extracting the
359    /// `ferritin-test-data` fixture it used to depend on (ferritin-100.10).
360    #[test]
361    #[ignore = "downloads zcpbx/ligandmpnn-weights from HuggingFace"]
362    fn test_load_model_proteinmpnn_loads_from_hub() {
363        exec_config()
364            .load_model(ModelTypes::ProteinMPNN)
365            .expect("ProteinMPNN should load from the hub");
366    }
367}
ferritin_plms/ligandmpnn/configs.rs

ferritin_plms/ligandmpnn/
configs.rs
