sai-pg 1.0.0__py3-none-any.whl

sai/utils/preprocessors/chunk_preprocessor.py

@@ -0,0 +1,152 @@
+# Copyright 2025 Xin Huang
+#
+# GNU General Public License v3.0
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, please see
+#
+#    https://www.gnu.org/licenses/gpl-3.0.en.html
+
+
+from typing import Any
+from sai.utils.generators import WindowGenerator
+from sai.utils.preprocessors import DataPreprocessor
+from .feature_preprocessor import FeaturePreprocessor
+
+
+class ChunkPreprocessor(DataPreprocessor):
+    """
+    Preprocesses VCF data in genomic windows and applies feature preprocessing.
+
+    This class generates genomic windows from a VCF file, processes them
+    with specified reference, target, and source individuals, and computes
+    feature vectors using the provided feature preprocessor.
+    """
+
+    def __init__(
+        self,
+        vcf_file: str,
+        ref_ind_file: str,
+        tgt_ind_file: str,
+        src_ind_file: str,
+        win_len: int,
+        win_step: int,
+        w: float,
+        y: list[float],
+        output_file: str,
+        stat_type: str,
+        anc_allele_file: str = None,
+        num_src: int = 1,
+    ):
+        """
+        Initializes a new instance of ChunkPreprocessor.
+
+        Parameters
+        ----------
+        vcf_file : str
+            Path to the VCF file to process.
+        ref_ind_file : str
+            Path to the file containing reference individual IDs.
+        tgt_ind_file : str
+            Path to the file containing target individual IDs.
+        src_ind_file : str
+            Path to the file containing source individual IDs.
+        win_len : int
+            Window length for generating genomic windows.
+        win_step : int
+            Step size for sliding windows across the genome.
+        w : float
+            Parameter w for feature vector computation.
+        y : list of float
+            List of y parameters for feature vector computation.
+        output_file : str
+            Path to the output file for storing feature vectors.
+        stat_type : str
+            Type of statistic to compute for feature vectors.
+        anc_allele_file : str, optional
+            Path to the ancestral allele file. If None, ancestral allele
+            information is considered unavailable.
+        num_src : int, optional
+            Number of source populations to use. Default is 1.
+        """
+        self.vcf_file = vcf_file
+        self.ref_ind_file = ref_ind_file
+        self.tgt_ind_file = tgt_ind_file
+        self.src_ind_file = src_ind_file
+        self.win_len = win_len
+        self.win_step = win_step
+        self.anc_allele_file = anc_allele_file
+        self.num_src = num_src
+
+        anc_allele_available = anc_allele_file is not None
+
+        self.feature_preprocessor = FeaturePreprocessor(
+            w=w,
+            y=y,
+            output_file=output_file,
+            stat_type=stat_type,
+            anc_allele_available=anc_allele_available,
+        )
+
+    def run(self, chr_name: str, start: int, end: int) -> list[dict[str, Any]]:
+        """
+        Runs the preprocessing pipeline on a specific chromosome region.
+
+        Generates genomic windows within the specified chromosome region,
+        processes each window to compute feature vectors, and aggregates the results.
+
+        Parameters
+        ----------
+        chr_name : str
+            Name of the chromosome to process.
+        start : int
+            Start position (1-based, inclusive) of the region to process.
+        end : int
+            End position (1-based, exclusive) of the region to process.
+
+        Returns
+        -------
+        list of dict of {str: Any}
+            A list of dictionaries containing computed feature vectors for each genomic window.
+        """
+        window_generator = WindowGenerator(
+            vcf_file=self.vcf_file,
+            chr_name=chr_name,
+            start=start,
+            end=end,
+            ref_ind_file=self.ref_ind_file,
+            tgt_ind_file=self.tgt_ind_file,
+            src_ind_file=self.src_ind_file,
+            win_len=self.win_len,
+            win_step=self.win_step,
+            anc_allele_file=self.anc_allele_file,
+            num_src=self.num_src,
+        )
+
+        items = []
+
+        for item in window_generator.get():
+            items.extend(self.feature_preprocessor.run(**item))
+
+        return items
+
+    def process_items(self, items: list[dict[str, Any]]) -> None:
+        """
+        Processes and writes computed feature vectors to the output.
+
+        Parameters
+        ----------
+        items : list of dict of {str: Any}
+            A list of dictionaries containing computed feature vectors for each genomic window.
+        """
+        self.feature_preprocessor.process_items(items)

sai/utils/preprocessors/data_preprocessor.py

