gsMap3D 0.1.0a1__py3-none-any.whl

gsMap/ldscore/io.py

@@ -0,0 +1,262 @@
+"""
+I/O utilities for loading genomic data in the LD score framework.
+
+This module provides optimized readers for PLINK binary files and omics feature matrices,
+leveraging pandas-plink for efficient I/O and standardizing data with NumPy.
+"""
+
+import logging
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+from pandas_plink import read_plink1_bin
+
+logger = logging.getLogger(__name__)
+
+
+class PlinkBEDReader:
+    """
+    Optimized reader for PLINK binary files using pandas-plink and Xarray.
+
+    Loads genotypes lazily via Dask/Xarray, applies matrix-based MAF calculation
+    and filtering, and converts to standardized NumPy arrays.
+
+    Attributes
+    ----------
+    bfile : str
+        Base filename prefix for PLINK files
+    G : xr.DataArray
+        The underlying xarray DataArray (samples x variants) containing genotypes (0, 1, 2, nan)
+    bim : pd.DataFrame
+        BIM file data (SNP information, filtered)
+    fam : pd.DataFrame
+        FAM file data (individual information)
+    m : int
+        Number of SNPs (after filtering)
+    n : int
+        Number of individuals
+    genotypes : np.ndarray
+        Pre-loaded and standardized genotype matrix (n_individuals, m_snps)
+    maf : np.ndarray
+        Minor allele frequency for each SNP
+    """
+
+    def __init__(
+        self,
+        bfile_prefix: str,
+        maf_min: float | None = None,
+        keep_snps: list[str] | None = None,
+        preload: bool = True
+    ):
+        """
+        Initialize PlinkBEDReader with optional filtering.
+
+        Parameters
+        ----------
+        bfile_prefix : str
+            PLINK file prefix (without .bed/.bim/.fam extension)
+        maf_min : float, optional
+            Minimum MAF threshold for SNP filtering (default: None, no filtering)
+        keep_snps : list[str], optional
+            List of SNP IDs to keep (default: None, keep all)
+        preload : bool, optional
+            Whether to pre-load and standardize genotypes into memory (default: True)
+        """
+        self.bfile = bfile_prefix
+
+        # Construct paths
+        bed_path = f"{bfile_prefix}.bed"
+        bim_path = f"{bfile_prefix}.bim"
+        fam_path = f"{bfile_prefix}.fam"
+
+        # Validate existence
+        if not (Path(bed_path).exists() and Path(bim_path).exists() and Path(fam_path).exists()):
+             raise FileNotFoundError(f"One or more PLINK files missing for prefix: {bfile_prefix}")
+
+        logger.info(f"Loading PLINK files from: {bfile_prefix}")
+
+        # Load using pandas-plink
+        # This returns an xarray DataArray with dask backing (lazy loading)
+        # Shape is (sample, variant)
+        self.G = read_plink1_bin(bed_path, bim_path, fam_path, verbose=False)
+
+        # Initial dimensions
+        self.n_original = self.G.sizes['sample']
+        self.m_original = self.G.sizes['variant']
+        self.snp_ids_original = pd.Index(self.G.snp.values)
+
+        logger.info(f"Loaded metadata: {self.m_original} SNPs × {self.n_original} individuals")
+
+        # Calculate MAF using matrix operations (lazy execution via dask)
+        logger.info("Calculating MAF ...")
+        self.maf = self._calculate_maf()
+
+        # Apply filters to the xarray DataArray
+        self._apply_filters(maf_min=maf_min, keep_snps=keep_snps)
+
+        # Apply Basic QC (Filter Monomorphic and All-Missing)
+        # This must happen before _sync_metadata so BIM reflects valid SNPs only
+        self._apply_basic_qc()
+
+        # Update dimensions after filtering
+        self.n = self.G.sizes['sample']
+        self.m = self.G.sizes['variant']
+
+        # Extract metadata DataFrames from xarray coordinates for compatibility
+        self._sync_metadata()
+
+        # Pre-load genotypes if requested
+        self.genotypes = None
+        if preload:
+            logger.info(f"Pre-loading and standardizing {self.m} SNPs...")
+            self.genotypes = self._load_and_standardize_all()
+            logger.info(f"✓ Genotypes ready: {self.genotypes.shape}")
+
+        logger.info(f"PlinkBEDReader initialized: {self.m} SNPs × {self.n} individuals")
+
+    def _calculate_maf(self) -> xr.DataArray:
+        """
+        Calculate Minor Allele Frequency using matrix operations on the DataArray.
+        """
+        # Calculate mean across samples (axis 0), ignoring NaNs
+        # The result 'freq_a1' represents the frequency of the A1 allele (coded as 2)
+        freq_a1 = self.G.mean(dim="sample", skipna=True) / 2.0
+
+        # Calculate MAF: min(f, 1-f)
+        maf = xr.ufuncs.minimum(freq_a1, 1.0 - freq_a1)
+
+        return maf.compute()
+
+    def _apply_filters(
+        self,
+        maf_min: float | None = None,
+        keep_snps: list[str] | None = None
+    ) -> None:
+        """
+        Apply SNP filters directly to the xarray DataArray.
+        """
+        # 1. Create a boolean mask for variants
+        variant_ids = self.G.variant.values
+        mask = np.ones(len(variant_ids), dtype=bool)
+
+        # 2. Apply MAF filter
+        if maf_min is not None and maf_min > 0:
+            maf_mask = (self.maf >= maf_min).values
+            mask &= maf_mask
+
+            n_removed_maf = np.sum(~maf_mask)
+            if n_removed_maf > 0:
+                logger.info(f"Filtered {n_removed_maf} SNPs with MAF < {maf_min}")
+
+        # 3. Apply Keep List filter
+        if keep_snps is not None:
+            current_snps = self.G.snp.values
+            keep_set = set(keep_snps)
+            snp_mask = np.isin(current_snps, list(keep_set))
+
+            mask &= snp_mask
+            n_removed_snp = np.sum(~snp_mask)
+            if n_removed_snp > 0:
+                logger.info(f"Filtered {n_removed_snp} SNPs not in keep list")
+
+        # 4. Filter the main DataArray
+        n_before = self.G.sizes['variant']
+        self.G = self.G.isel(variant=mask)
+        self.maf = self.maf[mask]
+
+        n_after = self.G.sizes['variant']
+        if n_before != n_after:
+            logger.info(f"Total SNPs filtered: {n_before - n_after}/{n_before}")
+
+    def _apply_basic_qc(self) -> None:
+        """
+        Filter out monomorphic variants (std=0) and variants with all missing values.
+        """
+        logger.info("Applying basic QC (removing monomorphic and all-missing variants)...")
+
+        # Calculate stats lazily via dask
+        # 1. Standard Deviation (skipna=True handles missing)
+        stds = self.G.std(dim="sample", skipna=True)
+
+        # 2. Count of non-missing values
+        counts = self.G.count(dim="sample")
+
+        # Compute to get numpy arrays for boolean masking
+        stds_val = stds.values
+        counts_val = counts.values
+
+        # Create masks
+        # Monomorphic: std == 0 (or very close to 0)
+        # All missing: count == 0
+        mask_polymorphic = (stds_val > 0)
+        mask_not_empty = (counts_val > 0)
+
+        mask = mask_polymorphic & mask_not_empty
+
+        n_removed = np.sum(~mask)
+        if n_removed > 0:
+            logger.info(f"QC: Filtered {n_removed} variants (monomorphic or all-missing)")
+
+            # Apply filter
+            self.G = self.G.isel(variant=mask)
+            self.maf = self.maf[mask]
+        else:
+            logger.info("QC: No monomorphic or all-missing variants found.")
+
+    def _sync_metadata(self):
+
+        self.bim = pd.DataFrame({
+            'CHR': self.G.chrom.values,
+            'SNP': self.G.snp.values,
+            'CM': self.G.cm.values,
+            'BP': self.G.pos.values,
+            'A1': self.G.a1.values,
+            'A2': self.G.a0.values,
+            'i': np.arange(self.m)
+        })
+        self.bim['MAF'] = self.maf.values
+
+        # pandas-plink stores FAM info in coordinates
+        self.fam = pd.DataFrame({
+            'fid': self.G.fid.values,
+            'iid': self.G.iid.values,
+            'father': self.G.father.values,
+            'mother': self.G.mother.values,
+            'gender': self.G.gender.values,
+            'trait': self.G.trait.values
+        })
+
+    def _load_and_standardize_all(self) -> np.ndarray:
+        """
+        Load genotypes from Dask into memory, convert to NumPy, and standardize.
+
+        We assume basic QC (monomorphic/all-missing removal) has already run.
+
+        Returns
+        -------
+        np.ndarray
+            Standardized genotype matrix of shape (n_individuals, m_snps)
+        """
+        logger.info("Reading filtered genotype matrix into memory...")
+        X = self.G.values.astype(np.float32)
+
+        # Compute stats (ignoring NaNs)
+        means = np.nanmean(X, axis=0)
+        stds = np.nanstd(X, axis=0)
+
+        # Impute missing values with column means
+        nan_mask = np.isnan(X)
+
+        # Broadcasting means: (m,) -> (1, m) to match (n, m) for X which is (n, m)
+        # Note: In X, rows=individuals, cols=snps. means shape is (m_snps,).
+        # We need to broadcast properly.
+        # X is (n, m), means is (m,). numpy broadcasts (m,) to (n, m) automatically on last dim.
+        X = np.where(nan_mask, means, X)
+
+        # Standardize: (X - mean) / std
+        # Standard broadcasting rules apply: (n, m) - (m,) -> (n, m)
+        X_std = (X - means) / stds
+
+        return X_std

