Files
rustytorch/demos/alphafold-shared/src/lib.rs
T
2026-03-04 00:08:42 +00:00

648 lines
20 KiB
Rust

//! Shared IPC types for AlphaFold-Lite protein structure prediction demo.
//!
//! This crate provides the data structures used for communication between
//! the Tauri frontend and the Rust backend for protein structure prediction.
use serde::{Deserialize, Serialize};
// ============================================================================
// Input Types
// ============================================================================
/// Request to predict protein structure from amino acid sequence.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PredictStructureRequest {
/// Amino acid sequence (single-letter codes: A, C, D, E, F, G, H, I, K, L, M, N, P, Q, R, S, T, V, W, Y)
pub sequence: String,
/// Optional name/identifier for the protein
pub name: Option<String>,
/// Model configuration
pub config: PredictionConfig,
}
/// Configuration for structure prediction.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PredictionConfig {
/// Number of recycles (refinement iterations)
pub num_recycles: u32,
/// Whether to use templates if available
pub use_templates: bool,
/// Number of structure samples to generate
pub num_samples: u32,
/// Random seed for reproducibility
pub seed: Option<u64>,
/// Model variant to use
pub model_variant: ModelVariant,
}
impl Default for PredictionConfig {
fn default() -> Self {
Self {
num_recycles: 3,
use_templates: false,
num_samples: 1,
seed: None,
model_variant: ModelVariant::Lite,
}
}
}
/// Model variant for structure prediction.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelVariant {
/// Lightweight model for fast predictions
Lite,
/// Standard model with balanced speed/accuracy
Standard,
/// High-accuracy model (slower)
Accurate,
}
// ============================================================================
// Output Types
// ============================================================================
/// Predicted protein structure.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProteinStructure {
/// Protein name/identifier
pub name: String,
/// Input amino acid sequence
pub sequence: String,
/// Number of residues
pub num_residues: usize,
/// 3D coordinates for each atom (N, CA, C, O, CB for each residue)
pub atom_coords: Vec<AtomCoord>,
/// Per-residue confidence scores (pLDDT: 0-100)
pub plddt_scores: Vec<f32>,
/// Predicted aligned error (PAE) matrix
pub pae_matrix: Option<Vec<Vec<f32>>>,
/// Overall model confidence
pub model_confidence: ModelConfidence,
/// Secondary structure assignment
pub secondary_structure: Vec<SecondaryStructure>,
/// Chain information
pub chains: Vec<ChainInfo>,
}
/// 3D coordinates for a single atom.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct AtomCoord {
/// Residue index (0-based)
pub residue_idx: usize,
/// Atom name (N, CA, C, O, CB, etc.)
pub atom_name: AtomName,
/// X coordinate (Angstroms)
pub x: f32,
/// Y coordinate (Angstroms)
pub y: f32,
/// Z coordinate (Angstroms)
pub z: f32,
/// B-factor (temperature factor / confidence)
pub b_factor: f32,
}
/// Standard backbone and CB atom names.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "UPPERCASE")]
pub enum AtomName {
/// Backbone nitrogen
N,
/// Alpha carbon
Ca,
/// Backbone carbonyl carbon
C,
/// Backbone oxygen
O,
/// Beta carbon (not present in glycine)
Cb,
}
impl AtomName {
/// Get the standard PDB atom name string.
pub fn as_pdb_str(&self) -> &'static str {
match self {
AtomName::N => "N",
AtomName::Ca => "CA",
AtomName::C => "C",
AtomName::O => "O",
AtomName::Cb => "CB",
}
}
}
/// Overall model confidence metrics.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelConfidence {
/// Average pLDDT score (0-100)
pub avg_plddt: f32,
/// Predicted template modeling score (pTM: 0-1)
pub ptm_score: f32,
/// Interface predicted template modeling score (ipTM: 0-1, for multimers)
pub iptm_score: Option<f32>,
/// Confidence category
pub category: ConfidenceCategory,
}
/// Confidence category based on pLDDT scores.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ConfidenceCategory {
/// Very high confidence (pLDDT > 90)
VeryHigh,
/// High confidence (70 < pLDDT <= 90)
High,
/// Low confidence (50 < pLDDT <= 70)
Low,
/// Very low confidence (pLDDT <= 50)
VeryLow,
}
impl ConfidenceCategory {
/// Determine category from average pLDDT score.
pub fn from_plddt(plddt: f32) -> Self {
if plddt > 90.0 {
ConfidenceCategory::VeryHigh
} else if plddt > 70.0 {
ConfidenceCategory::High
} else if plddt > 50.0 {
ConfidenceCategory::Low
} else {
ConfidenceCategory::VeryLow
}
}
/// Get the color for visualization (RGB hex).
pub fn color(&self) -> &'static str {
match self {
ConfidenceCategory::VeryHigh => "#0053D6", // Blue
ConfidenceCategory::High => "#65CBF3", // Cyan
ConfidenceCategory::Low => "#FFDB13", // Yellow
ConfidenceCategory::VeryLow => "#FF7D45", // Orange
}
}
}
/// Secondary structure assignment for a residue.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SecondaryStructure {
/// Alpha helix (H)
Helix,
/// Beta strand/sheet (E)
Strand,
/// Coil/loop (C)
Coil,
/// Turn (T)
Turn,
}
impl SecondaryStructure {
/// Get the DSSP single-letter code.
pub fn dssp_code(&self) -> char {
match self {
SecondaryStructure::Helix => 'H',
SecondaryStructure::Strand => 'E',
SecondaryStructure::Coil => 'C',
SecondaryStructure::Turn => 'T',
}
}
}
/// Information about a protein chain.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChainInfo {
/// Chain identifier (A, B, C, etc.)
pub chain_id: char,
/// Starting residue index (0-based)
pub start_residue: usize,
/// Ending residue index (exclusive)
pub end_residue: usize,
/// Sequence for this chain
pub sequence: String,
}
// ============================================================================
// Amino Acid Types
// ============================================================================
/// Standard amino acid.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AminoAcid {
Ala, // A - Alanine
Cys, // C - Cysteine
Asp, // D - Aspartic acid
Glu, // E - Glutamic acid
Phe, // F - Phenylalanine
Gly, // G - Glycine
His, // H - Histidine
Ile, // I - Isoleucine
Lys, // K - Lysine
Leu, // L - Leucine
Met, // M - Methionine
Asn, // N - Asparagine
Pro, // P - Proline
Gln, // Q - Glutamine
Arg, // R - Arginine
Ser, // S - Serine
Thr, // T - Threonine
Val, // V - Valine
Trp, // W - Tryptophan
Tyr, // Y - Tyrosine
Unk, // X - Unknown
}
impl AminoAcid {
/// Parse from single-letter code.
pub fn from_code(c: char) -> Option<Self> {
match c.to_ascii_uppercase() {
'A' => Some(AminoAcid::Ala),
'C' => Some(AminoAcid::Cys),
'D' => Some(AminoAcid::Asp),
'E' => Some(AminoAcid::Glu),
'F' => Some(AminoAcid::Phe),
'G' => Some(AminoAcid::Gly),
'H' => Some(AminoAcid::His),
'I' => Some(AminoAcid::Ile),
'K' => Some(AminoAcid::Lys),
'L' => Some(AminoAcid::Leu),
'M' => Some(AminoAcid::Met),
'N' => Some(AminoAcid::Asn),
'P' => Some(AminoAcid::Pro),
'Q' => Some(AminoAcid::Gln),
'R' => Some(AminoAcid::Arg),
'S' => Some(AminoAcid::Ser),
'T' => Some(AminoAcid::Thr),
'V' => Some(AminoAcid::Val),
'W' => Some(AminoAcid::Trp),
'Y' => Some(AminoAcid::Tyr),
'X' => Some(AminoAcid::Unk),
_ => None,
}
}
/// Get single-letter code.
pub fn code(&self) -> char {
match self {
AminoAcid::Ala => 'A',
AminoAcid::Cys => 'C',
AminoAcid::Asp => 'D',
AminoAcid::Glu => 'E',
AminoAcid::Phe => 'F',
AminoAcid::Gly => 'G',
AminoAcid::His => 'H',
AminoAcid::Ile => 'I',
AminoAcid::Lys => 'K',
AminoAcid::Leu => 'L',
AminoAcid::Met => 'M',
AminoAcid::Asn => 'N',
AminoAcid::Pro => 'P',
AminoAcid::Gln => 'Q',
AminoAcid::Arg => 'R',
AminoAcid::Ser => 'S',
AminoAcid::Thr => 'T',
AminoAcid::Val => 'V',
AminoAcid::Trp => 'W',
AminoAcid::Tyr => 'Y',
AminoAcid::Unk => 'X',
}
}
/// Get three-letter code.
pub fn code3(&self) -> &'static str {
match self {
AminoAcid::Ala => "ALA",
AminoAcid::Cys => "CYS",
AminoAcid::Asp => "ASP",
AminoAcid::Glu => "GLU",
AminoAcid::Phe => "PHE",
AminoAcid::Gly => "GLY",
AminoAcid::His => "HIS",
AminoAcid::Ile => "ILE",
AminoAcid::Lys => "LYS",
AminoAcid::Leu => "LEU",
AminoAcid::Met => "MET",
AminoAcid::Asn => "ASN",
AminoAcid::Pro => "PRO",
AminoAcid::Gln => "GLN",
AminoAcid::Arg => "ARG",
AminoAcid::Ser => "SER",
AminoAcid::Thr => "THR",
AminoAcid::Val => "VAL",
AminoAcid::Trp => "TRP",
AminoAcid::Tyr => "TYR",
AminoAcid::Unk => "UNK",
}
}
/// Get embedding index (0-20).
pub fn embedding_idx(&self) -> usize {
match self {
AminoAcid::Ala => 0,
AminoAcid::Cys => 1,
AminoAcid::Asp => 2,
AminoAcid::Glu => 3,
AminoAcid::Phe => 4,
AminoAcid::Gly => 5,
AminoAcid::His => 6,
AminoAcid::Ile => 7,
AminoAcid::Lys => 8,
AminoAcid::Leu => 9,
AminoAcid::Met => 10,
AminoAcid::Asn => 11,
AminoAcid::Pro => 12,
AminoAcid::Gln => 13,
AminoAcid::Arg => 14,
AminoAcid::Ser => 15,
AminoAcid::Thr => 16,
AminoAcid::Val => 17,
AminoAcid::Trp => 18,
AminoAcid::Tyr => 19,
AminoAcid::Unk => 20,
}
}
}
// ============================================================================
// Export Types
// ============================================================================
/// Request to export structure to a file format.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExportRequest {
/// Structure to export
pub structure: ProteinStructure,
/// Output format
pub format: ExportFormat,
}
/// Supported export formats.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ExportFormat {
/// Protein Data Bank format
Pdb,
/// MacroMolecular Crystallographic Information File
Mmcif,
/// JSON format
Json,
}
/// Export result containing the formatted data.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExportResult {
/// Formatted structure data
pub data: String,
/// Suggested filename
pub filename: String,
/// MIME type
pub mime_type: String,
}
// ============================================================================
// Validation
// ============================================================================
/// Validate an amino acid sequence.
pub fn validate_sequence(sequence: &str) -> Result<Vec<AminoAcid>, ValidationError> {
if sequence.is_empty() {
return Err(ValidationError::EmptySequence);
}
if sequence.len() > 2500 {
return Err(ValidationError::SequenceTooLong {
length: sequence.len(),
max: 2500,
});
}
let mut amino_acids = Vec::with_capacity(sequence.len());
for (i, c) in sequence.chars().enumerate() {
match AminoAcid::from_code(c) {
Some(aa) => amino_acids.push(aa),
None => {
return Err(ValidationError::InvalidResidue {
position: i,
char: c,
});
}
}
}
Ok(amino_acids)
}
/// Sequence validation error.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ValidationError {
/// Empty sequence provided
EmptySequence,
/// Sequence exceeds maximum length
SequenceTooLong { length: usize, max: usize },
/// Invalid residue character
InvalidResidue { position: usize, char: char },
}
impl std::fmt::Display for ValidationError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ValidationError::EmptySequence => write!(f, "Sequence cannot be empty"),
ValidationError::SequenceTooLong { length, max } => {
write!(f, "Sequence too long: {} residues (max {})", length, max)
}
ValidationError::InvalidResidue { position, char } => {
write!(f, "Invalid residue '{}' at position {}", char, position)
}
}
}
}
impl std::error::Error for ValidationError {}
// ============================================================================
// Sample Proteins
// ============================================================================
/// Sample protein for demonstration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SampleProtein {
/// Protein name
pub name: String,
/// UniProt ID or PDB ID
pub id: String,
/// Amino acid sequence
pub sequence: String,
/// Description
pub description: String,
/// Organism
pub organism: String,
/// Sequence length
pub length: usize,
}
/// Get list of sample proteins for the demo.
pub fn get_sample_proteins() -> Vec<SampleProtein> {
vec![
SampleProtein {
name: "Insulin".to_string(),
id: "P01308".to_string(),
sequence: "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKT".to_string(),
description: "Insulin precursor - regulates glucose metabolism".to_string(),
organism: "Homo sapiens".to_string(),
length: 54,
},
SampleProtein {
name: "Green Fluorescent Protein (GFP)".to_string(),
id: "1EMA".to_string(),
sequence: "MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK".to_string(),
description: "Fluorescent protein used as a reporter in biology".to_string(),
organism: "Aequorea victoria".to_string(),
length: 238,
},
SampleProtein {
name: "Ubiquitin".to_string(),
id: "P0CG48".to_string(),
sequence: "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG".to_string(),
description: "Highly conserved protein involved in protein degradation".to_string(),
organism: "Homo sapiens".to_string(),
length: 76,
},
SampleProtein {
name: "Lysozyme".to_string(),
id: "P00698".to_string(),
sequence: "KVFGRCELAAAMKRHGLDNYRGYSLGNWVCAAKFESNFNTQATNRNTDGSTDYGILQINSRWWCNDGRTPGSRNLCNIPCSALLSSDITASVNCAKKIVSDGNGMNAWVAWRNRCKGTDVQAWIRGCRL".to_string(),
description: "Enzyme that damages bacterial cell walls".to_string(),
organism: "Gallus gallus".to_string(),
length: 129,
},
SampleProtein {
name: "Myoglobin".to_string(),
id: "P02144".to_string(),
sequence: "MGLSDGEWQLVLNVWGKVEADIPGHGQEVLIRLFKGHPETLEKFDKFKHLKSEDEMKASEDLKKHGATVLTALGGILKKKGHHEAEIKPLAQSHATKHKIPVKYLEFISECIIQVLQSKHPGDFGADAQGAMNKALELFRKDMASNYKELGFQG".to_string(),
description: "Oxygen-binding protein in muscle tissue".to_string(),
organism: "Homo sapiens".to_string(),
length: 154,
},
]
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_amino_acid_from_code() {
assert_eq!(AminoAcid::from_code('A'), Some(AminoAcid::Ala));
assert_eq!(AminoAcid::from_code('a'), Some(AminoAcid::Ala));
assert_eq!(AminoAcid::from_code('Z'), None);
}
#[test]
fn test_amino_acid_codes() {
let aa = AminoAcid::Met;
assert_eq!(aa.code(), 'M');
assert_eq!(aa.code3(), "MET");
assert_eq!(aa.embedding_idx(), 10);
}
#[test]
fn test_validate_sequence_valid() {
let result = validate_sequence("ACDEFGHIKLMNPQRSTVWY");
assert!(result.is_ok());
assert_eq!(result.unwrap().len(), 20);
}
#[test]
fn test_validate_sequence_empty() {
let result = validate_sequence("");
assert!(matches!(result, Err(ValidationError::EmptySequence)));
}
#[test]
fn test_validate_sequence_invalid_char() {
let result = validate_sequence("ACBDE");
assert!(matches!(
result,
Err(ValidationError::InvalidResidue {
position: 2,
char: 'B'
})
));
}
#[test]
fn test_confidence_category() {
assert_eq!(
ConfidenceCategory::from_plddt(95.0),
ConfidenceCategory::VeryHigh
);
assert_eq!(
ConfidenceCategory::from_plddt(80.0),
ConfidenceCategory::High
);
assert_eq!(
ConfidenceCategory::from_plddt(60.0),
ConfidenceCategory::Low
);
assert_eq!(
ConfidenceCategory::from_plddt(40.0),
ConfidenceCategory::VeryLow
);
}
#[test]
fn test_prediction_config_default() {
let config = PredictionConfig::default();
assert_eq!(config.num_recycles, 3);
assert!(!config.use_templates);
assert_eq!(config.num_samples, 1);
}
#[test]
fn test_sample_proteins() {
let samples = get_sample_proteins();
assert!(!samples.is_empty());
for sample in &samples {
assert_eq!(sample.sequence.len(), sample.length);
assert!(validate_sequence(&sample.sequence).is_ok());
}
}
#[test]
fn test_atom_name_pdb_str() {
assert_eq!(AtomName::N.as_pdb_str(), "N");
assert_eq!(AtomName::Ca.as_pdb_str(), "CA");
assert_eq!(AtomName::C.as_pdb_str(), "C");
assert_eq!(AtomName::O.as_pdb_str(), "O");
assert_eq!(AtomName::Cb.as_pdb_str(), "CB");
}
#[test]
fn test_secondary_structure_dssp() {
assert_eq!(SecondaryStructure::Helix.dssp_code(), 'H');
assert_eq!(SecondaryStructure::Strand.dssp_code(), 'E');
assert_eq!(SecondaryStructure::Coil.dssp_code(), 'C');
assert_eq!(SecondaryStructure::Turn.dssp_code(), 'T');
}
#[test]
fn test_export_format_serialization() {
let format = ExportFormat::Pdb;
let json = serde_json::to_string(&format).unwrap();
assert_eq!(json, "\"pdb\"");
}
#[test]
fn test_model_variant_serialization() {
let variant = ModelVariant::Accurate;
let json = serde_json::to_string(&variant).unwrap();
assert_eq!(json, "\"accurate\"");
}
}