linkml-store 0.3.0__py3-none-any.whl

linkml_store/utils/embedding_utils.py

@@ -0,0 +1,299 @@
+"""Utilities for extracting and processing embeddings from indexed collections."""
+
+import logging
+from typing import Dict, List, Optional, Tuple, Union
+import numpy as np
+from dataclasses import dataclass, field
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class EmbeddingData:
+    """Container for embedding data from collections."""
+
+    vectors: np.ndarray
+    metadata: List[Dict]
+    collection_names: List[str]
+    collection_indices: List[int]
+    object_ids: List[str]
+
+    @property
+    def n_samples(self) -> int:
+        """Number of samples."""
+        return len(self.vectors)
+
+    @property
+    def n_dimensions(self) -> int:
+        """Number of dimensions in embeddings."""
+        return self.vectors.shape[1] if len(self.vectors.shape) > 1 else 0
+
+    def get_metadata_values(self, field: str) -> List:
+        """Extract values for a specific metadata field."""
+        return [m.get(field) for m in self.metadata]
+
+
+def extract_embeddings_from_collection(
+    collection,
+    index_name: str = None,
+    limit: Optional[int] = None,
+    include_metadata: bool = True,
+    metadata_fields: Optional[List[str]] = None
+) -> EmbeddingData:
+    """
+    Extract embeddings from an indexed collection.
+
+    Args:
+        collection: LinkML collection object
+        index_name: Name of the index to use (defaults to first available)
+        limit: Maximum number of embeddings to extract
+        include_metadata: Whether to include source object metadata
+        metadata_fields: Specific metadata fields to include (None = all)
+
+    Returns:
+        EmbeddingData object containing vectors and metadata
+    """
+    # Get the index name - handle collections without loaded indexers
+    if index_name is None:
+        # Try to find index collections directly
+        db = collection.parent
+        all_collections = db.list_collection_names()
+        # TODO: use the indexer metadata to find the index name
+        index_prefix = f"internal__index__{collection.alias}__"
+        index_collections = [c for c in all_collections if c.startswith(index_prefix)]
+
+        if not index_collections:
+            raise ValueError(f"Collection {collection.alias} has no indexes")
+
+        # Extract index name from first index collection
+        index_name = index_collections[0].replace(index_prefix, "")
+        if len(index_collections) > 1:
+            logger.warning(f"Multiple indexes found, using: {index_name}")
+
+    # Get the index collection
+    index_collection_name = f"internal__index__{collection.alias}__{index_name}"
+    index_collection = collection.parent.get_collection(index_collection_name)
+
+    # Query the index collection
+    query_result = index_collection.find(limit=limit)
+
+    if query_result.num_rows == 0:
+        raise ValueError(f"No indexed data found in {index_collection_name}")
+
+    vectors = []
+    metadata = []
+    object_ids = []
+
+    for row in query_result.rows:
+        # Extract vector (usually stored in __index__ field)
+        vector = row.get("__index__")
+        if vector is None:
+            logger.warning(f"No vector found for object {row.get('id')}")
+            continue
+
+        vectors.append(vector)
+
+        # Extract object ID
+        obj_id = row.get("id") or row.get("_id") or str(len(vectors))
+        object_ids.append(obj_id)
+
+        # Extract metadata
+        if include_metadata:
+            meta = {}
+            if metadata_fields:
+                # Only include specified fields
+                for field in metadata_fields:
+                    if field in row:
+                        meta[field] = row[field]
+            else:
+                # Include all fields except the vector
+                meta = {k: v for k, v in row.items() if k != "__index__"}
+            metadata.append(meta)
+
+    return EmbeddingData(
+        vectors=np.array(vectors),
+        metadata=metadata,
+        collection_names=[collection.alias] * len(vectors),
+        collection_indices=[0] * len(vectors),
+        object_ids=object_ids
+    )
+
+
+def extract_embeddings_from_multiple_collections(
+    database,
+    collection_names: List[str],
+    index_name: Optional[str] = None,
+    limit_per_collection: Optional[int] = None,
+    include_metadata: bool = True,
+    metadata_fields: Optional[List[str]] = None,
+    normalize: bool = False
+) -> EmbeddingData:
+    """
+    Extract embeddings from multiple collections.
+
+    Args:
+        database: LinkML database object
+        collection_names: List of collection names
+        index_name: Name of index to use (must be same across collections)
+        limit_per_collection: Max embeddings per collection
+        include_metadata: Whether to include source object metadata
+        metadata_fields: Specific metadata fields to include
+        normalize: Whether to normalize vectors to unit length
+
+    Returns:
+        Combined EmbeddingData object
+    """
+    all_vectors = []
+    all_metadata = []
+    all_collection_names = []
+    all_collection_indices = []
+    all_object_ids = []
+
+    for i, coll_name in enumerate(collection_names):
+        try:
+            collection = database.get_collection(coll_name)
+            data = extract_embeddings_from_collection(
+                collection,
+                index_name=index_name,
+                limit=limit_per_collection,
+                include_metadata=include_metadata,
+                metadata_fields=metadata_fields
+            )
+
+            all_vectors.append(data.vectors)
+            all_metadata.extend(data.metadata)
+            all_collection_names.extend([coll_name] * data.n_samples)
+            all_collection_indices.extend([i] * data.n_samples)
+            all_object_ids.extend(data.object_ids)
+
+        except Exception as e:
+            logger.error(f"Failed to extract embeddings from {coll_name}: {e}")
+            continue
+
+    if not all_vectors:
+        raise ValueError("No embeddings extracted from any collection")
+
+    combined_vectors = np.vstack(all_vectors)
+
+    # Normalize if requested
+    if normalize:
+        # Ensure vectors are float type for division
+        combined_vectors = combined_vectors.astype(np.float64)
+        norms = np.linalg.norm(combined_vectors, axis=1, keepdims=True)
+        # Avoid division by zero
+        norms = np.where(norms == 0, 1, norms)
+        combined_vectors = combined_vectors / norms
+
+    return EmbeddingData(
+        vectors=combined_vectors,
+        metadata=all_metadata,
+        collection_names=all_collection_names,
+        collection_indices=all_collection_indices,
+        object_ids=all_object_ids
+    )
+
+
+def sample_embeddings(
+    embedding_data: EmbeddingData,
+    n_samples: int = 1000,
+    method: str = "random",
+    random_state: Optional[int] = None
+) -> EmbeddingData:
+    """
+    Sample embeddings for visualization.
+
+    Args:
+        embedding_data: Original embedding data
+        n_samples: Number of samples to select
+        method: Sampling method ('random', 'uniform', 'density')
+        random_state: Random seed for reproducibility
+
+    Returns:
+        Sampled EmbeddingData object
+    """
+    if embedding_data.n_samples <= n_samples:
+        return embedding_data
+
+    if random_state is not None:
+        np.random.seed(random_state)
+
+    if method == "random":
+        indices = np.random.choice(
+            embedding_data.n_samples,
+            size=n_samples,
+            replace=False
+        )
+    elif method == "uniform":
+        # Sample uniformly across collections
+        indices = []
+        for coll_idx in set(embedding_data.collection_indices):
+            coll_mask = np.array(embedding_data.collection_indices) == coll_idx
+            coll_indices = np.where(coll_mask)[0]
+            n_from_coll = min(
+                len(coll_indices),
+                n_samples // len(set(embedding_data.collection_indices))
+            )
+            indices.extend(
+                np.random.choice(coll_indices, size=n_from_coll, replace=False)
+            )
+        indices = np.array(indices[:n_samples])
+    else:
+        raise ValueError(f"Unknown sampling method: {method}")
+
+    return EmbeddingData(
+        vectors=embedding_data.vectors[indices],
+        metadata=[embedding_data.metadata[i] for i in indices],
+        collection_names=[embedding_data.collection_names[i] for i in indices],
+        collection_indices=[embedding_data.collection_indices[i] for i in indices],
+        object_ids=[embedding_data.object_ids[i] for i in indices]
+    )
+
+
+def compute_embedding_statistics(embedding_data: EmbeddingData) -> Dict:
+    """
+    Compute statistics about embeddings.
+
+    Args:
+        embedding_data: Embedding data
+
+    Returns:
+        Dictionary of statistics
+    """
+    stats = {
+        "n_samples": embedding_data.n_samples,
+        "n_dimensions": embedding_data.n_dimensions,
+        "n_collections": len(set(embedding_data.collection_names)),
+        "collections": list(set(embedding_data.collection_names)),
+    }
+
+    # Per-collection counts
+    from collections import Counter
+    collection_counts = Counter(embedding_data.collection_names)
+    stats["samples_per_collection"] = dict(collection_counts)
+
+    # Vector statistics
+    if embedding_data.n_samples > 0:
+        stats["mean_norm"] = float(np.mean(np.linalg.norm(embedding_data.vectors, axis=1)))
+        stats["std_norm"] = float(np.std(np.linalg.norm(embedding_data.vectors, axis=1)))
+
+        # Compute average pairwise similarity (on sample if large)
+        sample_size = min(100, embedding_data.n_samples)
+        if sample_size > 1:
+            sample_indices = np.random.choice(
+                embedding_data.n_samples,
+                size=sample_size,
+                replace=False
+            )
+            sample_vectors = embedding_data.vectors[sample_indices]
+
+            # Normalize for cosine similarity
+            norms = np.linalg.norm(sample_vectors, axis=1, keepdims=True)
+            normalized = sample_vectors / (norms + 1e-10)
+            similarities = np.dot(normalized, normalized.T)
+
+            # Extract upper triangle (excluding diagonal)
+            upper_tri = similarities[np.triu_indices(sample_size, k=1)]
+            stats["mean_similarity"] = float(np.mean(upper_tri))
+            stats["std_similarity"] = float(np.std(upper_tri))
+
+    return stats

