rustytorch/demos/rtx-cellatlas-demo/src/graph_encoder.rs

//! Graph Neural Network encoder for spatial transcriptomics.
//!
//! Implements a GNN to capture spatial relationships between cells
//! in tissue sections.

use cellatlas_shared::{Cell, SpatialCoords};

/// Configuration for the Graph Encoder.
#[derive(Debug, Clone)]
pub struct GraphEncoderConfig {
    /// Input dimension (cell embedding size)
    pub input_dim: usize,
    /// Hidden dimension
    pub hidden_dim: usize,
    /// Output dimension
    pub output_dim: usize,
    /// Number of GNN layers
    pub num_layers: usize,
    /// Number of attention heads
    pub num_heads: usize,
    /// Distance threshold for building graph edges
    pub distance_threshold: f32,
    /// Maximum number of neighbors per cell
    pub max_neighbors: usize,
}

impl Default for GraphEncoderConfig {
    fn default() -> Self {
        Self {
            input_dim: 256,
            hidden_dim: 128,
            output_dim: 256,
            num_layers: 3,
            num_heads: 4,
            distance_threshold: 100.0, // pixels/microns
            max_neighbors: 15,
        }
    }
}

/// Graph Neural Network for spatial encoding.
#[derive(Debug)]
pub struct GraphEncoder {
    config: GraphEncoderConfig,
    layers: Vec<GraphAttentionLayer>,
    output_proj: OutputProjection,
}

impl GraphEncoder {
    /// Create a new Graph Encoder.
    #[must_use]
    pub fn new(config: GraphEncoderConfig) -> Self {
        let mut layers = Vec::with_capacity(config.num_layers);

        let first_layer =
            GraphAttentionLayer::new(config.input_dim, config.hidden_dim, config.num_heads);
        layers.push(first_layer);

        for _ in 1..config.num_layers {
            let layer =
                GraphAttentionLayer::new(config.hidden_dim, config.hidden_dim, config.num_heads);
            layers.push(layer);
        }

        let output_proj = OutputProjection::new(config.hidden_dim, config.output_dim);

        Self {
            config,
            layers,
            output_proj,
        }
    }

    /// Build spatial graph from cell coordinates.
    #[must_use]
    pub fn build_graph(&self, cells: &[Cell]) -> SpatialGraph {
        let spatial_cells: Vec<(usize, &SpatialCoords)> = cells
            .iter()
            .enumerate()
            .filter_map(|(i, c)| c.spatial_coords.as_ref().map(|s| (i, s)))
            .collect();

        let n = spatial_cells.len();
        let mut adjacency = vec![vec![]; n];
        let mut edge_weights = vec![vec![]; n];

        // Build k-NN graph with distance threshold
        for i in 0..n {
            let (_, coord_i) = &spatial_cells[i];
            let mut neighbors: Vec<(usize, f32)> = Vec::new();

            for j in 0..n {
                if i == j {
                    continue;
                }
                let (_, coord_j) = &spatial_cells[j];
                let dist = euclidean_distance(coord_i, coord_j);

                if dist <= self.config.distance_threshold {
                    neighbors.push((j, dist));
                }
            }

            // Sort by distance and keep top k
            neighbors.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
            neighbors.truncate(self.config.max_neighbors);

            // Convert distances to edge weights (inverse distance)
            for (j, dist) in neighbors {
                adjacency[i].push(j);
                edge_weights[i].push(1.0 / (dist + 1.0));
            }
        }

        SpatialGraph {
            cell_indices: spatial_cells.iter().map(|(i, _)| *i).collect(),
            adjacency,
            edge_weights,
        }
    }

    /// Encode cells using spatial context.
    #[must_use]
    pub fn encode(&self, cell_embeddings: &[Vec<f32>], graph: &SpatialGraph) -> Vec<Vec<f32>> {
        let n = graph.cell_indices.len();
        if n == 0 {
            return vec![];
        }

        // Get embeddings for spatial cells
        let mut hidden: Vec<Vec<f32>> = graph
            .cell_indices
            .iter()
            .map(|&i| {
                cell_embeddings
                    .get(i)
                    .cloned()
                    .unwrap_or_else(|| vec![0.0; self.config.input_dim])
            })
            .collect();

        // Apply GNN layers
        for layer in &self.layers {
            hidden = layer.forward(&hidden, &graph.adjacency, &graph.edge_weights);
        }

        // Output projection
        self.output_proj.forward(&hidden)
    }
}

/// Spatial graph structure.
#[derive(Debug, Clone)]
pub struct SpatialGraph {
    /// Cell indices in the original dataset
    pub cell_indices: Vec<usize>,
    /// Adjacency list (neighbors for each cell)
    pub adjacency: Vec<Vec<usize>>,
    /// Edge weights (for attention)
    pub edge_weights: Vec<Vec<f32>>,
}

