sai-pg 1.0.1__py3-none-any.whl → 1.1.0__py3-none-any.whl

tests/generators/test_window_generator.py

@@ -0,0 +1,164 @@
+# 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 pytest
+from sai.generators import WindowGenerator
+from sai.configs import PloidyConfig
+
+
+@pytest.fixture
+def test_generator():
+    # Initialize the WindowGenerator with actual data files
+    ploidy_config = PloidyConfig(
+        {
+            "ref": {"ref1": 2},
+            "tgt": {"tgt1": 2, "tgt2": 2},
+            "src": {"src1": 2, "src2": 2},
+        }
+    )
+
+    generator = WindowGenerator(
+        vcf_file="tests/data/test.data.vcf",
+        chr_name="21",
+        ref_ind_file="tests/data/test.ref.ind.list",
+        tgt_ind_file="tests/data/test.tgt.ind.list",
+        src_ind_file="tests/data/test.src.ind.list",
+        out_ind_file=None,
+        win_len=1000,  # Set window length as appropriate for testing
+        win_step=500,  # Set window step as appropriate for testing
+        ploidy_config=ploidy_config,
+    )
+    return generator
+
+
+def test_initialization(test_generator):
+    # Verify initialization parameters
+    assert test_generator.win_len == 1000
+    assert test_generator.win_step == 500
+
+
+def test_window_generator(test_generator):
+    # Collect data from generator
+    data_list = list(test_generator.get())
+
+    # Ensure windows were generated
+    assert len(data_list) == 380
+
+    # Inspect first window's contents for expected format and data
+    first_window = data_list[0]
+    assert "chr_name" in first_window
+    assert "start" in first_window
+    assert "end" in first_window
+    assert "ref_pop" in first_window
+    assert first_window["ref_pop"] == "ref1"
+    assert "tgt_pop" in first_window
+    assert "src_pop_list" in first_window
+    assert "ref_gts" in first_window
+    assert "tgt_gts" in first_window
+    assert "src_gts_list" in first_window
+    assert "ploidy_config" in first_window
+
+
+def test_none_window_generator(test_generator):
+    test_generator.ref_data = None
+
+    data_list = list(test_generator.get())
+    assert len(data_list) == 380
+
+    first_window = data_list[0]
+    assert first_window["ref_gts"] is None
+    assert first_window["tgt_gts"] is None
+    assert first_window["src_gts_list"] is None
+    assert first_window["ploidy_config"].get_ploidy("ref")[0] == 2
+    assert first_window["ploidy_config"].get_ploidy("tgt")[0] == 2
+    assert first_window["ploidy_config"].get_ploidy("src")[0] == 2
+    assert len(first_window["pos"]) == 0
+
+
+def test_len(test_generator):
+    # Check if __len__ provides a reasonable window count
+    assert len(test_generator) == 380
+
+
+@pytest.fixture
+def test_generator_two_sources():
+    # Initialize the WindowGenerator with num_src=2 for testing two-source combinations
+    ploidy_config = PloidyConfig(
+        {
+            "ref": {"ref1": 2},
+            "tgt": {"tgt1": 2, "tgt2": 2},
+            "src": {"src1": 2, "src2": 2},
+        }
+    )
+
+    generator = WindowGenerator(
+        vcf_file="tests/data/test.data.vcf",
+        chr_name="21",
+        ref_ind_file="tests/data/test.ref.ind.list",
+        tgt_ind_file="tests/data/test.tgt.ind.list",
+        src_ind_file="tests/data/test.src.ind.list",
+        out_ind_file=None,
+        win_len=1000,  # Set window length as appropriate for testing
+        win_step=500,  # Set window step as appropriate for testing
+        num_src=2,  # Set to 2 to test two-source combinations
+        ploidy_config=ploidy_config,
+    )
+    return generator
+
+
+def test_initialization_two_sources(test_generator_two_sources):
+    # Verify initialization parameters for two-source generator
+    assert test_generator_two_sources.win_len == 1000
+    assert test_generator_two_sources.win_step == 500
+    assert test_generator_two_sources.ploidy_config.get_ploidy("ref")[0] == 2
+    assert test_generator_two_sources.ploidy_config.get_ploidy("tgt")[0] == 2
+    assert test_generator_two_sources.ploidy_config.get_ploidy("src")[0] == 2
+    assert test_generator_two_sources.num_src == 2
+
+
+def test_window_generator_with_two_sources(test_generator_two_sources):
+    # Collect data from generator with two sources
+    data_list = list(test_generator_two_sources.get())
+
+    # Ensure windows were generated and have two sources in src_pop_list
+    assert len(data_list) > 0  # Ensure data was generated
+    first_window = data_list[0]
+
+    # Check keys in the first window
+    assert "start" in first_window
+    assert "end" in first_window
+    assert "ref_pop" in first_window
+    assert "tgt_pop" in first_window
+    assert "src_pop_list" in first_window
+    assert "ref_gts" in first_window
+    assert "tgt_gts" in first_window
+    assert "src_gts_list" in first_window
+    assert "ploidy_config" in first_window
+
+    # Verify that src_pop_list contains exactly two source populations
+    assert len(first_window["src_pop_list"]) == 2
+    assert (
+        len(first_window["src_gts_list"]) == 2
+    )  # Ensure two sets of genotypes in src_gts_list
+
+
+def test_len_two_sources(test_generator_two_sources):
+    # Check if __len__ provides a reasonable window count with two-source combinations
+    assert len(test_generator_two_sources) > 0  # Ensure it counts windows correctly

