From 01e95147aeda7fece5c03331e71a98296d75b46c Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 20 Oct 2026 23:13:24 +0200 Subject: [PATCH 1/5] Add section on looping over some dimensions only with vectorize=[dims] Co-Authored-By: Claude Opus 5.5 --- advanced/apply_ufunc/example-interp.ipynb | 75 +++++++++++++++++++++++ 1 file changed, 75 insertions(+) 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": {}, From 3a4482de7768a7ef590e636f4fb4cc6c4dfc8210 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 20 Oct 2026 23:19:17 +0200 Subject: [PATCH 2/5] Add vectorize=[dims] section to the automatic vectorization page Co-Authored-By: Claude Opus 5.5 --- .../automatic-vectorizing-numpy.ipynb | 77 ++++++++++++++++++- 1 file changed, 75 insertions(+), 2 deletions(-) diff --git a/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb b/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb index 46c4180a..f29bccb0 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,84 @@ " 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", + "`vectorize=True` loops over all loop dimensions, so `debug_interp` was called 12 times. But many functions can process a whole stack of vectors at once, as long as the other inputs are the same for all of them. For example, `scipy.interpolate.make_interp_spline` interpolates all vectors along the last axis at once, but only to a single set of new latitudes.\n", + "\n", + "Let's use different new latitudes at every time step. They vary along `time`, but not along `lon`, 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": "15", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from scipy.interpolate import make_interp_spline\n", + "\n", + "newlat_t = xr.DataArray(\n", + " np.linspace(20, 70, 100) + np.arange(4)[:, np.newaxis], dims=(\"time\", \"lat\")\n", + ")\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 = xr.apply_ufunc(\n", + " debug_interp_batched,\n", + " newlat_t,\n", + " air.lat,\n", + " air,\n", + " input_core_dims=[[\"lat\"], [\"lat\"], [\"lat\"]],\n", + " output_core_dims=[[\"lat\"]],\n", + " exclude_dims={\"lat\"},\n", + " vectorize=[\"time\"], # only loop over time\n", + ")\n", + "interped" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "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)`.\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": "17", + "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." ] } ], From 050d1159e84bb66a5c9ccb4caba9b00268807774 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 20 Oct 2026 23:20:36 +0200 Subject: [PATCH 3/5] Reuse the vectorize=True example for vectorize=[dims] Co-Authored-By: Claude Opus 5.5 --- .../automatic-vectorizing-numpy.ipynb | 56 +++++++++++++++---- 1 file changed, 44 insertions(+), 12 deletions(-) diff --git a/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb b/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb index f29bccb0..57eaa04e 100644 --- a/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb +++ b/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb @@ -344,9 +344,7 @@ "source": [ "## Looping over some dimensions only\n", "\n", - "`vectorize=True` loops over all loop dimensions, so `debug_interp` was called 12 times. But many functions can process a whole stack of vectors at once, as long as the other inputs are the same for all of them. For example, `scipy.interpolate.make_interp_spline` interpolates all vectors along the last axis at once, but only to a single set of new latitudes.\n", - "\n", - "Let's use different new latitudes at every time step. They vary along `time`, but not along `lon`, 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:" + "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:" ] }, { @@ -358,12 +356,45 @@ }, "outputs": [], "source": [ - "from scipy.interpolate import make_interp_spline\n", - "\n", "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", @@ -372,34 +403,35 @@ " return make_interp_spline(x, data, k=1, axis=-1)(xi)\n", "\n", "\n", - "interped = xr.apply_ufunc(\n", - " debug_interp_batched,\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\"},\n", + " exclude_dims={\"lat\"}, # dimensions allowed to change size. Must be set!\n", " vectorize=[\"time\"], # only loop over time\n", ")\n", - "interped" + "xr.testing.assert_allclose(interped, interped_batched)\n", + "interped_batched" ] }, { "cell_type": "markdown", - "id": "16", + "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)`.\n", + "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": "17", + "id": "19", "metadata": { "tags": [] }, From 585f838a26e8043ba3020746d7e73028eae343df Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 20 Oct 2026 23:22:06 +0200 Subject: [PATCH 4/5] Add performance comparison of vectorize=True and vectorize=[dims] Co-Authored-By: Claude Opus 5.5 --- .../automatic-vectorizing-numpy.ipynb | 83 +++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb b/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb index 57eaa04e..28f4fefb 100644 --- a/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb +++ b/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb @@ -435,6 +435,89 @@ "metadata": { "tags": [] }, + "source": [ + "### Performance\n", + "\n", + "The difference 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", From 88cd8759bd46693d030e144eb16578fa39700947 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 20 Oct 2026 23:22:59 +0200 Subject: [PATCH 5/5] Refer back to the np.vectorize performance warning Co-Authored-By: Claude Opus 5.5 --- advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb b/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb index 28f4fefb..64dbca57 100644 --- a/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb +++ b/advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb @@ -438,7 +438,7 @@ "source": [ "### Performance\n", "\n", - "The difference matters for larger arrays. Let's compare both approaches on the full dataset with 2920 time steps and 53 longitudes, without the `print` statements:" + "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:" ] }, {