gsMap3D 0.1.0a1__py3-none-any.whl

gsMap/config/latent2gene_config.py

@@ -0,0 +1,461 @@
+"""
+Configuration for latent to gene mapping.
+"""
+
+import logging
+from collections import OrderedDict
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+from typing import Annotated
+
+import typer
+import yaml
+
+from gsMap.config.base import ConfigWithAutoPaths
+from gsMap.config.utils import configure_jax_platform, process_h5ad_inputs, validate_h5ad_structure
+
+logger = logging.getLogger("gsMap.config")
+
+class DatasetType(str, Enum):
+    SCRNA_SEQ = 'scRNA'
+    SPATIAL_2D = 'spatial2D'
+    SPATIAL_3D = 'spatial3D'
+
+class MarkerScoreCrossSliceStrategy(str, Enum):
+    GLOBAL_POOL = 'global_pool'
+    PER_SLICE_POOL = 'per_slice_pool'
+    HIERARCHICAL_POOL = 'hierarchical_pool'
+
+@dataclass
+class LatentToGeneComputeConfig:
+    """Compute configuration for latent-to-gene step."""
+    __display_in_quick_mode_cli__ = True
+
+    use_gpu: Annotated[bool, typer.Option(
+        "--use-gpu/--no-gpu",
+        help="Use GPU for JAX computations (requires sufficient GPU memory)"
+    )] = True
+
+    memmap_tmp_dir: Annotated[Path | None, typer.Option(
+        help="Temporary directory for memory-mapped files to improve I/O performance on slow filesystems. "
+             "If provided, memory maps will be copied to this directory for faster random access during computation.",
+        exists=True,
+        file_okay=False,
+        dir_okay=True,
+        resolve_path=True
+    )] = None
+
+    # Batch sizes
+    rank_batch_size: int = 500
+
+    mkscore_batch_size: Annotated[int, typer.Option(
+        help="Number of cells per batch for marker score calculation. Reduce this value (e.g., 50) if encountering GPU OOM errors.",
+        min=10,
+        max=1000
+    )] = 500
+
+    find_homogeneous_batch_size: int = 100
+    rank_write_interval: int = 10
+
+    # Worker configurations
+    rank_read_workers: Annotated[int, typer.Option(
+        help="Number of parallel reader threads for rank memory map",
+        min=1,
+        max=50
+    )] = 16
+
+    mkscore_compute_workers: Annotated[int, typer.Option(
+        help="Number of parallel compute threads for marker score calculation",
+        min=1,
+        max=16
+    )] = 4
+
+    mkscore_write_workers: Annotated[int, typer.Option(
+        help="Number of parallel writer threads for marker scores",
+        min=1,
+        max=50
+    )] = 4
+
+    compute_input_queue_size: Annotated[int, typer.Option(
+        help="Maximum size of compute input queue (multiplier of mkscore_compute_workers)",
+        min=1,
+        max=10
+    )] = 5
+
+    writer_queue_size: Annotated[int, typer.Option(
+        help="Maximum size of writer input queue",
+        min=10,
+        max=500
+    )] = 100
+
+@dataclass
+class LatentToGeneCoreConfig:
+
+    dataset_type: Annotated[DatasetType, typer.Option(
+        help="Type of dataset: scRNA (uses KNN on latent space), spatial2D (2D spatial), or spatial3D (multi-slice)",
+        case_sensitive=False
+    )] = 'spatial2D'
+
+    # --------input h5ad file paths which have the latent representations
+    h5ad_path: Annotated[list[Path] | None, typer.Option(
+        help="Space-separated list of h5ad file paths. Sample names are derived from file names without suffix.",
+        exists=True,
+        file_okay=True,
+    )] = None
+
+    h5ad_yaml: Annotated[Path | None, typer.Option(
+        help="YAML file with sample names and h5ad paths",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+    )] = None
+
+    h5ad_list_file: Annotated[Path | None, typer.Option(
+        help="Each row is a h5ad file path, sample name is the file name without suffix",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+    )] = None
+
+    sample_h5ad_dict: OrderedDict | None = None
+
+    # --------input h5ad obs, obsm, layers keys
+
+    annotation: Annotated[str | None, typer.Option(
+        help="Cell type annotation in adata.obs to use. This would constrain finding homogeneous spots within each cell type"
+    )] = None
+
+    data_layer: Annotated[str, typer.Option(
+        help="Gene expression raw counts data layer in h5ad layers, e.g., 'count', 'counts'. Other wise use 'X' for adata.X"
+    )] = "X"
+
+    latent_representation_niche: Annotated[str | None, typer.Option(
+        help="Key for spatial niche embedding in obsm"
+    )] = None
+
+    latent_representation_cell: Annotated[str, typer.Option(
+        help="Key for cell identity embedding in obsm"
+    )] = "emb_cell"
+
+    spatial_key: Annotated[str, typer.Option(
+        help="Spatial key in adata.obsm"
+    )] = "spatial"
+
+    # --------parameters for finding homogeneous spots
+
+    spatial_neighbors: Annotated[int, typer.Option(
+        help="k1: Number of spatial neighbors in it's own slice for spatial dataset",
+        min=10,
+        max=5000
+    )] = 301
+
+    homogeneous_neighbors: Annotated[int, typer.Option(
+        help="k3: Number of homogeneous neighbors per cell (for spatial) or KNN neighbors (for scRNA-seq)",
+        min=1,
+        max=200
+    )] = 21
+
+    cell_embedding_similarity_threshold: Annotated[float, typer.Option(
+        help="Minimum similarity threshold for cell embedding.",
+        min=0.0,
+        max=1.0
+    )] = 0.0
+
+    spatial_domain_similarity_threshold: Annotated[float, typer.Option(
+        help="Minimum similarity threshold for spatial domain embedding.",
+        min=0.0,
+        max=1.0
+    )] = 0.6
+
+    no_expression_fraction: Annotated[bool, typer.Option(
+        "--no-expression-fraction",
+        help="Skip expression fraction filtering"
+    )] = False
+
+    # --------3D slice-aware neighbor search parameters
+    adjacent_slice_spatial_neighbors: Annotated[int, typer.Option(
+        help="Number of spatial neighbors to find on each adjacent slice for 3D data",
+        min=10,
+        max=2000
+    )] = 200
+
+    n_adjacent_slices: Annotated[int, typer.Option(
+        help="Number of adjacent slices to search above and below (± n_adjacent_slices) in 3D space for each focal spot. Padding will be applied automatically.",
+        min=0,
+        max=5
+    )] = 1
+
+    cross_slice_marker_score_strategy: Annotated[MarkerScoreCrossSliceStrategy, typer.Option(
+        help="Strategy for computing marker scores across slices in spatial3D datasets. "
+             "'global_pool': Select the top K most similar neighbors globally across all slices combined. "
+             "'per_slice_pool': Select a fixed number of neighbors (K) from each slice independently, then compute a single weighted average score from all selected neighbors. "
+             "'hierarchical_pool': Compute an independent marker score for each slice using its top K neighbors, then take the average of these per-slice scores.",
+        case_sensitive=False
+    )] = MarkerScoreCrossSliceStrategy.HIERARCHICAL_POOL
+
+    high_quality_neighbor_filter: Annotated[bool, typer.Option(
+        "--high-quality-neighbor-filter/--no-high-quality-filter",
+        help="Only find neighbors within high quality cells (requires High_quality column in obs)"
+    )] = False
+
+    fix_cross_slice_homogenous_neighbors: bool = False
+
+    # Performance options
+    # Minimum number of cells per cell type in that dataset to be used for finding homogeneous neighbors
+    min_cells_per_type: int | None = None
+
+    enable_profiling: bool = False
+
+
+@dataclass
+class LatentToGeneConfig(LatentToGeneComputeConfig, LatentToGeneCoreConfig, ConfigWithAutoPaths):
+    """Latent to Gene Configuration"""
+
+    @property
+    def total_homogeneous_neighbor_per_cell(self):
+        return self.homogeneous_neighbors * (1 + 2 * self.n_adjacent_slices)
+
+    def __post_init__(self):
+        """Initialize and validate configuration"""
+        super().__post_init__()
+
+        # Step 1: Configure JAX platform
+        configure_jax_platform(self.use_gpu)
+
+        # Step 2: Process and validate h5ad inputs
+        self._resolve_h5ad_inputs()
+
+        # Step 3: Configure dataset-specific parameters first
+        self._configure_dataset_parameters()
+
+        # Step 4: Set up validation fields and validate structure (after dataset config)
+        self._setup_and_validate_fields()
+
+        self.show_config(LatentToGeneConfig)
+
+    def _resolve_h5ad_inputs(self):
+        """Resolve h5ad inputs from various sources, prioritizing auto-detection."""
+
+        # Step 1: Try auto-detection first
+        self._auto_detect_h5ad_files()
+
+        # Step 2: If auto-detection didn't find anything, try explicit inputs
+        if not self.sample_h5ad_dict:
+            # Define input options
+            input_options = {
+                'h5ad_yaml': ('h5ad_yaml', 'yaml'),
+                'h5ad_path': ('h5ad_path', 'list'),
+                'h5ad_list_file': ('h5ad_list_file', 'file'),
+            }
+            self.sample_h5ad_dict = process_h5ad_inputs(self, input_options)
+
+        # Step 3: Validate at least one sample exists
+        if not self.sample_h5ad_dict or len(self.sample_h5ad_dict) == 0:
+            raise ValueError(
+                "No valid h5ad files found. Please provide one of: h5ad_yaml, h5ad_path, or h5ad_list_file, "
+                "or ensure find_latent representation has been run to allow auto-detection."
+            )
+
+        logger.info(f"Loaded and validated {len(self.sample_h5ad_dict)} samples")
+
+    def _auto_detect_h5ad_files(self):
+        """Auto-detect h5ad files from latent directory"""
+        if self.find_latent_metadata_path.exists():
+            import yaml
+            with open(self.find_latent_metadata_path) as f:
+                find_latent_metadata = yaml.safe_load(f)
+            self.sample_h5ad_dict = OrderedDict(
+                {sample_name: Path(latent_file)
+                 for sample_name, latent_file in
+                 find_latent_metadata['outputs']['latent_files'].items()
+                 })
+            # assert all files exist
+            for sample_name, latent_file in self.sample_h5ad_dict.items():
+                if not latent_file.exists():
+                    raise FileNotFoundError(f"Latent file not found for sample '{sample_name}': {latent_file}")
+            logger.info(
+                f"Auto-detected {len(self.sample_h5ad_dict)} samples from find_latent_metadata_path: {self.find_latent_metadata_path}")
+        else:
+            self.sample_h5ad_dict = OrderedDict()
+            latent_dir = self.latent_dir
+            logger.info(f"Auto-detecting h5ad files from latent directory: {latent_dir}")
+
+            # Look for latent files with different naming patterns
+            latent_files = list(latent_dir.glob("*_latent_adata.h5ad"))
+            if not latent_files:
+                latent_files = list(latent_dir.glob("*_add_latent.h5ad"))
+
+            if not latent_files:
+                return
+
+            # Extract sample names from file names
+            for latent_file in latent_files:
+                sample_name = self._extract_sample_name(latent_file)
+                self.sample_h5ad_dict[sample_name] = latent_file
+
+            # sort by sample name
+            self.sample_h5ad_dict = OrderedDict(sorted(self.sample_h5ad_dict.items()))
+
+        logger.info(f"Auto-detected {len(self.sample_h5ad_dict)} samples from latent directory")
+
+    def _extract_sample_name(self, latent_file):
+        """Extract sample name from latent file path"""
+        filename = latent_file.stem
+
+        # Remove known suffixes
+        suffixes_to_remove = ["_latent_adata", "_add_latent"]
+        for suffix in suffixes_to_remove:
+            if filename.endswith(suffix):
+                return filename[:-len(suffix)]
+
+        return filename
+
+    def _setup_and_validate_fields(self):
+        """Set up required/optional fields and validate h5ad structure"""
+        # Define required fields
+        required_fields = {
+            'latent_representation_cell': ('obsm', self.latent_representation_cell,
+                                           'Latent representation of cell identity'),
+            'spatial_key': ('obsm', self.spatial_key, 'Spatial key'),
+        }
+
+        # Add annotation as required if provided
+        if self.annotation:
+            required_fields['annotation'] = ('obs', self.annotation, 'Annotation')
+
+        # Add niche representation as required if provided
+        if self.latent_representation_niche:
+            required_fields['latent_representation_niche'] = (
+                'obsm',
+                self.latent_representation_niche,
+                'Latent representation of spatial niche'
+            )
+
+        # Add High_quality as required if find_neighbor_within_high_quality is enabled
+        if self.high_quality_neighbor_filter:
+            required_fields['High_quality'] = ('obs', 'High_quality', 'High quality cell indicator')
+
+        # Validate h5ad structure
+        validate_h5ad_structure(self.sample_h5ad_dict, required_fields)
+
+    def _configure_dataset_parameters(self):
+        """Configure parameters based on dataset type"""
+        self.min_cells_per_type = self.homogeneous_neighbors if self.min_cells_per_type is None else min(self.min_cells_per_type, self.homogeneous_neighbors)
+
+        if self.dataset_type == DatasetType.SPATIAL_2D:
+            self._configure_spatial_2d()
+        elif self.dataset_type == DatasetType.SPATIAL_3D:
+            self._configure_spatial_3d()
+        elif self.dataset_type == DatasetType.SCRNA_SEQ:
+            self._configure_scrna_seq()
+
+    def _configure_spatial_2d(self):
+        """Configure parameters for spatial 2D datasets"""
+        # spatial2D can have multiple slices but doesn't search across them
+        if self.n_adjacent_slices != 0:
+            self.n_adjacent_slices = 0
+            self.adjacent_slice_spatial_neighbors = 0
+            logger.info(
+                "Dataset type is spatial2D. This will only search homogeneous neighbors within each 2D slice (no cross-slice search). Setting adjacent_slices=0.")
+
+        if self.latent_representation_niche is None:
+            logger.warning("latent_representation_niche is not provided. Spatial domain similarity will not be used.")
+
+        assert self.homogeneous_neighbors <= self.spatial_neighbors, \
+            f"homogeneous_neighbors ({self.homogeneous_neighbors}) must be <= spatial_neighbors ({self.spatial_neighbors}) for spatial2D datasets"
+
+    def _configure_spatial_3d(self):
+        """Configure parameters for spatial 3D datasets"""
+        if self.n_adjacent_slices == 0:
+            raise ValueError(
+                "Dataset type is spatial3D, but adjacent_slices=0. "
+                "You must set adjacent_slices to 1 or higher to enable cross-slice search. "
+                "If you don't want cross-slice search, use dataset_type='spatial2D' instead."
+            )
+
+        if self.latent_representation_niche is None:
+            logger.warning("latent_representation_niche is not provided. Spatial domain similarity will not be used.")
+
+        assert self.adjacent_slice_spatial_neighbors <= self.spatial_neighbors, \
+            f"adjacent_slice_neighbors ({self.adjacent_slice_spatial_neighbors}) must be <= spatial_neighbors ({self.spatial_neighbors})"
+        assert self.homogeneous_neighbors <= self.adjacent_slice_spatial_neighbors, \
+            f"homogeneous_neighbors ({self.homogeneous_neighbors}) must be <= adjacent_slice_neighbors ({self.adjacent_slice_spatial_neighbors})"
+
+        n_slices = 1 + self.n_adjacent_slices  # only focal + above slices
+        assert n_slices <= len(self.sample_h5ad_dict), \
+            f"3D Cross slice search requires at least {n_slices} slices (1 focal + {self.n_adjacent_slices} above or {self.n_adjacent_slices} below). " \
+            f"Only {len(self.sample_h5ad_dict)} samples provided. Please provide more slices or reduce adjacent_slices."
+
+        logger.info(
+            f"Dataset type is spatial3D, using adjacent_slices={self.n_adjacent_slices} for cross-slice search")
+        logger.info("The Z axis order of slices is determined by the h5ad input order. Currently, the order is: ")
+        logger.info(f"{' -> '.join(list(self.sample_h5ad_dict.keys()))}")
+
+        homogeneous_neighbors = self.homogeneous_neighbors
+        n_adjacent_slices = self.n_adjacent_slices
+        # Check if we should use fix number of homogeneous neighbors per slice
+        if self.cross_slice_marker_score_strategy in [
+            MarkerScoreCrossSliceStrategy.PER_SLICE_POOL,
+            MarkerScoreCrossSliceStrategy.HIERARCHICAL_POOL
+        ]:
+
+            self.fix_cross_slice_homogenous_neighbors = True
+            logger.info(
+                f"Using {self.cross_slice_marker_score_strategy.value} strategy with fixed number of homogeneous neighbors per adjacent slice: {self.homogeneous_neighbors} per slice.")
+
+
+        elif self.cross_slice_marker_score_strategy == MarkerScoreCrossSliceStrategy.GLOBAL_POOL:
+            logger.info(
+                "Using global_pool strategy, will select top homogeneous neighbors from all adjacent slices based on similarity scores. Each adjacent slice can contribute variable number of homogeneous neighbors.")
+
+        logger.info(
+            f"Each focal cell will select {homogeneous_neighbors * (1 + 2 * n_adjacent_slices) = } total homogeneous neighbors across {(1 + 2 * n_adjacent_slices) = } slices.")
+
+    def _configure_scrna_seq(self):
+        """Configure parameters for scRNA-seq datasets"""
+        self.n_adjacent_slices = 0
+        self.spatial_key = None
+        self.latent_representation_niche = None
+
+def check_latent2gene_done(config: LatentToGeneConfig) -> bool:
+    """
+    Check if latent2gene step is done by verifying validity of metadata and output files.
+    """
+    from gsMap.latent2gene.memmap_io import MemMapDense
+
+    config._resolve_h5ad_inputs()
+
+    expected_outputs = {
+        "concatenated_latent_adata": Path(config.concatenated_latent_adata_path),
+        "rank_memmap": Path(config.rank_memmap_path),
+        "mean_frac": Path(config.mean_frac_path),
+        "marker_scores": Path(config.marker_scores_memmap_path),
+        "metadata": Path(config.latent2gene_metadata_path),
+    }
+
+    if not all(p.exists() for p in expected_outputs.values()):
+        return False
+
+    try:
+        # Check rank memmap completion
+        rank_memmap_complete, _ = MemMapDense.check_complete(expected_outputs["rank_memmap"])
+        if not rank_memmap_complete:
+            return False
+
+        # Check marker scores memmap completion
+        marker_scores_complete, _ = MemMapDense.check_complete(expected_outputs["marker_scores"])
+        if not marker_scores_complete:
+            return False
+
+        # Check metadata
+        with open(expected_outputs["metadata"]) as f:
+            metadata = yaml.unsafe_load(f)
+
+        if 'outputs' not in metadata:
+            return False
+
+        return True
+    except Exception as e:
+        logger.warning(f"Error checking latent2gene results: {e}")
+        return False