How to change certain values in an xarray depending on the coordinates

Question

I have a Data Array with 4, 5 or 6 representing the classes: bare soil, vegetation and water.

cube1 =

<xarray.DataArray 'SCL_20m' (time: 5, y: 3, x: 3)>
array([[[4., 6., 6.],
        [4., 4., 6.],
        [4., 4., 6.]],

       [[4., 6., 6.],
        [4., 4., 6.],
        [4., 4., 6.]],

       [[4., 6., 6.],
        [6., 6., 6.],
        [6., 6., 6.]],

       [[4., 6., 6.],
        [4., 4., 6.],
        [4., 4., 6.]],

       [[4., 6., 6.],
        [4., 4., 6.],
        [4., 4., 6.]]], dtype=float32)
Coordinates:
  * time         (time) datetime64[ns] 2020-06-07T10:12:20 ... 2020-06-14T10:...
  * y            (y) float64 6.607e+06 6.607e+06 6.607e+06
  * x            (x) float64 7.091e+05 7.092e+05 7.092e+05
    spatial_ref  int32 -2147483647

I am applying a particle filter pixel by pixel and I want to save a new xarray with the same coordinates as the original but with the values obtained after filtering. This is my idea:

xvalues = cube1.x.values
yvalues = cube1.y.values
filtered_cube = xr.ones_like(cube1)#copy the cube and then update time series after filtering
for xx in xvalues:
    for yy in yvalues:
        pixel_ts = cube1.sel(y=yy, x=xx, method="nearest")
        # Here I apply filtering to the pixel time series(pixel_ts) but for now let's mock the output array
        if yy==yvalues[0]:
            mock_array = np.array([0,0,0,0,0])
            # replace filtered_cube with mock_array when the coordinate y corresponds to yvalues[0]
            # HOW? 

I have tried to reproduce this example: https://stackoverflow.com/questions/49562588/how-can-i-replace-values-in-an-xarray-variable

But it didn't work, I think because it's done with a Dataset and I have a DataArray. Can anyone help me?

Answer 1

If I understand correctly, you want to replace the values for bare soil, vegetation and water everywhere where yyis the value of y.

If so, your mock_array might have the shape (5,3), for each of the 5 timesteps and 3 bare soil/vegetation/water values.

mock_array = np.arange(15).reshape((5,3))

mock_array

array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14]])

Values in cube1 can be updated by indexing by the yy position. If the index of yy were 0, then

cube1[:,0,:] = mock_array

cube1

<xarray.DataArray (time: 5, y: 3, x: 3)>
array([[[ 0.,  1.,  2.],
        [ 4.,  4.,  6.],
        [ 4.,  4.,  6.]],

       [[ 3.,  4.,  5.],
        [ 4.,  4.,  6.],
        [ 4.,  4.,  6.]],

       [[ 6.,  7.,  8.],
        [ 6.,  6.,  6.],
        [ 6.,  6.,  6.]],

       [[ 9., 10., 11.],
        [ 4.,  4.,  6.],
        [ 4.,  4.,  6.]],

       [[12., 13., 14.],
        [ 4.,  4.,  6.],
        [ 4.,  4.,  6.]]])
Coordinates:
  * y            (y) float64 6.607e+06 6.608e+06 6.609e+06
  * x            (x) float64 7.091e+05 7.092e+05 7.093e+05
  * time         (time) datetime64[ns] 2014-09-06 2014-09-07 ... 2014-09-10
    spatial_ref  int64 -2147483647

I don't have a more elegant suggestion to finding the positional index than

cube1.y.values.tolist().index(yy)

This will give you the index of yy. You can also put the whole statement straight into the previous one:

cube1[:,cube1.y.values.tolist().index(6607000),:] = mock_array

Answer 2

This may not answer the specific question posed here, but the beauty of xarray is that it provides some nice ways to avoid some of these pixel-by-pixel looping approaches.

For example, it may be easier in the long run to try and apply your particle filter function directly to each pixel in your data using xr.apply_ufunc:

import xarray as xr
import numpy as np
import pandas as pd

# Simulate example data
cube1 = xr.DataArray(data=np.array([[[4., 6., 6.], [4., 4., 6.], [4., 4., 6.]],
                                    [[4., 6., 6.], [4., 4., 6.], [4., 4., 6.]],
                                    [[4., 6., 6.], [6., 6., 6.], [6., 6., 6.]],
                                    [[4., 6., 6.], [4., 4., 6.], [4., 4., 6.]],
                                    [[4., 6., 6.], [4., 4., 6.], [4., 4., 6.]]]),
                     coords=[pd.date_range('2020-06-07', '2020-06-11'),
                             [6.607e+06, 6.607e+06, 6.607e+06],
                             [7.091e+05, 7.092e+05, 7.092e+05]],
                     dims=['time', 'y', 'x'])


# Create custom func that takes an array of observations through time,
# and returns a transformed array of equal size
def custom_func(x):
    return x + 1

# Apply custom function along the time dimension
cube2 = xr.apply_ufunc(custom_func,
                       cube1,
                       input_core_dims=[["time"]],
                       output_core_dims=[["time"]],
                       vectorize=True)
cube2.transpose('time', 'y', 'x')

This returns an array of the same shape as your original data, but the original pixel values have been replaced with the outputs of your custom_func (e.g. n + 1 in this example, but filtered values in your use case):