prsim_csv_to_zarr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import dask.dataframe as dd\n",
    "from distributed import Client\n",
    "import pandas as pd\n",
    "import dask\n",
    "import hvplot.xarray\n",
    "import s3fs\n",
    "import zarr\n",
    "import os\n",
    "from time import gmtime, strftime\n",
    "import shutil\n",
    "import glob\n",
    "import holoviews as hv\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dask Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(silence_logs=40)\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Source of data\n",
    "bucket = 's3://prsim/hydrogrammes-stochastiques/09995/' \n",
    "\n",
    "# Wasabi cloud storage configurations\n",
    "client_kwargs={'endpoint_url': 'https://s3.us-east-1.wasabisys.com'}\n",
    "config_kwargs = {'max_pool_connections': 30}\n",
    "storage_options = {'anon': True,\n",
    "                   \"client_kwargs\": client_kwargs}\n",
    "\n",
    "# List all files in bucket\n",
    "s3 = s3fs.S3FileSystem(client_kwargs=client_kwargs, \n",
    "                       config_kwargs=config_kwargs, \n",
    "                       anon=True)  # public read\n",
    "file_list = ['s3://{}'.format(each) for each in s3.ls(bucket)  if each.endswith('.csv')]\n",
    "\n",
    "# Data sink (store). Alternativaly, the sink could be directly in the cloud.\n",
    "store= os.environ['HOME'] + '/Documents/store'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Zarr storage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Conversion des csv vers des datasets Xarray via le lazy loading (dask.delayed). Mise en commun (xr.combine_nested) et dépot dans un dossier Zarr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def chunks(lst, n):\n",
    "    \"\"\"Yield successive n-sized chunks from lst.\"\"\"\n",
    "    for i in range(0, len(lst), n):\n",
    "        yield lst[i:i + n]\n",
    "        \n",
    "@dask.delayed\n",
    "def load(filename):\n",
    "    return dd.read_csv(filename,\n",
    "                       storage_options=storage_options).compute()\n",
    "\n",
    "@dask.delayed\n",
    "def process(df, sim_member):\n",
    "    data = xr.Dataset.from_dataframe(df)\n",
    "    data['sim_member'] = sim_member\n",
    "    return data.expand_dims('sim_member').set_coords('sim_member')\n",
    "\n",
    "@dask.delayed\n",
    "def combine(results):\n",
    "    return xr.combine_nested(results, 'sim_member')\n",
    "\n",
    "@dask.delayed\n",
    "def save(ds, store):\n",
    "    return ds.to_zarr(store, consolidated=True)\n",
    "\n",
    "# main function to load, process, combine and save multiple zarr datasets lazily\n",
    "def f(filenames, index):\n",
    "    results = []\n",
    "    for sim_member, filename in enumerate(filenames):\n",
    "        data = load(filename)\n",
    "        data = process(data, int(index*10000) + sim_member)\n",
    "        results.append(data)  \n",
    "    ds = combine(results)\n",
    "    ds = save(ds, store + '/part/' + str(index))\n",
    "    ds = dask.compute(ds, retries=10)\n",
    "\n",
    "# calling the main function with a generator\n",
    "for index, files in enumerate(chunks(file_list,10000)):\n",
    "    print(strftime(\"%Y-%m-%d %H:%M:%S\", gmtime()) + ' : ' + str(index))\n",
    "    f(files, index)\n",
    "\n",
    "# Merging to a final dataset\n",
    "ds_out = dask.delayed(xr.combine_nested)\\\n",
    "    ([dask.delayed(xr.open_zarr)(store_num) \n",
    "      for store_num in sorted(glob.glob(store + '/part/*'))],\n",
    "     'sim_member')\\\n",
    "        .compute()\n",
    "\n",
    "# intermediate step to allow rechunking\n",
    "for name in list(ds_out.keys()):\n",
    "    del ds_out[name].encoding['chunks']\n",
    "    \n",
    "# Save to final zarr store\n",
    "ds_out.chunk({'index':365,'sim_member':1000})\\\n",
    "    .to_zarr(store + '/final', consolidated=True)\n",
    "\n",
    "# Clean up partial files\n",
    "[shutil.rmtree(path) for path in sorted(glob.glob(store + '/part/*'))]\n",
    "shutil.rmtree(store + '/part/')\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# local\n",
    "ds = xr.open_zarr(store + '/final')\n",
    "\n",
    "## cloud\n",
    "# store_wasabi = s3fs.S3Map(root='s3://prsim/hydrogrammes-stochastiques/zarr/09995',\n",
    "#                           s3=s3,\n",
    "#                           check=False)\n",
    "#ds = xr.open_zarr(store_wasabi)\n",
    "\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "variables = [\"Dozois\",\"Lac Victoria et lac Granet\",\n",
    "             \"Rapide-7\",\"Rapide-2\",\"Riviere Kinojevis\",\"Lac des Quinze\"]\n",
    "\n",
    "da = ds[variables].to_array().sum('variable').load()\n",
    "\n",
    "da.max('sim_member').hvplot(ylabel='Débit', grid=True)*\\\n",
    "da.mean('sim_member').hvplot()*\\\n",
    "da.min('sim_member').hvplot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "dict_bassins = {'Angliers': [\"Dozois\",\"Lac Victoria et lac Granet\",\n",
    "                             \"Rapide-7\",\"Rapide-2\",\"Riviere Kinojevis\",\"Lac des Quinze\"],\n",
    "                'Rapide-2': [\"Dozois\",\"Lac Victoria et lac Granet\",\n",
    "                             \"Rapide-7\",\"Rapide-2\",],\n",
    "                'Rapide-7': [\"Dozois\",\"Lac Victoria et lac Granet\",\n",
    "                             \"Rapide-7\"],\n",
    "                'Lac Victoria et lac Granet': [\"Dozois\",\"Lac Victoria et lac Granet\"],\n",
    "                'Dozois': [\"Dozois\"]}\n",
    "graphs=[]\n",
    "\n",
    "for key, value in dict_bassins.items():\n",
    "    print(key)\n",
    "    da = ds[value].to_array()\\\n",
    "                 .sum('variable')\n",
    "    sim_member = da.where(da==da.max(), drop=True).sim_member\n",
    "    graphs.append(da.sel(sim_member=sim_member)[0,:]\\\n",
    "                    .hvplot(x='index', grid=True, ylabel='Débit', \n",
    "                            xlabel='sim member : ' + str(sim_member.values[0]),\n",
    "                            title=key, legend=False).opts(active_tools=['wheel_zoom','pan']))\n",
    "    graphs.append(ds[value].sel(sim_member=sim_member)\\\n",
    "    .drop('sim_member')\\\n",
    "    .hvplot(x='index', grid=True, ylabel='Débit', \n",
    "            xlabel='sim member : ' + str(sim_member.values[0]),\n",
    "            title=key, legend=True).opts(active_tools=['wheel_zoom','pan']))\n",
    "hv.Layout(graphs).opts(title=title).cols(2)\n",
    "\n",
    "title='Hydrogrammes cumulatifs maximal des séries stochastiques à chaque site'\n",
    "graph_layout = hv.Layout(graphs).opts(title=title).cols(2)\n",
    "hv.save(graph_layout, 'graph_layout.html', backend='bokeh')\n",
    "\n",
    "graph_layout"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hopig",
   "language": "python",
   "name": "hopig"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}