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

tests/stats/test_u_statistic.py

@@ -0,0 +1,233 @@
+# 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 numpy as np
+from sai.stats import UStatistic
+
+
+def test_UStatistic_compute_basic():
+    # Test data
+    ref_gts = np.array([[0, 0, 1], [0, 0, 0], [1, 1, 1]])
+    tgt_gts = np.array([[1, 1, 1], [1, 0, 0], [0, 1, 0]])
+    src_gts = np.array([[0, 0, 0], [1, 1, 1], [1, 0, 1]])
+    pos = np.array([0, 1, 2])
+    w, x, y = 0.5, 0.5, ("=", 0)
+
+    expected_result = 1  # Only the first site meets the criteria
+    expected_positions = np.array([0])
+
+    u_stat = UStatistic(
+        ref_gts=ref_gts,
+        tgt_gts=tgt_gts,
+        src_gts_list=[src_gts],
+        ref_ploidy=1,
+        tgt_ploidy=1,
+        src_ploidy_list=[1],
+    )
+    results = u_stat.compute(
+        pos=pos,
+        w=w,
+        x=x,
+        y_list=[y],
+        anc_allele_available=False,
+    )
+
+    assert results["name"] == "U"
+    assert (
+        results["value"] == expected_result
+    ), f"Expected {expected_result}, got {results['value']}"
+    assert np.array_equal(results["cdd_pos"], expected_positions)
+
+    w, x, y = 0.5, 0.5, ("=", 1)
+    expected_result = 0
+    expected_positions = np.array([])
+    results = u_stat.compute(
+        pos=pos,
+        w=w,
+        x=x,
+        y_list=[y],
+        anc_allele_available=True,
+    )
+
+    assert (
+        results["value"] == expected_result
+    ), f"Expected {expected_result}, got {results['value']}"
+    assert np.array_equal(results["cdd_pos"], expected_positions)
+
+
+def test_UStatistic_compute_no_match():
+    # Test data with no matching sites
+    ref_gts = np.array([[0, 1, 1], [1, 1, 1]])
+    tgt_gts = np.array([[0, 0, 0], [1, 0, 1]])
+    src_gts = np.array([[1, 1, 1], [1, 1, 1]])
+    pos = np.array([0, 1])
+
+    # Parameters for testing with no matching sites
+    w, x, y = 0.3, 0.5, ("=", 0)
+
+    expected_result = 0  # No sites meet the criteria
+    expected_positions = np.array([])
+
+    u_stat = UStatistic(
+        ref_gts=ref_gts,
+        tgt_gts=tgt_gts,
+        src_gts_list=[src_gts],
+        ref_ploidy=1,
+        tgt_ploidy=1,
+        src_ploidy_list=[1],
+    )
+    results = u_stat.compute(
+        pos=pos,
+        w=w,
+        x=x,
+        y_list=[y],
+        anc_allele_available=False,
+    )
+
+    assert (
+        results["value"] == expected_result
+    ), f"Expected {expected_result}, got {results['value']}"
+    assert np.array_equal(results["cdd_pos"], expected_positions)
+
+
+def test_UStatistic_compute_all_match():
+    # Test data where all sites meet the criteria
+    ref_gts = np.array([[0, 0, 0], [0, 0, 0]])
+    tgt_gts = np.array([[1, 1, 1], [1, 1, 1]])
+    src_gts = np.array([[0, 0, 0], [0, 0, 0]])
+    pos = np.array([0, 1])
+
+    # Parameters for testing all sites matching
+    w, x, y = 0.5, 0.5, ("=", 0)
+
+    expected_result = 2  # All sites meet the criteria
+    expected_positions = np.array([0, 1])
+
+    u_stat = UStatistic(
+        ref_gts=ref_gts,
+        tgt_gts=tgt_gts,
+        src_gts_list=[src_gts],
+        ref_ploidy=1,
+        tgt_ploidy=1,
+        src_ploidy_list=[1],
+    )
+    results = u_stat.compute(
+        pos=pos,
+        w=w,
+        x=x,
+        y_list=[y],
+        anc_allele_available=False,
+    )
+
+    assert (
+        results["value"] == expected_result
+    ), f"Expected {expected_result}, got {results['value']}"
+    assert np.array_equal(results["cdd_pos"], expected_positions)
+
+
+def test_UStatistic_compute_with_two_sources():
+    # Test data
+    ref_gts = np.array([[0, 0, 1], [0, 0, 0], [1, 1, 1]])
+    tgt_gts = np.array([[0, 1, 1], [0, 0, 1], [1, 1, 1]])
+    src_gts1 = np.array([[1, 1, 1], [0, 1, 1], [1, 1, 1]])
+    src_gts2 = np.array([[1, 1, 1], [0, 1, 1], [1, 1, 1]])
+    pos = np.array([0, 1, 2])
+    w, x, y_list = 0.5, 0.5, [("=", 1.0), ("=", 1.0)]
+
+    # Expected result: only loci where ref_freq < w, tgt_freq > x, src_freq1 == y1, and src_freq2 == y2
+    expected_result = 1  # Only one locus meets all criteria
+    expected_positions = np.array([0])
+
+    u_stat = UStatistic(
+        ref_gts=ref_gts,
+        tgt_gts=tgt_gts,
+        src_gts_list=[src_gts1, src_gts2],
+        ref_ploidy=1,
+        tgt_ploidy=1,
+        src_ploidy_list=[1, 1],
+    )
+    results = u_stat.compute(
+        pos=pos,
+        w=w,
+        x=x,
+        y_list=y_list,
+        anc_allele_available=False,
+    )
+
+    assert (
+        results["value"] == expected_result
+    ), f"Expected {expected_result}, got {results['value']}"
+    assert np.array_equal(results["cdd_pos"], expected_positions)
+
+
+def test_UStatistic_compute_with_mixed_ploidy():
+    ref_gts = np.array([[0, 1, 0], [0, 1, 0], [2, 1, 0]])
+    tgt_gts = np.array([[1, 1, 0], [1, 1, 1], [1, 1, 1]])
+    src_gts = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]])
+    pos = np.array([0, 1, 2])
+    w, x, y = 0.5, 0.5, ("=", 0)
+    expected_result = 2
+    expected_positions = np.array([0, 2])
+
+    u_stat = UStatistic(
+        ref_gts=ref_gts,
+        tgt_gts=tgt_gts,
+        src_gts_list=[src_gts],
+        ref_ploidy=3,
+        tgt_ploidy=1,
+        src_ploidy_list=[2],
+    )
+    results = u_stat.compute(
+        pos=pos,
+        w=w,
+        x=x,
+        y_list=[y],
+        anc_allele_available=False,
+    )
+
+    assert (
+        results["value"] == expected_result
+    ), f"Expected {expected_result}, got {results['value']}"
+    assert np.array_equal(results["cdd_pos"], expected_positions)
+
+
+def test_UStatistic_compute_with_missing_keys():
+    ref_gts = np.array([[0, 1, 0], [0, 1, 0], [2, 1, 0]])
+    tgt_gts = np.array([[1, 1, 0], [1, 1, 1], [1, 1, 1]])
+    src_gts = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]])
+    pos = np.array([0, 1, 2])
+    w, x, y = 0.5, 0.5, ("=", 0)
+
+    with pytest.raises(ValueError):
+        u_stat = UStatistic(
+            ref_gts=ref_gts,
+            tgt_gts=tgt_gts,
+            src_gts_list=[src_gts],
+            ref_ploidy=3,
+            tgt_ploidy=1,
+            src_ploidy_list=[2],
+        )
+        u_stat.compute(
+            pos=pos,
+            w=w,
+            x=x,
+            y_list=[y],
+        )

