Perform a calculation which takes a coordinate as an input in xarray non-iteratively

Question

I have an .nc file which contains meteorological data such as temperature on a time-lat-lon-level coordinate grid. From temperature, I'd like to calculate potential temperature using Poisson's equation

In order for that to work, I need to pull lev from my dataset and plug it into P along with the temperature from the same dataset. Using the for loop construction that came to mind,

merraLT['THETA'] = xr.zeros_like(merraLT.T) # creates empty variable field to place calculated potential temperature in
 for i in range(len(merraLT.lev)):
     merraLT['THETA'][:,i,:,:] = merraLT.T[:,i,:,:]*((1000/merraLT.lev[i])**(287/1004))

it takes half an hour to generate the potential temperature fields at each level. Is there a better way to accomplish such a thing?

edit: There is a one-line solution, but it still takes half an hour. Probably just because of the size of the data (reducing to 12 pressure levels from 42(!) brought it down to 5 minutes), but if anyone knows a better solution, please share.

merraLT['THETA'] = xr.zeros_like(merraLT.T) # creates empty variable field to place calculated potential temperature in
merraLT['THETA'] = merraLT.T*(1000/merraLT.lev)**(287/1004)

Answer 1

There is no need to use xr.zeros_like first, your one-line solution is the way to express this calculation using xarray.

As for why it takes 5 minutes, we can't possibly debug that without a lot more information about your data, how large it is, and how you are accessing it (is it coming over a network? Are you loading the entire contents of the file into memory? Is it the loading step or the calculation step that is taking all the time?).