Skip to content

Instantly share code, notes, and snippets.

@maneesh29s

maneesh29s/vectorized_functions_with_dask_cpu.ipynb

Created December 21, 2024 11:30
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save maneesh29s/a59a9f32db25d4d3c764e80c590026b8 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save maneesh29s/a59a9f32db25d4d3c764e80c590026b8 to your computer and use it in GitHub Desktop.
Download ZIP
Gist contains notebooks , which demonstrate how we can use JAX to use vectorized functions with dask arrays. Also contains Numba example for CPU.
Display the source blob
Display the rendered blob
Raw
vectorized_functions_with_dask_cpu.ipynb
{
"cells": [
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ[\"JAX_PLATFORMS\"]=\"cpu\"\n",
"\n",
"from jax import config\n",
"config.update(\"jax_traceback_filtering\", \"off\")"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"from numba import guvectorize, int64\n",
"import numpy as np\n",
"import xarray as xr\n",
"import dask.array\n",
"import jax"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Creating a local cluster"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"from dask.distributed import LocalCluster, WorkerPlugin"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"class EnvPlugin(WorkerPlugin):\n",
" def __init__(self):\n",
" pass\n",
"\n",
" def setup(self, worker):\n",
" import os\n",
" os.environ[\"JAX_PLATFORMS\"]=\"cpu\"\n",
"\n",
" from jax import config\n",
" config.update(\"jax_traceback_filtering\", \"off\")\n",
"\n",
" import numpy as np"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"cluster = LocalCluster(n_workers=4, threads_per_worker=1, scheduler_port=8786, dashboard_address=\":8787\")"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'tcp://127.0.0.1:33227': {'status': 'OK'},\n",
" 'tcp://127.0.0.1:33825': {'status': 'OK'},\n",
" 'tcp://127.0.0.1:37881': {'status': 'OK'},\n",
" 'tcp://127.0.0.1:38461': {'status': 'OK'}}"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"client = cluster.get_client()\n",
"client.register_plugin(EnvPlugin())\n",
"client.forward_logging()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Creating temp data"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"<xarray.backends.zarr.ZarrStore at 0x7fd880700640>"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"shape=(2048,1024,1024)\n",
"chunksizes=(64,1024,1024)\n",
"# # vis = xr.DataArray(np.random.random(shape).astype(dtype=np.float32), dims=[\"z\",\"y\",\"x\"])\n",
"# # flag = xr.DataArray(np.ones(shape, dtype=bool), dims=[\"z\",\"y\",\"x\"])\n",
"# # weights = xr.DataArray(np.random.random(shape).astype(dtype=np.float32), dims=[\"z\",\"y\",\"x\"])\n",
"# # kernel = xr.DataArray(np.random.random((16,16)).astype(dtype=np.float32), dims=[\"ky\",\"kx\"])\n",
"\n",
"vis = xr.DataArray(dask.array.random.random(shape, chunks=chunksizes,).astype(dtype=np.float32), dims=[\"z\",\"y\",\"x\"])\n",
"flag = xr.DataArray(dask.array.ones(shape, chunks=chunksizes, dtype=bool), dims=[\"z\",\"y\",\"x\"])\n",
"weights = xr.DataArray(dask.array.random.random(shape, chunks=chunksizes,).astype(dtype=np.float32), dims=[\"z\",\"y\",\"x\"])\n",
"\n",
"xds = xr.Dataset(dict(vis=vis, flag=flag, weights=weights))\n",
"xds.to_zarr(\"xds_big_rechunked.zarr\", mode='w')\n",
"# # xds.to_zarr(\"xds_small.zarr\", mode='w')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Read zarr dataset"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"xds = xr.open_zarr(\"xds_big.zarr\")"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<div><svg style=\"position: absolute; width: 0; height: 0; overflow: hidden\">\n",
"<defs>\n",
"<symbol id=\"icon-database\" viewBox=\"0 0 32 32\">\n",
"<path d=\"M16 0c-8.837 0-16 2.239-16 5v4c0 2.761 7.163 5 16 5s16-2.239 16-5v-4c0-2.761-7.163-5-16-5z\"></path>\n",
"<path d=\"M16 17c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n",
"<path d=\"M16 26c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n",
"</symbol>\n",
"<symbol id=\"icon-file-text2\" viewBox=\"0 0 32 32\">\n",
"<path d=\"M28.681 7.159c-0.694-0.947-1.662-2.053-2.724-3.116s-2.169-2.030-3.116-2.724c-1.612-1.182-2.393-1.319-2.841-1.319h-15.5c-1.378 0-2.5 1.121-2.5 2.5v27c0 1.378 1.122 2.5 2.5 2.5h23c1.378 0 2.5-1.122 2.5-2.5v-19.5c0-0.448-0.137-1.23-1.319-2.841zM24.543 5.457c0.959 0.959 1.712 1.825 2.268 2.543h-4.811v-4.811c0.718 0.556 1.584 1.309 2.543 2.268zM28 29.5c0 0.271-0.229 0.5-0.5 0.5h-23c-0.271 0-0.5-0.229-0.5-0.5v-27c0-0.271 0.229-0.5 0.5-0.5 0 0 15.499-0 15.5 0v7c0 0.552 0.448 1 1 1h7v19.5z\"></path>\n",
"<path d=\"M23 26h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
"<path d=\"M23 22h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
"<path d=\"M23 18h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
"</symbol>\n",
"</defs>\n",
"</svg>\n",
"<style>/* CSS stylesheet for displaying xarray objects in jupyterlab.\n",
" *\n",
" */\n",
"\n",
":root {\n",
" --xr-font-color0: var(--jp-content-font-color0, rgba(0, 0, 0, 1));\n",
" --xr-font-color2: var(--jp-content-font-color2, rgba(0, 0, 0, 0.54));\n",
" --xr-font-color3: var(--jp-content-font-color3, rgba(0, 0, 0, 0.38));\n",
" --xr-border-color: var(--jp-border-color2, #e0e0e0);\n",
" --xr-disabled-color: var(--jp-layout-color3, #bdbdbd);\n",
" --xr-background-color: var(--jp-layout-color0, white);\n",
" --xr-background-color-row-even: var(--jp-layout-color1, white);\n",
" --xr-background-color-row-odd: var(--jp-layout-color2, #eeeeee);\n",
"}\n",
"\n",
"html[theme=dark],\n",
"html[data-theme=dark],\n",
"body[data-theme=dark],\n",
"body.vscode-dark {\n",
" --xr-font-color0: rgba(255, 255, 255, 1);\n",
" --xr-font-color2: rgba(255, 255, 255, 0.54);\n",
" --xr-font-color3: rgba(255, 255, 255, 0.38);\n",
" --xr-border-color: #1F1F1F;\n",
" --xr-disabled-color: #515151;\n",
" --xr-background-color: #111111;\n",
" --xr-background-color-row-even: #111111;\n",
" --xr-background-color-row-odd: #313131;\n",
"}\n",
"\n",
".xr-wrap {\n",
" display: block !important;\n",
" min-width: 300px;\n",
" max-width: 700px;\n",
"}\n",
"\n",
".xr-text-repr-fallback {\n",
" /* fallback to plain text repr when CSS is not injected (untrusted notebook) */\n",
" display: none;\n",
"}\n",
"\n",
".xr-header {\n",
" padding-top: 6px;\n",
" padding-bottom: 6px;\n",
" margin-bottom: 4px;\n",
" border-bottom: solid 1px var(--xr-border-color);\n",
"}\n",
"\n",
".xr-header > div,\n",
".xr-header > ul {\n",
" display: inline;\n",
" margin-top: 0;\n",
" margin-bottom: 0;\n",
"}\n",
"\n",
".xr-obj-type,\n",
".xr-array-name {\n",
" margin-left: 2px;\n",
" margin-right: 10px;\n",
"}\n",
"\n",
".xr-obj-type {\n",
" color: var(--xr-font-color2);\n",
"}\n",
"\n",
".xr-sections {\n",
" padding-left: 0 !important;\n",
" display: grid;\n",
" grid-template-columns: 150px auto auto 1fr 20px 20px;\n",
"}\n",
"\n",
".xr-section-item {\n",
" display: contents;\n",
"}\n",
"\n",
".xr-section-item input {\n",
" display: none;\n",
"}\n",
"\n",
".xr-section-item input + label {\n",
" color: var(--xr-disabled-color);\n",
"}\n",
"\n",
".xr-section-item input:enabled + label {\n",
" cursor: pointer;\n",
" color: var(--xr-font-color2);\n",
"}\n",
"\n",
".xr-section-item input:enabled + label:hover {\n",
" color: var(--xr-font-color0);\n",
"}\n",
"\n",
".xr-section-summary {\n",
" grid-column: 1;\n",
" color: var(--xr-font-color2);\n",
" font-weight: 500;\n",
"}\n",
"\n",
".xr-section-summary > span {\n",
" display: inline-block;\n",
" padding-left: 0.5em;\n",
"}\n",
"\n",
".xr-section-summary-in:disabled + label {\n",
" color: var(--xr-font-color2);\n",
"}\n",
"\n",
".xr-section-summary-in + label:before {\n",
" display: inline-block;\n",
" content: '►';\n",
" font-size: 11px;\n",
" width: 15px;\n",
" text-align: center;\n",
"}\n",
"\n",
".xr-section-summary-in:disabled + label:before {\n",
" color: var(--xr-disabled-color);\n",
"}\n",
"\n",
".xr-section-summary-in:checked + label:before {\n",
" content: '▼';\n",
"}\n",
"\n",
".xr-section-summary-in:checked + label > span {\n",
" display: none;\n",
"}\n",
"\n",
".xr-section-summary,\n",
".xr-section-inline-details {\n",
" padding-top: 4px;\n",
" padding-bottom: 4px;\n",
"}\n",
"\n",
".xr-section-inline-details {\n",
" grid-column: 2 / -1;\n",
"}\n",
"\n",
".xr-section-details {\n",
" display: none;\n",
" grid-column: 1 / -1;\n",
" margin-bottom: 5px;\n",
"}\n",
"\n",
".xr-section-summary-in:checked ~ .xr-section-details {\n",
" display: contents;\n",
"}\n",
"\n",
".xr-array-wrap {\n",
" grid-column: 1 / -1;\n",
" display: grid;\n",
" grid-template-columns: 20px auto;\n",
"}\n",
"\n",
".xr-array-wrap > label {\n",
" grid-column: 1;\n",
" vertical-align: top;\n",
"}\n",
"\n",
".xr-preview {\n",
" color: var(--xr-font-color3);\n",
"}\n",
"\n",
".xr-array-preview,\n",
".xr-array-data {\n",
" padding: 0 5px !important;\n",
" grid-column: 2;\n",
"}\n",
"\n",
".xr-array-data,\n",
".xr-array-in:checked ~ .xr-array-preview {\n",
" display: none;\n",
"}\n",
"\n",
".xr-array-in:checked ~ .xr-array-data,\n",
".xr-array-preview {\n",
" display: inline-block;\n",
"}\n",
"\n",
".xr-dim-list {\n",
" display: inline-block !important;\n",
" list-style: none;\n",
" padding: 0 !important;\n",
" margin: 0;\n",
"}\n",
"\n",
".xr-dim-list li {\n",
" display: inline-block;\n",
" padding: 0;\n",
" margin: 0;\n",
"}\n",
"\n",
".xr-dim-list:before {\n",
" content: '(';\n",
"}\n",
"\n",
".xr-dim-list:after {\n",
" content: ')';\n",
"}\n",
"\n",
".xr-dim-list li:not(:last-child):after {\n",
" content: ',';\n",
" padding-right: 5px;\n",
"}\n",
"\n",
".xr-has-index {\n",
" font-weight: bold;\n",
"}\n",
"\n",
".xr-var-list,\n",
".xr-var-item {\n",
" display: contents;\n",
"}\n",
"\n",
".xr-var-item > div,\n",
".xr-var-item label,\n",
".xr-var-item > .xr-var-name span {\n",
" background-color: var(--xr-background-color-row-even);\n",
" margin-bottom: 0;\n",
"}\n",
"\n",
".xr-var-item > .xr-var-name:hover span {\n",
" padding-right: 5px;\n",
"}\n",
"\n",
".xr-var-list > li:nth-child(odd) > div,\n",
".xr-var-list > li:nth-child(odd) > label,\n",
".xr-var-list > li:nth-child(odd) > .xr-var-name span {\n",
" background-color: var(--xr-background-color-row-odd);\n",
"}\n",
"\n",
".xr-var-name {\n",
" grid-column: 1;\n",
"}\n",
"\n",
".xr-var-dims {\n",
" grid-column: 2;\n",
"}\n",
"\n",
".xr-var-dtype {\n",
" grid-column: 3;\n",
" text-align: right;\n",
" color: var(--xr-font-color2);\n",
"}\n",
"\n",
".xr-var-preview {\n",
" grid-column: 4;\n",
"}\n",
"\n",
".xr-index-preview {\n",
" grid-column: 2 / 5;\n",
" color: var(--xr-font-color2);\n",
"}\n",
"\n",
".xr-var-name,\n",
".xr-var-dims,\n",
".xr-var-dtype,\n",
".xr-preview,\n",
".xr-attrs dt {\n",
" white-space: nowrap;\n",
" overflow: hidden;\n",
" text-overflow: ellipsis;\n",
" padding-right: 10px;\n",
"}\n",
"\n",
".xr-var-name:hover,\n",
".xr-var-dims:hover,\n",
".xr-var-dtype:hover,\n",
".xr-attrs dt:hover {\n",
" overflow: visible;\n",
" width: auto;\n",
" z-index: 1;\n",
"}\n",
"\n",
".xr-var-attrs,\n",
".xr-var-data,\n",
".xr-index-data {\n",
" display: none;\n",
" background-color: var(--xr-background-color) !important;\n",
" padding-bottom: 5px !important;\n",
"}\n",
"\n",
".xr-var-attrs-in:checked ~ .xr-var-attrs,\n",
".xr-var-data-in:checked ~ .xr-var-data,\n",
".xr-index-data-in:checked ~ .xr-index-data {\n",
" display: block;\n",
"}\n",
"\n",
".xr-var-data > table {\n",
" float: right;\n",
"}\n",
"\n",
".xr-var-name span,\n",
".xr-var-data,\n",
".xr-index-name div,\n",
".xr-index-data,\n",
".xr-attrs {\n",
" padding-left: 25px !important;\n",
"}\n",
"\n",
".xr-attrs,\n",
".xr-var-attrs,\n",
".xr-var-data,\n",
".xr-index-data {\n",
" grid-column: 1 / -1;\n",
"}\n",
"\n",
"dl.xr-attrs {\n",
" padding: 0;\n",
" margin: 0;\n",
" display: grid;\n",
" grid-template-columns: 125px auto;\n",
"}\n",
"\n",
".xr-attrs dt,\n",
".xr-attrs dd {\n",
" padding: 0;\n",
" margin: 0;\n",
" float: left;\n",
" padding-right: 10px;\n",
" width: auto;\n",
"}\n",
"\n",
".xr-attrs dt {\n",
" font-weight: normal;\n",
" grid-column: 1;\n",
"}\n",
"\n",
".xr-attrs dt:hover span {\n",
" display: inline-block;\n",
" background: var(--xr-background-color);\n",
" padding-right: 10px;\n",
"}\n",
"\n",
".xr-attrs dd {\n",
" grid-column: 2;\n",
" white-space: pre-wrap;\n",
" word-break: break-all;\n",
"}\n",
"\n",
".xr-icon-database,\n",
".xr-icon-file-text2,\n",
".xr-no-icon {\n",
" display: inline-block;\n",
" vertical-align: middle;\n",
" width: 1em;\n",
" height: 1.5em !important;\n",
" stroke-width: 0;\n",
" stroke: currentColor;\n",
" fill: currentColor;\n",
"}\n",
"</style><pre class='xr-text-repr-fallback'>&lt;xarray.Dataset&gt; Size: 19GB\n",
"Dimensions: (z: 2048, y: 1024, x: 1024)\n",
"Dimensions without coordinates: z, y, x\n",
"Data variables:\n",
" flag (z, y, x) bool 2GB dask.array&lt;chunksize=(256, 1024, 1024), meta=np.ndarray&gt;\n",
" vis (z, y, x) float32 9GB dask.array&lt;chunksize=(256, 1024, 1024), meta=np.ndarray&gt;\n",
" weights (z, y, x) float32 9GB dask.array&lt;chunksize=(256, 1024, 1024), meta=np.ndarray&gt;</pre><div class='xr-wrap' style='display:none'><div class='xr-header'><div class='xr-obj-type'>xarray.Dataset</div></div><ul class='xr-sections'><li class='xr-section-item'><input id='section-39130aa5-b897-4b7e-a5b9-3cc701f60126' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-39130aa5-b897-4b7e-a5b9-3cc701f60126' class='xr-section-summary' title='Expand/collapse section'>Dimensions:</label><div class='xr-section-inline-details'><ul class='xr-dim-list'><li><span>z</span>: 2048</li><li><span>y</span>: 1024</li><li><span>x</span>: 1024</li></ul></div><div class='xr-section-details'></div></li><li class='xr-section-item'><input id='section-752376a3-440b-489e-bb72-2c6400984704' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-752376a3-440b-489e-bb72-2c6400984704' class='xr-section-summary' title='Expand/collapse section'>Coordinates: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-34fc62cd-1bca-4990-be43-9180c4fdcb59' class='xr-section-summary-in' type='checkbox' checked><label for='section-34fc62cd-1bca-4990-be43-9180c4fdcb59' class='xr-section-summary' >Data variables: <span>(3)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>flag</span></div><div class='xr-var-dims'>(z, y, x)</div><div class='xr-var-dtype'>bool</div><div class='xr-var-preview xr-preview'>dask.array&lt;chunksize=(256, 1024, 1024), meta=np.ndarray&gt;</div><input id='attrs-858e2723-75b2-4a8e-8f7d-ac8e212777ba' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-858e2723-75b2-4a8e-8f7d-ac8e212777ba' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-11ee8ac5-e63a-48a1-a109-1ac8789c8d5e' class='xr-var-data-in' type='checkbox'><label for='data-11ee8ac5-e63a-48a1-a109-1ac8789c8d5e' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><table>\n",
" <tr>\n",
" <td>\n",
" <table style=\"border-collapse: collapse;\">\n",
" <thead>\n",
" <tr>\n",
" <td> </td>\n",
" <th> Array </th>\n",
" <th> Chunk </th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" \n",
" <tr>\n",
" <th> Bytes </th>\n",
" <td> 2.00 GiB </td>\n",
" <td> 256.00 MiB </td>\n",
" </tr>\n",
" \n",
" <tr>\n",
" <th> Shape </th>\n",
" <td> (2048, 1024, 1024) </td>\n",
" <td> (256, 1024, 1024) </td>\n",
" </tr>\n",
" <tr>\n",
" <th> Dask graph </th>\n",
" <td colspan=\"2\"> 8 chunks in 2 graph layers </td>\n",
" </tr>\n",
" <tr>\n",
" <th> Data type </th>\n",
" <td colspan=\"2\"> bool numpy.ndarray </td>\n",
" </tr>\n",
" </tbody>\n",
" </table>\n",
" </td>\n",
" <td>\n",
" <svg width=\"190\" height=\"180\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
"\n",
" <!-- Horizontal lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
" <line x1=\"10\" y1=\"60\" x2=\"80\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Vertical lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"10\" y2=\"60\" style=\"stroke-width:2\" />\n",
" <line x1=\"18\" y1=\"8\" x2=\"18\" y2=\"68\" />\n",
" <line x1=\"27\" y1=\"17\" x2=\"27\" y2=\"77\" />\n",
" <line x1=\"36\" y1=\"26\" x2=\"36\" y2=\"86\" />\n",
" <line x1=\"45\" y1=\"35\" x2=\"45\" y2=\"95\" />\n",
" <line x1=\"54\" y1=\"44\" x2=\"54\" y2=\"104\" />\n",
" <line x1=\"62\" y1=\"52\" x2=\"62\" y2=\"112\" />\n",
" <line x1=\"71\" y1=\"61\" x2=\"71\" y2=\"121\" />\n",
" <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Colored Rectangle -->\n",
" <polygon points=\"10.0,0.0 80.58823529411765,70.58823529411765 80.58823529411765,130.58823529411765 10.0,60.0\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
"\n",
" <!-- Horizontal lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"70\" y2=\"0\" style=\"stroke-width:2\" />\n",
" <line x1=\"18\" y1=\"8\" x2=\"78\" y2=\"8\" />\n",
" <line x1=\"27\" y1=\"17\" x2=\"87\" y2=\"17\" />\n",
" <line x1=\"36\" y1=\"26\" x2=\"96\" y2=\"26\" />\n",
" <line x1=\"45\" y1=\"35\" x2=\"105\" y2=\"35\" />\n",
" <line x1=\"54\" y1=\"44\" x2=\"114\" y2=\"44\" />\n",
" <line x1=\"62\" y1=\"52\" x2=\"122\" y2=\"52\" />\n",
" <line x1=\"71\" y1=\"61\" x2=\"131\" y2=\"61\" />\n",
" <line x1=\"80\" y1=\"70\" x2=\"140\" y2=\"70\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Vertical lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
" <line x1=\"70\" y1=\"0\" x2=\"140\" y2=\"70\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Colored Rectangle -->\n",
" <polygon points=\"10.0,0.0 70.0,0.0 140.58823529411765,70.58823529411765 80.58823529411765,70.58823529411765\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
"\n",
" <!-- Horizontal lines -->\n",
" <line x1=\"80\" y1=\"70\" x2=\"140\" y2=\"70\" style=\"stroke-width:2\" />\n",
" <line x1=\"80\" y1=\"130\" x2=\"140\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Vertical lines -->\n",
" <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"130\" style=\"stroke-width:2\" />\n",
" <line x1=\"140\" y1=\"70\" x2=\"140\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Colored Rectangle -->\n",
" <polygon points=\"80.58823529411765,70.58823529411765 140.58823529411765,70.58823529411765 140.58823529411765,130.58823529411765 80.58823529411765,130.58823529411765\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
"\n",
" <!-- Text -->\n",
" <text x=\"110.588235\" y=\"150.588235\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >1024</text>\n",
" <text x=\"160.588235\" y=\"100.588235\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,160.588235,100.588235)\">1024</text>\n",
" <text x=\"35.294118\" y=\"115.294118\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(45,35.294118,115.294118)\">2048</text>\n",
"</svg>\n",
" </td>\n",
" </tr>\n",
"</table></div></li><li class='xr-var-item'><div class='xr-var-name'><span>vis</span></div><div class='xr-var-dims'>(z, y, x)</div><div class='xr-var-dtype'>float32</div><div class='xr-var-preview xr-preview'>dask.array&lt;chunksize=(256, 1024, 1024), meta=np.ndarray&gt;</div><input id='attrs-cc12d281-416a-4af8-a4da-e1d89d3bed1c' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-cc12d281-416a-4af8-a4da-e1d89d3bed1c' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-43fd465d-a757-4e02-9e84-0ee1ea214254' class='xr-var-data-in' type='checkbox'><label for='data-43fd465d-a757-4e02-9e84-0ee1ea214254' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><table>\n",
" <tr>\n",
" <td>\n",
" <table style=\"border-collapse: collapse;\">\n",
" <thead>\n",
" <tr>\n",
" <td> </td>\n",
" <th> Array </th>\n",
" <th> Chunk </th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" \n",
" <tr>\n",
" <th> Bytes </th>\n",
" <td> 8.00 GiB </td>\n",
" <td> 1.00 GiB </td>\n",
" </tr>\n",
" \n",
" <tr>\n",
" <th> Shape </th>\n",
" <td> (2048, 1024, 1024) </td>\n",
" <td> (256, 1024, 1024) </td>\n",
" </tr>\n",
" <tr>\n",
" <th> Dask graph </th>\n",
" <td colspan=\"2\"> 8 chunks in 2 graph layers </td>\n",
" </tr>\n",
" <tr>\n",
" <th> Data type </th>\n",
" <td colspan=\"2\"> float32 numpy.ndarray </td>\n",
" </tr>\n",
" </tbody>\n",
" </table>\n",
" </td>\n",
" <td>\n",
" <svg width=\"190\" height=\"180\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
"\n",
" <!-- Horizontal lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
" <line x1=\"10\" y1=\"60\" x2=\"80\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Vertical lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"10\" y2=\"60\" style=\"stroke-width:2\" />\n",
" <line x1=\"18\" y1=\"8\" x2=\"18\" y2=\"68\" />\n",
" <line x1=\"27\" y1=\"17\" x2=\"27\" y2=\"77\" />\n",
" <line x1=\"36\" y1=\"26\" x2=\"36\" y2=\"86\" />\n",
" <line x1=\"45\" y1=\"35\" x2=\"45\" y2=\"95\" />\n",
" <line x1=\"54\" y1=\"44\" x2=\"54\" y2=\"104\" />\n",
" <line x1=\"62\" y1=\"52\" x2=\"62\" y2=\"112\" />\n",
" <line x1=\"71\" y1=\"61\" x2=\"71\" y2=\"121\" />\n",
" <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Colored Rectangle -->\n",
" <polygon points=\"10.0,0.0 80.58823529411765,70.58823529411765 80.58823529411765,130.58823529411765 10.0,60.0\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
"\n",
" <!-- Horizontal lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"70\" y2=\"0\" style=\"stroke-width:2\" />\n",
" <line x1=\"18\" y1=\"8\" x2=\"78\" y2=\"8\" />\n",
" <line x1=\"27\" y1=\"17\" x2=\"87\" y2=\"17\" />\n",
" <line x1=\"36\" y1=\"26\" x2=\"96\" y2=\"26\" />\n",
" <line x1=\"45\" y1=\"35\" x2=\"105\" y2=\"35\" />\n",
" <line x1=\"54\" y1=\"44\" x2=\"114\" y2=\"44\" />\n",
" <line x1=\"62\" y1=\"52\" x2=\"122\" y2=\"52\" />\n",
" <line x1=\"71\" y1=\"61\" x2=\"131\" y2=\"61\" />\n",
" <line x1=\"80\" y1=\"70\" x2=\"140\" y2=\"70\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Vertical lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
" <line x1=\"70\" y1=\"0\" x2=\"140\" y2=\"70\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Colored Rectangle -->\n",
" <polygon points=\"10.0,0.0 70.0,0.0 140.58823529411765,70.58823529411765 80.58823529411765,70.58823529411765\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
"\n",
" <!-- Horizontal lines -->\n",
" <line x1=\"80\" y1=\"70\" x2=\"140\" y2=\"70\" style=\"stroke-width:2\" />\n",
" <line x1=\"80\" y1=\"130\" x2=\"140\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Vertical lines -->\n",
" <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"130\" style=\"stroke-width:2\" />\n",
" <line x1=\"140\" y1=\"70\" x2=\"140\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Colored Rectangle -->\n",
" <polygon points=\"80.58823529411765,70.58823529411765 140.58823529411765,70.58823529411765 140.58823529411765,130.58823529411765 80.58823529411765,130.58823529411765\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
"\n",
" <!-- Text -->\n",
" <text x=\"110.588235\" y=\"150.588235\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >1024</text>\n",
" <text x=\"160.588235\" y=\"100.588235\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,160.588235,100.588235)\">1024</text>\n",
" <text x=\"35.294118\" y=\"115.294118\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(45,35.294118,115.294118)\">2048</text>\n",
"</svg>\n",
" </td>\n",
" </tr>\n",
"</table></div></li><li class='xr-var-item'><div class='xr-var-name'><span>weights</span></div><div class='xr-var-dims'>(z, y, x)</div><div class='xr-var-dtype'>float32</div><div class='xr-var-preview xr-preview'>dask.array&lt;chunksize=(256, 1024, 1024), meta=np.ndarray&gt;</div><input id='attrs-4e5851b3-0cb0-4c5b-b4ac-900cf29042f9' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-4e5851b3-0cb0-4c5b-b4ac-900cf29042f9' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-1c9acc20-e670-4cfb-9fbd-be4ec5a958b1' class='xr-var-data-in' type='checkbox'><label for='data-1c9acc20-e670-4cfb-9fbd-be4ec5a958b1' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><table>\n",
" <tr>\n",
" <td>\n",
" <table style=\"border-collapse: collapse;\">\n",
" <thead>\n",
" <tr>\n",
" <td> </td>\n",
" <th> Array </th>\n",
" <th> Chunk </th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" \n",
" <tr>\n",
" <th> Bytes </th>\n",
" <td> 8.00 GiB </td>\n",
" <td> 1.00 GiB </td>\n",
" </tr>\n",
" \n",
" <tr>\n",
" <th> Shape </th>\n",
" <td> (2048, 1024, 1024) </td>\n",
" <td> (256, 1024, 1024) </td>\n",
" </tr>\n",
" <tr>\n",
" <th> Dask graph </th>\n",
" <td colspan=\"2\"> 8 chunks in 2 graph layers </td>\n",
" </tr>\n",
" <tr>\n",
" <th> Data type </th>\n",
" <td colspan=\"2\"> float32 numpy.ndarray </td>\n",
" </tr>\n",
" </tbody>\n",
" </table>\n",
" </td>\n",
" <td>\n",
" <svg width=\"190\" height=\"180\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
"\n",
" <!-- Horizontal lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
" <line x1=\"10\" y1=\"60\" x2=\"80\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Vertical lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"10\" y2=\"60\" style=\"stroke-width:2\" />\n",
" <line x1=\"18\" y1=\"8\" x2=\"18\" y2=\"68\" />\n",
" <line x1=\"27\" y1=\"17\" x2=\"27\" y2=\"77\" />\n",
" <line x1=\"36\" y1=\"26\" x2=\"36\" y2=\"86\" />\n",
" <line x1=\"45\" y1=\"35\" x2=\"45\" y2=\"95\" />\n",
" <line x1=\"54\" y1=\"44\" x2=\"54\" y2=\"104\" />\n",
" <line x1=\"62\" y1=\"52\" x2=\"62\" y2=\"112\" />\n",
" <line x1=\"71\" y1=\"61\" x2=\"71\" y2=\"121\" />\n",
" <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Colored Rectangle -->\n",
" <polygon points=\"10.0,0.0 80.58823529411765,70.58823529411765 80.58823529411765,130.58823529411765 10.0,60.0\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
"\n",
" <!-- Horizontal lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"70\" y2=\"0\" style=\"stroke-width:2\" />\n",
" <line x1=\"18\" y1=\"8\" x2=\"78\" y2=\"8\" />\n",
" <line x1=\"27\" y1=\"17\" x2=\"87\" y2=\"17\" />\n",
" <line x1=\"36\" y1=\"26\" x2=\"96\" y2=\"26\" />\n",
" <line x1=\"45\" y1=\"35\" x2=\"105\" y2=\"35\" />\n",
" <line x1=\"54\" y1=\"44\" x2=\"114\" y2=\"44\" />\n",
" <line x1=\"62\" y1=\"52\" x2=\"122\" y2=\"52\" />\n",
" <line x1=\"71\" y1=\"61\" x2=\"131\" y2=\"61\" />\n",
" <line x1=\"80\" y1=\"70\" x2=\"140\" y2=\"70\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Vertical lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
" <line x1=\"70\" y1=\"0\" x2=\"140\" y2=\"70\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Colored Rectangle -->\n",
" <polygon points=\"10.0,0.0 70.0,0.0 140.58823529411765,70.58823529411765 80.58823529411765,70.58823529411765\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
"\n",
" <!-- Horizontal lines -->\n",
" <line x1=\"80\" y1=\"70\" x2=\"140\" y2=\"70\" style=\"stroke-width:2\" />\n",
" <line x1=\"80\" y1=\"130\" x2=\"140\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Vertical lines -->\n",
" <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"130\" style=\"stroke-width:2\" />\n",
" <line x1=\"140\" y1=\"70\" x2=\"140\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Colored Rectangle -->\n",
" <polygon points=\"80.58823529411765,70.58823529411765 140.58823529411765,70.58823529411765 140.58823529411765,130.58823529411765 80.58823529411765,130.58823529411765\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
"\n",
" <!-- Text -->\n",
" <text x=\"110.588235\" y=\"150.588235\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >1024</text>\n",
" <text x=\"160.588235\" y=\"100.588235\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,160.588235,100.588235)\">1024</text>\n",
" <text x=\"35.294118\" y=\"115.294118\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(45,35.294118,115.294118)\">2048</text>\n",
"</svg>\n",
" </td>\n",
" </tr>\n",
"</table></div></li></ul></div></li><li class='xr-section-item'><input id='section-eb91d25a-339c-444f-8a9b-58f3e519199b' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-eb91d25a-339c-444f-8a9b-58f3e519199b' class='xr-section-summary' title='Expand/collapse section'>Indexes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-f9f59948-62ea-48dd-a2b6-4787ed8cac71' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-f9f59948-62ea-48dd-a2b6-4787ed8cac71' class='xr-section-summary' title='Expand/collapse section'>Attributes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><dl class='xr-attrs'></dl></div></li></ul></div></div>"
],
"text/plain": [
"<xarray.Dataset> Size: 19GB\n",
"Dimensions: (z: 2048, y: 1024, x: 1024)\n",
"Dimensions without coordinates: z, y, x\n",
"Data variables:\n",
" flag (z, y, x) bool 2GB dask.array<chunksize=(256, 1024, 1024), meta=np.ndarray>\n",
" vis (z, y, x) float32 9GB dask.array<chunksize=(256, 1024, 1024), meta=np.ndarray>\n",
" weights (z, y, x) float32 9GB dask.array<chunksize=(256, 1024, 1024), meta=np.ndarray>"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"xds"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<div><svg style=\"position: absolute; width: 0; height: 0; overflow: hidden\">\n",
"<defs>\n",
"<symbol id=\"icon-database\" viewBox=\"0 0 32 32\">\n",
"<path d=\"M16 0c-8.837 0-16 2.239-16 5v4c0 2.761 7.163 5 16 5s16-2.239 16-5v-4c0-2.761-7.163-5-16-5z\"></path>\n",
"<path d=\"M16 17c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n",
"<path d=\"M16 26c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n",
"</symbol>\n",
"<symbol id=\"icon-file-text2\" viewBox=\"0 0 32 32\">\n",
"<path d=\"M28.681 7.159c-0.694-0.947-1.662-2.053-2.724-3.116s-2.169-2.030-3.116-2.724c-1.612-1.182-2.393-1.319-2.841-1.319h-15.5c-1.378 0-2.5 1.121-2.5 2.5v27c0 1.378 1.122 2.5 2.5 2.5h23c1.378 0 2.5-1.122 2.5-2.5v-19.5c0-0.448-0.137-1.23-1.319-2.841zM24.543 5.457c0.959 0.959 1.712 1.825 2.268 2.543h-4.811v-4.811c0.718 0.556 1.584 1.309 2.543 2.268zM28 29.5c0 0.271-0.229 0.5-0.5 0.5h-23c-0.271 0-0.5-0.229-0.5-0.5v-27c0-0.271 0.229-0.5 0.5-0.5 0 0 15.499-0 15.5 0v7c0 0.552 0.448 1 1 1h7v19.5z\"></path>\n",
"<path d=\"M23 26h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
"<path d=\"M23 22h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
"<path d=\"M23 18h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
"</symbol>\n",
"</defs>\n",
"</svg>\n",
"<style>/* CSS stylesheet for displaying xarray objects in jupyterlab.\n",
" *\n",
" */\n",
"\n",
":root {\n",
" --xr-font-color0: var(--jp-content-font-color0, rgba(0, 0, 0, 1));\n",
" --xr-font-color2: var(--jp-content-font-color2, rgba(0, 0, 0, 0.54));\n",
" --xr-font-color3: var(--jp-content-font-color3, rgba(0, 0, 0, 0.38));\n",
" --xr-border-color: var(--jp-border-color2, #e0e0e0);\n",
" --xr-disabled-color: var(--jp-layout-color3, #bdbdbd);\n",
" --xr-background-color: var(--jp-layout-color0, white);\n",
" --xr-background-color-row-even: var(--jp-layout-color1, white);\n",
" --xr-background-color-row-odd: var(--jp-layout-color2, #eeeeee);\n",
"}\n",
"\n",
"html[theme=dark],\n",
"html[data-theme=dark],\n",
"body[data-theme=dark],\n",
"body.vscode-dark {\n",
" --xr-font-color0: rgba(255, 255, 255, 1);\n",
" --xr-font-color2: rgba(255, 255, 255, 0.54);\n",
" --xr-font-color3: rgba(255, 255, 255, 0.38);\n",
" --xr-border-color: #1F1F1F;\n",
" --xr-disabled-color: #515151;\n",
" --xr-background-color: #111111;\n",
" --xr-background-color-row-even: #111111;\n",
" --xr-background-color-row-odd: #313131;\n",
"}\n",
"\n",
".xr-wrap {\n",
" display: block !important;\n",
" min-width: 300px;\n",
" max-width: 700px;\n",
"}\n",
"\n",
".xr-text-repr-fallback {\n",
" /* fallback to plain text repr when CSS is not injected (untrusted notebook) */\n",
" display: none;\n",
"}\n",
"\n",
".xr-header {\n",
" padding-top: 6px;\n",
" padding-bottom: 6px;\n",
" margin-bottom: 4px;\n",
" border-bottom: solid 1px var(--xr-border-color);\n",
"}\n",
"\n",
".xr-header > div,\n",
".xr-header > ul {\n",
" display: inline;\n",
" margin-top: 0;\n",
" margin-bottom: 0;\n",
"}\n",
"\n",
".xr-obj-type,\n",
".xr-array-name {\n",
" margin-left: 2px;\n",
" margin-right: 10px;\n",
"}\n",
"\n",
".xr-obj-type {\n",
" color: var(--xr-font-color2);\n",
"}\n",
"\n",
".xr-sections {\n",
" padding-left: 0 !important;\n",
" display: grid;\n",
" grid-template-columns: 150px auto auto 1fr 20px 20px;\n",
"}\n",
"\n",
".xr-section-item {\n",
" display: contents;\n",
"}\n",
"\n",
".xr-section-item input {\n",
" display: none;\n",
"}\n",
"\n",
".xr-section-item input + label {\n",
" color: var(--xr-disabled-color);\n",
"}\n",
"\n",
".xr-section-item input:enabled + label {\n",
" cursor: pointer;\n",
" color: var(--xr-font-color2);\n",
"}\n",
"\n",
".xr-section-item input:enabled + label:hover {\n",
" color: var(--xr-font-color0);\n",
"}\n",
"\n",
".xr-section-summary {\n",
" grid-column: 1;\n",
" color: var(--xr-font-color2);\n",
" font-weight: 500;\n",
"}\n",
"\n",
".xr-section-summary > span {\n",
" display: inline-block;\n",
" padding-left: 0.5em;\n",
"}\n",
"\n",
".xr-section-summary-in:disabled + label {\n",
" color: var(--xr-font-color2);\n",
"}\n",
"\n",
".xr-section-summary-in + label:before {\n",
" display: inline-block;\n",
" content: '►';\n",
" font-size: 11px;\n",
" width: 15px;\n",
" text-align: center;\n",
"}\n",
"\n",
".xr-section-summary-in:disabled + label:before {\n",
" color: var(--xr-disabled-color);\n",
"}\n",
"\n",
".xr-section-summary-in:checked + label:before {\n",
" content: '▼';\n",
"}\n",
"\n",
".xr-section-summary-in:checked + label > span {\n",
" display: none;\n",
"}\n",
"\n",
".xr-section-summary,\n",
".xr-section-inline-details {\n",
" padding-top: 4px;\n",
" padding-bottom: 4px;\n",
"}\n",
"\n",
".xr-section-inline-details {\n",
" grid-column: 2 / -1;\n",
"}\n",
"\n",
".xr-section-details {\n",
" display: none;\n",
" grid-column: 1 / -1;\n",
" margin-bottom: 5px;\n",
"}\n",
"\n",
".xr-section-summary-in:checked ~ .xr-section-details {\n",
" display: contents;\n",
"}\n",
"\n",
".xr-array-wrap {\n",
" grid-column: 1 / -1;\n",
" display: grid;\n",
" grid-template-columns: 20px auto;\n",
"}\n",
"\n",
".xr-array-wrap > label {\n",
" grid-column: 1;\n",
" vertical-align: top;\n",
"}\n",
"\n",
".xr-preview {\n",
" color: var(--xr-font-color3);\n",
"}\n",
"\n",
".xr-array-preview,\n",
".xr-array-data {\n",
" padding: 0 5px !important;\n",
" grid-column: 2;\n",
"}\n",
"\n",
".xr-array-data,\n",
".xr-array-in:checked ~ .xr-array-preview {\n",
" display: none;\n",
"}\n",
"\n",
".xr-array-in:checked ~ .xr-array-data,\n",
".xr-array-preview {\n",
" display: inline-block;\n",
"}\n",
"\n",
".xr-dim-list {\n",
" display: inline-block !important;\n",
" list-style: none;\n",
" padding: 0 !important;\n",
" margin: 0;\n",
"}\n",
"\n",
".xr-dim-list li {\n",
" display: inline-block;\n",
" padding: 0;\n",
" margin: 0;\n",
"}\n",
"\n",
".xr-dim-list:before {\n",
" content: '(';\n",
"}\n",
"\n",
".xr-dim-list:after {\n",
" content: ')';\n",
"}\n",
"\n",
".xr-dim-list li:not(:last-child):after {\n",
" content: ',';\n",
" padding-right: 5px;\n",
"}\n",
"\n",
".xr-has-index {\n",
" font-weight: bold;\n",
"}\n",
"\n",
".xr-var-list,\n",
".xr-var-item {\n",
" display: contents;\n",
"}\n",
"\n",
".xr-var-item > div,\n",
".xr-var-item label,\n",
".xr-var-item > .xr-var-name span {\n",
" background-color: var(--xr-background-color-row-even);\n",
" margin-bottom: 0;\n",
"}\n",
"\n",
".xr-var-item > .xr-var-name:hover span {\n",
" padding-right: 5px;\n",
"}\n",
"\n",
".xr-var-list > li:nth-child(odd) > div,\n",
".xr-var-list > li:nth-child(odd) > label,\n",
".xr-var-list > li:nth-child(odd) > .xr-var-name span {\n",
" background-color: var(--xr-background-color-row-odd);\n",
"}\n",
"\n",
".xr-var-name {\n",
" grid-column: 1;\n",
"}\n",
"\n",
".xr-var-dims {\n",
" grid-column: 2;\n",
"}\n",
"\n",
".xr-var-dtype {\n",
" grid-column: 3;\n",
" text-align: right;\n",
" color: var(--xr-font-color2);\n",
"}\n",
"\n",
".xr-var-preview {\n",
" grid-column: 4;\n",
"}\n",
"\n",
".xr-index-preview {\n",
" grid-column: 2 / 5;\n",
" color: var(--xr-font-color2);\n",
"}\n",
"\n",
".xr-var-name,\n",
".xr-var-dims,\n",
".xr-var-dtype,\n",
".xr-preview,\n",
".xr-attrs dt {\n",
" white-space: nowrap;\n",
" overflow: hidden;\n",
" text-overflow: ellipsis;\n",
" padding-right: 10px;\n",
"}\n",
"\n",
".xr-var-name:hover,\n",
".xr-var-dims:hover,\n",
".xr-var-dtype:hover,\n",
".xr-attrs dt:hover {\n",
" overflow: visible;\n",
" width: auto;\n",
" z-index: 1;\n",
"}\n",
"\n",
".xr-var-attrs,\n",
".xr-var-data,\n",
".xr-index-data {\n",
" display: none;\n",
" background-color: var(--xr-background-color) !important;\n",
" padding-bottom: 5px !important;\n",
"}\n",
"\n",
".xr-var-attrs-in:checked ~ .xr-var-attrs,\n",
".xr-var-data-in:checked ~ .xr-var-data,\n",
".xr-index-data-in:checked ~ .xr-index-data {\n",
" display: block;\n",
"}\n",
"\n",
".xr-var-data > table {\n",
" float: right;\n",
"}\n",
"\n",
".xr-var-name span,\n",
".xr-var-data,\n",
".xr-index-name div,\n",
".xr-index-data,\n",
".xr-attrs {\n",
" padding-left: 25px !important;\n",
"}\n",
"\n",
".xr-attrs,\n",
".xr-var-attrs,\n",
".xr-var-data,\n",
".xr-index-data {\n",
" grid-column: 1 / -1;\n",
"}\n",
"\n",
"dl.xr-attrs {\n",
" padding: 0;\n",
" margin: 0;\n",
" display: grid;\n",
" grid-template-columns: 125px auto;\n",
"}\n",
"\n",
".xr-attrs dt,\n",
".xr-attrs dd {\n",
" padding: 0;\n",
" margin: 0;\n",
" float: left;\n",
" padding-right: 10px;\n",
" width: auto;\n",
"}\n",
"\n",
".xr-attrs dt {\n",
" font-weight: normal;\n",
" grid-column: 1;\n",
"}\n",
"\n",
".xr-attrs dt:hover span {\n",
" display: inline-block;\n",
" background: var(--xr-background-color);\n",
" padding-right: 10px;\n",
"}\n",
"\n",
".xr-attrs dd {\n",
" grid-column: 2;\n",
" white-space: pre-wrap;\n",
" word-break: break-all;\n",
"}\n",
"\n",
".xr-icon-database,\n",
".xr-icon-file-text2,\n",
".xr-no-icon {\n",
" display: inline-block;\n",
" vertical-align: middle;\n",
" width: 1em;\n",
" height: 1.5em !important;\n",
" stroke-width: 0;\n",
" stroke: currentColor;\n",
" fill: currentColor;\n",
"}\n",
"</style><pre class='xr-text-repr-fallback'>&lt;xarray.DataArray &#x27;vis&#x27; (z: 2048, y: 1024, x: 1024)&gt; Size: 9GB\n",
"dask.array&lt;open_dataset-vis, shape=(2048, 1024, 1024), dtype=float32, chunksize=(256, 1024, 1024), chunktype=numpy.ndarray&gt;\n",
"Dimensions without coordinates: z, y, x</pre><div class='xr-wrap' style='display:none'><div class='xr-header'><div class='xr-obj-type'>xarray.DataArray</div><div class='xr-array-name'>'vis'</div><ul class='xr-dim-list'><li><span>z</span>: 2048</li><li><span>y</span>: 1024</li><li><span>x</span>: 1024</li></ul></div><ul class='xr-sections'><li class='xr-section-item'><div class='xr-array-wrap'><input id='section-55f68d36-46e3-4461-ab06-83784aaa2fd2' class='xr-array-in' type='checkbox' checked><label for='section-55f68d36-46e3-4461-ab06-83784aaa2fd2' title='Show/hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-array-preview xr-preview'><span>dask.array&lt;chunksize=(256, 1024, 1024), meta=np.ndarray&gt;</span></div><div class='xr-array-data'><table>\n",
" <tr>\n",
" <td>\n",
" <table style=\"border-collapse: collapse;\">\n",
" <thead>\n",
" <tr>\n",
" <td> </td>\n",
" <th> Array </th>\n",
" <th> Chunk </th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" \n",
" <tr>\n",
" <th> Bytes </th>\n",
" <td> 8.00 GiB </td>\n",
" <td> 1.00 GiB </td>\n",
" </tr>\n",
" \n",
" <tr>\n",
" <th> Shape </th>\n",
" <td> (2048, 1024, 1024) </td>\n",
" <td> (256, 1024, 1024) </td>\n",
" </tr>\n",
" <tr>\n",
" <th> Dask graph </th>\n",
" <td colspan=\"2\"> 8 chunks in 2 graph layers </td>\n",
" </tr>\n",
" <tr>\n",
" <th> Data type </th>\n",
" <td colspan=\"2\"> float32 numpy.ndarray </td>\n",
" </tr>\n",
" </tbody>\n",
" </table>\n",
" </td>\n",
" <td>\n",
" <svg width=\"190\" height=\"180\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
"\n",
" <!-- Horizontal lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
" <line x1=\"10\" y1=\"60\" x2=\"80\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Vertical lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"10\" y2=\"60\" style=\"stroke-width:2\" />\n",
" <line x1=\"18\" y1=\"8\" x2=\"18\" y2=\"68\" />\n",
" <line x1=\"27\" y1=\"17\" x2=\"27\" y2=\"77\" />\n",
" <line x1=\"36\" y1=\"26\" x2=\"36\" y2=\"86\" />\n",
" <line x1=\"45\" y1=\"35\" x2=\"45\" y2=\"95\" />\n",
" <line x1=\"54\" y1=\"44\" x2=\"54\" y2=\"104\" />\n",
" <line x1=\"62\" y1=\"52\" x2=\"62\" y2=\"112\" />\n",
" <line x1=\"71\" y1=\"61\" x2=\"71\" y2=\"121\" />\n",
" <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Colored Rectangle -->\n",
" <polygon points=\"10.0,0.0 80.58823529411765,70.58823529411765 80.58823529411765,130.58823529411765 10.0,60.0\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
"\n",
" <!-- Horizontal lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"70\" y2=\"0\" style=\"stroke-width:2\" />\n",
" <line x1=\"18\" y1=\"8\" x2=\"78\" y2=\"8\" />\n",
" <line x1=\"27\" y1=\"17\" x2=\"87\" y2=\"17\" />\n",
" <line x1=\"36\" y1=\"26\" x2=\"96\" y2=\"26\" />\n",
" <line x1=\"45\" y1=\"35\" x2=\"105\" y2=\"35\" />\n",
" <line x1=\"54\" y1=\"44\" x2=\"114\" y2=\"44\" />\n",
" <line x1=\"62\" y1=\"52\" x2=\"122\" y2=\"52\" />\n",
" <line x1=\"71\" y1=\"61\" x2=\"131\" y2=\"61\" />\n",
" <line x1=\"80\" y1=\"70\" x2=\"140\" y2=\"70\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Vertical lines -->\n",
" <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
" <line x1=\"70\" y1=\"0\" x2=\"140\" y2=\"70\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Colored Rectangle -->\n",
" <polygon points=\"10.0,0.0 70.0,0.0 140.58823529411765,70.58823529411765 80.58823529411765,70.58823529411765\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
"\n",
" <!-- Horizontal lines -->\n",
" <line x1=\"80\" y1=\"70\" x2=\"140\" y2=\"70\" style=\"stroke-width:2\" />\n",
" <line x1=\"80\" y1=\"130\" x2=\"140\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Vertical lines -->\n",
" <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"130\" style=\"stroke-width:2\" />\n",
" <line x1=\"140\" y1=\"70\" x2=\"140\" y2=\"130\" style=\"stroke-width:2\" />\n",
"\n",
" <!-- Colored Rectangle -->\n",
" <polygon points=\"80.58823529411765,70.58823529411765 140.58823529411765,70.58823529411765 140.58823529411765,130.58823529411765 80.58823529411765,130.58823529411765\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
"\n",
" <!-- Text -->\n",
" <text x=\"110.588235\" y=\"150.588235\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >1024</text>\n",
" <text x=\"160.588235\" y=\"100.588235\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,160.588235,100.588235)\">1024</text>\n",
" <text x=\"35.294118\" y=\"115.294118\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(45,35.294118,115.294118)\">2048</text>\n",
"</svg>\n",
" </td>\n",
" </tr>\n",
"</table></div></div></li><li class='xr-section-item'><input id='section-acd33fa2-3990-4eac-b64a-152704743090' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-acd33fa2-3990-4eac-b64a-152704743090' class='xr-section-summary' title='Expand/collapse section'>Coordinates: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-9a02c254-6dc1-4c7c-91e8-d338450f411f' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-9a02c254-6dc1-4c7c-91e8-d338450f411f' class='xr-section-summary' title='Expand/collapse section'>Indexes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-3b91d017-231d-4a0f-aa7d-2b698d6cd2a0' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-3b91d017-231d-4a0f-aa7d-2b698d6cd2a0' class='xr-section-summary' title='Expand/collapse section'>Attributes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><dl class='xr-attrs'></dl></div></li></ul></div></div>"
],
"text/plain": [
"<xarray.DataArray 'vis' (z: 2048, y: 1024, x: 1024)> Size: 9GB\n",
"dask.array<open_dataset-vis, shape=(2048, 1024, 1024), dtype=float32, chunksize=(256, 1024, 1024), chunktype=numpy.ndarray>\n",
"Dimensions without coordinates: z, y, x"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"xds.vis"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Numpy operation"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"def np_operation(vis, weights, flag):\n",
" output = np.einsum(\"ij,jk->ik\", vis, vis)\n",
" output = output * weights\n",
" output = np.where(flag, output, 0.0)\n",
"\n",
" return output"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"result_np_vec = xr.apply_ufunc(\n",
" np_operation,\n",
" xds.vis,\n",
" xds.weights,\n",
" xds.flag,\n",
" input_core_dims=[['y','x'],['y','x'],['y','x'],],\n",
" output_core_dims=[['y','x']],\n",
" dask=\"parallelized\",\n",
" vectorize=True,\n",
" output_dtypes=[np.float32],\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"1min ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n"
]
}
],
"source": [
"%timeit -n 1 -r 1 result_np_vec.compute()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Jax operation on CPU"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"from functools import partial\n",
"import jax\n",
"import jax.numpy as jnp\n",
"\n",
"@jax.jit\n",
"@partial(jax.vmap, in_axes=(0,0,0), out_axes=0)\n",
"def jax_vmap_cpu_operation(vis, weights, flag):\n",
" output = jnp.einsum(\"ij,jk->ik\", vis, vis)\n",
" output = output * weights\n",
" output = jnp.where(flag, output, 0.0)\n",
"\n",
" return output"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"result_jax_vec = xr.apply_ufunc(\n",
" jax_vmap_cpu_operation,\n",
" xds.vis,\n",
" xds.weights,\n",
" xds.flag,\n",
" input_core_dims=[['y','x'],['y','x'],['y','x'],],\n",
" output_core_dims=[['y','x']],\n",
" dask=\"parallelized\",\n",
" vectorize=False,\n",
" output_dtypes=[np.float32],\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"28.9 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n"
]
}
],
"source": [
"%timeit -n 1 -r 1 result_jax_vec.compute()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# xr.testing.assert_allclose(result_jax_vec, result_np_vec)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Numba operation"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [],
"source": [
"import numba\n",
"\n",
"@numba.guvectorize([(numba.float32[:,::1],\n",
" numba.float32[:,::1],\n",
" numba.float32[:,::1],\n",
" numba.float32[:,::1],)],\n",
" '(ny,nx),(ny,nx),(ny,nx)->(ny,nx)')\n",
"def numba_operation(vis, weights, flag, output):\n",
" inter_out = vis @ vis\n",
" inter_out = inter_out * weights\n",
" inter_out = np.where(flag, inter_out, 0.0)\n",
"\n",
" output[:,:] = inter_out"
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {},
"outputs": [],
"source": [
"result_numba_vec = xr.apply_ufunc(\n",
" numba_operation,\n",
" xds.vis,\n",
" xds.weights,\n",
" xds.flag,\n",
" input_core_dims=[['y','x'],['y','x'],['y','x'],],\n",
" output_core_dims=[['y','x']],\n",
" dask=\"parallelized\",\n",
" vectorize=False,\n",
" output_dtypes=[np.float32],\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"38.7 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n"
]
}
],
"source": [
"%timeit -n 1 -r 1 result_numba_vec.compute()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# xr.testing.assert_allclose(result_jax_vec, result_numba_vec)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "spec_line",
"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.10.15"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Display the source blob
Display the rendered blob
Raw
vectorized_functions_with_dask_gpu.ipynb
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment