Stream: python-questions

Topic: using a variable to index a dimension in xarray

Isla Simpson (Nov 23 2021 at 02:00):

Hello, Can anyone help me with this? What I'd like to do is have a flexible subroutine that's doing something with an xarray data array along a dimension specified as input to this subroutine i.e., as a simple example would be something like...

def function(x, dim):
    y = x.isel(dim=slice(0,10))
    return y

Suppose x has dimensions [model, lat, lon] and I want to get the slice along the model dimension I'd call this using...

y = function(x, 'model')

But this doesn't work because it thinks I want to do the slicing along dimension "dim" as opposed to using the string that I've passed through the be variable "dim". Is there a way around this? Thanks in advance for any help!

Ryan May (Nov 23 2021 at 07:16):

What you want is Python's support for keyword argument unpacking from a dictionary using **:

def  function(x, dim):
    kwargs = {dim: slice(0, 10)}
    y = x.sel(**kwargs)
    return y

Isla Simpson (Nov 23 2021 at 14:52):

That's excellent! Thanks a lot!

Deepak Cherian (Nov 23 2021 at 17:58):

sel and isel also accept dictionaries so even y = x.sel(kwargs) will work

Isla Simpson (Nov 23 2021 at 18:03):

good to know. Thanks!

Last updated: Jan 30 2022 at 12:01 UTC