docs/source/perf_dask.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a64f4d84",
   "metadata": {},
   "source": [
    "# Performances with dask arrays\n",
    "\n",
    "We compare the performance for automatically generated dask arrays.\n",
    "This is a similar setup to the one used in the performance comparison with [numpy arrays](./perf_numpy.ipynb).\n",
    "Here we do not compare with the numpy histograms function but with `dask-histogram`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4c7de314",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boost_histogram as bh\n",
    "import dask_histogram as dh\n",
    "import dask.array as da\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import xarray_histogram as xh\n",
    "from numpy.testing import assert_allclose\n",
    "from xhistogram.xarray import histogram as xhistogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c8f07323",
   "metadata": {},
   "outputs": [],
   "source": [
    "nbins = (100, 100)\n",
    "ranges = ((-3, 3), (-3, 3))\n",
    "\n",
    "axes = [bh.axis.Regular(nbins[i], *ranges[i]) for i in range(2)]\n",
    "edges = [ax.edges for ax in axes]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9001f824",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(42)\n",
    "\n",
    "chunk_size = 1_000_000\n",
    "vals1d = da.random.normal(size=[10_000_000], chunks=(chunk_size)).astype(np.float32)\n",
    "da_1d = xr.DataArray(vals1d, name='test_1d')\n",
    "\n",
    "vals = da.random.normal(size=[2, 10_000_000], chunks=(1, chunk_size)).astype(np.float32)\n",
    "da_2d = [xr.DataArray(vals[i], name='test_2d_{}'.format(i)) for i in range(2)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9beee267",
   "metadata": {},
   "source": [
    "## One-dimensional histogram (flattened)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "15acf6cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "answer = dh.factory(vals1d, axes=[axes[0]]).values()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d5c6e5b",
   "metadata": {},
   "source": [
    "### dask-histogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e6de7eaa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "50 ms ± 3.2 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "h = dh.factory(vals1d, axes=[axes[0]])\n",
    "assert_allclose(h.values(), answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c82731c2",
   "metadata": {},
   "source": [
    "### xhistogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3626be45",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100 ms ± 2.67 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "assert(not da_1d._in_memory)\n",
    "h = xhistogram(da_1d, bins=edges[0], block_size=chunk_size)\n",
    "assert_allclose(h.values, answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62d7bff0",
   "metadata": {},
   "source": [
    "### xarray-histogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dfe15e06",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "52.1 ms ± 2.17 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "assert(not da_1d._in_memory)\n",
    "h = xh.histogram(da_1d, bins=[axes[0]])\n",
    "assert_allclose(h.values, answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4e5301e",
   "metadata": {},
   "source": [
    "## Two-dimensional (flattened)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bd42d17d",
   "metadata": {},
   "outputs": [],
   "source": [
    "answer = dh.factory(*vals, axes=axes).values()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73533a66",
   "metadata": {},
   "source": [
    "### dask-histogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7780a4f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "76.7 ms ± 3.82 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "h = dh.factory(*vals, axes=axes)\n",
    "assert_allclose(h.values(), answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b3fa37e",
   "metadata": {},
   "source": [
    "### xhistogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7954b1fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "162 ms ± 2.76 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "assert(not da_2d[0]._in_memory)\n",
    "h = xhistogram(*da_2d, bins=edges, block_size=chunk_size)\n",
    "assert_allclose(h.values, answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12348f56",
   "metadata": {},
   "source": [
    "### xarray-histogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "47ef400d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "80.4 ms ± 2.21 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "assert(not da_2d[0]._in_memory)\n",
    "h = xh.histogram(*da_2d, bins=axes)\n",
    "assert_allclose(h.values, answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "563d6547",
   "metadata": {},
   "source": [
    "## One-dimensional (along second dimension)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "636d5b20",
   "metadata": {},
   "outputs": [],
   "source": [
    "vals_1d_t = da.random.normal(size=[3, 10_000_000], chunks=(1, chunk_size))\n",
    "da_1d_t = xr.DataArray(vals_1d_t, name='test_1d', dims=['t', 'x'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9fcf8872",
   "metadata": {},
   "outputs": [],
   "source": [
    "answer = xhistogram(da_1d_t, bins=edges[0], dim=['x']).load()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea695e1e",
   "metadata": {},
   "source": [
    "### xhistogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "97e0efbf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "216 ms ± 9.32 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "h = xhistogram(da_1d_t, bins=edges[0], dim=['x'], block_size=chunk_size)\n",
    "assert_allclose(h.values, answer.values)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5e3769f",
   "metadata": {},
   "source": [
    "### xarray-histogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "9a19b59d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "84.2 ms ± 3.28 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "h = xh.histogram(da_1d_t, bins=[axes[0]], dims=['x'])\n",
    "assert_allclose(h.values.T, answer.values)"
   ]
  }
 ],
 "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}