Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,579 @@
|
||||
//! 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"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user