How to extract original coordinates from a rolling window?

[nepomnyi]

Hello dear Xarray team and community.

I haven't found an answer to my question, thus, I would like to treat it as a feature suggestion.

Unfortunately, this form doesn't allow me to attach a picture which would help me to explain my issue. So, please, bear with my written explanations.

I have a 2D velocity field. It looks like this one.

At every point is space, I have a velocity vector. Every space point is represented by an x and y coordinate. Every velocity vector is represented by its x component (let's call it u) and its y component (let's call it v).

I arrange my velocity field in an xr.Dataset in the following manner: I have coordinates x and y and I have two data variables u and v.

Consider a 9x9 window around some point in my velocity field. This point has coordinates (xP, yP) and velocity vector (uP, vP). I want to do some calculations within this window involving coordinates and velocity vectors of every point within this window. For the brevity of example, let's say I want to add (xP, yP) to the coordinates of every point within the window. Then I want to add (uP, vP) to the velocity vectors at every point within the window. Then I want to multiply coordinates by the velocity vectors at every point within the window. Then I want to add all the results within the window, thus obtaining one scalar value for the window.

I want to perform this procedure for every point in my original velocity field.

In order to code this procedure, I constructed a rolling window object, stacked the rolling windows, grouped the rolling windows and mapped my custom function to each group (i.e. to each rolling window).

.map sends a rolling window to my custom function. In my custom function, I extract velocity vectors and coordinates of every point within the rolling window.

From what I see, the extraction goes correctly for velocity vectors, but not for the coordinates. When - in my custom function - I attempt to see the coordinates of every point within my rolling window, I see an array of numbers both for x and y coordinates. For a 9x9 window this array looks like [0, 1, 2, 3, 4, 5, 6, 7, 8].

I think, that when the rolling window is created, it stores only the data variables and assigns new values to the coordinates within the rolling window.

My question:
I'm wondering, if I'm missing something, and there is a way to obtain the (x,y) coordinates of the original velocity field at every point within the rolling window.

Here's my custom function:

def ΓmovingWindowFunction(fWin, n): 
      """
      Parameters: fWin (xarray.Dataset) - a moving window of the dataset (fWin = field rolling window)
                         n (int) - the rolling window size (n=1 means a 3x3 rolling window)
      Returns: Γ1 (float) - the calculated value of the Γ1 function.
    """
    # n always points to the center of the rolling window. It gives me point P
     fWin['PMx'] = fWin['rollWx'] - fWin.coords['rollWx'][n] # x distance to every point from point P
     fWin['PMy'] = fWin['rollWy'] - fWin.coords['rollWy'][n] # y distance to every point from point P
     fWin['PM'] = (fWin['PMx']**2 + fWin['PMy']**2)**(0.5) # distance to every point from point P
     fWin['U'] = (fWin['u']**2 + fWin['v']**2)**(0.5) # total velocity at every point
     fWin['sinΘM_Γ1'] = (fWin['PMx']*fWin['v'] - fWin['PMy']*fWin['u']) / fWin['PM'] / fWin['U']

     Γ1 = fWin['sinΘM_Γ1'].sum() / (((2*n+1)**2)-1)

     return Γ1

Important note: see how I use Xarray in my calculations. It is incredibly slow! I had to rewrite my function to make use of numpy arrays. Which added a lot of additional pain. Xarray must be fast on its own, otherwise it induces great headache when I try to convert it to numpy. If I wanted to use numpy for my datasets I wouldn't adhere to Xarray.

But let's get back to my issue.

See, for instance, the line: fWin['PMx'] = fWin['rollWx'] - fWin.coords['rollWx'][n]. Here, rollWx refers to the x dimension of the rolling window. I expect it to give me the real coordinates (x,y) from the real velocity field. Instead, when I test it, I get arranged numbers from 0 to 8 for my 9x9 rolling window.

Here's how I apply my custom function to my dataset called field (here, (2*n+1) gives me the size of the rolling window):

    def Γpad(ds, n): 
          dsGroup = ds.groupby("gridcell")
          return dsGroup.map(ΓmovingWindowFunction, args=[n])

    rollingW = field.rolling({"x":(2*n+1), "y":(2*n+1), "t":1}, center=True)
    fieldRoll = rollingW.construct(x='rollWx', y='rollWy', t='rollWt') # I also have time dimension: velocity field at every time moment
    fieldStacked = fieldRoll.stack(gridcell=['x','y','t'])
    newArr = fieldStacked.map_blocks(Γpad, args=[n]).compute()  

Here's my dataset:

field:
 <xarray.Dataset>
Dimensions:  (x: 79, y: 63, t: 2)
Coordinates:
  * x        (x) float64 0.9412 1.882 2.824 3.765 ... 71.53 72.47 73.41 74.35
  * y        (y) float64 59.29 58.35 57.41 56.47 ... 3.765 2.824 1.882 0.9412
  * t        (t) float64 0.0 0.001
Data variables:
    u        (x, y, t) float64 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0
    v        (x, y, t) float64 -0.0 -0.0 -0.0 -0.0 -0.0 ... -0.0 -0.0 -0.0 -0.0
    chc      (x, y, t) float64 1.0 1.0 1.0 1.0 1.0 1.0 ... 1.0 1.0 1.0 1.0 1.0
    w        (x, y, t) float64 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0
    Γ1       (x, y, t) float64 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0
    uConv    (x, y, t) float64 nan nan nan nan nan nan ... nan nan nan nan nan
    vConv    (x, y, t) float64 nan nan nan nan nan nan ... nan nan nan nan nan
    UConv    (x, y, t) float64 nan nan nan nan nan nan ... nan nan nan nan nan
    U        (x, y, t) float64 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0
Attributes: (2)

Thank you in advance.
Ivan