doc/examples/apply_ufunc_vectorize_1d.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Applying unvectorized functions with `apply_ufunc`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example will illustrate how to conveniently apply an unvectorized function `func` to xarray objects using `apply_ufunc`. `func` expects 1D numpy arrays and returns a 1D numpy array. Our goal is to conveniently apply this function along a dimension of xarray objects that may or may not wrap dask arrays with a signature.\n",
    "\n",
    "We will illustrate this using `np.interp`: \n",
    "\n",
    "    Signature: np.interp(x, xp, fp, left=None, right=None, period=None)\n",
    "    Docstring:\n",
    "        One-dimensional linear interpolation.\n",
    "\n",
    "    Returns the one-dimensional piecewise linear interpolant to a function\n",
    "    with given discrete data points (`xp`, `fp`), evaluated at `x`.\n",
    "\n",
    "and write an `xr_interp` function with signature\n",
    "\n",
    "    xr_interp(xarray_object, dimension_name, new_coordinate_to_interpolate_to)\n",
    "\n",
    "### Load data\n",
    "\n",
    "First lets load an example dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:45:51.659160Z",
     "start_time": "2020-01-15T14:45:50.528742Z"
    }
   },
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "\n",
    "xr.set_options(display_style=\"html\")  # fancy HTML repr\n",
    "\n",
    "air = (\n",
    "    xr.tutorial.load_dataset(\"air_temperature\")\n",
    "    .air.sortby(\"lat\")  # np.interp needs coordinate in ascending order\n",
    "    .isel(time=slice(4), lon=slice(3))\n",
    ")  # choose a small subset for convenience\n",
    "air"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function we will apply is `np.interp` which expects 1D numpy arrays. This functionality is already implemented in xarray so we use that capability to make sure we are not making mistakes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:45:55.431708Z",
     "start_time": "2020-01-15T14:45:55.104701Z"
    }
   },
   "outputs": [],
   "source": [
    "newlat = np.linspace(15, 75, 100)\n",
    "air.interp(lat=newlat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define a function that works with one vector of data along `lat` at a time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:45:57.889496Z",
     "start_time": "2020-01-15T14:45:57.792269Z"
    }
   },
   "outputs": [],
   "source": [
    "def interp1d_np(data, x, xi):\n",
    "    return np.interp(xi, x, data)\n",
    "\n",
    "\n",
    "interped = interp1d_np(air.isel(time=0, lon=0), air.lat, newlat)\n",
    "expected = air.interp(lat=newlat)\n",
    "\n",
    "# no errors are raised if values are equal to within floating point precision\n",
    "np.testing.assert_allclose(expected.isel(time=0, lon=0).values, interped)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### No errors are raised so our interpolation is working.\n",
    "\n",
    "This function consumes and returns numpy arrays, which means we need to do a lot of work to convert the result back to an xarray object with meaningful metadata. This is where `apply_ufunc` is very useful.\n",
    "\n",
    "### `apply_ufunc`\n",
    "\n",
    "    Apply a vectorized function for unlabeled arrays on xarray objects.\n",
    "\n",
    "    The function will be mapped over the data variable(s) of the input arguments using \n",
    "    xarray’s standard rules for labeled computation, including alignment, broadcasting, \n",
    "    looping over GroupBy/Dataset variables, and merging of coordinates.\n",
    "    \n",
    "`apply_ufunc` has many capabilities but for simplicity this example will focus on the common task of vectorizing 1D functions over nD xarray objects. We will iteratively build up the right set of arguments to `apply_ufunc` and read through many error messages in doing so."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:45:59.768626Z",
     "start_time": "2020-01-15T14:45:59.543808Z"
    }
   },
   "outputs": [],
   "source": [
    "xr.apply_ufunc(\n",
    "    interp1d_np,  # first the function\n",
    "    air.isel(time=0, lon=0),  # now arguments in the order expected by 'interp1_np'\n",
    "    air.lat,\n",
    "    newlat,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`apply_ufunc` needs to know a lot of information about what our function does so that it can reconstruct the outputs. In this case, the size of dimension lat has changed and we need to explicitly specify that this will happen. xarray helpfully tells us that we need to specify the kwarg `exclude_dims`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `exclude_dims`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "```\n",
    "exclude_dims : set, optional\n",
    "        Core dimensions on the inputs to exclude from alignment and\n",
    "        broadcasting entirely. Any input coordinates along these dimensions\n",
    "        will be dropped. Each excluded dimension must also appear in\n",
    "        ``input_core_dims`` for at least one argument. Only dimensions listed\n",
    "        here are allowed to change size between input and output objects.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:46:02.187012Z",
     "start_time": "2020-01-15T14:46:02.105563Z"
    }
   },
   "outputs": [],
   "source": [
    "xr.apply_ufunc(\n",
    "    interp1d_np,  # first the function\n",
    "    air.isel(time=0, lon=0),  # now arguments in the order expected by 'interp1_np'\n",
    "    air.lat,\n",
    "    newlat,\n",
    "    exclude_dims=set((\"lat\",)),  # dimensions allowed to change size. Must be set!\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Core dimensions\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Core dimensions are central to using `apply_ufunc`. In our case, our function expects to receive a 1D vector along `lat` &#x2014; this is the dimension that is \"core\" to the function's functionality. Multiple core dimensions are possible. `apply_ufunc` needs to know which dimensions of each variable are core dimensions.\n",
    "\n",
    "    input_core_dims : Sequence[Sequence], optional\n",
    "        List of the same length as ``args`` giving the list of core dimensions\n",
    "        on each input argument that should not be broadcast. By default, we\n",
    "        assume there are no core dimensions on any input arguments.\n",
    "\n",
    "        For example, ``input_core_dims=[[], ['time']]`` indicates that all\n",
    "        dimensions on the first argument and all dimensions other than 'time'\n",
    "        on the second argument should be broadcast.\n",
    "\n",
    "        Core dimensions are automatically moved to the last axes of input\n",
    "        variables before applying ``func``, which facilitates using NumPy style\n",
    "        generalized ufuncs [2]_.\n",
    "        \n",
    "    output_core_dims : List[tuple], optional\n",
    "        List of the same length as the number of output arguments from\n",
    "        ``func``, giving the list of core dimensions on each output that were\n",
    "        not broadcast on the inputs. By default, we assume that ``func``\n",
    "        outputs exactly one array, with axes corresponding to each broadcast\n",
    "        dimension.\n",
    "\n",
    "        Core dimensions are assumed to appear as the last dimensions of each\n",
    "        output in the provided order.\n",
    "        \n",
    "Next we specify `\"lat\"` as `input_core_dims` on both `air` and `air.lat`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:46:05.031672Z",
     "start_time": "2020-01-15T14:46:04.947588Z"
    }
   },
   "outputs": [],
   "source": [
    "xr.apply_ufunc(\n",
    "    interp1d_np,  # first the function\n",
    "    air.isel(time=0, lon=0),  # now arguments in the order expected by 'interp1_np'\n",
    "    air.lat,\n",
    "    newlat,\n",
    "    input_core_dims=[[\"lat\"], [\"lat\"], []],\n",
    "    exclude_dims=set((\"lat\",)),  # dimensions allowed to change size. Must be set!\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "xarray is telling us that it expected to receive back a numpy array with 0 dimensions but instead received an array with 1 dimension corresponding to `newlat`. We can fix this by specifying `output_core_dims`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:46:09.325218Z",
     "start_time": "2020-01-15T14:46:09.303020Z"
    }
   },
   "outputs": [],
   "source": [
    "xr.apply_ufunc(\n",
    "    interp1d_np,  # first the function\n",
    "    air.isel(time=0, lon=0),  # now arguments in the order expected by 'interp1_np'\n",
    "    air.lat,\n",
    "    newlat,\n",
    "    input_core_dims=[[\"lat\"], [\"lat\"], []],  # list with one entry per arg\n",
    "    output_core_dims=[[\"lat\"]],\n",
    "    exclude_dims=set((\"lat\",)),  # dimensions allowed to change size. Must be set!\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we get some output! Let's check that this is right\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:46:11.295440Z",
     "start_time": "2020-01-15T14:46:11.226553Z"
    }
   },
   "outputs": [],
   "source": [
    "interped = xr.apply_ufunc(\n",
    "    interp1d_np,  # first the function\n",
    "    air.isel(time=0, lon=0),  # now arguments in the order expected by 'interp1_np'\n",
    "    air.lat,\n",
    "    newlat,\n",
    "    input_core_dims=[[\"lat\"], [\"lat\"], []],  # list with one entry per arg\n",
    "    output_core_dims=[[\"lat\"]],\n",
    "    exclude_dims=set((\"lat\",)),  # dimensions allowed to change size. Must be set!\n",
    ")\n",
    "interped[\"lat\"] = newlat  # need to add this manually\n",
    "xr.testing.assert_allclose(expected.isel(time=0, lon=0), interped)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "No errors are raised so it is right!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vectorization with `np.vectorize`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now our function currently only works on one vector of data which is not so useful given our 3D dataset.\n",
    "Let's try passing the whole dataset. We add a `print` statement so we can see what our function receives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:46:13.808646Z",
     "start_time": "2020-01-15T14:46:13.680098Z"
    }
   },
   "outputs": [],
   "source": [
    "def interp1d_np(data, x, xi):\n",
    "    print(f\"data: {data.shape} | x: {x.shape} | xi: {xi.shape}\")\n",
    "    return np.interp(xi, x, data)\n",
    "\n",
    "\n",
    "interped = xr.apply_ufunc(\n",
    "    interp1d_np,  # first the function\n",
    "    air.isel(\n",
    "        lon=slice(3), time=slice(4)\n",
    "    ),  # now arguments in the order expected by 'interp1_np'\n",
    "    air.lat,\n",
    "    newlat,\n",
    "    input_core_dims=[[\"lat\"], [\"lat\"], []],  # list with one entry per arg\n",
    "    output_core_dims=[[\"lat\"]],\n",
    "    exclude_dims=set((\"lat\",)),  # dimensions allowed to change size. Must be set!\n",
    ")\n",
    "interped[\"lat\"] = newlat  # need to add this manually\n",
    "xr.testing.assert_allclose(expected.isel(time=0, lon=0), interped)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's a hard-to-interpret error but our `print` call helpfully printed the shapes of the input data: \n",
    "\n",
    "    data: (10, 53, 25) | x: (25,) | xi: (100,)\n",
    "\n",
    "We see that `air` has been passed as a 3D numpy array which is not what `np.interp` expects. Instead we want loop over all combinations of `lon` and `time`; and apply our function to each corresponding vector of data along `lat`.\n",
    "`apply_ufunc` makes this easy by specifying `vectorize=True`:\n",
    "\n",
    "    vectorize : bool, optional\n",
    "        If True, then assume ``func`` only takes arrays defined over core\n",
    "        dimensions as input and vectorize it automatically with\n",
    "        :py:func:`numpy.vectorize`. This option exists for convenience, but is\n",
    "        almost always slower than supplying a pre-vectorized function.\n",
    "        Using this option requires NumPy version 1.12 or newer.\n",
    "        \n",
    "Also see the documentation for `np.vectorize`: https://docs.scipy.org/doc/numpy/reference/generated/numpy.vectorize.html. Most importantly\n",
    "\n",
    "    The vectorize function is provided primarily for convenience, not for performance. \n",
    "    The implementation is essentially a for loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:46:26.633233Z",
     "start_time": "2020-01-15T14:46:26.515209Z"
    }
   },
   "outputs": [],
   "source": [
    "def interp1d_np(data, x, xi):\n",
    "    print(f\"data: {data.shape} | x: {x.shape} | xi: {xi.shape}\")\n",
    "    return np.interp(xi, x, data)\n",
    "\n",
    "\n",
    "interped = xr.apply_ufunc(\n",
    "    interp1d_np,  # first the function\n",
    "    air,  # now arguments in the order expected by 'interp1_np'\n",
    "    air.lat,  # as above\n",
    "    newlat,  # as above\n",
    "    input_core_dims=[[\"lat\"], [\"lat\"], []],  # list with one entry per arg\n",
    "    output_core_dims=[[\"lat\"]],  # returned data has one dimension\n",
    "    exclude_dims=set((\"lat\",)),  # dimensions allowed to change size. Must be set!\n",
    "    vectorize=True,  # loop over non-core dims\n",
    ")\n",
    "interped[\"lat\"] = newlat  # need to add this manually\n",
    "xr.testing.assert_allclose(expected, interped)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This unfortunately is another cryptic error from numpy. \n",
    "\n",
    "Notice that `newlat` is not an xarray object. Let's add a dimension name `new_lat` and modify the call. Note this cannot be `lat` because xarray expects dimensions to be the same size (or broadcastable) among all inputs. `output_core_dims` needs to be modified appropriately. We'll manually rename `new_lat` back to `lat` for easy checking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:46:30.026663Z",
     "start_time": "2020-01-15T14:46:29.893267Z"
    }
   },
   "outputs": [],
   "source": [
    "def interp1d_np(data, x, xi):\n",
    "    print(f\"data: {data.shape} | x: {x.shape} | xi: {xi.shape}\")\n",
    "    return np.interp(xi, x, data)\n",
    "\n",
    "\n",
    "interped = xr.apply_ufunc(\n",
    "    interp1d_np,  # first the function\n",
    "    air,  # now arguments in the order expected by 'interp1_np'\n",
    "    air.lat,  # as above\n",
    "    newlat,  # as above\n",
    "    input_core_dims=[[\"lat\"], [\"lat\"], [\"new_lat\"]],  # list with one entry per arg\n",
    "    output_core_dims=[[\"new_lat\"]],  # returned data has one dimension\n",
    "    exclude_dims=set((\"lat\",)),  # dimensions allowed to change size. Must be a set!\n",
    "    vectorize=True,  # loop over non-core dims\n",
    ")\n",
    "interped = interped.rename({\"new_lat\": \"lat\"})\n",
    "interped[\"lat\"] = newlat  # need to add this manually\n",
    "xr.testing.assert_allclose(\n",
    "    expected.transpose(*interped.dims), interped  # order of dims is different\n",
    ")\n",
    "interped"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the printed input shapes are all 1D and correspond to one vector along the `lat` dimension.\n",
    "\n",
    "The result is now an xarray object with coordinate values copied over from `data`. This is why `apply_ufunc` is so convenient; it takes care of a lot of boilerplate necessary to apply functions that consume and produce numpy arrays to xarray objects.\n",
    "\n",
    "One final point: `lat` is now the *last* dimension in `interped`. This is a \"property\" of core dimensions: they are moved to the end before being sent to `interp1d_np` as was noted in the docstring for `input_core_dims`\n",
    "\n",
    "        Core dimensions are automatically moved to the last axes of input\n",
    "        variables before applying ``func``, which facilitates using NumPy style\n",
    "        generalized ufuncs [2]_."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parallelization with dask\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far our function can only handle numpy arrays. A real benefit of `apply_ufunc` is the ability to easily parallelize over dask chunks _when needed_. \n",
    "\n",
    "We want to apply this function in a vectorized fashion over each chunk of the dask array. This is possible using dask's `blockwise`, `map_blocks`, or `apply_gufunc`. Xarray's `apply_ufunc` wraps dask's `apply_gufunc` and asking it to map the function over chunks using `apply_gufunc` is as simple as specifying `dask=\"parallelized\"`. With this level of flexibility we need to provide dask with some extra information: \n",
    "  1. `output_dtypes`: dtypes of all returned objects, and \n",
    "  2. `output_sizes`: lengths of any new dimensions. \n",
    "  \n",
    "Here we need to specify `output_dtypes` since `apply_ufunc` can infer the size of the new dimension `new_lat` from the argument corresponding to the third element in `input_core_dims`. Here I choose the chunk sizes to illustrate that `np.vectorize` is still applied so that our function receives 1D vectors even though the blocks are 3D."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:48:42.469341Z",
     "start_time": "2020-01-15T14:48:42.344209Z"
    }
   },
   "outputs": [],
   "source": [
    "def interp1d_np(data, x, xi):\n",
    "    print(f\"data: {data.shape} | x: {x.shape} | xi: {xi.shape}\")\n",
    "    return np.interp(xi, x, data)\n",
    "\n",
    "\n",
    "interped = xr.apply_ufunc(\n",
    "    interp1d_np,  # first the function\n",
    "    air.chunk(\n",
    "        {\"time\": 2, \"lon\": 2}\n",
    "    ),  # now arguments in the order expected by 'interp1_np'\n",
    "    air.lat,  # as above\n",
    "    newlat,  # as above\n",
    "    input_core_dims=[[\"lat\"], [\"lat\"], [\"new_lat\"]],  # list with one entry per arg\n",
    "    output_core_dims=[[\"new_lat\"]],  # returned data has one dimension\n",
    "    exclude_dims=set((\"lat\",)),  # dimensions allowed to change size. Must be a set!\n",
    "    vectorize=True,  # loop over non-core dims\n",
    "    dask=\"parallelized\",\n",
    "    output_dtypes=[air.dtype],  # one per output\n",
    ").rename({\"new_lat\": \"lat\"})\n",
    "interped[\"lat\"] = newlat  # need to add this manually\n",
    "xr.testing.assert_allclose(expected.transpose(*interped.dims), interped)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Yay! our function is receiving 1D vectors, so we've successfully parallelized applying a 1D function over a block. If you have a distributed dashboard up, you should see computes happening as equality is checked.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### High performance vectorization: gufuncs, numba & guvectorize\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`np.vectorize` is a very convenient function but is unfortunately slow. It is only marginally faster than writing a for loop in Python and looping. A common way to get around this is to write a base interpolation function that can handle nD arrays in a compiled language like Fortran and then pass that to `apply_ufunc`.\n",
    "\n",
    "Another option is to use the numba package which provides a very convenient `guvectorize` decorator: https://numba.pydata.org/numba-doc/latest/user/vectorize.html#the-guvectorize-decorator\n",
    "\n",
    "Any decorated function gets compiled and will loop over any non-core dimension in parallel when necessary. We need to specify some extra information:\n",
    "\n",
    "   1. Our function cannot return a variable any more. Instead it must receive a variable (the last argument) whose contents the function will modify. So we change from `def interp1d_np(data, x, xi)` to `def interp1d_np_gufunc(data, x, xi, out)`. Our computed results must be assigned to `out`. All values of `out` must be assigned explicitly.\n",
    "   \n",
    "   2. `guvectorize` needs to know the dtypes of the input and output. This is specified in string form as the first argument. Each element of the tuple corresponds to each argument of the function. In this case, we specify `float64` for all inputs and outputs: `\"(float64[:], float64[:], float64[:], float64[:])\"` corresponding to `data, x, xi, out`\n",
    "   \n",
    "   3. Now we need to tell numba the size of the dimensions the function takes as inputs and returns as output i.e. core dimensions. This is done in symbolic form i.e. `data` and `x` are vectors of the same length, say `n`; `xi` and the output `out` have a different length, say `m`. So the second argument is (again as a string)\n",
    "         `\"(n), (n), (m) -> (m).\"` corresponding again to `data, x, xi, out`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:48:45.267633Z",
     "start_time": "2020-01-15T14:48:44.943939Z"
    }
   },
   "outputs": [],
   "source": [
    "from numba import float64, guvectorize\n",
    "\n",
    "\n",
    "@guvectorize(\"(float64[:], float64[:], float64[:], float64[:])\", \"(n), (n), (m) -> (m)\")\n",
    "def interp1d_np_gufunc(data, x, xi, out):\n",
    "    # numba doesn't really like this.\n",
    "    # seem to support fstrings so do it the old way\n",
    "    print(\n",
    "        \"data: \" + str(data.shape) + \" | x:\" + str(x.shape) + \" | xi: \" + str(xi.shape)\n",
    "    )\n",
    "    out[:] = np.interp(xi, x, data)\n",
    "    # gufuncs don't return data\n",
    "    # instead you assign to a the last arg\n",
    "    # return np.interp(xi, x, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The warnings are about object-mode compilation relating to the `print` statement. This means we don't get much speed up: https://numba.pydata.org/numba-doc/latest/user/performance-tips.html#no-python-mode-vs-object-mode. We'll keep the `print` statement temporarily to make sure that `guvectorize` acts like we want it to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:48:54.755405Z",
     "start_time": "2020-01-15T14:48:54.634724Z"
    }
   },
   "outputs": [],
   "source": [
    "interped = xr.apply_ufunc(\n",
    "    interp1d_np_gufunc,  # first the function\n",
    "    air.chunk(\n",
    "        {\"time\": 2, \"lon\": 2}\n",
    "    ),  # now arguments in the order expected by 'interp1_np'\n",
    "    air.lat,  # as above\n",
    "    newlat,  # as above\n",
    "    input_core_dims=[[\"lat\"], [\"lat\"], [\"new_lat\"]],  # list with one entry per arg\n",
    "    output_core_dims=[[\"new_lat\"]],  # returned data has one dimension\n",
    "    exclude_dims=set((\"lat\",)),  # dimensions allowed to change size. Must be a set!\n",
    "    # vectorize=True,  # not needed since numba takes care of vectorizing\n",
    "    dask=\"parallelized\",\n",
    "    output_dtypes=[air.dtype],  # one per output\n",
    ").rename({\"new_lat\": \"lat\"})\n",
    "interped[\"lat\"] = newlat  # need to add this manually\n",
    "xr.testing.assert_allclose(expected.transpose(*interped.dims), interped)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Yay! Our function is receiving 1D vectors and is working automatically with dask arrays. Finally let's comment out the print line and wrap everything up in a nice reusable function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-01-15T14:49:28.667528Z",
     "start_time": "2020-01-15T14:49:28.103914Z"
    }
   },
   "outputs": [],
   "source": [
    "from numba import float64, guvectorize\n",
    "\n",
    "\n",
    "@guvectorize(\n",
    "    \"(float64[:], float64[:], float64[:], float64[:])\",\n",
    "    \"(n), (n), (m) -> (m)\",\n",
    "    nopython=True,\n",
    ")\n",
    "def interp1d_np_gufunc(data, x, xi, out):\n",
    "    out[:] = np.interp(xi, x, data)\n",
    "\n",
    "\n",
    "def xr_interp(data, dim, newdim):\n",
    "    interped = xr.apply_ufunc(\n",
    "        interp1d_np_gufunc,  # first the function\n",
    "        data,  # now arguments in the order expected by 'interp1_np'\n",
    "        data[dim],  # as above\n",
    "        newdim,  # as above\n",
    "        input_core_dims=[[dim], [dim], [\"__newdim__\"]],  # list with one entry per arg\n",
    "        output_core_dims=[[\"__newdim__\"]],  # returned data has one dimension\n",
    "        exclude_dims=set((dim,)),  # dimensions allowed to change size. Must be a set!\n",
    "        # vectorize=True,  # not needed since numba takes care of vectorizing\n",
    "        dask=\"parallelized\",\n",
    "        output_dtypes=[\n",
    "            data.dtype\n",
    "        ],  # one per output; could also be float or np.dtype(\"float64\")\n",
    "    ).rename({\"__newdim__\": dim})\n",
    "    interped[dim] = newdim  # need to add this manually\n",
    "\n",
    "    return interped\n",
    "\n",
    "\n",
    "xr.testing.assert_allclose(\n",
    "    expected.transpose(*interped.dims),\n",
    "    xr_interp(air.chunk({\"time\": 2, \"lon\": 2}), \"lat\", newlat),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This technique is generalizable to any 1D function."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "nbsphinx": {
   "allow_errors": true
  },
  "org": null,
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}