Xarray apply function to every element of dataset

Question

I currently have to do some calculations on a netcdf dataset. For this, I have to apply a function to each non-NaN element.
Here is my current approach:

import xarray as xr
def calc_things(wind_speed):
 if np.isnan(wind_speed):
 return np.nan
 if wind_speed<3:
 return 0
 if wind_speed>11.3:
 return 100
 
 # do math stuff here
if __name__ == "__main__":
 with xr.open_dataset("input.nc") as df:
 wind_data = df['sfcWind']
 
 wind_applied = xr.apply_ufunc(calc_things, wind_data, vectorize=True)
 
 ds.to_netcdf("modified.nc")

Can somebody tell me if this is the right way to do it? I am also wondering if there is a better way to handle the NaN values instead of having an if-statement in the function.
All help is appreciated.

• Divyansh Gupta: If this snippet gives expected results, can you elaborate what you mean by the 'right' way to do things?
• ATYslh: Just because something works doesn't mean that it is the best way to do it. For example you can append to a numpy array, even though it is not encouraged due to performance.
• chrslg: Maybe you should try code review. On Stack Overflow we ask specific question on specific problem. Do you have a performance problem for example, sine you mentioned it?