dask-array-example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c34c6ec7-4202-4ca1-8318-badb72e62afb",
   "metadata": {},
   "source": [
    "# dask.array Example"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e66f068-f2dd-47f1-a292-708869d6a023",
   "metadata": {},
   "source": [
    "We're going to make our execution engine a local cluster from Dask `distributed` (so that we can view some of the nice diagnostics that are available)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e299cc03-a549-4463-9bb2-9153ff02b116",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from distributed import Client\n",
    "client = Client(\"127.0.0.1:61329\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5917d412-1175-4b46-b68e-e13be207a590",
   "metadata": {},
   "source": [
    "Standard import for `dask.array` is as `da`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffd4c136-4aa7-4cdd-bc26-193d29b2e168",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask.array as da"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7d6faf5-0a42-4c14-974a-1da5648133e2",
   "metadata": {},
   "source": [
    "Create a simple 2D guassian, chunked six times along each dimension:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2cc4576-09f7-4c2a-a6c7-f81a3bbd5c97",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = da.random.standard_normal(size=(1500, 1500), chunks=(250, 250))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0869653-710d-40ca-bac5-cd09c94ca29a",
   "metadata": {},
   "outputs": [],
   "source": [
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16ac24a6-ecc7-42e2-850a-b9a4dbc04c83",
   "metadata": {},
   "source": [
    "A second array; now in three dimensions, chunked in a similar way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7bf26d5-84fd-4c04-9457-74c4f8b5a78e",
   "metadata": {},
   "outputs": [],
   "source": [
    "y = da.random.standard_cauchy(size=(500, 1500, 1500), chunks=(250, 250, 250))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "456a6881-ea82-4b28-ab5b-0b26e20603e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5398e19d-fb86-49f8-a79d-63db36ca7724",
   "metadata": {},
   "source": [
    "Multiple operations in a single line:\n",
    "\n",
    "- Absolute value\n",
    "- Transpose\n",
    "- Subtraction\n",
    "- Sum reduction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb6e5128-58ee-4410-80ce-079cc65f9bc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = (da.fabs(y) - x.T).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2349d168-d9c7-40cb-8363-9f46c89539c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cec53328-ef19-443e-a243-3263b9c42966",
   "metadata": {},
   "source": [
    "Remember, at this point we've still only created tasks, no real calculations. We ask the client to execute the graph by calling `compute`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afdd0044-2fb9-4fc0-b88a-50ac3d89b92c",
   "metadata": {},
   "outputs": [],
   "source": [
    "z.visualize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c94b0de0-cb8f-4e08-ae2c-4f1be43713d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = client.compute(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8f58eec-7462-4733-99d0-c063f48a4bbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a82264d7-a7a4-4106-8d6c-9fef41ae820e",
   "metadata": {},
   "outputs": [],
   "source": [
    "z.result()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00f413bd-da33-4596-b618-3427fb325c56",
   "metadata": {},
   "source": [
    "A quick look at `dask.array` I/O:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75ce3584-71dd-44e3-9f23-2810291907a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = (da.fabs(y) - x.T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44c463fc-7ab6-416e-9665-fd02a1748080",
   "metadata": {},
   "outputs": [],
   "source": [
    "z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab4cccae-049e-4069-befb-d5fb30b223c5",
   "metadata": {},
   "source": [
    "We store the chunked array with Zarr (a format for the storage of chunked, compressed, N-dimensional arrays). `to_zarr` is one of a number of Dask methods that is _not_ lazy by default (using `compute=False` will provide a lazy evaluation mechanism)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "309cfce4-b5be-4ee9-9011-dff3a9364aa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "z.to_zarr(\"z.zarr\", overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48e83383-110e-424a-b528-6852fa3e8d27",
   "metadata": {},
   "source": [
    "Reading the stored data back into an array is indeed lazy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af20e345-c4b8-4e35-83f5-b73e77ce51b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls z.zarr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27a71b6b-d956-485a-a9e2-2ec51cc820b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = da.from_zarr(\"z.zarr\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71acb5f5-faf6-4060-92f9-30efcb777686",
   "metadata": {},
   "outputs": [],
   "source": [
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caab8b65-1565-4f62-af58-9e490f1c0db4",
   "metadata": {},
   "outputs": [],
   "source": [
    "z.sum().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32871e13-ba26-4fd8-9f94-16a464f54c30",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}