Faster way to assign individual value to a xarray dataset

user_11378446 0
Viewed 300

I realize that assigning individual values to a xarray dataset takes much more time than doing the same with a numpy array. Would there be a way to accelerate that ?

Here is everything that I tested by alternatively uncommenting each line in the loop:

import numpy as np
import xarray as xr
import datetime

levels = np.arange(0,3)
simNames = ['9airports_filter0dot7_v22']
airportList = ['Windhoek', 'Atlanta', 'Taipei']
numb_variables = 11

emptyDA = xr.DataArray(np.nan, 
                       coords = [simNames, airportList, np.arange(0, 20428), levels], 
                       dims   = ['simName', 'airport', 'profnum'           , 'level'])

ds = xr.Dataset({ 'iasi': emptyDA.copy(), 'IM':   emptyDA.copy(), 'IMS': emptyDA.copy(), 'err': emptyDA.copy(), 
                 'sigma': emptyDA.copy(), 'temp': emptyDA.copy(), 'dfs': emptyDA.copy(), 'ocf': emptyDA.copy(), 
                 'rcf':   emptyDA.copy(), 'time': emptyDA.copy().astype(dtype="datetime64[ns]"), 'surfPres': emptyDA.copy() })

mat3D = np.empty( shape=( len(airportList), 20428, len(levels) ) ) # 20428 is needed for the 9 airports
mat3D[:] = np.nan

mat4D = np.empty( shape=( 1, len(airportList), 20428, len(levels) ) ) # 20428 is needed for the 9 airports
mat4D[:] = np.nan

mat5D = np.empty( shape=( numb_variables, 1, len(airportList), 20428, len(levels) ) ) # 20428 is needed for the 9 airports
mat5D[:] = np.nan

begin_time = datetime.datetime.now()

for i in range(10000):
    ds['iasi'][0, 0, 0, 0] = 3.1416                                           # 1.08 sec
    # ds['iasi'].loc['9airports_filter0dot7_v22', 'Windhoek', 0, 0] = 3.1416    # 1.97 sec
    # ds['iasi'][0, 0, 0, 0].data = 3.1416                                      # 0.85 sec
    # ds['iasi'][0, 0, 0, 0].values = 3.1416                                    # 0.85 sec
    # ds.iasi[0, 0, 0, 0].values = 3.1416                                       # 0.88 sec
    # a = 3.1416                                                                # 0.0003
    # mat4D[0, 0, 0, 0] = 3.1416                                                # 0.0008 sec
    # mat3D[0, 0, 0] = 3.1416                                                   # 0.0008 sec
    # mat5D[0, 0, 0, 0, 0] = 3.1416                                             # 0.0009 sec

print(datetime.datetime.now() - begin_time)
python-xarray python
1 Answers

If you're going to be referencing the data using positional indices, you could definitely do the indexing within numpy, e.g. ds['iasi'].data[0, 0, 0, 0] = 3.1416. But the general assumption you should have is that xarray will have worse performance compared with numpy, as every operation you do in xarray is a bunch of python overhead, plus any pandas-based indexing you're doing, on top of the corresponding numpy operation.

That said, xarray works hard to add only a modest penalty on top of the underlying array operations. Importantly, the overhead will be felt most acutely when performing small operations. The larger the array, the more xarray's performance will approach numpy's, as the bulk of the work will be on the array itself, not in cueing it up with indexing, metadata management, and other elements that xarray adds on top.

To illustrate this, I've set up a side-by side comparison where we first initialize a (very large) empty numpy array, then progressively fill the first 1, then 10, then 100, then 1000, etc. elements until we've filled the entire (1000, 1000, 100) array with zeros:

In [4]: times = []
   ...: for sim in range(40):
   ...:     for i in range(1, 9):
   ...:         arr = np.empty(shape=(1000, 1000, 100), dtype='float64')
   ...:         # progressively fill in 10^i cells at each iteration
   ...:         start = time.time()
   ...:         arr.flat[:int(10**i)] = 0
   ...:         end = time.time()
   ...:         times.append([i, end - start])

Compare this to the analagous operation in xarray:

In [5]: times_xr = []
   ...: for sim in range(40):
   ...:     for i in range(1, 9):
   ...:         da = xr.DataArray(np.empty(shape=(1000, 1000, 100), dtype='float64'), dims=['x', 'y', 'z'])
   ...:         # progressively fill in 10^i cells at each iteration
   ...:         start = time.time()
   ...:         da.data.flat[:int(10**i)] = 0
   ...:         end = time.time()
   ...:         times_xr.append([i, end - start])

I only timed the data population step to be comparable to the case in your question, but note how (a) the time to populate the array scales far less than linearly with the size of the block being written for the first several orders of magnitude, and (b) xarray performance tracks numpy's fairly well, though there is a decent gap for the first few orders:

In [8]: data = pd.concat([
   ...:     pd.DataFrame(times, columns=['fill_order', 'time']).assign(package='numpy'),
   ...:     pd.DataFrame(times_xr, columns=['fill_order', 'time']).assign(package='xarray'),
   ...: ], axis=0)
   ...: ax = sns.lineplot(data=data, x='fill_order', y='time', hue='package')
   ...: ax.set_yscale('log')
   ...: ax.set_xlabel('log10 fill size')
   ...: ax.set_ylabel('time (s)')

figure showing xarray and numpy array fill performance, with numpy slightly outperforming xarray but both tracking each other pretty well at higher array sizes

user_38887190
Related