{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "170ee8d6",
   "metadata": {},
   "source": [
    "\n",
    "File: 03-wesml_solution.py\n",
    "\n",
    "\n",
    "Michel Bierlaire\n",
    "\n",
    "Tue Aug 05 2025, 11:50:56\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6564f731",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import os\n",
    "import pickle\n",
    "from typing import NamedTuple\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tqdm\n",
    "from IPython.core.display_functions import display\n",
    "from biogeme.biogeme import BIOGEME\n",
    "from biogeme.database import Database\n",
    "from biogeme.expressions import Variable, log\n",
    "from biogeme.models import loglogit\n",
    "from biogeme.results_processing import get_pandas_estimated_parameters\n",
    "from pandas import DataFrame, Series\n",
    "\n",
    "from spec_optima import Choice, av, v\n",
    "from stratified_sampling import Population, SamplingStrategy, SegmentSample\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "012c6bfe",
   "metadata": {},
   "source": [
    "The objective of this laboratory is to compare ESML and WESML for an endogenous sampling strategy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa0fb9a7",
   "metadata": {},
   "source": [
    "The specification of the true model is available in `spec_optima.py`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3d250be",
   "metadata": {},
   "source": [
    "A procedure to perform stratified sampling is implemented in the file `stratified_sampling.py`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65262907",
   "metadata": {},
   "source": [
    "# Sampling strategy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13c9f7db",
   "metadata": {},
   "source": [
    "Full population"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43b408e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "population_data_file = 'synthetic_population_with_choice.zip'\n",
    "population_data = pd.read_csv(population_data_file)\n",
    "population = Population(data=population_data, choice_variable='Choice')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdf2559f",
   "metadata": {},
   "source": [
    "Size of the sample and the population"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68637943",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_size = 200\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73d6bd80",
   "metadata": {},
   "source": [
    "We set the seed so that the results are reproducible"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c9e3f47",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(seed=9021967)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42df2113",
   "metadata": {},
   "source": [
    "We define a data structure to store the estimation results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08862d92",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EstimationResults(NamedTuple):\n",
    "    \"\"\"Data structure to store the estimation results.\"\"\"\n",
    "\n",
    "    esml: Series\n",
    "    wesml: Series\n",
    "    esml_corrected: Series\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e3f6ac8",
   "metadata": {},
   "source": [
    "## Endogenously stratified sample"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94e49a38",
   "metadata": {},
   "source": [
    "The strata are defined based on car availability and the chosen alternative.  As the choice is involved in the\n",
    "definition of strata, the stratified sampling is *endogenous*.\n",
    "Note that the segment where the car is available, and the car is chosen, is empty.\n",
    "Target shares in the sample: there are 5 non empty strata, taking 20% each"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f275d06",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def mask_car_available(dataframe: DataFrame):\n",
    "    return dataframe['CarAvail'] == 1\n",
    "\n",
    "\n",
    "car_available_segment = SegmentSample(\n",
    "    name=f'car_available',\n",
    "    mask_generator=mask_car_available,\n",
    "    target_share=0.6,\n",
    "    target_share_alternatives={0: 0.2, 1: 0.2, 2: 0.6},\n",
    ")\n",
    "\n",
    "\n",
    "def mask_car_unavailable(dataframe: DataFrame):\n",
    "    return dataframe['CarAvail'] != 1\n",
    "\n",
    "\n",
    "car_unavailable_segment = SegmentSample(\n",
    "    name=f'car_unavailable',\n",
    "    mask_generator=mask_car_unavailable,\n",
    "    target_share=0.4,\n",
    "    target_share_alternatives={0: 0.7, 1: 0, 2: 0.3},\n",
    ")\n",
    "\n",
    "\n",
    "ess_segments = [car_available_segment, car_unavailable_segment]\n",
    "\n",
    "\n",
    "ess_sampling = SamplingStrategy(\n",
    "    name='ess', population=population, sample_size=sample_size, segments=ess_segments\n",
    ")\n",
    "print(ess_sampling)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74221f6c",
   "metadata": {},
   "source": [
    "# Choice model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e6f8ee4",
   "metadata": {},
   "source": [
    "The specification of the choice model is available in the file <code>spec_optima</code>. Note that it is the exact\n",
    "same specification that has been used to generate the synthetic choices."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbc0c551",
   "metadata": {},
   "source": [
    "The following function extracts a sample from the population and estimates the choice model using standard ESML,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19b8ec7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_and_estimate(\n",
    "    sampling_strategy: SamplingStrategy,\n",
    ") -> EstimationResults:\n",
    "    \"\"\"Estimate the choice model with ESML, WESML, and ESML corrected.\n",
    "\n",
    "    # The \"corrected ESML\" includes the correction terms: log(R_n) = -log(W_n), where R_n is the sampling\n",
    "    # probability, and W_n is its inverse, the weight (the same as used in WESML).\n",
    "\n",
    "    The specification of the choice model is available in the file\n",
    "    spec_optima.py. Note that it is the exact same\n",
    "    specification that has been used to generate the synthetic\n",
    "    choices.\n",
    "\n",
    "    This function extracts a sample from the population and\n",
    "    estimates the choice model using this sample.\n",
    "\n",
    "    :param sampling sampling_strategy: the sampling strategy\n",
    "    :return: estimated parameters\n",
    "\n",
    "    \"\"\"\n",
    "    the_sample = sampling_strategy.sample()\n",
    "    database = Database(the_sample.name, the_sample.sample)\n",
    "    logprob = loglogit(v, av, Choice)\n",
    "    biogeme_esml = BIOGEME(database, logprob, generate_html=False, generate_yaml=False)\n",
    "    biogeme_esml.model_name = f'{the_sample.name}_esml'\n",
    "    results_esml: EstimationResults = biogeme_esml.quick_estimate()\n",
    "    pandas_results_esml = get_pandas_estimated_parameters(\n",
    "        estimation_results=results_esml\n",
    "    )\n",
    "\n",
    "    formulas = {\n",
    "        'log_like': logprob,\n",
    "        'weight': Variable(f'{sampling_strategy.column_weight}'),\n",
    "    }\n",
    "    biogeme_wesml = BIOGEME(\n",
    "        database, formulas, generate_html=False, generate_yaml=False\n",
    "    )\n",
    "    biogeme_wesml.model_name = f'{the_sample.name}_wesml'\n",
    "    results_wesml: EstimationResults = biogeme_wesml.quick_estimate()\n",
    "    pandas_results_wesml = get_pandas_estimated_parameters(\n",
    "        estimation_results=results_wesml\n",
    "    )\n",
    "\n",
    "    v_corrected = {\n",
    "        alt: utility\n",
    "        + log(Variable(f'{sampling_strategy.columns_correction_factor_prefix}_{alt}'))\n",
    "        for alt, utility in v.items()\n",
    "    }\n",
    "    logprob_corrected = loglogit(v_corrected, av, Choice)\n",
    "    biogeme_esml_corrected = BIOGEME(\n",
    "        database, logprob_corrected, generate_html=False, generate_yaml=False\n",
    "    )\n",
    "    biogeme_esml_corrected.model_name = f'{the_sample.name}_esml_corrected'\n",
    "    results_esml_corrected: EstimationResults = biogeme_esml_corrected.quick_estimate()\n",
    "    pandas_results_esml_corrected = get_pandas_estimated_parameters(\n",
    "        estimation_results=results_esml_corrected\n",
    "    )\n",
    "\n",
    "    return EstimationResults(\n",
    "        esml=pandas_results_esml['Value'],\n",
    "        wesml=pandas_results_wesml['Value'],\n",
    "        esml_corrected=pandas_results_esml_corrected['Value'],\n",
    "    )\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46f87146",
   "metadata": {},
   "source": [
    "# Code the experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b07c4791",
   "metadata": {},
   "source": [
    "We first load the true values of the parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91f5152a",
   "metadata": {},
   "outputs": [],
   "source": [
    "true_values_of_the_parameters = {\n",
    "    'asc_car': -1.8610422074360626,\n",
    "    'asc_car_german': -1.0188194471148966,\n",
    "    'asc_pt': -1.6550877823146726,\n",
    "    'asc_pt_car_unavail': 3.5259706461037417,\n",
    "    'asc_pt_german': 0.21100250245252905,\n",
    "    'beta_cost_car': -1.5653923336388653,\n",
    "    'beta_cost_pt': -0.45824318794797375,\n",
    "    'beta_dist': -1.3658487759164655,\n",
    "    'beta_dist_high_school': -1.0242967065023554,\n",
    "    'beta_dist_higher_education': -0.011682106903893624,\n",
    "    'beta_dist_university': 0.43619104548817356,\n",
    "    'beta_time': -1.6340227136243364,\n",
    "    'beta_time_others': -0.588784620423224,\n",
    "    'beta_time_part_time': -0.2778165351485947,\n",
    "    'beta_waiting': -0.22948530603741885,\n",
    "    'beta_waiting_not_work': -0.1501487521368328,\n",
    "}\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dbddf20",
   "metadata": {},
   "source": [
    "The following procedure implements the experiment. It repeats sampling and estimation, stores all estimation results.\n",
    "When it is done, the mean and the standard deviation is calculated for each parameter, as well as the *t*-statistic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10a82d0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_experiment(\n",
    "    sampling_strategy: SamplingStrategy,\n",
    "    repetitions: int,\n",
    ") -> tuple[DataFrame, DataFrame, DataFrame, DataFrame]:\n",
    "    \"\"\"Function that implements the experiment.\n",
    "\n",
    "    It repeats sampling and estimation, stores all estimation\n",
    "    results. When it is done, the mean and the standard deviation is\n",
    "    calculated for each parameter, as well as the *t*-statistic.\n",
    "\n",
    "    :param sampling_strategy: the sampling strategy\n",
    "    :param repetitions: number of times the experiment is performed\n",
    "\n",
    "    :return: four data frames: one with the complete estimation results\n",
    "        (the estimate of each parameter for each repetition) with ESML, WESML, and ESML corrected, and one\n",
    "        with a comparison between the true value and the average of\n",
    "        the estimated values.\n",
    "\n",
    "    \"\"\"\n",
    "    file_name = f'_03_{sampling_strategy.name}_{repetitions}.pickle'\n",
    "    if os.path.exists(file_name):\n",
    "        with open(file_name, 'rb') as f:\n",
    "            results_esml, results_wesml, results_esml_corrected, comparisons = (\n",
    "                pickle.load(f)\n",
    "            )\n",
    "            print(f'Loaded cached results from {file_name}')\n",
    "        return results_esml, results_wesml, results_esml_corrected, comparisons\n",
    "\n",
    "    list_of_results_esml = []\n",
    "    list_of_results_wesml = []\n",
    "    list_of_results_esml_corrected = []\n",
    "\n",
    "    for _ in tqdm.tqdm(range(repetitions)):\n",
    "        estimation_results = sample_and_estimate(sampling_strategy)\n",
    "        row_esml = estimation_results.esml\n",
    "        list_of_results_esml.append(row_esml.to_frame().T)\n",
    "\n",
    "        row_esml_corrected = estimation_results.esml_corrected\n",
    "        list_of_results_esml_corrected.append(row_esml_corrected.to_frame().T)\n",
    "\n",
    "        row_wesml = estimation_results.wesml\n",
    "        list_of_results_wesml.append(row_wesml.to_frame().T)\n",
    "\n",
    "    results_esml = pd.concat(list_of_results_esml, ignore_index=True)\n",
    "    results_esml_corrected = pd.concat(\n",
    "        list_of_results_esml_corrected, ignore_index=True\n",
    "    )\n",
    "    results_wesml = pd.concat(list_of_results_wesml, ignore_index=True)\n",
    "\n",
    "    comparisons = pd.DataFrame.from_dict(\n",
    "        true_values_of_the_parameters, orient='index', columns=['True']\n",
    "    )\n",
    "    # Set the columns of results to match the parameter names (comparisons.index)\n",
    "    results_esml.columns = comparisons.index\n",
    "    results_esml_corrected.columns = comparisons.index\n",
    "    results_wesml.columns = comparisons.index\n",
    "\n",
    "    comparisons['Estimated_esml'] = results_esml.mean()\n",
    "    comparisons['StdDev_esml'] = results_esml.std()\n",
    "    comparisons['t-test_esml'] = (\n",
    "        comparisons['Estimated_esml'] - comparisons['True']\n",
    "    ) / comparisons['StdDev_esml']\n",
    "\n",
    "    comparisons['Estimated_wesml'] = results_wesml.mean()\n",
    "    comparisons['StdDev_wesml'] = results_wesml.std()\n",
    "    comparisons['t-test_wesml'] = (\n",
    "        comparisons['Estimated_wesml'] - comparisons['True']\n",
    "    ) / comparisons['StdDev_wesml']\n",
    "\n",
    "    comparisons['Estimated_esml_corrected'] = results_esml_corrected.mean()\n",
    "    comparisons['StdDev_esml_corrected'] = results_esml_corrected.std()\n",
    "    comparisons['t-test_esml_corrected'] = (\n",
    "        comparisons['Estimated_esml_corrected'] - comparisons['True']\n",
    "    ) / comparisons['StdDev_esml_corrected']\n",
    "\n",
    "    with open(file_name, 'wb') as f:\n",
    "        pickle.dump(\n",
    "            (results_esml, results_wesml, results_esml_corrected, comparisons), f\n",
    "        )\n",
    "        print(f'Saved results to {file_name}')\n",
    "    return results_esml, results_wesml, results_esml_corrected, comparisons\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d47d1fa",
   "metadata": {},
   "source": [
    "# Run the experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "975a318c",
   "metadata": {},
   "outputs": [],
   "source": [
    "the_repetitions = 200\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebd2fde6",
   "metadata": {},
   "source": [
    "## Endogenous sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7f6fbaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "results_ess1_esml, results_ess1_wesml, results_ess1_esml_corrected, comparisons_ess1 = (\n",
    "    run_experiment(sampling_strategy=ess_sampling, repetitions=the_repetitions)\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c687996",
   "metadata": {},
   "source": [
    "Here are the comparisons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b9e1a26",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(comparisons_ess1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dec4cf2a",
   "metadata": {},
   "source": [
    "Analyze the above results.\n",
    "- What general observation can be made about the three estimators?\n",
    "- Focus now on the constants `asc_car` and `ASC_PT`. What observation can be made about the three estimates?\n",
    "- Evaluate the average precision of WESML and ESML (corrected). Which one is better?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25ace9eb",
   "metadata": {},
   "source": [
    "When performing the test of the hypothesis that the estimated value is equal to the true value, we observe the\n",
    "following:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa8e488a",
   "metadata": {},
   "source": [
    "From the column `t-test_esml`, corresponding to ESML, it is seen that all parameters are correctly recovered,\n",
    "except the constants:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf5f6da5",
   "metadata": {},
   "outputs": [],
   "source": [
    "asc_car_true = comparisons_ess1.loc['asc_car', 'True']\n",
    "asc_car_esml = comparisons_ess1.loc['asc_car', 'Estimated_esml']\n",
    "asc_car_esml_ttest = comparisons_ess1.loc['asc_car', 't-test_esml']\n",
    "print(f'asc_car true value: {asc_car_true:.3g}')\n",
    "print(f'asc_car ESML estimate: {asc_car_esml:.3g}')\n",
    "print(f'asc_car ESML estimate (t-test): {asc_car_esml_ttest:.3g}')\n",
    "\n",
    "asc_pt_true = comparisons_ess1.loc['asc_pt', 'True']\n",
    "asc_pt_esml = comparisons_ess1.loc['asc_pt', 'Estimated_esml']\n",
    "asc_pt_esml_ttest = comparisons_ess1.loc['asc_pt', 't-test_esml']\n",
    "print(f'asc_pt true value: {asc_pt_true:.3g}')\n",
    "print(f'asc_pt ESML estimate: {asc_pt_esml:.3g}')\n",
    "print(f'asc_pt ESML estimate (t-test): {asc_pt_esml_ttest:.3g}')\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5056528c",
   "metadata": {},
   "source": [
    "From the column `t-test_wesml`, corresponding to WESML, it is seen that all parameters are correctly recovered,\n",
    "including the constants."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "702a96ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "asc_car_true = comparisons_ess1.loc['asc_car', 'True']\n",
    "asc_car_wesml = comparisons_ess1.loc['asc_car', 'Estimated_wesml']\n",
    "asc_car_wesml_ttest = comparisons_ess1.loc['asc_car', 't-test_wesml']\n",
    "print(f'asc_car true value: {asc_car_true:.3g}')\n",
    "print(f'asc_car WESML estimate: {asc_car_wesml:.3g}')\n",
    "print(f'asc_car WESML estimate (t-test): {asc_car_wesml_ttest:.3g}')\n",
    "\n",
    "asc_pt_true = comparisons_ess1.loc['asc_pt', 'True']\n",
    "asc_pt_wesml = comparisons_ess1.loc['asc_pt', 'Estimated_wesml']\n",
    "asc_pt_wesml_ttest = comparisons_ess1.loc['asc_pt', 't-test_wesml']\n",
    "print(f'asc_pt true value: {asc_pt_true:.3g}')\n",
    "print(f'asc_pt WESML estimate: {asc_pt_wesml:.3g}')\n",
    "print(f'asc_pt WESML estimate (t-test): {asc_pt_wesml_ttest:.3g}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "087df4c9",
   "metadata": {},
   "source": [
    "From the column `t-test_esml_corrected`, corresponding to ESML with the correction of the constants, it is seen\n",
    "that all parameters are correctly recovered, including the (corrected) constants."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80aab990",
   "metadata": {},
   "outputs": [],
   "source": [
    "asc_car_true = comparisons_ess1.loc['asc_car', 'True']\n",
    "asc_car_esml_corrected = comparisons_ess1.loc['asc_car', 'Estimated_esml_corrected']\n",
    "asc_car_esml__corrected_ttest = comparisons_ess1.loc['asc_car', 't-test_esml_corrected']\n",
    "print(f'asc_car true value: {asc_car_true:.3g}')\n",
    "print(f'asc_car ESML estimate (corrected): {asc_car_esml_corrected:.3g}')\n",
    "print(f'asc_car ESML estimate (corrected, t-test): {asc_car_esml__corrected_ttest:.3g}')\n",
    "\n",
    "asc_pt_true = comparisons_ess1.loc['asc_pt', 'True']\n",
    "asc_pt_esml_corrected = comparisons_ess1.loc['asc_pt', 'Estimated_esml_corrected']\n",
    "asc_pt_esml_ttest_corrected = comparisons_ess1.loc['asc_pt', 't-test_esml_corrected']\n",
    "print(f'asc_pt true value: {asc_pt_true:.3g}')\n",
    "print(f'asc_pt ESML estimate (corrected): {asc_pt_esml_corrected:.3g}')\n",
    "print(f'asc_pt ESML estimate (corrected, t-test): {asc_pt_esml_ttest_corrected:.3g}')\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea005136",
   "metadata": {},
   "source": [
    "If we now compare the standard deviation for the two estimators, it appears that using ESML and correcting the\n",
    "constants provided more precise estimates than WESML."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dee38b3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "std_dev_wesml = comparisons_ess1['StdDev_wesml'].abs().mean()\n",
    "std_dev_esml_corrected = comparisons_ess1['StdDev_esml_corrected'].abs().mean()\n",
    "print(f'Std. dev. WESML: {std_dev_wesml:.3g}')\n",
    "print(f'Std. dev. ESML (corrected): {std_dev_esml_corrected:.3g}')"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}