{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "smaller-breeding",
   "metadata": {},
   "source": [
    "# Saving frames in parallel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "complimentary-symphony",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "from xmovie import Movie\n",
    "\n",
    "ds = xr.tutorial.open_dataset('air_temperature').isel(time=slice(0,200))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "prostate-cuisine",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating the movie object\n",
    "mov = Movie(ds.air, vmin=230, vmax=310)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd30e0fb-77fa-41b3-ba42-384f3b44fdd2",
   "metadata": {},
   "source": [
    "The creation of a movie can take quite long for datasets with many timesteps, creating many frames in a loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "weekly-connectivity",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Movie created at movie.mov\n",
      "CPU times: user 34.7 s, sys: 1.02 s, total: 35.7 s\n",
      "Wall time: 59.4 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "mov.save('movie.mov', overwrite_existing=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66dae1a7-f4f7-455e-9a47-8af1f600d122",
   "metadata": {},
   "source": [
    "You can speed up the frame creation by activating the `parallel` option. This will save the frames using dask.\n",
    "\n",
    "For this to work you need to chunk the input dataarray with a single step along the dimension that represent your frames (`framedim`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c44b3d04-349f-404b-a032-128e6f789922",
   "metadata": {},
   "outputs": [],
   "source": [
    "mov_parallel = Movie(ds.air.chunk({'time':1}), vmin=230, vmax=310)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "obvious-paintball",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Movie created at movie_parallel.mov\n",
      "CPU times: user 38.8 s, sys: 1.46 s, total: 40.3 s\n",
      "Wall time: 48.3 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "mov_parallel.save(\n",
    "    'movie_parallel.mov',\n",
    "    parallel=True,\n",
    "    overwrite_existing=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6dbd6fc-5b9c-424b-b198-9fb886d66d09",
   "metadata": {},
   "source": [
    "You can pass arguments to the dask `.compute()` call with `parallel_compute_kwargs` to tune for your particular setup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4d389008-c4db-4ce2-8d56-d561e59c8c9e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Movie created at movie_parallel.mov\n",
      "CPU times: user 4.84 s, sys: 249 ms, total: 5.09 s\n",
      "Wall time: 33.6 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "mov_parallel.save(\n",
    "    'movie_parallel_modified.mov',\n",
    "    parallel=True,\n",
    "    overwrite_existing=True,\n",
    "    parallel_compute_kwargs=dict(scheduler=\"processes\", num_workers=8)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cfe4ea3-1f6b-42ee-a346-f0e2042fa794",
   "metadata": {},
   "source": [
    "Thats not bad, a 50% time saving (keeping in mind that the time needed for the ffmpeg call is included). We expect speedups to be even bigger when using higher resolution data."
   ]
  }
 ],
 "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}