<jupyter_start><jupyter_text>Parallel processing with Dask [Dask](https://dask.org/) is a useful tool when working with large analyses (either in space or time) as it breaks data into manageable chunks that can be easily stored in memory. It can also use multiple computing cores to speed up computations. This has numerous benefits for analyses, which will be covered in this notebook.   This notebook covers how to enable Dask as part of loading data, which can allow you to analyse larger areas and longer time-spans without crashing the Digital Earth Australia environments, as well as potentially speeding up your calculations. Topics covered in this notebook include:1. The difference between the standard load command and loading with Dask.2. Enabling Dask and the Dask Dashboard.3. Setting chunk sizes for data loading.4. Loading data with Dask.5. Chaining operations together before loading any data and understanding task graphs.  To run this introduction to Dask, run all the cells in the notebook starting with the "Load packages" cell. For help with running notebook cells, refer back to the [Jupyter Notebooks notebook](01_Jupyter_notebooks.ipynb).  This cell imports the `datacube` package, which already includes Dask functionality. The `sys` package provides access to helpful support functions in the `dea_dask` module, specifically the `create_local_dask_cluster` function. The `dea_plotting` module contains the function `rgb` which will allow us to plot the data we load using Dask.<jupyter_code>import datacube

import sys
sys.path.append("./Scripts")
from dea_dask import create_local_dask_cluster
from dea_plotting import rgb<jupyter_output><empty_output><jupyter_text>Connect to the datacube<jupyter_code>dc = datacube.Datacube(app="dask_parallel_processing")<jupyter_output><empty_output><jupyter_text>Standard load: By default, the `datacube` library will **not** use Dask when loading data. This means that when `dc.load()` is used, all data relating to the load query will be requested and loaded into memory. For very large areas or long time spans, this can cause the Jupyter notebook to crash. For more information on how to use `dc.load()`, see the [Loading data](04_Loading_data.ipynb) notebook. Below, we show a standard load example:<jupyter_code>data = dc.load(product="ls8_nbart_geomedian_annual",
               measurements=["red", "green", "blue"],
               x=(128.95, 130.15),
               y=(-14.25, -15.40),
               time=("2018-01-01", "2018-12-31"))

data<jupyter_output><empty_output><jupyter_text>Enabling Dask: One of the major features of Dask is that it can take advantage of multiple CPU cores to speed up computations, which is known as distributed computing. This is good for situations where you need to do a lot of calculations on large datasets. To set up distributed computing with Dask, you need to first set up a Dask client using the function below:<jupyter_code>create_local_dask_cluster()<jupyter_output><empty_output><jupyter_text>A print out should appear, displaying information about the `Client` and the `Cluster`. For now, we're most interested in the hyperlink after the **Dashboard:** heading, which should look something like [/user/\/proxy/8787/status](), where [\]() is your username. This link provides a way for you to view how any computations you run are progressing. There are two ways to view the dashboard:1. Click the link, which will open a new tab in your browser2. Set up the dashboard inside the JupyterLab environment. We'll now cover how to do the second option.  Dask dashboard in JupyterLab: On the left-hand menu bar, click the Dask icon, as shown below:![Image](data/dask.png) Copy and paste the **Dashboard** link from the Client print out in the `DASK DASHBOARD URL` text box:![Image](data/dask_url_filled.png) If the url is valid, the buttons should go from grey to orange. Click the orange `PROGRESS` button on the dask panel, which will open a new tab inside JupyterLab. To view the Dask window and your active notebook at the same time, drag the new `Dask Progress` tab to the bottom of the screen. Now, when you do computations with Dask, you'll see the progress of the computations in this new Dask window.  Lazy load: When using Dask, the `dc.load()` function will switch from immediately loading the data to "lazy-loading" the data.  This means the data is only loaded when it is going to be used for a calculation, potentially saving time and memory. Lazy-loading changes the data structure returned from the `dc.load()` command: the returned `xarray.Dataset` will be comprised of `dask.array` objects. To request lazy-loaded data, add a `dask_chunks` parameter to your `dc.load()` call:<jupyter_code>lazy_data = dc.load(product="ls8_nbart_geomedian_annual",
                    measurements=["red", "green", "blue"],
                    x=(128.95, 130.15),
                    y=(-14.25, -15.40),
                    time=("2018-01-01", "2018-12-31"),
                    dask_chunks={"time": 1, "x": 3000, "y": 3000})