tests/multiprocessing/test_mp_manager.py

@@ -0,0 +1,92 @@
+# 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 sai.multiprocessing import mp_manager
+from sai.preprocessors import DataPreprocessor
+from sai.generators import DataGenerator
+
+
+class TmpDataPreprocessor(DataPreprocessor):
+    def run(self, rep):
+        return [rep]
+
+    def process_items(self, items):
+        print(items)
+
+
+class TmpDataGenerator(DataGenerator):
+    def __init__(self, start_rep=0, nrep=5):
+        self.start_rep = start_rep
+        self.nrep = nrep
+
+    def get(self):
+        for i in range(self.start_rep, self.start_rep + self.nrep):
+            yield {"rep": i}
+
+    def __len__(self):
+        return self.nrep
+
+
+class FailureDataPreprocessor(DataPreprocessor):
+    def run(self, rep):
+        raise Exception("Simulating failure by stopping.")
+
+    def process_items(self, items):
+        print(items)
+
+
+def test_mp_manager(capfd):
+    nprocess = 2
+    nrep = 5
+
+    data_processor = TmpDataPreprocessor()
+    generator = TmpDataGenerator(nrep=nrep)
+
+    mp_manager(
+        data_processor=data_processor, data_generator=generator, nprocess=nprocess
+    )
+
+    # Define the expected set of outputs
+    expected_set = {"[0, 1, 2, 3, 4]"}
+
+    # Capture the actual output and convert it to a set of strings
+    captured = capfd.readouterr()
+    actual_set = {captured.out.strip()}
+
+    # Compare the actual and expected sets
+    assert actual_set == expected_set, "The output does not match the expected results."
+
+
+def test_mp_manager_failure(capfd):
+    nprocess = 2
+    data_processor = FailureDataPreprocessor()
+    generator = TmpDataGenerator(nrep=5)
+
+    mp_manager(
+        data_processor=data_processor, data_generator=generator, nprocess=nprocess
+    )
+
+    # Use capfd to capture stdout and stderr
+    captured = capfd.readouterr()
+
+    # Assertions to verify expected output and behavior
+    assert "Simulating failure by stopping." in captured.err
+    assert "did not complete successfully. Initiating shutdown." in captured.out
+    assert "All workers are terminated." in captured.out

tests/multiprocessing/test_mp_pool.py

@@ -0,0 +1,79 @@
+# 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 pytest
+from sai.multiprocessing.mp_pool import mp_pool, mp_worker
+
+
+class DataPreprocessor:
+    """Mock class simulating a data processor."""
+
+    def run(self, x: int) -> int:
+        """Mock processing: returns the square of x."""
+        return x**2
+
+    def process_items(self, results: list) -> None:
+        """Stores the processed results for validation."""
+        self.final_results = results
+
+
+class DataGenerator:
+    """Mock class simulating a data generator."""
+
+    def __init__(self, data: list):
+        self.data = data
+
+    def get(self):
+        """Yields data in dictionary format."""
+        for x in self.data:
+            yield {"x": x}
+
+
+@pytest.mark.parametrize(
+    "params, expected",
+    [
+        ((DataPreprocessor(), {"x": 2}), 4),
+        ((DataPreprocessor(), {"x": 3}), 9),
+        ((DataPreprocessor(), {"x": 4}), 16),
+    ],
+)
+def test_mp_worker(params, expected):
+    """Tests mp_worker to ensure correct processing."""
+    assert mp_worker(params) == expected
+
+
+def test_mp_pool():
+    """Tests mp_pool to ensure parallel processing works correctly."""
+    data_processor = DataPreprocessor()
+    data_generator = DataGenerator([1, 2, 3, 4, 5])
+    nprocess = 2
+
+    # Run multiprocessing pool
+    mp_pool(data_processor, data_generator, nprocess)
+
+    # Validate results
+    assert hasattr(data_processor, "final_results")  # Ensure results are stored
+    assert sorted(data_processor.final_results) == [
+        1,
+        4,
+        9,
+        16,
+        25,
+    ]  # Check correctness

