docs/dask_averaging.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Averaging detector data with Dask\n",
    "\n",
    "We often want to average large detector data across trains, keeping the pulses within each train separate, so we have an average image for pulse 0, another for pulse 1, etc.\n",
    "\n",
    "This data may be too big to load into memory at once, but using [Dask](https://dask.org/) we can work with it like a numpy array. Dask takes care of splitting the job up into smaller pieces and assembling the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from karabo_data import open_run\n",
    "\n",
    "import dask.array as da\n",
    "from dask.distributed import Client, progress\n",
    "from dask_jobqueue import SLURMCluster\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we use [Dask-Jobqueue](https://jobqueue.dask.org/en/latest/) to talk to the Maxwell cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a847fc65922b49afbe5ee06080f513c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HTML(value='<h2>SLURMCluster</h2>'), HBox(children=(HTML(value='\\n<div>\\n  <style scoped>\\n    …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "partition = 'exfel'  # For EuXFEL staff\n",
    "#partition = 'upex'   # For users\n",
    "\n",
    "cluster = SLURMCluster(\n",
    "    queue=partition,\n",
    "    # Resources per SLURM job (per node, the way SLURM is configured on Maxwell)\n",
    "    # processes=16 runs 16 Dask workers in a job, so each worker has 1 core & 16 GB RAM.\n",
    "    processes=16, cores=16, memory='256GB',\n",
    ")\n",
    "\n",
    "# Get a notbook widget showing the cluster state\n",
    "cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Submit 2 SLURM jobs, for 32 Dask workers\n",
    "cluster.scale(32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the cluster is busy, you might need to wait a while for the jobs to start.\n",
    "The cluster widget above will update when they're running.\n",
    "\n",
    "Next, we'll set Dask up to use those workers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created dask client: <Client: scheduler='tcp://131.169.193.102:44986' processes=32 cores=32>\n"
     ]
    }
   ],
   "source": [
    "client = Client(cluster)\n",
    "print(\"Created dask client:\", client)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now Dask is ready, let's open the run we're going to operate on:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# of trains:    3299\n",
      "Duration:       0:05:29.800000\n",
      "First train ID: 517617973\n",
      "Last train ID:  517621271\n",
      "\n",
      "16 detector modules (SCS_DET_DSSC1M-1)\n",
      "  e.g. module SCS_DET_DSSC1M-1 0 : 128 x 512 pixels\n",
      "  75 frames per train, 247425 total frames\n",
      "\n",
      "3 instrument sources (excluding detectors):\n",
      "  - SA3_XTD10_XGM/XGM/DOOCS:output\n",
      "  - SCS_BLU_XGM/XGM/DOOCS:output\n",
      "  - SCS_UTC1_ADQ/ADC/1:network\n",
      "\n",
      "20 control sources:\n",
      "  - P_GATT\n",
      "  - SA3_XTD10_MONO/ENC/GRATING_AX\n",
      "  - SA3_XTD10_MONO/MDL/PHOTON_ENERGY\n",
      "  - SA3_XTD10_MONO/MOTOR/GRATINGS_X\n",
      "  - SA3_XTD10_MONO/MOTOR/GRATING_AX\n",
      "  - SA3_XTD10_MONO/MOTOR/HE_PM_X\n",
      "  - SA3_XTD10_MONO/MOTOR/LE_PM_X\n",
      "  - SA3_XTD10_VAC/DCTRL/AR_MODE_OK\n",
      "  - SA3_XTD10_VAC/DCTRL/D12_APERT_IN_OK\n",
      "  - SA3_XTD10_VAC/DCTRL/D6_APERT_IN_OK\n",
      "  - SA3_XTD10_VAC/DCTRL/N2_MODE_OK\n",
      "  - SA3_XTD10_VAC/GAUGE/G30470D_IN\n",
      "  - SA3_XTD10_VAC/GAUGE/G30480D_IN\n",
      "  - SA3_XTD10_VAC/GAUGE/G30490D_IN\n",
      "  - SA3_XTD10_VAC/GAUGE/G30510C\n",
      "  - SA3_XTD10_XGM/XGM/DOOCS\n",
      "  - SCS_BLU_XGM/XGM/DOOCS\n",
      "  - SCS_RR_UTC/MDL/BUNCH_DECODER\n",
      "  - SCS_RR_UTC/TSYS/TIMESERVER\n",
      "  - SCS_UTC1_ADQ/ADC/1\n",
      "\n"
     ]
    }
   ],
   "source": [
    "run = open_run(proposal=2212, run=103)\n",
    "run.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're working with data from the DSSC detector.\n",
    "In this run, it's recording 75 frames for each train - this is part of the info above.\n",
    "\n",
    "Now, we'll define how we're going to average over trains for each module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def average_module(modno, run, pulses_per_train=75):\n",
    "    source = f'SCS_DET_DSSC1M-1/DET/{modno}CH0:xtdf'\n",
    "    counts = run.get_data_counts(source, 'image.data')\n",
    "    \n",
    "    arr = run.get_dask_array(source, 'image.data')\n",
    "    # Make a new dimension for trains\n",
    "    arr_trains = arr.reshape(-1, pulses_per_train, 128, 512)\n",
    "    if modno == 0:\n",
    "        print(\"array shape:\", arr.shape)  # frames, dummy, 128, 512\n",
    "        print(\"Reshaped to:\", arr_trains.shape)\n",
    "\n",
    "    return arr_trains.mean(axis=0, dtype=np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "array shape: (247425, 1, 128, 512)\n",
      "Reshaped to: (3299, 75, 128, 512)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>,\n",
       " dask.array<mean_agg-aggregate, shape=(75, 128, 512), dtype=float32, chunksize=(75, 128, 512)>]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mod_averages = [\n",
    "    average_module(i, run, pulses_per_train=75)\n",
    "    for i in range(16)\n",
    "]\n",
    "\n",
    "mod_averages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr>\n",
       "<td>\n",
       "<table>\n",
       "  <thead>\n",
       "    <tr><td> </td><th> Array </th><th> Chunk </th></tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr><th> Bytes </th><td> 314.57 MB </td> <td> 19.66 MB </td></tr>\n",
       "    <tr><th> Shape </th><td> (16, 75, 128, 512) </td> <td> (1, 75, 128, 512) </td></tr>\n",
       "    <tr><th> Count </th><td> 2560 Tasks </td><td> 16 Chunks </td></tr>\n",
       "    <tr><th> Type </th><td> float32 </td><td> numpy.ndarray </td></tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</td>\n",
       "<td>\n",
       "<svg width=\"395\" height=\"116\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"31\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"0\" y1=\"25\" x2=\"31\" y2=\"25\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"0\" y2=\"25\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"1\" y1=\"0\" x2=\"1\" y2=\"25\" />\n",
       "  <line x1=\"3\" y1=\"0\" x2=\"3\" y2=\"25\" />\n",
       "  <line x1=\"5\" y1=\"0\" x2=\"5\" y2=\"25\" />\n",
       "  <line x1=\"7\" y1=\"0\" x2=\"7\" y2=\"25\" />\n",
       "  <line x1=\"9\" y1=\"0\" x2=\"9\" y2=\"25\" />\n",
       "  <line x1=\"11\" y1=\"0\" x2=\"11\" y2=\"25\" />\n",
       "  <line x1=\"13\" y1=\"0\" x2=\"13\" y2=\"25\" />\n",
       "  <line x1=\"15\" y1=\"0\" x2=\"15\" y2=\"25\" />\n",
       "  <line x1=\"17\" y1=\"0\" x2=\"17\" y2=\"25\" />\n",
       "  <line x1=\"19\" y1=\"0\" x2=\"19\" y2=\"25\" />\n",
       "  <line x1=\"21\" y1=\"0\" x2=\"21\" y2=\"25\" />\n",
       "  <line x1=\"23\" y1=\"0\" x2=\"23\" y2=\"25\" />\n",
       "  <line x1=\"25\" y1=\"0\" x2=\"25\" y2=\"25\" />\n",
       "  <line x1=\"27\" y1=\"0\" x2=\"27\" y2=\"25\" />\n",
       "  <line x1=\"29\" y1=\"0\" x2=\"29\" y2=\"25\" />\n",
       "  <line x1=\"31\" y1=\"0\" x2=\"31\" y2=\"25\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"0.000000,0.000000 31.635229,0.000000 31.635229,25.412617 0.000000,25.412617\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"15.817615\" y=\"45.412617\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >16</text>\n",
       "  <text x=\"51.635229\" y=\"12.706308\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(0,51.635229,12.706308)\">1</text>\n",
       "\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"101\" y1=\"0\" x2=\"124\" y2=\"23\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"101\" y1=\"42\" x2=\"124\" y2=\"66\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"101\" y1=\"0\" x2=\"101\" y2=\"42\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"124\" y1=\"23\" x2=\"124\" y2=\"66\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"101.000000,0.000000 124.877238,23.877238 124.877238,66.776033 101.000000,42.898796\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"101\" y1=\"0\" x2=\"221\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"124\" y1=\"23\" x2=\"244\" y2=\"23\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"101\" y1=\"0\" x2=\"124\" y2=\"23\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"221\" y1=\"0\" x2=\"244\" y2=\"23\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"101.000000,0.000000 221.000000,0.000000 244.877238,23.877238 124.877238,23.877238\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"124\" y1=\"23\" x2=\"244\" y2=\"23\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"124\" y1=\"66\" x2=\"244\" y2=\"66\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"124\" y1=\"23\" x2=\"124\" y2=\"66\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"244\" y1=\"23\" x2=\"244\" y2=\"66\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"124.877238,23.877238 244.877238,23.877238 244.877238,66.776033 124.877238,66.776033\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"184.877238\" y=\"86.776033\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >512</text>\n",
       "  <text x=\"264.877238\" y=\"45.326636\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,264.877238,45.326636)\">128</text>\n",
       "  <text x=\"102.938619\" y=\"74.837414\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(45,102.938619,74.837414)\">75</text>\n",
       "</svg>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<stack, shape=(16, 75, 128, 512), dtype=float32, chunksize=(1, 75, 128, 512)>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Stack the averages into a single array\n",
    "all_average = da.stack(mod_averages)\n",
    "all_average"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Dask shows us what shape the result array will be, but so far, no real computation has happened.\n",
    "Now that we've defined what we want, let's tell Dask to compute it.\n",
    "\n",
    "This will take a minute or two. If you're running it, scroll up to the Dask cluster widget and click the status link to see what it's doing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 20.8 s, sys: 2.6 s, total: 23.4 s\n",
      "Wall time: 1min 42s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "all_average_arr = all_average.compute()  # Get a concrete numpy array for the result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`all_average_arr` is a regular numpy array with our results. Here are the values from the corner of module 0, frame 0:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[48.822674 50.983025 44.953014 44.08245  45.056988]\n",
      " [45.8251   49.183388 46.39982  43.371628 47.53501 ]\n",
      " [51.03395  46.02243  44.92058  50.966656 42.918762]\n",
      " [43.190662 49.961502 44.23007  43.252197 47.663536]\n",
      " [48.844803 51.489845 50.45438  46.305546 47.51258 ]]\n"
     ]
    }
   ],
   "source": [
    "print(all_average_arr[0, 0, :5, :5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please shut down the cluster (or scale it down to 0 workers) if you won't be using it for a while.\n",
    "This releases the resources for other people."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()\n",
    "cluster.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}