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

tests/preprocessors/test_feature_preprocessor.py

@@ -0,0 +1,223 @@
+# 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
+import sai.stats
+from sai.generators import WindowGenerator
+from sai.preprocessors import FeaturePreprocessor
+from sai.configs import PloidyConfig, StatConfig
+
+
+@pytest.fixture
+def feature_preprocessor():
+    # Create an instance of FeaturePreprocessor with thresholds and temporary output file
+    stat_config = StatConfig(
+        {
+            "U": {
+                "ref": {"ref1": 0.3},
+                "tgt": {"tgt1": 0.5},
+                "src": {"src1": "=1", "src2": "=1"},
+            },
+            "Q": {
+                "ref": {"ref1": 0.3},
+                "tgt": {"tgt1": 0.95},
+                "src": {"src1": "=0.2", "src2": "=0.4"},
+            },
+        }
+    )
+
+    return FeaturePreprocessor(
+        output_file="test_output.tsv",
+        stat_config=stat_config,
+    )
+
+
+def test_run(feature_preprocessor):
+    # Define mock input data to test the run method
+    chr_name = "21"
+    ref_pop = "ref1"
+    tgt_pop = "tgt1"
+    src_pop_list = ["src1", "src2"]
+    start, end = 1000, 2000
+
+    # Create mock genotype data
+    ref_gts = np.array([[0, 0, 1], [1, 1, 0], [0, 1, 1]])
+    tgt_gts = np.array([[0, 1, 1], [1, 1, 1], [0, 0, 1]])
+    src_gts_list = [
+        np.array([[0, 0, 0], [1, 0, 0], [1, 1, 1]]),
+        np.array([[1, 1, 1], [0, 1, 1], [0, 0, 1]]),
+    ]
+
+    pos = np.array([100, 200, 300])
+
+    ploidy_config = PloidyConfig(
+        {
+            "ref": {"ref1": 1},
+            "tgt": {"tgt1": 1},
+            "src": {"src1": 1},
+        }
+    )
+
+    # Run the method
+    result = feature_preprocessor.run(
+        chr_name=chr_name,
+        ref_pop=ref_pop,
+        tgt_pop=tgt_pop,
+        src_pop_list=src_pop_list,
+        start=start,
+        end=end,
+        pos=pos,
+        ref_gts=ref_gts,
+        tgt_gts=tgt_gts,
+        src_gts_list=src_gts_list,
+        ploidy_config=ploidy_config,
+    )
+
+    # Check that the result contains the expected keys
+    assert result[0]["chr_name"] == chr_name
+    assert result[0]["start"] == start
+    assert result[0]["end"] == end
+    assert result[0]["ref_pop"] == ref_pop
+    assert result[0]["tgt_pop"] == tgt_pop
+    assert result[0]["src_pop_list"] == src_pop_list
+    assert "U" in result[0]
+    assert "Q" in result[0]
+
+    result = feature_preprocessor.run(
+        chr_name=chr_name,
+        ref_pop=ref_pop,
+        tgt_pop=tgt_pop,
+        src_pop_list=src_pop_list,
+        start=start,
+        end=end,
+        pos=pos,
+        ref_gts=None,
+        tgt_gts=None,
+        src_gts_list=None,
+        ploidy_config=None,
+    )
+
+    assert np.isnan(result[0]["U"])
+    assert np.isnan(result[0]["Q"])
+
+
+def test_process_items(feature_preprocessor, tmp_path):
+    # Generate a temporary output file path using tmp_path
+    temp_output = tmp_path / "test_output.tsv"
+    feature_preprocessor.output_file = str(temp_output)
+
+    # Define a mock items dictionary to test process_items method
+    items = {
+        "chr_name": "21",
+        "start": 1000,
+        "end": 2000,
+        "ref_pop": "ref1",
+        "tgt_pop": "tgt1",
+        "src_pop_list": ["src1", "src2"],
+        "nsnps": 10,
+        "U": 5,
+        "Q": 0.8,
+        "cdd_pos": {
+            "U": np.array([]),
+            "Q": np.array([]),
+        },
+    }
+
+    # Run the process_items method
+    feature_preprocessor.process_items([items])
+
+    # Check the output file content
+    with open(temp_output, "r") as f:
+        lines = f.readlines()
+        assert len(lines) == 1  # Ensure only one line is written
+        expected_output = "21\t1000\t2000\tref1\ttgt1\tsrc1,src2\t10\t5\t0.8\n"
+        assert lines[0] == expected_output
+
+
+@pytest.fixture
+def example_data():
+    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"
+
+
+def test_run_from_file(example_data, tmp_path):
+    # Set up the WindowGenerator
+    ploidy_config = PloidyConfig(
+        {
+            "ref": {"AFR": 2},
+            "tgt": {"CHB": 2},
+            "src": {"Nean": 2},
+        }
+    )
+
+    generator = WindowGenerator(
+        vcf_file=pytest.example_vcf,
+        chr_name=21,
+        ref_ind_file=pytest.example_ref_ind_list,
+        tgt_ind_file=pytest.example_tgt_ind_list,
+        src_ind_file=pytest.example_src_ind_list,
+        out_ind_file=None,
+        win_len=6666,
+        win_step=6666,
+        ploidy_config=ploidy_config,
+    )
+
+    # Create a temporary output file path using tmp_path
+    temp_output_file = tmp_path / "output.tsv"
+
+    # Initialize the FeaturePreprocessor with the temporary output file path
+    stat_config = StatConfig(
+        {"U": {"ref": {"AFR": 0.3}, "tgt": {"CHB": 0.5}, "src": {"Nean": "=1"}}}
+    )
+
+    preprocessor = FeaturePreprocessor(
+        output_file=str(temp_output_file),
+        stat_config=stat_config,
+    )
+
+    # Run the generator and preprocessor
+    for window_data in generator.get():
+        items = preprocessor.run(
+            chr_name=window_data["chr_name"],
+            ref_pop=window_data["ref_pop"],
+            tgt_pop=window_data["tgt_pop"],
+            src_pop_list=window_data["src_pop_list"],
+            start=window_data["start"],
+            end=window_data["end"],
+            pos=window_data["pos"],
+            ref_gts=window_data["ref_gts"],
+            tgt_gts=window_data["tgt_gts"],
+            src_gts_list=window_data["src_gts_list"],
+            ploidy_config=window_data["ploidy_config"],
+        )
+        preprocessor.process_items(items)
+        assert items[0]["U"] == 3
+
+    # Check that the temporary file was created and is not empty
+    assert temp_output_file.exists()
+    assert temp_output_file.stat().st_size > 0
+
+    # Optionally, open and inspect contents of temp_output_file here if necessary
+    with open(temp_output_file, "r") as f:
+        lines = f.readlines()
+        assert len(lines) > 0  # Ensure some content was written

