10_render_frames.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sqlite3\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import time\n",
    "import math\n",
    "import pickle\n",
    "import multiprocessing as mp\n",
    "import gc\n",
    "\n",
    "from datetime import datetime, timedelta\n",
    "\n",
    "! mkdir -p out/frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_data():\n",
    "    with sqlite3.connect(\"data.sqlite3\") as db:\n",
    "        return pd.read_sql_query(f\"\"\"\n",
    "            -- Dedupe rows in case scrape was executed multiple times\n",
    "            with max_scrape as (\n",
    "                select id, max(scrapeTime) as scrapeTime\n",
    "                from earthquakes\n",
    "                group by id\n",
    "            )\n",
    "            select time, latitude, longitude, depth, mag\n",
    "            from earthquakes\n",
    "            inner join max_scrape using (id, scrapeTime)\n",
    "            where status <> 'deleted'\n",
    "        \"\"\", db, parse_dates=None)\n",
    "\n",
    "df = read_data()\n",
    "df['datetime'] = pd.to_datetime(df.time)\n",
    "df['unix'] = (df.datetime.astype(np.int64) / 1e9).astype(np.int64)\n",
    "df['energy'] = 10 ** (1.5 * (df.mag + 3.2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Targeting 2min@60fps\n",
    "target_fps = 60\n",
    "target_len_minutes = 2\n",
    "\n",
    "# Time range to draw\n",
    "start_time = datetime(2010, 1, 1, 0, 0).timestamp()\n",
    "end_time = datetime(2019, 7, 8, 0, 0).timestamp()\n",
    "\n",
    "# Number of frames for quakes to fade away\n",
    "fadeaway_frames = 90\n",
    "\n",
    "target_frame_count = target_fps * target_len_minutes * 60\n",
    "total_data_unix = end_time - start_time\n",
    "unix_per_frame = int(total_data_unix / target_frame_count)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "min_depth = 1\n",
    "max_depth = 800\n",
    "\n",
    "def draw_data(end_unix):\n",
    "    import os\n",
    "    # Bypass annoying JupyterLab issue with Conda\n",
    "    os.environ.setdefault('PROJ_LIB', '/home/randombk/.conda/envs/data-analytics/share/proj')\n",
    "    \n",
    "    # Avoid matplotlib multiprocessing errors by importing inside the function\n",
    "    import matplotlib.pyplot as plt\n",
    "    from mpl_toolkits.basemap import Basemap\n",
    "    from pandas.plotting import register_matplotlib_converters\n",
    "    register_matplotlib_converters()\n",
    "    \n",
    "    frame_start_time = time.time()\n",
    "    \n",
    "    # Prepare main map data\n",
    "    data = df[\n",
    "        (df.unix < end_unix) & \n",
    "        (df.unix > end_unix - unix_per_frame * fadeaway_frames)\n",
    "    ].copy()\n",
    "    data['opacity'] = 1 - ((end_unix - data.unix) / unix_per_frame) / fadeaway_frames\n",
    "    \n",
    "    # Prepare bottom heatmap data\n",
    "    heatmap_data = df[(df.unix < end_unix) & df.depth.notnull() & (df.mag > 3)]\n",
    "    \n",
    "    # Prepare plot\n",
    "    fig = plt.figure(figsize=(19, 12))\n",
    "    ax_map = fig.add_axes([0.05, 0.2, 0.9, 0.8], xticklabels=[], yticklabels=[], ylim=(-1.2, 1.2))\n",
    "    ax_depth_heatmap = fig.add_axes(\n",
    "        [0.05, 0.05, 0.9, 0.15],\n",
    "        xlim=(datetime.utcfromtimestamp(start_time), datetime.utcfromtimestamp(end_time)),\n",
    "        ylim=(-max_depth, min_depth),\n",
    "        ylabel=\"Epicenter Depth (km)\"\n",
    "    )\n",
    "\n",
    "    title = [\n",
    "        \"Global Earthquakes Greater Than 2.5 Magnitude\",\n",
    "        datetime.utcfromtimestamp(end_unix).strftime('%Y-%m-%d %H:%M:%S UTC'),\n",
    "    ]\n",
    "    ax_map.set_title(\"\\n\".join(title))\n",
    "\n",
    "    # Draw background map\n",
    "    basemap = Basemap(projection='robin',lon_0=0, ax=ax_map)\n",
    "    basemap.drawcoastlines()\n",
    "    basemap.drawparallels(np.arange(-90, 90, 30),labels=[1, 0, 0, 0])\n",
    "    basemap.drawmeridians(np.arange(basemap.lonmin,basemap.lonmax+30,60),labels=[0,0,0,1])\n",
    "    \n",
    "    # Draw earthquakes\n",
    "    quake_x, quake_y = basemap(data.longitude.to_list(), data.latitude.to_list())\n",
    "    quake_colors = np.zeros((len(data), 4))\n",
    "    quake_colors[:,0] = np.clip(np.exp(data.mag / 10) - 1, 0, 1)\n",
    "    quake_colors[:,1] = np.clip(2 - np.exp((data.mag - 2) / 10), 0.2, 0.5)\n",
    "    quake_colors[:,2] = np.clip(2 - np.exp((data.mag - 2) / 10), 0, 0.7)\n",
    "    quake_colors[:,3] = data.opacity\n",
    "    quake_sizes = (data.energy/1.0e7)**0.5 / 10\n",
    "    ax_map.scatter(quake_x, quake_y, s=quake_sizes, color=quake_colors)\n",
    "    \n",
    "    # Draw heatmap\n",
    "    heatmap_colors = np.zeros((len(heatmap_data), 4))\n",
    "    heatmap_colors[:,0] = np.log(np.clip(heatmap_data.depth, min_depth, max_depth)) / np.log(max_depth)\n",
    "    heatmap_colors[:,2] = 1.0 - np.log(np.clip(heatmap_data.depth, min_depth, max_depth)) / np.log(max_depth)\n",
    "    heatmap_colors[:,3] = 0.3\n",
    "    heatmap_sizes = (heatmap_data.energy/1.0e7)**0.5 / 100\n",
    "    ax_depth_heatmap.scatter(\n",
    "        heatmap_data.datetime, \n",
    "        -np.clip(heatmap_data.depth, min_depth, max_depth), \n",
    "        s=heatmap_sizes,\n",
    "        color=heatmap_colors\n",
    "    )\n",
    "\n",
    "\n",
    "    # Draw Metrics\n",
    "    metrics = [\n",
    "        f\"Total Events: {len(heatmap_data)}\",\n",
    "        f\"Mag 7+: {len(heatmap_data[heatmap_data.mag >= 7])}\",\n",
    "        f\"Mag 6+: {len(heatmap_data[heatmap_data.mag >= 6])}\",\n",
    "#         f\"Mag 5-5.9: {len(heatmap_data[(heatmap_data.mag >= 5) & (heatmap_data.mag < 6)])}\",\n",
    "#         f\"Mag 2.5-4.9: {len(heatmap_data[heatmap_data.mag < 5])}\",\n",
    "    ]\n",
    "    for idx, val in enumerate(metrics):\n",
    "        ax_map.text(\n",
    "            -600000, 17000000 - 320000 * idx, \n",
    "            val,\n",
    "            fontsize='large',\n",
    "        )\n",
    "\n",
    "    # Draw Magnitude Label\n",
    "    ax_map.text(\n",
    "        -600000, 500000, \n",
    "        \"Magnitude\",\n",
    "        rotation=90,\n",
    "        fontsize='x-large'\n",
    "    )\n",
    "    legend_mags = np.array([8, 7, 6, 5, 4, 3])\n",
    "    legend_xpos = [1300000, 3200000, 4200000, 4800000, 5200000, 5600000]\n",
    "    legend_ypos = [1500000, 850000, 620000, 500000, 455000, 450000]\n",
    "    legend_energies = 10 ** (1.5 * (legend_mags + 3.2))\n",
    "    legend_mag_sizes = (legend_energies/1.0e7)**0.5 / 10\n",
    "    for i in range(len(legend_mags)):\n",
    "        ax_map.scatter(\n",
    "            legend_xpos[i],\n",
    "            legend_ypos[i],\n",
    "            s=legend_mag_sizes[i],\n",
    "            c = [(\n",
    "                np.clip(np.exp(legend_mags[i] / 10) - 1, 0, 1),\n",
    "                np.clip(2 - np.exp((legend_mags[i] - 2) / 10), 0.2, 0.5),\n",
    "                np.clip(2 - np.exp((legend_mags[i] - 2) / 10), 0, 0.7)\n",
    "            )],\n",
    "        )\n",
    "        ax_map.annotate(str(legend_mags[i]), (legend_xpos[i], 0), ha='center')\n",
    "    \n",
    "    # Draw footnotes\n",
    "    footnotes = [\n",
    "        'Created by David Li',\n",
    "        '/u/dlp_randombk',\n",
    "        '',\n",
    "        'Source: U.S. Geological Survey',\n",
    "        'Accessed 2019-07-07',\n",
    "    ]\n",
    "    for idx, val in enumerate(footnotes):\n",
    "        ax_map.text(\n",
    "            33500000, 1320000 - 320000 * idx, \n",
    "            val,\n",
    "            fontsize='large',\n",
    "            horizontalalignment='right'\n",
    "        )\n",
    "\n",
    "    fig.savefig(f\"out/frames/{int(end_unix)}.png\", dpi=160, transparent=False)\n",
    "    \n",
    "    fig.clf()\n",
    "    plt.close(fig)\n",
    "    plt.close()\n",
    "    del data\n",
    "    del heatmap_data\n",
    "    gc.collect()\n",
    "    \n",
    "    return (end_unix, time.time() - frame_start_time)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using 24 cores...\n"
     ]
    }
   ],
   "source": [
    "frame_times = []\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    import warnings\n",
    "    warnings.filterwarnings('ignore')\n",
    "    \n",
    "    # Prepare list of jobs\n",
    "    jobs = range(int(start_time), int(end_time + unix_per_frame), unix_per_frame)\n",
    "    print(f\"Using {mp.cpu_count()} cores...\")\n",
    "\n",
    "    # There's a memory leak somewhere. We'll just hack our way around it.\n",
    "    with mp.Pool(processes=mp.cpu_count(), maxtasksperchild=4) as pool:\n",
    "        frame_times = pool.map(draw_data, jobs, chunksize=16)\n",
    "        with open('out/frame_times.pickle', 'wb') as fp:\n",
    "            pickle.dump(frame_times, fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# draw_data(end_time)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Data Analytics",
   "language": "python",
   "name": "data-analytics"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}