title: 'Parallelized' apply_ufunc for scripy.interpolate.griddata
body: <p>I'm working with large files from an ocean model with an unstructered grid. For instance, variable flow velocity <code>ux</code> with dimensions <code>(194988, 1009, 20)</code> for respectively: 'nFlowElement' (name unstructered grid element), 'time' and laydim (depth dimension). I'd like to interpolate these results to a structured grid with dimensions <code>(600, 560, 1009, 20) </code>for respectively: latitude, longitude, time and laydim. For this I am using <code>scipy.interpolate.griddata</code>. As these dataarrays are too large to load into your working memory at once, I am trying to work with 'chunks' (dask). Unfortunately, I bump into problems when trying to use apply_ufunc with setting: <code>dask = 'parallelized'</code>.</p>
  <p>For smaller computational domains (smaller nFlowElement dimension) I ám still able to load the dataarray in my work memory. Then, the following code gives me the wanted result:</p>
  <pre><code>def interp_to_grid(u,xc,yc,xint,yint):
      print(u.shape,xc.shape,xint.shape)
      ug = griddata((xc,yc),u,(xint,yint), method='nearest', fill_value=np.nan)
      return ug 
  
  uxg = xr.apply_ufunc(interp_to_grid,
                       ux, xc, yc, xint, yint,
                       dask = 'allowed',
                       input_core_dims=[['nFlowElem','time','laydim'],['nFlowElem'],['nFlowElem'],['dim_0','dim_1'],['dim_0','dim_1']],
                       output_core_dims=[['dim_0','dim_1','time','laydim']],
                       output_dtypes = [xr.DataArray]
                       )
  </code></pre>
  <p>Notice that in the function interp_to_grid the input variables have the following dimensions:</p>
  <ul>
  <li><code>u</code> (i.e. ux, the original flow velocity output): (194988, 1009, 20) for (nFlowElem, time, laydim)</li>
  <li><code>xc,yc</code> (the latitude and longitude coordinates associated with these 194988 elements) so both (194988,)</li>
  <li><code>xint, yint</code> (the structured grid coordinates to which I would like to interpolate the data): both are (600, 560) for (dim_0,dim_1)
  Notice that scipy.interpolate.griddata does not require me to loop over the time and laydim dimension (as formulated in the code above). For this it is criticial to feed <code>griddata</code> the dimensions in the right order ('time' and 'laydim' last). The interpolated result, uxg, has dimensions (600, 560, 1009, 20) - as wanted and expected.</li>
  </ul>
  <p>However, for much larger spatial domains it is required to work with dask = 'parallelized', because these input dataarrays can nolonger be loaded into my working memory. I have tried to apply chunks over the time dimension, but also over the nFlowElement dimension. I am aware that it is not possible to chunk over core dimensions.</p>
  <p><strong>This is one of my &quot;parallel&quot; attempts (with chunks along the time dim):</strong></p>
  <p>Input ux:</p>
  <pre><code>&lt;xarray.DataArray 'ucx' (nFlowElem: 194988, time: 1009, laydim: 20)&gt;
  dask.array&lt;transpose, shape=(194988, 1009, 20), dtype=float64, chunksize=(194988, 10, 20), chunktype=numpy.ndarray&gt;
  Coordinates:
      FlowElem_xcc  (nFlowElem) float64 dask.array&lt;chunksize=(194988,), meta=np.ndarray&gt;
      FlowElem_ycc  (nFlowElem) float64 dask.array&lt;chunksize=(194988,), meta=np.ndarray&gt;
    * time          (time) datetime64[ns] 2014-09-17 ... 2014-10-01
  Dimensions without coordinates: nFlowElem, laydim
  Attributes:
      standard_name:  eastward_sea_water_velocity
      long_name:      velocity on flow element center, x-component
      units:          m s-1
      grid_mapping:   wgs84
  </code></pre>
  <p>Apply_func:</p>
  <pre><code>uxg = xr.apply_ufunc(interp_to_grid,
                       ux, xc, yc, xint, yint,
                       dask = 'parallelized',
                       input_core_dims=[['nFlowElem'],['nFlowElem'],['nFlowElem'],['dim_0','dim_1'],['dim_0','dim_1']],
                       output_core_dims=[['dim_0','dim_1']],
                       output_dtypes = [xr.DataArray],
                       )
  </code></pre>
  <p>Gives error:</p>
  <pre><code>  File &quot;interpnd.pyx&quot;, line 78, in scipy.interpolate.interpnd.NDInterpolatorBase.__init__
  
    File &quot;interpnd.pyx&quot;, line 192, in scipy.interpolate.interpnd._check_init_shape
  
  ValueError: different number of values and points
  </code></pre>
  <p>I have played around a lot with changing the core dimensions in apply_ufunc and the dimension along which to chunk. Also I have tried to manually change the order of dimensions of dataarray <code>u</code> which is 'fed to' griddata (in <code>interp_to_grid</code>).</p>
  <p>Any advice is very welcome!
  Best Wishes,
  Luka</p>
  <p><strong>### UPDATE 1: ###</strong></p>
  <p>I have also tried to chunk over the <code>nFlowElem</code> dimension instead of 'time'.
  Then I define the input and output core dimensions as written below:</p>
  <pre><code>input_core_dims=[['time','laydim'],[],[],['dim_0','dim_1'],['dim_0','dim_1']],
  output_core_dims=[['dim_0','dim_1','time','laydim']],
  </code></pre>
  <p>Now I &quot;only&quot; get an error concerning the size of my output (which is nolonger chunked).
  My output is nolonger chunked, because nFlowElem is lost after the interpolation.
  Consequently, the (600,560,1009,20) dataarray is too large for my work memory.</p>
  <p>Is there perhaps a way to have apply_ufunc chunk your output along any of the other dimensions?</p>
  <p><strong>### UPDATE 2: ###</strong></p>
  <p>Alternatively, I have tried to chunk over the time dimension (50 time steps) and I have removed all input/output core dimensions.</p>
  <p>And if I then define interp_to_grid as follows (to get the right input dimensions):</p>
  <pre><code>def interp_to_grid(u,xc,yc,xint,yint):
  
      xc = xc[:,0,0,0,0]
      yc = yc[:,0,0,0,0]
      u = u[:,:,:,0,0]
      print(u.shape,xc.shape,xint.shape) #input shape
  
      ug = griddata((xc,yc),u,(xint,yint), method='nearest', fill_value=np.nan)
      print(ug.shape) #output shape
  
      return ug 
  </code></pre>
  <p>I do get the right dimensions for my <code>interp_to_grid</code>-input (see first line) and output (see second line). However, I get the <code>ValueError: axes don't match array</code>:</p>
  <pre><code>(194988, 50, 20) (194988,) (600, 560)
  (600, 560, 50, 20)
  Traceback (most recent call last):
  
    File &quot;&lt;ipython-input-11-6b65fe3dba5b&gt;&quot;, line 1, in &lt;module&gt;
      uxg.loc[dict(time=np.datetime64('2014-09-18'),laydim=19)].plot()
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\xarray\plot\plot.py&quot;, line 444, in __call__
      return plot(self._da, **kwargs)
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\xarray\plot\plot.py&quot;, line 160, in plot
      darray = darray.squeeze().compute()
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\xarray\core\dataarray.py&quot;, line 899, in compute
      return new.load(**kwargs)
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\xarray\core\dataarray.py&quot;, line 873, in load
      ds = self._to_temp_dataset().load(**kwargs)
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\xarray\core\dataset.py&quot;, line 798, in load
      evaluated_data = da.compute(*lazy_data.values(), **kwargs)
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\dask\base.py&quot;, line 565, in compute
      results = schedule(dsk, keys, **kwargs)
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\dask\threaded.py&quot;, line 84, in get
      **kwargs
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\dask\local.py&quot;, line 487, in get_async
      raise_exception(exc, tb)
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\dask\local.py&quot;, line 317, in reraise
      raise exc
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\dask\local.py&quot;, line 222, in execute_task
      result = _execute_task(task, data)
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\dask\core.py&quot;, line 121, in _execute_task
      return func(*(_execute_task(a, cache) for a in args))
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\dask\core.py&quot;, line 121, in &lt;genexpr&gt;
      return func(*(_execute_task(a, cache) for a in args))
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\dask\core.py&quot;, line 121, in _execute_task
      return func(*(_execute_task(a, cache) for a in args))
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\dask\optimization.py&quot;, line 963, in __call__
      return core.get(self.dsk, self.outkey, dict(zip(self.inkeys, args)))
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\dask\core.py&quot;, line 151, in get
      result = _execute_task(task, cache)
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\dask\core.py&quot;, line 121, in _execute_task
      return func(*(_execute_task(a, cache) for a in args))
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\dask\utils.py&quot;, line 35, in apply
      return func(*args, **kwargs)
  
    File &quot;&lt;__array_function__ internals&gt;&quot;, line 6, in transpose
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\numpy\core\fromnumeric.py&quot;, line 658, in transpose
      return _wrapfunc(a, 'transpose', axes)
  
    File &quot;C:\Users\920507\Anaconda3\envs\dfm_tools_env\lib\site-packages\numpy\core\fromnumeric.py&quot;, line 58, in _wrapfunc
      return bound(*args, **kwds)
  
  ValueError: axes don't match array
  </code></pre>