popgen and windowing

sgkit/stats/popgen.py

@@ -2,10 +2,10 @@
 import itertools
 from typing import Hashable, Optional, Sequence, Tuple, Union
 
-import dask.array as da
 import numpy as np
 from xarray import Dataset
 
+import sgkit.distarray as da
 from sgkit.cohorts import _cohorts_to_array
 from sgkit.stats.utils import assert_array_shape
 from sgkit.utils import (

sgkit/tests/test_popgen.py

@@ -1,7 +1,6 @@
 import itertools
 
 import allel
-import dask.array as da
 import msprime  # type: ignore
 import numpy as np
 import pytest
@@ -11,6 +10,7 @@
 from hypothesis import given, settings
 from hypothesis import strategies as st
 
+import sgkit.distarray as da
 from sgkit import (
     Fst,
     Garud_H,

sgkit/tests/test_window.py

@@ -1,12 +1,12 @@
 import re
 
 import allel
-import dask.array as da
 import numpy as np
 import pandas as pd
 import pytest
 import xarray as xr
 
+import sgkit.distarray as da
 from sgkit import (
     simulate_genotype_call_dataset,
     window_by_interval,

sgkit/window.py

@@ -1,9 +1,9 @@
 from typing import Any, Callable, Hashable, Iterable, Optional, Tuple, Union
 
-import dask.array as da
 import numpy as np
 from xarray import Dataset
 
+import sgkit.distarray as da
 from sgkit import variables
 from sgkit.model import get_contigs, num_contigs
 from sgkit.utils import conditional_merge_datasets, create_dataset