Rename window to window_by_index, and add window_by_position

Author: tomwhite

docs/getting_started.rst

@@ -231,7 +231,7 @@ then with windows.
 
     # Define windows of size 20 variants. This creates a new dimension called `windows`, and
     # some new variables for internal use.
-    ds = sg.window(ds, 20)
+    ds = sg.window(ds, size=20)
 
     # The diversity statistic is now computed for every window
     sg.diversity(ds, merge=False)
@@ -269,7 +269,7 @@ example shows how to give cohorts names.
 .. ipython:: python
     :okwarning:
 
-    ds = sg.window(ds, 20)
+    ds = sg.window(ds, size=20)
     ds = sg.Fst(ds)
 
     cohort_names = ["Africa", "Asia", "Europe"]
@@ -310,7 +310,7 @@ Xarray and Pandas operations in a single pipeline:
         # Apply filter to include variants present across > 80% of samples
         .pipe(lambda ds: ds.sel(variants=ds.variant_call_rate > .8))
         # Create windows of size 20 variants
-        .pipe(lambda ds: sg.window(ds, 20))
+        .pipe(lambda ds: sg.window(ds, size=20))
         # Assign a "cohort" variable that splits samples into two groups
         .assign(sample_cohort=np.repeat([0, 1], ds.dims['samples'] // 2))
         # Compute Fst between the groups

sgkit/tests/test_ld.py

@@ -156,7 +156,7 @@ def test_vs_skallel(args):
 
     ds = simulate_genotype_call_dataset(n_variant=x.shape[0], n_sample=x.shape[1])
     ds["dosage"] = (["variants", "samples"], da.asarray(x).rechunk({0: chunks}))
-    ds = window(ds, size, step)
+    ds = window(ds, size=size, step=step)
 
     ldm = ld_matrix(ds, threshold=threshold)
     has_duplicates = ldm.compute().duplicated(subset=["i", "j"]).any()

sgkit/tests/test_window.py

@@ -13,7 +13,8 @@
     _get_windows,
     has_windows,
     moving_statistic,
-    window,
+    window_by_index,
+    window_by_position,
 )
 
 
@@ -71,23 +72,23 @@ def test_moving_statistic__min_chunksize_smaller_than_size():
         moving_statistic(values, np.sum, size=3, step=3, dtype=values.dtype)
 
 
-def test_window():
+def test_window_by_index():
     ds = simulate_genotype_call_dataset(n_variant=10, n_sample=3, seed=0)
     assert not has_windows(ds)
-    ds = window(ds, 2, 2)
+    ds = window_by_index(ds, size=2, step=2)
     assert has_windows(ds)
     np.testing.assert_equal(ds[window_contig].values, [0, 0, 0, 0, 0])
     np.testing.assert_equal(ds[window_start].values, [0, 2, 4, 6, 8])
     np.testing.assert_equal(ds[window_stop].values, [2, 4, 6, 8, 10])
 
     with pytest.raises(MergeWarning):
-        window(ds, 2, 2)
+        window_by_index(ds, size=2, step=2)
 
 
-def test_window__default_step():
+def test_window_by_index__default_step():
     ds = simulate_genotype_call_dataset(n_variant=10, n_sample=3, seed=0)
     assert not has_windows(ds)
-    ds = window(ds, 2)
+    ds = window_by_index(ds, size=2)
     assert has_windows(ds)
     np.testing.assert_equal(ds[window_contig].values, [0, 0, 0, 0, 0])
     np.testing.assert_equal(ds[window_start].values, [0, 2, 4, 6, 8])
@@ -120,13 +121,13 @@ def test_window__default_step():
         ),
     ],
 )
