notebooks/covid-19-average-lockdown.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/oughtinc/ergo/blob/notebooks-readme/covid-19-average-lockdown.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "yElDWWhHmM4b"
   },
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 84
    },
    "colab_type": "code",
    "id": "o8yH0BYFmLqH",
    "outputId": "02145d85-3fd2-4914-db14-722561f0dfee"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
      "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
      "    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n",
      "  Building wheel for ergo (PEP 517) ... \u001b[?25l\u001b[?25hdone\n"
     ]
    }
   ],
   "source": [
    "!pip install --quiet poetry  # Fixes https://github.com/python-poetry/poetry/issues/532\n",
    "!pip install --quiet git+https://github.com/oughtinc/ergo.git\n",
    "!pip install --quiet pendulum requests\n",
    "!pip install --quiet torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 50
    },
    "colab_type": "code",
    "id": "42KVEbypmg-k",
    "outputId": "636b1af6-11c6-4d16-ddaa-a167df115675"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The google.colab.data_table extension is already loaded. To reload it, use:\n",
      "  %reload_ext google.colab.data_table\n"
     ]
    }
   ],
   "source": [
    "%load_ext google.colab.data_table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YvWn9iJqmilB"
   },
   "outputs": [],
   "source": [
    "import ergo\n",
    "import pendulum\n",
    "import requests\n",
    "import torch\n",
    "import pandas as pd\n",
    "\n",
    "from typing import List\n",
    "from pendulum import Date, Duration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rMz4hIhedK-_"
   },
   "outputs": [],
   "source": [
    "def get_questions_for_cat(cat):\n",
    "  r = requests.get(f\"https://pandemic.metaculus.com/api2/questions/?search=cat:{cat}\")\n",
    "  return r.json()\n",
    "\n",
    "def get_question_ids_for_cat(cat):\n",
    "  qs = get_questions_for_cat(cat)\n",
    "  return [q[\"id\"] for q in qs[\"results\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "PFAueEx3cgsD"
   },
   "outputs": [],
   "source": [
    "def show_related_questions(related_question_ids):\n",
    "  related_questions = []\n",
    "  for id in related_question_ids:\n",
    "    # try-except there b/c some \"questions\" aren't really questions and throw erros w/ metaculus\n",
    "    # like https://pandemic.metaculus.com/questions/3957/lockdown-series-when-will-life-return-to-normal-ish/\n",
    "    try:\n",
    "      rq = metaculus.get_question(id) \n",
    "      related_questions.append(rq)\n",
    "    except:\n",
    "      pass\n",
    "\n",
    "  def needs_attention(related_question):\n",
    "    if not question.my_predictions:\n",
    "      return True\n",
    "    else:\n",
    "      period = pendulum.from_timestamp(\n",
    "          related_question.prediction_timeseries[-1][\"t\"]) - \\\n",
    "          pendulum.from_timestamp(question.my_predictions[\"predictions\"][-1][\"t\"]\n",
    "      )\n",
    "\n",
    "    return period.in_seconds() > 1\n",
    "\n",
    "  def direction(q):\n",
    "    i = 1\n",
    "    have_gone_far_back_enough = True\n",
    "\n",
    "    while (have_gone_far_back_enough):\n",
    "      try:\n",
    "        period = pendulum.from_timestamp(\n",
    "            q.prediction_timeseries[-i][\"t\"]) - \\\n",
    "            pendulum.from_timestamp(question.my_predictions[\"predictions\"][-1][\"t\"]\n",
    "        )\n",
    "      except:\n",
    "        i = i-1\n",
    "        break\n",
    "\n",
    "      if (period.in_seconds() < 1):\n",
    "        have_gone_far_back_enough = False\n",
    "      else:\n",
    "        i = i+1\n",
    "\n",
    "    if (isinstance(q.prediction_timeseries[-i][\"community_prediction\"], float)):\n",
    "      old = q.prediction_timeseries[-i][\"community_prediction\"]\n",
    "      new = q.prediction_timeseries[-1][\"community_prediction\"]\n",
    "      return old - new\n",
    "      # if (q.prediction_timeseries[-2][\"community_prediction\"] > q.prediction_timeseries[-1][\"community_prediction\"]):\n",
    "      #   return new - old\n",
    "      # elif (q.prediction_timeseries[-2][\"community_prediction\"] > q.prediction_timeseries[-1][\"community_prediction\"]):\n",
    "      #   return \"+\"\n",
    "      # else:\n",
    "      #   return \"=\"\n",
    "    else:\n",
    "      old = q.prediction_timeseries[-i][\"community_prediction\"][\"q2\"]\n",
    "      new = q.prediction_timeseries[-1][\"community_prediction\"][\"q2\"]\n",
    "      return old - new\n",
    "      # if (q.prediction_timeseries[-2][\"community_prediction\"][\"q2\"] > q.prediction_timeseries[-1][\"community_prediction\"][\"q2\"]):\n",
    "      #   return \"-\"\n",
    "      # elif (q.prediction_timeseries[-2][\"community_prediction\"][\"q2\"] > q.prediction_timeseries[-1][\"community_prediction\"][\"q2\"]):\n",
    "      #   return \"+\"\n",
    "      # else:\n",
    "      #   return \"=\"\n",
    "\n",
    "  related_questions_data = [\n",
    "    [\n",
    "     related_question.id, \n",
    "     needs_attention(related_question), \n",
    "     float(direction(related_question)),\n",
    "     related_question.data[\"possibilities\"][\"scale\"][\"min\"],\n",
    "     related_question.data[\"possibilities\"][\"scale\"][\"max\"],\n",
    "     (pendulum.period(pendulum.parse(related_question.data[\"possibilities\"][\"scale\"][\"min\"]), pendulum.parse(related_question.data[\"possibilities\"][\"scale\"][\"max\"])) * float(direction(related_question))).in_hours(),\n",
    "     related_question.title, \n",
    "     f\"https://pandemic.metaculus.com{related_question.page_url}\"]\n",
    "    for related_question\n",
    "    in related_questions\n",
    "  ]\n",
    "\n",
    "  print(related_questions[0].data[\"possibilities\"][\"scale\"][\"max\"])\n",
    "  print(related_questions[0].data[\"possibilities\"][\"scale\"][\"min\"])\n",
    "\n",
    "  df = pd.DataFrame(related_questions_data, columns=[\"id\", \"changed\", \"degree of change\", \"min\", \"max\", \"period\", \"title\", \"link\"])\n",
    "\n",
    "  return df.set_index(\"id\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6yZvXpjXfIOW"
   },
   "outputs": [],
   "source": [
    "def relate_questions_in_cat(cat):\n",
    "  related_question_ids = get_question_ids_for_cat(cat)\n",
    "  return show_related_questions(related_question_ids)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "DfU-SaV5mqqe"
   },
   "source": [
    "# Questions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "LQq8q8U3m5im"
   },
   "source": [
    "Here is the question we want to forecast:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BCJwi2JVmwEv"
   },
   "outputs": [],
   "source": [
    "question_data = {\n",
    "    \"id\": 3925,\n",
    "    \"name\": \"How many days will the average American spend under lockdown between 2020-03-25 and 2020-04-24\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "colab_type": "code",
    "id": "gs9BAHl8M28N",
    "outputId": "fc6f9dd5-69bb-4207-e411-4b18f3af1d9b"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.module+javascript": "\n      import \"https://ssl.gstatic.com/colaboratory/data_table/a6224c040fa35dcf/data_table.js\";\n\n      window.createDataTable({\n        data: [[{\n            'v': 3925,\n            'f': \"3925\",\n        },\n\"How many days will the average American spend under lockdown between 2020-03-25 and 2020-04-24\"]],\n        columns: [[\"number\", \"id\"], [\"string\", \"name\"]],\n        columnOptions: [{\"width\": \"1px\", \"className\": \"index_column\"}],\n        rowsPerPage: 25,\n        helpUrl: \"https://colab.research.google.com/notebooks/data_table.ipynb\",\n        suppressOutputScrolling: true,\n        minimumWidth: undefined,\n      });\n    ",
      "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 th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>3925</th>\n",
       "      <td>How many days will the average American spend ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                   name\n",
       "id                                                     \n",
       "3925  How many days will the average American spend ..."
      ]
     },
     "execution_count": 108,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metaculus = ergo.Metaculus(api_domain=\"pandemic\")\n",
    "metaculus.login_via_username_and_password(username=\"oughtpublic\", password=\"123456\")\n",
    "question = metaculus.get_question(question_data[\"id\"], name=question_data[\"name\"])\n",
    "\n",
    "df = pd.DataFrame([[question.id, question.name]], columns=[\"id\", \"name\"])\n",
    "df.set_index(\"id\")\n",
    "\n",
    "# p = pendulum.instance(question.last_activity_time) - pendulum.from_timestamp(question.prediction_timeseries[-1][\"t\"])\n",
    "# p.in_seconds() < 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hZR00LQjU21V"
   },
   "source": [
    "# Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "A7RN2JHaU6Ie"
   },
   "source": [
    "Data: https://www.nytimes.com/interactive/2020/us/coronavirus-stay-at-home-order.html\n",
    "\n",
    "Manually copied in on 2020-04-05\n",
    "\n",
    "Updated on 2020-04-08:\n",
    "\n",
    "\n",
    "*   Added three regions to Oklahoma: Claremore, Moore, and Sallisaw.\n",
    "*   Added South Carolina state-wide lockdown\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "c9eb4mdcU6N5"
   },
   "outputs": [],
   "source": [
    "M = 1000000\n",
    "\n",
    "USA_data = {\n",
    "    \"start\": False,\n",
    "    \"pop\": 327.2 * M,\n",
    "    \"regions\": {\n",
    "        \"Alabama\": {\n",
    "            \"start\": pendulum.Date(2020, 4, 4),\n",
    "            \"pop\": 4.9 * M\n",
    "        },\n",
    "        \"Alaska\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 31),\n",
    "            \"pop\": 0.737 * M\n",
    "        },\n",
    "        \"Arizona\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 31),\n",
    "            \"pop\": 7.2 * M\n",
    "        },\n",
    "        \"California\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 19),\n",
    "            \"pop\": 39.6 * M\n",
    "        },\n",
    "        \"Colorado\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 26),\n",
    "            \"pop\": 5.7 * M\n",
    "        },\n",
    "        \"Connecticut\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 23),\n",
    "            \"pop\": 3.6 * M\n",
    "        },\n",
    "        \"Delaware\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 24),\n",
    "            \"pop\": 0.973 * M\n",
    "        },\n",
    "        \"District of Columbia\": {\n",
    "            \"start\": pendulum.Date(2020, 4, 1),\n",
    "            \"pop\": 0.702 * M\n",
    "        },\n",
    "        \"Florida\": {\n",
    "            \"start\": pendulum.Date(2020, 4, 3),\n",
    "            \"pop\": 21.5 * M\n",
    "        },\n",
    "        \"Georgia\": {\n",
    "            \"start\": pendulum.Date(2020, 4, 3),\n",
    "            \"pop\": 10.6 * M\n",
    "        },\n",
    "        \"Hawaii\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 25),\n",
    "            \"pop\": 1.4 * M\n",
    "        },\n",
    "        \"Idaho\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 25),\n",
    "            \"pop\": 1.8 * M\n",
    "        },\n",
    "        \"Illinois\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 21),\n",
    "            \"pop\": 12.7 * M\n",
    "        },\n",
    "        \"Indiana\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 24),\n",
    "            \"pop\": 6.7 * M\n",
    "        },\n",
    "        \"Kansas\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 30),\n",
    "            \"pop\": 2.9 * M\n",
    "        },\n",
    "        \"Kentucky\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 26),\n",
    "            \"pop\": 4.5 * M\n",
    "        },\n",
    "        \"Louisiana\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 23),\n",
    "            \"pop\": 4.6 * M\n",
    "        },\n",
    "        \"Maine\": {\n",
    "            \"start\": pendulum.Date(2020, 4, 2),\n",
    "            \"pop\": 1.3 * M\n",
    "        },\n",
    "        \"Maryland\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 30),\n",
    "            \"pop\": 6 * M\n",
    "        },\n",
    "        \"Massachusetts\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 24),\n",
    "            \"pop\": 6.9 * M\n",
    "        },\n",
    "        \"Michigan\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 24),\n",
    "            \"pop\": 10 * M\n",
    "        },\n",
    "        \"Minnesota\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 27),\n",
    "            \"pop\": 5.6 * M\n",
    "        },\n",
    "        \"Mississippi\": {\n",
    "            \"start\": pendulum.Date(2020, 4, 3),\n",
    "            \"pop\": 3 * M\n",
    "        },\n",
    "        \"Missouri\": {\n",
    "            \"start\": pendulum.Date(2020, 4, 6),\n",
    "            \"pop\": 6.1 * M\n",
    "        },\n",
    "        \"Montana\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 28),\n",
    "            \"pop\": 1.1 * M\n",
    "        },\n",
    "        \"Nevada\": {\n",
    "            \"start\": pendulum.Date(2020, 4, 1),\n",
    "            \"pop\": 3.1 * M\n",
    "        },\n",
    "        \"New Hampshire\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 27),\n",
    "            \"pop\": 1.4 * M\n",
    "        },\n",
    "        \"New Jersey\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 21),\n",
    "            \"pop\": 8.9 * M\n",
    "        },\n",
    "        \"New Mexico\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 24),\n",
    "            \"pop\": 2.1 * M\n",
    "        },\n",
    "        \"New York\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 22),\n",
    "            \"pop\": 19.5 * M\n",
    "        },\n",
    "        \"North Carolina\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 30),\n",
    "            \"pop\": 10.4 * M\n",
    "        },\n",
    "        \"Ohio\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 23),\n",
    "            \"pop\": 11.7 * M\n",
    "        },\n",
    "        \"Oklahoma\": {\n",
    "            \"start\": False,\n",
    "            \"pop\": 3.9 * M,\n",
    "            \"regions\": {\n",
    "                \"Claremore\": {\n",
    "                    \"start\": pendulum.Date(2020, 4, 6),\n",
    "                    \"pop\": 0.019 * M\n",
    "                },\n",
    "                \"Edmond\": {\n",
    "                    \"start\": pendulum.Date(2020, 3, 30),\n",
    "                    \"pop\": 0.093 * M\n",
    "                },\n",
    "                \"Moore\": {\n",
    "                    \"start\": pendulum.Date(2020, 4, 4),\n",
    "                    \"pop\": 0.062 * M\n",
    "                },\n",
    "                \"Norman\": {\n",
    "                    \"start\": pendulum.Date(2020, 3, 25),\n",
    "                    \"pop\": 0.123 * M\n",
    "                },\n",
    "                \"Oklahoma City\": {\n",
    "                    \"start\": pendulum.Date(2020, 3, 28),\n",
    "                    \"pop\": 0.649 * M\n",
    "                },\n",
    "                \"Sallisaw\": {\n",
    "                    \"start\": pendulum.Date(2020, 4, 4),\n",
    "                    \"pop\": 0.009 * M\n",
    "                },\n",
    "                \"Stillwater\": {\n",
    "                    \"start\": pendulum.Date(2020, 3, 30),\n",
    "                    \"pop\": 0.05 * M\n",
    "                },\n",
    "                \"Tulsa\": {\n",
    "                    \"start\": pendulum.Date(2020, 3, 28),\n",
    "                    \"pop\": 0.401 * M\n",
    "                }\n",
    "            }\n",
    "        },\n",
    "        \"Oregon\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 23),\n",
    "            \"pop\": 4.2 * M\n",
    "        },\n",
    "        \"Pennsylvania\": {\n",
    "            \"start\": pendulum.Date(2020, 4, 1),\n",
    "            \"pop\": 12.8 * M\n",
    "        },\n",
    "        \"Puerto Rico\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 15),\n",
    "            \"pop\": 3.2 * M\n",
    "        },\n",
    "        \"Rhode Island\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 28),\n",
    "            \"pop\": 1.1 * M\n",
    "        },\n",
    "        \"South Carolina\": {\n",
    "            \"start\": pendulum.Date(2020, 4, 7),\n",
    "            \"pop\": 5.1 * M,\n",
    "            \"regions\": {\n",
    "                \"Charleston\": {\n",
    "                    \"start\": pendulum.Date(2020, 3, 26),\n",
    "                    \"pop\": 0.136 * M\n",
    "                },\n",
    "                \"Columbia\": {\n",
    "                    \"start\": pendulum.Date(2020, 3, 29),\n",
    "                    \"pop\": 0.133 * M\n",
    "                }\n",
    "            }\n",
    "        },\n",
    "        \"Tennessee\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 31),\n",
    "            \"pop\": 6.8 * M\n",
    "        },\n",
    "        \"Texas\": {\n",
    "            \"start\": pendulum.Date(2020, 4, 2),\n",
    "            \"pop\": 29 * M\n",
    "        },\n",
    "        \"Utah\": {\n",
    "            \"start\": False,\n",
    "            \"pop\": 3.2 * M,\n",
    "            \"regions\": {\n",
    "                \"Davis County\": {\n",
    "                    \"start\": pendulum.Date(2020, 4, 1),\n",
    "                    \"pop\": 0.352 * M\n",
    "                },\n",
    "                \"Salt Lake County\": {\n",
    "                    \"start\": pendulum.Date(2020, 3, 30),\n",
    "                    \"pop\": 1.2 * M\n",
    "                },\n",
    "                \"Summit County\": {\n",
    "                    \"start\": pendulum.Date(2020, 3, 27),\n",
    "                    \"pop\": 0.042 * M\n",
    "                }\n",
    "            }\n",
    "        },\n",
    "        \"Vermont\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 25),\n",
    "            \"pop\": 0.626 * M\n",
    "        },\n",
    "        \"Virginia\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 30),\n",
    "            \"pop\": 8.5 * M\n",
    "        },\n",
    "        \"Washington\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 23),\n",
    "            \"pop\": 7.5 * M\n",
    "        },\n",
    "        \"West Virginia\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 24),\n",
    "            \"pop\": 1.8 * M\n",
    "        },\n",
    "        \"Wisconsin\": {\n",
    "            \"start\": pendulum.Date(2020, 3, 25),\n",
    "            \"pop\": 5.8 * M\n",
    "        },\n",
    "        \"Wyoming\": {\n",
    "            \"start\": False,\n",
    "            \"pop\": 0.578 * M,\n",
    "            \"regions\": {\n",
    "                \"Jackson\": {\n",
    "                    \"start\": pendulum.Date(2020, 3, 28),\n",
    "                    \"pop\": 0.01 * M\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "kshvxjMkjPn9"
   },
   "source": [
    "#Assumptions\n",
    "\n",
    "* Every place that is currently on lockdown had no subregion on lockdown prior to implementing the region-wide lockdown. We know this if false, but it's simple.\n",
    "\n",
    "* Every place that is currently on lockdown will remain on lockdown through at least Apirl 25.\n",
    "\n",
    "* Any state that isn't fullly locked down will have a 1% of entering a lockdown every day."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "W53hmaDJrn3Q"
   },
   "outputs": [],
   "source": [
    "chance_of_full_lockdown_transition = 0.01"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "PRBe5Kt3F7YM"
   },
   "source": [
    "# Related Questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 584
    },
    "colab_type": "code",
    "id": "IdSIUWqeF-Vj",
    "outputId": "93370fd2-fae8-4328-edbd-68c637da4fd4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-12-27\n",
      "2020-03-27\n"
     ]
    },
    {
     "data": {
      "application/vnd.google.colaboratory.module+javascript": "\n      import \"https://ssl.gstatic.com/colaboratory/data_table/a6224c040fa35dcf/data_table.js\";\n\n      window.createDataTable({\n        data: [[{\n            'v': 3935,\n            'f': \"3935\",\n        },\ntrue,\n{\n            'v': -0.0006099999999999994,\n            'f': \"-0.0006099999999999994\",\n        },\n\"2020-03-27\",\n\"2021-12-27\",\n{\n            'v': -9,\n            'f': \"-9\",\n        },\n\"When will the current lockdown in the UK be lifted or reduced in intensity?\",\n\"https://pandemic.metaculus.com/questions/3935/when-will-the-current-lockdown-in-the-uk-be-lifted-or-reduced-in-intensity/\"],\n [{\n            'v': 3936,\n            'f': \"3936\",\n        },\nfalse,\n{\n            'v': 0.0,\n            'f': \"0.0\",\n        },\n\"2020-03-27\",\n\"2021-12-27\",\n{\n            'v': 0,\n            'f': \"0\",\n        },\n\"When will the current lockdown in Spain be lifted or reduced in intensity?\",\n\"https://pandemic.metaculus.com/questions/3936/when-will-the-current-lockdown-in-spain-be-lifted-or-reduced-in-intensity/\"],\n [{\n            'v': 3937,\n            'f': \"3937\",\n        },\nfalse,\n{\n            'v': 0.0,\n            'f': \"0.0\",\n        },\n\"2020-03-27\",\n\"2021-12-27\",\n{\n            'v': 0,\n            'f': \"0\",\n        },\n\"When will the current lockdown in Italy be lifted or reduced in intensity?\",\n\"https://pandemic.metaculus.com/questions/3937/when-will-the-current-lockdown-in-italy-be-lifted-or-reduced-in-intensity/\"],\n [{\n            'v': 3938,\n            'f': \"3938\",\n        },\nfalse,\n{\n            'v': 0.0,\n            'f': \"0.0\",\n        },\n\"2020-03-27\",\n\"2021-12-27\",\n{\n            'v': 0,\n            'f': \"0\",\n        },\n\"When will the current partial lockdown in the State of S\\u00e3o Paulo (Brazil) be lifted or reduced in intensity?\",\n\"https://pandemic.metaculus.com/questions/3938/when-will-the-current-partial-lockdown-in-the-state-of-s%25C3%25A3o-paulo-brazil-be-lifted-or-reduced-in-intensity/\"],\n [{\n            'v': 3939,\n            'f': \"3939\",\n        },\nfalse,\n{\n            'v': 0.0,\n            'f': \"0.0\",\n        },\n\"2020-03-27\",\n\"2021-12-27\",\n{\n            'v': 0,\n            'f': \"0\",\n        },\n\"When will the current lockdown in the state of California be lifted or reduced in intensity?\",\n\"https://pandemic.metaculus.com/questions/3939/when-will-the-current-lockdown-in-the-state-of-california-be-lifted-or-reduced-in-intensity/\"],\n [{\n            'v': 3940,\n            'f': \"3940\",\n        },\ntrue,\n{\n            'v': -0.0005899999999999933,\n            'f': \"-0.0005899999999999933\",\n        },\n\"2020-03-27\",\n\"2021-12-27\",\n{\n            'v': -9,\n            'f': \"-9\",\n        },\n\"When will the current partial lockdown in the Province of Punjab (Pakistan) be lifted or reduced in intensity?\",\n\"https://pandemic.metaculus.com/questions/3940/when-will-the-current-partial-lockdown-in-the-province-of-punjab-pakistan-be-lifted-or-reduced-in-intensity/\"],\n [{\n            'v': 3941,\n            'f': \"3941\",\n        },\nfalse,\n{\n            'v': 0.0,\n            'f': \"0.0\",\n        },\n\"2020-03-27\",\n\"2021-12-27\",\n{\n            'v': 0,\n            'f': \"0\",\n        },\n\"When will the current lockdown in Poland be lifted or reduced in intensity?\",\n\"https://pandemic.metaculus.com/questions/3941/when-will-the-current-lockdown-in-poland-be-lifted-or-reduced-in-intensity/\"],\n [{\n            'v': 3942,\n            'f': \"3942\",\n        },\ntrue,\n{\n            'v': -0.000899999999999998,\n            'f': \"-0.000899999999999998\",\n        },\n\"2020-03-27\",\n\"2021-12-27\",\n{\n            'v': -13,\n            'f': \"-13\",\n        },\n\"When will the current total lockdown in South Africa be lifted or reduced in intensity?\",\n\"https://pandemic.metaculus.com/questions/3942/when-will-the-current-total-lockdown-in-south-africa-be-lifted-or-reduced-in-intensity/\"],\n [{\n            'v': 3946,\n            'f': \"3946\",\n        },\ntrue,\n{\n            'v': -0.0005299999999999888,\n            'f': \"-0.0005299999999999888\",\n        },\n\"2020-03-27\",\n\"2021-12-27\",\n{\n            'v': -8,\n            'f': \"-8\",\n        },\n\"When will the current lockdown in India be lifted or reduced in intensity?\",\n\"https://pandemic.metaculus.com/questions/3946/when-will-the-current-lockdown-in-india-be-lifted-or-reduced-in-intensity/\"],\n [{\n            'v': 3948,\n            'f': \"3948\",\n        },\ntrue,\n{\n            'v': -0.0006400000000000017,\n            'f': \"-0.0006400000000000017\",\n        },\n\"2020-03-27\",\n\"2021-12-27\",\n{\n            'v': -9,\n            'f': \"-9\",\n        },\n\"When will the current lockdown in France be lifted or reduced in intensity?\",\n\"https://pandemic.metaculus.com/questions/3948/when-will-the-current-lockdown-in-france-be-lifted-or-reduced-in-intensity/\"],\n [{\n            'v': 3951,\n            'f': \"3951\",\n        },\nfalse,\n{\n            'v': 0.0,\n            'f': \"0.0\",\n        },\n\"2020-03-27\",\n\"2021-12-27\",\n{\n            'v': 0,\n            'f': \"0\",\n        },\n\"When will the current lockdown in New Zealand be lifted or reduced in intensity?\",\n\"https://pandemic.metaculus.com/questions/3951/when-will-the-current-lockdown-in-new-zealand-be-lifted-or-reduced-in-intensity/\"],\n [{\n            'v': 3952,\n            'f': \"3952\",\n        },\ntrue,\n{\n            'v': -0.0003599999999999992,\n            'f': \"-0.0003599999999999992\",\n        },\n\"2020-03-27\",\n\"2021-12-27\",\n{\n            'v': -5,\n            'f': \"-5\",\n        },\n\"When will the current total lockdown in the State of New York be lifted or reduced in intensity?\",\n\"https://pandemic.metaculus.com/questions/3952/when-will-the-current-total-lockdown-in-the-state-of-new-york-be-lifted-or-reduced-in-intensity/\"]],\n        columns: [[\"number\", \"id\"], [\"string\", \"changed\"], [\"number\", \"degree of change\"], [\"string\", \"min\"], [\"string\", \"max\"], [\"number\", \"period\"], [\"string\", \"title\"], [\"string\", \"link\"]],\n        columnOptions: [{\"width\": \"1px\", \"className\": \"index_column\"}],\n        rowsPerPage: 25,\n        helpUrl: \"https://colab.research.google.com/notebooks/data_table.ipynb\",\n        suppressOutputScrolling: true,\n        minimumWidth: undefined,\n      });\n    ",
      "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 th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>changed</th>\n",
       "      <th>degree of change</th>\n",
       "      <th>min</th>\n",
       "      <th>max</th>\n",
       "      <th>period</th>\n",
       "      <th>title</th>\n",
       "      <th>link</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>3935</th>\n",
       "      <td>True</td>\n",
       "      <td>-0.00061</td>\n",
       "      <td>2020-03-27</td>\n",
       "      <td>2021-12-27</td>\n",
       "      <td>-9</td>\n",
       "      <td>When will the current lockdown in the UK be li...</td>\n",
       "      <td>https://pandemic.metaculus.com/questions/3935/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3936</th>\n",
       "      <td>False</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>2020-03-27</td>\n",
       "      <td>2021-12-27</td>\n",
       "      <td>0</td>\n",
       "      <td>When will the current lockdown in Spain be lif...</td>\n",
       "      <td>https://pandemic.metaculus.com/questions/3936/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3937</th>\n",
       "      <td>False</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>2020-03-27</td>\n",
       "      <td>2021-12-27</td>\n",
       "      <td>0</td>\n",
       "      <td>When will the current lockdown in Italy be lif...</td>\n",
       "      <td>https://pandemic.metaculus.com/questions/3937/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3938</th>\n",
       "      <td>False</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>2020-03-27</td>\n",
       "      <td>2021-12-27</td>\n",
       "      <td>0</td>\n",
       "      <td>When will the current partial lockdown in the ...</td>\n",
       "      <td>https://pandemic.metaculus.com/questions/3938/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3939</th>\n",
       "      <td>False</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>2020-03-27</td>\n",
       "      <td>2021-12-27</td>\n",
       "      <td>0</td>\n",
       "      <td>When will the current lockdown in the state of...</td>\n",
       "      <td>https://pandemic.metaculus.com/questions/3939/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3940</th>\n",
       "      <td>True</td>\n",
       "      <td>-0.00059</td>\n",
       "      <td>2020-03-27</td>\n",
       "      <td>2021-12-27</td>\n",
       "      <td>-9</td>\n",
       "      <td>When will the current partial lockdown in the ...</td>\n",
       "      <td>https://pandemic.metaculus.com/questions/3940/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3941</th>\n",
       "      <td>False</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>2020-03-27</td>\n",
       "      <td>2021-12-27</td>\n",
       "      <td>0</td>\n",
       "      <td>When will the current lockdown in Poland be li...</td>\n",
       "      <td>https://pandemic.metaculus.com/questions/3941/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3942</th>\n",
       "      <td>True</td>\n",
       "      <td>-0.00090</td>\n",
       "      <td>2020-03-27</td>\n",
       "      <td>2021-12-27</td>\n",
       "      <td>-13</td>\n",
       "      <td>When will the current total lockdown in South ...</td>\n",
       "      <td>https://pandemic.metaculus.com/questions/3942/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3946</th>\n",
       "      <td>True</td>\n",
       "      <td>-0.00053</td>\n",
       "      <td>2020-03-27</td>\n",
       "      <td>2021-12-27</td>\n",
       "      <td>-8</td>\n",
       "      <td>When will the current lockdown in India be lif...</td>\n",
       "      <td>https://pandemic.metaculus.com/questions/3946/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3948</th>\n",
       "      <td>True</td>\n",
       "      <td>-0.00064</td>\n",
       "      <td>2020-03-27</td>\n",
       "      <td>2021-12-27</td>\n",
       "      <td>-9</td>\n",
       "      <td>When will the current lockdown in France be li...</td>\n",
       "      <td>https://pandemic.metaculus.com/questions/3948/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3951</th>\n",
       "      <td>False</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>2020-03-27</td>\n",
       "      <td>2021-12-27</td>\n",
       "      <td>0</td>\n",
       "      <td>When will the current lockdown in New Zealand ...</td>\n",
       "      <td>https://pandemic.metaculus.com/questions/3951/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3952</th>\n",
       "      <td>True</td>\n",
       "      <td>-0.00036</td>\n",
       "      <td>2020-03-27</td>\n",
       "      <td>2021-12-27</td>\n",
       "      <td>-5</td>\n",
       "      <td>When will the current total lockdown in the St...</td>\n",
       "      <td>https://pandemic.metaculus.com/questions/3952/...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      changed  ...                                               link\n",
       "id             ...                                                   \n",
       "3935     True  ...  https://pandemic.metaculus.com/questions/3935/...\n",
       "3936    False  ...  https://pandemic.metaculus.com/questions/3936/...\n",
       "3937    False  ...  https://pandemic.metaculus.com/questions/3937/...\n",
       "3938    False  ...  https://pandemic.metaculus.com/questions/3938/...\n",
       "3939    False  ...  https://pandemic.metaculus.com/questions/3939/...\n",
       "3940     True  ...  https://pandemic.metaculus.com/questions/3940/...\n",
       "3941    False  ...  https://pandemic.metaculus.com/questions/3941/...\n",
       "3942     True  ...  https://pandemic.metaculus.com/questions/3942/...\n",
       "3946     True  ...  https://pandemic.metaculus.com/questions/3946/...\n",
       "3948     True  ...  https://pandemic.metaculus.com/questions/3948/...\n",
       "3951    False  ...  https://pandemic.metaculus.com/questions/3951/...\n",
       "3952     True  ...  https://pandemic.metaculus.com/questions/3952/...\n",
       "\n",
       "[12 rows x 7 columns]"
      ]
     },
     "execution_count": 125,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relate_questions_in_cat(\"internal--lockdown-series\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-QBpZp2StFrg"
   },
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "F2gOVOp79H-I"
   },
   "outputs": [],
   "source": [
    "def get_pop_in_lockdown_on_date(region, date):\n",
    "  is_in_future = date > pendulum.now().date()\n",
    "\n",
    "  if (region[\"start\"] and region[\"start\"] < date):\n",
    "    return region[\"pop\"]\n",
    "  elif (is_in_future and ergo.flip(chance_of_full_lockdown_transition)):\n",
    "    region[\"start\"] = pendulum.now().date()\n",
    "    return region[\"pop\"]\n",
    "  elif (\"regions\" in region):\n",
    "    pop_in_lockdown = 0\n",
    "    for region_key in region[\"regions\"]:\n",
    "      sub_region = region[\"regions\"][region_key]\n",
    "      pop_in_lockdown += get_pop_in_lockdown_on_date(sub_region, date)\n",
    "    return pop_in_lockdown\n",
    "  else:\n",
    "    return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wQh2G6IZruY8"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def get_avg_pop_in_lockdown_over_range(region, start_date, end_date):\n",
    "  period = pendulum.period(start_date, end_date)\n",
    "  arr = np.array([])\n",
    "  for dt in period.range(\"days\"):\n",
    "    arr = np.append(arr, get_pop_in_lockdown_on_date(region, dt))\n",
    "  return np.mean(arr)\n",
    "\n",
    "def get_avg_proportion_in_lockdown_over_range(region, start_date, end_date):\n",
    "  return get_avg_pop_in_lockdown_over_range(region, start_date, end_date) / region[\"pop\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "S2_i58Wnrymu"
   },
   "source": [
    "We need to make a deep copy of the region data for each model so that model-specific changes don't carry over to the next run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OUxF_TKufS1v"
   },
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "def model():\n",
    "  usa_region = copy.deepcopy(USA_data)\n",
    "  avg = get_avg_proportion_in_lockdown_over_range(usa_region, pendulum.Date(2020, 3, 25), pendulum.Date(2020, 4, 25))\n",
    "  ergo.tag(torch.Tensor([avg*32]), question.name) # 3/25 to 4/25 inclusive is 32 days\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "K1J12dOLggZs"
   },
   "outputs": [],
   "source": [
    "samples = ergo.run(lambda: model(), num_samples=10)\n",
    "\n",
    "samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xr94tXrFODuh"
   },
   "source": [
    "# Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "UfGruRDTOgYC"
   },
   "source": [
    "Histogram:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "FQf1IjKYOiqh"
   },
   "outputs": [],
   "source": [
    "samples.hist(column=question.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "D9x4kW64OEMN"
   },
   "source": [
    "Summary stats:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YI7IsER5yXrM"
   },
   "outputs": [],
   "source": [
    "samples.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "TiV8zqstRfpW"
   },
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "mQDqb5FKcIaj"
   },
   "source": [
    "# Submit predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Ov__zuhYvlS3"
   },
   "source": [
    "Convert samples to Metaculus distributions and visualize:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LiwHiT3WezEu"
   },
   "outputs": [],
   "source": [
    "if question.name in samples:\n",
    "  question.show_submission(samples[question.name])\n",
    "else:\n",
    "  print(f\"No predictions for {question.name}\")\n",
    "print(\"\\n\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "iC_rceYye_lx"
   },
   "source": [
    "If everything looks good, submit the predictions!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LzwGTAFZcIhw"
   },
   "outputs": [],
   "source": [
    "# def submit_all():\n",
    "#   for question in questions:\n",
    "#     if question.name in samples:\n",
    "#       try:\n",
    "#         params = question.submit_from_samples(samples[question.name])\n",
    "#         print(f\"Submitted for {question.name}\")\n",
    "#         print(f\"https://pandemic.metaculus.com{question.page_url}\")\n",
    "#       except requests.exceptions.HTTPError as e:\n",
    "#         print(f\"Couldn't make prediction for {question.name} -- maybe this question is now closed? See error below.\")\n",
    "#         print(e)\n",
    "#     else:\n",
    "#       print(f\"No predictions for {question.name}\")\n",
    "\n",
    "# submit_all()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hCZ-mVvtlLil"
   },
   "source": [
    "# To do"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hqZ81Z_JlLpf"
   },
   "source": [
    "- Add to dos"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "DfU-SaV5mqqe",
    "hZR00LQjU21V"
   ],
   "include_colab_link": true,
   "name": "How long will the average American spend under lockdown?",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}