rustytorch/demos/rtx-alphafold-demo/src/confidence.rs

//! Confidence prediction module for pLDDT and PAE scores.

use alphafold_shared::ConfidenceCategory;

/// Confidence predictor for structure quality assessment.
#[derive(Debug, Clone)]
pub struct ConfidencePredictor {
    /// Weights for pLDDT prediction
    plddt_weights: Vec<f32>,
    /// Weights for PAE prediction
    pae_weights: Vec<f32>,
}

impl Default for ConfidencePredictor {
    fn default() -> Self {
        Self::new()
    }
}

impl ConfidencePredictor {
    /// Create a new confidence predictor.
    #[must_use]
    pub fn new() -> Self {
        use rand::SeedableRng;
        use rand_distr::{Distribution, Normal};

        let mut rng = rand::rngs::StdRng::seed_from_u64(789);
        let normal = Normal::new(0.0_f32, 0.1).unwrap();

        let plddt_weights: Vec<f32> = (0..256).map(|_| normal.sample(&mut rng)).collect();
        let pae_weights: Vec<f32> = (0..256).map(|_| normal.sample(&mut rng)).collect();

        Self {
            plddt_weights,
            pae_weights,
        }
    }

    /// Predict pLDDT scores from embeddings.
    #[must_use]
    pub fn predict_plddt(&self, embeddings: &[Vec<f32>]) -> Vec<f32> {
        embeddings
            .iter()
            .map(|emb| {
                let mut score = 0.0_f32;
                for (i, &w) in self.plddt_weights.iter().enumerate() {
                    if i < emb.len() {
                        score += emb[i] * w;
                    }
                }
                // Sigmoid and scale to 0-100
                (1.0 / (1.0 + (-score).exp())) * 100.0
            })
            .collect()
    }

    /// Predict PAE matrix from pair embeddings.
    #[must_use]
    pub fn predict_pae(&self, pair_embeddings: &[Vec<Vec<f32>>]) -> Vec<Vec<f32>> {
        let seq_len = pair_embeddings.len();
        let mut pae = vec![vec![0.0_f32; seq_len]; seq_len];

        for i in 0..seq_len {
            for j in 0..seq_len {
                let emb = &pair_embeddings[i][j];
                let mut score = 0.0_f32;
                for (k, &w) in self.pae_weights.iter().enumerate() {
                    if k < emb.len() {
                        score += emb[k] * w;
                    }
                }
                // Softplus and scale to 0-31.75Å
                pae[i][j] = (1.0 + score.exp()).ln() * 5.0;
            }
        }

        pae
    }
}

/// Calculate confidence statistics from pLDDT scores.
#[derive(Debug, Clone)]
pub struct ConfidenceStats {
    /// Average pLDDT
    pub mean: f32,
    /// Median pLDDT
    pub median: f32,
    /// Minimum pLDDT
    pub min: f32,
    /// Maximum pLDDT
    pub max: f32,
    /// Standard deviation
    pub std: f32,
    /// Percentage of residues in each category
    pub category_percentages: CategoryPercentages,
}

/// Percentage of residues in each confidence category.
#[derive(Debug, Clone)]
pub struct CategoryPercentages {
    pub very_high: f32,
    pub high: f32,
    pub low: f32,
    pub very_low: f32,
}

impl ConfidenceStats {
    /// Calculate statistics from pLDDT scores.
    #[must_use]
    pub fn from_plddt(scores: &[f32]) -> Self {
        if scores.is_empty() {
            return Self {
                mean: 0.0,
                median: 0.0,
                min: 0.0,
                max: 0.0,
                std: 0.0,
                category_percentages: CategoryPercentages {
                    very_high: 0.0,
                    high: 0.0,
                    low: 0.0,
                    very_low: 0.0,
                },
            };
        }

        let n = scores.len() as f32;

        // Mean
        let mean = scores.iter().sum::<f32>() / n;

        // Sorted for median
        let mut sorted = scores.to_vec();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());

        let median = if scores.len().is_multiple_of(2) {
            f32::midpoint(sorted[scores.len() / 2 - 1], sorted[scores.len() / 2])
        } else {
            sorted[scores.len() / 2]
        };

        let min = sorted[0];
        let max = sorted[sorted.len() - 1];

        // Standard deviation
        let variance = scores.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / n;
        let std = variance.sqrt();

        // Category percentages
        let mut very_high = 0;
        let mut high = 0;
        let mut low = 0;
        let mut very_low = 0;

