gsMap3D 0.1.0a1__py3-none-any.whl

gsMap/config/spatial_ldsc_config.py

@@ -0,0 +1,334 @@
+"""
+Configuration for spatial LD score regression.
+"""
+
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Annotated, Literal
+
+import typer
+import yaml
+
+from gsMap.config.base import ConfigWithAutoPaths
+
+logger = logging.getLogger("gsMap.config")
+
+@dataclass
+class SpatialLDSCComputeConfig:
+    """Compute configuration for spatial LDSC step."""
+    __display_in_quick_mode_cli__ = True
+
+    use_gpu: Annotated[bool, typer.Option(
+        "--use-gpu/--no-gpu",
+        help="Use GPU for JAX-accelerated spatial LDSC implementation"
+    ), {"__display_in_quick_mode_cli__": True}] = 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
+
+    ldsc_read_workers: Annotated[int, typer.Option(
+        help="Number of read workers",
+        min=1
+    )] = 10
+
+    ldsc_compute_workers: Annotated[int, typer.Option(
+        help="Number of compute workers for LDSC regression",
+        min=1
+    )] = 10
+
+    spots_per_chunk_quick_mode: Annotated[int, typer.Option(
+        help="Number of spots per chunk in quick mode",
+        min=1
+    )] = 50
+
+@dataclass
+class GWASSumstatsConfig:
+    """Configuration for GWAS summary statistics."""
+    trait_name: Annotated[str | None, typer.Option(
+        help="Name of the trait for GWAS analysis"
+    )] = None
+
+    sumstats_file: Annotated[Path | None, typer.Option(
+        help="Path to GWAS summary statistics file",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        resolve_path=True
+    )] = None
+
+    sumstats_config_file: Annotated[Path | None, typer.Option(
+        help="Path to sumstats config file",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        resolve_path=True
+    )] = None
+
+    sumstats_config_dict: dict[str, Path] = field(default_factory=dict)
+
+    @property
+    def trait_name_list(self) -> list[str]:
+        """Return the list of trait names to process."""
+        return list(self.sumstats_config_dict.keys())
+
+    def __post_init__(self):
+        self._init_sumstats()
+
+    def _init_sumstats(self):
+        """
+        Process sumstats input options and populate sumstats_config_dict.
+
+        Args:
+            self: Configuration object with sumstats_file, sumstats_config_file, and trait_name fields.
+
+        Returns:
+            Dict mapping trait names to sumstats file paths.
+        """
+        sumstats_config_dict = {}
+
+        if self.sumstats_file is None and self.sumstats_config_file is None:
+            raise ValueError("One of sumstats_file and sumstats_config_file must be provided.")
+        if self.sumstats_file is not None and self.sumstats_config_file is not None:
+            raise ValueError(
+                "Only one of sumstats_file and sumstats_config_file must be provided."
+            )
+        if self.sumstats_file is not None and self.trait_name is None:
+            raise ValueError("trait_name must be provided if sumstats_file is provided.")
+        if self.sumstats_config_file is not None and self.trait_name is not None:
+            raise ValueError(
+                "trait_name must not be provided if sumstats_config_file is provided."
+            )
+        # load the sumstats self file
+        if self.sumstats_config_file is not None:
+            # get the directory of the config file to resolve relative paths
+            config_dir = Path(self.sumstats_config_file).parent
+            with open(self.sumstats_config_file) as f:
+                config_loaded = yaml.load(f, Loader=yaml.FullLoader)
+            for _trait_name, sumstats_file in config_loaded.items():
+                s_path = Path(sumstats_file)
+                if not s_path.is_absolute():
+                    s_path = config_dir / s_path
+                sumstats_config_dict[_trait_name] = s_path.resolve()
+        # load the sumstats file
+        elif self.sumstats_file is not None:
+            sumstats_config_dict[self.trait_name] = Path(self.sumstats_file)
+        else:
+            raise ValueError("One of sumstats_file and sumstats_config_file must be provided.")
+
+        for sumstats_file in sumstats_config_dict.values():
+            assert Path(sumstats_file).exists(), f"{sumstats_file} does not exist."
+
+        self.sumstats_config_dict = sumstats_config_dict
+
+
+@dataclass
+class SpatialLDSCCoreConfig(GWASSumstatsConfig):
+    """Core configuration for spatial LDSC."""
+    w_ld_dir: Annotated[Path | None, typer.Option(
+        help="Directory containing the weights files (w_ld)",
+        exists=True,
+        file_okay=False,
+        dir_okay=True,
+        resolve_path=True
+    )] = None
+
+    additional_baseline_h5ad_path_list: Annotated[list[Path], typer.Option(
+        help="List of additional baseline h5ad paths",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        resolve_path=True
+    )] = field(default_factory=list)
+
+
+    chisq_max: Annotated[int | None, typer.Option(
+        help="Maximum chi-square value"
+    )] = None
+
+    cell_indices_range: Annotated[tuple[int, int] | None, typer.Option(
+        help="0-based range [start, end) of cell indices to process"
+    )] = None
+
+    sample_filter: Annotated[str | None, typer.Option(
+        help="Filter processing to a specific sample"
+    )] = None
+
+    n_blocks: Annotated[int, typer.Option(
+        help="Number of jackknife blocks",
+        min=1
+    )] = 200
+
+    # spots_per_chunk_quick_mode is inherited from SpatialLDSCComputeConfig
+
+    snp_gene_weight_adata_path: Annotated[Path | None, typer.Option(
+        help="Path to the SNP-gene weight matrix (H5AD format)",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        resolve_path=True
+    )] = None
+
+    # use_gpu is inherited from SpatialLDSCComputeConfig
+
+    marker_score_feather_path: Annotated[Path | None, typer.Option(
+        help="Path to marker score feather file",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        resolve_path=True
+    )] = None
+
+    marker_score_h5ad_path: Annotated[Path | None, typer.Option(
+        help="Path to marker score h5ad file",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        resolve_path=True
+    )] = None
+
+    marker_score_format: Annotated[Literal["memmap", "feather", "h5ad"] | None, typer.Option(
+        help="Format of marker scores"
+    )] = None
+
+
+
+@dataclass
+class SpatialLDSCConfig(SpatialLDSCCoreConfig, SpatialLDSCComputeConfig, ConfigWithAutoPaths):
+    """Spatial LDSC Configuration"""
+
+    def __post_init__(self):
+        super().__post_init__()
+
+
+        # Import here to avoid circular imports
+        from gsMap.config.utils import configure_jax_platform, get_anndata_shape
+
+        # Configure JAX platform if use_gpu is enabled
+        configure_jax_platform(self.use_gpu)
+
+        # Auto-detect marker_score_format if not specified
+        if self.marker_score_format is None:
+            if self.marker_score_feather_path is not None:
+                self.marker_score_format = "feather"
+                logger.info("Auto-detected marker_score_format as 'feather' based on marker_score_feather_path")
+            elif self.marker_score_h5ad_path is not None:
+                self.marker_score_format = "h5ad"
+                logger.info("Auto-detected marker_score_format as 'h5ad' based on marker_score_h5ad_path")
+            else:
+                self.marker_score_format = "memmap"
+                logger.info("Using default marker_score_format 'memmap'")
+
+        # Validate cell_indices_range is 0-based
+        if self.cell_indices_range is not None:
+            # Validate exclusivity between sample_filter and cell_indices_range
+
+            if self.sample_filter is not None:
+                raise ValueError(
+                    "Only one of sample_filter or cell_indices_range can be provided, not both. "
+                    "Use sample_filter to filter by sample, or cell_indices_range to process specific cell indices."
+                )
+
+            start, end = self.cell_indices_range
+
+            # Check that indices are 0-based
+            if start < 0:
+                raise ValueError(f"cell_indices_range start must be >= 0, got {start}")
+            if start == 1:
+                logger.warning(
+                    "cell_indices_range appears to be 1-based (start=1). "
+                    "Please ensure indices are 0-based. Adjusting start to 0."
+                )
+                start = 0
+
+            # Check that start < end
+            if start >= end:
+                raise ValueError(f"cell_indices_range start ({start}) must be less than end ({end})")
+
+            # Validate against actual data shape based on marker score format
+            if self.marker_score_format == "memmap":
+                # For memmap format, check the concatenated latent adata
+                adata_path = Path(self.workdir) / self.project_name / "latent2gene" / "concatenated_latent_adata.h5ad"
+                shape = get_anndata_shape(str(adata_path))
+                if shape is not None:
+                    n_obs, _ = shape
+                    if end > n_obs:
+                        logger.warning(
+                            f"cell_indices_range end ({end}) exceeds number of observations ({n_obs}). "
+                            f"Setting end to {n_obs}."
+                        )
+                        end = n_obs
+            elif self.marker_score_format == "h5ad":
+                # For h5ad format, check the provided h5ad path
+                adata_path = Path(self.marker_score_h5ad_path)
+                assert adata_path.exists(), f"Marker score h5ad not found at {adata_path}."
+                shape = get_anndata_shape(str(adata_path))
+                if shape is not None:
+                    n_obs, _ = shape
+                    if end > n_obs:
+                        logger.warning(
+                            f"cell_indices_range end ({end}) exceeds number of observations ({n_obs}). "
+                            f"Setting end to {n_obs}."
+                        )
+                        end = n_obs
+            elif self.marker_score_format == "feather":
+                # For feather format, validate the path exists
+                feather_path = Path(self.marker_score_feather_path)
+                assert feather_path.exists(), f"Marker score feather file not found at {feather_path}."
+
+                # Use pyarrow to get the number of rows and validate end
+                try:
+                    import pyarrow.feather as feather
+                    # Read metadata without loading full data
+                    feather_table = feather.read_table(str(feather_path), memory_map=True, columns=[])
+                    n_obs = feather_table.num_rows
+                    if end > n_obs:
+                        logger.warning(
+                            f"cell_indices_range end ({end}) exceeds number of rows ({n_obs}) in feather file. "
+                            f"Setting end to {n_obs}."
+                        )
+                        end = n_obs
+                except ImportError:
+                    logger.warning(
+                        "pyarrow not available. Cannot validate cell_indices_range against feather file. "
+                        "Install pyarrow to enable validation."
+                    )
+                except Exception as e:
+                    logger.warning(f"Could not read feather file metadata: {e}")
+
+            # Update cell_indices_range with validated values
+            self.cell_indices_range = (start, end)
+            logger.info(f"Processing cell_indices_range: [{start}, {end})")
+
+        if self.snp_gene_weight_adata_path is None:
+            raise ValueError("snp_gene_weight_adata_path must be provided.")
+
+        # Handle w_ld_dir
+        if self.w_ld_dir is None:
+            w_ld_dir = Path(self.ldscore_save_dir) / "w_ld"
+            if w_ld_dir.exists():
+                self.w_ld_dir = w_ld_dir
+                logger.info(f"Using weights directory generated in the generate_ldscore step: {self.w_ld_dir}")
+            else:
+                raise ValueError(
+                    "No w_ld_dir provided and no weights directory found in generate_ldscore output. "
+                    "Either provide --w-ld-dir or run generate_ldscore first."
+                )
+        else:
+            logger.info(f"Using provided weights directory: {self.w_ld_dir}")
+
+        self.show_config(SpatialLDSCConfig)
+
+
+def check_spatial_ldsc_done(config: SpatialLDSCConfig, trait_name: str) -> bool:
+    """
+    Check if spatial_ldsc step is done for a specific trait.
+    """
+    result_file = config.get_ldsc_result_file(trait_name)
+    return result_file.exists()

