如何用numba加速xarray计算？

Question

I am trying to better understand numba decorators especially guvectorize.

I tried to start here . In particular at step 15 at the very bottom.

I tried to modify this to calculate wind speed.

Here is what I got:

import numpy as np
import xarray as xr
import datetime
import glob
import dask

import sys
import os
import tempfile

from numba import float64, guvectorize, vectorize, njit

import time as t

@guvectorize(
    "(float64, float64, float64)",
    "(), () -> ()",
    nopython=True,
)
def calcWindspeed_ufunc(u, v, out):
        out = np.sqrt( u**2 + v**2 )


def calcWindspeed(u, v):

    return xr.apply_ufunc(calcWindspeed_ufunc, u, v,
                         input_core_dims=[[],[]],
                         output_core_dims=[[]],
                         # vectorize=True,
                         dask="parallelized",
                         output_dtypes=[u.dtype])


def main():

    nlon = 120
    nlat = 100
    ntime = 3650
    lon = np.linspace(129.4, 153.75, nlon)
    lat = np.linspace(-43.75, -10.1, nlat)
    time = np.linspace(0, 365, ntime)

    #< Create random data
    u = 10 * np.random.rand(len(time), len(lat), len(lon))
    u = xr.Dataset({"u": (["time", "lat", "lon"], u)},coords={"time": time, "lon": lon, "lat": lat})
    u = u.chunk({'time':365})
    u = u['u']
    v = u.copy()


    start = t.time()
    ws_xr = np.sqrt( u**2 + v**2 ).load()
    end = t.time()
    print('It took xarray {} seconds!'.format(end-start))

    start = t.time()
    ws_ufunc = calcWindspeed(u, v).load()
    end = t.time()
    print('It took numba {} seconds!'.format(end-start))

    # Difference of the output
    print( (ws_xr-ws_ufunc).max() )


if __name__ == '__main__':
    import dask.distributed
    import sys

    # Get the number of CPUS in the job and start a dask.distributed cluster
    mem          = 190
    cores        = 4
    memory_limit = '{}gb'.format(int(max(mem/cores, 4)))
    client       = dask.distributed.Client(n_workers=cores, threads_per_worker=1, memory_limit=memory_limit, local_dir=tempfile.mkdtemp())


    #< Print client summary
    print('### Client summary')
    print(client)
    print('\n\n')

    #< Call the main function
    main()

    #< Close the client
    client.shutdown()

This works technically (it runs) but the output is wrong. The difference between both calculations should be close to 0 but in my case is 14.

I don't understand what I am doing wrong.

Thank you for your help!

Answer 1

A couple of thoughts:

• There's no need to use numba if it's just calling numpy. Numba runs compiled code, but the current example doesn't really have any code...
• If you're using this to run over multiple dimensions, you can do that with xr.apply_ufunc alone
• If you want others to engage with the example, could you slim it down to its minimum size? Currently there's dask, xarray, numba — if you cut those out does the diff still hold?

As a reference, here are some functions I wrote using xarray and numbahttps://github.com/shoyer/numbagg/blob/master/numbagg/moving.py