linkml_store/utils/enrichment_analyzer.py

@@ -0,0 +1,217 @@
+from collections import Counter
+from typing import Dict, List
+
+import numpy as np
+import pandas as pd
+from pydantic import BaseModel
+from scipy import stats
+
+from linkml_store.api import Collection
+
+
+class EnrichedCategory(BaseModel):
+    """
+    Information about a category enriched in a sample
+    """
+
+    category: str
+    fold_change: float
+    original_p_value: float
+    adjusted_p_value: float
+
+
+class EnrichmentAnalyzer:
+    def __init__(self, df: pd.DataFrame, sample_key: str, classification_key: str):
+        """
+        Initialize the analyzer with a DataFrame and key column names.
+        Precomputes category frequencies for the entire dataset.
+
+        Args:
+            df: DataFrame containing the data
+            sample_key: Column name for sample IDs
+            classification_key: Column name for category lists
+        """
+        self.df = df
+        self.sample_key = sample_key
+        self.classification_key = classification_key
+
+        # Precompute global category statistics
+        self.global_stats = self._compute_global_stats()
+
+        # Cache for sample-specific category counts
+        self.sample_cache: Dict[str, Counter] = {}
+
+    @classmethod
+    def from_collection(cls, collection: Collection, sample_key: str, classification_key: str) -> "EnrichmentAnalyzer":
+        """
+        Initialize the analyzer with a Collection and key column names.
+        Precomputes category frequencies for the entire dataset.
+
+        Args:
+            collection: Collection containing the data
+            sample_key: Column name for sample IDs
+            classification_key: Column name for category lists
+        """
+        column_atts = [sample_key, classification_key]
+        results = collection.find(select_cols=column_atts, limit=-1)
+        df = results.rows_dataframe
+        ea = cls(df, sample_key=sample_key, classification_key=classification_key)
+        return ea
+
+    def _compute_global_stats(self) -> Dict[str, int]:
+        """
+        Compute global category frequencies across all samples.
+        Returns a dictionary of category -> count
+        """
+        global_counter = Counter()
+
+        # Flatten all categories and count
+        for categories in self.df[self.classification_key]:
+            if isinstance(categories, list):
+                global_counter.update(categories)
+            else:
+                # Handle case where categories might be a string
+                global_counter.update([categories])
+
+        return global_counter
+
+    @property
+    def sample_ids(self) -> List[str]:
+        df = self.df
+        return df[self.sample_key].unique().tolist()
+
+    def _get_sample_stats(self, sample_id: str) -> Counter:
+        """
+        Get category frequencies for a specific sample.
+        Uses caching to avoid recomputation.
+        """
+        if sample_id in self.sample_cache:
+            return self.sample_cache[sample_id]
+
+        sample_data = self.df[self.df[self.sample_key] == sample_id]
+        if sample_data.empty:
+            raise KeyError(f"Sample ID '{sample_id}' not found")
+        sample_data = sample_data.dropna()
+        # if sample_data.empty:
+        #    raise ValueError(f"Sample ID '{sample_id}' has missing values after dropping NA")
+        counter = Counter()
+
+        for categories in sample_data[self.classification_key]:
+            if isinstance(categories, list):
+                counter.update(categories)
+            else:
+                counter.update([categories])
+
+        self.sample_cache[sample_id] = counter
+        return counter
+
+    def find_enriched_categories(
+        self,
+        sample_id: str,
+        min_occurrences: int = 5,
+        p_value_threshold: float = 0.05,
+        multiple_testing_correction: str = "bh",
+    ) -> List[EnrichedCategory]:
+        """
+        Find categories that are enriched in the given sample.
+
+        Args:
+            sample_id: ID of the sample to analyze
+            min_occurrences: Minimum number of occurrences required for a category
+            p_value_threshold: P-value threshold for significance
+
+        Returns:
+            List of tuples (category, fold_change, p_value) sorted by significance
+        """
+        sample_stats = self._get_sample_stats(sample_id)
+        total_sample_annotations = sum(sample_stats.values())
+        total_global_annotations = sum(self.global_stats.values())
+
+        results = []
+
+        for category, sample_count in sample_stats.items():
+            global_count = self.global_stats[category]
+
+            # Skip rare categories
+            if global_count < min_occurrences:
+                continue
+
+            # Calculate fold change
+            sample_freq = sample_count / total_sample_annotations
+            global_freq = global_count / total_global_annotations
+            fold_change = sample_freq / global_freq if global_freq > 0 else float("inf")
+
+            # Perform Fisher's exact test
+            contingency_table = np.array(
+                [
+                    [sample_count, global_count - sample_count],
+                    [
+                        total_sample_annotations - sample_count,
+                        total_global_annotations - total_sample_annotations - (global_count - sample_count),
+                    ],
+                ]
+            )
+
+            _, p_value = stats.fisher_exact(contingency_table)
+
+            if p_value < p_value_threshold:
+                results.append((category, fold_change, p_value))
+
+        if not results:
+            return results
+
+        # Sort by p-value
+        results.sort(key=lambda x: x[2])
+
+        # Apply multiple testing correction
+        categories, fold_changes, p_values = zip(*results)
+
+        if multiple_testing_correction.lower() == "bonf":
+            # Bonferroni correction
+            n_tests = len(self.global_stats)  # Total number of categories tested
+            adjusted_p_values = [min(1.0, p * n_tests) for p in p_values]
+
+        elif multiple_testing_correction.lower() == "bh":
+            # Benjamini-Hochberg correction
+            n = len(p_values)
+            sorted_indices = np.argsort(p_values)
+            sorted_p_values = np.array(p_values)[sorted_indices]
+
+            # Calculate BH adjusted p-values
+            adjusted_p_values = np.zeros(n)
+            for i, p in enumerate(sorted_p_values):
+                adjusted_p_values[i] = p * n / (i + 1)
+
+            # Ensure monotonicity
+            for i in range(n - 2, -1, -1):
+                adjusted_p_values[i] = min(adjusted_p_values[i], adjusted_p_values[i + 1])
+
+            # Restore original order
+            inverse_indices = np.argsort(sorted_indices)
+            adjusted_p_values = adjusted_p_values[inverse_indices]
+
+            # Ensure we don't exceed 1.0
+            adjusted_p_values = np.minimum(adjusted_p_values, 1.0)
+
+        else:
+            # No correction
+            adjusted_p_values = p_values
+
+        # Filter by adjusted p-value threshold and create final results
+        # Create EnrichedCategory objects
+        final_results = [
+            EnrichedCategory(category=cat, fold_change=fc, original_p_value=p, adjusted_p_value=adj_p)
+            for cat, fc, p, adj_p in zip(categories, fold_changes, p_values, adjusted_p_values)
+            if adj_p < p_value_threshold
+        ]
+
+        # Sort by adjusted p-value
+        final_results.sort(key=lambda x: x.adjusted_p_value)
+        return final_results
+
+
+# Example usage:
+# analyzer = EnrichmentAnalyzer(df, 'sample_id', 'categories')
+# enriched = analyzer.find_enriched_categories('sample1')
+# for category, fold_change, p_value in enriched:
+#     print(f"{category}: {fold_change:.2f}x enrichment (p={p_value:.2e})")

linkml_store/utils/file_utils.py

@@ -0,0 +1,37 @@
+import logging
+import shutil
+import tempfile
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+
+def safe_remove_directory(dir_path: Path, no_backup: bool = False) -> Optional[Path]:
+    # Ensure the directory exists
+    if not dir_path.exists():
+        raise FileNotFoundError(f"Directory does not exist: {dir_path}")
+    try:
+        if no_backup:
+            # Move to a temporary directory instead of permanent removal
+            with tempfile.TemporaryDirectory() as tmpdir:
+                tmp_path = Path(tmpdir) / dir_path.name
+                shutil.move(str(dir_path), str(tmp_path))
+                logger.info(f"Directory moved to temporary location: {tmp_path}")
+                # The directory will be automatically removed when exiting the context manager
+            return None
+        else:
+            # Create a backup directory name with timestamp
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            backup_dir = dir_path.with_name(f"{dir_path.name}_backup_{timestamp}")
+
+            # Move the directory to the backup location
+            shutil.move(str(dir_path), str(backup_dir))
+            logger.info(f"Directory backed up to: {backup_dir}")
+            return backup_dir
+
+    except Exception as e:
+        logger.error(f"An error occurred: {e}")
+        return None