tests/registries/test_registries.py

@@ -0,0 +1,74 @@
+# 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.registries import (
+    StatRegistry,
+)
+
+
+@pytest.mark.parametrize(
+    "RegistryClass",
+    [
+        StatRegistry,
+    ],
+)
+def test_register_and_get(RegistryClass):
+    registry = RegistryClass()
+
+    class DummyClass:
+        pass
+
+    registry.register("dummy")(DummyClass)
+    assert "dummy" in registry.list_registered()
+
+    retrieved = registry.get("dummy")
+    assert retrieved is DummyClass
+
+
+@pytest.mark.parametrize(
+    "RegistryClass",
+    [
+        StatRegistry,
+    ],
+)
+def test_duplicate_registration(RegistryClass):
+    registry = RegistryClass()
+
+    class DummyClass:
+        pass
+
+    registry.register("dummy")(DummyClass)
+
+    with pytest.raises(ValueError):
+        registry.register("dummy")(DummyClass)
+
+
+@pytest.mark.parametrize(
+    "RegistryClass",
+    [
+        StatRegistry,
+    ],
+)
+def test_get_unregistered(RegistryClass):
+    registry = RegistryClass()
+
+    with pytest.raises(KeyError):
+        registry.get("nonexistent")

tests/stats/test_danc_statistic.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
+
+
+import numpy as np
+from sai.stats import DancStatistic
+
+
+def test_DancStatistic_compute():
+    ref_gts = np.array([[0, 0], [0, 0], [1, 1]])  # Referece population
+    tgt_gts = np.array([[1, 0], [0, 1], [0, 1]])  # Taget population
+    src_gts = np.array([[0, 1], [1, 0], [1, 0]])  # Source population
+    # out_gts = None  # No outgroup provided
+
+    # D ancestral
+    # baaa - abaa = 0.25 - 0.5 = -0.25
+    # baaa + abaa = 0.25 + 0.5 = 0.75
+    # (baaa - abaa) / (baaa + abaa) = -1/3
+
+    danc_stat = DancStatistic(
+        ref_gts=ref_gts,
+        tgt_gts=tgt_gts,
+        src_gts_list=[src_gts],
+        ref_ploidy=1,
+        tgt_ploidy=1,
+        src_ploidy_list=[1],
+    )
+    results = danc_stat.compute()
+
+    expected_result = -1 / 3
+
+    assert results["name"] == "Danc"
+    assert np.isclose(
+        results["value"][0], expected_result
+    ), f"Expected {expected_result}, but got {results['value'][0]}"

tests/stats/test_dd_statistic.py

@@ -0,0 +1,45 @@
+# 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 sai.stats import DdStatistic
+
+
+def test_DdStatistic_compute():
+    ref_gts = np.array([[1, 1], [0, 0]])
+    tgt_gts = np.array([[1, 0], [0, 1]])
+    src_gts = np.array([[0, 1], [1, 1]])
+
+    dd_stat = DdStatistic(
+        ref_gts=ref_gts,
+        tgt_gts=tgt_gts,
+        src_gts_list=[src_gts],
+        ref_ploidy=1,
+        tgt_ploidy=1,
+        src_ploidy_list=[1],
+    )
+    results = dd_stat.compute()
+
+    expected_result = 0.5
+
+    assert results["name"] == "DD"
+    assert np.isclose(
+        results["value"][0], expected_result
+    ), f"Expected {expected_result}, but got {results['value'][0]}"

tests/stats/test_df_statistic.py

@@ -0,0 +1,73 @@
+# 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 sai.stats import DfStatistic
+
+
+def test_DfStatistic_compute():
+    ref_gts = np.array([[0, 0], [0, 0], [1, 1]])  # Referece population
+    tgt_gts = np.array([[1, 0], [0, 1], [0, 1]])  # Taget population
+    src_gts = np.array([[0, 1], [1, 0], [1, 0]])  # Source population
+    # out_gts = None  # No outgroup provided
+
+    # ref_freq = [0, 0, 1]
+    # tgt_freq = [0.5, 0.5, 0.5]
+    # src_freq = [0.5, 0.5, 0.5]
+    # out_freq = [0, 0, 0]
+
+    # pattern: 'abba'
+    # site 0: (1-0)*0.5*0.5*(1-0) = 0.25
+    # site 1: (1-0)*0.5*0.5*(1-0) = 0.25
+    # site 2: (1-1)*0.5*0.5*(1-0) = 0
+    # sum = 0.5
+
+    # pattern: 'baba'
+    # site 0: 0*(1-0.5)*0.5*(1-0) = 0
+    # site 1: 0*(1-0.5)*0.5*(1-0) = 0
+    # site 2: 1*(1-0.5)*0.5*(1-0) = 0.25
+    # sum = 0.25
+
+    # pattern: 'bbaa'
+    # site 0: 0*0.5*(1-0.5)*(1-0) = 0
+    # site 1: 0*0.5*(1-0.5)*(1-0) = 0
+    # site 2: 1*0.5*(1-0.5)*(1-0) = 0.25
+    # sum = 0.25
+
+    # abba - baba = 0.5 - 0.25 = 0.25
+    # abba + baba + 2 * bbaa = 0.5 + 0.25 + 2*0.25 = 1.25
+
+    # Call the function with the test input
+    df_stat = DfStatistic(
+        ref_gts=ref_gts,
+        tgt_gts=tgt_gts,
+        src_gts_list=[src_gts],
+        ref_ploidy=1,
+        tgt_ploidy=1,
+        src_ploidy_list=[1],
+    )
+    results = df_stat.compute()
+
+    # Check the result
+    expected_result = 0.2
+    assert results["name"] == "df"
+    assert np.isclose(
+        results["value"][0], expected_result
+    ), f"Expected {expected_result}, but got {results['value'][0]}"

