Concise way to filter data in xarray

Question

I need to apply a very simple 'match statement' to the values in an xarray array:

Where the value > 0, make 2
Where the value == 0, make 0
Where the value is NaN, make NaN

Here's my current solution. I'm using NaNs, .fillna, & type coercion in lieu of 2d indexing.

valid = date_by_items.notnull()
positive = date_by_items > 0
positive = positive * 2
result = positive.fillna(0.).where(valid)
result

This changes this:

In [20]: date_by_items = xr.DataArray(np.asarray((list(range(3)) * 10)).reshape(6,5), dims=('date','item'))
    ...: date_by_items
    ...: 
Out[20]: 
<xarray.DataArray (date: 6, item: 5)>
array([[0, 1, 2, 0, 1],
       [2, 0, 1, 2, 0],
       [1, 2, 0, 1, 2],
       [0, 1, 2, 0, 1],
       [2, 0, 1, 2, 0],
       [1, 2, 0, 1, 2]])
Coordinates:
  * date     (date) int64 0 1 2 3 4 5
  * item     (item) int64 0 1 2 3 4

... to this:

Out[22]: 
<xarray.DataArray (date: 6, item: 5)>
array([[ 0.,  2.,  2.,  0.,  2.],
       [ 2.,  0.,  2.,  2.,  0.],
       [ 2.,  2.,  0.,  2.,  2.],
       [ 0.,  2.,  2.,  0.,  2.],
       [ 2.,  0.,  2.,  2.,  0.],
       [ 2.,  2.,  0.,  2.,  2.]])
Coordinates:
  * date     (date) int64 0 1 2 3 4 5
  * item     (item) int64 0 1 2 3 4

While in pandas df[df>0] = 2 would be enough. Surely I'm doing something pedestrian and there's an terser way?

Maximilian · Accepted Answer · 2018-03-23 22:58:15Z

17

xarray now supports .where(condition, other), so this is now valid:

result = date_by_items.where(date_by_items > 0, 2)

edited Mar 23, 2018 at 22:58

answered Mar 23, 2018 at 5:18

Maximilian

8,7403 gold badges57 silver badges73 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

shoyer · Accepted Answer · 2016-08-10 22:30:16Z

5

If you are happy to load your data in-memory as a NumPy array, you can modify the DataArray values in place with NumPy:

date_by_items.values[date_by_items.values > 0] = 2

The cleanest way to handle this would be if xarray supported the other argument to where, but we haven't implemented that yet (hopefully soon -- the groundwork has been laid!). When that works, you'll be able to write date_by_items.where(date_by_items > 0, 2).

Either way, you'll need to do this twice to apply both your criteria.

answered Aug 10, 2016 at 22:30

shoyer

9,6831 gold badge41 silver badges57 bronze badges

Comments

Dmitry Deryabin · Accepted Answer · 2023-08-03 15:09:25Z

1

Another concise way would be to do date_by_items.values[date_by_items.values > 0] = 2

answered Aug 3, 2023 at 15:09

Dmitry Deryabin

1,5882 gold badges14 silver badges27 bronze badges

Comments

Stijn · Accepted Answer · 2022-06-23 16:07:30Z

You can use the where(condition, other) method indeed. But be aware that the other argument will be used where the condition is false. So the behavior in the other answers is incorrect, as they will put a 2 where date_by_items > 0 does not hold.

>>> date = list(range(0,6))
>>> item = list(range(0,5))
>>> date_by_items = xr.DataArray(np.asarray((list(range(3)) * 10)).reshape(6,5), coords=[date, item], dims=('date','item'))
>>> date_by_items
<xarray.DataArray (date: 6, item: 5)>
array([[0, 1, 2, 0, 1],
       [2, 0, 1, 2, 0],
       [1, 2, 0, 1, 2],
       [0, 1, 2, 0, 1],
       [2, 0, 1, 2, 0],
       [1, 2, 0, 1, 2]])
Coordinates:
  * date     (date) int64 0 1 2 3 4 5
  * item     (item) int64 0 1 2 3 4


>>> date_by_items.where(date_by_items > 0, 2)  # wrong behavior
<xarray.DataArray (date: 6, item: 5)>
array([[2, 1, 2, 2, 1],
       [2, 2, 1, 2, 2],
       [1, 2, 2, 1, 2],
       [2, 1, 2, 2, 1],
       [2, 2, 1, 2, 2],
       [1, 2, 2, 1, 2]])
Coordinates:
  * date     (date) int64 0 1 2 3 4 5
  * item     (item) int64 0 1 2 3 4

Instead, when you want the requested behavior, you either have to invert the condition or use the xarray.where(condition, x, y) method instead.

>>> date_by_items.where(date_by_items <= 0, 2)  # inverted condition
<xarray.DataArray (date: 6, item: 5)>
array([[0, 2, 2, 0, 2],
       [2, 0, 2, 2, 0],
       [2, 2, 0, 2, 2],
       [0, 2, 2, 0, 2],
       [2, 0, 2, 2, 0],
       [2, 2, 0, 2, 2]])
Coordinates:
  * date     (date) int64 0 1 2 3 4 5
  * item     (item) int64 0 1 2 3 4

>>> xarray.where(date_by_items > 0, 2, date_by_items)
<xarray.DataArray (date: 6, item: 5)>
array([[0, 2, 2, 0, 2],
       [2, 0, 2, 2, 0],
       [2, 2, 0, 2, 2],
       [0, 2, 2, 0, 2],
       [2, 0, 2, 2, 0],
       [2, 2, 0, 2, 2]])
Coordinates:
  * date     (date) int64 0 1 2 3 4 5
  * item     (item) int64 0 1 2 3 4

Collectives™ on Stack Overflow

Concise way to filter data in xarray

4 Answers 4

Comments

Comments

Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

4 Answers 4

Comments

Comments

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related