ferritin_plms/esmfold2/
input_types.rs

1//! Input builder types for ESMFold2 structure prediction.
2//!
3//! This module provides typed builder structs for constructing the
4//! [`StructurePredictionInput`] required by the ESMFold2 API. It mirrors
5//! the Python `esm.models.esmfold2` SDK, supporting protein chains,
6//! DNA chains with optional chemical modifications, and small-molecule
7//! ligands specified by CCD (Chemical Component Dictionary) codes.
8//!
9//! # Example
10//! ```rust
11//! use ferritin_plms::esmfold2::input_types::{
12//!     DNAInput, LigandInput, ProteinInput, StructurePredictionInput,
13//! };
14//!
15//! let protein = ProteinInput::new("A", "MKTAYIAK").unwrap();
16//! let ligand  = LigandInput::from_ccd("L", "ATP");
17//! let input   = StructurePredictionInput::new()
18//!     .add_protein(protein)
19//!     .add_ligand(ligand);
20//!
21//! assert_eq!(input.num_chains(), 2);
22//! ```
23
24use std::fmt;
25
26/// Standard 20 amino acids plus common ambiguous / non-standard codes.
27///
28/// Uppercase and lowercase are both accepted.  Includes:
29/// - Standard: `ACDEFGHIKLMNPQRSTVWY`
30/// - Ambiguous: `BJOUXZ`  (B=Asp/Asn, J=Leu/Ile, O=Pyrrolysine, U=Selenocysteine,
31///   X=unknown, Z=Glu/Gln)
32/// - Gap `-` and stop `*`
33const PROTEIN_ALPHABET: &str = "ACDEFGHIKLMNPQRSTVWYBJOUXZacdefghiklmnpqrstvwybjouxz-*";
34
35/// IUPAC DNA nucleotide alphabet including ambiguous base codes.
36const DNA_ALPHABET: &str = "ACGTNRYSWKMBDHVacgtnryswkmbdhv";
37
38// ── ProteinInput ─────────────────────────────────────────────────────────────
39
40/// A single protein chain for structure prediction.
41#[derive(Debug, Clone)]
42pub struct ProteinInput {
43    /// Chain identifier (e.g. `"A"`).
44    pub id: String,
45    /// Amino acid sequence in single-letter code.
46    pub sequence: String,
47}
48
49impl ProteinInput {
50    /// Construct a new `ProteinInput`, validating the sequence alphabet.
51    ///
52    /// Returns `Err` if the sequence is empty or contains an unrecognised character.
53    pub fn new(id: impl Into<String>, sequence: impl Into<String>) -> Result<Self, String> {
54        let id = id.into();
55        let sequence = sequence.into();
56        validate_protein_sequence(&sequence)?;
57        Ok(Self { id, sequence })
58    }
59
60    /// Length of the amino acid sequence.
61    pub fn len(&self) -> usize {
62        self.sequence.len()
63    }
64
65    /// Returns `true` if the sequence is empty.
66    pub fn is_empty(&self) -> bool {
67        self.sequence.is_empty()
68    }
69}
70
71impl fmt::Display for ProteinInput {
72    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
73        write!(f, "Protein[{}](len={})", self.id, self.sequence.len())
74    }
75}
76
77// ── Modification ─────────────────────────────────────────────────────────────
78
79/// A chemical modification at a specific position within a nucleic-acid chain.
80#[derive(Debug, Clone)]
81pub struct Modification {
82    /// 1-based position in the sequence.
83    pub position: usize,
84    /// CCD code identifying the modification (e.g. `"5MC"` for 5-methylcytosine).
85    pub ccd: String,
86}
87
88impl Modification {
89    /// Construct a new `Modification`.
90    pub fn new(position: usize, ccd: impl Into<String>) -> Self {
91        Self {
92            position,
93            ccd: ccd.into(),
94        }
95    }
96}
97
98impl fmt::Display for Modification {
99    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
100        write!(f, "Mod[{}@pos{}]", self.ccd, self.position)
101    }
102}
103
104// ── DNAInput ─────────────────────────────────────────────────────────────────
105
106/// A single DNA chain, optionally carrying chemical modifications.
107#[derive(Debug, Clone)]
108pub struct DNAInput {
109    /// Chain identifier.
110    pub id: String,
111    /// DNA nucleotide sequence (IUPAC single-letter codes).
112    pub sequence: String,
113    /// Chemical modifications along the chain (may be empty).
114    pub modifications: Vec<Modification>,
115}
116
117impl DNAInput {
118    /// Construct a new `DNAInput` without modifications.
119    ///
120    /// Returns `Err` if the sequence is empty or contains an invalid nucleotide.
121    pub fn new(id: impl Into<String>, sequence: impl Into<String>) -> Result<Self, String> {
122        let id = id.into();
123        let sequence = sequence.into();
124        validate_dna_sequence(&sequence)?;
125        Ok(Self {
126            id,
127            sequence,
128            modifications: Vec::new(),
129        })
130    }
131
132    /// Construct a new `DNAInput` with a pre-built list of modifications.
133    pub fn with_modifications(
134        id: impl Into<String>,
135        sequence: impl Into<String>,
136        modifications: Vec<Modification>,
137    ) -> Result<Self, String> {
138        let id = id.into();
139        let sequence = sequence.into();
140        validate_dna_sequence(&sequence)?;
141        Ok(Self {
142            id,
143            sequence,
144            modifications,
145        })
146    }
147
148    /// Append a modification and return `self` (builder-style).
149    pub fn add_modification(mut self, modification: Modification) -> Self {
150        self.modifications.push(modification);
151        self
152    }
153
154    /// Length of the nucleotide sequence.
155    pub fn len(&self) -> usize {
156        self.sequence.len()
157    }
158
159    /// Returns `true` if the sequence is empty.
160    pub fn is_empty(&self) -> bool {
161        self.sequence.is_empty()
162    }
163}
164
165impl fmt::Display for DNAInput {
166    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
167        write!(
168            f,
169            "DNA[{}](len={}, mods={})",
170            self.id,
171            self.sequence.len(),
172            self.modifications.len()
173        )
174    }
175}
176
177// ── LigandInput ──────────────────────────────────────────────────────────────
178
179/// A small-molecule ligand described by one or more CCD codes.
180#[derive(Debug, Clone)]
181pub struct LigandInput {
182    /// Ligand identifier.
183    pub id: String,
184    /// Ordered list of CCD codes composing the ligand.
185    pub ccd: Vec<String>,
186}
187
188impl LigandInput {
189    /// Construct a `LigandInput` from an explicit list of CCD codes.
190    pub fn new(id: impl Into<String>, ccd: Vec<String>) -> Self {
191        Self { id: id.into(), ccd }
192    }
193
194    /// Convenience constructor for a single-component ligand.
195    pub fn from_ccd(id: impl Into<String>, ccd: impl Into<String>) -> Self {
196        Self {
197            id: id.into(),
198            ccd: vec![ccd.into()],
199        }
200    }
201
202    /// Number of CCD components in this ligand.
203    pub fn num_components(&self) -> usize {
204        self.ccd.len()
205    }
206}
207
208impl fmt::Display for LigandInput {
209    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
210        write!(f, "Ligand[{}](ccd={:?})", self.id, self.ccd)
211    }
212}
213
214// ── ChainInput ───────────────────────────────────────────────────────────────
215
216/// One chain in a multi-entity structure prediction request.
217#[derive(Debug, Clone)]
218pub enum ChainInput {
219    Protein(ProteinInput),
220    DNA(DNAInput),
221    Ligand(LigandInput),
222}
223
224impl ChainInput {
225    /// Returns the chain identifier regardless of variant.
226    pub fn id(&self) -> &str {
227        match self {
228            ChainInput::Protein(p) => &p.id,
229            ChainInput::DNA(d) => &d.id,
230            ChainInput::Ligand(l) => &l.id,
231        }
232    }
233}
234
235impl fmt::Display for ChainInput {
236    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
237        match self {
238            ChainInput::Protein(p) => write!(f, "{p}"),
239            ChainInput::DNA(d) => write!(f, "{d}"),
240            ChainInput::Ligand(l) => write!(f, "{l}"),
241        }
242    }
243}
244
245impl From<ProteinInput> for ChainInput {
246    fn from(p: ProteinInput) -> Self {
247        ChainInput::Protein(p)
248    }
249}
250
251impl From<DNAInput> for ChainInput {
252    fn from(d: DNAInput) -> Self {
253        ChainInput::DNA(d)
254    }
255}
256
257impl From<LigandInput> for ChainInput {
258    fn from(l: LigandInput) -> Self {
259        ChainInput::Ligand(l)
260    }
261}
262
263// ── StructurePredictionInput ─────────────────────────────────────────────────
264
265/// The top-level input for an ESMFold2 structure prediction request.
266///
267/// Chains are added one at a time via the builder methods.
268#[derive(Debug, Clone, Default)]
269pub struct StructurePredictionInput {
270    /// Ordered list of chains included in the prediction.
271    pub sequences: Vec<ChainInput>,
272}
273
274impl StructurePredictionInput {
275    /// Create an empty prediction input.
276    pub fn new() -> Self {
277        Self::default()
278    }
279
280    /// Append any `ChainInput` variant (builder-style, consumes and returns `self`).
281    pub fn add_chain(mut self, chain: impl Into<ChainInput>) -> Self {
282        self.sequences.push(chain.into());
283        self
284    }
285
286    /// Append a protein chain.
287    pub fn add_protein(self, protein: ProteinInput) -> Self {
288        self.add_chain(protein)
289    }
290
291    /// Append a DNA chain.
292    pub fn add_dna(self, dna: DNAInput) -> Self {
293        self.add_chain(dna)
294    }
295
296    /// Append a ligand.
297    pub fn add_ligand(self, ligand: LigandInput) -> Self {
298        self.add_chain(ligand)
299    }
300
301    /// Total number of chains (protein + DNA + ligand).
302    pub fn num_chains(&self) -> usize {
303        self.sequences.len()
304    }
305}
306
307impl fmt::Display for StructurePredictionInput {
308    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
309        write!(f, "StructurePredictionInput(chains=[")?;
310        for (i, chain) in self.sequences.iter().enumerate() {
311            if i > 0 {
312                write!(f, ", ")?;
313            }
314            write!(f, "{chain}")?;
315        }
316        write!(f, "])")
317    }
318}
319
320// ── Validation ───────────────────────────────────────────────────────────────
321
322/// Validate that every character in a protein sequence belongs to the accepted alphabet.
323pub fn validate_protein_sequence(sequence: &str) -> Result<(), String> {
324    if sequence.is_empty() {
325        return Err("Protein sequence must not be empty".to_string());
326    }
327    for (i, ch) in sequence.char_indices() {
328        if !PROTEIN_ALPHABET.contains(ch) {
329            return Err(format!(
330                "Invalid amino acid '{}' at position {} in protein sequence",
331                ch,
332                i + 1
333            ));
334        }
335    }
336    Ok(())
337}
338
339/// Validate that every character in a DNA sequence belongs to the accepted alphabet.
340pub fn validate_dna_sequence(sequence: &str) -> Result<(), String> {
341    if sequence.is_empty() {
342        return Err("DNA sequence must not be empty".to_string());
343    }
344    for (i, ch) in sequence.char_indices() {
345        if !DNA_ALPHABET.contains(ch) {
346            return Err(format!(
347                "Invalid nucleotide '{}' at position {} in DNA sequence",
348                ch,
349                i + 1
350            ));
351        }
352    }
353    Ok(())
354}
355
356// ── Tests ─────────────────────────────────────────────────────────────────────
357
358#[cfg(test)]
359mod tests {
360    use super::*;
361
362    #[test]
363    fn test_protein_input_valid() {
364        let p = ProteinInput::new("A", "ACDEFGHIKLMNPQRSTVWY").unwrap();
365        assert_eq!(p.id, "A");
366        assert_eq!(p.len(), 20);
367        assert!(!p.is_empty());
368    }
369
370    #[test]
371    fn test_protein_input_ambiguous_codes() {
372        // B, J, O, U, X, Z are all acceptable
373        ProteinInput::new("A", "BJOUXZ").expect("ambiguous codes should be valid");
374    }
375
376    #[test]
377    fn test_protein_input_invalid_char() {
378        let err = ProteinInput::new("A", "ACDE1FGHIK").unwrap_err();
379        assert!(err.contains("'1'"));
380    }
381
382    #[test]
383    fn test_protein_input_empty() {
384        assert!(ProteinInput::new("A", "").is_err());
385    }
386
387    #[test]
388    fn test_dna_input_valid() {
389        let d = DNAInput::new("B", "ACGTNRYSW").unwrap();
390        assert_eq!(d.id, "B");
391        assert_eq!(d.len(), 9);
392    }
393
394    #[test]
395    fn test_dna_input_invalid_char() {
396        let err = DNAInput::new("B", "ACGT1").unwrap_err();
397        assert!(err.contains("'1'"));
398    }
399
400    #[test]
401    fn test_dna_input_with_modification() {
402        let d = DNAInput::new("B", "ACGT")
403            .unwrap()
404            .add_modification(Modification::new(2, "5MC"));
405        assert_eq!(d.modifications.len(), 1);
406        assert_eq!(d.modifications[0].ccd, "5MC");
407    }
408
409    #[test]
410    fn test_ligand_from_ccd() {
411        let l = LigandInput::from_ccd("L", "ATP");
412        assert_eq!(l.id, "L");
413        assert_eq!(l.ccd, vec!["ATP".to_string()]);
414        assert_eq!(l.num_components(), 1);
415    }
416
417    #[test]
418    fn test_structure_prediction_input_builder() {
419        let protein = ProteinInput::new("A", "MKTAYIAK").unwrap();
420        let dna = DNAInput::new("B", "ACGTACGT").unwrap();
421        let ligand = LigandInput::from_ccd("L", "ATP");
422
423        let input = StructurePredictionInput::new()
424            .add_protein(protein)
425            .add_dna(dna)
426            .add_ligand(ligand);
427
428        assert_eq!(input.num_chains(), 3);
429        assert_eq!(input.sequences[0].id(), "A");
430        assert_eq!(input.sequences[1].id(), "B");
431        assert_eq!(input.sequences[2].id(), "L");
432    }
433
434    #[test]
435    fn test_display_protein() {
436        let p = ProteinInput::new("A", "MKTAYIAK").unwrap();
437        assert_eq!(format!("{p}"), "Protein[A](len=8)");
438    }
439
440    #[test]
441    fn test_display_dna() {
442        let d = DNAInput::new("B", "ACGT").unwrap();
443        assert_eq!(format!("{d}"), "DNA[B](len=4, mods=0)");
444    }
445
446    #[test]
447    fn test_display_ligand() {
448        let l = LigandInput::from_ccd("L", "ATP");
449        assert_eq!(format!("{l}"), "Ligand[L](ccd=[\"ATP\"])");
450    }
451
452    #[test]
453    fn test_display_structure_prediction_input() {
454        let input = StructurePredictionInput::new()
455            .add_protein(ProteinInput::new("A", "MKTAYIAK").unwrap());
456        let s = format!("{input}");
457        assert!(s.contains("StructurePredictionInput"));
458        assert!(s.contains("Protein[A]"));
459    }
460
461    #[test]
462    fn test_chain_input_id() {
463        let chain: ChainInput = ProteinInput::new("X", "ACDE").unwrap().into();
464        assert_eq!(chain.id(), "X");
465    }
466}
ferritin_plms/esmfold2/input_types.rs

ferritin_plms/esmfold2/
input_types.rs