lazy_data<jupyter_output><empty_output><jupyter_text>The function should return much faster, as it is not reading any data from disk.  Dask chunks: After adding the `dask_chunks` parameter to `dc.load()`, the lazy-loaded data contains `dask.array` objects with the `chunksize` listed. The `chunksize` should match the `dask_chunks` parameter originally passed to `dc.load()`.  Dask works by breaking up large datasets into chunks, which can be read individually. You may specify the number of pixels in each chunk for each dataset dimension. For example, we passed the following chunk definition to `dc.load()`:```dask_chunks = {"time": 1, "x": 3000, "y": 3000}``` This definition tells Dask to cut the data into chunks containing 3000 pixels in the `x` and `y` dimensions and one measurement in the `time` dimension. For DEA data, we always set `"time": 1` in the `dask_chunk` definition, since the data files only span a single time. If a chunk size is not provided for a given dimension, or if it set to `-1`, then the chunk will be set to the size of the array in that dimension. This means all the data in that dimension will be loaded at once, rather than being broken into smaller chunks.  Viewing Dask chunks: To get a visual intuition for how the data has been broken into chunks, we can use the `.data` attribute provided by `xarray`. This attribute can be used on individual measurements from the lazy-loaded data. When used in a Jupyter Notebook, it provides a table summarising the size of individual chunks and the number of chunks needed. An example is shown below, using the `red` measurement from the lazy-loaded data:<jupyter_code>lazy_data.red.data<jupyter_output><empty_output><jupyter_text>From the Chunk column of the table, we can see that the data has been broken into four chunks, with each chunk having a shape of `(1 time, 3000 pixels, 3000 pixels)` and taking up a maximum of 18.00MB of memory. Comparing this with the Array column, using Dask means that we can load four lots of less than 18.00MB, rather than one lot of 54.63MB.  This is valuable when it comes to working with large areas or time-spans, as the entire array may not always fit into the memory available. Breaking large datasets into chunks and loading chunks one at a time means that you can do computations over large areas without crashing the DEA environment.  Loading lazy data: When working with lazy-loaded data, you have to specifically ask Dask to read and load data when you want to use it. Until you do this, the lazy-loaded dataset only knows where the data is, not its values. To load the data from disk, call `.load()` on the `DataArray` or `Dataset`. If you opened the Dask progress window, you should see the computation proceed there.<jupyter_code>loaded_data = lazy_data.load()

loaded_data<jupyter_output><empty_output><jupyter_text>The Daks arrays constructed by the lazy load, e.g.:```red      (time, y, x) uint16 dask.array``` have now been replaced with actual numbers:``` red      (time, y, x) uint16 10967 11105 10773 10660 ... 12431 12410 12313 ``` After applying the `.load()` command, the lazy-loaded data is the same as the data loaded from the first query. We can now plot the data to verify it loaded correctly:<jupyter_code>rgb(loaded_data, bands=["red", "green", "blue"])<jupyter_output><empty_output><jupyter_text>Lazy operations: In addition to breaking data into smaller chunks that fit in memory, Dask has another advantage in that it can track how you want to work with the data, then only perform the necessary operations later. We'll now explore how to do this by calculating the normalised difference vegetation index (NDVI) for our data. To do this, we'll perform the lazy-load operation again, this time adding the near-infrared band (`nir`) to the `dc.load()` command:<jupyter_code>lazy_data = dc.load(product="ls8_nbart_geomedian_annual",
                    measurements=["red", "green", "blue", "nir"],
                    x=(128.95, 130.15),
                    y=(-14.25, -15.40),
                    time=("2018-01-01", "2018-12-31"),
                    dask_chunks={"time": 1, "x": 3000, "y": 3000})

lazy_data<jupyter_output><empty_output><jupyter_text>Task graphs: When using lazy-loading, Dask breaks up the loading operation into a series of steps. A useful way to visualise the steps is the task graph, which can be accessed by adding the `.visualize()` method to a `.data` call:<jupyter_code>lazy_data.red.data.visualize()<jupyter_output><empty_output><jupyter_text>**The task graph is read from bottom to top.**1. The rectangles at the bottom of the graph are the database entries describing the files that need to be read to load the data.2. Above the rectangles are individual load commands that will do the reading.There is one for each chunk.The arrows describe which files need to be read for each operation: the chunk on the left needs data from all four database entries (i.e. four arrows), whereas the chunk on the right only needs data from one.3. At the very top are the indexes of the chunks that will make up the final array.  Adding more tasks: The power of this method comes from chaining tasks together before loading the data. This is because Dask will only load the data that is required by the final operation in the chain. We can demonstrate this by requesting only a small portion of the red band. If we do this for the lazy-loaded data, we can view the new task graph:<jupyter_code>extract_from_red = lazy_data.red[:, 100:200, 100:200]
extract_from_red.data.visualize()<jupyter_output><empty_output><jupyter_text>Notice that the new task `getitem` has been added, and that it only applies to the left-most chunk. If we call `.load()` on the `extract_from_red` Dask array, Dask trace the operation back through the graph to find only the relevant data. This can save both memory and time. We can establish that the above operation yields the same result as loading the data without Dask and subsetting it by running the command below:<jupyter_code>lazy_red_subset = extract_from_red.load()
data_red_subset = data.red[:, 100:200, 100:200]

