{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Lazy evaluation on Dask arrays\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "If you are unfamiliar with Dask, read\n", "[Parallel computing with Dask](https://docs.xarray.dev/en/stable/user-guide/dask.html)\n", "in Xarray documentation first.\n", "\n", "Recall that the regridding process is divided in two steps : computing the\n", "weights and applying the weights. Dask support is much more advanced for the\n", "latter, and this what the first part of this notebook is about.\n", "\n", "Dask allows [lazy evaluation](https://en.wikipedia.org/wiki/Lazy_evaluation) and\n", "[out-of-core computing](https://en.wikipedia.org/wiki/External_memory_algorithm),\n", "to allow processing large volumes of data with limited memory. You may also get\n", "a speed-up by parallelizing the process in some cases, but a general rule of\n", "thumb is that if the data fits in memory, regridding will be faster without\n", "dask.\n" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "%matplotlib inline\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import dask.array as da # need to have dask.array installed, although not directly using it here.\n", "import xarray as xr\n", "import xesmf as xe" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## A simple example\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Prepare input data\n" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "
<xarray.Dataset>\n",
       "Dimensions:  (lat: 25, time: 2920, lon: 53)\n",
       "Coordinates:\n",
       "  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0\n",
       "  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0\n",
       "  * time     (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00\n",
       "Data variables:\n",
       "    air      (time, lat, lon) float32 dask.array<chunksize=(500, 25, 53), meta=np.ndarray>\n",
       "Attributes:\n",
       "    Conventions:  COARDS\n",
       "    title:        4x daily NMC reanalysis (1948)\n",
       "    description:  Data is from NMC initialized reanalysis\\n(4x/day).  These a...\n",
       "    platform:     Model\n",
       "    references:   http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly...
" ], "text/plain": [ "\n", "Dimensions: (lat: 25, time: 2920, lon: 53)\n", "Coordinates:\n", " * lat (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0\n", " * lon (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0\n", " * time (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00\n", "Data variables:\n", " air (time, lat, lon) float32 dask.array\n", "Attributes:\n", " Conventions: COARDS\n", " title: 4x daily NMC reanalysis (1948)\n", " description: Data is from NMC initialized reanalysis\\n(4x/day). These a...\n", " platform: Model\n", " references: http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly..." ] }, "execution_count": 2, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ds = xr.tutorial.open_dataset(\"air_temperature\", chunks={\"time\": 500})\n", "ds" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Frozen({'time': (500, 500, 500, 500, 500, 420), 'lat': (25,), 'lon': (53,)})" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ds.chunks" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", "
\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
Array Chunk
Bytes 14.76 MiB 2.53 MiB
Shape (2920, 25, 53) (500, 25, 53)
Dask graph 6 chunks in 2 graph layers
Data type float32 numpy.ndarray
\n", " 		\n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " 53\n", " 25\n", " 2920\n", "\n", "
" ], "text/plain": [ "dask.array" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ds[\"air\"].data" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Build regridder\n" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "xESMF Regridder \n", "Regridding algorithm: bilinear \n", "Weight filename: bilinear_25x53_59x87.nc \n", "Reuse pre-computed weights? False \n", "Input grid shape: (25, 53) \n", "Output grid shape: (59, 87) \n", "Periodic in longitude? False" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ds_out = xr.Dataset(\n", " {\n", " \"lat\": ([\"lat\"], np.arange(16, 75, 1.0)),\n", " \"lon\": ([\"lon\"], np.arange(200, 330, 1.5)),\n", " }\n", ")\n", "\n", "regridder = xe.Regridder(ds, ds_out, \"bilinear\")\n", "regridder" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Apply to xarray Dataset/DataArray\n" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "CPU times: user 8.17 ms, sys: 2.32 ms, total: 10.5 ms\n", "Wall time: 10.2 ms\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "
<xarray.Dataset>\n",
       "Dimensions:  (time: 2920, lat: 59, lon: 87)\n",
       "Coordinates:\n",
       "  * time     (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00\n",
       "  * lat      (lat) float64 16.0 17.0 18.0 19.0 20.0 ... 70.0 71.0 72.0 73.0 74.0\n",
       "  * lon      (lon) float64 200.0 201.5 203.0 204.5 ... 324.5 326.0 327.5 329.0\n",
       "Data variables:\n",
       "    air      (time, lat, lon) float32 dask.array<chunksize=(500, 59, 87), meta=np.ndarray>\n",
       "Attributes:\n",
       "    regrid_method:  bilinear
" ], "text/plain": [ "\n", "Dimensions: (time: 2920, lat: 59, lon: 87)\n", "Coordinates:\n", " * time (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00\n", " * lat (lat) float64 16.0 17.0 18.0 19.0 20.0 ... 70.0 71.0 72.0 73.0 74.0\n", " * lon (lon) float64 200.0 201.5 203.0 204.5 ... 324.5 326.0 327.5 329.0\n", "Data variables:\n", " air (time, lat, lon) float32 dask.array\n", "Attributes:\n", " regrid_method: bilinear" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# only build the dask graph; actual computation happens later when calling compute()\n", "%time ds_out = regridder(ds)\n", "ds_out" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", "
\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
Array Chunk
Bytes 57.18 MiB 9.79 MiB
Shape (2920, 59, 87) (500, 59, 87)
Dask graph 6 chunks in 8 graph layers
Data type float32 numpy.ndarray
\n", " 		\n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " 87\n", " 59\n", " 2920\n", "\n", "
" ], "text/plain": [ "dask.array" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ds_out[\"air\"].data" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "CPU times: user 756 ms, sys: 155 ms, total: 911 ms\n", "Wall time: 600 ms\n" ] } ], "source": [ "%time result = ds_out['air'].compute() # actually applies regridding" ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "(numpy.ndarray, (2920, 59, 87))" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "type(result.data), result.data.shape" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Chunking behaviour\n", "\n", "xESMF will adjust its default behaviour according to the input data. On spatial\n", "dimensions where the data has only one chunk, the output of a `Regridder` call\n", "will also have only one chunk, no matter the new dimension size. This like the\n", "previous example.\n", "\n", "However, if the input has more than one chunk along a spatial dimension, then\n", "the regridder will try to preserve the chunk size. When upscaling data, this\n", "means the number of chunks increases and with it the number of dask tasks added\n", "to the graph. This can actually decrease performance if the graph becomes too\n", "large, filled up with many small tasks.\n", "\n", "One can always override xESMF's default behaviour by passing `output_chunks` to\n", "the `Regridder` call.\n", "\n", "In the example below, the input has three chunks along `lon`:\n" ] }, { "cell_type": "code", "execution_count": 13, "metadata": { "scrolled": true }, "outputs": [ { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", "
\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
Array Chunk
Bytes 14.76 MiB 6.96 MiB
Shape (2920, 25, 53) (2920, 25, 25)
Dask graph 3 chunks in 3 graph layers
Data type float32 numpy.ndarray
\n", " 		\n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " 53\n", " 25\n", " 2920\n", "\n", "
" ], "text/plain": [ "dask.array" ] }, "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ds_3lon = ds.chunk({\"lat\": 25, \"lon\": 25, \"time\": -1})\n", "ds_3lon.air.data" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "In this case, the output DataArray will have the same chunk size on longitude,\n", "but still only one chunk along latitude.\n" ] }, { "cell_type": "code", "execution_count": 14, "metadata": { "scrolled": true }, "outputs": [ { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", "
\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
Array Chunk
Bytes 57.18 MiB 16.43 MiB
Shape (2920, 59, 87) (2920, 59, 25)
Dask graph 4 chunks in 10 graph layers
Data type float32 numpy.ndarray
\n", " 		\n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " 87\n", " 59\n", " 2920\n", "\n", "
" ], "text/plain": [ "dask.array" ] }, "execution_count": 14, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ds_spatial_out = regridder(ds_spatial) # Regridding ds_spatial\n", "ds_spatial_out[\"air\"].data" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Unless the `output_chunks` argument is passed to the `regridder`\n" ] }, { "cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", "
\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
Array Chunk
Bytes 57.18 MiB 1.11 MiB
Shape (2920, 59, 87) (2920, 10, 10)
Dask graph 54 chunks in 10 graph layers
Data type float32 numpy.ndarray
\n", " 		\n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", "\n", " \n", " 87\n", " 59\n", " 2920\n", "\n", "
" ], "text/plain": [ "dask.array" ] }, "execution_count": 15, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ds_spatial_out = regridder(ds_spatial, output_chunks={\"lat\": 10, \"lon\": 10})\n", "ds_spatial_out[\"air\"].data" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Parallel weight generation with Dask\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Dask can also be used to build the regridder and compute its weights in\n", "parallel. To do so, xESMF uses the chunks on the destination grid and computes\n", "subsets of weights on each chunk in parallel.\n", "\n", "This feature is currently in an experimental state and it will force dask to use\n", "processes to parallelize the computation. Moreover, it is slower than then\n", "normal method abd thus is it _only_ useful if the **destination** grid does not\n", "fit in memory. Recall that the parallization is done over chunks of the\n", "destination grid and each iteration will need to load the source grid in memory.\n", "\n", "For a more performant way to generate weights in parallel, it might be better to\n", "use `ESMF` directly instead, assuming you have an MPI-enabled version. See the\n", "\"Solving large problems using HPC\" notebook.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Parallel weight generation example\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Prepare input data\n" ] }, { "cell_type": "code", "execution_count": 16, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "
<xarray.Dataset>\n",
       "Dimensions:  (lat: 25, time: 2920, lon: 53)\n",
       "Coordinates:\n",
       "  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0\n",
       "  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0\n",
       "  * time     (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00\n",
       "Data variables:\n",
       "    air      (time, lat, lon) float32 dask.array<chunksize=(500, 25, 53), meta=np.ndarray>\n",
       "Attributes:\n",
       "    Conventions:  COARDS\n",
       "    title:        4x daily NMC reanalysis (1948)\n",
       "    description:  Data is from NMC initialized reanalysis\\n(4x/day).  These a...\n",
       "    platform:     Model\n",
       "    references:   http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly...
" ], "text/plain": [ "\n", "Dimensions: (lat: 25, time: 2920, lon: 53)\n", "Coordinates:\n", " * lat (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0\n", " * lon (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0\n", " * time (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00\n", "Data variables:\n", " air (time, lat, lon) float32 dask.array\n", "Attributes:\n", " Conventions: COARDS\n", " title: 4x daily NMC reanalysis (1948)\n", " description: Data is from NMC initialized reanalysis\\n(4x/day). These a...\n", " platform: Model\n", " references: http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly..." ] }, "execution_count": 16, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ds = xr.tutorial.open_dataset(\"air_temperature\", chunks={\"time\": 500})\n", "ds" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Prepare output dataset and chunk it\n" ] }, { "cell_type": "code", "execution_count": 17, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Frozen({'time': (36,), 'y': (50, 50, 50, 50, 5), 'x': (50, 50, 50, 50, 50, 25)})" ] }, "execution_count": 17, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ds_out = xr.tutorial.open_dataset(\"rasm\")\n", "ds_out = ds_out.chunk({\"y\": 50, \"x\": 50})\n", "ds_out.chunks" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Create regridder, generating the weights in parallel\n" ] }, { "cell_type": "code", "execution_count": 18, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "[WARNING] yaksa: 10 leaked handle pool objects\n", "[WARNING] yaksa: 10 leaked handle pool objects\n", "[WARNING] yaksa: 10 leaked handle pool objects\n", "[WARNING] yaksa: 10 leaked handle pool objects\n" ] }, { "data": { "text/plain": [ "xESMF Regridder \n", "Regridding algorithm: bilinear \n", "Weight filename: bilinear_25x53_205x275.nc \n", "Reuse pre-computed weights? False \n", "Input grid shape: (25, 53) \n", "Output grid shape: (205, 275) \n", "Periodic in longitude? False" ] }, "execution_count": 18, "metadata": {}, "output_type": "execute_result" } ], "source": [ "para_regridder = xe.Regridder(ds, ds_out, \"bilinear\", parallel=True)\n", "para_regridder" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Attempting to build the Regridder using the option `parallel=True` with either\n", "`reuse_weights=True` or with `weights != None` will produce a warning. In both\n", "cases, since the weights are already generated, the regridder will be built\n", "skipping the parallel part.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Using a mask to chunk an empty Dataset\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "If the destination grid has no variables and contains 1D lat/lon coordinates,\n", "using xarray's `.chunk()` method will not work\n" ] }, { "cell_type": "code", "execution_count": 19, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "
<xarray.Dataset>\n",
       "Dimensions:  (lat: 59, lon: 87)\n",
       "Coordinates:\n",
       "  * lat      (lat) float64 16.0 17.0 18.0 19.0 20.0 ... 70.0 71.0 72.0 73.0 74.0\n",
       "  * lon      (lon) float64 200.0 201.5 203.0 204.5 ... 324.5 326.0 327.5 329.0\n",
       "Data variables:\n",
       "    *empty*
" ], "text/plain": [ "\n", "Dimensions: (lat: 59, lon: 87)\n", "Coordinates:\n", " * lat (lat) float64 16.0 17.0 18.0 19.0 20.0 ... 70.0 71.0 72.0 73.0 74.0\n", " * lon (lon) float64 200.0 201.5 203.0 204.5 ... 324.5 326.0 327.5 329.0\n", "Data variables:\n", " *empty*" ] }, "execution_count": 19, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ds_out = xr.Dataset(\n", " {\n", " \"lat\": ([\"lat\"], np.arange(16, 75, 1.0), {\"units\": \"degrees_north\"}),\n", " \"lon\": ([\"lon\"], np.arange(200, 330, 1.5), {\"units\": \"degrees_east\"}),\n", " }\n", ")\n", "ds_out" ] }, { "cell_type": "code", "execution_count": 20, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Frozen({})" ] }, "execution_count": 20, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ds_out.chunk({\"lat\": 25, \"lon\": 25})\n", "ds_out.chunks" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To deal with this issue, we can create a `mask` and add it to `ds_out`. Using a\n", "boolean mask ensures `ds_out` is not bloated by data and setting the mask to be\n", "`True` everywhere will not affect regridding.\n" ] }, { "cell_type": "code", "execution_count": 21, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Frozen({'lat': (25, 25, 9), 'lon': (25, 25, 25, 12)})" ] }, "execution_count": 21, "metadata": {}, "output_type": "execute_result" } ], "source": [ "mask = da.ones((ds_out.lat.size, ds_out.lon.size), dtype=bool, chunks=(25, 25))\n", "ds_out[\"mask\"] = (ds_out.dims, mask)\n", "\n", "# Now we check the chunks of ds_out\n", "ds_out.chunks" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "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.11.7" } }, "nbformat": 4, "nbformat_minor": 4 }