groupby on coordinate (not on dimension)

[shcrela]

Hi All,
Sorry if this question was asked before, I wasn't able to find it.
So I have a DataArray with two dimensions, 'points' and 'wavelengths'.
On the 'points' dimension I have multiple coordinates (Time, X, Y, Z, etc.) and on 'wavelengths' dimension only one coordinate.
Now I want to group by unique values on the 'Y' coordinate and get the new DataArray with calculated median value in each group.
Ok, so I did:
da.groupby('Y', restore_coord_dims=True).median()
The problem with this is that I loose all the other coordinates which lie on the same dimension.
I would like keep the median value of each group for the remaining coordinates (Time, X, Z, etc.).
Can someone instruct me on how to do this?
Thanks !
Dejan.

[benbovy]

Hi @shcrela,

You could first convert the coordinates to data variables:

da.reset_coords(["Time", "X", "Z"]).groupby("Y").median()

Then use set_coords() if you want it promoted as coordinates again.

[shcrela]

Thanks a lot @benbovy !

In order to get back the same kind of object, I ended up with something like:

def m2s(da, gc="Y"):
    da.name = "Intensities"
    newda = da.reset_coords(set(da.coords.keys()) - {gc}).groupby(gc).median()
    newda = newda.set_coords(list(da.coords.keys()))
    newda = newda.to_array()
    newda = newda.squeeze("variable", drop=True)
    newda = newda.swap_dims({str(gc): "points"})
    newda.attrs = da.attrs
    return newda

Even tough the above function is clearly not optimized, it came to ma as a surprise that the "manual" workaround I used before I read your answer, did not behave all that worse for da.shape = (20536, 1001):

def map2series(da, gc="Y"):
    c_dict = dict()
    for c in da.coords:
        c_dict[c] = []  # a list for each coordinate so that we can append 1 by 1

    for val in np.unique(da[gc].data):  # For each unique value of gc
        # We create a temporary DataArray
        temp_da = da[da[gc] == val]
        for c in da.coords:
            if c == "Time":  # For Time coord, we take the first value
                c_dict[c].append(temp_da[c].data[0])
            else:  # For other coords, we calculate the median
                c_dict[c].append(np.median(temp_da[c].data))

    coord_dict = dict()
    for k, v in c_dict.items():
        coord_dict[k] = ("points", v, da[k].attrs)

    newda = xr.DataArray(data=da.groupby(gc).median().data,
                          dims=("points", "RamanShifts"),
                          coords={**coord_dict,
                                  "shifts": ("RamanShifts",
                                             da.shifts.data, {"units": "1/cm"})
                                 },
                          attrs=da.attrs)
    return newda

m2s: 343 ms ± 4.61 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
map2series: 419 ms ± 7.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

[dcherian]

median runs in a for loop so yes it shoud be pretty similar to what you have.