Initial commit

demos/rtx-drugbinder-demo/src/molecule_generator.rs

@@ -0,0 +1,741 @@
+//! Diffusion-based molecule generator.
+//!
+//! Generates novel drug-like molecules using diffusion models.
+
+use drugbinder_shared::{
+    AffinityClass, AffinityPrediction, Atom, Bond, BondType, Coordinate3D, Element,
+    GeneratedMolecule, GenerationConfig, GenerationResult, GenerationStats, Hybridization,
+    MolecularProperties, Molecule, ProteinTarget,
+};
+
+use crate::DrugBinderError;
+
+/// Configuration for molecule generator.
+#[derive(Debug, Clone)]
+pub struct MoleculeGeneratorConfig {
+    /// Hidden dimension
+    pub hidden_dim: usize,
+    /// Number of diffusion steps
+    pub num_steps: usize,
+    /// Maximum atoms
+    pub max_atoms: usize,
+    /// Temperature for sampling
+    pub temperature: f32,
+    /// Target affinity
+    pub target_affinity: Option<f32>,
+    /// Enforce Lipinski
+    pub lipinski_constraints: bool,
+}
+
+impl Default for MoleculeGeneratorConfig {
+    fn default() -> Self {
+        Self {
+            hidden_dim: 128,
+            num_steps: 100,
+            max_atoms: 50,
+            temperature: 1.0,
+            target_affinity: Some(8.0),
+            lipinski_constraints: true,
+        }
+    }
+}
+
+impl From<&GenerationConfig> for MoleculeGeneratorConfig {
+    fn from(config: &GenerationConfig) -> Self {
+        Self {
+            hidden_dim: 128,
+            num_steps: 100,
+            max_atoms: 50,
+            temperature: 1.0,
+            target_affinity: config.target_affinity,
+            lipinski_constraints: config.lipinski_constraints,
+        }
+    }
+}
+
+/// Molecule generator using diffusion.
+#[derive(Debug)]
+pub struct MoleculeGenerator {
+    config: MoleculeGeneratorConfig,
+    denoiser: GraphDenoiser,
+    scheduler: DDPMScheduler,
+    property_predictor: PropertyPredictor,
+}
+
+impl MoleculeGenerator {
+    /// Create a new molecule generator.
+    #[must_use] 
+    pub fn new(config: MoleculeGeneratorConfig) -> Self {
+        let denoiser = GraphDenoiser::new(config.hidden_dim, config.max_atoms);
+        let scheduler = DDPMScheduler::new(config.num_steps);
+        let property_predictor = PropertyPredictor::new();
+
+        Self {
+            config,
+            denoiser,
+            scheduler,
+            property_predictor,
+        }
+    }
+
+    /// Generate molecules for a target.
+    pub fn generate(
+        &self,
+        target: &ProteinTarget,
+        config: &GenerationConfig,
+    ) -> Result<GenerationResult, DrugBinderError> {
+        use rand::SeedableRng;
+
+        // Seed based on target
+        let seed: u64 = target.sequence.len() as u64 * 1337;
+        let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
+
+        let mut molecules = Vec::new();
+        let mut valid_count = 0;
+        let mut unique_smiles = std::collections::HashSet::new();
+
+        for i in 0..config.num_molecules {
+            // Generate one molecule
+            match self.generate_single(target, i, &mut rng) {
+                Some(mol) => {
+                    // Check validity
+                    if self.is_valid(&mol.molecule) {
+                        valid_count += 1;
+
+                        // Check uniqueness
+                        if unique_smiles.insert(mol.molecule.smiles.clone()) {
+                            // Check constraints
+                            if !config.lipinski_constraints || mol.molecule.properties.lipinski_pass
+                            {
+                                molecules.push(mol);
+                            }
+                        }
+                    }
+                }
+                None => continue,
+            }
+        }
+
+        // Calculate stats
+        let high_affinity = molecules
+            .iter()
+            .filter(|m| m.predicted_binding.pkd > 7.0)
+            .count();
+
+        let stats = GenerationStats {
+            total_generated: config.num_molecules,
+            valid: valid_count,
+            unique: unique_smiles.len(),
+            novel: molecules.len(), // All generated are "novel"
+            high_affinity,
+        };
+
+        Ok(GenerationResult { molecules, stats })
+    }
+
+    fn generate_single<R: rand::Rng>(
+        &self,
+        target: &ProteinTarget,
+        index: usize,
+        rng: &mut R,
+    ) -> Option<GeneratedMolecule> {
+        use rand_distr::{Distribution, Normal, Uniform};
+
+        let normal = Normal::new(0.0_f32, 1.0).ok()?;
+        let uniform_atoms = Uniform::new(8, self.config.max_atoms.min(30)).ok()?;
+
+        // Sample number of atoms
+        let num_atoms = uniform_atoms.sample(rng);
+
+        // Generate through diffusion (simplified)
+        let mut atom_features: Vec<Vec<f32>> = (0..num_atoms)
+            .map(|_| {
+                (0..self.config.hidden_dim)
+                    .map(|_| normal.sample(rng))
+                    .collect()
+            })
+            .collect();
+
+        // Denoise
+        for t in (0..self.config.num_steps).rev() {
+            atom_features = self.denoiser.denoise(&atom_features, t, target);
+        }
+
+        // Decode to molecule
+        let smiles = self.decode_to_smiles(&atom_features, index, rng);
+        let atoms = self.decode_atoms(&atom_features);
+        let bonds = self.generate_bonds(&atoms);
+
+        // Calculate molecular weight
+        let mw: f32 = atoms.iter().map(|a| element_mass(&a.element)).sum();
+
+        // Generate properties
+        let properties = self.property_predictor.predict(&atoms, &bonds);
+
+        // Generate 3D coordinates
+        let coords = self.generate_coordinates(&atoms, rng);
+
+        let molecule = Molecule {
+            id: format!("gen_{index}"),
+            name: Some(format!("Generated_{index}")),
+            smiles,
+            molecular_weight: mw,
+            num_atoms: atoms.len(),
+            num_bonds: bonds.len(),
+            atoms,
+            bonds,
+            coordinates_3d: Some(coords),
+            properties,
+        };
+
+        // Predict binding
+        let pkd = self.predict_affinity(&molecule, target);
+        let predicted_binding = AffinityPrediction {
+            pkd,
+            pic50: pkd - 0.3,
+            delta_g: -1.364 * pkd,
+            uncertainty: 0.5 + normal.sample(rng).abs() * 0.3,
+            affinity_class: AffinityClass::from_pkd(pkd),
+        };
+
+        // Calculate novelty and SA score
+        let novelty = 0.7 + normal.sample(rng).abs() * 0.25;
+        let sa_score = 2.0 + normal.sample(rng).abs() * 3.0;
+
+        Some(GeneratedMolecule {
+            molecule,
+            predicted_binding,
+            novelty: novelty.clamp(0.0, 1.0),
+            sa_score: sa_score.clamp(1.0, 10.0),
+        })
+    }
+
+    fn decode_to_smiles<R: rand::Rng>(
+        &self,
+        features: &[Vec<f32>],
+        index: usize,
+        _rng: &mut R,
+    ) -> String {
+        // Generate a plausible SMILES from features (simplified)
+        let scaffolds = [
+            "c1ccccc1",         // Benzene
+            "c1ccc2ccccc2c1",   // Naphthalene
+            "c1ccncc1",         // Pyridine
+            "c1cnc2ccccc2n1",   // Quinazoline
+            "C1CCCCC1",         // Cyclohexane
+            "C1CCNCC1",         // Piperidine
+            "c1cc2ccccc2[nH]1", // Indole
+        ];
+
+        let substituents = [
+            "C", "CC", "CCC", "C(=O)O", "C(=O)N", "O", "N", "F", "Cl", "OC", "NC", "C(C)C",
+            "C(=O)OC", "NC(=O)C",
+        ];
+
+        let scaffold = scaffolds[index % scaffolds.len()];
+        let num_subs = (features.len() / 10).clamp(1, 3);
+
+        let mut smiles = scaffold.to_string();
+        for i in 0..num_subs {
+            let sub = substituents[(index + i * 7) % substituents.len()];
+            smiles = format!("{smiles}({sub})");
+        }
+
+        smiles
+    }
+
+    fn decode_atoms(&self, features: &[Vec<f32>]) -> Vec<Atom> {
+        features
+            .iter()
+            .enumerate()
+            .map(|(i, feat)| {
+                // Determine element from feature vector
+                let element = self.feature_to_element(feat);
+
+                Atom {
+                    index: i,
+                    element,
+                    formal_charge: 0,
+                    num_hydrogens: self.implicit_hydrogens(&element),
+                    is_aromatic: feat.first().is_some_and(|&x| x > 0.0),
+                    hybridization: Hybridization::Sp3,
+                }
+            })
+            .collect()
+    }
+
+    fn feature_to_element(&self, feat: &[f32]) -> Element {
+        // Map feature to element based on feature pattern
+        let sum: f32 = feat.iter().take(10).sum();
+
+        if sum < -2.0 {
+            Element::N
+        } else if sum < 0.0 {
+            Element::O
+        } else if sum < 2.0 {
+            Element::C
+        } else if sum < 3.0 {
+            Element::S
+        } else if sum < 4.0 {
+            Element::F
+        } else {
+            Element::C // Default to carbon
+        }
+    }
+
+    fn implicit_hydrogens(&self, element: &Element) -> u8 {
+        match element {
+            Element::C => 4,
+            Element::N => 3,
+            Element::O => 2,
+            Element::S => 2,
+            Element::F | Element::Cl | Element::Br | Element::I => 1,
+            _ => 0,
+        }
+    }
+
+    fn generate_bonds(&self, atoms: &[Atom]) -> Vec<Bond> {
+        // Generate plausible bonds
+        let mut bonds = Vec::new();
+
+        for i in 0..atoms.len().saturating_sub(1) {
+            let bond_type = if atoms[i].is_aromatic && atoms[i + 1].is_aromatic {
+                BondType::Aromatic
+            } else {
+                BondType::Single
+            };
+
+            bonds.push(Bond {
+                atom1: i,
+                atom2: i + 1,
+                bond_type,
+                is_conjugated: atoms[i].is_aromatic,
+                is_in_ring: i < 6, // First 6 atoms often in ring
+            });
+        }
+
+        // Close ring if aromatic
+        if atoms.len() >= 6 && atoms[0].is_aromatic {
+            bonds.push(Bond {
+                atom1: atoms.len() - 1,
+                atom2: 0,
+                bond_type: BondType::Aromatic,
+                is_conjugated: true,
+                is_in_ring: true,
+            });
+        }
+
+        bonds
+    }
+
+    fn generate_coordinates<R: rand::Rng>(&self, atoms: &[Atom], rng: &mut R) -> Vec<Coordinate3D> {
+        use rand_distr::{Distribution, Normal};
+
+        let normal = Normal::new(0.0_f32, 1.5).unwrap();
+
+        atoms
+            .iter()
+            .enumerate()
+            .map(|(i, _)| {
+                // Simple 3D placement
+                let angle = (i as f32) * 2.0 * std::f32::consts::PI / atoms.len() as f32;
+                Coordinate3D {
+                    x: angle.cos() * 1.5 + normal.sample(rng) * 0.3,
+                    y: angle.sin() * 1.5 + normal.sample(rng) * 0.3,
+                    z: normal.sample(rng) * 0.5,
+                }
+            })
+            .collect()
+    }
+
+    fn predict_affinity(&self, molecule: &Molecule, _target: &ProteinTarget) -> f32 {
+        // Simple affinity prediction based on molecular properties
+        let base = 5.0;
+
+        // Adjust based on properties
+        let mw_bonus = if molecule.molecular_weight < 500.0 {
+            0.5
+        } else {
+            -0.5
+        };
+        let hbd_bonus = if molecule.properties.hbd <= 5 {
+            0.3
+        } else {
+            -0.3
+        };
+        let logp_bonus = if molecule.properties.log_p < 5.0 {
+            0.4
+        } else {
+            -0.4
+        };
+        let ring_bonus = f32::from(molecule.properties.num_aromatic_rings) * 0.3;
+
+        (base + mw_bonus + hbd_bonus + logp_bonus + ring_bonus).clamp(3.0, 10.0)
+    }
+
+    fn is_valid(&self, molecule: &Molecule) -> bool {
+        // Basic validity checks
+        !molecule.atoms.is_empty()
+            && molecule.molecular_weight > 50.0
+            && molecule.molecular_weight < 1000.0
+    }
+}
+
+/// Graph denoiser for diffusion.
+#[derive(Debug)]
+struct GraphDenoiser {
+    hidden_dim: usize,
+    max_atoms: usize,
+    weights: Vec<Vec<f32>>,
+}
+
+impl GraphDenoiser {
+    fn new(hidden_dim: usize, max_atoms: usize) -> Self {
+        use rand::SeedableRng;
+        use rand_distr::{Distribution, Normal};
+
+        let mut rng = rand::rngs::StdRng::seed_from_u64(42);
+        let std = (2.0 / hidden_dim as f32).sqrt();
+        let normal = Normal::new(0.0_f32, std).unwrap();
+
+        let weights: Vec<Vec<f32>> = (0..hidden_dim)
+            .map(|_| (0..hidden_dim).map(|_| normal.sample(&mut rng)).collect())
+            .collect();
+
+        Self {
+            hidden_dim,
+            max_atoms,
+            weights,
+        }
+    }
+
+    fn denoise(
+        &self,
+        atom_features: &[Vec<f32>],
+        timestep: usize,
+        _target: &ProteinTarget,
+    ) -> Vec<Vec<f32>> {
+        let alpha = 1.0 - (timestep as f32 / 100.0);
+
+        atom_features
+            .iter()
+            .map(|feat| {
+                // Simple denoising: reduce noise component
+                let mut denoised = vec![0.0; self.hidden_dim];
+
+                for (i, &x) in feat.iter().enumerate() {
+                    if i < self.hidden_dim {
+                        // Apply learned transformation
+                        for (j, &w) in self.weights[i].iter().enumerate() {
+                            denoised[j] += x * w * alpha;
+                        }
+                    }
+                }
+
+                // Add residual
+                for (i, &x) in feat.iter().enumerate().take(self.hidden_dim) {
+                    denoised[i] += x * (1.0 - alpha);
+                }
+
+                denoised
+            })
+            .collect()
+    }
+}
+
+/// DDPM scheduler.
+#[derive(Debug)]
+struct DDPMScheduler {
+    num_steps: usize,
+    betas: Vec<f32>,
+    alphas: Vec<f32>,
+    alpha_cumprod: Vec<f32>,
+}
+
+impl DDPMScheduler {
+    fn new(num_steps: usize) -> Self {
+        // Linear beta schedule
+        let beta_start = 0.0001;
+        let beta_end = 0.02;
+
+        let betas: Vec<f32> = (0..num_steps)
+            .map(|i| beta_start + (beta_end - beta_start) * (i as f32) / (num_steps as f32))
+            .collect();
+
+        let alphas: Vec<f32> = betas.iter().map(|&b| 1.0 - b).collect();
+
+        let mut alpha_cumprod = Vec::with_capacity(num_steps);
+        let mut prod = 1.0;
+        for &a in &alphas {
+            prod *= a;
+            alpha_cumprod.push(prod);
+        }
+
+        Self {
+            num_steps,
+            betas,
+            alphas,
+            alpha_cumprod,
+        }
+    }
+
+    #[allow(dead_code)]
+    fn get_alpha(&self, t: usize) -> f32 {
+        self.alphas.get(t).copied().unwrap_or(1.0)
+    }
+
+    #[allow(dead_code)]
+    fn get_alpha_cumprod(&self, t: usize) -> f32 {
+        self.alpha_cumprod.get(t).copied().unwrap_or(1.0)
+    }
+}
+
+/// Property predictor for generated molecules.
+#[derive(Debug)]
+struct PropertyPredictor;
+
+impl PropertyPredictor {
+    fn new() -> Self {
+        Self
+    }
+
+    fn predict(&self, atoms: &[Atom], bonds: &[Bond]) -> MolecularProperties {
+        // Calculate properties from structure
+
+        // Count HBD (NH, OH)
+        let hbd: u8 = atoms
+            .iter()
+            .filter(|a| matches!(a.element, Element::N | Element::O) && a.num_hydrogens > 0)
+            .count() as u8;
+
+        // Count HBA (N, O)
+        let hba: u8 = atoms
+            .iter()
+            .filter(|a| matches!(a.element, Element::N | Element::O))
+            .count() as u8;
+
+        // Estimate LogP
+        let carbon_count = atoms.iter().filter(|a| a.element == Element::C).count() as f32;
+        let nitrogen_count = atoms.iter().filter(|a| a.element == Element::N).count() as f32;
+        let oxygen_count = atoms.iter().filter(|a| a.element == Element::O).count() as f32;
+
+        let log_p = carbon_count * 0.5 - nitrogen_count * 0.8 - oxygen_count * 1.0;
+
+        // Count rotatable bonds
+        let rotatable_bonds: u8 = bonds
+            .iter()
+            .filter(|b| b.bond_type == BondType::Single && !b.is_in_ring)
+            .count() as u8;
+
+        // Count rings
+        let ring_bonds = bonds.iter().filter(|b| b.is_in_ring).count();
+        let num_rings = (ring_bonds / 5).max(1) as u8;
+
+        // Count aromatic rings
+        let aromatic_bonds = bonds
+            .iter()
+            .filter(|b| b.bond_type == BondType::Aromatic)
+            .count();
+        let num_aromatic_rings = (aromatic_bonds / 5) as u8;
+
+        // Calculate TPSA (simplified)
+        let tpsa = nitrogen_count * 26.0 + oxygen_count * 20.0;
+
+        // Calculate molecular weight
+        let mw: f32 = atoms.iter().map(|a| element_mass(&a.element)).sum();
+
+        // Check Lipinski's Rule of Five
+        let lipinski_pass = mw <= 500.0 && log_p <= 5.0 && hbd <= 5 && hba <= 10;
+
+        // Calculate QED (simplified)
+        let qed = calculate_qed(mw, log_p, hbd, hba, tpsa, rotatable_bonds);
+
+        MolecularProperties {
+            log_p,
+            hbd,
+            hba,
+            tpsa,
+            rotatable_bonds,
+            num_rings,
+            num_aromatic_rings,
+            lipinski_pass,
+            qed,
+        }
+    }
+}
+
+fn element_mass(element: &Element) -> f32 {
+    match element {
+        Element::H => 1.008,
+        Element::C => 12.011,
+        Element::N => 14.007,
+        Element::O => 15.999,
+        Element::F => 18.998,
+        Element::P => 30.974,
+        Element::S => 32.065,
+        Element::Cl => 35.453,
+        Element::Br => 79.904,
+        Element::I => 126.90,
+        Element::Na => 22.990,
+        Element::Mg => 24.305,
+        Element::K => 39.098,
+        Element::Ca => 40.078,
+        Element::Fe => 55.845,
+        Element::Zn => 65.38,
+        Element::Cu => 63.546,
+        Element::Other => 12.0,
+    }
+}
+
+fn calculate_qed(mw: f32, log_p: f32, hbd: u8, hba: u8, tpsa: f32, rotatable_bonds: u8) -> f32 {
+    // Simplified QED calculation
+    let mw_score = gaussian(mw, 350.0, 100.0);
+    let logp_score = gaussian(log_p, 2.5, 1.5);
+    let hbd_score = gaussian(f32::from(hbd), 1.0, 2.0);
+    let hba_score = gaussian(f32::from(hba), 4.0, 3.0);
+    let tpsa_score = gaussian(tpsa, 70.0, 30.0);
+    let rotb_score = gaussian(f32::from(rotatable_bonds), 3.0, 3.0);
+
+    // Geometric mean
+    let product = mw_score * logp_score * hbd_score * hba_score * tpsa_score * rotb_score;
+    product.powf(1.0 / 6.0)
+}
+
+fn gaussian(x: f32, mean: f32, std: f32) -> f32 {
+    (-0.5 * ((x - mean) / std).powi(2)).exp()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_generator_creation() {
+        let config = MoleculeGeneratorConfig::default();
+        let generator = MoleculeGenerator::new(config);
+        assert_eq!(generator.config.num_steps, 100);
+    }
+
+    #[test]
+    fn test_generate_molecules() {
+        let config = MoleculeGeneratorConfig::default();
+        let generator = MoleculeGenerator::new(config);
+
+        let targets = drugbinder_shared::get_sample_targets();
+        let target = &targets[0];
+
+        let gen_config = GenerationConfig {
+            num_molecules: 5,
+            ..Default::default()
+        };
+
+        let result = generator.generate(target, &gen_config);
+        assert!(result.is_ok());
+
+        let generation = result.unwrap();
+        assert!(generation.stats.valid > 0);
+    }
+
+    #[test]
+    fn test_scheduler() {
+        let scheduler = DDPMScheduler::new(100);
+        assert_eq!(scheduler.num_steps, 100);
+        assert!(scheduler.get_alpha(0) > scheduler.get_alpha(99));
+    }
+
+    #[test]
+    fn test_property_predictor() {
+        let predictor = PropertyPredictor::new();
+
+        let atoms = vec![
+            Atom {
+                index: 0,
+                element: Element::C,
+                formal_charge: 0,
+                num_hydrogens: 3,
+                is_aromatic: false,
+                hybridization: Hybridization::Sp3,
+            },
+            Atom {
+                index: 1,
+                element: Element::O,
+                formal_charge: 0,
+                num_hydrogens: 1,
+                is_aromatic: false,
+                hybridization: Hybridization::Sp3,
+            },
+        ];
+
+        let bonds = vec![Bond {
+            atom1: 0,
+            atom2: 1,
+            bond_type: BondType::Single,
+            is_conjugated: false,
+            is_in_ring: false,
+        }];
+
+        let props = predictor.predict(&atoms, &bonds);
+        assert!(props.hbd >= 1); // OH is HBD
+        assert!(props.hba >= 1); // O is HBA
+    }
+
+    #[test]
+    fn test_element_mass() {
+        assert!((element_mass(&Element::C) - 12.011).abs() < 0.01);
+        assert!((element_mass(&Element::O) - 15.999).abs() < 0.01);
+    }
+
+    #[test]
+    fn test_qed_calculation() {
+        // Drug-like molecule
+        let qed1 = calculate_qed(350.0, 2.5, 2, 4, 70.0, 3);
+        // Non-drug-like
+        let qed2 = calculate_qed(700.0, 7.0, 8, 15, 200.0, 15);
+
+        assert!(qed1 > qed2); // Drug-like should have higher QED
+    }
+
+    #[test]
+    fn test_denoiser() {
+        let denoiser = GraphDenoiser::new(64, 30);
+
+        let features: Vec<Vec<f32>> = (0..10).map(|_| vec![0.1; 64]).collect();
+
+        let targets = drugbinder_shared::get_sample_targets();
+        let denoised = denoiser.denoise(&features, 50, &targets[0]);
+
+        assert_eq!(denoised.len(), 10);
+        assert_eq!(denoised[0].len(), 64);
+    }
+
+    #[test]
+    fn test_lipinski_check() {
+        let predictor = PropertyPredictor::new();
+
+        // Small drug-like molecule
+        let atoms: Vec<Atom> = (0..15)
+            .map(|i| Atom {
+                index: i,
+                element: if i < 10 { Element::C } else { Element::O },
+                formal_charge: 0,
+                num_hydrogens: if i < 10 { 2 } else { 1 },
+                is_aromatic: false,
+                hybridization: Hybridization::Sp3,
+            })
+            .collect();
+
+        let bonds: Vec<Bond> = (0..14)
+            .map(|i| Bond {
+                atom1: i,
+                atom2: i + 1,
+                bond_type: BondType::Single,
+                is_conjugated: false,
+                is_in_ring: false,
+            })
+            .collect();
+
+        let props = predictor.predict(&atoms, &bonds);
+        // Should pass Lipinski for small molecule
+        assert!(props.lipinski_pass || props.log_p > 5.0);
+    }
+}