tests/test___main__.py

@@ -0,0 +1,51 @@
+# 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 unittest.mock import patch, MagicMock
+from sai.__main__ import main
+
+
+@patch("sai.__main__._sai_cli_parser")  # Mock _sai_cli_parser to control its output
+@patch(
+    "sai.__main__._set_sigpipe_handler"
+)  # Mock _set_sigpipe_handler as it doesn’t need testing
+def test_main(mock_set_sigpipe_handler, mock_sai_cli_parser):
+    # Mock parser and its return values
+    mock_parser = MagicMock()
+    mock_args = MagicMock()
+    mock_args.runner = MagicMock()
+
+    # Configure _sai_cli_parser to return the mock parser
+    mock_sai_cli_parser.return_value = mock_parser
+    # Configure the mock parser to return mock_args when parse_args is called
+    mock_parser.parse_args.return_value = mock_args
+
+    # Call the main function with a test argument list
+    test_args = ["score", "--vcf", "tests/data/example.vcf", "--chr-name", "chr1"]
+    main(test_args)
+
+    # Check if _set_sigpipe_handler was called
+    mock_set_sigpipe_handler.assert_called_once()
+
+    # Verify parse_args was called with test_args
+    mock_parser.parse_args.assert_called_once_with(test_args)
+
+    # Ensure runner was called with the parsed arguments
+    mock_args.runner.assert_called_once_with(mock_args)