-def test_window__multiple_contigs(
+def test_window_by_index__multiple_contigs(
     n_variant, n_contig, window_contigs_exp, window_starts_exp, window_stops_exp
 ):
     ds = simulate_genotype_call_dataset(
         n_variant=n_variant, n_sample=1, n_contig=n_contig
     )
-    ds = window(ds, 2, 2)
+    ds = window_by_index(ds, size=2, step=2)
     np.testing.assert_equal(ds[window_contig].values, window_contigs_exp)
     np.testing.assert_equal(ds[window_start].values, window_starts_exp)
     np.testing.assert_equal(ds[window_stop].values, window_stops_exp)
@@ -193,3 +194,43 @@ def test_get_chunked_windows(
     )
     np.testing.assert_equal(rel_window_starts_actual, rel_window_starts_exp)
     np.testing.assert_equal(windows_per_chunk_actual, windows_per_chunk_exp)
+
+
+def test_window_by_position():
+    ds = simulate_genotype_call_dataset(n_variant=5, n_sample=3, seed=0)
+    assert not has_windows(ds)
+    ds["variant_position"] = (
+        ["variants"],
+        np.array([1, 4, 6, 8, 12]),
+    )
+    ds = window_by_position(ds, size=5)
+    assert has_windows(ds)
+    np.testing.assert_equal(ds[window_contig].values, [0, 0, 0, 0, 0])
+    np.testing.assert_equal(ds[window_start].values, [0, 1, 2, 3, 4])
+    np.testing.assert_equal(ds[window_stop].values, [2, 4, 4, 5, 5])
+
+
+def test_window_by_position__multiple_contigs():
+    ds = simulate_genotype_call_dataset(n_variant=10, n_sample=3, n_contig=2)
+    ds["variant_position"] = (
+        ["variants"],
+        np.array([1, 4, 6, 8, 12, 1, 21, 25, 40, 55]),
+    )
+    ds = window_by_position(ds, size=10)
+    assert has_windows(ds)
+    np.testing.assert_equal(ds[window_contig].values, [0, 0, 0, 0, 0, 1, 1, 1, 1, 1])
+    np.testing.assert_equal(ds[window_start].values, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+    np.testing.assert_equal(ds[window_stop].values, [4, 5, 5, 5, 5, 6, 8, 8, 9, 10])
+
+
+def test_window_by_position__offset():
+    ds = simulate_genotype_call_dataset(n_variant=10, n_sample=3, n_contig=2)
+    ds["variant_position"] = (
+        ["variants"],
+        np.array([1, 4, 6, 8, 12, 1, 21, 25, 40, 55]),
+    )
+    ds = window_by_position(ds, size=10, offset=-5)
+    assert has_windows(ds)
+    np.testing.assert_equal(ds[window_contig].values, [0, 0, 0, 0, 0, 1, 1, 1, 1, 1])
+    np.testing.assert_equal(ds[window_start].values, [0, 0, 0, 1, 3, 5, 6, 6, 8, 9])
+    np.testing.assert_equal(ds[window_stop].values, [2, 4, 4, 5, 5, 6, 8, 8, 9, 10])

sgkit/window.py

@@ -1,9 +1,10 @@
-from typing import Any, Callable, Iterable, Optional, Tuple, Union
+from typing import Any, Callable, Hashable, Iterable, Optional, Tuple, Union
 
 import dask.array as da
 import numpy as np
 from xarray import Dataset
 
+from sgkit import variables
 from sgkit.utils import conditional_merge_datasets, create_dataset
 from sgkit.variables import window_contig, window_start, window_stop
 
@@ -12,27 +13,32 @@
 # Window definition (user code)
 
 
-def window(
+def window_by_index(
     ds: Dataset,
+    *,
     size: int,
     step: Optional[int] = None,
+    variant_contig: Hashable = variables.variant_contig,
     merge: bool = True,
 ) -> Dataset:
-    """Add fixed-size windowing information to a dataset.
+    """Add window information to a dataset, measured by number of variants.
 
     Windows are defined over the ``variants`` dimension, and are
     used by some downstream functions to calculate statistics for
-    each window.
+    each window. Windows never span contigs.
 
     Parameters
     ----------
     ds
         Genotype call dataset.
     size
-        The window size (number of variants).
+        The window size, measured by number of variants.
     step
         The distance (number of variants) between start positions of windows.
         Defaults to ``size``.
+    variant_contig
+        Name of variable containing variant contig indexes.
+        Defined by :data:`sgkit.variables.variant_contig_spec`.
     merge
         If True (the default), merge the input dataset and the computed
         output variables into a single dataset, otherwise return only
@@ -43,12 +49,84 @@ def window(
     -------
     A dataset containing the following variables:
 
+    - :data:`sgkit.variables.window_contig_spec` (windows):
+      The index values of window contigs.
     - :data:`sgkit.variables.window_start_spec` (windows):
       The index values of window start positions.
     - :data:`sgkit.variables.window_stop_spec` (windows):
       The index values of window stop positions.
     """
     step = step or size
+    return _window_per_contig(ds, variant_contig, merge, _get_windows, size, step)
+
+
+window = window_by_index
+
+
+def window_by_position(
+    ds: Dataset,
+    *,
+    size: int,
+    offset: int = 0,
+    variant_contig: Hashable = variables.variant_contig,
+    variant_position: Hashable = variables.variant_position,
+    merge: bool = True,
+) -> Dataset:
+    """Add window information to a dataset, measured by distance along the genome.
+
+    Windows are defined over the ``variants`` dimension, and are
+    used by some downstream functions to calculate statistics for
+    each window. Windows never span contigs.
+
+    Parameters
+    ----------
+    ds
+        Genotype call dataset.
+    size
+        The window size, measured by number of base pairs.
+    offset
+        The window offset, measured by number of base pairs. Defaults
+        to no offset. For centered windows, use a negative offset that
+        is half the window size.
+    variant_contig
+        Name of variable containing variant contig indexes.
+        Defined by :data:`sgkit.variables.variant_contig_spec`.
+    variant_position
+        Name of variable containing variant positions.
+        Must be monotonically increasing within a contig.
+        Defined by :data:`sgkit.variables.variant_position_spec`.
+    merge
+        If True (the default), merge the input dataset and the computed
+        output variables into a single dataset, otherwise return only
+        the computed output variables.
+        See :ref:`dataset_merge` for more details.
+
+    Returns
+    -------
+    A dataset containing the following variables:
+
+    - :data:`sgkit.variables.window_contig_spec` (windows):
+      The index values of window contigs.
+    - :data:`sgkit.variables.window_start_spec` (windows):
+      The index values of window start positions.
+    - :data:`sgkit.variables.window_stop_spec` (windows):
+      The index values of window stop positions.
+    """
+
+    positions = ds[variant_position].values
+    return _window_per_contig(
+        ds, variant_contig, merge, _get_windows_by_position, size, offset, positions
+    )
+
+
+def _window_per_contig(
+    ds: Dataset,
+    variant_contig: Hashable,
+    merge: bool,
+    windowing_fn: Callable[..., Any],
+    *args: Any,
+    **kwargs: Any,
+) -> Dataset:
     n_variants = ds.dims["variants"]
     n_contigs = len(ds.attrs["contigs"])
     contig_ids = np.arange(n_contigs)
@@ -60,7 +138,9 @@ def window(
     contig_window_starts = []
     contig_window_stops = []
     for i in range(n_contigs):
-        starts, stops = _get_windows(contig_bounds[i], contig_bounds[i + 1], size, step)
+        starts, stops = windowing_fn(
+            contig_bounds[i], contig_bounds[i + 1], *args, **kwargs
+        )
         contig_window_starts.append(starts)
         contig_window_stops.append(stops)
         contig_window_contigs.append(np.full_like(starts, i))
@@ -97,6 +177,18 @@ def _get_windows(
     return window_starts, window_stops
 
 
+def _get_windows_by_position(
+    start: int, stop: int, size: int, offset: int, positions: ArrayLike
+) -> Tuple[ArrayLike, ArrayLike]:
+    contig_pos = positions[start:stop]
+    if offset == 0:
+        window_starts = np.arange(contig_pos.shape[0]) + start
+    else:
+        window_starts = np.searchsorted(contig_pos, contig_pos + offset) + start
+    window_stops = np.searchsorted(contig_pos, contig_pos + offset + size) + start
+    return window_starts, window_stops
+
+
 # Computing statistics for windows (internal code)