diff --git a/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb b/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb index 46c4180a..64dbca57 100644 --- a/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb +++ b/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb @@ -46,6 +46,7 @@ "1. Understand the difference between core dimensions and loop dimensions\n", "1. Understand vectorization\n", "1. Learn how to apply such functions without loops using `apply_ufunc` by providing the `vectorize` keyword argument.\n", + "1. Learn how to only loop over some of the loop dimensions.\n", "\n", "## Core dimensions and looping\n", "\n", @@ -331,12 +332,199 @@ " 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", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "14", + "metadata": { + "tags": [] + }, + "source": [ + "## Looping over some dimensions only\n", + "\n", + "Let's use different new latitudes at every time step. They now vary along `time`, but not along `lon`. The same call as above still works:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "newlat_t = xr.DataArray(\n", + " np.linspace(20, 70, 100) + np.arange(4)[:, np.newaxis], dims=(\"time\", \"lat\")\n", + ")\n", + "\n", + "interped = xr.apply_ufunc(\n", + " debug_interp, # first the function\n", + " newlat_t,\n", + " air.lat,\n", + " air,\n", + " input_core_dims=[[\"lat\"], [\"lat\"], [\"lat\"]],\n", + " output_core_dims=[[\"lat\"]],\n", + " exclude_dims={\"lat\"}, # dimensions allowed to change size. Must be set!\n", + " vectorize=True,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": { + "tags": [] + }, + "source": [ + "`debug_interp` is again called 12 times, but the 3 calls for the `lon` values of each `time` all use the same new latitudes. Many functions can process a whole stack of vectors at once in that case. For example, `scipy.interpolate.make_interp_spline` interpolates all vectors along the last axis at once, to a single set of new latitudes.\n", + "\n", + "So we only need to loop over `time`. Instead of `True`, we pass the dimensions to loop over to `vectorize`. All other loop dimensions, here `lon`, are passed to our function in bulk:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from scipy.interpolate import make_interp_spline\n", + "\n", + "\n", + "def debug_interp_batched(xi, x, data):\n", + " print(f\"data: {data.shape} | x: {x.shape} | xi: {xi.shape}\")\n", + " # xi does not vary along lon, so it has size 1 along that dimension\n", + " xi = xi.reshape(xi.shape[-1])\n", + " return make_interp_spline(x, data, k=1, axis=-1)(xi)\n", + "\n", + "\n", + "interped_batched = xr.apply_ufunc(\n", + " debug_interp_batched, # only the function and vectorize changed\n", + " newlat_t,\n", + " air.lat,\n", + " air,\n", + " input_core_dims=[[\"lat\"], [\"lat\"], [\"lat\"]],\n", + " output_core_dims=[[\"lat\"]],\n", + " exclude_dims={\"lat\"}, # dimensions allowed to change size. Must be set!\n", + " vectorize=[\"time\"], # only loop over time\n", + ")\n", + "xr.testing.assert_allclose(interped, interped_batched)\n", + "interped_batched" + ] + }, + { + "cell_type": "markdown", + "id": "18", + "metadata": { + "tags": [] + }, + "source": [ + "Now `debug_interp_batched` is called only 4 times, once per `time`, and receives all 3 `lon` at once: `data` has shape `(lon, lat)`. The result is the same.\n", + "\n", + "Loop dimensions that are passed in bulk follow numpy's broadcasting rules. `xi` does not vary along `lon`, so it has size 1 along that dimension, and our function has to handle that." + ] + }, + { + "cell_type": "markdown", + "id": "19", + "metadata": { + "tags": [] + }, + "source": [ + "### Performance\n", + "\n", + "Remember the warning above that `numpy.vectorize` is essentially a for loop? Every call of our function adds Python overhead, so the number of loop iterations matters for larger arrays. Let's compare both approaches on the full dataset with 2920 time steps and 53 longitudes, without the `print` statements:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "20", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "air_full = xr.tutorial.load_dataset(\"air_temperature\").air.sortby(\"lat\")\n", + "newlat_full = xr.DataArray(\n", + " np.linspace(20, 70, 100) + np.linspace(0, 3, air_full.sizes[\"time\"])[:, np.newaxis],\n", + " dims=(\"time\", \"lat\"),\n", + ")\n", "\n", + "\n", + "def interp_np(xi, x, data):\n", + " return np.interp(xi, x, data)\n", + "\n", + "\n", + "def interp_batched(xi, x, data):\n", + " return make_interp_spline(x, data, k=1, axis=-1)(xi.reshape(xi.shape[-1]))\n", + "\n", + "\n", + "def interp_full(func, vectorize):\n", + " return xr.apply_ufunc(\n", + " func,\n", + " newlat_full,\n", + " air_full.lat,\n", + " air_full,\n", + " input_core_dims=[[\"lat\"], [\"lat\"], [\"lat\"]],\n", + " output_core_dims=[[\"lat\"]],\n", + " exclude_dims={\"lat\"},\n", + " vectorize=vectorize,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "%timeit -n 1 -r 3 interp_full(interp_np, vectorize=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "22", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "%timeit -n 1 -r 3 interp_full(interp_batched, vectorize=[\"time\"])" + ] + }, + { + "cell_type": "markdown", + "id": "23", + "metadata": { + "tags": [] + }, + "source": [ + "Looping over `time` only is about 3 times faster: `interp_batched` is called 2920 times instead of 2920 × 53 = 154,760 times, even though a single call of `make_interp_spline` is slower than a single call of `np.interp`." + ] + }, + { + "cell_type": "markdown", + "id": "24", + "metadata": { + "tags": [] + }, + "source": [ "## Conclusion\n", "This is why `apply_ufunc` is so convenient; it takes care of a lot of code necessary to apply functions that consume and produce numpy arrays to xarray objects.\n", "\n", - "The `vectorize` keyword argument, when set to True, will use `numpy.vectorize` to apply the function by looping over the \"loop dimensions\" --- dimensions that are not the core dimensions for the applied function." + "The `vectorize` keyword argument, when set to True, will use `numpy.vectorize` to apply the function by looping over the \"loop dimensions\" --- dimensions that are not the core dimensions for the applied function.\n", + "\n", + "If the function can handle some loop dimensions itself, pass the names of the dimensions to loop over instead, e.g. `vectorize=[\"time\"]`, to avoid unnecessary loop iterations." ] } ], diff --git a/advanced/apply_ufunc/example-interp.ipynb b/advanced/apply_ufunc/example-interp.ipynb index 57318fd1..2eae95df 100644 --- a/advanced/apply_ufunc/example-interp.ipynb +++ b/advanced/apply_ufunc/example-interp.ipynb @@ -40,6 +40,7 @@ "1. Handling core dimensions of the output with `output_core_dims`\n", "1. Handling core dimensions that change size using `exclude_dims`\n", "1. Automatic vectorizing or looping over dimensions that are not core dimensions using `vectorize=True`\n", + "1. Only looping over some of these dimensions using `vectorize=[\"time\"]`\n", "1. Automatically parallelization with dask arrays using `dask=\"parallelized\"`\n", "1. High-performance vectorization with numba and `vectorize=False`.\n", "\n", @@ -546,6 +547,80 @@ "\n" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Looping over some dimensions only" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "`vectorize=True` calls our function once for every combination of the non-core dimensions, here 12 times for 4 `time` and 3 `lon` values. That is slow for large arrays, but often not necessary: many functions can process a whole stack of vectors at once, as long as the other inputs are the same for all of them.\n", + "\n", + "Let's interpolate to different latitudes at every time step. The new latitudes now vary along `time`, but not along `lon`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "newlat_t = xr.DataArray(\n", + " np.linspace(20, 70, 100) + np.arange(4)[:, np.newaxis],\n", + " dims=(\"time\", \"new_lat\"),\n", + " coords={\"time\": air.time},\n", + ")\n", + "expected_t = air.interp(lat=newlat_t).drop_vars(\"lat\")\n", + "newlat_t" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Instead of `True`, we can pass the dimensions to loop over to `vectorize`, here `time`. All other non-core dimensions, here `lon`, are then passed to our function in bulk, so it has to handle extra leading dimensions following numpy's broadcasting rules. `np.interp` only supports 1D data, so we use scipy's `make_interp_spline` with `k=1` (linear) instead, which interpolates all vectors along the last axis at once:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from scipy.interpolate import make_interp_spline\n", + "\n", + "\n", + "def interp1d_batched(data, x, xi):\n", + " print(f\"data: {data.shape} | x: {x.shape} | xi: {xi.shape}\")\n", + " # xi does not vary along lon, so it has size 1 along that dimension\n", + " xi = xi.reshape(xi.shape[-1])\n", + " return make_interp_spline(x, data, k=1, axis=-1)(xi)\n", + "\n", + "\n", + "interped_t = xr.apply_ufunc(\n", + " interp1d_batched,\n", + " air,\n", + " air.lat,\n", + " newlat_t,\n", + " input_core_dims=[[\"lat\"], [\"lat\"], [\"new_lat\"]],\n", + " output_core_dims=[[\"new_lat\"]],\n", + " exclude_dims={\"lat\"},\n", + " vectorize=[\"time\"], # only loop over time\n", + ")\n", + "xr.testing.assert_allclose(expected_t.transpose(*interped_t.dims), interped_t)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Our function is now called only 4 times, once per `time`, and receives all `lon` at once: `data` has shape `(lon, lat)`. `xi` keeps a size-1 axis for `lon`, because non-core dimensions follow numpy's broadcasting rules, so we drop it before interpolating. This also works with `dask=\"parallelized\"`, where the loop runs within each block." + ] + }, { "cell_type": "markdown", "metadata": {},