gsMap/ldscore/mapping.py

@@ -0,0 +1,262 @@
+import logging
+
+import numpy as np
+import pandas as pd
+import pyranges as pr
+import scipy.sparse
+
+logger = logging.getLogger(__name__)
+
+
+def create_snp_feature_map(
+        bim_df: pd.DataFrame,
+        mapping_type: str,
+        mapping_data: pd.DataFrame | dict[str, str],
+        feature_window_size: int = 0,
+        strategy: str = "score",
+) -> tuple[scipy.sparse.csr_matrix, list[str], pd.DataFrame | None]:
+    """
+    Create a sparse mapping matrix assigning each SNP in the BIM file to feature indices.
+
+    Returns:
+    1. Sparse matrix where rows correspond to SNPs (in bim_df order) and columns correspond to features.
+    2. List of feature names corresponding to indices 0 to F-1.
+    3. (Optional) DataFrame containing the curated SNP-feature mappings (for BED type).
+
+    Feature Indexing Scheme
+    -----------------------
+    - Mapped features get indices: 0, 1, 2, ..., F-1
+    - Unmapped bin gets index: F
+    - Total columns: F + 1
+    - Values: Score from BED file (if available, else 1.0)
+
+    Parameters
+    ----------
+    bim_df : pd.DataFrame
+        PLINK BIM file data (must have 'CHR', 'SNP', 'BP').
+    mapping_type : str
+        'bed' or 'dict'.
+    mapping_data : Union[pd.DataFrame, Dict[str, str]]
+        Mapping source data. For 'bed' type, this should be a DataFrame read from
+        a standard BED file using pr.read_bed() with columns:
+        - Chromosome (chr)
+        - Start (0-based start position)
+        - End (1-based end position)
+        - Feature (name/identifier from 4th column)
+        - Score (optional, from 5th column)
+        - Strand (optional, from 6th column)
+
+        Note: BED files should be in standard BED6 format WITHOUT a header line.
+    feature_window_size : int
+        Window extension (for 'bed').
+    strategy : str
+        'score', 'tss', 'center', or 'allow_repeat'.
+        - 'allow_repeat': A SNP can map to multiple features (values are summed or kept).
+        - 'score': Keep mapping with highest score per SNP.
+        - 'tss': Keep mapping closest to TSS per SNP.
+        - 'center': Keep mapping closest to the center of the feature interval per SNP.
+
+    Returns
+    -------
+    mapping_matrix : scipy.sparse.csr_matrix
+        Sparse matrix of shape (M, F+1).
+    feature_names : List[str]
+        List of feature names corresponding to indices 0 to F-1.
+    mapping_df : Optional[pd.DataFrame]
+        DataFrame with columns [SNP, Feature, ...] showing final mappings. None if using 'dict'.
+    """
+    m_ref = len(bim_df)
+    unique_feature_names = []
+    curated_mapping_df = None
+
+    # Prepare basic SNP info for joining
+    # We assign an integer index to every SNP in BIM to construct the sparse matrix later
+    bim_df = bim_df.copy()
+    bim_df['snp_row_idx'] = np.arange(m_ref)
+
+    # Intermediate storage for sparse construction
+    row_indices = []
+    col_indices = []
+    data_values = []
+
+    # === STRATEGY A: Dictionary Mapping ===
+    if mapping_type == 'dict':
+        if not isinstance(mapping_data, dict):
+            raise ValueError("mapping_data must be a dictionary when mapping_type='dict'")
+
+        # 1. Identify all unique features
+        unique_feature_names = sorted(set(mapping_data.values()))
+        feature_to_idx = {f: i for i, f in enumerate(unique_feature_names)}
+        n_features = len(unique_feature_names)
+
+        # 2. Map Dictionary Values to Indices
+        # Filter mapping data to only include SNPs present in BIM to save memory/time
+        bim_snps = set(bim_df['SNP'])
+        valid_mapping = {k: v for k, v in mapping_data.items() if k in bim_snps}
+
+        # 3. Create Sparse Entries
+        # Create a Series for mapping: SNP -> Feature Index
+        snp_to_feat_idx = pd.Series(valid_mapping).map(feature_to_idx)
+
+        # Map BIM SNPs to feature indices
+        mapped_feat_indices = bim_df['SNP'].map(snp_to_feat_idx)
+
+        # Drop NaNs (Unmapped)
+        valid_mask = mapped_feat_indices.notna()
+
+        if valid_mask.any():
+            row_indices = bim_df.loc[valid_mask, 'snp_row_idx'].values
+            col_indices = mapped_feat_indices[valid_mask].values.astype(int)
+            data_values = np.ones(len(row_indices), dtype=np.float32)
+
+    # === STRATEGY B: BED/Genomic Interval Mapping ===
+    elif mapping_type == 'bed':
+        if not isinstance(mapping_data, pd.DataFrame):
+            raise ValueError("mapping_data must be a DataFrame when mapping_type='bed'")
+
+        df_features = mapping_data.copy()
+        required_cols = ['Feature', 'Chromosome', 'Start', 'End']
+        if not all(c in df_features.columns for c in required_cols):
+            raise ValueError(f"BED DataFrame missing required columns: {required_cols}")
+
+        # 1. Identify all unique features
+        unique_feature_names = sorted(df_features['Feature'].unique())
+        feature_to_idx = {f: i for i, f in enumerate(unique_feature_names)}
+        n_features = len(unique_feature_names)
+
+        # 2. Prepare BIM for PyRanges
+        bim_pr_df = bim_df[['CHR', 'BP', 'SNP', 'snp_row_idx']].rename(columns={
+            'CHR': 'Chromosome',
+            'BP': 'Start'
+        })
+        bim_pr_df['End'] = bim_pr_df['Start'] + 1
+
+        # Clean chromosome names (remove 'chr' prefix if present)
+        # PLINK BIM files typically use numeric chromosome identifiers
+        bim_pr_df['Chromosome'] = bim_pr_df['Chromosome'].astype(str).str.replace('chr', '', case=False)
+        df_features['Chromosome'] = df_features['Chromosome'].astype(str).str.replace('chr', '', case=False)
+
+        pr_bim = pr.PyRanges(bim_pr_df)
+
+        # 3. Pre-calculate TSS
+        if strategy == 'tss':
+            if 'Strand' not in df_features.columns:
+                raise ValueError("Strategy 'tss' requires 'Strand' column in mapping data.")
+            df_features['RefTSS'] = np.where(
+                df_features['Strand'] == '+',
+                df_features['Start'],
+                df_features['End']
+            )
+
+        # 4. Apply Window Size
+        df_features['Start'] = np.maximum(0, df_features['Start'] - feature_window_size)
+        df_features['End'] = df_features['End'] + feature_window_size
+
+        pr_features = pr.PyRanges(df_features)
+
+        # 5. Join
+        # Columns from pr_features (Right) that overlap with pr_bim (Left) get suffix '_b'
+        # Overlapping cols usually: Start, End.
+        # pr_bim (SNP) Start is 'Start'. pr_features (Window) Start is 'Start_b'.
+        joined = pr_bim.join(pr_features, apply_strand_suffix=False).df
+
+        logger.info(f"Initial SNP-feature overlaps: {len(joined)} pairs")
+        if not joined.empty:
+            logger.info(f"  Unique SNPs in overlaps: {joined['SNP'].nunique()} | Unique Features in overlaps: {joined['Feature'].nunique()}")
+
+        if not joined.empty:
+            # 6. Resolve Conflicts / Filter
+            if strategy == 'score':
+                if 'Score' not in joined.columns:
+                    raise ValueError("Strategy 'score' requires 'Score' column in mapping data.")
+                joined = joined.sort_values(by=['SNP', 'Score'], ascending=[True, False])
+                joined = joined.drop_duplicates(subset=['SNP'], keep='first')
+
+            elif strategy == 'tss':
+                joined['distance_to_tss'] = np.abs(joined['Start'] - joined['RefTSS'])
+                joined = joined.sort_values(by=['SNP', 'distance_to_tss'], ascending=[True, True])
+                joined = joined.drop_duplicates(subset=['SNP'], keep='first')
+
+            elif strategy == 'center':
+                # Calculate center of the feature interval (the window around feature)
+                # 'Start_b' and 'End_b' are the feature window coordinates from PyRanges join
+                joined['feature_center'] = (joined['Start_b'] + joined['End_b']) / 2.0
+                # Calculate distance from SNP position ('Start') to center
+                joined['distance_to_center'] = np.abs(joined['Start'] - joined['feature_center'])
+                # Sort and pick closest
+                joined = joined.sort_values(by=['SNP', 'distance_to_center'], ascending=[True, True])
+                joined = joined.drop_duplicates(subset=['SNP'], keep='first')
+
+            elif strategy == 'allow_repeat':
+                # No de-duplication. One SNP can map to multiple features.
+                pass
+
+            # 7. Prepare Sparse Data
+            # Row indices come from BIM 'snp_row_idx' which is preserved in join
+            row_indices = joined['snp_row_idx'].values
+            col_indices = joined['Feature'].map(feature_to_idx).values
+
+            logger.info(f"Final SNP-feature pairs after strategy '{strategy}': {len(row_indices)}")
+
+            # Data values: Use Score if available and requested, otherwise 1.0
+            if 'Score' in joined.columns and strategy in ['score', 'allow_repeat']:
+                # Use provided score
+                data_values = joined['Score'].values.astype(np.float32)
+            else:
+                # Default to 1.0 for geometric strategies (tss, center) or missing score
+                data_values = np.ones(len(row_indices), dtype=np.float32)
+
+            # Save the result for output
+            # We filter columns to make it cleaner
+            output_cols = ['SNP', 'Chromosome', 'Start', 'Feature']
+            if 'Score' in joined.columns: output_cols.append('Score')
+            if 'distance_to_tss' in joined.columns: output_cols.append('distance_to_tss')
+            if 'distance_to_center' in joined.columns: output_cols.append('distance_to_center')
+
+            # Ensure columns exist before selecting
+            output_cols = [c for c in output_cols if c in joined.columns]
+            curated_mapping_df = joined[output_cols].copy()
+            curated_mapping_df.rename(columns={'Start': 'SNP_BP'}, inplace=True)
+
+    else:
+        raise ValueError(f"Unknown mapping_type: {mapping_type}")
+
+    # === Final Matrix Construction ===
+
+    # Determine Unmapped SNPs
+    # We create a matrix of shape (M, F+1). The last column (index F) is for unmapped.
+    # Any row_idx NOT present in row_indices gets a 1.0 in the last column.
+
+    mapped_rows_set = set(row_indices)
+    all_rows = set(range(m_ref))
+    unmapped_rows = list(all_rows - mapped_rows_set)
+
+    if unmapped_rows:
+        unmapped_rows = np.array(unmapped_rows, dtype=int)
+        unmapped_cols = np.full(len(unmapped_rows), n_features, dtype=int)  # Last column index
+        unmapped_data = np.ones(len(unmapped_rows), dtype=np.float32)
+
+        # Append unmapped data
+        if len(row_indices) > 0:
+            row_indices = np.concatenate([row_indices, unmapped_rows])
+            col_indices = np.concatenate([col_indices, unmapped_cols])
+            data_values = np.concatenate([data_values, unmapped_data])
+        else:
+            row_indices = unmapped_rows
+            col_indices = unmapped_cols
+            data_values = unmapped_data
+
+    else:
+        # If all SNPs mapped, we technically don't need to append anything,
+        # but we still need the matrix shape to include the extra column.
+        pass
+
+    # Construct CSR Matrix
+    # Shape: (m_ref, n_features + 1)
+    mapping_matrix = scipy.sparse.csr_matrix(
+        (data_values, (row_indices, col_indices)),
+        shape=(m_ref, n_features + 1),
+        dtype=np.float32
+    )
+
+    return mapping_matrix, unique_feature_names, curated_mapping_df