09_dask_s3_joined.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5f7188d-d90b-4f91-a93b-a2462a7b31e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client, LocalCluster\n",
    "from dask_gateway import Gateway\n",
    "import dask.dataframe as dd\n",
    "import dask\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from utils.const import S3_JOINED_FILEPATH\n",
    "from utils.dask_metric_funcs import root_mean_squared_error, relative_bias, r_squared, kling_gupta_efficiency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9f8df61-f92f-4886-90c6-39665b4416f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "GROUPBY_FIELDS = [\"primary_location_id\", \"configuration\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6335cd02-ac62-491e-9eb1-267c7fa51280",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = LocalCluster()\n",
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "791ff76a-27ef-486e-b2c7-bcadf0a0732a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cluster.close()\n",
    "# client.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a5a3ac0-98c2-4dd3-b691-1941ae913c7b",
   "metadata": {},
   "source": [
    "**Gateway cluster**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf02d864-4132-4c18-9800-cadf46705b66",
   "metadata": {},
   "outputs": [],
   "source": [
    "gateway = Gateway()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db01e61f-57ab-4c99-8739-fb6b96352d6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "options = gateway.cluster_options()\n",
    "options.worker_cores = 1\n",
    "options.worker_memory = 4\n",
    "# options  # should show interactive widget to select cores, etc\n",
    "\n",
    "cluster = gateway.new_cluster(options)\n",
    "client = cluster.get_client()\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d081e402-41d6-4e6a-a9a3-135fcae723ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d91c4e9-6064-4908-bbe7-cc07b9bf1dea",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.scale(8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4492243-dc4b-4066-a6c2-f5fd1b8a7732",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.close()\n",
    "client.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3964e483-1bbc-4446-8a00-2e270e8e2902",
   "metadata": {},
   "source": [
    "NOTE: We don't see improvement using a distributed cluster (actually worse?), probably because the dataset is not large enough to warrant its use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1c418b5-0973-44ba-9335-b6331bd6b1af",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddf = dd.read_parquet(\n",
    "    S3_JOINED_FILEPATH,\n",
    "    columns=[\"primary_value\", \"secondary_value\", \"primary_location_id\", \"configuration\", \"absolute_difference\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb736f25-e704-4c39-b6c6-5e6644f52168",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddf_tmp = ddf.copy()\n",
    "grouped = ddf_tmp.groupby(GROUPBY_FIELDS)[[\"primary_value\", \"secondary_value\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de38abad-5e82-4dc1-9064-087d631b55c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "rmse_srs = grouped.apply(root_mean_squared_error, meta=pd.Series(dtype=np.float32, name='root_mean_squared_error')).compute()\n",
    "bias_srs = grouped.apply(relative_bias, meta=pd.Series(dtype=np.float32, name='relative_bias')).compute()\n",
    "rsquared_srs = grouped.apply(r_squared, meta=pd.Series(dtype=np.float32, name='r_squared')).compute()\n",
    "kge_srs = grouped.apply(kling_gupta_efficiency, meta=pd.Series(dtype=np.float32, name='kling_gupta_efficiency')).compute()\n",
    "\n",
    "metrics_df = pd.concat([rmse_srs, bias_srs, rsquared_srs, kge_srs], axis=1)\n",
    "metrics_df = metrics_df.reset_index()\n",
    "metrics_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09c62a76-f352-44bb-becf-131ceb689d96",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics_df.to_parquet(\"/data/benchmarks/teehr-benchmark-202404/results/dask_s3_joined_results.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fd5d523-db29-48f6-9c35-a09f5ed349e1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f4a93d5-f352-43bb-80f9-3266423eea3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%time\n",
    "# rmse_srs = rmse(ddf, groupby_fields=GROUPBY_FIELDS)\n",
    "# bias_srs = relative_bias(ddf, groupby_fields=GROUPBY_FIELDS)\n",
    "# r2_srs = r_squared(ddf, groupby_fields=GROUPBY_FIELDS)\n",
    "# kge_srs = kge(ddf, groupby_fields=GROUPBY_FIELDS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "356e9484-da6a-44b2-914b-12d8c1a79bcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ddf = ddf.repartition(npartitions=ddf.npartitions // 10)  # Reducing partions helps eliminate shuffling but reduces performance of calculations\n",
    "\n",
    "# ddf.npartitions\n",
    "\n",
    "# ddf = ddf.persist()  # if on a distributed system\n",
    "\n",
    "# %%time\n",
    "# tmp_df = ddf.head()\n",
    "# tmp_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a49342a-b02e-4ce7-862f-c4493de36761",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}