CODEWITHSUNDEEP
pythonmatplotlibpython-xarray
smillerc
smillerc

Reputation: 168

Plotting xarray datasets with variable coordinates

I'm trying to use xarray to plot data on a variable grid. The grid that my data is stored on is changing over time, but keeps the same dimensions.

I'd like to be able to plot 1d slices of it at a given time. A toy example of what I'm trying to do is shown below.

import xarray as xr
import numpy as np
import matplotlib.pyplot as plt

time = [0.1, 0.2] # i.e. time in seconds

# a 1d grid changes over time, but keeps the same dims
radius = np.array([np.arange(3),
                   np.arange(3)*1.2])

velocity = np.sin(radius) # make some random velocity field

ds = xr.Dataset({'velocity': (['time', 'radius'],  velocity)},
            coords={'r': (['time','radius'], radius), 
                    'time': time})

If I try to plot it at different times, i.e.

ds.sel(time=0.1)['velocity'].plot()
ds.sel(time=0.2)['velocity'].plot()
plt.show()

xarray plot version

But I'd like it to replicate the behavior that I can do explicitly using matplotlib. Here it properly plots the velocity against the radius at that time.

plt.plot(radius[0], velocity[0])
plt.plot(radius[1], velocity[1])
plt.show()

proper plot version

I may be using xarray wrong, but it should be plotting the velocity against the proper value of radius at that time.

Am I setting up the Dataset wrong or using the plot/index feature wrong?

Upvotes: 3

Views: 2797

Answers (1)

jhamman
jhamman

Reputation: 6434

I agree this behavior is unexpected but its not entirely a bug.

Looking at the variables you're trying to plot:

da = ds.sel(time=0.2)['velocity']
print(da)

yields:

<xarray.DataArray 'velocity' (radius: 3)>
array([ 0.      ,  0.932039,  0.675463])
Coordinates:
    r        (radius) float64 0.0 1.2 2.4
    time     float64 0.2
Dimensions without coordinates: radius

What we see is that there is not a coordinate variable named radius which is what xarray is looking for when making its x coordinate for the plots you've shown above. In your case, you need a work simple around where we rename the 1-D coordinate variable the same name as the dimension:

for time in [0.1, 0.2]:
    ds.sel(time=time)['velocity'].rename({'r': 'radius'}).plot(label=time)

plt.legend()
plt.title('example for SO')

enter image description here

Upvotes: 1

Related Questions