notebooks/introduction/4_model_customization.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install necessary requirements\n",
    "\n",
    "# If you run this notebook on Google Colab, or in standalone mode, you need to install the required packages.\n",
    "# Uncomment the following lines:\n",
    "\n",
    "# !pip install choice-learn\n",
    "\n",
    "# If you run the notebook within the GitHub repository, you need to run the following lines, that can skipped otherwise:\n",
    "import os\n",
    "import sys\n",
    "\n",
    "sys.path.append(\"../../\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# Remove GPU use\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to customization\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/artefactory/choice-learn/blob/main/notebooks/introduction/4_model_customization.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Choice-Learn package aims at providing structure and helpful functions in order to design any choice model. The main idea is to write the utility function and let the package work its magic.\n",
    "It is recommended to read the [data tutorial](./2_data_handling.ipynb) before to understand the ChoiceDataset class.\n",
    "\n",
    "## Summary\n",
    "- [BaseClass: ChoiceModel](#baseclass-choicemodel)    \n",
    "    - [EndPoints](#the-different-endpoints)\n",
    "    - [Parameters](#parameters)\n",
    "    - [SubClassing](#subclassing)\n",
    "- [Example 1: Rewriting Conditional Logit as custom model](#example-1-rewriting-the-conditional-mnl-on-modecanada)\n",
    "- [Example 2: Defining a non-linear utility function with TensorFlow](#example-2-defining-a-non-linear-utility-function-with-tensorflow)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BaseClass: ChoiceModel\n",
    "\n",
    "Choice-Learn models are built on the ChoiceModel base class and most of them follow the same structure.\\\n",
    "In this tutorial, we will delve into the details of modelling and the possibilities of the package. In particular we will see how Choice-Learn helps for manual formulation of a choice model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's import ChoiceModel\n",
    "\n",
    "from choice_learn.models.base_model import ChoiceModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The different EndPoints\n",
    "\n",
    "The ChoiceModel class revolves around several methods that are shared by most models:\n",
    "- <ins>Model Specification</ins>\\\n",
    "    *.\\_\\_init\\_\\_()* and/or *.instantiate()* are used to specify the form of the model\n",
    "\n",
    "- <ins>Model Estimation</ins>\\\n",
    "    *.fit()* uses a ChoiceDataset to find the best values for the different trainable weights\n",
    "    \n",
    "- <ins>Use of the model</ins>\\\n",
    "    *.evaluate()* can be used to estimate the negative log likelihood of the model's choice probabilities compared to the ground truth from a ChoiceDataset\\\n",
    "    *.predict_probas()* can be used to predict the model's choice probabilities related to a ChoiceDataset\\\n",
    "    *.compute_batch_utility()* can be used to predict a batch items utilities\n",
    "\n",
    "### Parameters\n",
    "\n",
    "A few parameters are shared through the ChoiceModel class and can be changed. A full list is [available](../../choice_learn/models/base_model.py), here are the most useful:\n",
    "\n",
    "- **optimizer**: Name of the optimizer to use. Default is lbfgs\n",
    "    - Non-stochastic: It is recommended to use them - and in particular lbfgs - for smaller datasets and models. It is faster but needs all data in memory, therefore the batch_size argument is not used. More info on the TensorFlow [documentation](https://www.tensorflow.org/probability/examples/Optimizers_in_TensorFlow_Probability).\n",
    "    - Stochastic Gradient Descent optimizers - such as Adam. They will lead to slower convergence but work well with batching. List is [here](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers).\n",
    "- **batch_size**: Data batch size to use when stochastic gradient descent optimizer is used. Default is 32.\n",
    "- **lr:** Learning Rate of the optimizer to use when stochastic gradient descent optimizer is used. Default is 0.001.\n",
    "- **epochs:** Max number of iterations before stopping optimization. Default is 1000.\n",
    "\n",
    "## Subclassing\n",
    "\n",
    "Inheritance is used for better code formatting in Choice-Learn. It is also optimized to let anyone *easily* define its own utility model. The idea is that by subclassing ChoiceModel one only needs to define the utility function with TensorFlow for it to work.\\\n",
    "The advantages are twofold:\n",
    "- It needs little time. An example will follow to show you how it can be done in a few minutes.\n",
    "- It is possible to use non-linear formulations of the utility. As long as it is written with [TensorFlow operations](https://www.tensorflow.org/api_docs/python/tf/math), Choice-Learn and TensorFlow handle the optimization. For the more adventurers, you can even [define your own operations](https://www.tensorflow.org/api_docs/python/tf/custom_gradient) as long as you provide the gradients."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 1: Rewriting the conditional MNL on ModeCanada\n",
    "We download the ModeCanada dataset as a ChoiceDataset, see [here](./2_data_handling.ipynb) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from choice_learn.datasets import load_modecanada\n",
    "\n",
    "dataset = load_modecanada(as_frame=False, preprocessing=\"tutorial\", add_items_one_hot=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will subclass the parent class ChoiceModel that we need to import. It mainly works with TensorFlow as a backend, it is thus recommended to use  their operation as much as possible. Most NumPy operations have a TensorFlow equivalent. You can look at the documentation [here](https://www.tensorflow.org/api_docs/python/tf).\n",
    "\n",
    "For our custom model to work, we need to specify:\n",
    "- Weights initialization in __init__()\n",
    "- the utility function in compute_batch_utility()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from choice_learn.models.base_model import ChoiceModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### *Utility formulation*\n",
    "\n",
    "Following the Conditional Logit tutorial we want to estimate the following utility function:\n",
    "$$\n",
    "U(i, s) = \\beta^{inter}_i + \\beta^{price} \\cdot price(i, s) + \\beta^{freq} \\cdot freq(i, s) + \\beta^{ovt} \\cdot ovt(i, s) + \\beta^{income}_i \\cdot income(s) + \\beta^{ivt}_i \\cdot ivt(i, t) + \\epsilon(i, t)\n",
    "$$\n",
    "You can check the cLogit example for more details\n",
    "\n",
    "### *Coefficients Initialization*\n",
    "\n",
    "Following our utility formula we need four coefficients vectors:\n",
    "- $\\beta^{inter}$ has 3 values\n",
    "- $\\beta^{price}$, $\\beta^{freq}$, $\\beta^{ovt}$ are regrouped and each has one value, shared by all items\n",
    "- $\\beta^{income}$ has 3 values\n",
    "- $\\beta^{ivt}$ has 4 values\n",
    "\n",
    "### *Utility Computation*\n",
    "\n",
    "In the method compute_utility, we need to define how to estimate each item utility for each choice using  the features and initialized weights.\n",
    "The arguments of the function are a batch of each features type of the ChoiceDataset class:\n",
    "\n",
    "| Order | Argument | shape | Features for ModeCanada| \n",
    "|---|---|---|---|\n",
    "| 2 | shared_features_by_choice | (batch_size, n_shared_features) | Customer Income | \n",
    "| 3 | items_features_by_choice | (batch_size, n_items, n_items_features) | Cost, Freq, Ivt, Ovt values of each mode | \n",
    "| 4 | available_items_by_choice | (batch_size, n_items) | Not Used | \n",
    "| 5 | choices | (batch_size, ) | Not Used | \n",
    "\n",
    "batch_size represents the number of choices given in the batch.\n",
    "The method needs to return the utilities, in the form of a matrix of shape (n_choices, n_items), representing the utility of each item for each choice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can verify the names and order of the features:\n",
    "print(dataset.shared_features_by_choice_names, dataset.items_features_by_choice_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomCanadaConditionalLogit(ChoiceModel):\n",
    "    \"\"\"Conditional Logit following for ModeCanada.\n",
    "\n",
    "    Arguments:\n",
    "    ----------\n",
    "    optimizer : str\n",
    "        tf.keras.optimizer to use for training, default is Adam\n",
    "    lr: float\n",
    "        learning rate for optimizer, default is 1e-3\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        add_exit_choice=False, # Whether to add exit choice with utility=1\n",
    "        optimizer=\"lbfgs\", # Optimizer to use\n",
    "        tolerance=1e-8, # Absolute function tolerance for optimality if lbfgs is used\n",
    "        lr=0.001, # learning rate if stochastic gradient descent optimizer\n",
    "        epochs=1000, # maximum number of epochs\n",
    "        batch_size=32, # batch size if stochastic gradient descent optimizer\n",
    "    ):\n",
    "        \"\"\"Model coefficients instantiation.\"\"\"\n",
    "        super().__init__(add_exit_choice=add_exit_choice,\n",
    "                         optimizer=optimizer,\n",
    "                         tolerance=tolerance,\n",
    "                         lr=lr,\n",
    "                         epochs=epochs,\n",
    "                         batch_size=batch_size)\n",
    "\n",
    "        # Create model weights. Basically is one weight by feature + one for intercept\n",
    "        self.beta_inter = tf.Variable(tf.random_normal_initializer(0.0, 0.02, seed=42)(shape=(1, 3)),\n",
    "                                 name=\"beta_inter\")\n",
    "        self.beta_freq_cost_ovt = tf.Variable(\n",
    "            tf.random_normal_initializer(0.0, 0.02, seed=42)(shape=(1, 3)),\n",
    "            name=\"beta_freq_cost_ovt\"\n",
    "            )\n",
    "        self.beta_income = tf.Variable(tf.random_normal_initializer(0.0, 0.02, seed=42)(shape=(1, 3)),\n",
    "                             name=\"beta_income\")\n",
    "        self.beta_ivt = tf.Variable(tf.random_normal_initializer(0.0, 0.02, seed=42)(shape=(1, 4)),\n",
    "                               name=\"beta_ivt\")\n",
    "\n",
    "    # Do not forget to add them to the list of trainable_weights, it is mandatory !\n",
    "    @property\n",
    "    def trainable_weights(self):\n",
    "        \"\"\"Do not forget to add the weights to the list of trainable_weights.\n",
    "        \n",
    "        It is needed to use the @property definition as here.\n",
    "\n",
    "        Return:\n",
    "        -------\n",
    "        list:\n",
    "            list of tf.Variable to be optimized\n",
    "        \"\"\"\n",
    "        return [self.beta_inter, self.beta_freq_cost_ovt, self.beta_income, self.beta_ivt]\n",
    "\n",
    "\n",
    "    def compute_batch_utility(self,\n",
    "                              shared_features_by_choice,\n",
    "                              items_features_by_choice,\n",
    "                              available_items_by_choice,\n",
    "                              choices):\n",
    "        \"\"\"Method that defines how the model computes the utility of a product.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        shared_features_by_choice : tuple of np.ndarray (choices_features)\n",
    "            a batch of shared features\n",
    "            Shape must be (n_choices, n_shared_features)\n",
    "        items_features_by_choice : tuple of np.ndarray (choices_items_features)\n",
    "            a batch of items features\n",
    "            Shape must be (n_choices, n_items_features)\n",
    "        available_items_by_choice : np.ndarray\n",
    "            A batch of items availabilities\n",
    "            Shape must be (n_choices, n_items)\n",
    "        choices_batch : np.ndarray\n",
    "            Choices\n",
    "            Shape must be (n_choices, )\n",
    "\n",
    "        Returns:\n",
    "        --------\n",
    "        np.ndarray\n",
    "            Utility of each product for each choice.\n",
    "            Shape must be (n_choices, n_items)\n",
    "        \"\"\"\n",
    "        _ = (available_items_by_choice, choices)  # Avoid unused variable warning\n",
    "\n",
    "        # Adding the 0 value intercept of first item to get the right shape\n",
    "        full_beta_inter = tf.concat([tf.constant([[.0]]), self.beta_inter], axis=-1)\n",
    "        # Concatenation to reach right shape for dot product\n",
    "        full_beta_income = tf.concat([tf.constant([[.0]]), self.beta_income], axis=-1)  # shape = (1, n_items)\n",
    "\n",
    "        items_ivt_by_choice = items_features_by_choice[:, :, 3] # shape = (n_choices, n_items, )\n",
    "        items_cost_freq_ovt_by_choice = items_features_by_choice[:, :, :3 ]# shape = (n_choices, n_items, 3)\n",
    "        u_cost_freq_ovt = tf.squeeze(tf.tensordot(items_cost_freq_ovt_by_choice,\n",
    "                                                  tf.transpose(self.beta_freq_cost_ovt), axes=1)) # shape = (n_choices, n_items)\n",
    "        u_ivt = tf.multiply(items_ivt_by_choice, self.beta_ivt) # shape = (n_choices, n_items)\n",
    "\n",
    "        u_income = tf.tensordot(shared_features_by_choice, full_beta_income, axes=1)  # shape = (n_choices, n_items)\n",
    "\n",
    "        # Reshaping the intercept that is constant over all choices (n_items, ) -> (n_choices, n_items)\n",
    "        u_intercept = tf.concat([full_beta_inter] * (u_income.shape[0]), axis=0)\n",
    "        return u_intercept + u_cost_freq_ovt + u_income + u_ivt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.items_features_by_choice[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = CustomCanadaConditionalLogit()\n",
    "history = model.fit(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decomposition of the utility operations\n",
    "\n",
    "#### <ins>*Intercept*</ins>\n",
    "\n",
    "- $U_{inter}[air, s] = \\beta^{inter}_{air} = 0$\n",
    "- $U_{inter}[bus, s] = \\beta^{inter}_{bus}$\n",
    "- $U_{inter}[car, s] = \\beta^{inter}_{car}$\n",
    "- $U_{inter}[train, s] = \\beta^{inter}_{train}$\n",
    "\n",
    "$\\beta^{inter} = \\left(\\begin{array}{c} \n",
    "0 \\\\\n",
    "\\beta^{inter}_{bus} \\\\\n",
    "\\beta^{inter}_{car} \\\\\n",
    "\\beta^{inter}_{train} \\\\\n",
    "\\end{array}\\right)$\n",
    "\n",
    "$U_{inter} = \\beta^{inter.T}$\n",
    "\n",
    "#### <ins>*Price, Freq, OVT*</ins>\n",
    "- $U_{price, freq, ovt}[air, s] = \\beta^{price} \\cdot price[air, s] + \\beta^{freq} \\cdot freq[air, s] + \\beta^{ovt} \\cdot ovt[air, s]$\n",
    "- $U_{price, freq, ovt}[bus, s] = \\beta^{price} \\cdot price[bus, s] + \\beta^{freq} \\cdot freq[bus, s] + \\beta^{ovt} \\cdot ovt[bus, s]$\n",
    "- $U_{price, freq, ovt}[car, s] = \\beta^{price} \\cdot price[car, s) + \\beta^{freq} \\cdot freq[car, s] + \\beta^{ovt} \\cdot ovt(car, s]$\n",
    "- $U_{price, freq, ovt}[train, s] = \\beta^{price} \\cdot price[train, s] + \\beta^{freq} \\cdot freq[train, s] + \\beta^{ovt} \\cdot ovt[train, s]$\n",
    "\n",
    "$\\beta^{price, freq, ovt} = \\left(\\begin{array}{c} \n",
    "\\beta^{price} \\\\\n",
    "\\beta^{freq} \\\\\n",
    "\\beta^{ovt} \\\\\n",
    "\\end{array}\\right)$ and $items\\_feature\\_by\\_choice[0, :3] = \\left(\\begin{array}{ccc} \n",
    "price[air, 0] & freq[air, 0] & ovt[air, 0] \\\\\n",
    "price[bus, 0] & freq[bus, 0] & ovt[bus, 0] \\\\\n",
    "price[car, 0] & freq[car, 0] & ovt[car, 0] \\\\\n",
    "price[train, 0] & freq[train, 0] & ovt[train, 0] \\\\\n",
    "\\end{array}\\right)$\n",
    "\n",
    "$U_{price, freq, ovt} = \\beta^{price, freq, ovt .T} \\cdot items\\_feature\\_by\\_choice[:, :3]$\n",
    "\n",
    "Note that in the matrix we didn't illustrate the choices dimension, explaining the [0, :3] -> [:, :3].\n",
    "items_features_by_choice[:, :3] has a shape of (batch_size, 4, 3) and $ \\beta^{price, freq, ovt}$ a shape of (1, 3).\n",
    "Resulting $U_{price, freq, ovt} $ has therefore a shape of (batch_size, 4)\n",
    "\n",
    "#### <ins>*IVT*</ins>\n",
    "- $U_{ivt}[air, s] = \\beta^{ivt}_{air} \\cdot ivt[air, s]$\n",
    "- $U_{ivt}[bus, s] = \\beta^{ivt}_{bus} \\cdot ivt[bus, s]$\n",
    "- $U_{ivt}[car, s] = \\beta^{ivt}_{car} \\cdot ivt[car, s]$\n",
    "- $U_{ivt}[train, s] = \\beta^{ivt}_{train} \\cdot ivt[train, s]$\n",
    "\n",
    "$\\beta^{ivt} = \\left(\\begin{array}{c} \n",
    "\\beta^{ivt}_{air} \\\\\n",
    "\\beta^{ivt}_{bus} \\\\\n",
    "\\beta^{ivt}_{car}\\\\\n",
    "\\beta^{ivt}_{train} \\\\\n",
    "\\end{array}\\right)$\\\n",
    "and\\\n",
    "$items\\_features\\_by\\_choice[:, 3] = \\left(\\begin{array}{cccc} \n",
    "ivt[0, air] & ivt[0, bus] & ivt[0, car] & ivt[0,train] \\\\\n",
    "ivt[1, air] & ivt[1, bus] & ivt[1, car] & ivt[1,train] \\\\\n",
    "... & ... & ... & ... \\\\\n",
    "ivt[batch\\_size, air] & ivt[batch\\_size, bus] & ivt[batch\\_size, car] & ivt[batch\\_size,train] \\\\\n",
    "\\end{array}\\right)$\n",
    "\n",
    "\n",
    "$U_{ivt} = \\beta^{ivt} * items\\_features\\_by\\_choice[:, 3]$ of shape (batch_size, 4)\n",
    "\n",
    "#### <ins>*Income*</ins>\n",
    "- $U_{income}[air, s] = \\beta^{income}_{air} \\cdot income[s]$\n",
    "- $U_{income}[bus, s] = \\beta^{income}_{bus} \\cdot income[s]$\n",
    "- $U_{income}[car, s] = \\beta^{income}_{car} \\cdot income[s]$\n",
    "- $U_{income}[train, s] = \\beta^{income}_{train} \\cdot income[s]$\n",
    "\n",
    "$\\beta^{income} = \\left(\\begin{array}{c} \n",
    "\\beta^{income}_{air} \\\\\n",
    "\\beta^{income}_{bus} \\\\\n",
    "\\beta^{income}_{car}\\\\\n",
    "\\beta^{income}_{train} \\\\\n",
    "\\end{array}\\right)$ and $shared\\_features = \\left(\\begin{array}{c} \n",
    "income[0] \\\\\n",
    "income[1] \\\\\n",
    "... \\\\\n",
    "income[batch\\_size] \\\\\n",
    "\\end{array}\\right)$\n",
    "\n",
    "$U_{income} = \\beta^{income .T} \\cdot shared\\_features$\n",
    "\n",
    "By concatenating batch_size times $U_{inter}$ over the choices we obtain 4 matrixes of shape (batch_size, 4).\n",
    "\n",
    "The final utility is then:\n",
    "$U = U_{inter} + U_{price, freq, ovt} + U_{ivt} + U_{income}$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results\n",
    "We can now test that we obtain the same results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "keep_output": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<tf.Variable 'beta_inter:0' shape=(1, 3) dtype=float32, numpy=array([[0.6983521, 1.8441089, 3.2741907]], dtype=float32)>\n",
      "<tf.Variable 'beta_freq_cost_ovt:0' shape=(1, 3) dtype=float32, numpy=array([[-0.03333881,  0.09252932, -0.0430035 ]], dtype=float32)>\n",
      "<tf.Variable 'beta_income:0' shape=(1, 3) dtype=float32, numpy=array([[-0.08908677, -0.02799308, -0.03814653]], dtype=float32)>\n",
      "<tf.Variable 'beta_ivt:0' shape=(1, 4) dtype=float32, numpy=\n",
      "array([[ 0.05950957, -0.0067836 , -0.00646028, -0.00145035]],\n",
      "      dtype=float32)>\n"
     ]
    }
   ],
   "source": [
    "print(model.trainable_weights[0])\n",
    "print(model.trainable_weights[1])\n",
    "print(model.trainable_weights[2])\n",
    "print(model.trainable_weights[3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The coefficients are organized differently but reach the same values. It is also the case for negative log-lilkelihood:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "keep_output": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total Neg LikeliHood; tf.Tensor(1874.363, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "print(\"Total Neg LikeliHood;\", model.evaluate(dataset) * len(dataset))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 2: Defining a non-linear utility function with TensorFlow\n",
    "\n",
    "In this example we have used a simple linear function for utility computation. We could use any function we would like. Particularly we can use neural networks and activation functions to add non-linearities.\n",
    "\n",
    "A simple example would be:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.layers import Dense\n",
    "\n",
    "class NeuralNetUtility(ChoiceModel):\n",
    "    def __init__(self, n_neurons, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.n_neurons = n_neurons\n",
    "\n",
    "        # Items Features Layer\n",
    "        self.dense_items_features = Dense(units=n_neurons, activation=\"elu\")\n",
    "\n",
    "        # Shared Features Layer\n",
    "        self.dense_shared_features = Dense(units=n_neurons, activation=\"elu\")\n",
    "\n",
    "        # Third layer: embeddings to utility (dense representation of features > U)\n",
    "        self.final_layer = Dense(units=1, activation=\"linear\")\n",
    "\n",
    "    # We do not forget to specify self.trainable_weights with all coefficients that need to be estimated.\n",
    "    # Small trick using @property to acces future weights of layers\n",
    "    # that have not been instantiated yet !\n",
    "    @property\n",
    "    def trainable_weights(self):\n",
    "        \"\"\"Endpoint to acces model's trainable_weights.\n",
    "\n",
    "        Returns:\n",
    "        --------\n",
    "        list\n",
    "            list of trainable_weights\n",
    "        \"\"\"\n",
    "        return self.dense_items_features.trainable_variables\\\n",
    "              + self.dense_shared_features.trainable_variables\\\n",
    "                  + self.final_layer.trainable_variables\n",
    "\n",
    "    def compute_batch_utility(self,\n",
    "                              shared_features_by_choice,\n",
    "                              items_features_by_choice,\n",
    "                              available_items_by_choice,\n",
    "                              choices):\n",
    "        \"\"\"Computes batch utility from features.\"\"\"\n",
    "        _, _ = available_items_by_choice, choices\n",
    "        # We apply the neural network to all items_features_by_choice for all the items\n",
    "        # We then concatenate the utilities of each item of shape (n_choices, 1) into a single one of shape (n_choices, n_items)\n",
    "        shared_features_embeddings = self.dense_shared_features(shared_features_by_choice)\n",
    "\n",
    "        items_features_embeddings = []\n",
    "        for i in range(items_features_by_choice[0].shape[1]):\n",
    "            # Utility is Dense(embeddings sum)\n",
    "            item_embedding = shared_features_embeddings + self.dense_items_features(items_features_by_choice[:, i])\n",
    "            items_features_embeddings.append(self.final_layer(item_embedding))\n",
    "\n",
    "        # Concatenation to get right shape (n_choices, n_items, )\n",
    "        item_utility_by_choice = tf.concat(items_features_embeddings, axis=1)\n",
    "\n",
    "        return item_utility_by_choice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = NeuralNetUtility(n_neurons=10, optimizer=\"Adam\", epochs=200)\n",
    "history = model.fit(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.evaluate(dataset) * len(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want more complex examples, you can look at the following implementations:\n",
    "- [RUMnet](../../choice_learn/models/rumnet.py)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}