480 lines
14 KiB
Rust
480 lines
14 KiB
Rust
//! AlphaFold-Lite protein structure prediction demo.
|
|
//!
|
|
//! This crate implements a lightweight protein structure prediction model
|
|
//! inspired by `AlphaFold`, using geometric deep learning and attention mechanisms.
|
|
//!
|
|
//! # Architecture
|
|
//!
|
|
//! The model consists of:
|
|
//! - **Embedding Layer**: Amino acid embeddings with positional encoding
|
|
//! - **Evoformer**: Attention over sequences and pairs
|
|
//! - **Structure Module**: SE(3)-equivariant attention for 3D coordinates
|
|
//! - **Confidence Head**: pLDDT prediction
|
|
//!
|
|
//! # Example
|
|
//!
|
|
//! ```ignore
|
|
//! use rtx_alphafold_demo::{AlphaFoldLite, predict_structure};
|
|
//!
|
|
//! let sequence = "MKFLILLFNILCLFPVLAADNHGVGPQGAS";
|
|
//! let structure = predict_structure(sequence, Default::default()).await?;
|
|
//! println!("Predicted {} atoms with avg pLDDT: {:.1}",
|
|
//! structure.atom_coords.len(),
|
|
//! structure.model_confidence.avg_plddt);
|
|
//! ```
|
|
|
|
#![allow(missing_docs)] // Demo crate - documentation not required for all items
|
|
|
|
pub mod confidence;
|
|
pub mod encoder;
|
|
pub mod evoformer;
|
|
pub mod pdb_export;
|
|
pub mod sample_data;
|
|
pub mod structure_module;
|
|
|
|
use alphafold_shared::{
|
|
AtomCoord, AtomName, ChainInfo, ConfidenceCategory, ExportFormat, ExportResult,
|
|
ModelConfidence, PredictStructureRequest, PredictionConfig, ProteinStructure,
|
|
SecondaryStructure, validate_sequence,
|
|
};
|
|
use thiserror::Error;
|
|
|
|
/// Errors that can occur during structure prediction.
|
|
#[derive(Debug, Error)]
|
|
pub enum AlphaFoldError {
|
|
/// Invalid amino acid sequence
|
|
#[error("Invalid sequence: {0}")]
|
|
InvalidSequence(#[from] alphafold_shared::ValidationError),
|
|
|
|
/// Model inference error
|
|
#[error("Inference error: {0}")]
|
|
InferenceError(String),
|
|
|
|
/// Export error
|
|
#[error("Export error: {0}")]
|
|
ExportError(String),
|
|
}
|
|
|
|
/// Main entry point for structure prediction.
|
|
pub async fn predict_structure(
|
|
request: PredictStructureRequest,
|
|
) -> Result<ProteinStructure, AlphaFoldError> {
|
|
let sequence = request.sequence.trim().to_uppercase();
|
|
let name = request.name.unwrap_or_else(|| "protein".to_string());
|
|
|
|
// Validate sequence
|
|
let amino_acids = validate_sequence(&sequence)?;
|
|
let num_residues = amino_acids.len();
|
|
|
|
tracing::info!(
|
|
"Predicting structure for {} ({} residues)",
|
|
name,
|
|
num_residues
|
|
);
|
|
|
|
// For demo purposes, generate a plausible structure
|
|
// In production, this would run the actual neural network
|
|
let (atom_coords, plddt_scores) = generate_demo_structure(&sequence, &request.config);
|
|
|
|
// Assign secondary structure based on patterns
|
|
let secondary_structure = assign_secondary_structure(&plddt_scores);
|
|
|
|
// Calculate model confidence
|
|
let avg_plddt = plddt_scores.iter().sum::<f32>() / plddt_scores.len() as f32;
|
|
let model_confidence = ModelConfidence {
|
|
avg_plddt,
|
|
ptm_score: calculate_ptm_score(&plddt_scores),
|
|
iptm_score: None,
|
|
category: ConfidenceCategory::from_plddt(avg_plddt),
|
|
};
|
|
|
|
// Create chain info
|
|
let chains = vec![ChainInfo {
|
|
chain_id: 'A',
|
|
start_residue: 0,
|
|
end_residue: num_residues,
|
|
sequence: sequence.clone(),
|
|
}];
|
|
|
|
Ok(ProteinStructure {
|
|
name,
|
|
sequence,
|
|
num_residues,
|
|
atom_coords,
|
|
plddt_scores,
|
|
pae_matrix: Some(generate_pae_matrix(num_residues)),
|
|
model_confidence,
|
|
secondary_structure,
|
|
chains,
|
|
})
|
|
}
|
|
|
|
/// Export structure to requested format.
|
|
pub fn export_structure(
|
|
structure: &ProteinStructure,
|
|
format: ExportFormat,
|
|
) -> Result<ExportResult, AlphaFoldError> {
|
|
let (data, extension, mime_type) = match format {
|
|
ExportFormat::Pdb => (pdb_export::to_pdb(structure), "pdb", "chemical/x-pdb"),
|
|
ExportFormat::Mmcif => (pdb_export::to_mmcif(structure), "cif", "chemical/x-mmcif"),
|
|
ExportFormat::Json => (
|
|
serde_json::to_string_pretty(structure)
|
|
.map_err(|e| AlphaFoldError::ExportError(e.to_string()))?,
|
|
"json",
|
|
"application/json",
|
|
),
|
|
};
|
|
|
|
let filename = format!(
|
|
"{}_{}.{}",
|
|
structure.name.replace(' ', "_").to_lowercase(),
|
|
chrono_lite_timestamp(),
|
|
extension
|
|
);
|
|
|
|
Ok(ExportResult {
|
|
data,
|
|
filename,
|
|
mime_type: mime_type.to_string(),
|
|
})
|
|
}
|
|
|
|
/// Generate a demo structure with realistic geometry.
|
|
fn generate_demo_structure(
|
|
sequence: &str,
|
|
config: &PredictionConfig,
|
|
) -> (Vec<AtomCoord>, Vec<f32>) {
|
|
use rand::SeedableRng;
|
|
use rand_distr::{Distribution, Normal};
|
|
|
|
let seed = config.seed.unwrap_or(42);
|
|
let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
|
|
|
|
let num_residues = sequence.len();
|
|
let mut atom_coords = Vec::with_capacity(num_residues * 5);
|
|
let mut plddt_scores = Vec::with_capacity(num_residues);
|
|
|
|
// Standard bond lengths (Angstroms)
|
|
const N_CA_BOND: f32 = 1.458;
|
|
const CA_C_BOND: f32 = 1.523;
|
|
const C_N_BOND: f32 = 1.329;
|
|
const C_O_BOND: f32 = 1.231;
|
|
const CA_CB_BOND: f32 = 1.521;
|
|
|
|
// Backbone dihedral angles for secondary structure
|
|
let phi_helix = -57.0_f32.to_radians();
|
|
let psi_helix = -47.0_f32.to_radians();
|
|
let phi_strand = -139.0_f32.to_radians();
|
|
let psi_strand = 135.0_f32.to_radians();
|
|
|
|
// Generate backbone coordinates
|
|
let mut prev_c = [0.0_f32, 0.0, 0.0];
|
|
let mut prev_n = [-C_N_BOND, 0.0, 0.0];
|
|
|
|
let noise = Normal::new(0.0_f32, 0.1).unwrap();
|
|
let plddt_noise = Normal::new(0.0_f32, 5.0).unwrap();
|
|
|
|
for (i, residue) in sequence.chars().enumerate() {
|
|
// Determine secondary structure propensity
|
|
let (phi, psi) = if matches!(residue, 'A' | 'E' | 'L' | 'M' | 'K' | 'R') {
|
|
// Helix-favoring residues
|
|
(phi_helix, psi_helix)
|
|
} else if matches!(residue, 'V' | 'I' | 'Y' | 'F' | 'W' | 'T') {
|
|
// Strand-favoring residues
|
|
(phi_strand, psi_strand)
|
|
} else {
|
|
// Mix
|
|
let t = (i as f32 / num_residues as f32) * std::f32::consts::TAU;
|
|
if t.sin() > 0.0 {
|
|
(phi_helix, psi_helix)
|
|
} else {
|
|
(phi_strand, psi_strand)
|
|
}
|
|
};
|
|
|
|
// Calculate N position
|
|
let n_pos = if i == 0 {
|
|
[0.0, 0.0, 0.0]
|
|
} else {
|
|
let direction = normalize([
|
|
prev_c[0] - prev_n[0],
|
|
prev_c[1] - prev_n[1],
|
|
prev_c[2] - prev_n[2],
|
|
]);
|
|
[
|
|
prev_c[0] + direction[0] * C_N_BOND + noise.sample(&mut rng),
|
|
prev_c[1] + direction[1] * C_N_BOND + noise.sample(&mut rng),
|
|
prev_c[2] + direction[2] * C_N_BOND + noise.sample(&mut rng),
|
|
]
|
|
};
|
|
|
|
// Calculate CA position
|
|
let ca_pos = [
|
|
n_pos[0] + phi.cos() * N_CA_BOND,
|
|
n_pos[1] + phi.sin() * N_CA_BOND,
|
|
n_pos[2] + noise.sample(&mut rng),
|
|
];
|
|
|
|
// Calculate C position
|
|
let c_pos = [
|
|
ca_pos[0] + psi.cos() * CA_C_BOND,
|
|
ca_pos[1] + psi.sin() * CA_C_BOND,
|
|
ca_pos[2] + noise.sample(&mut rng),
|
|
];
|
|
|
|
// Calculate O position (roughly perpendicular to CA-C bond)
|
|
let o_pos = [
|
|
c_pos[0] + C_O_BOND * 0.866,
|
|
c_pos[1] - C_O_BOND * 0.5,
|
|
c_pos[2] + noise.sample(&mut rng),
|
|
];
|
|
|
|
// Calculate CB position (not for glycine)
|
|
let cb_pos = if residue == 'G' {
|
|
None
|
|
} else {
|
|
Some([
|
|
ca_pos[0] - CA_CB_BOND * 0.5,
|
|
ca_pos[1] + CA_CB_BOND * 0.866,
|
|
ca_pos[2] + noise.sample(&mut rng),
|
|
])
|
|
};
|
|
|
|
// Calculate pLDDT (higher in structured regions)
|
|
let base_plddt = if i < 5 || i >= num_residues - 5 {
|
|
// Lower confidence at termini
|
|
60.0
|
|
} else if matches!(residue, 'P' | 'G') {
|
|
// Lower for flexible residues
|
|
65.0
|
|
} else {
|
|
// Higher for structured regions
|
|
85.0
|
|
};
|
|
let plddt = (base_plddt + plddt_noise.sample(&mut rng)).clamp(20.0, 100.0);
|
|
plddt_scores.push(plddt);
|
|
|
|
// Add atoms
|
|
atom_coords.push(AtomCoord {
|
|
residue_idx: i,
|
|
atom_name: AtomName::N,
|
|
x: n_pos[0],
|
|
y: n_pos[1],
|
|
z: n_pos[2],
|
|
b_factor: plddt,
|
|
});
|
|
|
|
atom_coords.push(AtomCoord {
|
|
residue_idx: i,
|
|
atom_name: AtomName::Ca,
|
|
x: ca_pos[0],
|
|
y: ca_pos[1],
|
|
z: ca_pos[2],
|
|
b_factor: plddt,
|
|
});
|
|
|
|
atom_coords.push(AtomCoord {
|
|
residue_idx: i,
|
|
atom_name: AtomName::C,
|
|
x: c_pos[0],
|
|
y: c_pos[1],
|
|
z: c_pos[2],
|
|
b_factor: plddt,
|
|
});
|
|
|
|
atom_coords.push(AtomCoord {
|
|
residue_idx: i,
|
|
atom_name: AtomName::O,
|
|
x: o_pos[0],
|
|
y: o_pos[1],
|
|
z: o_pos[2],
|
|
b_factor: plddt,
|
|
});
|
|
|
|
if let Some(cb) = cb_pos {
|
|
atom_coords.push(AtomCoord {
|
|
residue_idx: i,
|
|
atom_name: AtomName::Cb,
|
|
x: cb[0],
|
|
y: cb[1],
|
|
z: cb[2],
|
|
b_factor: plddt,
|
|
});
|
|
}
|
|
|
|
prev_n = n_pos;
|
|
prev_c = c_pos;
|
|
}
|
|
|
|
(atom_coords, plddt_scores)
|
|
}
|
|
|
|
/// Assign secondary structure based on confidence scores and patterns.
|
|
fn assign_secondary_structure(plddt_scores: &[f32]) -> Vec<SecondaryStructure> {
|
|
let mut ss = Vec::with_capacity(plddt_scores.len());
|
|
|
|
for (i, &plddt) in plddt_scores.iter().enumerate() {
|
|
let structure = if plddt > 85.0 {
|
|
// High confidence suggests ordered structure
|
|
if i % 7 < 4 {
|
|
SecondaryStructure::Helix
|
|
} else {
|
|
SecondaryStructure::Strand
|
|
}
|
|
} else if plddt > 70.0 {
|
|
SecondaryStructure::Turn
|
|
} else {
|
|
SecondaryStructure::Coil
|
|
};
|
|
ss.push(structure);
|
|
}
|
|
|
|
ss
|
|
}
|
|
|
|
/// Calculate predicted TM-score from pLDDT scores.
|
|
fn calculate_ptm_score(plddt_scores: &[f32]) -> f32 {
|
|
if plddt_scores.is_empty() {
|
|
return 0.0;
|
|
}
|
|
|
|
// Approximate pTM from pLDDT (empirical relationship)
|
|
let avg_plddt = plddt_scores.iter().sum::<f32>() / plddt_scores.len() as f32;
|
|
(avg_plddt / 100.0).powf(2.0).clamp(0.0, 1.0)
|
|
}
|
|
|
|
/// Generate PAE matrix (predicted aligned error).
|
|
fn generate_pae_matrix(num_residues: usize) -> Vec<Vec<f32>> {
|
|
use rand::SeedableRng;
|
|
use rand_distr::{Distribution, Normal};
|
|
|
|
let mut rng = rand::rngs::StdRng::seed_from_u64(42);
|
|
let noise = Normal::new(0.0_f32, 2.0).unwrap();
|
|
|
|
let mut matrix = Vec::with_capacity(num_residues);
|
|
|
|
for i in 0..num_residues {
|
|
let mut row = Vec::with_capacity(num_residues);
|
|
for j in 0..num_residues {
|
|
// Lower PAE for nearby residues
|
|
let distance = (i as i32 - j as i32).unsigned_abs() as f32;
|
|
let base_pae = if distance < 5.0 {
|
|
2.0
|
|
} else if distance < 15.0 {
|
|
5.0
|
|
} else {
|
|
10.0
|
|
};
|
|
let pae = (base_pae + noise.sample(&mut rng)).clamp(0.5, 31.0);
|
|
row.push(pae);
|
|
}
|
|
matrix.push(row);
|
|
}
|
|
|
|
matrix
|
|
}
|
|
|
|
/// Normalize a 3D vector.
|
|
fn normalize(v: [f32; 3]) -> [f32; 3] {
|
|
let len = (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt();
|
|
if len < 1e-6 {
|
|
[1.0, 0.0, 0.0]
|
|
} else {
|
|
[v[0] / len, v[1] / len, v[2] / len]
|
|
}
|
|
}
|
|
|
|
/// Generate a simple timestamp without chrono dependency.
|
|
fn chrono_lite_timestamp() -> String {
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
|
let secs = SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap_or_default()
|
|
.as_secs();
|
|
format!("{secs}")
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use alphafold_shared::ModelVariant;
|
|
|
|
#[tokio::test]
|
|
async fn test_predict_structure() {
|
|
let request = PredictStructureRequest {
|
|
sequence: "MKFLILLFNILCLFPVLAADNHGVGPQGAS".to_string(),
|
|
name: Some("test_protein".to_string()),
|
|
config: PredictionConfig::default(),
|
|
};
|
|
|
|
let result = predict_structure(request).await;
|
|
assert!(result.is_ok());
|
|
|
|
let structure = result.unwrap();
|
|
assert_eq!(structure.num_residues, 30);
|
|
assert_eq!(structure.plddt_scores.len(), 30);
|
|
assert!(!structure.atom_coords.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_predict_structure_invalid_sequence() {
|
|
let request = PredictStructureRequest {
|
|
sequence: "INVALID123".to_string(),
|
|
name: None,
|
|
config: PredictionConfig::default(),
|
|
};
|
|
|
|
let result = predict_structure(request).await;
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_export_json() {
|
|
let structure = ProteinStructure {
|
|
name: "test".to_string(),
|
|
sequence: "AAA".to_string(),
|
|
num_residues: 3,
|
|
atom_coords: vec![],
|
|
plddt_scores: vec![90.0, 85.0, 80.0],
|
|
pae_matrix: None,
|
|
model_confidence: ModelConfidence {
|
|
avg_plddt: 85.0,
|
|
ptm_score: 0.72,
|
|
iptm_score: None,
|
|
category: ConfidenceCategory::High,
|
|
},
|
|
secondary_structure: vec![
|
|
SecondaryStructure::Helix,
|
|
SecondaryStructure::Helix,
|
|
SecondaryStructure::Coil,
|
|
],
|
|
chains: vec![],
|
|
};
|
|
|
|
let result = export_structure(&structure, ExportFormat::Json);
|
|
assert!(result.is_ok());
|
|
let export = result.unwrap();
|
|
assert!(export.data.contains("\"name\": \"test\""));
|
|
assert!(export.filename.ends_with(".json"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_normalize() {
|
|
let v = [3.0, 4.0, 0.0];
|
|
let n = normalize(v);
|
|
let len = (n[0] * n[0] + n[1] * n[1] + n[2] * n[2]).sqrt();
|
|
assert!((len - 1.0).abs() < 0.001);
|
|
}
|
|
|
|
#[test]
|
|
fn test_ptm_score_calculation() {
|
|
let high_plddt = vec![95.0; 100];
|
|
let ptm = calculate_ptm_score(&high_plddt);
|
|
assert!(ptm > 0.8);
|
|
|
|
let low_plddt = vec![40.0; 100];
|
|
let ptm = calculate_ptm_score(&low_plddt);
|
|
assert!(ptm < 0.2);
|
|
}
|
|
}
|