tests/parsers/test_argument_validation.py

@@ -0,0 +1,133 @@
+# 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 argparse
+import pytest
+from sai.parsers.argument_validation import positive_int
+from sai.parsers.argument_validation import positive_number
+from sai.parsers.argument_validation import between_zero_and_one
+from sai.parsers.argument_validation import existed_file
+from sai.parsers.argument_validation import validate_stat_type
+
+
+def test_positive_int():
+    # Valid positive integer
+    assert positive_int("5") == 5
+
+    # Not a positive integer (zero)
+    with pytest.raises(argparse.ArgumentTypeError, match="0 is not a positive integer"):
+        positive_int("0")
+
+    # Negative integer
+    with pytest.raises(
+        argparse.ArgumentTypeError, match="-1 is not a positive integer"
+    ):
+        positive_int("-1")
+
+    # Non-integer string
+    with pytest.raises(argparse.ArgumentTypeError, match="abc is not a valid integer"):
+        positive_int("abc")
+
+
+def test_positive_number():
+    # Valid positive number
+    assert positive_number("3.14") == 3.14
+
+    # Not a positive number (zero)
+    with pytest.raises(argparse.ArgumentTypeError, match="0 is not a positive number"):
+        positive_number("0")
+
+    # Negative number
+    with pytest.raises(
+        argparse.ArgumentTypeError, match="-2.5 is not a positive number"
+    ):
+        positive_number("-2.5")
+
+    # Non-numeric string
+    with pytest.raises(argparse.ArgumentTypeError, match="xyz is not a valid number"):
+        positive_number("xyz")
+
+
+def test_between_zero_and_one():
+    # Values within range
+    assert between_zero_and_one("0.5") == 0.5
+    assert between_zero_and_one("0") == 0
+    assert between_zero_and_one("1") == 1
+
+    # Values out of range
+    with pytest.raises(argparse.ArgumentTypeError, match="1.5 is not between 0 and 1"):
+        between_zero_and_one("1.5")
+
+    with pytest.raises(argparse.ArgumentTypeError, match="-0.1 is not between 0 and 1"):
+        between_zero_and_one("-0.1")
+
+    # Non-numeric string
+    with pytest.raises(
+        argparse.ArgumentTypeError, match="not_a_number is not a valid number"
+    ):
+        between_zero_and_one("not_a_number")
+
+
+def test_existed_file(tmp_path):
+    # Create a temporary file for testing
+    temp_file = tmp_path / "temp.txt"
+    temp_file.write_text("This is a test file.")
+
+    # Validate an existing file path
+    assert existed_file(str(temp_file)) == str(temp_file)
+
+    # Validate a non-existent file path
+    with pytest.raises(
+        argparse.ArgumentTypeError, match="non_existent_file is not found"
+    ):
+        existed_file("non_existent_file")
+
+
+def test_valid_inputs():
+    assert validate_stat_type("U50") == "U50"
+    assert validate_stat_type("Q05") == "Q05"
+    assert validate_stat_type("Q95") == "Q95"
+    assert validate_stat_type("Q99") == "Q99"
+
+
+def test_invalid_inputs():
+    with pytest.raises(argparse.ArgumentTypeError):
+        validate_stat_type("U")
+
+    with pytest.raises(argparse.ArgumentTypeError):
+        validate_stat_type("Q")
+
+    with pytest.raises(argparse.ArgumentTypeError):
+        validate_stat_type("Q5")
+
+    with pytest.raises(argparse.ArgumentTypeError):
+        validate_stat_type("U100")
+
+    with pytest.raises(argparse.ArgumentTypeError):
+        validate_stat_type("Q100")
+
+    with pytest.raises(argparse.ArgumentTypeError):
+        validate_stat_type("X50")
+
+    with pytest.raises(argparse.ArgumentTypeError):
+        validate_stat_type("Qabc")
+
+    with pytest.raises(argparse.ArgumentTypeError):
+        validate_stat_type("")

tests/parsers/test_outlier_parser.py

