In [1]:

    

import numpy as np
import dask.async
import dask.array as da
import sys
sys.path.insert(0, '../..')
%load_ext autoreload
%autoreload 1
%aimport allel
%aimport allel.model
%aimport allel.model.dask
allel.__version__


    

    
Out[1]:





'0.20.0.feature_dask'

In [3]:

    

g = allel.GenotypeArray([
    [[0, 0], [0, 1], [-1, -1]],
    [[0, 2], [1, 1], [-1, -1]],
    [[1, 0], [2, 1], [-1, -1]],
    [[2, 2], [-1, -1], [-1, -1]],
    [[-1, -1], [-1, -1], [-1, -1]]
], dtype='i1')
g


    

    
Out[3]:





GenotypeArray((5, 3, 2), dtype=int8)



0
1
2




0
0/0
0/1
./.


1
0/2
1/1
./.


2
1/0
2/1
./.


3
2/2
./.
./.


4
./.
./.
./.

In [5]:

    

sel0 = [0, 2]
sel1 = [0, 2]
expect = g.subset(sel0, sel1)
expect


    

    
Out[5]:





GenotypeArray((2, 2, 2), dtype=int8)



0
1




0
0/0
./.


1
1/0
./.

In [6]:

    

gd = da.from_array(g, chunks=(2, 2, None))
gd


    

    
Out[6]:





dask.array<from-ar..., shape=(5, 3, 2), dtype=int8, chunksize=(2, 2, 2)>

In [10]:

    

gd[sel0][:, sel1]


    

    
Out[10]:





dask.array<getitem..., shape=(2, 2, 2), dtype=int8, chunksize=(1, 1, 2)>

In [11]:

    

actual = allel.GenotypeArray(gd[sel0][:, sel1].compute())
actual


    

    
Out[11]:





GenotypeArray((2, 2, 2), dtype=int8)



0
1




0
0/0
./.


1
1/0
./.