@@ -0,0 +1,94 @@
+# Copyright 2025 Xin Huang
+#
+# GNU General Public License v3.0
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, please see
+#
+#    https://www.gnu.org/licenses/gpl-3.0.en.html
+
+
+from typing import Any
+from abc import ABC, abstractmethod
+
+
+class DataPreprocessor(ABC):
+    """
+    Abstract base class for preprocessing genomic data.
+
+    This class defines a common interface for various data preprocessing operations,
+    such as filtering, normalization, and transformation of genomic data. Subclasses
+    should implement specific methods to handle data processing tasks, ensuring a
+    consistent way to run operations and manage the output of processed data.
+
+    Methods:
+    --------
+    run(**kwargs) -> Any:
+        Execute the core data processing task. Subclasses must define this method
+        to carry out specific preprocessing tasks such as filtering, normalization,
+        or transformation. This method should return the processed data, which
+        will then be handled by the main process to manage further steps or output.
+
+    process_items(items, **kwargs) -> None:
+        Handle the output or further processing of data once the `run` method
+        has completed. This allows subclasses to define how processed data
+        should be managed, such as saving results to a file, database, or converting
+        the data to a specific format for future analysis.
+    """
+
+    @abstractmethod
+    def run(self, **kwargs) -> Any:
+        """
+        Abstract method to run the preprocessing operations.
+
+        Subclasses must implement this method to perform specific preprocessing
+        tasks based on the initialized parameters and any additional keyword
+        arguments.
+
+        Parameters:
+        -----------
+        **kwargs : dict
+            Additional keyword arguments that may be required for specific
+            preprocessing operations.
+
+        Returns:
+        --------
+        processed_data : Any
+            The result of the preprocessing task, which can be further handled
+            by the `process_items` method.
+        """
+        pass
+
+    @abstractmethod
+    def process_items(self, items: Any, **kwargs) -> None:
+        """
+        Abstract method to handle the output or post-processing of data.
+
+        Subclasses must implement this method to define how the processed data
+        should be managed. This could include saving the data to a file,
+        transforming it into a new format, or preparing it for the next step
+        of analysis.
+
+        Parameters:
+        -----------
+        items : Any
+            The processed data returned by the `run` method, which will be managed
+            or output according to the logic defined in this method.
+
+        **kwargs : dict
+            Additional keyword arguments that can be used for customizing the
+            output process. For example, this may include options like `output_file`
+            to specify where the data should be saved or other settings to control
+            the output format.
+        """
+        pass

sai/utils/preprocessors/feature_preprocessor.py

