osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
586 lines
19 KiB
Rust
586 lines
19 KiB
Rust
//! Cell type annotation algorithms.
|
|
//!
|
|
//! Implements marker-based annotation, transfer learning, and
|
|
//! foundation model-based cell type prediction.
|
|
|
|
use crate::CellAtlasError;
|
|
use cellatlas_shared::{
|
|
AnalysisConfig, AnnotationMethod, Cell, CellType, CellTypeAnnotation, ClusterInfo,
|
|
MarkerGeneSet, get_common_markers,
|
|
};
|
|
|
|
/// Annotate cells with cell types.
|
|
pub fn annotate_cells(
|
|
cells: &[Cell],
|
|
clusters: &[ClusterInfo],
|
|
config: &AnalysisConfig,
|
|
) -> Result<(Vec<Cell>, Vec<CellTypeAnnotation>), CellAtlasError> {
|
|
match config.annotation.method {
|
|
AnnotationMethod::MarkerBased => {
|
|
let markers = config
|
|
.annotation
|
|
.marker_genes
|
|
.clone()
|
|
.unwrap_or_else(get_common_markers);
|
|
marker_based_annotation(cells, clusters, &markers)
|
|
}
|
|
AnnotationMethod::TransferLearning => {
|
|
transfer_learning_annotation(cells, clusters, &config.annotation.reference)
|
|
}
|
|
AnnotationMethod::FoundationModel => foundation_model_annotation(cells, clusters),
|
|
}
|
|
}
|
|
|
|
/// Marker-based cell type annotation.
|
|
///
|
|
/// Assigns cell types based on expression of known marker genes.
|
|
fn marker_based_annotation(
|
|
cells: &[Cell],
|
|
_clusters: &[ClusterInfo],
|
|
marker_sets: &[MarkerGeneSet],
|
|
) -> Result<(Vec<Cell>, Vec<CellTypeAnnotation>), CellAtlasError> {
|
|
let mut annotated_cells = cells.to_vec();
|
|
let mut cell_type_counts: std::collections::HashMap<String, CellTypeStats> =
|
|
std::collections::HashMap::new();
|
|
|
|
// Build gene name to index map (simulated)
|
|
let gene_map = build_gene_map();
|
|
|
|
// Annotate each cell
|
|
for cell in &mut annotated_cells {
|
|
let best_annotation = find_best_cell_type(cell, marker_sets, &gene_map);
|
|
cell.cell_type = best_annotation.clone();
|
|
|
|
if let Some(ct) = best_annotation {
|
|
let stats = cell_type_counts
|
|
.entry(ct.name.clone())
|
|
.or_insert_with(|| CellTypeStats {
|
|
count: 0,
|
|
total_confidence: 0.0,
|
|
clusters: std::collections::HashSet::new(),
|
|
});
|
|
stats.count += 1;
|
|
stats.total_confidence += ct.confidence;
|
|
if let Some(cluster_id) = cell.cluster_id {
|
|
stats.clusters.insert(cluster_id);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Build cell type annotations summary
|
|
let total_cells = annotated_cells.len();
|
|
let cell_types: Vec<CellTypeAnnotation> = cell_type_counts
|
|
.into_iter()
|
|
.map(|(name, stats)| CellTypeAnnotation {
|
|
name,
|
|
n_cells: stats.count,
|
|
percentage: (stats.count as f32 / total_cells as f32) * 100.0,
|
|
avg_confidence: stats.total_confidence / stats.count as f32,
|
|
clusters: stats.clusters.into_iter().collect(),
|
|
})
|
|
.collect();
|
|
|
|
Ok((annotated_cells, cell_types))
|
|
}
|
|
|
|
/// Transfer learning annotation using a reference dataset.
|
|
fn transfer_learning_annotation(
|
|
cells: &[Cell],
|
|
_clusters: &[ClusterInfo],
|
|
reference: &Option<String>,
|
|
) -> Result<(Vec<Cell>, Vec<CellTypeAnnotation>), CellAtlasError> {
|
|
// For demo, simulate transfer learning
|
|
let reference_name = reference.as_deref().unwrap_or("Human Cell Atlas");
|
|
tracing::info!("Using reference dataset: {}", reference_name);
|
|
|
|
// Create reference cell type signatures (simulated)
|
|
let reference_signatures = get_reference_signatures();
|
|
|
|
let mut annotated_cells = cells.to_vec();
|
|
let mut cell_type_counts: std::collections::HashMap<String, CellTypeStats> =
|
|
std::collections::HashMap::new();
|
|
|
|
for cell in &mut annotated_cells {
|
|
// Match cell to reference signatures
|
|
let best_match = match_to_reference(cell, &reference_signatures);
|
|
cell.cell_type = Some(best_match.clone());
|
|
|
|
let stats = cell_type_counts
|
|
.entry(best_match.name.clone())
|
|
.or_insert_with(|| CellTypeStats {
|
|
count: 0,
|
|
total_confidence: 0.0,
|
|
clusters: std::collections::HashSet::new(),
|
|
});
|
|
stats.count += 1;
|
|
stats.total_confidence += best_match.confidence;
|
|
if let Some(cluster_id) = cell.cluster_id {
|
|
stats.clusters.insert(cluster_id);
|
|
}
|
|
}
|
|
|
|
let total_cells = annotated_cells.len();
|
|
let cell_types: Vec<CellTypeAnnotation> = cell_type_counts
|
|
.into_iter()
|
|
.map(|(name, stats)| CellTypeAnnotation {
|
|
name,
|
|
n_cells: stats.count,
|
|
percentage: (stats.count as f32 / total_cells as f32) * 100.0,
|
|
avg_confidence: stats.total_confidence / stats.count as f32,
|
|
clusters: stats.clusters.into_iter().collect(),
|
|
})
|
|
.collect();
|
|
|
|
Ok((annotated_cells, cell_types))
|
|
}
|
|
|
|
/// Foundation model-based annotation using `CellAtlas` transformer.
|
|
fn foundation_model_annotation(
|
|
cells: &[Cell],
|
|
clusters: &[ClusterInfo],
|
|
) -> Result<(Vec<Cell>, Vec<CellTypeAnnotation>), CellAtlasError> {
|
|
use crate::cell_transformer::{CellTransformer, CellTransformerConfig};
|
|
|
|
tracing::info!("Using CellAtlas foundation model for annotation");
|
|
|
|
// Initialize transformer
|
|
let config = CellTransformerConfig {
|
|
num_genes: 2000,
|
|
hidden_dim: 128,
|
|
num_heads: 4,
|
|
num_layers: 4,
|
|
dropout: 0.1,
|
|
max_seq_len: 512,
|
|
};
|
|
let transformer = CellTransformer::new(config);
|
|
|
|
// Encode cells
|
|
let gene_indices: Vec<Vec<usize>> = cells
|
|
.iter()
|
|
.map(|c| c.expression.gene_indices.clone())
|
|
.collect();
|
|
let values: Vec<Vec<f32>> = cells.iter().map(|c| c.expression.values.clone()).collect();
|
|
|
|
let embeddings = transformer.encode(&gene_indices, &values);
|
|
|
|
// Classify based on embeddings (simulated classifier)
|
|
let cell_type_predictions = classify_embeddings(&embeddings, clusters);
|
|
|
|
let mut annotated_cells = cells.to_vec();
|
|
let mut cell_type_counts: std::collections::HashMap<String, CellTypeStats> =
|
|
std::collections::HashMap::new();
|
|
|
|
for (i, cell) in annotated_cells.iter_mut().enumerate() {
|
|
if i < cell_type_predictions.len() {
|
|
let prediction = &cell_type_predictions[i];
|
|
cell.cell_type = Some(prediction.clone());
|
|
|
|
let stats = cell_type_counts
|
|
.entry(prediction.name.clone())
|
|
.or_insert_with(|| CellTypeStats {
|
|
count: 0,
|
|
total_confidence: 0.0,
|
|
clusters: std::collections::HashSet::new(),
|
|
});
|
|
stats.count += 1;
|
|
stats.total_confidence += prediction.confidence;
|
|
if let Some(cluster_id) = cell.cluster_id {
|
|
stats.clusters.insert(cluster_id);
|
|
}
|
|
}
|
|
}
|
|
|
|
let total_cells = annotated_cells.len();
|
|
let cell_types: Vec<CellTypeAnnotation> = cell_type_counts
|
|
.into_iter()
|
|
.map(|(name, stats)| CellTypeAnnotation {
|
|
name,
|
|
n_cells: stats.count,
|
|
percentage: (stats.count as f32 / total_cells as f32) * 100.0,
|
|
avg_confidence: stats.total_confidence / stats.count as f32,
|
|
clusters: stats.clusters.into_iter().collect(),
|
|
})
|
|
.collect();
|
|
|
|
Ok((annotated_cells, cell_types))
|
|
}
|
|
|
|
/// Helper struct for tracking cell type statistics.
|
|
#[derive(Debug)]
|
|
struct CellTypeStats {
|
|
count: usize,
|
|
total_confidence: f32,
|
|
clusters: std::collections::HashSet<usize>,
|
|
}
|
|
|
|
/// Build gene name to index map.
|
|
fn build_gene_map() -> std::collections::HashMap<String, usize> {
|
|
let common_genes = vec![
|
|
"CD3D", "CD3E", "CD3G", "CD4", "CD8A", "CD8B", "CD19", "MS4A1", "CD79A", "CD79B", "CD14",
|
|
"LYZ", "CST3", "FCGR3A", "GNLY", "NKG7", "KLRD1", "PRF1", "FCER1A", "CD1C", "CLEC10A",
|
|
"IL7R", "CCR7", "S100A4", "PPBP", "PF4", "GP9", "HBA1", "HBA2", "HBB",
|
|
];
|
|
|
|
common_genes
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, &gene)| (gene.to_string(), i))
|
|
.collect()
|
|
}
|
|
|
|
/// Find best cell type for a cell based on marker expression.
|
|
fn find_best_cell_type(
|
|
cell: &Cell,
|
|
marker_sets: &[MarkerGeneSet],
|
|
gene_map: &std::collections::HashMap<String, usize>,
|
|
) -> Option<CellType> {
|
|
use rand::SeedableRng;
|
|
use rand_distr::{Distribution, Normal};
|
|
|
|
let mut rng = rand::rngs::StdRng::seed_from_u64(cell.id.len() as u64);
|
|
let noise = Normal::new(0.0_f32, 0.1).unwrap();
|
|
|
|
let mut best_score = 0.0_f32;
|
|
let mut best_type: Option<CellType> = None;
|
|
|
|
for marker_set in marker_sets {
|
|
// Count positive marker expression
|
|
let mut positive_count = 0;
|
|
let mut positive_total = 0.0;
|
|
|
|
for marker in &marker_set.positive_markers {
|
|
if let Some(&gene_idx) = gene_map.get(marker)
|
|
&& gene_idx < cell.expression.gene_indices.len()
|
|
{
|
|
let expr = cell.expression.get(gene_idx);
|
|
if expr > 0.0 {
|
|
positive_count += 1;
|
|
positive_total += expr;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Check negative markers
|
|
let mut negative_count = 0;
|
|
for marker in &marker_set.negative_markers {
|
|
if let Some(&gene_idx) = gene_map.get(marker) {
|
|
let expr = cell.expression.get(gene_idx);
|
|
if expr > 0.0 {
|
|
negative_count += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Calculate score
|
|
let positive_score = if marker_set.positive_markers.is_empty() {
|
|
0.5
|
|
} else {
|
|
positive_count as f32 / marker_set.positive_markers.len() as f32
|
|
};
|
|
|
|
let negative_penalty = if marker_set.negative_markers.is_empty() {
|
|
0.0
|
|
} else {
|
|
negative_count as f32 / marker_set.negative_markers.len() as f32 * 0.5
|
|
};
|
|
|
|
let score = (positive_score - negative_penalty + noise.sample(&mut rng)).clamp(0.0, 1.0);
|
|
|
|
if score > best_score {
|
|
best_score = score;
|
|
best_type = Some(CellType {
|
|
name: marker_set.cell_type.clone(),
|
|
confidence: score,
|
|
lineage: vec!["Immune".to_string(), marker_set.cell_type.clone()],
|
|
markers: marker_set.positive_markers.clone(),
|
|
});
|
|
}
|
|
}
|
|
|
|
best_type
|
|
}
|
|
|
|
/// Get reference cell type signatures for transfer learning.
|
|
fn get_reference_signatures() -> Vec<ReferenceSignature> {
|
|
vec![
|
|
ReferenceSignature {
|
|
cell_type: "CD4+ T cell".to_string(),
|
|
marker_weights: vec![
|
|
("CD3D".to_string(), 1.0),
|
|
("CD4".to_string(), 1.0),
|
|
("IL7R".to_string(), 0.8),
|
|
],
|
|
lineage: vec![
|
|
"Immune".to_string(),
|
|
"T cell".to_string(),
|
|
"CD4+ T cell".to_string(),
|
|
],
|
|
},
|
|
ReferenceSignature {
|
|
cell_type: "CD8+ T cell".to_string(),
|
|
marker_weights: vec![
|
|
("CD3D".to_string(), 1.0),
|
|
("CD8A".to_string(), 1.0),
|
|
("GZMK".to_string(), 0.7),
|
|
],
|
|
lineage: vec![
|
|
"Immune".to_string(),
|
|
"T cell".to_string(),
|
|
"CD8+ T cell".to_string(),
|
|
],
|
|
},
|
|
ReferenceSignature {
|
|
cell_type: "B cell".to_string(),
|
|
marker_weights: vec![
|
|
("CD19".to_string(), 1.0),
|
|
("MS4A1".to_string(), 1.0),
|
|
("CD79A".to_string(), 0.9),
|
|
],
|
|
lineage: vec!["Immune".to_string(), "B cell".to_string()],
|
|
},
|
|
ReferenceSignature {
|
|
cell_type: "Monocyte".to_string(),
|
|
marker_weights: vec![
|
|
("CD14".to_string(), 1.0),
|
|
("LYZ".to_string(), 0.9),
|
|
("CST3".to_string(), 0.8),
|
|
],
|
|
lineage: vec![
|
|
"Immune".to_string(),
|
|
"Myeloid".to_string(),
|
|
"Monocyte".to_string(),
|
|
],
|
|
},
|
|
ReferenceSignature {
|
|
cell_type: "NK cell".to_string(),
|
|
marker_weights: vec![
|
|
("GNLY".to_string(), 1.0),
|
|
("NKG7".to_string(), 1.0),
|
|
("KLRD1".to_string(), 0.8),
|
|
],
|
|
lineage: vec!["Immune".to_string(), "NK cell".to_string()],
|
|
},
|
|
ReferenceSignature {
|
|
cell_type: "Dendritic cell".to_string(),
|
|
marker_weights: vec![("FCER1A".to_string(), 1.0), ("CD1C".to_string(), 0.9)],
|
|
lineage: vec![
|
|
"Immune".to_string(),
|
|
"Myeloid".to_string(),
|
|
"Dendritic cell".to_string(),
|
|
],
|
|
},
|
|
]
|
|
}
|
|
|
|
/// Reference cell type signature.
|
|
#[derive(Debug)]
|
|
struct ReferenceSignature {
|
|
cell_type: String,
|
|
marker_weights: Vec<(String, f32)>,
|
|
lineage: Vec<String>,
|
|
}
|
|
|
|
/// Match a cell to reference signatures.
|
|
fn match_to_reference(cell: &Cell, references: &[ReferenceSignature]) -> CellType {
|
|
use rand::SeedableRng;
|
|
use rand_distr::{Distribution, Normal};
|
|
|
|
let mut rng = rand::rngs::StdRng::seed_from_u64(cell.id.len() as u64 * 17);
|
|
let noise = Normal::new(0.0_f32, 0.1).unwrap();
|
|
|
|
let gene_map = build_gene_map();
|
|
|
|
let mut best_score = 0.0_f32;
|
|
let mut best_ref = &references[0];
|
|
|
|
for reference in references {
|
|
let mut score = 0.0_f32;
|
|
let mut total_weight = 0.0_f32;
|
|
|
|
for (gene, weight) in &reference.marker_weights {
|
|
if let Some(&gene_idx) = gene_map.get(gene) {
|
|
let expr = cell.expression.get(gene_idx);
|
|
if expr > 0.0 {
|
|
score += weight;
|
|
}
|
|
total_weight += weight;
|
|
}
|
|
}
|
|
|
|
let normalized_score = if total_weight > 0.0 {
|
|
score / total_weight
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
let noisy_score = (normalized_score + noise.sample(&mut rng)).clamp(0.0, 1.0);
|
|
|
|
if noisy_score > best_score {
|
|
best_score = noisy_score;
|
|
best_ref = reference;
|
|
}
|
|
}
|
|
|
|
CellType {
|
|
name: best_ref.cell_type.clone(),
|
|
confidence: best_score,
|
|
lineage: best_ref.lineage.clone(),
|
|
markers: best_ref
|
|
.marker_weights
|
|
.iter()
|
|
.map(|(g, _)| g.clone())
|
|
.collect(),
|
|
}
|
|
}
|
|
|
|
/// Classify cell embeddings using a simple classifier.
|
|
fn classify_embeddings(embeddings: &[Vec<f32>], _clusters: &[ClusterInfo]) -> Vec<CellType> {
|
|
use rand::SeedableRng;
|
|
use rand_distr::{Distribution, Normal};
|
|
|
|
let cell_types = [
|
|
("CD4+ T cell", vec!["Immune", "T cell", "CD4+ T cell"]),
|
|
("CD8+ T cell", vec!["Immune", "T cell", "CD8+ T cell"]),
|
|
("B cell", vec!["Immune", "B cell"]),
|
|
("Monocyte", vec!["Immune", "Myeloid", "Monocyte"]),
|
|
("NK cell", vec!["Immune", "NK cell"]),
|
|
("Dendritic cell", vec!["Immune", "Myeloid", "DC"]),
|
|
];
|
|
|
|
let mut rng = rand::rngs::StdRng::seed_from_u64(42);
|
|
let confidence_noise = Normal::new(0.0_f32, 0.1).unwrap();
|
|
|
|
embeddings
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, emb)| {
|
|
// Simple classification based on embedding features
|
|
let sum: f32 = emb.iter().sum();
|
|
let type_idx = (i + (sum.abs() as usize)) % cell_types.len();
|
|
let (name, lineage) = &cell_types[type_idx];
|
|
|
|
CellType {
|
|
name: name.to_string(),
|
|
confidence: (0.75 + confidence_noise.sample(&mut rng)).clamp(0.5, 0.99),
|
|
lineage: lineage
|
|
.iter()
|
|
.map(std::string::ToString::to_string)
|
|
.collect(),
|
|
markers: vec![],
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use cellatlas_shared::{AnnotationConfig, Embedding2D, QualityMetrics, SparseExpression};
|
|
|
|
fn create_test_cell(id: &str, gene_indices: Vec<usize>, values: Vec<f32>) -> Cell {
|
|
Cell {
|
|
id: id.to_string(),
|
|
barcode: None,
|
|
expression: SparseExpression {
|
|
gene_indices,
|
|
values: values.clone(),
|
|
num_genes: 100,
|
|
},
|
|
cell_type: None,
|
|
state: None,
|
|
qc_metrics: QualityMetrics {
|
|
n_genes: values.len(),
|
|
total_counts: values.iter().sum(),
|
|
pct_mito: 3.0,
|
|
pct_ribo: 10.0,
|
|
doublet_score: None,
|
|
},
|
|
spatial_coords: None,
|
|
cluster_id: Some(0),
|
|
embedding: Some(Embedding2D { x: 0.0, y: 0.0 }),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_marker_based_annotation() {
|
|
let cells = vec![
|
|
create_test_cell("t_cell", vec![0, 1, 2], vec![5.0, 4.0, 3.0]),
|
|
create_test_cell("b_cell", vec![6, 7, 8], vec![6.0, 5.0, 4.0]),
|
|
];
|
|
|
|
let clusters = vec![ClusterInfo {
|
|
id: 0,
|
|
n_cells: 2,
|
|
markers: vec!["CD3D".to_string()],
|
|
centroid: None,
|
|
}];
|
|
|
|
let markers = get_common_markers();
|
|
|
|
let result = marker_based_annotation(&cells, &clusters, &markers);
|
|
assert!(result.is_ok());
|
|
|
|
let (annotated, cell_types) = result.unwrap();
|
|
assert_eq!(annotated.len(), 2);
|
|
assert!(!cell_types.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_transfer_learning_annotation() {
|
|
let cells = vec![
|
|
create_test_cell("cell_1", vec![0, 1], vec![5.0, 4.0]),
|
|
create_test_cell("cell_2", vec![6, 7], vec![6.0, 5.0]),
|
|
];
|
|
|
|
let clusters = vec![ClusterInfo {
|
|
id: 0,
|
|
n_cells: 2,
|
|
markers: vec![],
|
|
centroid: None,
|
|
}];
|
|
|
|
let result = transfer_learning_annotation(&cells, &clusters, &None);
|
|
assert!(result.is_ok());
|
|
|
|
let (annotated, cell_types) = result.unwrap();
|
|
assert!(annotated.iter().all(|c| c.cell_type.is_some()));
|
|
assert!(!cell_types.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_foundation_model_annotation() {
|
|
let cells = vec![
|
|
create_test_cell("cell_1", vec![0, 1, 2], vec![1.0, 2.0, 1.5]),
|
|
create_test_cell("cell_2", vec![3, 4, 5], vec![2.0, 1.0, 0.5]),
|
|
];
|
|
|
|
let clusters = vec![ClusterInfo {
|
|
id: 0,
|
|
n_cells: 2,
|
|
markers: vec![],
|
|
centroid: None,
|
|
}];
|
|
|
|
let result = foundation_model_annotation(&cells, &clusters);
|
|
assert!(result.is_ok());
|
|
|
|
let (annotated, cell_types) = result.unwrap();
|
|
assert_eq!(annotated.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_reference_signatures() {
|
|
let signatures = get_reference_signatures();
|
|
assert!(!signatures.is_empty());
|
|
assert!(signatures.iter().any(|s| s.cell_type.contains("T cell")));
|
|
}
|
|
|
|
#[test]
|
|
fn test_build_gene_map() {
|
|
let map = build_gene_map();
|
|
assert!(map.contains_key("CD3D"));
|
|
assert!(map.contains_key("CD19"));
|
|
assert!(map.contains_key("GNLY"));
|
|
}
|
|
}
|