<jupyter_start><jupyter_text>Example 1<jupyter_code>g = allel.GenotypeArray([[[0, 0], [0, 1], [0, 1]],
                         [[0, 1], [1, 1], [1, 1]],
                         [[1, 1], [1, 1], [0, 0]],
                         [[1, 1], [1, 1], [0, 0]],
                         [[1, 1], [0, 1], [-1, -1]]], dtype='i1')
g
# the mapping array is effectively a look-up table, telling how to transform
# integer values in each row
mapping = np.array([[0, 1],  # no transformation
                    [1, 0],  # 0->1, 1->0
                    [0, 0],  # 0->0, 1->0
                    [1, 0],  # 0->1, 1->0
                    [0, 1]], # no transformation
                   dtype='i1')
# the only shape constraint is that size of first dimension must match
assert g.shape[0] == mapping.shape[0]
# this is the pure numpy implementation
expect = g.map_alleles(mapping)
expect
chunks_dim0 = 2
gd = da.from_array(g, chunks=(chunks_dim0, 2, None))
gd
md = da.from_array(mapping, chunks=(chunks_dim0, None))  # N.B., first dimension chunk size matches gd1
md
def f(block, bmapping):
    return allel.GenotypeArray(block).map_alleles(bmapping[:, 0, :])
gmapped1 = da.map_blocks(f, gd, md[:, None, :], chunks=gd.chunks)
gmapped1
actual = allel.GenotypeArray(gmapped1.compute())
actual
assert np.array_equal(expect, actual)<jupyter_output><empty_output><jupyter_text>Example 2<jupyter_code>g = allel.GenotypeArray([
    [[0, 0], [0, 1], [0, 0], [0, 1], [-1, -1]],
    [[0, 2], [1, 1], [0, 2], [1, 1], [-1, -1]],
    [[0, 0], [0, 1], [0, 0], [0, 1], [-1, -1]],
    [[0, 2], [1, 1], [0, 2], [1, 1], [-1, -1]],
    [[1, 0], [2, 1], [1, 0], [2, 1], [-1, -1]],
    [[2, 2], [-1, -1], [2, 2], [-1, -1], [-1, -1]],
    [[-1, -1], [-1, -1], [-1, -1], [-1, -1], [-1, -1]]
], dtype='i1')
g
mapping = np.array([[0, 1, 2],
                    [2, 0, 1],
                    [1, 2, 0],
                    [2, 1, 0],
                    [1, 2, 0],
                    [2, 1, 0],
                    [2, 0, 1]], dtype=np.int8)
expect = g.map_alleles(mapping)
expect
gd = da.from_array(g, chunks=(chunks_dim0, 2, None))
gd
md = da.from_array(mapping, chunks=(chunks_dim0, None))
md
def f(block, bmapping):
    return allel.GenotypeArray(block).map_alleles(bmapping[:, 0, :])
res = da.map_blocks(f, gd, md[:, None, :], chunks=gd.chunks)
res
actual = allel.GenotypeArray(res.compute())
actual
assert np.array_equal(expect, actual)<jupyter_output><empty_output><jupyter_text>Alternative solution<jupyter_code>def ff(block, bmapping):
    return allel.GenotypeArray(block[:, :, :, 0]).map_alleles(bmapping[:, 0, 0, :])
res2 = da.map_blocks(ff, gd[:, :, :, None], md[:, None, None, :], chunks=gd.chunks, drop_dims=3)
res2
actual2 = res2.compute()
actual2.shape
assert np.array_equal(expect, actual2)<jupyter_output><empty_output><jupyter_text>Example 3<jupyter_code>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
mapping = np.array([[0, 1, 2],
                    [2, 0, 1],
                    [1, 2, 0],
                    [2, 1, 0],
                    [2, 0, 1]], dtype=np.int8)
expect = [[[0, 0], [0, 1], [-1, -1]],
          [[2, 1], [0, 0], [-1, -1]],
          [[2, 1], [0, 2], [-1, -1]],
          [[0, 0], [-1, -1], [-1, -1]],
          [[-1, -1], [-1, -1], [-1, -1]]]
assert np.array_equal(expect, g.map_alleles(mapping))
gd = da.from_array(g, chunks=(2, 2, None))
gd
md = da.from_array(mapping, chunks=(2, None))
md
res = da.map_blocks(f, gd, md[:, None, :], chunks=gd.chunks)
res
actual = allel.GenotypeArray(res.compute())
actual
assert np.array_equal(expect, actual)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.23323537269569863, 'fraction_numerical': 0.06839433609404222, 'mean_word_length': 2.5121951219512195, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '507410', 'license': 'MIT', 'n_tokens_mistral': 1612, 'n_tokens_neox': 1500, 'n_words': 399, 'path': 'notebooks/sandbox/dask_map_alleles.ipynb', 'repo_name': 'mufernando/scikit-allel', 'size': '3743'}