Overlapping Groups

This post is motivated by the problem of computing the Meridional Overturning Circulation. One of the steps is a binned average over latitude, over regions of the World Ocean. Commonly we want to average globally, as well as over the Atlantic, and the Indo-Pacific. Generally group-by problems involve non-overlapping groups. In this example, the “global” group overlaps with the “Indo-Pacific” and “Atlantic” groups. Below we consider a simplified version of this problem.

Consider the following labels:

import numpy as np
import xarray as xr

from flox.xarray import xarray_reduce

labels = xr.DataArray(
    [1, 2, 3, 1, 2, 3, 0, 0, 0],
    dims="x",
    name="label",
)
labels

<xarray.DataArray 'label' (x: 9)> Size: 72B
array([1, 2, 3, 1, 2, 3, 0, 0, 0])
Dimensions without coordinates: x

xarray.DataArray

'label'

• x: 9

• 1 2 3 1 2 3 0 0 0

  array([1, 2, 3, 1, 2, 3, 0, 0, 0])

These labels are non-overlapping. So when we reduce this data array over those labels along x

da = xr.ones_like(labels)
da

<xarray.DataArray 'label' (x: 9)> Size: 72B
array([1, 1, 1, 1, 1, 1, 1, 1, 1])
Dimensions without coordinates: x

xarray.DataArray

'label'

• x: 9

• 1 1 1 1 1 1 1 1 1

  array([1, 1, 1, 1, 1, 1, 1, 1, 1])

we get (note the reduction over x is implicit here):

xarray_reduce(da, labels, func="sum")

<xarray.DataArray 'label' (label: 4)> Size: 32B
array([3, 2, 2, 2])
Coordinates:
  * label    (label) int64 32B 0 1 2 3

xarray.DataArray

'label'

• label: 4

• 3 2 2 2

  array([3, 2, 2, 2])

• Coordinates: (1)
  • label
    (label)
    int64
    0 1 2 3

    array([0, 1, 2, 3])

Now let’s also calculate the sum where labels is either 1 or 2. We could easily compute this using the grouped result but here we use this simple example for illustration. The trick is to add a new dimension with new labels (here 4) in the appropriate locations.

# assign 4 where label == 1 or 2, and -1 otherwise
newlabels = xr.where(labels.isin([1, 2]), 4, -1)

# concatenate along a new dimension y;
# note y is not present on da
expanded = xr.concat([labels, newlabels], dim="y")
expanded

<xarray.DataArray 'label' (y: 2, x: 9)> Size: 144B
array([[ 1,  2,  3,  1,  2,  3,  0,  0,  0],
       [ 4,  4, -1,  4,  4, -1, -1, -1, -1]])
Dimensions without coordinates: y, x

xarray.DataArray

'label'

• y: 2
• x: 9

• 1 2 3 1 2 3 0 0 0 4 4 -1 4 4 -1 -1 -1 -1

  array([[ 1,  2,  3,  1,  2,  3,  0,  0,  0],
         [ 4,  4, -1,  4,  4, -1, -1, -1, -1]])

Now we reduce over x and the new dimension y (again implicitly) to get the appropriate sum under label=4 (and label=-1). We can discard the value accumulated under label=-1 later.

xarray_reduce(da, expanded, func="sum")

<xarray.DataArray 'label' (label: 6)> Size: 48B
array([5, 3, 2, 2, 2, 4])
Coordinates:
  * label    (label) int64 48B -1 0 1 2 3 4

xarray.DataArray

'label'

• label: 6

• 5 3 2 2 2 4

  array([5, 3, 2, 2, 2, 4])

• Coordinates: (1)
  • label
    (label)
    int64
    -1 0 1 2 3 4

    array([-1,  0,  1,  2,  3,  4])

This way we compute all the reductions we need, in a single pass over the data.

This technique generalizes to more complicated aggregations. The trick is to

• generate appropriate labels

• concatenate these new labels along a new dimension (y) absent on the object being reduced (da), and

• reduce over that new dimension in addition to any others.