print(f"The loaded arrays match: {lazy_red_subset.equals(data_red_subset)}")<jupyter_output><empty_output><jupyter_text>Since the arrays are the same, it is worth using lazy-loading to chain operations together, then calling `.load()` only when you're ready to get the answer. This saves time and memory, since Dask will only load the input data that is required to get the final output.  In this example, the lazy-load only needed to load a small section of the `red` band, whereas the original load to get `data` had to load the `red`, `green` and `blue` bands, then subset the `red` band, meaning time was spent loading data that wasn't used.  Multiple tasks: The power of using lazy-loading in Dask is that you can continue to chain operations together until you are ready to get the answer. Here, we chain multiple steps together to calculate a new band for our array. Specifically, we use the `red` and `nir` bands to calculate the [Normalised Difference Vegetation Index](https://en.wikipedia.org/wiki/Normalized_difference_vegetation_index) (NDVI), which can help identify areas of healthy vegetation. For remote sensing data such as satellite imagery, it is defined as:$$\begin{aligned}\text{NDVI} & = \frac{(\text{NIR} - \text{Red})}{(\text{NIR} + \text{Red})}, \\\end{aligned}$$ where $\text{NIR}$ is the near-infrared band of the data, and $\text{Red}$ is the red band:<jupyter_code>lazy_data["ndvi"] = (lazy_data.nir - lazy_data.red) / (lazy_data.nir + lazy_data.red)<jupyter_output><empty_output><jupyter_text>Doing this adds the new `ndvi` Dask array to the `lazy_data` dataset:<jupyter_code>lazy_data<jupyter_output><empty_output><jupyter_text>Now that the operation is defined, we can view its task graph:<jupyter_code>lazy_data.ndvi.data.visualize()<jupyter_output><empty_output><jupyter_text>Reading the graph bottom-to-top, we can see the equation taking place. The `add` and `sub` are performed on each band before being divided (`truediv`). We can see how each output chunk is independent from the others.  This means we could calculate each chunk without ever having to load all of the bands into memory at the same time. Finally, we can calculate the NDVI values and load the result into memory by calling the `.load()` command. We'll store the result in the `ndvi_load` variable:<jupyter_code>ndvi_load = lazy_data.ndvi.load()
ndvi_load<jupyter_output><empty_output><jupyter_text>We can plot a subset of our NDVI data to verify that Dask correctly calculated our sequence of tasks. Vegetation should appear in green (indicating high NDVI values), and non-vegetation in red (indicating low NDVI values, e.g. over water):<jupyter_code>ndvi_load[:, 2000:4000, 2000:4000].plot(cmap="RdYlGn")<jupyter_output><empty_output><jupyter_text>Note that running the `.load()` command also modifies the `ndvi` entry in the `lazy_load` dataset:<jupyter_code>lazy_data<jupyter_output><empty_output><jupyter_text>You can see that `ndvi` is a number, whereas all the other variables are Dask arrays.  Keeping variables as Dask arrays: If you wanted to calculate the NDVI values, but leave `ndvi` as a dask array in `lazy_load`, you can use the `.compute()` command instead. To demonstrate this, we first redefine the `ndvi` variable so that it becomes a Dask array again:<jupyter_code>lazy_data["ndvi"] = (lazy_data.nir - lazy_data.red) / (lazy_data.nir + lazy_data.red)
lazy_data<jupyter_output><empty_output><jupyter_text>Now, we perform the same steps as before to calculate NDVI, but use `.compute()` instead of `.load()`:<jupyter_code>ndvi_compute = lazy_data.ndvi.compute()
ndvi_compute<jupyter_output><empty_output><jupyter_text>You can see that the values have been calculated, but as shown below, the `ndvi` variable is kept as a Dask array.<jupyter_code>lazy_data<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.07753492413608716, 'fraction_numerical': 0.01711115567141234, 'mean_word_length': 4.759045419553503, 'pattern_counts': {'":': 10, '<': 85, '<?xml version=': 0, '>': 85, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '21', 'dir': 'jupyter-structured-clean-dedup', 'id': '416890', 'license': 'Apache-2.0', 'n_tokens_mistral': 4416, 'n_tokens_neox': 4140, 'n_words': 2018, 'path': 'dea_materials/day2/8_dask_parallel_processing.ipynb', 'repo_name': 'ANU-WALD/dea_course', 'size': '14961'}