gsMap/config/utils.py

@@ -0,0 +1,286 @@
+"""
+Utility functions for gsMap configuration and data validation.
+"""
+
+import logging
+from collections import OrderedDict
+from pathlib import Path
+
+import h5py
+import yaml
+
+logger = logging.getLogger("gsMap.config")
+
+def configure_jax_platform(use_accelerator: bool = True):
+    """Configure JAX platform based on availability of accelerators.
+
+    Args:
+        use_accelerator: If True, try to use GPU then TPU if available,
+                        otherwise fall back to CPU. If False, force CPU usage.
+
+    Raises:
+        ImportError: If JAX is not installed.
+    """
+    try:
+        import os
+        os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+        logging.getLogger("jax._src.xla_bridge").setLevel(logging.ERROR)
+
+        import jax
+        from jax import config as jax_config
+
+        if not use_accelerator:
+            jax_config.update('jax_platform_name', 'cpu')
+            logger.info("JAX configured to use CPU for computations (accelerators disabled)")
+            return
+
+        # Priority list for accelerators
+        platforms = ['gpu', 'tpu']
+        configured = False
+
+        for platform in platforms:
+            try:
+                devices = jax.devices(platform)
+                if len(devices) > 0:
+                    jax_config.update('jax_platform_name', platform)
+                    logger.info(f"JAX configured to use {platform.upper()} for computations.")
+                    logger.info(f" ({len(devices)} device(s) detected). Device {devices[0]} will be used.")
+                    configured = True
+                    break
+            except (RuntimeError, ValueError, Exception):
+                continue
+
+        if not configured:
+            jax_config.update('jax_platform_name', 'cpu')
+            logger.info("No GPU or TPU detected, JAX configured to use CPU for computations")
+
+    except ImportError:
+        raise ImportError(
+            "JAX is required but not installed. Please install JAX by running: "
+            "pip install jax jaxlib (for CPU) or see JAX documentation for GPU/TPU installation."
+        ) from None
+
+def get_anndata_shape(h5ad_path: str):
+    """Get the shape (n_obs, n_vars) of an AnnData file without loading it."""
+    with h5py.File(h5ad_path, 'r') as f:
+        # 1. Verify it's a valid AnnData file by checking metadata
+        if f.attrs.get('encoding-type') != 'anndata':
+            logger.error(f"File '{h5ad_path}' does not appear to be a valid AnnData file.")
+            return None
+
+        # 2. Determine n_obs and n_vars from the primary metadata sources
+        if 'obs' not in f or 'var' not in f:
+            logger.error("AnnData file is missing 'obs' or 'var' group.")
+            return None
+
+        # Get the name of the index column from attributes
+        obs_index_key = f['obs'].attrs.get('_index', None)
+        var_index_key = f['var'].attrs.get('_index', None)
+
+        if not obs_index_key or obs_index_key not in f['obs']:
+            logger.error("Could not determine index for 'obs'.")
+            return None
+        if not var_index_key or var_index_key not in f['var']:
+            logger.error("Could not determine index for 'var'.")
+            return None
+
+        # The shape is the length of these index arrays
+        obs_obj = f['obs'][obs_index_key]
+        if isinstance(obs_obj, h5py.Group):
+            obs_obj = obs_obj['categories']
+        n_obs = obs_obj.shape[0]
+
+        var_obj  = f['var'][var_index_key]
+        if isinstance(var_obj, h5py.Group):
+            var_obj = var_obj['categories']
+        n_vars = var_obj.shape[0]
+
+        return n_obs, n_vars
+
+def inspect_h5ad_structure(filename):
+    """
+    Inspect the structure of an h5ad file without loading data.
+
+    Returns dict with keys present in each slot.
+    """
+    structure = {}
+
+    with h5py.File(filename, 'r') as f:
+        # Check main slots
+        slots = ['obs', 'var', 'obsm', 'varm', 'obsp', 'varp', 'uns', 'layers', 'X', 'raw']
+
+        for slot in slots:
+            if slot in f:
+                if slot in ['obsm', 'varm', 'obsp', 'varp', 'layers', 'uns']:
+                    # These are groups containing multiple keys
+                    structure[slot] = list(f[slot].keys())
+                elif slot in ['obs', 'var']:
+                    # These are dataframes - get column names
+                    if 'column-order' in f[slot].attrs:
+                        structure[slot] = list(f[slot].attrs['column-order'])
+                    else:
+                        structure[slot] = list(f[slot].keys())
+                else:
+                    # X, raw - just note they exist
+                    structure[slot] = True
+
+    return structure
+
+def validate_h5ad_structure(sample_h5ad_dict, required_fields, optional_fields=None):
+    """
+    Validate h5ad files have required structure.
+
+    Args:
+        sample_h5ad_dict: OrderedDict of {sample_name: h5ad_path}
+        required_fields: Dict of {field_name: (slot, field_key, error_msg_template)}
+            e.g., {'spatial': ('obsm', 'spatial', 'Spatial key')}
+        optional_fields: Dict of {field_name: (slot, field_key)} for fields to warn about
+
+    Returns:
+        None, raises ValueError if required fields are missing
+    """
+    for sample_name, h5ad_path in sample_h5ad_dict.items():
+        if not h5ad_path.exists():
+            raise FileNotFoundError(f"H5AD file not found for sample '{sample_name}': {h5ad_path}")
+
+        # Inspect h5ad structure
+        structure = inspect_h5ad_structure(h5ad_path)
+
+        # Check required fields
+        for field_name, (slot, field_key, error_msg) in required_fields.items():
+            if field_key is None:  # Skip if field not specified
+                continue
+
+            # Special handling for data_layer
+            if field_name == 'data_layer' and field_key != 'X':
+                if 'layers' not in structure or field_key not in structure.get('layers', []):
+                    raise ValueError(
+                        f"Data layer '{field_key}' not found in layers for sample '{sample_name}'. "
+                        f"Available layers: {structure.get('layers', [])}"
+                    )
+            elif field_name == 'data_layer' and field_key == 'X':
+                if 'X' not in structure:
+                    raise ValueError(f"X matrix not found in h5ad file for sample '{sample_name}'")
+            else:
+                # Standard validation for obsm, obs, etc.
+                if slot not in structure or field_key not in structure.get(slot, []):
+                    available = structure.get(slot, [])
+                    raise ValueError(
+                        f"{error_msg} '{field_key}' not found in {slot} for sample '{sample_name}'. "
+                        f"Available keys in {slot}: {available}"
+                    )
+
+        # Check optional fields (warn only)
+        if optional_fields:
+            for field_name, (slot, field_key) in optional_fields.items():
+                if field_key is None:  # Skip if field not specified
+                    continue
+
+                if slot not in structure or field_key not in structure.get(slot, []):
+                    available = structure.get(slot, [])
+                    logger.warning(
+                        f"Optional field '{field_key}' not found in {slot} for sample '{sample_name}'. "
+                        f"Available keys in {slot}: {available}"
+                    )
+
+def process_h5ad_inputs(config, input_options):
+    """
+    Process h5ad input options and create sample_h5ad_dict.
+
+    Args:
+        config: Configuration object with h5ad input fields
+        input_options: Dict mapping option names to (field_name, processing_type)
+            e.g., {'h5ad_yaml': ('h5ad_yaml', 'yaml'),
+                   'h5ad': ('h5ad', 'list'),
+                   'h5ad_list_file': ('h5ad_list_file', 'file')}
+
+    Returns:
+        OrderedDict of {sample_name: h5ad_path}
+    """
+
+    if config.sample_h5ad_dict  is not None and len(config.sample_h5ad_dict) > 0:
+        logger.info("Using pre-defined sample_h5ad_dict from configuration")
+        return OrderedDict(config.sample_h5ad_dict)
+
+    sample_h5ad_dict = OrderedDict()
+
+    # Check which options are provided
+    options_provided = []
+    for option_name, (field_name, _) in input_options.items():
+        if hasattr(config, field_name) and getattr(config, field_name):
+            options_provided.append(option_name)
+
+    # Ensure at most one option is provided
+    if len(options_provided) > 1:
+        raise AssertionError(f"At most one input option can be provided. Got {len(options_provided)}: {', '.join(options_provided)}. " f"Please provide only one of: {', '.join(input_options.keys())}")
+
+    # Process the provided input option
+    for option_name, (field_name, processing_type) in input_options.items():
+        field_value = getattr(config, field_name, None)
+        if not field_value:
+            continue
+
+        if processing_type == 'yaml':
+            logger.info(f"Using {option_name}: {field_value}")
+            yaml_file_path = Path(field_value)
+            yaml_parent_dir = yaml_file_path.parent
+            with open(field_value) as f:
+                h5ad_data = yaml.safe_load(f)
+                for sample_name, h5ad_path in h5ad_data.items():
+                    h5ad_path = Path(h5ad_path)
+                    # Resolve relative paths relative to yaml file location
+                    if not h5ad_path.is_absolute():
+                        h5ad_path = yaml_parent_dir / h5ad_path
+                    sample_h5ad_dict[sample_name] = h5ad_path
+
+        elif processing_type == 'list':
+            logger.info(f"Using {option_name} with {len(field_value)} files")
+            for h5ad_path in field_value:
+                h5ad_path = Path(h5ad_path)
+                sample_name = h5ad_path.stem
+                if sample_name in sample_h5ad_dict:
+                    logger.warning(f"Duplicate sample name: {sample_name}, will be overwritten")
+                sample_h5ad_dict[sample_name] = h5ad_path
+
+        elif processing_type == 'file':
+            logger.info(f"Using {option_name}: {field_value}")
+            list_file_path = Path(field_value)
+            list_file_parent_dir = list_file_path.parent
+            with open(field_value) as f:
+                for line in f:
+                    line = line.strip()
+                    if line:  # Skip empty lines
+                        h5ad_path = Path(line)
+                        # Resolve relative paths relative to list file location
+                        if not h5ad_path.is_absolute():
+                            h5ad_path = list_file_parent_dir / h5ad_path
+                        sample_name = h5ad_path.stem
+                        if sample_name in sample_h5ad_dict:
+                            logger.warning(f"Duplicate sample name: {sample_name}, will be overwritten")
+                        sample_h5ad_dict[sample_name] = h5ad_path
+        break
+
+    return sample_h5ad_dict
+
+def verify_homolog_file_format(config):
+    if config.homolog_file is not None:
+        logger.info(
+            f"User provided homolog file to map gene names to human: {config.homolog_file}"
+        )
+        # check the format of the homolog file
+        with open(config.homolog_file) as f:
+            first_line = f.readline().strip()
+            _n_col = len(first_line.split())
+            if _n_col != 2:
+                raise ValueError(
+                    f"Invalid homolog file format. Expected 2 columns, first column should be other species gene name, second column should be human gene name. "
+                    f"Got {_n_col} columns in the first line."
+                )
+            else:
+                first_col_name, second_col_name = first_line.split()
+                config.species = first_col_name
+                logger.info(
+                    f"Homolog file provided and will map gene name from column1:{first_col_name} to column2:{second_col_name}"
+                )
+    else:
+        logger.info("No homolog file provided. Run in human mode.")

gsMap/find_latent/__init__.py

@@ -0,0 +1,3 @@
+from .find_latent_representation import run_find_latent_representation
+
+__all__ = ['run_find_latent_representation']