tests/test_sai.py

@@ -0,0 +1,102 @@
+# 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 pandas as pd
+import sai.stats
+from sai.sai import score, outlier
+
+
+@pytest.fixture
+def example_data(tmp_path):
+    # Define example file paths
+    pytest.example_vcf = "tests/data/example.vcf"
+    pytest.example_config = "tests/data/test_sai.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,
+        "output_file": str(temp_output_file),
+        "output_dir": tmp_path,
+        "config": pytest.example_config,
+    }
+
+
+def test_score(example_data):
+    # Run score function and capture output
+    score(
+        vcf_file=example_data["vcf_file"],
+        chr_name="21",
+        win_len=6666,
+        win_step=6666,
+        anc_allele_file=None,
+        output_file=example_data["output_file"],
+        config=example_data["config"],
+        num_workers=1,
+    )
+
+    # Read the generated output file and validate contents
+    df = pd.read_csv(example_data["output_file"], sep="\t")
+
+    col_name = [col for col in df.columns if col.startswith("Q")][0]
+
+    assert df[col_name].iloc[0] == 0.9, "Unexpected value in 'Q' column"
+
+
+def test_score_mixed_ploidy(example_data):
+    score(
+        vcf_file="tests/data/test.mixed.ploidy.data.vcf",
+        chr_name="21",
+        win_len=50000,
+        win_step=50000,
+        anc_allele_file=None,
+        output_file=example_data["output_file"],
+        config="tests/data/test_mixed_ploidy.config.yaml",
+        num_workers=1,
+    )
+
+    df = pd.read_csv(example_data["output_file"], sep="\t")
+
+    assert df["U"].iloc[0] == 0, "Unexpected value in 'U' column"
+    assert df["U"].iloc[1] == 1, "Unexpected value in 'U' column"
+
+
+def test_outlier(example_data):
+    output_prefix = f"{example_data['output_dir']}/outliers"
+    outlier(
+        score_file="tests/data/test.q.scores",
+        output_prefix=output_prefix,
+        quantile=0.25,
+    )
+
+    df = pd.read_csv(f"{output_prefix}.Q.0.25.outliers.tsv", sep="\t")
+    assert df["Q"].iloc[0] == 0.7
+
+    outlier(
+        score_file="tests/data/test.q.scores",
+        output_prefix=output_prefix,
+        quantile=0.75,
+    )
+
+    df = pd.read_csv(f"{output_prefix}.Q.0.75.outliers.tsv", sep="\t")
+    assert df["Q"].iloc[0] == 1.0
+    assert df["Q"].iloc[1] == 1.0