RL-w-fft-comparison.ipynb

RL-w-fft-w-dask.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "bokeh.server.protocol_handler - ERROR - error handling message\n",
      " message: Message 'PATCH-DOC' content: {'references': [], 'events': [{'kind': 'ModelChanged', 'model': {'id': '1159'}, 'attr': 'text_align', 'new': {'value': 'right'}}, {'kind': 'ModelChanged', 'model': {'id': '1160'}, 'attr': 'text_align', 'new': {'value': 'right'}}, {'kind': 'ModelChanged', 'model': {'id': '1104'}, 'attr': 'reset_start', 'new': -8}, {'kind': 'ModelChanged', 'model': {'id': '1104'}, 'attr': 'reset_end', 'new': 0}]} \n",
      " error: DeserializationError(\"Enum('left', 'right', 'center') expected str, got {'value': 'right'} of type dict\")\n",
      "Traceback (most recent call last):\n",
      "  File \"/datasets/bzaitlen/miniconda3/envs/20200616/lib/python3.7/site-packages/bokeh/server/protocol_handler.py\", line 90, in handle\n",
      "    work = await handler(message, connection)\n",
      "  File \"/datasets/bzaitlen/miniconda3/envs/20200616/lib/python3.7/site-packages/bokeh/server/session.py\", line 67, in _needs_document_lock_wrapper\n",
      "    result = func(self, *args, **kwargs)\n",
      "  File \"/datasets/bzaitlen/miniconda3/envs/20200616/lib/python3.7/site-packages/bokeh/server/session.py\", line 261, in _handle_patch\n",
      "    message.apply_to_document(self.document, self)\n",
      "  File \"/datasets/bzaitlen/miniconda3/envs/20200616/lib/python3.7/site-packages/bokeh/protocol/messages/patch_doc.py\", line 100, in apply_to_document\n",
      "    doc._with_self_as_curdoc(lambda: doc.apply_json_patch(self.content, setter))\n",
      "  File \"/datasets/bzaitlen/miniconda3/envs/20200616/lib/python3.7/site-packages/bokeh/document/document.py\", line 1150, in _with_self_as_curdoc\n",
      "    return f()\n",
      "  File \"/datasets/bzaitlen/miniconda3/envs/20200616/lib/python3.7/site-packages/bokeh/protocol/messages/patch_doc.py\", line 100, in <lambda>\n",
      "    doc._with_self_as_curdoc(lambda: doc.apply_json_patch(self.content, setter))\n",
      "  File \"/datasets/bzaitlen/miniconda3/envs/20200616/lib/python3.7/site-packages/bokeh/document/document.py\", line 411, in apply_json_patch\n",
      "    patched_obj.set_from_json(attr, value, models=references, setter=setter)\n",
      "  File \"/datasets/bzaitlen/miniconda3/envs/20200616/lib/python3.7/site-packages/bokeh/core/has_props.py\", line 341, in set_from_json\n",
      "    descriptor.set_from_json(self, json, models, setter)\n",
      "  File \"/datasets/bzaitlen/miniconda3/envs/20200616/lib/python3.7/site-packages/bokeh/core/property/descriptors.py\", line 611, in set_from_json\n",
      "    self.property.from_json(json, models),\n",
      "  File \"/datasets/bzaitlen/miniconda3/envs/20200616/lib/python3.7/site-packages/bokeh/core/property/bases.py\", line 463, in from_json\n",
      "    raise DeserializationError(\"%s expected %s, got %s of type %s\" % (self, expected, json, type(json).__name__))\n",
      "bokeh.core.property.bases.DeserializationError: Enum('left', 'right', 'center') expected str, got {'value': 'right'} of type dict\n"
     ]
    }
   ],
   "source": [
    "from dask.distributed import Client\n",
    "from dask_cuda import LocalCUDACluster\n",
    "from dask_cuda.initialize import initialize\n",
    "from dask.distributed import wait\n",
    "\n",
    "import skimage.io\n",
    "\n",
    "import dask.array as da\n",
    "\n",
    "import cupy as cp\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Configurations\n",
    "protocol = \"ucx\"\n",
    "interface = \"enp134s0f1\"    # DGX-2\n",
    "enable_tcp_over_ucx = True\n",
    "enable_nvlink = True\n",
    "enable_infiniband = False\n",
    "\n",
    "initialize(\n",
    "    create_cuda_context=True,\n",
    "    enable_tcp_over_ucx=enable_tcp_over_ucx,\n",
    "    enable_infiniband=enable_infiniband,\n",
    "    enable_nvlink=enable_nvlink,\n",
    ")\n",
    "\n",
    "cluster = LocalCUDACluster(local_directory=\"/tmp/bzaitlen\",   \n",
    "                           protocol=protocol,\n",
    "                           interface=interface,\n",
    "                           enable_tcp_over_ucx=enable_tcp_over_ucx,\n",
    "                           enable_infiniband=enable_infiniband,\n",
    "                           enable_nvlink=enable_nvlink,\n",
    "                           rmm_pool_size=\"26GB\",\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table style=\"border: 2px solid white;\">\n",
       "<tr>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Client</h3>\n",
       "<ul style=\"text-align: left; list-style: none; margin: 0; padding: 0;\">\n",
       "  <li><b>Scheduler: </b>ucx://172.22.1.27:49346</li>\n",
       "  <li><b>Dashboard: </b><a href='http://172.22.1.27:8787/status' target='_blank'>http://172.22.1.27:8787/status</a></li>\n",
       "</ul>\n",
       "</td>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Cluster</h3>\n",
       "<ul style=\"text-align: left; list-style:none; margin: 0; padding: 0;\">\n",
       "  <li><b>Workers: </b>16</li>\n",
       "  <li><b>Cores: </b>16</li>\n",
       "  <li><b>Memory: </b>1.62 TB</li>\n",
       "</ul>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<Client: 'ucx://172.22.1.27:49346' processes=16 threads=16, memory=1.62 TB>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read in data and visualize slice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = skimage.io.imread('/public/NVMICROSCOPY/y1z1_C1_A.tif')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0, 180)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arr = sample[1, :, :]\n",
    "arr.min(), arr.max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(950, 2048, 2048)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/datasets/bzaitlen/miniconda3/envs/20200616/lib/python3.7/site-packages/skimage/io/_plugins/matplotlib_plugin.py:150: UserWarning: Low image data range; displaying image with stretched contrast.\n",
      "  lo, hi, cmap = _get_display_range(image)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f94c4716890>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {