diff --git a/ruranges/numpy.py b/ruranges/numpy.py index d300cf4..b6b0923 100644 --- a/ruranges/numpy.py +++ b/ruranges/numpy.py @@ -1476,6 +1476,10 @@ def _dispatch_unary( if prefix == "complement_numpy": chrom_ids = _normalize_groups(extra_kw["chrom_len_ids"], name="chrom_len_ids") + if groups is None: + if len(chrom_ids) != 1: + raise ValueError("`groups=None` requires exactly one chromosome in `chrom_len_ids`.") + groups = np.full(len(starts_arr), chrom_ids[0], dtype=np.uint32) extra_kw["chrom_len_ids"] = chrom_ids chrom_lens = _as_1d_array(extra_kw["chrom_lens"], name="chrom_lens") pos_extra_arrays.append(("chrom_lens", chrom_lens)) diff --git a/tests/test_complement_implicit_groups.py b/tests/test_complement_implicit_groups.py new file mode 100644 index 0000000..7f4e89f --- /dev/null +++ b/tests/test_complement_implicit_groups.py @@ -0,0 +1,27 @@ +import numpy as np +import pytest + +from ruranges.numpy import complement + + +@pytest.mark.parametrize("chromosome", [0, 1, 23, 24, 2**32 - 1]) +@pytest.mark.parametrize("include_first", [False, True]) +def test_implicit_group_uses_the_supplied_chromosome(chromosome, include_first): + kwargs = dict( + starts=np.array([30, 10]), + ends=np.array([40, 20]), + chrom_len_ids=np.array([chromosome], dtype=np.uint64), + chrom_lens=np.array([100]), + include_first_interval=include_first, + ) + implicit = complement(groups=None, **kwargs) + explicit = complement(groups=np.full(2, chromosome, dtype=np.uint32), **kwargs) + for actual, expected in zip(implicit, explicit): + np.testing.assert_array_equal(actual, expected) + n = 3 if include_first else 2 + np.testing.assert_array_equal(implicit[0], [chromosome] * n) + np.testing.assert_array_equal(implicit[1], [0, 20, 40] if include_first else [20, 40]) + np.testing.assert_array_equal(implicit[2], [10, 30, 100] if include_first else [30, 100]) + assert implicit[0].dtype == np.uint32 + # The attribution index identifies the group, not a specific interval. + assert np.all(implicit[3] < len(kwargs["starts"]))