impl SpatialGraph {
    /// Get the number of nodes.
    #[must_use]
    pub fn num_nodes(&self) -> usize {
        self.cell_indices.len()
    }

    /// Get the number of edges.
    #[must_use]
    pub fn num_edges(&self) -> usize {
        self.adjacency.iter().map(std::vec::Vec::len).sum()
    }
}

/// Graph Attention Layer.
#[derive(Debug)]
struct GraphAttentionLayer {
    input_dim: usize,
    output_dim: usize,
    num_heads: usize,
    head_dim: usize,
    // Attention weights
    w_query: Vec<Vec<f32>>,
    w_key: Vec<Vec<f32>>,
    w_value: Vec<Vec<f32>>,
    w_out: Vec<Vec<f32>>,
    // Layer norm
    ln_scale: Vec<f32>,
    ln_bias: Vec<f32>,
}

impl GraphAttentionLayer {
    fn new(input_dim: usize, output_dim: usize, num_heads: usize) -> Self {
        use rand::SeedableRng;
        use rand_distr::{Distribution, Normal};

        let mut rng = rand::rngs::StdRng::seed_from_u64(42);
        let std = (2.0 / input_dim as f32).sqrt();
        let normal = Normal::new(0.0_f32, std).unwrap();

        let head_dim = output_dim / num_heads;

        let mut init_weights = |rows: usize, cols: usize| -> Vec<Vec<f32>> {
            (0..rows)
                .map(|_| (0..cols).map(|_| normal.sample(&mut rng)).collect())
                .collect()
        };

        Self {
            input_dim,
            output_dim,
            num_heads,
            head_dim,
            w_query: init_weights(input_dim, output_dim),
            w_key: init_weights(input_dim, output_dim),
            w_value: init_weights(input_dim, output_dim),
            w_out: init_weights(output_dim, output_dim),
            ln_scale: vec![1.0; output_dim],
            ln_bias: vec![0.0; output_dim],
        }
    }

    fn forward(
        &self,
        hidden: &[Vec<f32>],
        adjacency: &[Vec<usize>],
        edge_weights: &[Vec<f32>],
    ) -> Vec<Vec<f32>> {
        let n = hidden.len();
        let scale = (self.head_dim as f32).sqrt();

        // Compute Q, K, V for all nodes
        let queries = self.matmul(hidden, &self.w_query);
        let keys = self.matmul(hidden, &self.w_key);
        let values = self.matmul(hidden, &self.w_value);

        // Graph attention
        let mut output = vec![vec![0.0; self.output_dim]; n];

        for i in 0..n {
            let neighbors = &adjacency[i];
            if neighbors.is_empty() {
                // No neighbors, just use self
                output[i] = values[i].clone();
                continue;
            }

            let weights = &edge_weights[i];

            // Compute attention scores
            let scores: Vec<f32> = neighbors
                .iter()
                .zip(weights.iter())
                .map(|(&j, &w)| {
                    let qk: f32 = (0..self.output_dim)
                        .map(|d| queries[i][d] * keys[j][d])
                        .sum();
                    qk / scale + w.ln() // Add edge weight as bias
                })
                .collect();

            // 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::<f32>() + 1e-8;
            let attn_weights: Vec<f32> = exp_scores.iter().map(|s| s / sum_exp).collect();

            // Weighted sum of neighbor values
            for (idx, j) in neighbors.iter().enumerate() {
                let j = *j;
                for d in 0..self.output_dim {
                    output[i][d] += attn_weights[idx] * values[j][d];
                }
            }
        }

        // Output projection + layer norm
        let projected = self.matmul(&output, &self.w_out);
        self.layer_norm(&projected)
    }

    fn matmul(&self, a: &[Vec<f32>], b: &[Vec<f32>]) -> Vec<Vec<f32>> {
        let m = a.len();
        let n = b[0].len();
        let k = a[0].len().min(b.len());

        let mut result = vec![vec![0.0; n]; m];
        for i in 0..m {
            for j in 0..n {
                for l in 0..k {
                    result[i][j] += a[i][l] * b[l][j];
                }
            }
        }
        result
    }

    fn layer_norm(&self, hidden: &[Vec<f32>]) -> Vec<Vec<f32>> {
        hidden
            .iter()
            .map(|row| {
                let mean: f32 = row.iter().sum::<f32>() / row.len() as f32;
                let var: f32 =
                    row.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / row.len() as f32;
                let std = (var + 1e-5).sqrt();
                row.iter()
                    .enumerate()
                    .map(|(i, &x)| (x - mean) / std * self.ln_scale[i] + self.ln_bias[i])
                    .collect()
            })
            .collect()
    }
}

