Initial commit

demos/rtx-alphafold-demo/src/evoformer.rs

@@ -0,0 +1,220 @@
+//! Evoformer module - attention over sequences and residue pairs.
+//!
+//! The Evoformer processes sequence and pair representations through
+//! alternating attention mechanisms to capture evolutionary relationships.
+
+use crate::encoder::{EMBEDDING_DIM, PairEmbedding, SequenceEmbedding};
+
+/// Number of attention heads.
+pub const NUM_HEADS: usize = 8;
+
+/// Evoformer block for processing sequence and pair representations.
+#[derive(Debug, Clone)]
+pub struct EvoformerBlock {
+    /// Attention head dimension
+    head_dim: usize,
+    /// Number of attention heads
+    num_heads: usize,
+    /// Query projection weights
+    wq: Vec<Vec<f32>>,
+    /// Key projection weights
+    wk: Vec<Vec<f32>>,
+    /// Value projection weights
+    wv: Vec<Vec<f32>>,
+    /// Output projection weights
+    wo: Vec<Vec<f32>>,
+}
+
+impl Default for EvoformerBlock {
+    fn default() -> Self {
+        Self::new(NUM_HEADS)
+    }
+}
+
+impl EvoformerBlock {
+    /// Create a new Evoformer block.
+    #[must_use] 
+    pub fn new(num_heads: usize) -> Self {
+        use rand::SeedableRng;
+        use rand_distr::{Distribution, Normal};
+
+        let head_dim = EMBEDDING_DIM / num_heads;
+        let mut rng = rand::rngs::StdRng::seed_from_u64(123);
+        let normal = Normal::new(0.0_f32, 0.02).unwrap();
+
+        let mut init_weights = |rows: usize, cols: usize| -> Vec<Vec<f32>> {
+            (0..rows)
+                .map(|_| (0..cols).map(|_| normal.sample(&mut rng)).collect())
+                .collect()
+        };
+
+        Self {
+            head_dim,
+            num_heads,
+            wq: init_weights(EMBEDDING_DIM, EMBEDDING_DIM),
+            wk: init_weights(EMBEDDING_DIM, EMBEDDING_DIM),
+            wv: init_weights(EMBEDDING_DIM, EMBEDDING_DIM),
+            wo: init_weights(EMBEDDING_DIM, EMBEDDING_DIM),
+        }
+    }
+
+    /// Apply row-wise (sequence) attention.
+    pub fn row_attention(&self, seq_emb: &mut SequenceEmbedding, pair_emb: &PairEmbedding) {
+        let seq_len = seq_emb.sequence_length;
+
+        // Compute Q, K, V projections
+        let queries = self.project(&seq_emb.embeddings, &self.wq);
+        let keys = self.project(&seq_emb.embeddings, &self.wk);
+        let values = self.project(&seq_emb.embeddings, &self.wv);
+
+        // Multi-head attention
+        let scale = 1.0 / (self.head_dim as f32).sqrt();
+
+        for i in 0..seq_len {
+            let mut output = vec![0.0_f32; EMBEDDING_DIM];
+
+            for head in 0..self.num_heads {
+                let head_start = head * self.head_dim;
+                let head_end = head_start + self.head_dim;
+
+                // Compute attention scores
+                let mut scores = Vec::with_capacity(seq_len);
+                for j in 0..seq_len {
+                    let mut score = 0.0_f32;
+                    for k in head_start..head_end {
+                        score += queries[i][k] * keys[j][k];
+                    }
+                    // Add pair bias from pair representation
+                    if let Some(pair) = pair_emb.get(i, j) {
+                        score += pair.iter().take(self.head_dim).sum::<f32>() * 0.1;
+                    }
+                    scores.push(score * scale);
+                }
+
+                // Softmax
+                let max_score = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);
+                let exp_scores: Vec<f32> = scores.iter().map(|s| (s - max_score).exp()).collect();
+                let sum_exp: f32 = exp_scores.iter().sum();
+                let attention: Vec<f32> = exp_scores.iter().map(|e| e / sum_exp).collect();
+
+                // Weighted sum of values
+                for j in 0..seq_len {
+                    for k in head_start..head_end {
+                        output[k] += attention[j] * values[j][k];
+                    }
+                }
+            }
+
+            // Output projection and residual
+            let projected = self.project_single(&output, &self.wo);
+            for (k, &p) in projected.iter().enumerate() {
+                seq_emb.embeddings[i][k] += p * 0.1; // Scaled residual
+            }
+        }
+    }
+
+    /// Apply column-wise attention (for MSA, simplified here).
+    pub fn column_attention(&self, seq_emb: &mut SequenceEmbedding) {
+        // Simplified: just apply layer norm
+        seq_emb.layer_norm(1e-5);
+    }
+
+    /// Apply triangle attention for pair updates.
+    pub fn triangle_attention(&self, pair_emb: &mut PairEmbedding) {
+        let seq_len = pair_emb.sequence_length;
+
+        // Simplified triangle multiplication
+        for i in 0..seq_len {
+            for j in 0..seq_len {
+                let mut update = vec![0.0_f32; pair_emb.pair_dim];
+
+                // Outgoing edges
+                for k in 0..seq_len {
+                    if let (Some(ik), Some(kj)) = (pair_emb.get(i, k), pair_emb.get(k, j)) {
+                        for d in 0..pair_emb.pair_dim {
+                            update[d] += ik[d] * kj[d] * 0.01;
+                        }
+                    }
+                }
+
+                // Apply update
+                for d in 0..pair_emb.pair_dim {
+                    pair_emb.embeddings[i][j][d] += update[d];
+                }
+            }
+        }
+    }
+
+    /// Project embeddings through weight matrix.
+    fn project(&self, embeddings: &[Vec<f32>], weights: &[Vec<f32>]) -> Vec<Vec<f32>> {
+        embeddings
+            .iter()
+            .map(|emb| self.project_single(emb, weights))
+            .collect()
+    }
+
+    /// Project a single embedding.
+    fn project_single(&self, embedding: &[f32], weights: &[Vec<f32>]) -> Vec<f32> {
+        let mut output = vec![0.0_f32; weights.len()];
+        for (i, row) in weights.iter().enumerate() {
+            for (j, &w) in row.iter().enumerate() {
+                if j < embedding.len() {
+                    output[i] += embedding[j] * w;
+                }
+            }
+        }
+        output
+    }
+}
+
+/// Evoformer stack with multiple blocks.
+#[derive(Debug, Clone)]
+pub struct Evoformer {
+    /// Stack of Evoformer blocks
+    blocks: Vec<EvoformerBlock>,
+}
+
+impl Evoformer {
+    /// Create a new Evoformer with specified number of blocks.
+    #[must_use] 
+    pub fn new(num_blocks: usize) -> Self {
+        let blocks = (0..num_blocks).map(|_| EvoformerBlock::default()).collect();
+        Self { blocks }
+    }
+
+    /// Process sequence and pair representations through all blocks.
+    pub fn forward(&self, seq_emb: &mut SequenceEmbedding, pair_emb: &mut PairEmbedding) {
+        for block in &self.blocks {
+            block.row_attention(seq_emb, pair_emb);
+            block.column_attention(seq_emb);
+            block.triangle_attention(pair_emb);
+            pair_emb.outer_product_update(seq_emb);
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::encoder::ProteinEncoder;
+
+    #[test]
+    fn test_evoformer_block() {
+        let block = EvoformerBlock::default();
+        assert_eq!(block.num_heads, NUM_HEADS);
+        assert_eq!(block.head_dim, EMBEDDING_DIM / NUM_HEADS);
+    }
+
+    #[test]
+    fn test_evoformer_forward() {
+        let encoder = ProteinEncoder::new();
+        let mut seq_emb = encoder.encode("ACDEF");
+        let mut pair_emb = encoder.encode_pairs(5);
+
+        let evoformer = Evoformer::new(2);
+        evoformer.forward(&mut seq_emb, &mut pair_emb);
+
+        // Check that embeddings were modified
+        assert_eq!(seq_emb.sequence_length, 5);
+    }
+}