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

798 lines
22 KiB
Rust

//! Shared IPC types for DrugBinder drug-target binding affinity demo.
//!
//! This crate provides data structures for communication between
//! the Tauri frontend and Rust backend for drug-target binding prediction.
use serde::{Deserialize, Serialize};
// ============================================================================
// Molecule Types
// ============================================================================
/// A small molecule drug candidate.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Molecule {
/// Unique identifier
pub id: String,
/// Common name
pub name: Option<String>,
/// SMILES representation
pub smiles: String,
/// Molecular weight (g/mol)
pub molecular_weight: f32,
/// Number of atoms
pub num_atoms: usize,
/// Number of bonds
pub num_bonds: usize,
/// Atom information
pub atoms: Vec<Atom>,
/// Bond information
pub bonds: Vec<Bond>,
/// 3D coordinates (if available)
pub coordinates_3d: Option<Vec<Coordinate3D>>,
/// Molecular properties
pub properties: MolecularProperties,
}
/// Atom in a molecule.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Atom {
/// Atom index
pub index: usize,
/// Element symbol
pub element: Element,
/// Formal charge
pub formal_charge: i8,
/// Number of hydrogens
pub num_hydrogens: u8,
/// Is aromatic
pub is_aromatic: bool,
/// Hybridization
pub hybridization: Hybridization,
}
/// Chemical element.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "UPPERCASE")]
pub enum Element {
H,
C,
N,
O,
F,
P,
S,
Cl,
Br,
I,
Na,
Mg,
K,
Ca,
Fe,
Zn,
Cu,
Other,
}
impl Element {
/// Get atomic number.
pub fn atomic_number(&self) -> u8 {
match self {
Element::H => 1,
Element::C => 6,
Element::N => 7,
Element::O => 8,
Element::F => 9,
Element::P => 15,
Element::S => 16,
Element::Cl => 17,
Element::Br => 35,
Element::I => 53,
Element::Na => 11,
Element::Mg => 12,
Element::K => 19,
Element::Ca => 20,
Element::Fe => 26,
Element::Zn => 30,
Element::Cu => 29,
Element::Other => 0,
}
}
}
/// Atom hybridization state.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Hybridization {
Sp,
Sp2,
Sp3,
Sp3d,
Sp3d2,
Other,
}
/// Bond in a molecule.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Bond {
/// First atom index
pub atom1: usize,
/// Second atom index
pub atom2: usize,
/// Bond type
pub bond_type: BondType,
/// Is conjugated
pub is_conjugated: bool,
/// Is in ring
pub is_in_ring: bool,
}
/// Bond type.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum BondType {
Single,
Double,
Triple,
Aromatic,
Other,
}
/// 3D coordinate.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct Coordinate3D {
pub x: f32,
pub y: f32,
pub z: f32,
}
/// Molecular properties (Lipinski's Rule of Five, etc.).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MolecularProperties {
/// Octanol-water partition coefficient
pub log_p: f32,
/// Number of hydrogen bond donors
pub hbd: u8,
/// Number of hydrogen bond acceptors
pub hba: u8,
/// Topological polar surface area (Ų)
pub tpsa: f32,
/// Number of rotatable bonds
pub rotatable_bonds: u8,
/// Number of rings
pub num_rings: u8,
/// Number of aromatic rings
pub num_aromatic_rings: u8,
/// Passes Lipinski's Rule of Five
pub lipinski_pass: bool,
/// QED (Quantitative Estimate of Drug-likeness)
pub qed: f32,
}
// ============================================================================
// Protein Types
// ============================================================================
/// Protein target for binding.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProteinTarget {
/// Unique identifier (e.g., UniProt ID)
pub id: String,
/// Protein name
pub name: String,
/// Organism
pub organism: String,
/// Sequence
pub sequence: String,
/// PDB ID (if structure available)
pub pdb_id: Option<String>,
/// Binding pockets
pub pockets: Vec<BindingPocket>,
}
/// Binding pocket on a protein.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BindingPocket {
/// Pocket ID
pub id: usize,
/// Pocket name
pub name: String,
/// Residue indices
pub residues: Vec<usize>,
/// Pocket center
pub center: Coordinate3D,
/// Pocket volume (ų)
pub volume: f32,
/// Druggability score (0-1)
pub druggability: f32,
}
// ============================================================================
// Binding Prediction Types
// ============================================================================
/// Request to predict binding affinity.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BindingPredictionRequest {
/// Molecule to evaluate
pub molecule: MoleculeInput,
/// Target protein
pub target: TargetInput,
/// Configuration
pub config: PredictionConfig,
}
/// Molecule input (can be SMILES or full molecule).
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum MoleculeInput {
/// SMILES string
Smiles {
smiles: String,
name: Option<String>,
},
/// Full molecule object
Full(Molecule),
}
/// Target input.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum TargetInput {
/// Protein ID (look up from database)
Id { protein_id: String },
/// Protein sequence
Sequence {
sequence: String,
name: Option<String>,
},
/// Full protein object
Full(ProteinTarget),
}
/// Prediction configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PredictionConfig {
/// Model to use
pub model: BindingModel,
/// Predict binding pose
pub predict_pose: bool,
/// Number of poses to generate
pub num_poses: usize,
/// Pocket ID to dock into (if known)
pub pocket_id: Option<usize>,
}
impl Default for PredictionConfig {
fn default() -> Self {
Self {
model: BindingModel::GraphTransformer,
predict_pose: true,
num_poses: 5,
pocket_id: None,
}
}
}
/// Binding affinity model.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BindingModel {
/// Graph neural network
GraphNN,
/// Transformer-based
GraphTransformer,
/// 3D CNN
Cnn3D,
/// Ensemble
Ensemble,
}
/// Binding prediction result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BindingPrediction {
/// Predicted binding affinity (pKd or pIC50)
pub affinity: AffinityPrediction,
/// Binding poses
pub poses: Vec<BindingPose>,
/// Interaction analysis
pub interactions: Vec<ProteinLigandInteraction>,
/// Confidence score
pub confidence: f32,
/// Model explanation
pub explanation: Option<PredictionExplanation>,
}
/// Affinity prediction with uncertainty.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AffinityPrediction {
/// Predicted pKd (negative log of Kd in molar)
pub pkd: f32,
/// Predicted pIC50
pub pic50: f32,
/// Predicted ΔG (kcal/mol)
pub delta_g: f32,
/// Uncertainty (standard deviation)
pub uncertainty: f32,
/// Affinity class
pub affinity_class: AffinityClass,
}
/// Affinity classification.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AffinityClass {
/// Very high affinity (pKd > 9)
VeryHigh,
/// High affinity (pKd 7-9)
High,
/// Medium affinity (pKd 5-7)
Medium,
/// Low affinity (pKd 3-5)
Low,
/// Very low affinity (pKd < 3)
VeryLow,
}
impl AffinityClass {
/// Get affinity class from pKd value.
pub fn from_pkd(pkd: f32) -> Self {
if pkd > 9.0 {
AffinityClass::VeryHigh
} else if pkd > 7.0 {
AffinityClass::High
} else if pkd > 5.0 {
AffinityClass::Medium
} else if pkd > 3.0 {
AffinityClass::Low
} else {
AffinityClass::VeryLow
}
}
}
/// Predicted binding pose.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BindingPose {
/// Pose ID
pub id: usize,
/// Ligand coordinates in binding pose
pub ligand_coords: Vec<Coordinate3D>,
/// Docking score
pub score: f32,
/// RMSD from input (if applicable)
pub rmsd: Option<f32>,
/// Energy breakdown
pub energy_terms: EnergyTerms,
}
/// Energy terms for binding.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnergyTerms {
/// van der Waals
pub vdw: f32,
/// Electrostatics
pub electrostatic: f32,
/// Hydrogen bonding
pub hbond: f32,
/// Hydrophobic
pub hydrophobic: f32,
/// Desolvation
pub desolvation: f32,
/// Entropy penalty
pub entropy: f32,
/// Total
pub total: f32,
}
/// Protein-ligand interaction.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProteinLigandInteraction {
/// Interaction type
pub interaction_type: InteractionType,
/// Protein residue
pub residue: String,
/// Ligand atom index
pub ligand_atom: usize,
/// Distance (Å)
pub distance: f32,
/// Strength (0-1)
pub strength: f32,
}
/// Type of interaction.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum InteractionType {
HydrogenBond,
HydrophobicContact,
PiStacking,
PiCation,
SaltBridge,
HalogenBond,
MetalCoordination,
}
/// Prediction explanation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PredictionExplanation {
/// Important atoms (by contribution)
pub important_atoms: Vec<AtomContribution>,
/// Important residues
pub important_residues: Vec<ResidueContribution>,
/// Feature importance
pub feature_importance: Vec<FeatureImportance>,
}
/// Atom contribution to prediction.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AtomContribution {
pub atom_index: usize,
pub contribution: f32,
}
/// Residue contribution to prediction.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResidueContribution {
pub residue: String,
pub contribution: f32,
}
/// Feature importance.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeatureImportance {
pub feature_name: String,
pub importance: f32,
}
// ============================================================================
// Virtual Screening Types
// ============================================================================
/// Virtual screening request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScreeningRequest {
/// Molecules to screen
pub molecules: Vec<MoleculeInput>,
/// Target
pub target: TargetInput,
/// Configuration
pub config: ScreeningConfig,
}
/// Screening configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScreeningConfig {
/// Number of top compounds to return
pub top_k: usize,
/// Minimum affinity threshold (pKd)
pub min_affinity: Option<f32>,
/// Filter by Lipinski's rules
pub lipinski_filter: bool,
/// Sort by
pub sort_by: ScreeningSortBy,
}
impl Default for ScreeningConfig {
fn default() -> Self {
Self {
top_k: 100,
min_affinity: Some(5.0),
lipinski_filter: true,
sort_by: ScreeningSortBy::Affinity,
}
}
}
/// Screening sort criteria.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ScreeningSortBy {
Affinity,
Confidence,
QED,
LipinskiScore,
}
/// Screening result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScreeningResult {
/// Screened compounds with predictions
pub hits: Vec<ScreeningHit>,
/// Total molecules screened
pub total_screened: usize,
/// Molecules passing filters
pub num_passed_filters: usize,
/// Processing time (ms)
pub processing_time_ms: u64,
}
/// Single screening hit.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScreeningHit {
/// Rank
pub rank: usize,
/// Molecule
pub molecule: Molecule,
/// Binding prediction
pub prediction: BindingPrediction,
}
// ============================================================================
// Molecule Generation Types
// ============================================================================
/// Request to generate molecules.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerationRequest {
/// Target protein
pub target: TargetInput,
/// Generation method
pub method: GenerationMethod,
/// Configuration
pub config: GenerationConfig,
}
/// Generation method.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum GenerationMethod {
/// Diffusion-based generation
Diffusion,
/// Reinforcement learning
ReinforcementLearning,
/// VAE-based
VariationalAutoencoder,
/// Fragment-based
FragmentGrowing,
}
/// Generation configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerationConfig {
/// Number of molecules to generate
pub num_molecules: usize,
/// Target affinity (pKd)
pub target_affinity: Option<f32>,
/// Scaffold to use
pub scaffold: Option<String>,
/// Enforce Lipinski's rules
pub lipinski_constraints: bool,
/// Maximum molecular weight
pub max_mw: f32,
}
impl Default for GenerationConfig {
fn default() -> Self {
Self {
num_molecules: 100,
target_affinity: Some(8.0),
scaffold: None,
lipinski_constraints: true,
max_mw: 500.0,
}
}
}
/// Generation result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerationResult {
/// Generated molecules
pub molecules: Vec<GeneratedMolecule>,
/// Generation statistics
pub stats: GenerationStats,
}
/// Generated molecule with metadata.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GeneratedMolecule {
/// The molecule
pub molecule: Molecule,
/// Predicted binding
pub predicted_binding: AffinityPrediction,
/// Novelty score (vs training set)
pub novelty: f32,
/// Synthetic accessibility score (1-10)
pub sa_score: f32,
}
/// Generation statistics.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerationStats {
/// Total generated
pub total_generated: usize,
/// Valid molecules
pub valid: usize,
/// Unique molecules
pub unique: usize,
/// Novel molecules
pub novel: usize,
/// High affinity (pKd > 7)
pub high_affinity: usize,
}
// ============================================================================
// Sample Data
// ============================================================================
/// Sample molecules for demo.
pub fn get_sample_molecules() -> Vec<Molecule> {
vec![
create_sample_molecule("aspirin", "Aspirin", "CC(=O)OC1=CC=CC=C1C(=O)O", 180.16),
create_sample_molecule(
"ibuprofen",
"Ibuprofen",
"CC(C)CC1=CC=C(C=C1)C(C)C(=O)O",
206.28,
),
create_sample_molecule(
"caffeine",
"Caffeine",
"CN1C=NC2=C1C(=O)N(C(=O)N2C)C",
194.19,
),
create_sample_molecule(
"acetaminophen",
"Acetaminophen",
"CC(=O)NC1=CC=C(C=C1)O",
151.16,
),
]
}
fn create_sample_molecule(id: &str, name: &str, smiles: &str, mw: f32) -> Molecule {
Molecule {
id: id.to_string(),
name: Some(name.to_string()),
smiles: smiles.to_string(),
molecular_weight: mw,
num_atoms: count_atoms(smiles),
num_bonds: count_atoms(smiles).saturating_sub(1),
atoms: vec![],
bonds: vec![],
coordinates_3d: None,
properties: MolecularProperties {
log_p: 1.5,
hbd: 1,
hba: 3,
tpsa: 60.0,
rotatable_bonds: 2,
num_rings: 1,
num_aromatic_rings: 1,
lipinski_pass: true,
qed: 0.7,
},
}
}
fn count_atoms(smiles: &str) -> usize {
// Simplified atom counting
smiles.chars().filter(|c| c.is_uppercase()).count()
}
/// Sample protein targets.
pub fn get_sample_targets() -> Vec<ProteinTarget> {
vec![
ProteinTarget {
id: "P00918".to_string(),
name: "Carbonic anhydrase 2".to_string(),
organism: "Homo sapiens".to_string(),
sequence: "MSHHWGYGKHNGPEHWHKDFPIAKGERQSPVDIDTHTAKYDPSLKPLSVSYDQATSLRILNNGHAFNVEFDDSQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVLGIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPPLLECVTWIVLKEPISVSSEQVLKFRKLNFNGEGEPEELMVDNWRPAQPLKNRQIKASFK".to_string(),
pdb_id: Some("1CA2".to_string()),
pockets: vec![
BindingPocket {
id: 0,
name: "Active site".to_string(),
residues: vec![91, 92, 94, 96, 119, 143, 198, 199, 200],
center: Coordinate3D { x: 12.5, y: 8.3, z: 15.2 },
volume: 350.0,
druggability: 0.85,
},
],
},
ProteinTarget {
id: "P00533".to_string(),
name: "Epidermal growth factor receptor".to_string(),
organism: "Homo sapiens".to_string(),
sequence: "MRPSGTAGAALLALLAALCPASRALEEKKVCQGTSNKLTQLGTFEDHFLSLQRMFNNCEVVLGNLEITYVQRNYDLSFLKTIQEVAGYVLIALNTVERIPLENLQIIRGNMYYENSYALAVLSNYDANKTGLKELPMRNLQEILHGAVRFSNNPALCNVESIQWRDIVSSDFLSNMSMDFQNHLGSCQKCDPSCPNGSCWGAGEENCQKLTKIICAQQCSGRCRGKSPSDCCHNQCAAGCTGPRESDCLVCRKFRDEATCKDTCPPLMLYNPTTYQMDVNPEGKYSFGATCVKKCPRNYVVTDHGSCVRACGADSYEMEEDGVRKCKKCEGPCRKVCNGIGIGEFKDSLSINATNIKHFKNCTSISGDLHILPVAFRGDSFTHTPPLDPQELDILKTVKEITGFLLIQAWPENRTDLHAFENLEIIRGRTKQHGQFSLAVVSLNITSLGLRSLKEISDGDVIISGNKNLCYANTINWKKLFGTSGQKTKIISNRGENSCKATGQVCHALCSPEGCWGPEPRDCVSCRNVSRGRECVDKCNLLEGEPREFVENSECIQCHPECLPQAMNITCTGRGPDNCIQCAHYIDGPHCVKTCPAGVMGENNTLVWKYADAGHVCHLCHPNCTYGCTGPGLEGCPTNGPKIPS".to_string(),
pdb_id: Some("1M17".to_string()),
pockets: vec![
BindingPocket {
id: 0,
name: "ATP binding site".to_string(),
residues: vec![718, 719, 721, 726, 745, 790, 791, 792, 793, 854, 855],
center: Coordinate3D { x: 25.0, y: 18.5, z: 42.0 },
volume: 480.0,
druggability: 0.92,
},
],
},
]
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_element_atomic_number() {
assert_eq!(Element::C.atomic_number(), 6);
assert_eq!(Element::N.atomic_number(), 7);
assert_eq!(Element::O.atomic_number(), 8);
}
#[test]
fn test_affinity_class() {
assert_eq!(AffinityClass::from_pkd(10.0), AffinityClass::VeryHigh);
assert_eq!(AffinityClass::from_pkd(8.0), AffinityClass::High);
assert_eq!(AffinityClass::from_pkd(6.0), AffinityClass::Medium);
assert_eq!(AffinityClass::from_pkd(4.0), AffinityClass::Low);
assert_eq!(AffinityClass::from_pkd(2.0), AffinityClass::VeryLow);
}
#[test]
fn test_sample_molecules() {
let molecules = get_sample_molecules();
assert_eq!(molecules.len(), 4);
assert!(
molecules
.iter()
.any(|m| m.name == Some("Aspirin".to_string()))
);
}
#[test]
fn test_sample_targets() {
let targets = get_sample_targets();
assert!(!targets.is_empty());
assert!(targets.iter().any(|t| t.name.contains("Carbonic")));
}
#[test]
fn test_prediction_config_default() {
let config = PredictionConfig::default();
assert_eq!(config.model, BindingModel::GraphTransformer);
assert!(config.predict_pose);
assert_eq!(config.num_poses, 5);
}
#[test]
fn test_screening_config_default() {
let config = ScreeningConfig::default();
assert_eq!(config.top_k, 100);
assert!(config.lipinski_filter);
}
#[test]
fn test_generation_config_default() {
let config = GenerationConfig::default();
assert_eq!(config.num_molecules, 100);
assert!(config.lipinski_constraints);
assert_eq!(config.max_mw, 500.0);
}
#[test]
fn test_molecule_serialization() {
let mol = &get_sample_molecules()[0];
let json = serde_json::to_string(mol).unwrap();
assert!(json.contains("aspirin"));
}
#[test]
fn test_coordinate_3d() {
let coord = Coordinate3D {
x: 1.0,
y: 2.0,
z: 3.0,
};
let json = serde_json::to_string(&coord).unwrap();
let parsed: Coordinate3D = serde_json::from_str(&json).unwrap();
assert!((parsed.x - 1.0).abs() < 0.001);
}
}