B-with-dask.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abd20a30-199e-42a2-8764-94bf43e732e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ftplib import FTP, error_perm, error_temp\n",
    "from pathlib import Path\n",
    "from itertools import product\n",
    "\n",
    "\n",
    "from tqdm import tqdm\n",
    "import s3fs\n",
    "from cloudside.asos import MetarParser\n",
    "from cloudside.asos import FIVEMIN, _find_reset_time, _process_precip\n",
    "\n",
    "from dask import bag\n",
    "from distributed import Client\n",
    "from coiled import Cluster\n",
    "\n",
    "# %% Create a cluster and client with coiled\n",
    "cluster = Cluster(name=\"asos-mirror\", n_workers=20, package_sync=True, show_widget=False)\n",
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "649431bc-02f5-4755-9d6a-2bdd547bad6c",
   "metadata": {},
   "source": [
    "## B - With Dask\n",
    "\n",
    "* (Same) Function to list all of the files available for a given station-year (e.g., 2020 at KPDX)\n",
    "* (Slightly Different) Function to transfer individual files from FTP to S3\n",
    "* (Different) Function process the precip data of a dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5453e98-894b-4ea0-be8a-903bf8d4a6d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "fs = s3fs.S3FileSystem()\n",
    "\n",
    "def get_all_files(station, year):\n",
    "    with FTP(\"ftp.ncei.noaa.gov\") as ftp:\n",
    "        ftp.login(passwd=email)\n",
    "        try:\n",
    "            dat_files = ftp.nlst(f\"/pub/data/asos-fivemin/6401-{year}/*{station}*\")\n",
    "        except error_perm as e:\n",
    "            print(e)\n",
    "    return dat_files\n",
    "\n",
    "\n",
    "def transfer(dat_in, dat_out_uri):\n",
    "    # don't try to transfer something that is already in S3\n",
    "    if not fs.exists(dat_out_uri) or (fs.size(dat_out_uri) == 0):\n",
    "        with fs.open(dat_out_uri, \"w\") as dat_out, \\\n",
    "             FTP(\"ftp.ncei.noaa.gov\") as ftp:\n",
    "            ftp.login(passwd=email)\n",
    "            try:\n",
    "                ftp.retrlines(\n",
    "                    f\"RETR {dat_in}\", lambda x: dat_out.write(x + \"\\n\")\n",
    "                )\n",
    "                status = \"transferred\"\n",
    "                file_size = fs.size(dat_out_uri)\n",
    "            except Exception as e:\n",
    "                file_size = None\n",
    "                status = str(e)\n",
    "    else:\n",
    "        file_size = fs.size(dat_out_uri)\n",
    "        status = \"exists already\"\n",
    "\n",
    "    return {\n",
    "        \"source\": dat_in,\n",
    "        \"file\": dat_out_uri,\n",
    "        \"status\": status,\n",
    "        \"size\": file_size\n",
    "    }\n",
    "\n",
    "\n",
    "def precip_process(df):\n",
    "    data = df.groupby(\"datetime\").last().sort_index().resample(FIVEMIN).asfreq()\n",
    "\n",
    "    rt = _find_reset_time(data[\"raw_precipitation\"])\n",
    "    precip = _process_precip(data, rt, \"raw_precipitation\")\n",
    "    return data.assign(precip=precip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a8267387-7404-4cca-8249-6f8c8c586442",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.61 s, sys: 422 ms, total: 2.03 s\n",
      "Wall time: 4min 54s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead tr th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe thead tr:last-of-type th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <th colspan=\"2\" halign=\"left\">temperature</th>\n",
       "      <th colspan=\"2\" halign=\"left\">wind_speed</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <th>min</th>\n",
       "      <th>max</th>\n",
       "      <th>min</th>\n",
       "      <th>max</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>station</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>KATL</th>\n",
       "      <td>-10.6</td>\n",
       "      <td>37.8</td>\n",
       "      <td>0.0</td>\n",
       "      <td>39.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KBHM</th>\n",
       "      <td>-17.2</td>\n",
       "      <td>39.4</td>\n",
       "      <td>0.0</td>\n",
       "      <td>35.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KLAS</th>\n",
       "      <td>-11.7</td>\n",
       "      <td>46.7</td>\n",
       "      <td>0.0</td>\n",
       "      <td>39.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KLAX</th>\n",
       "      <td>3.9</td>\n",
       "      <td>40.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>37.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KLGA</th>\n",
       "      <td>-17.2</td>\n",
       "      <td>38.3</td>\n",
       "      <td>0.0</td>\n",
       "      <td>39.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KORD</th>\n",
       "      <td>-24.4</td>\n",
       "      <td>36.1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>60.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KPDX</th>\n",
       "      <td>-11.7</td>\n",
       "      <td>46.7</td>\n",
       "      <td>0.0</td>\n",
       "      <td>94.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KSEA</th>\n",
       "      <td>-8.3</td>\n",
       "      <td>42.2</td>\n",
       "      <td>0.0</td>\n",
       "      <td>38.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KSFO</th>\n",
       "      <td>2.2</td>\n",
       "      <td>40.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>45.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        temperature       wind_speed      \n",
       "                min   max        min   max\n",
       "station                                   \n",
       "KATL          -10.6  37.8        0.0  39.0\n",
       "KBHM          -17.2  39.4        0.0  35.0\n",
       "KLAS          -11.7  46.7        0.0  39.0\n",
       "KLAX            3.9  40.0        0.0  37.0\n",
       "KLGA          -17.2  38.3        0.0  39.0\n",
       "KORD          -24.4  36.1        0.0  60.0\n",
       "KPDX          -11.7  46.7        0.0  94.0\n",
       "KSEA           -8.3  42.2        0.0  38.0\n",
       "KSFO            2.2  40.0        0.0  45.0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "# %% set up the data we're going to download\n",
    "email = \"paul+faa@coiled.io\"\n",
    "stations = [\n",
    "    \"KATL\", \"KPDX\", \"KLAS\",\n",
    "    \"KBHM\", \"KLGA\", \"KSEA\",\n",
    "    \"KORD\", \"KLAX\", \"KSFO\",\n",
    "]\n",
    "years = list(range(2016, 2022))\n",
    "\n",
    "station_years = product(stations, years)\n",
    "all_files = (\n",
    "    bag.from_sequence(station_years)\n",
    "       .map(lambda station_year: get_all_files(*station_year))\n",
    "       .compute()\n",
    ") # returns a list of file lists\n",
    "\n",
    "flattened = []\n",
    "for station_set in all_files:\n",
    "    flattened.extend(station_set)\n",
    "\n",
    "s3_base = \"oss-shared-scratch/paul-scratch/faa-scratch-demo/raw\"\n",
    "in_out = [\n",
    "    {\"dat_in\": f, \"dat_out_uri\": f\"{s3_base}/{Path(f).name}\"}\n",
    "    for f in flattened\n",
    "]\n",
    "\n",
    "# %%\n",
    "status_futures = client.map(lambda x: transfer(**x), in_out, retries=5)\n",
    "statuses = client.gather(status_futures)\n",
    "\n",
    "parsed = (\n",
    "    bag.read_text(f\"s3://{s3_base}/*.dat\")\n",
    "        .map(lambda x: MetarParser(x, strict=False).asos_dict(null_sky_as=None))\n",
    "        .to_dataframe()\n",
    "        .map_partitions(precip_process)\n",
    ")\n",
    "\n",
    "summary = (\n",
    "    parsed.loc[lambda df: df[\"temperature\"].gt(-25)]\n",
    "        .groupby(by=\"station\")[[\"temperature\", \"wind_speed\"]]\n",
    "        .agg([\"min\", \"max\"])\n",
    ")\n",
    "\n",
    "summary.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d7e32c5-4874-444d-9dd1-111874cf4163",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()\n",
    "cluster.close()"
   ]
  }
 ],
 "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}