universal_annotator.py

# COPYRIGHT NOTICE
# This file is part of the "Universal Biomedical Skills" project.
# Copyright (c) 2026 MD BABU MIA, PhD <md.babu.mia@mssm.edu>
# All Rights Reserved.
#
# This code is proprietary and confidential.
# Unauthorized copying of this file, via any medium is strictly prohibited.
#
# Provenance: Authenticated by MD BABU MIA

import numpy as np
import pandas as pd
from typing import Optional, List, Union

try:
    import scanpy as sc
except ImportError:
    class MockScanpy:
        class tl:
            @staticmethod
            def score_genes(adata, gene_list, score_name):
                # Mock scoring: random scores
                adata.obs[score_name] = np.random.rand(len(adata.obs))
            
            @staticmethod
            def rank_genes_groups(adata, groupby, method):
                # Mock ranking
                adata.uns['rank_genes_groups'] = {
                    'names': {
                        group: [f"Gene_{i}" for i in range(10)] 
                        for group in adata.obs[groupby].unique()
                    }
                }

        class pp:
            @staticmethod
            def normalize_total(adata, target_sum): pass
            @staticmethod
            def log1p(adata): pass

        class AnnData:
            def __init__(self, X):
                self.X = X
                self.obs = pd.DataFrame(index=[f"Cell_{i}" for i in range(X.shape[0])])
                self.var = pd.DataFrame(index=[f"Gene_{i}" for i in range(X.shape[1])])
                self.var_names = self.var.index
                self.uns = {}

    sc = MockScanpy()
    print("Warning: 'scanpy' not found. Using MockScanpy for demonstration.")

class UniversalAnnotator:
    """
    A unified interface for Single-Cell Cell Type Annotation.
    Wraps multiple backend strategies (Markers, Reference, LLM).
    
    Demonstration of how to integrate tools like CellTypist and Seurat-like scoring
    into a Python class.
    """
    
    def __init__(self, adata: sc.AnnData):
        self.adata = adata
        
    def annotate_marker_based(self, marker_dict: dict, method: str = 'scoring'):
        """
        Classic annotation using known marker genes.
        Similar to Seurat's AddModuleScore or AUCell.
        
        Args:
            marker_dict: {'T-cell': ['CD3D', 'CD3E'], 'B-cell': ['CD79A']}
        """
        print(f"Annotating using {len(marker_dict)} cell types via Marker Genes...")
        
        for cell_type, markers in marker_dict.items():
            # Check which markers exist in dataset
            valid_markers = [m for m in markers if m in self.adata.var_names]
            if not valid_markers:
                print(f"Warning: No markers found for {cell_type}")
                continue
                
            # Calculate score (simple mean expression for demo)
            # In prod: use sc.tl.score_genes
            sc.tl.score_genes(self.adata, valid_markers, score_name=f"score_{cell_type}")
            
        # Assign max score as label
        score_cols = [f"score_{ct}" for ct in marker_dict.keys()]
        scores = self.adata.obs[score_cols]
        self.adata.obs['predicted_cell_type'] = scores.idxmax(axis=1).str.replace('score_', '')
        print("Marker-based annotation complete.")

    def annotate_with_celltypist(self, model_name: str = 'Immune_All_Low.pkl'):
        """
        Wrapper for CellTypist (Deep Learning based).
        Requires 'celltypist' library installed.
        """
        try:
            import celltypist
            print(f"Running CellTypist with model: {model_name}")
            
            # Normalize if needed
            if 'log1p' not in self.adata.uns:
                sc.pp.normalize_total(self.adata, target_sum=1e4)
                sc.pp.log1p(self.adata)
                
            predictions = celltypist.annotate(
                self.adata, 
                model=model_name, 
                majority_voting=True
            )
            
            self.adata.obs['celltypist_prediction'] = predictions.predicted_labels['predicted_labels']
            print("CellTypist annotation complete.")
            
        except ImportError:
            print("Error: 'celltypist' library not found. Please install via pip.")

    def annotate_with_llm(self, cluster_col: str, marker_num: int = 10, api_key: Optional[str] = None):
        """
        Simulation of LLM-based annotation (like mLLMCelltype).
        Extracts top markers per cluster and prompts an LLM to identify the cell type.
        """
        print("Extracting markers for LLM annotation...")
        
        # 1. Rank genes
        sc.tl.rank_genes_groups(self.adata, groupby=cluster_col, method='wilcoxon')
        
        # 2. Construct Prompt for each cluster
        groups = self.adata.obs[cluster_col].unique()
        for group in groups:
            # Get top 10 genes
            top_genes = self.adata.uns['rank_genes_groups']['names'][group][:marker_num]
            genes_str = ", ".join(top_genes)
            
            prompt = (
                f"Identify the cell type based on these top marker genes from a human tissue sample: "
                f"[{genes_str}]. Return only the cell type name."
            )
            
            # Mock LLM call
            print(f"Cluster {group} Prompt: {prompt}")
            # prediction = call_llm(prompt) 
            # self.adata.obs.loc[self.adata.obs[cluster_col] == group, 'llm_cell_type'] = prediction

# --- Example Usage ---
if __name__ == "__main__":
    # Mock AnnData
    print("Initializing Mock Single-Cell Data...")
    adata = sc.AnnData(np.random.rand(100, 50))
    # In real Scanpy, var_names is an Index. In our Mock, it's an Index too.
    # We can just assign to it directly or use a list.
    new_names = [f"Gene_{i}" for i in range(50)]
    new_names[0:3] = ['CD3D', 'CD3E', 'CD79A']
    adata.var_names = pd.Index(new_names)
    
    annotator = UniversalAnnotator(adata)
    
    # 1. Marker Based
    markers = {
        'T-cell': ['CD3D', 'CD3E'],
        'B-cell': ['CD79A']
    }
    annotator.annotate_marker_based(markers)
    print(adata.obs[['predicted_cell_type']].head())

    # 2. CellTypist (Will fail gracefully if not installed)
    annotator.annotate_with_celltypist()

__AUTHOR_SIGNATURE__ = "9a7f3c2e-MD-BABU-MIA-2026-MSSM-SECURE"