@@ -0,0 +1,53 @@
+# 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 pytest
+import argparse
+from sai.parsers.outlier_parser import add_outlier_parser
+
+
+@pytest.fixture
+def parser():
+    # Initialize the argument parser with a subparser for the 'outlier' command
+    main_parser = argparse.ArgumentParser()
+    subparsers = main_parser.add_subparsers(dest="command")
+    add_outlier_parser(subparsers)
+    return main_parser
+
+
+def test_add_outlier_parser(parser):
+    # Simulate command-line arguments to parse
+    args = parser.parse_args(
+        [
+            "outlier",
+            "--score",
+            "tests/data/example.scores",
+            "--output-prefix",
+            "output/test_outliers",
+            "--quantile",
+            "0.95",
+        ]
+    )
+
+    # Validate parsed arguments
+    assert args.command == "outlier"
+    assert args.score == "tests/data/example.scores"
+    assert args.output_prefix == "output/test_outliers"
+    assert args.quantile == 0.95

tests/parsers/test_score_parser.py

@@ -0,0 +1,63 @@
+# 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 pytest
+import argparse
+from sai.parsers.score_parser import add_score_parser
+
+
+@pytest.fixture
+def parser():
+    # Initialize the argument parser with a subparser for the 'score' command
+    main_parser = argparse.ArgumentParser()
+    subparsers = main_parser.add_subparsers(dest="command")
+    add_score_parser(subparsers)
+    return main_parser
+
+
+def test_add_score_parser(parser):
+    # Simulate command-line arguments to parse
+    args = parser.parse_args(
+        [
+            "score",
+            "--vcf",
+            "tests/data/example.vcf",
+            "--chr-name",
+            "chr1",
+            "--win-len",
+            "50000",
+            "--win-step",
+            "10000",
+            "--output",
+            "output/results.tsv",
+            "--config",
+            "tests/data/test_config.yaml",
+        ]
+    )
+
+    # Validate parsed arguments
+    assert args.command == "score"
+    assert args.vcf == "tests/data/example.vcf"
+    assert args.chr_name == "chr1"
+    assert args.win_len == 50000
+    assert args.win_step == 10000
+    assert args.anc_alleles is None
+    assert args.output == "output/results.tsv"
+    assert args.config == "tests/data/test_config.yaml"

tests/preprocessors/test_chunk_preprocessor.py

@@ -0,0 +1,79 @@
+# 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 pytest
+import yaml
+import sai.stats
+from sai.preprocessors import ChunkPreprocessor
+from sai.configs import PloidyConfig, StatConfig
+
+
+@pytest.fixture
+def example_data(tmp_path):
+    # Define example file paths
+    pytest.example_vcf = "tests/data/example.vcf"
+    pytest.example_ref_ind_list = "tests/data/example.ref.ind.list"
+    pytest.example_tgt_ind_list = "tests/data/example.tgt.ind.list"
+    pytest.example_src_ind_list = "tests/data/example.src.ind.list"
+    pytest.example_config = "tests/data/example.config.yaml"
+
+    # Create a temporary output file path for the score function
+    temp_output_file = tmp_path / "output.tsv"
+
+    return {
+        "vcf_file": pytest.example_vcf,
+        "ref_ind_file": pytest.example_ref_ind_list,
+        "tgt_ind_file": pytest.example_tgt_ind_list,
+        "src_ind_file": pytest.example_src_ind_list,
+        "config": pytest.example_config,
+        "output_file": str(temp_output_file),
+        "output_dir": tmp_path,
+    }
+
+
+def test_chunk_preprocessor(example_data):
+    with open(example_data["config"], "r") as f:
+        config = yaml.safe_load(f)
+
+    stat_config = StatConfig(config["statistics"])
+    ploidy_config = PloidyConfig(config["ploidies"])
+
+    preprocessor = ChunkPreprocessor(
+        vcf_file=example_data["vcf_file"],
+        ref_ind_file=example_data["ref_ind_file"],
+        tgt_ind_file=example_data["tgt_ind_file"],
+        src_ind_file=example_data["src_ind_file"],
+        out_ind_file=None,
+        win_len=6666,
+        win_step=6666,
+        num_src=1,
+        anc_allele_file=None,
+        output_file=example_data["output_file"],
+        stat_config=stat_config,
+        ploidy_config=ploidy_config,
+    )
+
+    results = preprocessor.run(
+        chr_name="21",
+        start=0,
+        end=6666,
+    )
+
+    assert results[0]["Q"] == 0.9