Plot slice through 3-D xarray.Dataset

Question

I would like to plot a slice through a 3-D xarray.Dataset as illustrated below:

This is how I started off (for simplicity, using the tutorial air_temperature dataset)

# import packages
import xarray as xr
import numpy as np

# load dataset
ds = xr.tutorial.open_dataset("air_temperature")

# get slice values
tgt_lon = xr.DataArray(np.linspace(220, 280, num=15), dims="lon")
tgt_lat = xr.DataArray(np.linspace(30, 50, num=15), dims="lat")

# crop to region of interest - this works fine
da = ds.sel(lon=tgt_lon, 
            lat=tgt_lat, 
            method="nearest")

Now we still have a 3-D xarray.Dataset. For a 2-D plot, we want to stack longitude and latitude to a distance array relative to a zero point (here: corner x0/y0 as shown in figure).

# zero point: lat_min/lon_min
lon_orig = da.lon.min().values
lat_orig = da.lat.min().values

# stack longitude and latitude -- gives tuple for dist
sta = da.stack(dist=('lon', 'lat')) 

... and then we could compute the distance relative to lon_orig/lat_orig by looping over sta. For the actual plotting (using contourf), we would then flatten the arrays to 1-D arrays. This all seems a bit tedious and I was wondering if I am missing the obvious way of doing this?

Ultimately, we want three 1-D arrays with the same shape for plotting:

• z (for this dataset: time)
• distance (relative to x0/y0, here computed from longitude and latitude)
• the actual data (here: air temperature)

Thanks!

Answer 1

Setting up my Example

First of all I am going to use a different example since I want to know what I am supposed to get as a result. My example is a 100x100x100 array containing a ball of radius 50.

I have adopted the code from this answer to plot it.

n = 100

def distance_from_center(x,y,z):
    V = np.stack([x-n/2,y-n/2,z-n/2])
    return np.linalg.norm(V, axis=0)

ball = np.fromfunction(distance_from_center, (n, n, n), dtype='float')
ball = (ball > 50).astype('int')

Elementary Solution

You can do such a thing in 4 lines of code. So I will do that first before I get more fancy. My first slice is from corner to corner. Like this:

I have marked where I will slice in green.

t = np.arange(n)
slice = ball[t,t,:]
plt.matshow(slice.T, cmap='gray')
plt.gca().set_aspect(1/2**(1/2))

Notice here that I have to set the aspect ratio to one over square root of two since the 50 pixels in the diagonal direction are a longer distance than the 50 in the ordinary direction. The picture looks like expected. There is white space on both sides since the corners of the array provide space for it but we still get a circle.

You can do other cuts like this but you're always thinking about how to get whole number coordinates.

More Sophisticated Solution

The alternative is to just take any coordinates and interpolate between points where you don't hit points perfectly.

import itertools

class Image_knn():
    def fit(self, image):
        self.image = image.astype('float')

    def predict(self, x, y, z):
        image = self.image
        weights_x = [(1-(x % 1)).reshape(-1), (x % 1).reshape(-1)]
        weights_y = [(1-(y % 1)).reshape(-1), (y % 1).reshape(-1)]
        weights_z = [(1-(z % 1)).reshape(-1), (z % 1).reshape(-1)]
        start_x = np.floor(x)
        start_y = np.floor(y)
        start_z = np.floor(z)

        return sum([image[np.clip(np.floor(start_x + x), 0, image.shape[0]-1).astype('int'),
                          np.clip(np.floor(start_y + y), 0, image.shape[1]-1).astype('int'),
                          np.clip(np.floor(start_z + z), 0, image.shape[1]-1).astype('int')]
                    *weights_x[x]*weights_y[y]*weights_z[z]
                    for x,y,z in itertools.product(range(2),range(2),range(2))])

    
image_model = Image_knn()
image_model.fit(ball)

fig, ax = plt.subplots(nrows=3,ncols=3)
ax = ax.reshape(-1)

I am going to give an example using different slices where one direction is the z direction and the other is the one containing the bottom left corner and varying points on the bottom right. First the 3d plot indicating where the cuts will be:

Now I calculate the coordinates for the cuts and plot them. This time I can keep the aspect ratio since I use more points for the direction that is longer. Notice that I put the coordinates of the 2nd point in the heading of the plot.

fig, ax = plt.subplots(nrows=3,ncols=3)
ax = ax.reshape(-1)

for i,x in enumerate(np.linspace(0,100,9)):
    p = np.array([0,100])
    q = np.array([100,100-x])
    distance = np.round(np.linalg.norm(q-p)).astype('int')
    t = np.linspace(0,1,distance)
    xy = t.reshape(-1,1)*q+(1-t).reshape(-1,1)*p
    r,s = np.meshgrid(np.arange(100), np.arange(distance))
    x = xy[s][:,:,0].reshape(-1)
    y = xy[s][:,:,1].reshape(-1)
    z = r.reshape(-1)

    out = image_model.predict(x,y,z)
    ax[i].imshow(out.reshape(distance, 100).T, cmap='gray',vmin=0,vmax=1)
    ax[i].set_title(tuple(q.astype('int')))