tests/stats/test_dplus_statistic.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 numpy as np
+from sai.stats import DplusStatistic
+
+
+def test_DplusStatistic_compute():
+    ref_gts = np.array([[0, 0], [0, 0], [1, 1]])  # Referece population
+    tgt_gts = np.array([[1, 0], [0, 1], [0, 1]])  # Taget population
+    src_gts = np.array([[0, 1], [1, 0], [1, 0]])  # Source population
+    # out_gts = None  # No outgroup provided
+
+    # ref_freq = [0, 0, 1]
+    # tgt_freq = [0.5, 0.5, 0.5]
+    # src_freq = [0.5, 0.5, 0.5]
+    # out_freq = [0, 0, 0]
+
+    # pattern: 'abba'
+    # site 0: (1-0)*0.5*0.5*(1-0) = 0.25
+    # site 1: (1-0)*0.5*0.5*(1-0) = 0.25
+    # site 2: (1-1)*0.5*0.5*(1-0) = 0
+    # sum = 0.5
+
+    # pattern: 'baba'
+    # site 0: 0*(1-0.5)*0.5*(1-0) = 0
+    # site 1: 0*(1-0.5)*0.5*(1-0) = 0
+    # site 2: 1*(1-0.5)*0.5*(1-0) = 0.25
+    # sum = 0.25
+
+    # pattern: 'baaa'
+    # site 0: 0*(1-0.5)*(1-0.5)*(1-0) = 0
+    # site 1: 0*(1-0.5)*(1-0.5)*(1-0) = 0
+    # site 2: 1*(1-0.5)*(1-0.5)*(1-0) = 0.25
+    # sum = 0.25
+
+    # pattern: 'abaa'
+    # site 0: (1-0)*0.5*(1-0.5)*(1-0) = 0.25
+    # site 1: (1-0)*0.5*(1-0.5)*(1-0) = 0.25
+    # site 2: (1-1)*0.5*(1-0.5)*(1-0) = 0
+    # sume = 0.5
+
+    # abba - baba + baaa - abaa = 0.5 - 0.25 + 0.25 - 0.5 = 0
+    # abba + baba + baaa + abaa = 0.5 + 0.25 + 0.25 + 0.5 = 1.5
+
+    # Call the function with the test input
+    dplus_stat = DplusStatistic(
+        ref_gts=ref_gts,
+        tgt_gts=tgt_gts,
+        src_gts_list=[src_gts],
+        ref_ploidy=1,
+        tgt_ploidy=1,
+        src_ploidy_list=[1],
+    )
+    results = dplus_stat.compute()
+
+    # Check the result
+    expected_result = 0
+    assert results["name"] == "Dplus"
+    assert np.isclose(
+        results["value"][0], expected_result
+    ), f"Expected {expected_result}, but got {results['value'][0]}"

tests/stats/test_fd_statistic.py

@@ -0,0 +1,68 @@
+# 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 sai.stats import FdStatistic
+
+
+def test_FdStatistic_compute():
+    ref_gts = np.array([[0, 1], [1, 0], [0, 1]])  # Reference population
+    tgt_gts = np.array([[1, 0], [0, 1], [1, 0]])  # Target population
+    src_gts = np.array([[1, 1], [1, 1], [1, 1]])  # Source population
+    # out_gts = None  # No outgroup provided
+
+    # ref_freq = [0.5, 0.5, 0.5]
+    # tgt_freq = [0.5, 0.5, 0.5]
+    # src_freq = [1, 1, 1]
+    # out_freq = [0, 0, 0]
+
+    # dnr_freq = src_freq = [1, 1, 1]
+    # pattern: 'abba'
+    # site 0: (1-0.5)*1*1*(1-0) = 0.5
+    # site 1: (1-0.5)*1*1*(1-0) = 0.5
+    # site 2: (1-0.5)*1*1*(1-0) = 0.5
+    # sum = 1.5
+
+    # pattern: 'baba'
+    # site 0: 0.5*(1-1)*1*(1-0) = 0
+    # site 1: 0.5*(1-1)*1*(1-0) = 0
+    # site 2: 0.5*(1-1)*1*(1-0) = 0
+    # sum = 0
+
+    # abba_d - baba_d = 1.5
+
+    # Call the function with the test input
+    fd_stat = FdStatistic(
+        ref_gts=ref_gts,
+        tgt_gts=tgt_gts,
+        src_gts_list=[src_gts],
+        ref_ploidy=1,
+        tgt_ploidy=1,
+        src_ploidy_list=[1],
+    )
+
+    results = fd_stat.compute()
+
+    # Check the result
+    expected_result = 0
+    assert results["name"] == "fd"
+    assert np.isclose(
+        results["value"][0], expected_result
+    ), f"Expected {expected_result}, but got {results['value'][0]}"