Stream: python-questions

Topic: Plotting POP bottom velocity using xarray

Mira Berdahl (Mar 21 2024 at 18:45):

Hi,
I'm trying to use KMT to plot the bottom velocity field of some POP output. In the past I've been just extracting a specific level with xarray (mfds.UVEL.isel(z_t=0), for example) but that won't work because KMT is a 2D field and z_t is 1D. I presume this is simple, but can anyone provide an example or guidance on this?
Thanks in advance.

Michael Levy (Mar 22 2024 at 14:25):

@Kristen Krumhardt or @Lev Romashkov might have a better way of doing this, but for temperature I think you want something like

TEMP_BOTTOM = xr.full_like(ds['TEMP'].isel(z_t=0), np.nan)
for k in range(60):
    TEMP_BOTTOM.data = np.where(ds['KMT'].data == k+1, ds['TEMP'].isel(z_t=k).data, TEMP_BOTTOM.data)

and for UVEL you would use

UVEL_BOTTOM = xr.full_like(ds['UVEL'].isel(z_t=0), np.nan)
for k in range(60):
    UVEL_BOTTOM.data = np.where(ds['KMU'].data == k+1, ds['UVEL'].isel(z_t=k).data, UVEL_BOTTOM.data)

Note the subtle difference between using KMT in the logical expression ds['KMT'].data == k+1 for temperature and KMU in ds['KMU'].data == k+1 for UVEL -- the tracers are computed at cell centers, while velocities are computed at cell vertices, and the bathymetry is slightly different between the two.

Kristen Krumhardt (Mar 22 2024 at 14:31):

Here's another way, but sometimes it takes forever or gets dask errors, depending on how big your dataset is (this example is for variables TEMP and POC_FLUX_IN):

def field_at_bottom(da):
    """return a field indexed at the model's bottom layer"""

    tmp_bot = xr.DataArray(np.ones(da[:, 0, :, :].shape) * np.nan,
                           dims=tuple(da.dims[i] for i in [0, 2, 3]),
                           coords={c: da.coords[c] for c in ['time']},
                          )

    assert KMT.shape == da.shape[-2:]

    for j in range(len(da.nlat)):
        for i in range(len(da.nlon)):
            if KMT[j, i] > 0:
                k = int(KMT[j, i] - 1)
                tmp_bot.values[:, j, i] = da[:, k, j, i]
    return tmp_bot

for v in ['TEMP', 'POC_FLUX_IN',]:
    template = dsets[v][v][:, 0, :, :].drop('z_t')
    dso[f'{v}_bottom'] = xr.map_blocks(
        field_at_bottom, dsets[v][v],
        template=template
    ).compute()

Mira Berdahl (Mar 22 2024 at 22:33):

@Michael Levy and @Kristen Krumhardt Thanks for the help here! These work.

Fred Castruccio (Apr 01 2024 at 21:33):

Mira,
Here is a third option.
You can simply use isel like you did to get the surface velocity (z_t=0) except that you are going to use KMU to select the bottom cells.
See example below.

uvel = ds.UVEL
kmu = ds.KMU.astype('int') - 1
Ubot = UVEL.isel(z_t=kmu)

Last updated: May 16 2025 at 17:14 UTC