<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 This notebook has some profiling of Dask used to make a selection along both first and second axes of a large-ish multidimensional array. The use case is making selections of genotype data, e.g., as required for making a web-browser for genotype data as in www.malariagen.net/apps/ag1000g.
Step1:  Real data
Step2:  Synthetic data
<ASSISTANT_TASK:>
Python Code:
import zarr; print('zarr', zarr.__version__)
import dask; print('dask', dask.__version__)
import dask.array as da
import numpy as np
Explanation: This notebook has some profiling of Dask used to make a selection along both first and second axes of a large-ish multidimensional array. The use case is making selections of genotype data, e.g., as required for making a web-browser for genotype data as in www.malariagen.net/apps/ag1000g.
End of explanation
# here's the real data
callset = zarr.open_group('/kwiat/2/coluzzi/ag1000g/data/phase1/release/AR3.1/variation/main/zarr2/zstd/ag1000g.phase1.ar3',
                          mode='r')
callset
# here's the array we're going to work with
g = callset['3R/calldata/genotype']
g
# wrap as dask array with very simple chunking of first dim only
%time gd = da.from_array(g, chunks=(g.chunks[0], None, None))
gd
# load condition used to make selection on first axis
dim0_condition = callset['3R/variants/FILTER_PASS'][:]
dim0_condition.shape, dim0_condition.dtype, np.count_nonzero(dim0_condition)
# invent a random selection for second axis
dim1_indices = sorted(np.random.choice(765, size=100, replace=False))
# setup the 2D selection - this is the slow bit
%time gd_sel = gd[dim0_condition][:, dim1_indices]
gd_sel
# now load a slice from this new selection - quick!
%time gd_sel[1000000:1100000].compute(optimize_graph=False)
# what's taking so long?
import cProfile
cProfile.run('gd[dim0_condition][:, dim1_indices]', sort='time')
cProfile.run('gd[dim0_condition][:, dim1_indices]', sort='cumtime')
Explanation: Real data
End of explanation
# create a synthetic dataset for profiling
a = zarr.array(np.random.randint(-1, 4, size=(20000000, 200, 2), dtype='i1'),
               chunks=(10000, 100, 2), compressor=zarr.Blosc(cname='zstd', clevel=1, shuffle=2))
a
# create a synthetic selection for first axis
c = np.random.randint(0, 2, size=a.shape[0], dtype=bool)
# create a synthetic selection for second axis
s = sorted(np.random.choice(a.shape[1], size=100, replace=False))
%time d = da.from_array(a, chunks=(a.chunks[0], None, None))
d
%time ds = d[c][:, s]
cProfile.run('d[c][:, s]', sort='time')
%time ds[1000000:1100000].compute(optimize_graph=False)
# problem is in fact just the dim0 selection
cProfile.run('d[c]', sort='time')
Explanation: Synthetic data
End of explanation
<END_TASK>