Simulated-Data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-25T09:11:17.321955Z",
     "start_time": "2023-01-25T09:11:17.289189Z"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "units = range(1, 500+1)\n",
    "\n",
    "np.random.seed(1)\n",
    "\n",
    "data = pd.DataFrame(dict(\n",
    "    store = np.repeat(units, 4),\n",
    "    unit_fe = np.repeat(np.random.normal(0, 5, size=len(units)), 4),\n",
    "    weeks_to_xmas = np.tile([3,2,1,0], len(units)),\n",
    "    avg_week_sales = np.repeat(np.random.gamma(20, 1, size=len(units)), 4).round(2),\n",
    ")\n",
    ").assign(\n",
    "    is_on_sale = lambda d: (\n",
    "         (d[\"unit_fe\"] > 0)*np.random.binomial(1, 0.2, d.shape[0])\n",
    "        | (d[\"avg_week_sales\"] > np.random.gamma(25, 1, size=len(units)*4))\n",
    "        | (np.random.normal(d[\"weeks_to_xmas\"], 2) > 2)*np.random.binomial(1, 0.7, size=d.shape[0]) # xmas\n",
    "    )*1,\n",
    ").assign(\n",
    "    y0 = lambda d: (-10 + 10*d[\"unit_fe\"] + d[\"avg_week_sales\"] \n",
    "                    + 2*d[\"avg_week_sales\"]*d[\"weeks_to_xmas\"])\n",
    ").assign(\n",
    "    y1 = lambda d: d[\"y0\"] + 50\n",
    ").assign(\n",
    "    tau = lambda d: d[\"y1\"] - d[\"y0\"],\n",
    "    weekly_amount_sold = lambda d: np.random.normal(np.where(d[\"is_on_sale\"] == 1, d[\"y1\"], d[\"y0\"]), 10).clip(0, np.inf).round(2)\n",
    ").drop(columns=[\"unit_fe\", \"y0\", \"y1\", \"tau\"])\n",
    "\n",
    "data.to_csv(\"./causal-inference-in-python/data/xmas_sales.csv\", index=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-24T19:58:50.923765Z",
     "start_time": "2023-01-24T19:58:50.362564Z"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "pd.set_option('display.max_rows', 5)\n",
    "\n",
    "np.random.seed(123)\n",
    "data = (\n",
    "    pd.read_csv(\"./causal-inference-in-python/data/online_classroom.csv\")\n",
    "    .assign(cross_sell_email = lambda d: np.select(\n",
    "        [d[\"format_ol\"].astype(bool), d[\"format_blended\"].astype(bool)],\n",
    "        [\"no_email\", \"long\"],\n",
    "        default=\"short\")\n",
    "           )\n",
    "    .assign(age = lambda d: np.random.gamma(1.5, 5, d.shape[0]).astype(int)+14)\n",
    "    .assign(conversion = lambda d: ((d[\"falsexam\"]-d[\"age\"]) > 70).astype(int))\n",
    "    .drop(columns=[\"asian\", \"black\", \"hawaiian\", \"hispanic\", \"unknown\", \"white\", \"format_ol\", \"format_blended\", \"falsexam\"])\n",
    ")\n",
    "\n",
    "data.to_csv(\"./causal-inference-in-python/data/cross_sell_email.csv\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-24T19:58:50.937159Z",
     "start_time": "2023-01-24T19:58:50.925222Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "data = (\n",
    "    pd.read_csv(\"./causal-inference-in-python/data/online_classroom.csv\")\n",
    "    .assign(recommender = lambda d: np.select(\n",
    "        [d[\"format_ol\"].astype(bool), d[\"format_blended\"].astype(bool)],\n",
    "        [\"benchmark\", \"benchmark\"],\n",
    "        default=\"challenger\"))\n",
    "    .assign(age = lambda d: np.random.gamma(1.5, 5, d.shape[0]).astype(int)+14)\n",
    "    .assign(tenure = lambda d: np.random.gamma(1, 1.1, d.shape[0]).astype(int).clip(0, 4))\n",
    "    .assign(watch_time = lambda d: np.random.normal(3*(d[\"recommender\"]==\"challenger\") + 10*d[\"tenure\"] - 0.8*d[\"age\"] + 0.001*d[\"age\"]**2, 11))\n",
    "    .assign(watch_time = lambda d: ((d[\"watch_time\"]- d[\"watch_time\"].min())/(d[\"watch_time\"].max()+20 - d[\"watch_time\"].min())*6).round(2))\n",
    "    .drop(columns=[\"asian\", \"gender\", \"black\", \"hawaiian\", \"hispanic\", \"unknown\", \"white\", \"format_ol\", \"format_blended\", \"falsexam\"])\n",
    ")\n",
    "\n",
    "data.to_csv(\"./causal-inference-in-python/data/rec_ab_test.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-24T19:58:51.452368Z",
     "start_time": "2023-01-24T19:58:50.940799Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "risk_data = (\n",
    "    pd.read_csv(\"./causal-inference-in-python/data/wage.csv\")\n",
    "    .sample(50000, replace=True)\n",
    "    .assign(wage = lambda d: np.random.normal(100 + d[\"wage\"], 50).round(-1))\n",
    "    .assign(credit_score1 = lambda d: np.random.normal(100 + 0.1*d[\"wage\"] + d[\"urban\"] + d[\"educ\"] + 2*d[\"exper\"], 50).round(-1))\n",
    "    .assign(credit_score2 = lambda d: np.random.normal(100 + 0.05*d[\"wage\"] + d[\"hours\"] + d[\"age\"] + 2*d[\"IQ\"], 50).round(-1))\n",
    "    \n",
    "    .assign(credit_score1 = lambda d: np.round(1000*(d[\"credit_score1\"] - d[\"credit_score1\"].min() + 20)/(d[\"credit_score1\"].max() - d[\"credit_score1\"].min()), 0))\n",
    "    .assign(credit_score2 = lambda d: np.round(1000*(d[\"credit_score2\"] - d[\"credit_score2\"].min() + 20)/(d[\"credit_score2\"].max() - d[\"credit_score2\"].min()), 0))\n",
    "\n",
    "    \n",
    "    .assign(credit_limit = lambda d: np.random.normal(500 \n",
    "                                                      + 1.5*pd.IntervalIndex(pd.qcut(d[\"credit_score1\"], 10)).mid\n",
    "                                                      + 1.5*pd.IntervalIndex(pd.qcut(d[\"wage\"], 10)).mid,\n",
    "                                                      1000).round(-2).clip(200, np.inf))\n",
    "    .assign(default = lambda d: (np.random.normal(+0\n",
    "                                                  - 0.1*d[\"credit_score1\"] \n",
    "                                                  - 0.5*d[\"credit_score2\"] \n",
    "                                                  + 0.05*d[\"IQ\"] \n",
    "                                                  + 0.005*d[\"credit_limit\"] \n",
    "                                                  -0.17*d[\"wage\"] \n",
    "                                                  \n",
    "                                                  , 375) > -50\n",
    "                                ).astype(int))\n",
    "    .drop(columns=[\"IQ\", \"lhwage\", \"tenure\", \"black\", \"hours\", \"sibs\", \"south\", \"urban\", \"brthord\", \"meduc\", \"feduc\", \"age\"])\n",
    "    .reset_index(drop=True)\n",
    ")\n",
    "\n",
    "spend_data = (risk_data\n",
    "              .assign(spend = lambda d: (np.random.normal(50 \n",
    "                                                          + 120*np.power(d[\"credit_limit\"], 1/2.5)\n",
    "                                                          - d[\"credit_score1\"] \n",
    "                                                          + 1.2 * d[\"wage\"]\n",
    "                                                          - 10*d[\"exper\"]\n",
    "                                                          - 10*d[\"educ\"]\n",
    "                                                          + 400*d[\"married\"]\n",
    "                                                         ).astype(int)))\n",
    "              .drop(columns=[\"default\"])\n",
    "             )\n",
    "\n",
    "risk_data.to_csv(\"./causal-inference-in-python/data/risk_data.csv\", index=False)\n",
    "spend_data.to_csv(\"./causal-inference-in-python/data/spend_data.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-24T19:58:51.509074Z",
     "start_time": "2023-01-24T19:58:51.453916Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "risk_data_rnd = (\n",
    "    pd.read_csv(\"./causal-inference-in-python/data/wage.csv\")\n",
    "    .sample(10000, replace=True)\n",
    "    .assign(wage = lambda d: np.random.normal(100 + d[\"wage\"], 50).round(-1))\n",
    "    .assign(credit_score1 = lambda d: np.random.normal(100 + 0.1*d[\"wage\"] + d[\"urban\"] + d[\"educ\"] + 2*d[\"exper\"], 50).round(-1))\n",
    "    .assign(credit_score2 = lambda d: np.random.normal(100 + 0.05*d[\"wage\"] + d[\"hours\"] + d[\"age\"] + 2*d[\"IQ\"], 50).round(-1))\n",
    "    \n",
    "    .assign(credit_score1 = lambda d: np.round(1000*(d[\"credit_score1\"] - d[\"credit_score1\"].min() + 20)/(d[\"credit_score1\"].max() - d[\"credit_score1\"].min()), 0))\n",
    "    .assign(credit_score2 = lambda d: np.round(1000*(d[\"credit_score2\"] - d[\"credit_score2\"].min() + 20)/(d[\"credit_score2\"].max() - d[\"credit_score2\"].min()), 0))\n",
    "    \n",
    "    .assign(credit_score1_buckets = lambda d: ((d[\"credit_score1\"] / 200).round(0) * 200).astype(int))\n",
    "    \n",
    "    .assign(credit_limit = lambda d: (np.random.beta(1 + d[\"credit_score1_buckets\"]/100, 6)*10000).round(-2))\n",
    "    .assign(default = lambda d: (np.random.normal(+0\n",
    "                                                  - 0.1*d[\"credit_score1\"] \n",
    "                                                  - 0.5*d[\"credit_score2\"] \n",
    "                                                  + 0.05*d[\"IQ\"] \n",
    "                                                  + 0.5*d[\"sibs\"]\n",
    "                                                  - 0.5*d[\"feduc\"]\n",
    "                                                  + 0.005*d[\"credit_limit\"] \n",
    "                                                  -0.17*d[\"wage\"] \n",
    "                                                  \n",
    "                                                  , 100) > -300\n",
    "                                ).astype(int))\n",
    "    .drop(columns=[\"IQ\", \"lhwage\", \"tenure\", \"black\", \"hours\", \"sibs\", \"south\", \"urban\", \"brthord\", \"meduc\", \"feduc\", \"age\"])\n",
    "    .reset_index(drop=True)\n",
    ")\n",
    "\n",
    "risk_data_rnd.to_csv(\"./causal-inference-in-python/data/risk_data_rnd.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-25T09:16:08.938041Z",
     "start_time": "2023-01-25T09:16:08.903066Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "spend_data_rnd = (\n",
    "    pd.read_csv(\"./causal-inference-in-python/data/wage.csv\")\n",
    "    .sample(500, replace=True)\n",
    "    .assign(wage = lambda d: np.random.normal(100 + d[\"wage\"], 50).round(-1))\n",
    "    .assign(credit_score1 = lambda d: np.random.normal(100, 50, len(d)).round(-1))\n",
    "    \n",
    "    .assign(credit_score1 = lambda d: np.round(1000*(d[\"credit_score1\"] - d[\"credit_score1\"].min() + 20)/(d[\"credit_score1\"].max() - d[\"credit_score1\"].min()), 0))\n",
    "    .assign(credit_score1_buckets = lambda d: ((d[\"credit_score1\"] / 200).round(0) * 200).astype(int))\n",
    "    \n",
    "    .assign(credit_limit = lambda d: (np.random.beta(1 + d[\"credit_score1_buckets\"]/100, 6)*10000).round(-2))\n",
    "    .assign(spend = lambda d: (np.random.normal(500 \n",
    "                                                + 20*np.power(d[\"credit_limit\"], 1/2)\n",
    "                                                + 1.2 * d[\"wage\"]\n",
    "                                                - 10*d[\"exper\"]\n",
    "                                                - 10*d[\"educ\"]\n",
    "                                                + 400*d[\"married\"], 600).astype(int)))\n",
    "    .drop(columns=[\"IQ\", \"lhwage\", \"tenure\", \"black\", \"hours\", \"sibs\", \"south\", \"urban\", \"brthord\", \"meduc\", \"feduc\", \"age\"])\n",
    "    .reset_index(drop=True)\n",
    ")\n",
    "\n",
    "spend_data_rnd.to_csv(\"./causal-inference-in-python/data/spend_data_rnd.csv\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-12T18:04:52.239145Z",
     "start_time": "2023-03-12T18:04:49.773194Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "np.random.seed(123)\n",
    "\n",
    "df = (pd.read_csv(\"https://raw.githubusercontent.com/matheusfacure/python-causality-handbook/master/causal-inference-for-the-brave-and-true/data/learning_mindset.csv\")\n",
    "      .rename(columns={\"schoolid\": \"departament_id\", \"achievement_score\": \"engagement_score\", \"success_expect\": \"tenure\",\n",
    "                       \"school_achievement\": \"last_engagement_score\", \"school_urbanicity\": \"role\", \"school_poverty\": \"department_score\",\n",
    "                       \"school_size\": \"department_size\", \"ethnicity\": \"n_of_reports\"})\n",
    "      # reduce overlapp for better examples\n",
    "      .assign(intervention = lambda d: (d[\"intervention\"].astype(bool) | (np.random.normal(d[\"last_engagement_score\"]) > 1.65)).astype(int))\n",
    "      .assign(intervention = lambda d: (d[\"intervention\"].astype(bool) | (np.random.normal(d[\"tenure\"], 2) > 7)).astype(int))\n",
    "      .assign(n_of_reports = lambda d: d[\"n_of_reports\"].clip(0, 8))\n",
    "      .assign(department_size = lambda d: d.groupby(\"departament_id\")[\"n_of_reports\"].transform(sum))\n",
    "      .assign(last_engagement_score = lambda d: np.random.normal(d[\"last_engagement_score\"]))\n",
    "      .drop(columns=[\"frst_in_family\", \"school_mindset\", \"school_ethnic_minority\"]))\n",
    "\n",
    "\n",
    "df.to_csv(\"./causal-inference-in-python/data/management_training.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-12T18:04:52.296332Z",
     "start_time": "2023-03-12T18:04:52.248495Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "\n",
    "n = 10000\n",
    "\n",
    "x1 = np.random.uniform(-1, 1, n)\n",
    "x2 = np.random.uniform(-1, 1, n)\n",
    "t = np.random.normal(7.5 - 1*(x1+x2), 2).clip(1, 14).round(1)\n",
    "y = np.random.normal((23 - 0.8*t - 8*(x1 + x2))).clip(1, np.inf).round(0)\n",
    "\n",
    "df_cont_t = pd.DataFrame(dict(ml_1=x1, ml_2=x2, interest=t, duration=y))\n",
    "\n",
    "df_cont_t.to_csv(\"./causal-inference-in-python/data/interest_rate.csv\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-26T13:24:24.810974Z",
     "start_time": "2023-06-26T13:24:24.700439Z"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from pandas.tseries import holiday\n",
    "\n",
    "np.random.seed(123)\n",
    "\n",
    "\n",
    "month_fe = {1:  23,\n",
    "            2:  11,\n",
    "            3:  10,\n",
    "            4:  9,\n",
    "            5:  4,\n",
    "            6:  2,\n",
    "            7:  10,\n",
    "            8:  5,\n",
    "            9:  10,\n",
    "            10: 20,\n",
    "            11: 25,\n",
    "            12: 30}\n",
    "\n",
    "week_fe = {0:  2,\n",
    "           1:  2,\n",
    "           2:  3,\n",
    "           3:  7,\n",
    "           4:  15,\n",
    "           5:  12,\n",
    "           6:  5}\n",
    "\n",
    "store_fe = {0:  20,\n",
    "            1:  14,\n",
    "            2:  3,\n",
    "            3:  10,\n",
    "            4:  9,\n",
    "            5:  12,\n",
    "            6:  5}\n",
    "\n",
    "day = pd.Series(pd.date_range(\"2016-01-01\", \"2019-01-01\"))\n",
    "weekend = day.dt.weekday >= 5\n",
    "is_holiday = day.isin(holiday.USFederalHolidayCalendar().holidays(\"2016-01-01\", \"2019-01-01\"))\n",
    "is_dec = day.dt.month == 12\n",
    "is_nov = day.dt.month == 11\n",
    "\n",
    "time_data = pd.DataFrame(dict(\n",
    "    day=day,\n",
    "    month = day.dt.month,\n",
    "    weekday = day.dt.weekday,\n",
    "    weekend=weekend,\n",
    "    is_holiday = is_holiday,\n",
    "    is_dec = day.dt.month == 12,\n",
    "    is_nov = day.dt.month == 11,\n",
    ")).assign(\n",
    "    month_fe = lambda d: d[\"month\"].map(month_fe),\n",
    "    week_fe = lambda d: d[\"weekday\"].map(week_fe),\n",
    ")\n",
    "\n",
    "data=(pd.concat([time_data.assign(rest_id=i) for i in range(len(store_fe))])\n",
    "      .assign(store_fe = lambda d: d[\"rest_id\"].map(store_fe))\n",
    "      .assign(competitors_price = lambda d: (np.random.beta(5+d[\"is_holiday\"]+d[\"is_dec\"]*2+d[\"is_nov\"], 0.5*d[\"store_fe\"], len(d))*9 + 1).round(2))\n",
    "      .assign(sales0 = lambda d: np.random.poisson(d[\"month_fe\"] + d[\"store_fe\"] + d[\"week_fe\"]\n",
    "                                                    + 3*d[\"competitors_price\"] \n",
    "                                                    + 10*d[\"is_holiday\"] \n",
    "                                                    + 5*d[\"weekend\"]))\n",
    "      .assign(effect = lambda d: np.random.poisson(150+d[\"month_fe\"] - d[\"store_fe\"] + d[\"week_fe\"]\n",
    "                                                    - 40*np.sqrt(d[\"competitors_price\"])\n",
    "                                                    + 5*d[\"is_holiday\"] \n",
    "                                                    - 3*d[\"weekend\"]))\n",
    "      .assign(discounts = lambda d: ((np.random.beta(1, 3, size=len(d)) * 50)/5).astype(int)*5)\n",
    "      .assign(sales = lambda d: d[\"sales0\"] + 0.5*d[\"effect\"]*d[\"discounts\"])\n",
    "      [[\"rest_id\", \"day\", \"month\", \"weekday\", \"weekend\", \"is_holiday\", \"is_dec\", \"is_nov\", \"competitors_price\",\n",
    "        \"discounts\", \"sales\"]])\n",
    "\n",
    "\n",
    "data.to_csv(\"./causal-inference-in-python/data/daily_restaurant_sales.csv\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-26T13:24:28.018989Z",
     "start_time": "2023-06-26T13:24:27.911853Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(1)\n",
    "\n",
    "categs = {\"vehicle\": 0.1,\n",
    "          \"food\": 1,\n",
    "          \"beverage\": 1,\n",
    "          \"art\": 1,\n",
    "          \"baby\": 1,\n",
    "          \"personal_care\": 1,\n",
    "          \"toys\": 1,\n",
    "          \"clothing\": 2,\n",
    "          \"decor\": 1,\n",
    "          \"cell_phones\": 3,\n",
    "          \"construction\": 1,\n",
    "          \"home_appliances\": 1,\n",
    "          \"electronics\": 2,\n",
    "          \"sports\": 1,\n",
    "          \"tools\": 1,\n",
    "          \"games\": 2,\n",
    "          \"industry\": 1,\n",
    "          \"pc\": 2,\n",
    "          \"jewel\": 1,\n",
    "          \"books\": 1,\n",
    "          \"music_books_movies\": 1,\n",
    "          \"health\": 2,\n",
    "         }\n",
    "\n",
    "\n",
    "n=10000\n",
    "\n",
    "age = (18+np.random.beta(2, 7, n)*90).round(0)\n",
    "tenure = np.random.exponential(0.5, n).round(0)\n",
    "ammount_spent = (np.random.exponential(1, n)*100 + np.random.binomial(1, 0.01, n)*np.random.uniform(500, 50000, n)).round(2)\n",
    "\n",
    "categ_purchage = {cat: np.random.poisson(l, n) for cat, l in categs.items()}\n",
    "\n",
    "X = pd.DataFrame(categ_purchage).assign(\n",
    "    age=age,\n",
    "    tenure=tenure,\n",
    "    ammount_spent=ammount_spent,\n",
    ")\n",
    "\n",
    "mkt_email_rnd = np.random.binomial(1, 0.5, n) \n",
    "\n",
    "coefs = np.concatenate([\n",
    "    np.random.uniform(-1, 1, len(categs)),\n",
    "    np.array([1, 20, 0.4])\n",
    "])\n",
    "y0 = np.random.exponential(X.values.dot(coefs)) \n",
    "\n",
    "cate_coef = np.concatenate([\n",
    "    np.random.uniform(0, 100, len(categs)),\n",
    "    np.array([-1, -20, 0])\n",
    "])\n",
    "\n",
    "tau = np.random.normal(X.values.dot(cate_coef))\n",
    "\n",
    "data_rnd = X.assign(\n",
    "    mkt_email = mkt_email_rnd,\n",
    "    next_mnth_pv = y0 + tau*mkt_email_rnd\n",
    ")[[\"mkt_email\", \"next_mnth_pv\", \"age\", \"tenure\", \"ammount_spent\"] + list(categs.keys())].round(2)\n",
    "\n",
    "\n",
    "data_rnd.to_csv(\"./causal-inference-in-python/data/email_rnd_data.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-26T13:24:30.494095Z",
     "start_time": "2023-06-26T13:24:28.482269Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(2)\n",
    "\n",
    "n=300000\n",
    "\n",
    "age = (18+np.random.beta(2, 7, n)*90).round(0)\n",
    "tenure = np.random.exponential(0.5, n).round(0)\n",
    "ammount_spent = (np.random.exponential(1, n)*100 + np.random.binomial(1, 0.01, n)*np.random.uniform(500, 50000, n)).round(2)\n",
    "\n",
    "categ_purchage = {cat: np.random.poisson(l, n) for cat, l in categs.items()}\n",
    "\n",
    "X = pd.DataFrame(categ_purchage).assign(\n",
    "    age=age,\n",
    "    tenure=tenure,\n",
    "    ammount_spent=ammount_spent,\n",
    ")\n",
    "\n",
    "coefs_ps = np.concatenate([\n",
    "    np.random.uniform(-1, 1, len(categs)),\n",
    "    np.array([-1, 40, 0.4])\n",
    "])\n",
    "\n",
    "\n",
    "mkt_email_biased = (np.random.normal(X.values.dot(coefs_ps), 2000) > 0).astype(int)\n",
    "\n",
    "y0 = np.random.exponential(X.values.dot(coefs)) \n",
    "tau = np.random.normal(X.values.dot(cate_coef))\n",
    "\n",
    "data_biased = X.assign(\n",
    "    mkt_email = mkt_email_biased,\n",
    "    next_mnth_pv = y0 + tau*mkt_email_biased\n",
    ")[[\"mkt_email\", \"next_mnth_pv\", \"age\", \"tenure\", \"ammount_spent\"] + list(categs.keys())].round(2)\n",
    "\n",
    "data_biased.to_csv(\"./causal-inference-in-python/data/email_obs_data.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-26T13:24:33.155603Z",
     "start_time": "2023-06-26T13:24:31.504219Z"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "from pandas.tseries import holiday\n",
    "\n",
    "np.random.seed(123)\n",
    "\n",
    "\n",
    "month_fe = {1:  23,\n",
    "            2:  11,\n",
    "            3:  10,\n",
    "            4:  9,\n",
    "            5:  4,\n",
    "            6:  2,\n",
    "            7:  10,\n",
    "            8:  5,\n",
    "            9:  10,\n",
    "            10: 20,\n",
    "            11: 25,\n",
    "            12: 30}\n",
    "\n",
    "week_fe = {0:  2,\n",
    "           1:  2,\n",
    "           2:  3,\n",
    "           3:  7,\n",
    "           4:  15,\n",
    "           5:  12,\n",
    "           6:  5}\n",
    "\n",
    "store_fe = {0:  20,\n",
    "            1:  14,\n",
    "            2:  3,\n",
    "            3:  10,\n",
    "            4:  9,\n",
    "            5:  12,\n",
    "            6:  5}\n",
    "\n",
    "day = pd.Series(pd.date_range(\"2016-01-01\", \"2019-01-01\"))\n",
    "weekend = day.dt.weekday >= 5\n",
    "is_holiday = day.isin(holiday.USFederalHolidayCalendar().holidays(\"2016-01-01\", \"2019-01-01\"))\n",
    "is_dec = day.dt.month == 12\n",
    "is_nov = day.dt.month == 11\n",
    "\n",
    "time_data = pd.DataFrame(dict(\n",
    "    day=day,\n",
    "    month = day.dt.month,\n",
    "    weekday = day.dt.weekday,\n",
    "    weekend=weekend,\n",
    "    is_holiday = is_holiday,\n",
    "    is_dec = day.dt.month == 12,\n",
    "    is_nov = day.dt.month == 11,\n",
    ")).assign(\n",
    "    month_fe = lambda d: d[\"month\"].map(month_fe),\n",
    "    week_fe = lambda d: d[\"weekday\"].map(week_fe),\n",
    ")\n",
    "\n",
    "data_cont=(pd.concat([time_data.assign(rest_id=i) for i in range(len(store_fe))])\n",
    "      .assign(store_fe = lambda d: d[\"rest_id\"].map(store_fe))\n",
    "      .assign(competitors_price = lambda d: (np.random.beta(5+d[\"is_holiday\"]+d[\"is_dec\"]*2+d[\"is_nov\"], 0.5*d[\"store_fe\"], len(d))*9 + 1).round(2))\n",
    "      .assign(sales0 = lambda d: np.random.poisson(d[\"month_fe\"] + d[\"store_fe\"] + d[\"week_fe\"]\n",
    "                                                    + 3*d[\"competitors_price\"] \n",
    "                                                    + 10*d[\"is_holiday\"] \n",
    "                                                    + 5*d[\"weekend\"]))\n",
    "      .assign(effect = lambda d: np.random.poisson(150+d[\"month_fe\"] - d[\"store_fe\"] + d[\"week_fe\"]\n",
    "                                                    - 40*np.sqrt(d[\"competitors_price\"])\n",
    "                                                    + 5*d[\"is_holiday\"] \n",
    "                                                    - 3*d[\"weekend\"]))\n",
    "      .assign(discounts = lambda d: ((np.random.beta(1, 3, size=len(d)) * 50)/5).astype(int)*5)\n",
    "      .assign(sales = lambda d: d[\"sales0\"] + 0.5*d[\"effect\"]*d[\"discounts\"])\n",
    "      [[\"rest_id\", \"day\", \"month\", \"weekday\", \"weekend\", \"is_holiday\", \"is_dec\", \"is_nov\", \"competitors_price\",\n",
    "        \"discounts\", \"sales\"]])\n",
    "\n",
    "\n",
    "data_cont.to_csv(\"./causal-inference-in-python/data/discount_data.csv\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-26T13:24:36.557346Z",
     "start_time": "2023-06-26T13:24:36.438111Z"
    }
   },
   "outputs": [],
   "source": [
    "date = pd.date_range(\"2021-05-01\", \"2021-07-31\", freq=\"D\")\n",
    "cohorts = pd.to_datetime([\"2021-05-15\", \"2021-06-04\", \"2021-06-20\"]).date\n",
    "poss_regions = [\"S\", \"N\", \"W\", \"E\"]\n",
    "\n",
    "reg_ps = dict(zip(poss_regions,    [.3, .6, .7, .8]))\n",
    "reg_fe = dict(zip(poss_regions,    [20,  16,  8,  2]))\n",
    "reg_trend = dict(zip(poss_regions, [0,  0.2,  .4,  .6]))\n",
    "\n",
    "units = np.array(range(1, 200+1))\n",
    "\n",
    "np.random.seed(123)\n",
    "\n",
    "unit_reg = np.random.choice(poss_regions, len(units))\n",
    "exp_trend = np.random.exponential(0.01, len(units))\n",
    "treated_unit = np.random.binomial(1, np.vectorize(reg_ps.__getitem__)(unit_reg))\n",
    "\n",
    "# staggered addopton dgp\n",
    "df = pd.DataFrame(dict(\n",
    "    date = np.tile(date.date, len(units)),\n",
    "    city = np.repeat(units, len(date)),\n",
    "    region = np.repeat(unit_reg, len(date)),\n",
    "    treated_unit = np.repeat(treated_unit, len(date)),\n",
    "    cohort = np.repeat(np.random.choice(cohorts, len(units)), len(date)),\n",
    "    eff_heter = np.repeat(np.random.exponential(1, size=len(units)), len(date)),\n",
    "    \n",
    "    unit_fe = np.repeat(np.random.normal(0, 2, size=len(units)), len(date)),\n",
    "    time_fe = np.tile(np.random.normal(size=len(date)), len(units)),\n",
    "    week_day = np.tile(date.weekday, len(units)),\n",
    "    w_seas = np.tile(abs(5-date.weekday) % 7, len(units)),\n",
    ")).assign(\n",
    "    reg_fe = lambda d: d[\"region\"].map(reg_fe), \n",
    "    reg_trend = lambda d: d[\"region\"].map(reg_trend), \n",
    "    reg_ps = lambda d: d[\"region\"].map(reg_ps), \n",
    "    trend = lambda d: (d[\"date\"] - d[\"date\"].min()).dt.days,\n",
    "    day = lambda d: (d[\"date\"] - d[\"date\"].min()).dt.days,\n",
    "    cohort = lambda d: np.where(d[\"treated_unit\"] == 1, d[\"cohort\"], pd.to_datetime(\"2100-01-01\")),\n",
    ").assign(\n",
    "    treated = lambda d: ((d[\"date\"] >= d[\"cohort\"]) & d[\"treated_unit\"] == 1).astype(int),\n",
    ").assign(\n",
    "    y0 = lambda d: np.round(10 \n",
    "                            + d[\"treated_unit\"]\n",
    "                            + d[\"reg_trend\"]*d[\"trend\"]/2\n",
    "                            + d[\"unit_fe\"] \n",
    "                            + 0.4*d[\"time_fe\"] \n",
    "                            + 2*d[\"reg_fe\"]\n",
    "                            + d[\"w_seas\"]/5, 0),\n",
    ").assign(\n",
    "#     y0 = lambda d: np.round(d[\"y0\"] + d.groupby(\"city\")[\"y0\"].shift(1).fillna(0)*0.2, 0)\n",
    ").assign(\n",
    "    y1 = lambda d: d[\"y0\"] + np.minimum(0.2*(np.maximum(0, (pd.to_datetime(d[\"date\"]) - d[\"cohort\"]).dt.days)), 1)*d[\"eff_heter\"]*2\n",
    ").assign(\n",
    "    tau = lambda d: d[\"y1\"] - d[\"y0\"],\n",
    "    downloads = lambda d: np.where(d[\"treated\"] == 1, d[\"y1\"], d[\"y0\"]) + np.random.normal(0,.7,len(d)),\n",
    "#     date = lambda d: pd.to_datetime(d[\"date\"]),\n",
    ").round({\"downloads\": 0})\n",
    "\n",
    "# df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-26T13:24:38.347953Z",
     "start_time": "2023-06-26T13:24:38.141867Z"
    }
   },
   "outputs": [],
   "source": [
    "# filter only 1st block\n",
    "reg_filter = [\"N\", \"S\", \"E\", \"W\"]\n",
    "\n",
    "mkt_data_all = (df\n",
    "                .query(\"region.isin(@reg_filter)\")\n",
    "                .query(\"date.astype('string') <= '2021-06-01'\")\n",
    "#                 .assign(cohort = lambda d: np.where(d[\"cohort\"].astype(str) <= \"2021-06-01\", d[\"cohort\"], \"2050-01-01\"))\n",
    "                .drop(columns=[\"reg_fe\", \"time_fe\", \"cohort\", \"w_seas\", \"week_day\", \"reg_trend\", \"trend\", \"day\", \"unit_fe\",\n",
    "                               \"y0\", \"y1\", \"eff_heter\", \"reg_ps\", \"treated_unit\",\n",
    "                              ])\n",
    "                .assign(post=lambda d: (d[\"date\"].astype(str) >= \"2021-05-15\").astype(int))\n",
    "                .assign(treated=lambda d: d.groupby(\"city\")[\"treated\"].transform(max)))\n",
    "\n",
    "\n",
    "mkt_data = (mkt_data_all.query(\"region=='S'\"))\n",
    "\n",
    "mkt_data_cohorts = (df\n",
    "                    .assign(post=lambda d: (d[\"date\"] >= d[\"cohort\"]).astype(int))\n",
    "                    .assign(treated=lambda d: d.groupby(\"city\")[\"treated\"].transform(max))\n",
    "                    .drop(columns=[\"reg_fe\", \"time_fe\", \"w_seas\", \"week_day\", \"reg_trend\", \"trend\", \"day\", \"unit_fe\",\n",
    "                               \"y0\", \"y1\", \"eff_heter\"\n",
    "                                   , \"reg_ps\", \"treated_unit\"]))\n",
    "\n",
    "mkt_data_cohorts.to_csv(\"./causal-inference-in-python/data/offline_mkt_staggered.csv\", index=False)\n",
    "\n",
    "mkt_data_all.to_csv(\"./causal-inference-in-python/data/short_offline_mkt_all_regions.csv\", index=False)\n",
    "\n",
    "mkt_data.to_csv(\"./causal-inference-in-python/data/short_offline_mkt_south.csv\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-24T19:59:00.262446Z",
     "start_time": "2023-01-24T19:59:00.189182Z"
    }
   },
   "outputs": [],
   "source": [
    "from statsmodels.tsa.arima_process import ArmaProcess\n",
    "\n",
    "def simulate_series(dates,\n",
    "                    name,\n",
    "                    pop,\n",
    "                    pop_pct,\n",
    "                    ws_w,\n",
    "                    ms_w,\n",
    "                    ys_w,\n",
    "                    ar_w,\n",
    "                    trend_w,\n",
    "                    noise_w,\n",
    "                    ar=1,\n",
    "                    ma=7,\n",
    "                   ):\n",
    "    \n",
    "    ws = abs(5-dates.weekday) % 7\n",
    "    ms = abs(dates.day-10)\n",
    "    ys = abs(6-dates.month) % 12\n",
    "    \n",
    "    arma = ArmaProcess(ar, ma).generate_sample(len(dates))\n",
    "    trend = np.arange(0, len(dates))\n",
    "    noise = np.random.normal(0, 1, len(dates))\n",
    "    \n",
    "    \n",
    "    comps =  (ws, ms, ys, arma, trend, noise)\n",
    "    comp_noise = np.random.uniform(0.95, 1.05, len(comps))\n",
    "    \n",
    "    coefs = (ws_w, ms_w, ys_w, ar_w, trend_w, noise_w)\n",
    "    \n",
    "    result = sum(c*w*noise\n",
    "                 for c, w, noise\n",
    "                 in zip(comps, coefs, comp_noise))*pop*pop_pct*0.005\n",
    "    \n",
    "    return pd.Series(result, name=name).clip(0, np.inf).round(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-24T19:59:00.271093Z",
     "start_time": "2023-01-24T19:59:00.263454Z"
    }
   },
   "outputs": [],
   "source": [
    "br_cities = pd.read_csv(\"./causal-inference-in-python/data/br_cities.csv\")\n",
    "\n",
    "np.random.seed(3)\n",
    "\n",
    "states = br_cities[\"state\"].unique()\n",
    "pop_pct = np.random.beta(5, 10, len(states))*0.1\n",
    "trends = np.random.beta(9, 5, len(states))*0.5-0.25\n",
    "week_sas = np.random.uniform(2, 4, len(states))\n",
    "m_sas = np.random.uniform(0, 1, len(states))\n",
    "y_sas = np.random.uniform(0, 2, len(states))\n",
    "\n",
    "states_params = pd.DataFrame(dict(\n",
    "    pop_pct=pop_pct,\n",
    "    state=states,\n",
    "    trends=trends,\n",
    "    week_sas=week_sas,\n",
    "    m_sas=m_sas,\n",
    "    y_sas=y_sas,\n",
    "))\n",
    "\n",
    "# states_params.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-24T19:59:00.279476Z",
     "start_time": "2023-01-24T19:59:00.272275Z"
    }
   },
   "outputs": [],
   "source": [
    "br_cities = pd.read_csv(\"./causal-inference-in-python/data/br_cities.csv\").merge(states_params, on=\"state\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-24T19:59:00.456232Z",
     "start_time": "2023-01-24T19:59:00.280636Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(12345)\n",
    "\n",
    "date = pd.date_range(\"2022-03-01\", \"2022-06-30\", freq=\"D\")\n",
    "\n",
    "series = [\n",
    "    simulate_series(date,\n",
    "                    name=city[\"city\"],\n",
    "                    pop=city[\"pop\"]*city[\"pop_pct\"],\n",
    "                    pop_pct=city[\"pop_pct\"],\n",
    "                    ws_w=city[\"week_sas\"],\n",
    "                    ms_w=city[\"m_sas\"],\n",
    "                    ys_w=city[\"y_sas\"],\n",
    "                    ar_w=0.1,\n",
    "                    trend_w=city[\"trends\"],\n",
    "                    noise_w=1000/np.sqrt(city[\"pop\"]*city[\"pop_pct\"])\n",
    "                   ).round(0).to_frame().assign(\n",
    "        population=city[\"pop\"], city=city[\"city\"], state=city[\"state\"], date=date,\n",
    "    ).rename(columns={city[\"city\"]:\"app_download\"})\n",
    "    for _, city in br_cities.sort_values(\"pop\", ascending=False).head(50).iterrows()\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-24T19:59:00.481510Z",
     "start_time": "2023-01-24T19:59:00.457362Z"
    }
   },
   "outputs": [],
   "source": [
    "features = pd.concat([pd.DataFrame(dict(\n",
    "    activity=(ArmaProcess([1, 0.9, 0.8, 0.7], 7).generate_sample(len(date))+row[\"hdi\"]).round(0)*3,\n",
    "    date=date,\n",
    "    state=row[\"state\"],\n",
    "    city=row[\"city\"])) \n",
    "                      for _, row in br_cities.sort_values(\"pop\", ascending=False).head(50).iterrows()])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-24T19:59:00.486234Z",
     "start_time": "2023-01-24T19:59:00.482811Z"
    }
   },
   "outputs": [],
   "source": [
    "def simulate_effect(df, city_col, date_col, y_col, city, start_at, window, effect_pct):\n",
    "    \n",
    "    window_mask = (df[date_col] >= start_at) & (df[city_col] == city)\n",
    "    \n",
    "    def rbf(x, center, sigma):\n",
    "        return np.exp(-(x-center)**2 / sigma**2)\n",
    "\n",
    "    kernel = rbf((df[date_col] - pd.to_datetime(start_at)).dt.days - window//2 - 1, 0, window/2)\n",
    "    \n",
    "    \n",
    "    y = np.where(window_mask, df[y_col] + df[y_col]*effect_pct*kernel, df[y_col])\n",
    "    \n",
    "    return df.assign(**{y_col: y})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-24T19:59:00.542421Z",
     "start_time": "2023-01-24T19:59:00.487636Z"
    }
   },
   "outputs": [],
   "source": [
    "df = pd.concat(series)\n",
    "\n",
    "df = simulate_effect(df, \"city\", \"date\", \"app_download\", \"sao_paulo\", \"2022-05-01\", 40, 0.4)\n",
    "df = simulate_effect(df, \"city\", \"date\", \"app_download\", \"joao_pessoa\", \"2022-05-01\", 40, 0.5)\n",
    "df = simulate_effect(df, \"city\", \"date\", \"app_download\", \"porto_alegre\", \"2022-05-01\", 40, 0.6)\n",
    "\n",
    "df = df.assign(\n",
    "    post = (df[\"date\"] >= \"2022-05-01\")*1,\n",
    "    treated = (df[\"city\"].isin([\"sao_paulo\", \"joao_pessoa\", \"porto_alegre\"]))*1,\n",
    ")\n",
    "\n",
    "df.to_csv(\"./causal-inference-in-python/data/online_mkt.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-24T19:59:00.602696Z",
     "start_time": "2023-01-24T19:59:00.543585Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "\n",
    "df_norm = df.assign(app_download_pct = 100*df[\"app_download\"]/df[\"population\"])\n",
    "\n",
    "\n",
    "df_norm_cov = (df_norm\n",
    "               .merge(br_cities[[\"city\", \"state\", \"hdi\"]])\n",
    "               .assign(comp_download_pct = lambda d: 100*(0.8*d[\"app_download\"]+d[\"hdi\"])/d[\"population\"])\n",
    "               .drop(columns=[\"hdi\"]))\n",
    "\n",
    "df_norm_cov.to_csv(\"./causal-inference-in-python/data/online_mkt_cov.csv\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-06T12:48:17.398154Z",
     "start_time": "2023-02-06T12:48:17.373468Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from matplotlib import pyplot as plt\n",
    "from statsmodels.tsa.arima_process import ArmaProcess\n",
    "\n",
    "T = 120\n",
    "m = 2\n",
    "p = 0.5\n",
    "\n",
    "\n",
    "def y_given_d(d, effect_params=[3,2,1], T=T, seed=None):\n",
    "    np.random.seed(seed)\n",
    "    x = np.arange(1, T+1)\n",
    "    return (np.log(x+1)\n",
    "            + 2*np.sin(x*2*np.pi/24)\n",
    "            + np.convolve(~d.astype(bool), effect_params)[:-(len(effect_params)-1)]\n",
    "            + np.random.normal(0, 1, T)\n",
    "#             + ArmaProcess([3,2,], 3).generate_sample(T)\n",
    "           ).round(2)\n",
    "\n",
    "\n",
    "# correct tau = 3+2+1=6\n",
    "effect_params = [3,2,1]\n",
    "\n",
    "\n",
    "np.random.seed(123)\n",
    "\n",
    "df_sb_every = pd.DataFrame(\n",
    "    dict(\n",
    "        d=np.random.binomial(1, 0.5, T),\n",
    "    )\n",
    ").assign(\n",
    "    delivery_time = lambda d: y_given_d(d[\"d\"], seed=1),\n",
    "    delivery_time_1 = lambda d: y_given_d(np.ones(T), seed=1),\n",
    "    delivery_time_0 = lambda d: y_given_d(np.zeros(T), seed=1),\n",
    ").assign(\n",
    "    tau = lambda d: d[\"delivery_time_1\"] - d[\"delivery_time_0\"]\n",
    ")\n",
    "\n",
    "\n",
    "df_sb_every.to_csv(\"./causal-inference-in-python/data/sb_exp_every.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-06T12:37:42.355708Z",
     "start_time": "2023-02-06T12:37:42.333596Z"
    }
   },
   "outputs": [],
   "source": [
    "def gen_d(rand_points, p):\n",
    "    result = [np.random.binomial(1, p)]\n",
    "    \n",
    "    for t in rand_points[1:]:\n",
    "        result.append(np.random.binomial(1, p)*t + (1-t)*result[-1])\n",
    "    \n",
    "    return np.array(result)\n",
    "\n",
    "\n",
    "m = 2\n",
    "T = 120\n",
    "p = 0.5\n",
    "\n",
    "n = T//m \n",
    "\n",
    "rand_points_opt = np.isin(\n",
    "    np.arange(1, T+1),\n",
    "    [1] + [i*m+1 for i in range(2, int(n)-1)]\n",
    ")\n",
    "\n",
    "np.random.seed(1)\n",
    "\n",
    "d_opt = gen_d(rand_points_opt, p)\n",
    "y_opt = y_given_d(d_opt, T=T, seed=1)\n",
    "\n",
    "pd.DataFrame(dict(rand_points=rand_points_opt,\n",
    "                  d=d_opt,\n",
    "                  delivery_time=y_opt)).to_csv(\"./causal-inference-in-python/data/sb_exp_opt.csv\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 11"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-24T19:59:00.655823Z",
     "start_time": "2023-01-24T19:59:00.603890Z"
    }
   },
   "outputs": [],
   "source": [
    "n=10000\n",
    "np.random.seed(123)\n",
    "\n",
    "# confounders\n",
    "age = 18 + np.random.normal(24, 4, n).round(1)\n",
    "income = 500 + np.random.gamma(1, age * 100, n).round(0)\n",
    "credit_score = (np.random.beta(1,3, n)*1000).round(0)\n",
    "\n",
    "\n",
    "u = np.random.uniform(-1, 1, n)\n",
    "\n",
    "categs = [\"never-taker\", \"complier\"]\n",
    "\n",
    "cutoff=0.6\n",
    "e = (1/(1+np.exp(-(u-0.05*age+0.01*credit_score))))\n",
    "cust_categ = np.select([e <= cutoff, e > cutoff], categs)\n",
    "\n",
    "# plt.hist(e)\n",
    "\n",
    "# Instrument\n",
    "prime_elegible = np.random.binomial(1, 0.5, n)\n",
    "choose_prime = (\n",
    "    ((cust_categ == \"complier\") & (prime_elegible == 1))\n",
    ").astype(int)\n",
    "\n",
    "# outcome\n",
    "group_effect = np.select([cust_categ == \"complier\", cust_categ == \"never-taker\"],\n",
    "                         [700, 200])\n",
    "\n",
    "pv_0 = np.random.normal(200 + 0.4*income + 10*age - 500*u, 200).round(2)\n",
    "pv_1 = pv_0 + group_effect\n",
    "\n",
    "tau = pv_1 - pv_0\n",
    "pv = pv_0 + group_effect*choose_prime\n",
    "\n",
    "\n",
    "\n",
    "df = pd.DataFrame(dict(categ=cust_categ,\n",
    "                       age=age,\n",
    "                       income=income,\n",
    "                       credit_score=credit_score,\n",
    "                       prime_elegible=prime_elegible,\n",
    "                       prime_card=choose_prime,\n",
    "                       pv=pv,\n",
    "                       tau=tau,))[\n",
    "    ['age', 'income', 'credit_score', 'prime_elegible',\n",
    "     'prime_card', 'pv', 'tau', 'categ']\n",
    "]\n",
    "\n",
    "df.to_csv(\"./causal-inference-in-python/data/prime_card.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:root] *",
   "language": "python",
   "name": "conda-root-py"
  },
  "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.12"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}