dsgrid-toolkit 0.3.3__cp313-cp313-win_amd64.whl

dsgrid/notebooks/connect_to_dsgrid_registry.ipynb

@@ -0,0 +1,949 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "93935e00",
+   "metadata": {},
+   "source": [
+    "## 1. initialize"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "59c46803",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-28T22:18:35.922502Z",
+     "start_time": "2022-06-28T22:18:35.770022Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "import os\n",
+    "import getpass\n",
+    "import shutil\n",
+    "\n",
+    "from pyspark.sql import SparkSession\n",
+    "from pyspark import SparkConf, SparkContext"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e604690d",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-28T22:18:37.067792Z",
+     "start_time": "2022-06-28T22:18:37.059186Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "os.environ"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5c7e730c",
+   "metadata": {},
+   "source": [
+    "## 2. start spark cluster"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c0fac86c",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-28T22:18:38.952702Z",
+     "start_time": "2022-06-28T22:18:38.947792Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# tweak setting here:\n",
+    "def init_spark(cluster=None, name=\"dsgrid\", tz=\"UTC\"):\n",
+    "    \"\"\"Initialize a SparkSession.\"\"\"\n",
+    "    conf = SparkConf().setAppName(name)\n",
+    "\n",
+    "    if cluster is None:\n",
+    "        spark = SparkSession.builder.master(\"local\").appName(name).getOrCreate()\n",
+    "    elif cluster == \"AWS\":\n",
+    "        pass\n",
+    "        # does not need to setMaster for AWS cluster\n",
+    "    else:\n",
+    "        conf = conf.setMaster(cluster)\n",
+    "    conf = conf.setAll(\n",
+    "        [\n",
+    "            #             (\"spark.sql.shuffle.partitions\", \"200\"),\n",
+    "            #             (\"spark.executor.instances\", \"7\"),\n",
+    "            #             (\"spark.executor.cores\", \"5\"),\n",
+    "            #             (\"spark.executor.memory\", \"10g\"),\n",
+    "            #             (\"spark.driver.memory\", \"10g\"),\n",
+    "            #             (\"spark.dynamicAllocation.enabled\", True),\n",
+    "            #             (\"spark.shuffle.service.enabled\", True),\n",
+    "            (\"spark.sql.session.timeZone\", tz),\n",
+    "        ]\n",
+    "    )\n",
+    "    spark = SparkSession.builder.config(conf=conf).getOrCreate()\n",
+    "    return spark"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "45a0e446",
+   "metadata": {},
+   "source": [
+    "To launch a standalone cluster or a cluster on Kestrel, follow **instructions** here: \\\n",
+    "https://github.com/dsgrid/dsgrid/tree/main/dev#spark-standalone-cluster\n",
+    "\n",
+    "accordingly, uncomment and update the cluster name below:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fd3c099a",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-28T22:18:59.645718Z",
+     "start_time": "2022-06-28T22:18:41.443768Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "main_tz = \"EST\"  # <--- UTC, EST\n",
+    "\n",
+    "### STAND-ALONE CLUSTER\n",
+    "# cluster = \"spark://lliu2-34727s:7077\"\n",
+    "# name = \"stand-alone\"\n",
+    "\n",
+    "### CLUSTER ON HPC - Type in nodename\n",
+    "# NODENAME = \"r103u23\" # <--- change after deploying cluster\n",
+    "# cluster = f\"spark://{NODENAME}.ib0.cm.hpc.nrel.gov:7077\"\n",
+    "# name = \"HPC\"\n",
+    "\n",
+    "### CLUSTER ON HPC - Get cluster from file dropped by prep_spark_cluster_notebook.py\n",
+    "# import toml\n",
+    "# config = toml.load(\"cluster.toml\")\n",
+    "# cluster = config[\"cluster\"]\n",
+    "# name = \"HPC\"\n",
+    "\n",
+    "### LOCAL MODE\n",
+    "# cluster = None\n",
+    "# name = \"local\"\n",
+    "\n",
+    "### AWS MODE\n",
+    "cluster = \"AWS\"\n",
+    "name = \"AWS\"\n",
+    "\n",
+    "# Initialize\n",
+    "spark = init_spark(cluster, \"dsgrid-load\", tz=main_tz)\n",
+    "\n",
+    "# get Spark Context UI\n",
+    "sc = spark.sparkContext\n",
+    "sc"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5c91e3b2",
+   "metadata": {},
+   "source": [
+    "#### The *Spark UI* above works only for local mode. For HPC cluster Spark UI, use:\n",
+    "http://localhost:8080"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1f4214f4",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-28T22:18:59.689409Z",
+     "start_time": "2022-06-28T22:18:59.647919Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "for x in sorted(sc.getConf().getAll()):\n",
+    "    print(x)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b33cfadc",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d8b37296",
+   "metadata": {},
+   "source": [
+    "## 3. dsgrid"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aa24da30",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-28T22:18:59.889562Z",
+     "start_time": "2022-06-28T22:18:59.691274Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "from IPython.core.display import display, HTML\n",
+    "\n",
+    "display(HTML(\"<style>.container { width:100% !important; }</style>\"))\n",
+    "import pandas as pd\n",
+    "\n",
+    "pd.set_option(\"display.max_rows\", 20)\n",
+    "# import plotly\n",
+    "# pd.options.plotting.backend = \"plotly\"\n",
+    "import numpy as np\n",
+    "import itertools\n",
+    "import pytz\n",
+    "from datetime import datetime, timedelta\n",
+    "\n",
+    "from semver import VersionInfo\n",
+    "from pydantic import ValidationError\n",
+    "import pyspark.sql.functions as F\n",
+    "import pyspark.sql.types as sparktypes"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "239217ef",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-28T22:19:00.152343Z",
+     "start_time": "2022-06-28T22:18:59.891604Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "from dsgrid.common import LOCAL_REGISTRY\n",
+    "from dsgrid.registry.registry_manager import RegistryManager\n",
+    "from dsgrid.utils.files import load_data\n",
+    "from dsgrid.utils.spark import create_dataframe, read_dataframe, get_unique_values\n",
+    "from dsgrid.dimension.base_models import DimensionType\n",
+    "from dsgrid.dataset.dataset import Dataset\n",
+    "from dsgrid.project import Project\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "02cec913",
+   "metadata": {},
+   "source": [
+    "## 3.1. Check dsgrid registry"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2ce24d62",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-28T22:19:05.931097Z",
+     "start_time": "2022-06-28T22:19:05.927089Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "## sync registry and then load offline\n",
+    "# LOCAL_REGISTRY = \"s3://nrel-dsgrid-registry-archive\"\n",
+    "registry_path = os.getenv(\"DSGRID_REGISTRY_PATH\", default=LOCAL_REGISTRY)\n",
+    "registry_path"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f15cdc99",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-28T22:19:40.741402Z",
+     "start_time": "2022-06-28T22:19:11.930311Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "sync_and_pull = True  # <--- registry config only\n",
+    "if sync_and_pull:\n",
+    "    print(f\"syncing registry: {registry_path}\")\n",
+    "    RegistryManager.load(registry_path, offline_mode=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bd2c937c",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-28T22:19:41.397174Z",
+     "start_time": "2022-06-28T22:19:40.743809Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# ETH@Review: Were you intending to write something to the right of the arrow?\n",
+    "offline_mode = True  # <---\n",
+    "\n",
+    "registry_mgr = RegistryManager.load(registry_path, offline_mode=offline_mode)\n",
+    "project_mgr = registry_mgr.project_manager\n",
+    "dataset_mgr = registry_mgr.dataset_manager\n",
+    "dim_map_mgr = registry_mgr.dimension_mapping_manager\n",
+    "dim_mgr = registry_mgr.dimension_manager\n",
+    "# ETH@Review: This line seems out of place. Or change \"Loading\" to \"Loaded\"?\n",
+    "print(f\"Loaded dsgrid registry at: {registry_path}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23c1b83f",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-28T22:19:48.637994Z",
+     "start_time": "2022-06-28T22:19:41.399044Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "project_mgr.show(max_width=30, drop_fields=[\"Date\", \"Submitter\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "37b93d4f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d4815918",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:06.629410Z",
+     "start_time": "2022-06-18T01:51:06.627155Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# %%timeit\n",
+    "# ## Dan's test\n",
+    "# from dsgrid.config.time_dimension_base_config import TimeDimensionBaseConfig\n",
+    "\n",
+    "# i = 0\n",
+    "# for d_id in registry_mgr.dimension_manager._id_to_type:\n",
+    "#     config = registry_mgr.dimension_manager.get_by_id(d_id)\n",
+    "#     if not isinstance(config, TimeDimensionBaseConfig):\n",
+    "#         config.get_records_dataframe().count()\n",
+    "#         i += 1\n",
+    "\n",
+    "# print(i)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "484d7f4c",
+   "metadata": {},
+   "source": [
+    "## 3.2. Load Project\n",
+    "This section is mostly exploratory (For *Section 4. Queries*, only need to load project) \n",
+    "\n",
+    "####  Some user criteria:\n",
+    "At the projects, I want to be able to:\n",
+    "- Examine what's available in the project:\n",
+    "    * Show project dimensions by type, show resolution by type - I don't care: base/supplemental, mappings, id\n",
+    "    * Get unique records by dimension/resolution\n",
+    "    * Get unique records by selected dimension sets\n",
+    "    * Show mapped dataset\n",
+    "    * Show unit (or select a unit of analysis) and fuel types\n",
+    "- Make queries using:\n",
+    "    * Project dimensions + fuel types + time resolutions\n",
+    "    * Get all types of statistics (max, mean, min, percentiles, count, sum)\n",
+    "    \n",
+    "- dataset level: never mapped, think TEMPO,\n",
+    "- interface to allow for query optimization\n",
+    "    \n",
+    "#### Notes:\n",
+    " * Project_manager has access to all other managers.\n",
+    " * Each manager has the responsiblity to retrieve configs\n",
+    " * Access ConfigModel from configs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8083ad86",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:08.162466Z",
+     "start_time": "2022-06-18T01:51:06.631419Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# load projct\n",
+    "project_id = \"dsgrid_conus_2022\"  # <---\n",
+    "project = project_mgr.load_project(project_id)\n",
+    "\n",
+    "print(\"project loaded\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cd51fbba",
+   "metadata": {},
+   "source": [
+    "## 3.3. Load Project Datasets"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "89660b28",
+   "metadata": {},
+   "source": [
+    "### 3.3.3. TEMPO\n",
+    "\n",
+    "load and check tempo dataset here"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13295901",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:13.121796Z",
+     "start_time": "2022-06-18T01:51:08.166523Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "dataset_id = \"tempo_conus_2022\"  # <----\n",
+    "project.load_dataset(dataset_id)\n",
+    "tempo = project.get_dataset(dataset_id)\n",
+    "print(\"tempo dataset loaded\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fbce31e1",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:13.126258Z",
+     "start_time": "2022-06-18T01:51:13.123686Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "### TO BE DELETED ###\n",
+    "tempo_load_data_lookup = tempo.load_data_lookup\n",
+    "tempo_load_data = tempo.load_data\n",
+    "\n",
+    "# file = \"/scratch/dthom/tempo_load_data3.parquet\" # <---\n",
+    "# tempo_load_data = spark.read.parquet(file)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9211fdc6",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:13.697017Z",
+     "start_time": "2022-06-18T01:51:13.128635Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "tempo_mapped_load_data_lookup = tempo._handler._remap_dimension_columns(tempo_load_data_lookup)\n",
+    "tempo_mapped_load_data = tempo._handler._remap_dimension_columns(tempo_load_data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8037d63a",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:13.701676Z",
+     "start_time": "2022-06-18T01:51:13.699034Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "del tempo_load_data_lookup\n",
+    "del tempo_load_data"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2719ca95",
+   "metadata": {},
+   "source": [
+    "## 4. Queries\n",
+    "### Query util functions"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3a58e461",
+   "metadata": {},
+   "source": [
+    "### 4.1. Hourly electricity consumption by *scenario, model_year, and ReEDS PCA*"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "95324872",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:14.175522Z",
+     "start_time": "2022-06-18T01:51:13.703388Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "### all_enduses-totelectric_enduses map\n",
+    "\n",
+    "dim_map_id = \"conus-2022-detailed-end-uses-kwh__all-electric-end-uses__c4149547-1209-4ce3-bb4c-3ab292067e8a\"  # <---\n",
+    "electric_enduses_map = dim_map_mgr.get_by_id(dim_map_id).get_records_dataframe()\n",
+    "\n",
+    "### get all project electric end uses\n",
+    "electric_enduses = (\n",
+    "    electric_enduses_map.filter(\"to_id is not NULL\")\n",
+    "    .select(\"from_id\")\n",
+    "    .toPandas()[\"from_id\"]\n",
+    "    .to_list()\n",
+    ")\n",
+    "electric_enduses"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fc6700c0",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:14.483047Z",
+     "start_time": "2022-06-18T01:51:14.177069Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "### county-to-PCA map\n",
+    "dim_map_id = \"us_counties_2020_l48__reeds_pca__fcc554e1-87c9-483f-89e3-a0df9563cf89\"  # <---\n",
+    "county_to_pca_map = dim_map_mgr.get_by_id(dim_map_id).get_records_dataframe()\n",
+    "county_to_pca_map.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "16a76b85",
+   "metadata": {},
+   "source": [
+    "### 4.1.3. TEMPO\n",
+    "query TEMPO data here"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c590dc5e",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:15.726162Z",
+     "start_time": "2022-06-18T01:51:14.485411Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "## Load timezone map (not registered)\n",
+    "timezone_file = \"s3://nrel-dsgrid-int-scratch/scratch-lliu2/county_fip_to_local_prevailing_time.csv\"  # \"/scratch/lliu2/project_county_timezone/county_fip_to_local_prevailing_time.csv\"\n",
+    "tz_map = spark.read.csv(timezone_file, header=True)\n",
+    "tz_map = tz_map.withColumn(\"from_fraction\", F.lit(1))\n",
+    "tz_map.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "87a88b0c",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:15.733246Z",
+     "start_time": "2022-06-18T01:51:15.728420Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "### get electric end uses for transportation\n",
+    "tra_elec_enduses = [col for col in tempo_mapped_load_data.columns if col in electric_enduses]\n",
+    "tra_elec_enduses"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4970e36b",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:15.737460Z",
+     "start_time": "2022-06-18T01:51:15.735225Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "### TO BE DELETED\n",
+    "# tempo_mapped_load_data_lookup = tempo_mapped_load_data_lookup.filter(\"id in ('1621180393', '770011011', '1058530452')\")\n",
+    "# tempo_mapped_load_data = tempo_mapped_load_data.filter(\"id in ('1621180393', '770011011', '1058530452')\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "54adf133",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:15.762916Z",
+     "start_time": "2022-06-18T01:51:15.739144Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "## 0. consolidate load_data: get total hourly electricity consumption by id\n",
+    "# make get_time_cols accessible at dataset level\n",
+    "tra_elec_kwh = tempo_mapped_load_data.select(\n",
+    "    \"id\",\n",
+    "    \"day_of_week\",\n",
+    "    \"hour\",\n",
+    "    \"month\",\n",
+    "    sum([F.col(col) for col in tra_elec_enduses]).alias(\"electricity\"),\n",
+    ")\n",
+    "# tra_elec_kwh.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "14e42de4",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:15.831343Z",
+     "start_time": "2022-06-18T01:51:15.765116Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "## 1. map load_data_lookup to timezone\n",
+    "load_data_lookup = (\n",
+    "    tempo_mapped_load_data_lookup.filter(\"id is not NULL\")\n",
+    "    .select(\"sector\", \"scenario\", \"model_year\", \"geography\", \"id\", \"fraction\")\n",
+    "    .join(\n",
+    "        tz_map,\n",
+    "        on=F.col(\"geography\") == tz_map.from_id,\n",
+    "        how=\"left\",\n",
+    "    )\n",
+    "    .drop(\"from_id\")\n",
+    "    .withColumnRenamed(\"to_id\", \"timezone\")\n",
+    ")\n",
+    "\n",
+    "## combine fraction\n",
+    "nonfraction_cols = [x for x in load_data_lookup.columns if x not in {\"fraction\", \"from_fraction\"}]\n",
+    "load_data_lookup = load_data_lookup.fillna(1, subset=[\"from_fraction\"]).selectExpr(\n",
+    "    *nonfraction_cols, \"fraction*from_fraction AS fraction\"\n",
+    ")\n",
+    "# load_data_lookup.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b1edc12e",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:15.897261Z",
+     "start_time": "2022-06-18T01:51:15.833472Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "## 2. join load_data and lookup\n",
+    "tra_elec_kwh = load_data_lookup.join(\n",
+    "    tra_elec_kwh,\n",
+    "    on=\"id\",\n",
+    "    how=\"left\",\n",
+    ").drop(\"id\")\n",
+    "\n",
+    "tra_elec_kwh = tra_elec_kwh.groupBy(\n",
+    "    \"sector\",\n",
+    "    \"scenario\",\n",
+    "    \"geography\",\n",
+    "    \"model_year\",\n",
+    "    \"timezone\",\n",
+    "    \"day_of_week\",\n",
+    "    \"month\",\n",
+    "    \"hour\",\n",
+    ").agg(F.sum(F.col(\"fraction\") * F.col(\"electricity\")).alias(\"electricity\"))\n",
+    "\n",
+    "## cache df\n",
+    "# tra_elec_kwh = tra_elec_kwh.cache()\n",
+    "# tra_elec_kwh.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3b9756a8",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:17.848803Z",
+     "start_time": "2022-06-18T01:51:15.899510Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "year = 2012  # <--- weather year\n",
+    "sys_tz = TimeZone.EST.tz\n",
+    "timezones_local = [TimeZone.EPT, TimeZone.CPT, TimeZone.MPT, TimeZone.PPT]\n",
+    "\n",
+    "## 3. create range of model_year\n",
+    "model_time_pd = []\n",
+    "for tz in timezones_local:\n",
+    "    model_time_df = pd.DataFrame()\n",
+    "    # create time range in local time\n",
+    "    model_time_df[\"timestamp\"] = pd.date_range(\n",
+    "        start=datetime(year=int(year), month=1, day=1, hour=0),\n",
+    "        end=datetime(year=int(year), month=12, day=31, hour=23),\n",
+    "        tz=tz.tz,\n",
+    "        freq=\"H\",\n",
+    "    )\n",
+    "    model_time_df[\"timezone\"] = tz.value\n",
+    "    model_time_df[\"day_of_week\"] = model_time_df[\"timestamp\"].dt.day_of_week.astype(str)\n",
+    "    model_time_df[\"month\"] = model_time_df[\"timestamp\"].dt.month.astype(str)\n",
+    "    model_time_df[\"hour\"] = model_time_df[\"timestamp\"].dt.hour.astype(str)\n",
+    "\n",
+    "    # convert to main timezone\n",
+    "    model_time_df[\"timestamp\"] = model_time_df[\"timestamp\"].dt.tz_convert(sys_tz)\n",
+    "    # wrap time to year\n",
+    "    model_time_df[\"timestamp\"] = model_time_df[\"timestamp\"].apply(lambda x: x.replace(year=year))\n",
+    "\n",
+    "    model_time_pd.append(model_time_df)\n",
+    "\n",
+    "model_time_pd = pd.concat(model_time_pd, axis=0, ignore_index=True)\n",
+    "model_time_pd[\"timestamp\"] = (\n",
+    "    model_time_pd[\"timestamp\"].dt.tz_localize(None).astype(str)\n",
+    ")  # conver timestamp to str, this is important!\n",
+    "print(model_time_pd)\n",
+    "\n",
+    "# convert to spark df\n",
+    "schema = sparktypes.StructType(\n",
+    "    [\n",
+    "        sparktypes.StructField(\"timestamp\", sparktypes.StringType(), False),\n",
+    "        sparktypes.StructField(\"timezone\", sparktypes.StringType(), False),\n",
+    "        sparktypes.StructField(\"day_of_week\", sparktypes.StringType(), False),\n",
+    "        sparktypes.StructField(\"month\", sparktypes.StringType(), False),\n",
+    "        sparktypes.StructField(\"hour\", sparktypes.StringType(), False),\n",
+    "    ]\n",
+    ")\n",
+    "model_time = spark.createDataFrame(model_time_pd, schema=schema)\n",
+    "\n",
+    "## covert timestamp from str to timestamp\n",
+    "model_time = model_time.withColumn(\n",
+    "    \"timestamp\",\n",
+    "    F.from_unixtime(\n",
+    "        F.unix_timestamp(F.col(\"timestamp\"), \"yyyy-MM-dd HH:mm:ss\"), \"yyyy-MM-dd HH:mm:ss\"\n",
+    "    ),\n",
+    ")\n",
+    "model_time = model_time.withColumn(\"timestamp\", F.to_timestamp(\"timestamp\"))\n",
+    "model_time = model_time.cache()\n",
+    "\n",
+    "print(model_time.printSchema())\n",
+    "print(model_time.count())\n",
+    "model_time.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "554bc22f",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:17.875693Z",
+     "start_time": "2022-06-18T01:51:17.851121Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "## 4. expand to model_years\n",
+    "tra_elec_kwh = model_time.join(\n",
+    "    tra_elec_kwh, on=[\"timezone\", \"day_of_week\", \"month\", \"hour\"], how=\"right\"\n",
+    ").drop(\"day_of_week\", \"month\", \"hour\")\n",
+    "\n",
+    "## cache df\n",
+    "# tra_elec_kwh = tra_elec_kwh.cache()\n",
+    "# tra_elec_kwh.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "734e3970",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T01:51:17.915745Z",
+     "start_time": "2022-06-18T01:51:17.878004Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "# 5. map load_data_lookup to PCA\n",
+    "tra_elec_kwh = (\n",
+    "    tra_elec_kwh.join(\n",
+    "        county_to_pca_map, on=F.col(\"geography\") == county_to_pca_map.from_id, how=\"left\"\n",
+    "    )\n",
+    "    .drop(\"from_id\")\n",
+    "    .drop(\"geography\")\n",
+    "    .withColumnRenamed(\"to_id\", \"geography\")\n",
+    "    .groupBy(\"sector\", \"scenario\", \"geography\", \"model_year\", \"timestamp\")\n",
+    "    .agg(F.sum(\"electricity\").alias(\"electricity\"))\n",
+    ")\n",
+    "\n",
+    "# tra_elec_kwh.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13233566",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T03:28:13.295454Z",
+     "start_time": "2022-06-18T01:51:17.917568Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "### 6. save as partitions\n",
+    "tra_output_file = \"s3://nrel-dsgrid-int-scratch/scratch-lliu2/tempo_projections.parquet\"  # Path(f\"/scratch/{getpass.getuser()}/tempo_projections.parquet\")\n",
+    "\n",
+    "# # refresh file dir\n",
+    "if Path(tra_output_file).exists():\n",
+    "    shutil.rmtree(tra_output_file)\n",
+    "\n",
+    "if Path(tra_output_file).exists():\n",
+    "    raise ValueError(\n",
+    "        f\"file: {tra_output_file} already exist. `shutile.rmtree(tra_output_file)` to override.\"\n",
+    "    )\n",
+    "\n",
+    "tra_elec_kwh.sort(\"scenario\", \"model_year\", \"geography\", \"timestamp\").repartition(\n",
+    "    \"scenario\", \"model_year\"\n",
+    ").write.partitionBy(\"scenario\", \"model_year\").option(\"path\", tra_output_file).saveAsTable(\n",
+    "    \"tra_elec_kwh\", format=\"parquet\"\n",
+    ")\n",
+    "\n",
+    "print(\"tra_elec_kwh saved\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "853ffc6d",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T03:28:13.301224Z",
+     "start_time": "2022-06-18T03:28:13.297395Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# %%time\n",
+    "# ########## load transportation projection data ###########\n",
+    "# tra_output_file = \"s3://nrel-dsgrid-int-scratch/scratch-lliu2/tempo_projections.parquet\" #Path(f\"/scratch/{getpass.getuser()}/tempo_projections.parquet\")\n",
+    "\n",
+    "# if Path(tra_output_file).exists():\n",
+    "#     tra_elec_kwh = read_dataframe(tra_output_file)\n",
+    "#     print(\"tra_elec_kwh loaded\")\n",
+    "# else:\n",
+    "#     print(f\"tra_output_file={tra_output_file} does not exist\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "28f6277a",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-06-18T05:24:12.421483Z",
+     "start_time": "2022-06-18T04:03:24.836808Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "ts = tra_elec_kwh.groupBy(\"timestamp\").count().orderBy(\"timestamp\").toPandas()\n",
+    "ts"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "87d8b42a",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.8.10 ('dsgrid')",
+   "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.8.10"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "2458d4f391e03ccae12714782d51aa387d09e7b7a16d6832b1f2bffaf5a9bcc2"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}