/// Output projection layer.
#[derive(Debug)]
struct OutputProjection {
    weights: Vec<Vec<f32>>,
}

impl OutputProjection {
    fn new(input_dim: usize, output_dim: usize) -> Self {
        use rand::SeedableRng;
        use rand_distr::{Distribution, Normal};

        let mut rng = rand::rngs::StdRng::seed_from_u64(789);
        let std = (2.0 / input_dim as f32).sqrt();
        let normal = Normal::new(0.0_f32, std).unwrap();

        let weights: Vec<Vec<f32>> = (0..input_dim)
            .map(|_| (0..output_dim).map(|_| normal.sample(&mut rng)).collect())
            .collect();

        Self { weights }
    }

    fn forward(&self, hidden: &[Vec<f32>]) -> Vec<Vec<f32>> {
        let m = hidden.len();
        let n = self.weights[0].len();
        let k = hidden[0].len().min(self.weights.len());

        let mut result = vec![vec![0.0; n]; m];
        for i in 0..m {
            for j in 0..n {
                for l in 0..k {
                    result[i][j] += hidden[i][l] * self.weights[l][j];
                }
            }
        }
        result
    }
}

/// Compute Euclidean distance between two spatial coordinates.
fn euclidean_distance(a: &SpatialCoords, b: &SpatialCoords) -> f32 {
    let dx = a.x - b.x;
    let dy = a.y - b.y;
    let dz = match (a.z, b.z) {
        (Some(az), Some(bz)) => az - bz,
        _ => 0.0,
    };
    (dx * dx + dy * dy + dz * dz).sqrt()
}

#[cfg(test)]
mod tests {
    use super::*;
    use cellatlas_shared::{QualityMetrics, SparseExpression};

    fn create_test_cell(id: &str, x: f32, y: f32) -> Cell {
        Cell {
            id: id.to_string(),
            barcode: None,
            expression: SparseExpression {
                gene_indices: vec![0, 1, 2],
                values: vec![1.0, 2.0, 1.5],
                num_genes: 100,
            },
            cell_type: None,
            state: None,
            qc_metrics: QualityMetrics {
                n_genes: 1000,
                total_counts: 5000.0,
                pct_mito: 3.0,
                pct_ribo: 10.0,
                doublet_score: None,
            },
            spatial_coords: Some(SpatialCoords {
                x,
                y,
                z: None,
                section_id: Some(0),
            }),
            cluster_id: None,
            embedding: None,
        }
    }

    #[test]
    fn test_graph_encoder_creation() {
        let config = GraphEncoderConfig::default();
        let encoder = GraphEncoder::new(config);
        assert_eq!(encoder.config.num_layers, 3);
    }

    #[test]
    fn test_build_graph() {
        let config = GraphEncoderConfig {
            distance_threshold: 50.0,
            max_neighbors: 3,
            ..Default::default()
        };
        let encoder = GraphEncoder::new(config);

        let cells = vec![
            create_test_cell("cell_0", 0.0, 0.0),
            create_test_cell("cell_1", 10.0, 10.0),
            create_test_cell("cell_2", 100.0, 100.0),
            create_test_cell("cell_3", 5.0, 5.0),
        ];

        let graph = encoder.build_graph(&cells);
        assert_eq!(graph.num_nodes(), 4);
        // cell_0, cell_1, cell_3 should be connected (within 50 units)
        // cell_2 is far away
        assert!(!graph.adjacency[0].is_empty());
    }

    #[test]
    fn test_encode_with_graph() {
        let config = GraphEncoderConfig {
            input_dim: 32,
            hidden_dim: 16,
            output_dim: 32,
            num_layers: 2,
            num_heads: 2,
            distance_threshold: 100.0,
            max_neighbors: 5,
        };
        let encoder = GraphEncoder::new(config);

        let cells = vec![
            create_test_cell("cell_0", 0.0, 0.0),
            create_test_cell("cell_1", 10.0, 10.0),
            create_test_cell("cell_2", 20.0, 20.0),
        ];

        let cell_embeddings: Vec<Vec<f32>> = (0..3)
            .map(|_| (0..32).map(|i| i as f32 * 0.1).collect())
            .collect();

        let graph = encoder.build_graph(&cells);
        let encoded = encoder.encode(&cell_embeddings, &graph);

        assert_eq!(encoded.len(), 3);
        for emb in &encoded {
            assert_eq!(emb.len(), 32);
        }
    }

    #[test]
    fn test_euclidean_distance() {
        let a = SpatialCoords {
            x: 0.0,
            y: 0.0,
            z: None,
            section_id: None,
        };
        let b = SpatialCoords {
            x: 3.0,
            y: 4.0,
            z: None,
            section_id: None,
        };
        assert!((euclidean_distance(&a, &b) - 5.0).abs() < 0.001);
    }
}