How to set the coordinates of the output of xarray.assign?

user_3696204 0
Viewed 564

I've been trying to create two new variables based on the latitude coordinate of a data point in an xarray dataset. However, I can only seem to assign new coordinates. The data set looks like this:

<xarray.Dataset>
Dimensions:  (lon: 360, lat: 180, time: 412)
Coordinates:
  * lon      (lon) float64 0.5 1.5 2.5 3.5 4.5 ... 355.5 356.5 357.5 358.5 359.5
  * lat      (lat) float64 -89.5 -88.5 -87.5 -86.5 -85.5 ... 86.5 87.5 88.5 89.5
  * time     (time) datetime64[ns] 1981-09-01 1981-10-01 ... 2015-12-01
Data variables:
    evapr    (time, lat, lon) float32 ...
    lhtfl    (time, lat, lon) float32 ...
...

What I've attempted so far is this:

def get_latitude_band(latitude):
    return np.select(
        condlist=
        [abs(latitude) < 23.45,
         abs(latitude) < 35,
         abs(latitude) < 66.55],
        choicelist=
        ["tropical",
         "sub_tropical",
         "temperate"],
        
        default="frigid"
    )

def get_hemisphere(latitude):
    return np.select(
        [latitude > 0, latitude <=0],
        ["north", "south"]
    )

    
mhw_data = mhw_data \
    .assign(climate_zone=get_latitude_band(mhw_data.lat)) \
    .assign(hemisphere=get_hemisphere(mhw_data.lat)) \
    .reset_index(["hemisphere", "climate_zone"]) \
    .reset_coords()
            
print(mhw_data)

Which is getting me close:

<xarray.Dataset>
Dimensions:        (lon: 360, lat: 180, time: 412, hemisphere: 180, climate_zone: 180)
Coordinates:
  * lon            (lon) float64 0.5 1.5 2.5 3.5 4.5 ... 356.5 357.5 358.5 359.5
  * lat            (lat) float64 -89.5 -88.5 -87.5 -86.5 ... 86.5 87.5 88.5 89.5
  * time           (time) datetime64[ns] 1981-09-01 1981-10-01 ... 2015-12-01
Dimensions without coordinates: hemisphere, climate_zone
Data variables:
    evapr          (time, lat, lon) float32 ...
    lhtfl          (time, lat, lon) float32 ...
    ...
    hemisphere_    (hemisphere) object 'south' 'south' ... 'north' 'north'
    climate_zone_  (climate_zone) object 'frigid' 'frigid' ... 'frigid' 'frigid'
...

However, I want to then stack the DataSet and convert it to a DataFrame. I am unable to do so, and I think it is because the new variables hemisphere_ and climate_zone_ do not have time, lat, lon coordinates:

stacked = mhw_data[var].stack(dim=["lon", "lat", "time"]).to_pandas().T

results in a KeyError on "lon".

So my question is: How do I assign new variables to the xarray dataset that maintain the original coordinates of time, lat and lon?

python-xarray pandas python
2 Answers

To assign a new variable or coordinate, xarray needs to know what the dimensions are called. There are a number of ways to define a DataArray or Coordinate, but the one closest to what you're currently using is to provide a tuple of (dim_names, array):

mhw_data = mhw_data.assign_coords(
    climate_zone=(('lat', ), get_latitude_band(mhw_data.lat)),
    hemisphere=(('lat', ), get_hemisphere(mhw_data.lat)),
)

Here I'm using da.assign_coords, which will define climate_zone and hemisphere as non-dimension coordinates, which you can think of as additional metadata about latitude and about your data, but which aren't proper data in themselves. This will also allow them to be preserved when sending individual arrays to pandas.

For stacking, converting to pandas will stack automatically. The following will return a DataFrame with variables/non-dimension coordinates as columns and dimensions as a MultiIndex:

stacked = mhw_data.to_dataframe()

Alternatively, if you want a Series indexed by (lat, lon, time) for just one of these coordinates you can always use expand_dims:

(
    mhw_data.climate_zone
    .expand_dims(lon=mhw_data.lon, time=mhw_data.time)
    .to_series()
)
user_38887190

The two possible solutions I've worked out for myself are as follow:

  1. First, stack the xarray data into pandas DataFrames, and then create new columns:
df = None
variables = list(mhw_data.data_vars)

for var in tqdm(variables): 
    
    stacked = mhw_data[var].stack(dim=["lon", "lat", "time"]).to_pandas().T
    if df is None:
        df = stacked
    else:
        df = pd.concat([df, stacked], axis=1)

df.reset_index(inplace=True)
df.columns = list(mhw_data.variables)

df["climate_zone"] = df["lat"].swifter.apply(get_latitude_band)
df["hemisphere"] = df["lat"].swifter.apply(get_hemisphere)
  1. Create new xarray.DataArrays for each variable you'd like to add, then add them to the dataset:
# calculate climate zone and hemisphere from latitude. 
latitudes = mhw_data.lat.values.reshape(-1, 1)

zones = get_latitude_band(latitudes)
hemispheres = get_hemisphere(latitudes)

# Take advantage of numpy broadcasting to get our data to lign up with the xarray shape. 
shape = tuple(mhw_data.sizes.values())
base = np.zeros(shape)

zones = zones + base
hemispheres = hemispheres + base

# finally, create two new DataArrays and assign them as variables in the dataset. 
zone_xarray = xr.DataArray(data=zones, coords=mhw_data.coords, dims=mhw_data.dims)
hemi_xarray = xr.DataArray(data=hemispheres, coords=mhw_data.coords, dims=mhw_data.dims)

mhw_data["zone"] = zone_xarray
mhw_data["hemisphere"] = hemi_xarray

# ... call the code to stack and convert to pandas (shown in method 1) ...#

My intuition is that method 1 is faster and more memory efficient because there are no repeated values that need broadcasting into a large 3-dimensional array. I did not test this, however.

Also, my intuition is that there is a less cumbersome xarray native way of accomplishing the same goal, but I could not find it.

One thing is certain, method 1 is far more concise due to the fact that there is no need to create intermediate arrays or reshape data.

user_36962040
Related