Function test_anom_mean

dask/tests/test_order.py:1433–1495

()

def test_anom_mean():
    np = pytest.importorskip("numpy")
    xr = pytest.importorskip("xarray")

    import dask.array as da

    data = da.random.random((200, 1), chunks=(1, -1))

    ngroups = 5
    arr = xr.DataArray(
        data,
        dims=["time", "x"],
        coords={"day": ("time", np.arange(data.shape[0]) % ngroups)},
    )

    clim = arr.groupby("day").mean(dim="time")
    anom = arr.groupby("day") - clim
    anom_mean = anom.mean(dim="time")
    graph = collections_to_expr([anom_mean]).optimize().__dask_graph__()
    _, dependents = get_deps(graph)
    diags, pressure = diagnostics(graph)
    # Encoding the "best" ordering for this graph is tricky. When inspecting the
    # visualization, one sees that there are small, connected tree-like steps at
    # the beginning (tested well below in anom_mean_raw) followed by a
    # concat+transpose per group (see ngroups above). This transpose task fans
    # out into many (20-30) getitem tasks that are tiny and feed into a
    # `mean_chunk` which is the primary reducer in this graph. Therefore we want
    # to run those as quickly as possible.
    # This is difficult to assert on but the pressure is an ok-ish proxy
    assert max(pressure) <= 178
    from collections import defaultdict

    count_dependents = defaultdict(set)
    for k in dict(graph).keys():
        count_dependents[len(dependents[k])].add(k)

    # array-taker has the most dependents, but it's not what we want to look at
    n_splits = sorted(count_dependents)[-2]
    # There is a transpose/stack group that is splitting into many tasks
    # see https://github.com/dask/dask/pull/10660#discussion_r1420571664
    # the name depends on the version of xarray
    assert n_splits > 30  # at time of writing 40
    transpose_tasks = count_dependents[n_splits]
    transpose_metrics = {k: diags[k] for k in transpose_tasks}
    assert len(transpose_metrics) == ngroups, {key_split(k) for k in diags}
    # This is a pretty tightly connected graph overall and we'll have to hold
    # many tasks in memory until this can complete. However, we should ensure
    # that we get to the mean_chunks asap while the transposes are released
    # quickly.
    # If this breaks, I suggest to visually inspect the graph and run the above
    # on a single threaded LocalCluster and verify that the progress is indeed
    # in five steps (i.e. five groups)
    ages_mean_chunks = {k: v.age for k, v in diags.items() if "mean_chunk" in k[0]}
    avg_age_mean_chunks = sum(ages_mean_chunks.values()) / len(ages_mean_chunks)
    max_age_mean_chunks = max(ages_mean_chunks.values())
    ages_transpose = {k: v.age for k, v in transpose_metrics.items()}
    assert max_age_mean_chunks > 900
    assert avg_age_mean_chunks > 100