        for &score in scores {
            match ConfidenceCategory::from_plddt(score) {
                ConfidenceCategory::VeryHigh => very_high += 1,
                ConfidenceCategory::High => high += 1,
                ConfidenceCategory::Low => low += 1,
                ConfidenceCategory::VeryLow => very_low += 1,
            }
        }

        let category_percentages = CategoryPercentages {
            very_high: very_high as f32 / n * 100.0,
            high: high as f32 / n * 100.0,
            low: low as f32 / n * 100.0,
            very_low: very_low as f32 / n * 100.0,
        };

        Self {
            mean,
            median,
            min,
            max,
            std,
            category_percentages,
        }
    }
}

/// Analyze PAE matrix for domain boundaries.
#[must_use]
pub fn analyze_pae_domains(pae: &[Vec<f32>], threshold: f32) -> Vec<DomainBoundary> {
    let seq_len = pae.len();
    if seq_len < 10 {
        return vec![];
    }

    let mut boundaries = Vec::new();
    let mut in_domain = false;
    let mut domain_start = 0;

    for i in 0..seq_len {
        // Average PAE for residue i to nearby residues
        let mut local_pae = 0.0_f32;
        let mut count = 0;

        for j in i.saturating_sub(5)..=(i + 5).min(seq_len - 1) {
            local_pae += pae[i][j];
            count += 1;
        }
        local_pae /= count as f32;

        if !in_domain && local_pae < threshold {
            // Start of a domain
            in_domain = true;
            domain_start = i;
        } else if in_domain && local_pae >= threshold {
            // End of a domain
            in_domain = false;
            if i - domain_start >= 10 {
                boundaries.push(DomainBoundary {
                    start: domain_start,
                    end: i,
                    avg_pae: calculate_block_pae(pae, domain_start, i),
                });
            }
        }
    }

    // Close last domain if still in one
    if in_domain && seq_len - domain_start >= 10 {
        boundaries.push(DomainBoundary {
            start: domain_start,
            end: seq_len,
            avg_pae: calculate_block_pae(pae, domain_start, seq_len),
        });
    }

    boundaries
}

/// Domain boundary information.
#[derive(Debug, Clone)]
pub struct DomainBoundary {
    /// Start residue index
    pub start: usize,
    /// End residue index (exclusive)
    pub end: usize,
    /// Average PAE within domain
    pub avg_pae: f32,
}

/// Calculate average PAE within a block.
fn calculate_block_pae(pae: &[Vec<f32>], start: usize, end: usize) -> f32 {
    let mut sum = 0.0_f32;
    let mut count = 0;

    for i in start..end {
        for j in start..end {
            sum += pae[i][j];
            count += 1;
        }
    }

    if count > 0 { sum / count as f32 } else { 0.0 }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_confidence_stats() {
        let scores = vec![95.0, 85.0, 75.0, 65.0, 55.0];
        let stats = ConfidenceStats::from_plddt(&scores);

        assert!((stats.mean - 75.0).abs() < 0.01);
        assert_eq!(stats.min, 55.0);
        assert_eq!(stats.max, 95.0);
    }

    #[test]
    fn test_confidence_predictor() {
        let predictor = ConfidencePredictor::new();
        let embeddings = vec![vec![0.1_f32; 256]; 10];

        let plddt = predictor.predict_plddt(&embeddings);
        assert_eq!(plddt.len(), 10);
        assert!(plddt.iter().all(|&x| x >= 0.0 && x <= 100.0));
    }

    #[test]
    fn test_category_percentages() {
        let scores = vec![95.0, 95.0, 80.0, 60.0, 40.0]; // 2 very high, 1 high, 1 low, 1 very low
        let stats = ConfidenceStats::from_plddt(&scores);

        assert!((stats.category_percentages.very_high - 40.0).abs() < 0.01);
        assert!((stats.category_percentages.high - 20.0).abs() < 0.01);
        assert!((stats.category_percentages.low - 20.0).abs() < 0.01);
        assert!((stats.category_percentages.very_low - 20.0).abs() < 0.01);
    }

    #[test]
    fn test_analyze_pae_domains() {
        // Create a PAE matrix with two domains
        let seq_len = 50;
        let mut pae = vec![vec![15.0_f32; seq_len]; seq_len];

        // First domain (0-20): low PAE
        for i in 0..20 {
            for j in 0..20 {
                pae[i][j] = 3.0;
            }
        }

        // Second domain (30-50): low PAE
        for i in 30..50 {
            for j in 30..50 {
                pae[i][j] = 4.0;
            }
        }

        let domains = analyze_pae_domains(&pae, 10.0);
        assert_eq!(domains.len(), 2);
    }
}