@@ -0,0 +1,211 @@
+# Copyright 2025 Xin Huang
+#
+# GNU General Public License v3.0
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, please see
+#
+#    https://www.gnu.org/licenses/gpl-3.0.en.html
+
+
+import numpy as np
+from typing import Any
+from sai.stats.features import calc_u, calc_q
+from sai.utils.preprocessors import DataPreprocessor
+
+
+class FeaturePreprocessor(DataPreprocessor):
+    """
+    A preprocessor subclass for generating feature vectors from genomic data.
+
+    This class extends DataPreprocessor to include additional functionality for creating
+    feature vectors based on genomic variants, reference and target individual genotypes,
+    and window-based genomic statistics.
+    """
+
+    def __init__(
+        self,
+        w: float,
+        y: list[float],
+        output_file: str,
+        stat_type: str,
+        anc_allele_available: bool = False,
+    ):
+        """
+        Initializes FeatureVectorsPreprocessor with specific frequency thresholds
+        and output file for storing generated feature vectors.
+
+        Parameters
+        ----------
+        w : float
+            Frequency threshold for `calc_u` and `calc_q`.
+        y : list[float]
+            List of frequency thresholds for `calc_u` and `calc_q`.
+        output_file : str
+            Path to the output file to save processed feature vectors.
+        stat_type: str,
+            Specifies the type of statistic to compute.
+            - "UXX" (e.g., "U50", "U90") : Compute the U statistic using `calc_u()`.
+            - "QXX" (e.g., "Q95", "Q50") : Compute the Q statistic using `calc_q()`,
+        anc_allele_available: bool, optional
+            If True, ancestral allele information is available.
+            If False, ancestral allele information is unavailable.
+            Default is False.
+
+        Raises
+        ------
+        ValueError
+            If `stat_type` is not in a valid format. Must be either: 'UXX' or 'QXX'.
+        """
+        self.w = w
+        self.y = y
+        self.output_file = output_file
+        self.anc_allele_available = anc_allele_available
+        if not (
+            len(stat_type) == 3
+            and stat_type[0] in {"U", "Q"}
+            and stat_type[1:].isdigit()
+        ):
+            raise ValueError(
+                f"Invalid stat_type format: {stat_type}. Expected format 'UXX' or 'QXX' (e.g., 'U50' or 'Q95')."
+            )
+        self.stat_prefix = stat_type[0]
+        self.threshold = int(stat_type[1:]) / 100
+
+    def run(
+        self,
+        chr_name: str,
+        ref_pop: str,
+        tgt_pop: str,
+        src_pop_list: list[str],
+        start: int,
+        end: int,
+        pos: np.ndarray,
+        ref_gts: np.ndarray,
+        tgt_gts: np.ndarray,
+        src_gts_list: list[np.ndarray],
+        ploidy: int,
+    ) -> list[dict[str, Any]]:
+        """
+        Generates feature vectors for a specified genomic window.
+
+        Parameters
+        ----------
+        chr_name : str
+            Chromosome name.
+        ref_pop : str
+            Reference population name.
+        tgt_pop : str
+            Target population name.
+        src_pop_list : list[str]
+            List of source population names.
+        start : int
+            Start position of the genomic window.
+        end : int
+            End position of the genomic window.
+        pos : np.ndarray
+            A 1D numpy array where each element represents the genomic position.
+        ref_gts : np.ndarray
+            Genotype data for the reference population.
+        tgt_gts : np.ndarray
+            Genotype data for the target population.
+        src_gts_list : list[np.ndarray]
+            List of genotype arrays for each source population.
+        ploidy: int
+            Ploidy of the genome.
+
+        Returns
+        -------
+        list[dict[str, Any]]
+            A list containing a dictionary of calculated feature vectors for the genomic window.
+        """
+        items = {
+            "chr_name": chr_name,
+            "start": start,
+            "end": end,
+            "ref_pop": ref_pop,
+            "tgt_pop": tgt_pop,
+            "src_pop_list": src_pop_list,
+            "nsnps": len(pos),
+        }
+
+        if (
+            (ref_gts is None)
+            or (tgt_gts is None)
+            or (src_gts_list is None)
+            or (ploidy is None)
+        ):
+            items["statistic"] = np.nan
+            items["candidates"] = np.array([])
+        elif self.stat_prefix == "U":
+            items["statistic"], items["candidates"] = calc_u(
+                ref_gts=ref_gts,
+                tgt_gts=tgt_gts,
+                src_gts_list=src_gts_list,
+                pos=pos,
+                w=self.w,
+                x=self.threshold,
+                y_list=self.y,
+                ploidy=ploidy,
+                anc_allele_available=self.anc_allele_available,
+            )
+        elif self.stat_prefix == "Q":
+            items["statistic"], items["candidates"] = calc_q(
+                ref_gts=ref_gts,
+                tgt_gts=tgt_gts,
+                src_gts_list=src_gts_list,
+                pos=pos,
+                w=self.w,
+                y_list=self.y,
+                quantile=self.threshold,
+                ploidy=ploidy,
+                anc_allele_available=self.anc_allele_available,
+            )
+        else:
+            raise ValueError(
+                f"Invalid stat_type: {self.stat_type}. Must be 'U' or 'QXX' (e.g., 'Q95')."
+            )
+
+        return [items]
+
+    def process_items(self, items: list[dict[str, Any]]) -> None:
+        """
+        Processes and writes a single dictionary of feature vectors to the output file.
+
+        Parameters
+        ----------
+        items : dict[str, Any]
+            A dictionary containing feature vectors for a genomic window.
+        """
+        with open(
+            self.output_file, "a"
+        ) as f:  # Open in append mode for continuous writing
+            lines = []
+            for item in items:
+                src_pop_str = ",".join(item["src_pop_list"])
+                candidates = (
+                    "NA"
+                    if item["candidates"].size == 0
+                    else ",".join(
+                        f"{item['chr_name']}:{pos}" for pos in item["candidates"]
+                    )
+                )
+
+                line = (
+                    f"{item['chr_name']}\t{item['start']}\t{item['end']}\t"
+                    f"{item['ref_pop']}\t{item['tgt_pop']}\t{src_pop_str}\t"
+                    f"{item['nsnps']}\t{item['statistic']}\t{candidates}\n"
+                )
+                lines.append(line)
+
+            f.writelines(lines)