{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "94b13e2f",
   "metadata": {},
   "source": [
    "\n",
    "File: 01-stratified_sampling.py\n",
    "\n",
    "\n",
    "Michel Bierlaire\n",
    "\n",
    "Mon Aug 04 2025, 10:27:55\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a528dbe6",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from pandas import DataFrame, Series\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad4a163f",
   "metadata": {},
   "source": [
    "The objective of this laboratory is to experiment various sampling strategies.\n",
    "\n",
    "A synthetic population composed of 1 million individuals has been generated based on the Optima database,\n",
    "containing real observations, real individuals, with observed values of all the variables.\n",
    "Each synthetic individual was generated as follows:\n",
    "\n",
    "- One real individual was randomly selected.\n",
    "- The socio-economic characteristics of the synthetic individual are the same as those of the real individuals.\n",
    "- The attributes of each alternative are calculated from the real observed attributes, modified with a random perturbation.\n",
    "- The chosen alternative is generated by simulation using the \"true\" choice model specified above, with the \"true\" value of the parameters.\n",
    "\n",
    "The scripts that have been used are\n",
    "- `generate_population_variables.py` for the generation of the dependent variables,\n",
    "- `generate_population_choices.py` for the generation of the choices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9537bb6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "population_data_file = 'synthetic_population_with_choice.zip'\n",
    "population = pd.read_csv(population_data_file)\n",
    "population.describe()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cf61512",
   "metadata": {},
   "source": [
    "Calculate the population size, and set the sample size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d2190ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "population_size = population.shape[0]\n",
    "print(f'Population size: {population_size:_}')\n",
    "choice_variable = 'Choice'\n",
    "choice_set = np.sort(population[choice_variable].unique())\n",
    "print(f'{choice_set=}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cad78db",
   "metadata": {},
   "source": [
    "Sample size used for these experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88b97ee8",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_size = 1000\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d06955f4",
   "metadata": {},
   "source": [
    "We describe the procedure to select a sample from that population using a stratified sample, where the strata are\n",
    "defined based on the trip purpose, the language and the chosen alternative. Then, you will be asked to implement a\n",
    "similar procedure for other sampling strategies."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9b6da49",
   "metadata": {},
   "source": [
    "## Step 1: Stratification."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "550e0894",
   "metadata": {},
   "source": [
    "For each segment, we define a function that generates a Pandas mask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f422c96",
   "metadata": {},
   "source": [
    "In the population, there are two trip purposes: 1: 'work', 2: 'not_work'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "422bdf6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(population['TripPurpose'].unique())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61072773",
   "metadata": {},
   "source": [
    "There are also two linguistic regions: 1: 'french', 2: 'german'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13a8d060",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(population['LangCode'].unique())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83ef9ced",
   "metadata": {},
   "source": [
    "The strata are defined based on the trip purpose, the language and the chosen alternative.   As the choice is\n",
    "involved in the definition of strata, the stratified sampling is *endogenous*."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39ab3121",
   "metadata": {},
   "source": [
    "We define first the strata based on exogenous variables (trip purpose and language)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3305126",
   "metadata": {},
   "source": [
    "Trip purpose = work (1), language = French (1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3eb37b9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mask_work_french(data_frame: DataFrame) -> Series:\n",
    "    \"\"\"Generate a mask for trip purpose = work (1), language = French (1)\n",
    "\n",
    "    :param data_frame: data frame where the mask is applied\n",
    "    \"\"\"\n",
    "    return (data_frame['TripPurpose'] == 1) & (data_frame['LangCode'] == 1)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9338e275",
   "metadata": {},
   "source": [
    "Check the number and share of individuals in the population in this segment.\n",
    "Using the notation in the lecture, the share is $W(x) = N(x) / N$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfe2a248",
   "metadata": {},
   "outputs": [],
   "source": [
    "population_work_french = mask_work_french(population).sum()\n",
    "population_share_work_french = population_work_french / population_size\n",
    "print(f'Size segment work, French in the population:  {population_work_french}')\n",
    "print(\n",
    "    f'Share segment work, French in the population: W(x) = {100*population_share_work_french:.3g}%'\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "553e7f82",
   "metadata": {},
   "source": [
    "Trip purpose = work (1), language = German (2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ed05652",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mask_work_german(data_frame: DataFrame) -> Series:\n",
    "    \"\"\"Generate a mask for trip purpose = work (1), language = German (2)\n",
    "\n",
    "    :param data_frame: data frame where the mask is applied\n",
    "    \"\"\"\n",
    "    return (data_frame['TripPurpose'] == 1) & (data_frame['LangCode'] == 2)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59e73fc6",
   "metadata": {},
   "source": [
    "Check the number of individuals in the population in this segment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65d190c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "population_work_german = mask_work_german(population).sum()\n",
    "population_share_work_german = population_work_german / population_size\n",
    "print(f'Size segment work, German in the population:  {population_work_german}')\n",
    "print(\n",
    "    f'Share segment work, German in the population: W(x) = {100*population_share_work_german:.3g}%'\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "297ca5dc",
   "metadata": {},
   "source": [
    "Trip purpose = non work (2), language = French (1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd7599ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mask_non_work_french(data_frame: DataFrame) -> Series:\n",
    "    \"\"\"Generate a mask for trip purpose = non work (2), language = French (1)\n",
    "\n",
    "    :param data_frame: data frame where the mask is applied\n",
    "    \"\"\"\n",
    "    return (data_frame['TripPurpose'] == 2) & (data_frame['LangCode'] == 1)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52927d90",
   "metadata": {},
   "source": [
    "Check the number of individuals in the population in this segment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8884987e",
   "metadata": {},
   "outputs": [],
   "source": [
    "population_non_work_french = mask_non_work_french(population).sum()\n",
    "population_share_non_work_french = population_non_work_french / population_size\n",
    "print(f'Size segment non work, French in the population:  {population_non_work_french}')\n",
    "print(\n",
    "    f'Share segment non work, French in the population: W(x) = {100*population_share_non_work_french:.3g}%'\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a29cf837",
   "metadata": {},
   "source": [
    "Trip purpose = non work (2), language = German (2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b819299",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mask_non_work_german(data_frame: DataFrame) -> Series:\n",
    "    \"\"\"Generate a mask for trip purpose = non work (1), language = German (2)\n",
    "\n",
    "    :param data_frame: data frame where the mask is applied\n",
    "    \"\"\"\n",
    "    return (data_frame['TripPurpose'] == 2) & (data_frame['LangCode'] == 2)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c95d1669",
   "metadata": {},
   "source": [
    "Check the number of individuals in the population in this segment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e7d075b",
   "metadata": {},
   "outputs": [],
   "source": [
    "population_non_work_german = mask_non_work_german(population).sum()\n",
    "population_share_non_work_german = population_non_work_german / population_size\n",
    "print(f'Size segment non work, German in the population:  {population_non_work_german}')\n",
    "print(\n",
    "    f'Share segment non work, German in the population: W(x) = {100*population_share_non_work_german:.3g}%'\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77e0a3e1",
   "metadata": {},
   "source": [
    "We can verify that it is indeed covering the whole population"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a7c927c",
   "metadata": {},
   "outputs": [],
   "source": [
    "total = (\n",
    "    population_share_work_french\n",
    "    + population_share_work_german\n",
    "    + population_share_non_work_french\n",
    "    + population_share_non_work_german\n",
    ")\n",
    "print(f'Total share: {100*total:.3g}%')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fadb816f",
   "metadata": {},
   "source": [
    "## Step 2: Target shares."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95dadf77",
   "metadata": {},
   "source": [
    "We define the target shares of each segment in the sample."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "836debee",
   "metadata": {},
   "source": [
    "Using the notation from the lecture, those are $H(x)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e2f3fd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_work_french = 0.25\n",
    "target_work_german = 0.25\n",
    "target_non_work_french = 0.25\n",
    "target_non_work_german = 0.25\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c94dc422",
   "metadata": {},
   "source": [
    "## Step 3: Endogenous stratification."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e34014f1",
   "metadata": {},
   "source": [
    "Within each exogenous segment, we stratify according to the alternative"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "362e6af2",
   "metadata": {},
   "source": [
    "Work, French"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7fe84d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mask_work_french_choice(data_frame: DataFrame, alt_id: int) -> Series:\n",
    "    \"\"\"Generate a mask for trip purpose = work, language = French and choice = alt_id\n",
    "\n",
    "    :param data_frame: data frame where the mask is applied\n",
    "    \"\"\"\n",
    "    exogenous_mask = mask_work_french(data_frame)\n",
    "    choice_mask = data_frame['Choice'] == alt_id\n",
    "    return exogenous_mask & choice_mask\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ace19d4f",
   "metadata": {},
   "source": [
    "Check the number of individuals in the population in this segment.\n",
    "In the notation of the course, this corresponds to $W(i, x)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9f71d2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "population_work_french_choice = {}\n",
    "population_share_work_french_choice = {}\n",
    "for alt_id in choice_set:\n",
    "    population_work_french_choice[alt_id] = mask_work_french_choice(\n",
    "        data_frame=population, alt_id=alt_id\n",
    "    ).sum()\n",
    "    population_share_work_french_choice[alt_id] = (\n",
    "        population_work_french_choice[alt_id] / population_size\n",
    "    )\n",
    "    print(\n",
    "        f'Size segment work, French, choice={alt_id} in the population:  {population_work_french_choice[alt_id]}'\n",
    "    )\n",
    "    print(\n",
    "        f'Share segment work, French, choice={alt_id} in the population: W(i, x) = '\n",
    "        f'{100*population_share_work_french_choice[alt_id]:.3g}%'\n",
    "    )\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cd3ce81",
   "metadata": {},
   "source": [
    "Work, German"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b66c7155",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mask_work_german_choice(data_frame: DataFrame, alt_id: int) -> Series:\n",
    "    \"\"\"Generate a mask for trip purpose = work, language = German and choice = alt_id\n",
    "\n",
    "    :param data_frame: data frame where the mask is applied\n",
    "    \"\"\"\n",
    "    exogenous_mask = mask_work_german(data_frame)\n",
    "    choice_mask = data_frame['Choice'] == alt_id\n",
    "    return exogenous_mask & choice_mask\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cde1272",
   "metadata": {},
   "source": [
    "Check the number of individuals in the population in this segment.\n",
    "In the notation of the course, this corresponds to $W(i, x)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5df33ddc",
   "metadata": {},
   "outputs": [],
   "source": [
    "population_work_german_choice = {}\n",
    "population_share_work_german_choice = {}\n",
    "for alt_id in choice_set:\n",
    "    population_work_german_choice[alt_id] = mask_work_german_choice(\n",
    "        data_frame=population, alt_id=alt_id\n",
    "    ).sum()\n",
    "    population_share_work_german_choice[alt_id] = (\n",
    "        population_work_german_choice[alt_id] / population_size\n",
    "    )\n",
    "    print(\n",
    "        f'Size segment work, German, choice={alt_id} in the population:  {population_work_german_choice[alt_id]}'\n",
    "    )\n",
    "    print(\n",
    "        f'Share segment work, German, choice={alt_id} in the population: W(i, x) = '\n",
    "        f'{100*population_share_work_german_choice[alt_id]:.3g}%'\n",
    "    )\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3ab1c6d",
   "metadata": {},
   "source": [
    "Non work, French"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d402dbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mask_non_work_french_choice(data_frame: DataFrame, alt_id: int) -> Series:\n",
    "    \"\"\"Generate a mask for trip purpose = non work, language = French and choice = alt_id\n",
    "\n",
    "    :param data_frame: data frame where the mask is applied\n",
    "    \"\"\"\n",
    "    exogenous_mask = mask_non_work_french(data_frame)\n",
    "    choice_mask = data_frame['Choice'] == alt_id\n",
    "    return exogenous_mask & choice_mask\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36353cfb",
   "metadata": {},
   "source": [
    "Check the number of individuals in the population in this segment.\n",
    "In the notation of the course, this corresponds to $W(i, x)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2010a43",
   "metadata": {},
   "outputs": [],
   "source": [
    "population_non_work_french_choice = {}\n",
    "population_share_non_work_french_choice = {}\n",
    "for alt_id in choice_set:\n",
    "    population_non_work_french_choice[alt_id] = mask_non_work_french_choice(\n",
    "        data_frame=population, alt_id=alt_id\n",
    "    ).sum()\n",
    "    population_share_non_work_french_choice[alt_id] = (\n",
    "        population_non_work_french_choice[alt_id] / population_size\n",
    "    )\n",
    "    print(\n",
    "        f'Size segment non work, French, choice={alt_id} in the population:  {population_non_work_french_choice[alt_id]}'\n",
    "    )\n",
    "    print(\n",
    "        f'Share segment non_work, French, choice={alt_id} in the population: W(i, x) = '\n",
    "        f'{100*population_share_non_work_french_choice[alt_id]:.3g}%'\n",
    "    )\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e20e398",
   "metadata": {},
   "source": [
    "Non work, German"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9366228e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mask_non_work_german_choice(data_frame: DataFrame, alt_id: int) -> Series:\n",
    "    \"\"\"Generate a mask for trip purpose = non work, language = German and choice = alt_id\n",
    "\n",
    "    :param data_frame: data frame where the mask is applied\n",
    "    \"\"\"\n",
    "    exogenous_mask = mask_non_work_german(data_frame)\n",
    "    choice_mask = data_frame['Choice'] == alt_id\n",
    "    return exogenous_mask & choice_mask\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81d40713",
   "metadata": {},
   "source": [
    "Check the number of individuals in the population in this segment.\n",
    "In the notation of the course, this corresponds to $W(i, x)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5d93968",
   "metadata": {},
   "outputs": [],
   "source": [
    "population_non_work_german_choice = {}\n",
    "population_share_non_work_german_choice = {}\n",
    "for alt_id in choice_set:\n",
    "    population_non_work_german_choice[alt_id] = mask_non_work_german_choice(\n",
    "        data_frame=population, alt_id=alt_id\n",
    "    ).sum()\n",
    "    population_share_non_work_german_choice[alt_id] = (\n",
    "        population_non_work_german_choice[alt_id] / population_size\n",
    "    )\n",
    "    print(\n",
    "        f'Size segment non work, German, choice={alt_id} in the population:  {population_non_work_german_choice[alt_id]}'\n",
    "    )\n",
    "    print(\n",
    "        f'Share segment non work, German, choice={alt_id} in the population: W(i, x) = '\n",
    "        f'{100*population_share_non_work_german_choice[alt_id]:.3g}%'\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac11f1f4",
   "metadata": {},
   "source": [
    "## Step 4: Target shares for each alternative"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "390b3f4a",
   "metadata": {},
   "source": [
    "We define the target shares of each alternative within each segment."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a07b53a8",
   "metadata": {},
   "source": [
    "In the notation used in the lecture, we provide $H(i | x)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78cda953",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "target_work_french_choice = {0: 1 / 3, 1: 1 / 3, 2: 1 / 3}\n",
    "target_work_german_choice = {0: 1 / 3, 1: 1 / 3, 2: 1 / 3}\n",
    "target_non_work_french_choice = {0: 1 / 3, 1: 1 / 3, 2: 1 / 3}\n",
    "target_non_work_german_choice = {0: 1 / 3, 1: 1 / 3, 2: 1 / 3}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "276f9cb7",
   "metadata": {},
   "source": [
    "## Step 5: Sampling probability\n",
    "\n",
    "We calculate the sampling probability for each segment, using the following formula:\n",
    "$$R(i, x) = \\frac{H(i, x) N_s}{ W(i, x) N} = \\frac{H(i | x) H(x) N_s}{ W(i , x) N} $$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bf6e0c3",
   "metadata": {},
   "source": [
    "Work, French"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9afe00a",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampling_work_french_choice = {}\n",
    "for alt_id in choice_set:\n",
    "    sampling_work_french_choice[alt_id] = (\n",
    "        target_work_french_choice[alt_id]\n",
    "        * target_work_french\n",
    "        * sample_size\n",
    "        / (population_share_work_french_choice[alt_id] * population_size)\n",
    "    )\n",
    "    print(\n",
    "        f'Sampling probability for work, French and choice = {alt_id}: {sampling_work_french_choice[alt_id]:.3g}'\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3424199",
   "metadata": {},
   "source": [
    "Work, German"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d2172c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampling_work_german_choice = {}\n",
    "for alt_id in choice_set:\n",
    "    sampling_work_german_choice[alt_id] = (\n",
    "        target_work_german_choice[alt_id]\n",
    "        * target_work_german\n",
    "        * sample_size\n",
    "        / (population_share_work_german_choice[alt_id] * population_size)\n",
    "    )\n",
    "    print(\n",
    "        f'Sampling probability for work, German and choice = {alt_id}: {sampling_work_german_choice[alt_id]:.3g}'\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67b6d255",
   "metadata": {},
   "source": [
    "Non work, French"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2984919d",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampling_non_work_french_choice = {}\n",
    "for alt_id in choice_set:\n",
    "    sampling_non_work_french_choice[alt_id] = (\n",
    "        target_non_work_french_choice[alt_id]\n",
    "        * target_non_work_french\n",
    "        * sample_size\n",
    "        / (population_share_non_work_french_choice[alt_id] * population_size)\n",
    "    )\n",
    "    print(\n",
    "        f'Sampling probability for non work, French and choice = {alt_id}: {sampling_non_work_french_choice[alt_id]:.3g}'\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baa89b08",
   "metadata": {},
   "source": [
    "Non work, German"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11365242",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampling_non_work_german_choice = {}\n",
    "for alt_id in choice_set:\n",
    "    sampling_non_work_german_choice[alt_id] = (\n",
    "        target_non_work_german_choice[alt_id]\n",
    "        * target_non_work_german\n",
    "        * sample_size\n",
    "        / (population_share_non_work_german_choice[alt_id] * population_size)\n",
    "    )\n",
    "    print(\n",
    "        f'Sampling probability for non work, German and choice = {alt_id}: {sampling_non_work_german_choice[alt_id]:.3g}'\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bd74fb1",
   "metadata": {},
   "source": [
    "## Step 6: Sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a48cd3a",
   "metadata": {},
   "source": [
    "To perform the sampling, we include two new columns in the data frame. The first one associates each individual with\n",
    "her sampling probability. The second generates a random draw for each individual."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ec824ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "column_sampling = 'sampling_proba'\n",
    "\n",
    "for alt_id in choice_set:\n",
    "    population.loc[mask_work_french_choice(population, alt_id), column_sampling] = (\n",
    "        sampling_work_french_choice[alt_id]\n",
    "    )\n",
    "    population.loc[mask_work_german_choice(population, alt_id), column_sampling] = (\n",
    "        sampling_work_german_choice[alt_id]\n",
    "    )\n",
    "    population.loc[mask_non_work_french_choice(population, alt_id), column_sampling] = (\n",
    "        sampling_non_work_french_choice[alt_id]\n",
    "    )\n",
    "    population.loc[mask_non_work_german_choice(population, alt_id), column_sampling] = (\n",
    "        sampling_non_work_german_choice[alt_id]\n",
    "    )\n",
    "\n",
    "column_random_draw = 'random_draw'\n",
    "population[column_random_draw] = np.random.uniform(0, 1, size=population_size)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0722f3ef",
   "metadata": {},
   "source": [
    "We keep only the individuals such that the random draw is lesser than the sampling  probability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "402afc67",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = population[population[column_random_draw] < population[column_sampling]]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf5e9cf5",
   "metadata": {},
   "source": [
    "## Step 7: Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61b14b58",
   "metadata": {},
   "source": [
    "We now investigate the composition of the sample."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c29df5f",
   "metadata": {},
   "source": [
    "Sample size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "917e58de",
   "metadata": {},
   "outputs": [],
   "source": [
    "actual_sample_size = len(sample)\n",
    "print(f'Sample size: {actual_sample_size}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a319fff1",
   "metadata": {},
   "source": [
    "We now investigate each segment defining the strategy. For each of them, we compare the actual share in the\n",
    "sample with the target share."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0943ba44",
   "metadata": {},
   "outputs": [],
   "source": [
    "for alt_id in choice_set:\n",
    "    the_sample_segment = mask_work_french_choice(data_frame=sample, alt_id=alt_id).sum()\n",
    "    target = target_work_french * target_work_french_choice[alt_id]\n",
    "    print(\n",
    "        f'Work, French, alt {alt_id}: {the_sample_segment} observations '\n",
    "        f'[{100* the_sample_segment/actual_sample_size:.3g}%]. Target: {100*target:.3g}%'\n",
    "    )\n",
    "\n",
    "for alt_id in choice_set:\n",
    "    the_sample_segment = mask_work_german_choice(data_frame=sample, alt_id=alt_id).sum()\n",
    "    target = target_work_german * target_work_german_choice[alt_id]\n",
    "    print(\n",
    "        f'Work, German, alt {alt_id}: {the_sample_segment} observations '\n",
    "        f'[{100* the_sample_segment/actual_sample_size:.3g}%]. Target: {100*target:.3g}%'\n",
    "    )\n",
    "\n",
    "for alt_id in choice_set:\n",
    "    the_sample_segment = mask_non_work_french_choice(\n",
    "        data_frame=sample, alt_id=alt_id\n",
    "    ).sum()\n",
    "    target = target_non_work_french * target_non_work_french_choice[alt_id]\n",
    "    print(\n",
    "        f'Non work, French, alt {alt_id}: {the_sample_segment} observations '\n",
    "        f'[{100* the_sample_segment/actual_sample_size:.3g}%]. Target: {100*target:.3g}%'\n",
    "    )\n",
    "\n",
    "for alt_id in choice_set:\n",
    "    the_sample_segment = mask_non_work_german_choice(\n",
    "        data_frame=sample, alt_id=alt_id\n",
    "    ).sum()\n",
    "    target = target_non_work_german * target_non_work_german_choice[alt_id]\n",
    "    print(\n",
    "        f'Non work, German, alt {alt_id}: {the_sample_segment} observations '\n",
    "        f'[{100* the_sample_segment/actual_sample_size:.3g}%]. Target: {100*target:.3g}%'\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a2be5e0",
   "metadata": {},
   "source": [
    "# Other sampling strategies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0087262a",
   "metadata": {},
   "source": [
    "Now, adapt the procedure to sample using the following strategies:\n",
    "- a simple random sample,\n",
    "- a stratified sample, where the strata are defined based on trip purpose and car availability\n",
    "\n",
    "For each stratified sample strategy, the number of individuals in each stratum in the sample should be the same."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb612243",
   "metadata": {},
   "source": [
    "Note that the variables are coded as follows:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "196c54eb",
   "metadata": {},
   "source": [
    "Trip purpose: 1: work, 2: non work"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28da56df",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(population['TripPurpose'].unique())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1880e2b9",
   "metadata": {},
   "source": [
    "Car availability: 1: available, 3: non available"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecb26d8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(population['CarAvail'].unique())"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}