notebooks/03 Monte Carlo Tree Search on the PeMS-Bay Dataset.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uRK534z7pYNU"
      },
      "source": [
        "<center>\n",
        "    <h1>Verbal Explanation of Spatial Temporal GNNs for Traffic Forecasting</h1>\n",
        "    <h2>Monte Carlo Tree Search on the PeMS-Bay dataset</h2>\n",
        "</center>\n",
        "\n",
        "---\n",
        "\n",
        "In this notebook the explanations of the predicted events of the *STGNN* on the *PeMS-Bay* dataset are computed by means of an *Explainer*. The explainer is an MCTS algorithm that performs a localized search on the input graph of the instance to be explained. It returns the important subgraph which serves as the explanation of the predicted selected event outputed by the STGNN. The search space is cut by a transparent *global heuristic* that exploits traffic domain-specific characteristics and limits the search on a subgraph of the input graph composed of the top scoring nodes according to the global heuristic. The explainer is inspired by the paper *Explaining temporal graph models through an explorer-navigator framework* <a name=\"cite_paper\"></a>[<sup>[1]</sup>](#note_paper).\n",
        "\n",
        "Firstly, grid search is applied to the training set in order to tune the hyperparameters of the explainer, afterwards, the important subgraphs of each dataset are extracted and the results are evaluated.\n",
        "\n",
        "For more detailed informations about the used functions, look into the corresponding docstrings inside the python files, inside the `src` folder.\n",
        "\n",
        "---\n",
        "<small>\n",
        "\n",
        "<a name=\"note_paper\"></a>[1]\n",
        "W. Xia et al. “Explaining temporal graph models through an explorer-navigator framework”. In: *The\n",
        "Eleventh International Conference on Learning Representations*, 2023.\n",
        "URL: https://openreview.net/forum?id=BR_ZhvcYbGJ.\n",
        "</small>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "6kbAULJKy-wu"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "import os\n",
        "\n",
        "# Set the main path in the root folder of the project.\n",
        "sys.path.append(os.path.join('..'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HGIZtEY6y-wv",
        "outputId": "99c3e01a-c95e-4f12-a777-39f2ad07ae22"
      },
      "outputs": [],
      "source": [
        "# Settings for autoreloading.\n",
        "%load_ext autoreload\n",
        "%autoreload 2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "obBqbIuGy-wv"
      },
      "outputs": [],
      "source": [
        "from src.utils.seed import set_random_seed\n",
        "\n",
        "# Set the random seed for deterministic operations.\n",
        "SEED = 42\n",
        "set_random_seed(SEED)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uHGuhgDmy-wv",
        "outputId": "be71123b-5936-490e-d0e7-cc64643f41a4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "The selected device is: \"cuda\"\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "\n",
        "# Set the device for training and querying the model.\n",
        "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "print(f'The selected device is: \"{DEVICE}\"')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xhVgau7-y-ww"
      },
      "source": [
        "# 1 Loading the Data\n",
        "In this section the data is loaded"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "HlTWno-1y-wx"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "BASE_DATA_DIR = os.path.join('..', 'data', 'pems-bay')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "BbiHBGWSy-wx"
      },
      "outputs": [],
      "source": [
        "import pickle\n",
        "with open(os.path.join(BASE_DATA_DIR, 'processed', 'scaler.pkl'), 'rb') as f:\n",
        "    scaler = pickle.load(f)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "nb6d7kVZy-wx"
      },
      "outputs": [],
      "source": [
        "from src.spatial_temporal_gnn.model import SpatialTemporalGNN\n",
        "from src.data.data_extraction import get_adjacency_matrix\n",
        "\n",
        "# Get the adjacency matrix\n",
        "adj_matrix_structure = get_adjacency_matrix(\n",
        "    os.path.join(BASE_DATA_DIR, 'raw', 'adj_mx_pems_bay.pkl'))\n",
        "\n",
        "# Get the header of the adjacency matrix, the node indices and the\n",
        "# matrix itself.\n",
        "header, node_ids_dict, adj_matrix = adj_matrix_structure\n",
        "\n",
        "# Get the STGNN and load the checkpoints.\n",
        "spatial_temporal_gnn = SpatialTemporalGNN(9, 1, 12, 12, adj_matrix, DEVICE, 64)\n",
        "\n",
        "stgnn_checkpoints_path = os.path.join('..', 'models', 'checkpoints',\n",
        "                                      'st_gnn_pems_bay.pth')\n",
        "\n",
        "stgnn_checkpoints = torch.load(stgnn_checkpoints_path)\n",
        "spatial_temporal_gnn.load_state_dict(stgnn_checkpoints['model_state_dict'])\n",
        "\n",
        "# Set the STGNN in evaluation mode.\n",
        "spatial_temporal_gnn.eval();"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "YIB5Kn57y-wx"
      },
      "outputs": [],
      "source": [
        "from src.data.data_extraction import get_locations_dataframe\n",
        "\n",
        "# Get the dataframe containing the latitude and longitude of each sensor.\n",
        "locations_df = get_locations_dataframe(\n",
        "    os.path.join(BASE_DATA_DIR, 'raw', 'graph_sensor_locations_pems_bay.csv'),\n",
        "    has_header=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "kwVU6H96y-wy"
      },
      "outputs": [],
      "source": [
        "# Get the node positions dictionary.\n",
        "node_pos_dict = { i: id for id, i in node_ids_dict.items() }"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "x-lj4muty-wy"
      },
      "outputs": [],
      "source": [
        "import pickle\n",
        "\n",
        "# Get the data scaler.\n",
        "with open(os.path.join(BASE_DATA_DIR, 'processed', 'scaler.pkl'), 'rb') as f:\n",
        "    scaler = pickle.load(f)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "kNItiWwGy-wy"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import numpy as np\n",
        "\n",
        "# Get the data and the values predicted by the STGNN.\n",
        "x_train = np.load(os.path.join(BASE_DATA_DIR, 'explainable', 'x_train.npy'))\n",
        "y_train = np.load(os.path.join(BASE_DATA_DIR, 'explainable', 'y_train.npy'))\n",
        "x_train_time = np.load(os.path.join(BASE_DATA_DIR, 'explainable', 'x_train_time.npy'))\n",
        "y_train_time = np.load(os.path.join(BASE_DATA_DIR, 'explainable', 'y_train_time.npy'))\n",
        "\n",
        "x_val = np.load(os.path.join(BASE_DATA_DIR, 'explainable', 'x_val.npy'))\n",
        "y_val = np.load(os.path.join(BASE_DATA_DIR, 'explainable', 'y_val.npy'))\n",
        "x_val_time = np.load(os.path.join(BASE_DATA_DIR, 'explainable', 'x_val_time.npy'))\n",
        "y_val_time = np.load(os.path.join(BASE_DATA_DIR, 'explainable', 'y_val_time.npy'))\n",
        "\n",
        "x_test = np.load(os.path.join(BASE_DATA_DIR, 'explainable', 'x_test.npy'))\n",
        "y_test = np.load(os.path.join(BASE_DATA_DIR, 'explainable', 'y_test.npy'))\n",
        "x_test_time = np.load(os.path.join(BASE_DATA_DIR, 'explainable', 'x_test_time.npy'))\n",
        "y_test_time = np.load(os.path.join(BASE_DATA_DIR, 'explainable', 'y_test_time.npy'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3uvxhafzpYNa"
      },
      "source": [
        "# 2 Explainer\n",
        "In this section the explainer algorithm is applied through a grid search on the train dataset, then the results are evaluated.\n",
        "\n",
        "The explainer firstly computes a global heuristic $c_{v_i, v_j}^t$ among all input nodes $v_i$ at timesteps $t$ and the nodes of the event to explain $v_j$ as:\n",
        "$$ c_{v_i, v_j}^t = \\frac{\\Delta t}{\\Delta d} + \\frac{\\exp (- (\\Delta t + \\Delta d))}{\\Delta s} $$\n",
        "Where $\\Delta t$ is the temporal distance between $t$ and the last timestep of $v_j$, $\\Delta d$ their spatial distance and $\\Delta s$ the abdolute difference between the speed of $v_i$ at time $t$ and the average speed of $v_j$ in the event to explain.\n",
        "\n",
        "Afterwards, it performs a MCTS using the input subgraph composed of top scoring nodes by the heuristic as the root of the tree. It performs a series of rollout that involve choosing a path from the root to a leaf node. At each branching, a node of the graph described in the tree node is removed and a child tree node is obtained. The leaf node respects a sparsity threshold, as its graph has a limited pre-defined number of nodes. In practice, given node of the tree $\\mathcal{N}_i$:\n",
        "    \n",
        "1. A new tree node is expanded from $\\mathcal{N}_i$ according to an action. The action consists in removing a specific traffic node in the graph described in $\\mathcal{N}_i$.\n",
        "2. A child node is selected from $\\mathcal{N}_i$ among the expanded ones. In particular, it is selected the child node that maximizes exploration and exploitation.\n",
        "3. If $\\mathcal{N}_i$ is a leaf node, its reward is computed using the STGNN $\\Phi$.\n",
        "4. If $\\mathcal{N}_i$ is a leaf node, its reward is backpropagated from it up to the path nodes until the root is reached to update information useful for exploration and exploitation in the tree search.\n",
        "\n",
        "Ultimately, the final explanation outcome defined as the important subgraph is determined by the subgraph of the leaf node that attains the highest reward and which meets the specified sparsity threshold."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HDbJTJFypYNa"
      },
      "source": [
        "## 2.1 Spatial Distance Matrix\n",
        "Firstly, a speed distance matrix is built in order to compute the shortest distance in kms for each pair of nodes. This matrix is built by firstly finding the shortest path between couple of nodes through the distance matrix $A$ and then computing the length of the path in kms.\n",
        "\n",
        "Given a path $[n_1, ..., n_m]$, the distance among $n_1$ and $n_m$ is obtained as:\n",
        "$$ \\sum_{i}^{m -1} \\text{geodesic}(n_i, n_{i+1})$$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bCEWyqCvy-wz"
      },
      "outputs": [],
      "source": [
        "from src.data.data_processing import get_distance_matrix\n",
        "\n",
        "if not os.path.exists(\n",
        "    os.path.join(BASE_DATA_DIR, 'processed', 'distance_matrix.npy')):\n",
        "    # Build the distance matrix between the nodes.\n",
        "    distance_matrix = get_distance_matrix(\n",
        "        adj_matrix,\n",
        "        locations_df,\n",
        "        node_pos_dict)\n",
        "\n",
        "    # Save the distance matrix.\n",
        "    np.save(os.path.join(BASE_DATA_DIR, 'processed', 'distance_matrix.npy'),\n",
        "            distance_matrix)\n",
        "\n",
        "else:\n",
        "    # Load the distance matrix.\n",
        "    distance_matrix = np.load(\n",
        "        os.path.join(BASE_DATA_DIR, 'processed', 'distance_matrix.npy'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MKsZPATSpYNa"
      },
      "source": [
        "## 2.2 Grid Search\n",
        "\n",
        "Grid search is performed on a portion of the train dataset in order to find the best set of hyperparameters of the Explainer:\n",
        "* `explanation size factor`: multiplied by the number of the nodes of the event to explain, defines the size of the important subgraph found by the explainer;\n",
        "* `cut size factor`: multiplied by `explanation size factor`, defines the number of top scoring nodes of the input graph according to the global heuristic that have to be used to cut the search;\n",
        "* `exploration weight`: parameter to weight the exploration in the MCTS;\n",
        "* `n rollouts`: the number of simulations to perform in the MCTS.\n",
        "* `remove value`: The value used to flag the missing data speed in the important subgraph.\n",
        "\n",
        "The results are evaluated in terms of *Fidelity*$^-$ and *average time* to output the subgraph.\n",
        "\n",
        "* **Fidelity**$^\\textbf{-}$\n",
        "    $$\n",
        "        \\text{Fidelity}^- = \\frac{1}{N}  \\sum_{i=1}^N (\\Phi(\\mathcal{G}_i) - \\Phi({\\mathcal{G}}_i^m))\n",
        "    $$\n",
        "    with $N$ the number of instances, $\\Phi$ the STGNN, $\\mathcal{G}_i$ the $i^th$ input graph and $\\mathcal{G}_i^m$ the important subgraph.\n",
        "    The error $\\Phi(\\mathcal{G}_i) - \\Phi({\\mathcal{G}}_i^m)$ is computed using *MAE*, *RMSE* and *MAPE*.\n",
        "\n",
        "* **Average time:**\n",
        "    It measures the average time employed to output the important subgraph by for each instance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "y70iKYfv2P9r",
        "outputId": "4df2da48-66bd-4046-b003-42073a82b371"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Testing: cut_size_factor: 2 explanation_size_factor: 2 exploration_weight: 5 n_rollouts: 30 remove_value: -60\n",
            "[100/100] - 685s - MAE: { severe_congestion 13.7 -congestion 3.07 -free_flow 2.42 - total: 6.37 } - RMSE: { severe_congestion 14 -congestion 3.7 -free_flow 2.61 - total: 6.74 } - MAPE: { severe_congestion 66.7% -congestion 6.08% -free_flow 3.58% - total: 25.6% } - Average time: 6.85s \n",
            "\n",
            "Testing: cut_size_factor: 2 explanation_size_factor: 2 exploration_weight: 5 n_rollouts: 50 remove_value: -60\n",
            "[100/100] - 1113s - MAE: { severe_congestion 12.9 -congestion 2.9 -free_flow 2.23 - total: 6 } - RMSE: { severe_congestion 13.3 -congestion 3.39 -free_flow 2.41 - total: 6.34 } - MAPE: { severe_congestion 63.5% -congestion 5.75% -free_flow 3.29% - total: 24.3% } - Average time: 11.1s \n",
            "\n",
            "Testing: cut_size_factor: 2 explanation_size_factor: 2 exploration_weight: 10 n_rollouts: 30 remove_value: -60\n",
            "[100/100] - 605s - MAE: { severe_congestion 13.3 -congestion 2.92 -free_flow 2.4 - total: 6.18 } - RMSE: { severe_congestion 13.7 -congestion 3.41 -free_flow 2.62 - total: 6.55 } - MAPE: { severe_congestion 65.1% -congestion 5.78% -free_flow 3.55% - total: 25% } - Average time: 6.05s \n",
            "\n",
            "Testing: cut_size_factor: 2 explanation_size_factor: 2 exploration_weight: 10 n_rollouts: 50 remove_value: -60\n",
            "[100/100] - 981s - MAE: { severe_congestion 13 -congestion 2.82 -free_flow 2.28 - total: 6 } - RMSE: { severe_congestion 13.3 -congestion 3.36 -free_flow 2.46 - total: 6.34 } - MAPE: { severe_congestion 64.6% -congestion 5.54% -free_flow 3.36% - total: 24.6% } - Average time: 9.81s \n",
            "\n",
            "Testing: cut_size_factor: 2 explanation_size_factor: 2 exploration_weight: 20 n_rollouts: 30 remove_value: -60\n",
            "[100/100] - 603s - MAE: { severe_congestion 13.9 -congestion 2.95 -free_flow 2.35 - total: 6.37 } - RMSE: { severe_congestion 14.3 -congestion 3.46 -free_flow 2.56 - total: 6.73 } - MAPE: { severe_congestion 67.4% -congestion 5.83% -free_flow 3.48% - total: 25.7% } - Average time: 6.03s \n",
            "\n",
            "Testing: cut_size_factor: 2 explanation_size_factor: 2 exploration_weight: 20 n_rollouts: 50 remove_value: -60\n",
            "[100/100] - 955s - MAE: { severe_congestion 13.1 -congestion 2.66 -free_flow 2.26 - total: 5.98 } - RMSE: { severe_congestion 13.5 -congestion 3.17 -free_flow 2.44 - total: 6.33 } - MAPE: { severe_congestion 64.8% -congestion 5.27% -free_flow 3.33% - total: 24.7% } - Average time: 9.55s \n",
            "\n",
            "Testing: cut_size_factor: 2 explanation_size_factor: 3 exploration_weight: 5 n_rollouts: 30 remove_value: -60\n",
            "[100/100] - 769s - MAE: { severe_congestion 13.3 -congestion 3.22 -free_flow 2.36 - total: 6.27 } - RMSE: { severe_congestion 13.7 -congestion 3.78 -free_flow 2.53 - total: 6.63 } - MAPE: { severe_congestion 65% -congestion 6.39% -free_flow 3.5% - total: 25.1% } - Average time: 7.69s \n",
            "\n",
            "Testing: cut_size_factor: 2 explanation_size_factor: 3 exploration_weight: 5 n_rollouts: 50 remove_value: -60\n",
            "[100/100] - 1249s - MAE: { severe_congestion 13.1 -congestion 2.45 -free_flow 2.23 - total: 5.91 } - RMSE: { severe_congestion 13.5 -congestion 2.93 -free_flow 2.43 - total: 6.25 } - MAPE: { severe_congestion 64.9% -congestion 4.88% -free_flow 3.3% - total: 24.5% } - Average time: 12.5s \n",
            "\n",
            "Testing: cut_size_factor: 2 explanation_size_factor: 3 exploration_weight: 10 n_rollouts: 30 remove_value: -60\n",
            "[100/100] - 736s - MAE: { severe_congestion 13.7 -congestion 2.98 -free_flow 2.35 - total: 6.32 } - RMSE: { severe_congestion 14.2 -congestion 3.49 -free_flow 2.51 - total: 6.68 } - MAPE: { severe_congestion 67.8% -congestion 5.95% -free_flow 3.47% - total: 25.9% } - Average time: 7.36s \n",
            "\n",
            "Testing: cut_size_factor: 2 explanation_size_factor: 3 exploration_weight: 10 n_rollouts: 50 remove_value: -60\n",
            "[100/100] - 1239s - MAE: { severe_congestion 12.7 -congestion 2.53 -free_flow 2.24 - total: 5.79 } - RMSE: { severe_congestion 13.1 -congestion 3.02 -free_flow 2.43 - total: 6.15 } - MAPE: { severe_congestion 62.9% -congestion 4.99% -free_flow 3.32% - total: 23.9% } - Average time: 12.4s \n",
            "\n",
            "Testing: cut_size_factor: 2 explanation_size_factor: 3 exploration_weight: 20 n_rollouts: 30 remove_value: -60\n",
            "[100/100] - 747s - MAE: { severe_congestion 13.7 -congestion 3.22 -free_flow 2.34 - total: 6.42 } - RMSE: { severe_congestion 14.2 -congestion 3.76 -free_flow 2.53 - total: 6.79 } - MAPE: { severe_congestion 67.4% -congestion 6.36% -free_flow 3.46% - total: 26% } - Average time: 7.47s \n",
            "\n",
            "Testing: cut_size_factor: 2 explanation_size_factor: 3 exploration_weight: 20 n_rollouts: 50 remove_value: -60\n",
            "[100/100] - 1246s - MAE: { severe_congestion 13 -congestion 2.6 -free_flow 2.23 - total: 5.92 } - RMSE: { severe_congestion 13.4 -congestion 3.09 -free_flow 2.39 - total: 6.25 } - MAPE: { severe_congestion 63.6% -congestion 5.15% -free_flow 3.3% - total: 24.2% } - Average time: 12.5s \n",
            "\n"
          ]
        }
      ],
      "source": [
        "from src.explanation.monte_carlo.explanation import apply_grid_search\n",
        "\n",
        "apply_grid_search(\n",
        "    x_train[::10],\n",
        "    y_train[::10],\n",
        "    distance_matrix,\n",
        "    spatial_temporal_gnn,\n",
        "    scaler,\n",
        "    n_rollouts_list=[30, 50],\n",
        "    explanation_size_factor_list=[2, 3],\n",
        "    cut_size_factor_list=[2],\n",
        "    exploration_weight_list=[5, 10, 20],\n",
        "    remove_value_list=[-60])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZXjjftGCpYNb"
      },
      "source": [
        "The best hyperparameters are expressed below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "qVtOLnQ68GDf"
      },
      "outputs": [],
      "source": [
        "CUT_SIZE_FACTOR = 2\n",
        "EXPLANATION_SIZE_FACTOR = 3\n",
        "EXPLORATION_WEIGHT = 10\n",
        "N_ROLLOUTS = 50\n",
        "REMOVE_VALUE = -60"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1sBTwdLGpYNb"
      },
      "source": [
        "## 2.3 Computing the explanations and Saving the Results\n",
        "In this section the explanations of the predicted events of the train, validation and test datasets are computed using the selected hyperparameters.\n",
        "\n",
        "The results are evaluated in terms of *Fidelity*$^-$ using *MAE*, *RMSE* and *MAPE*; *Average time*; *Fidelity*$^+$ using *MAE*, *RMSE* and *MAPE*; and *Sparsity*.\n",
        "\n",
        "*Fidelity*$^-$ and *Average time* have already been introduced. For what concerns *Fidelity*$^+$ and *Sparsity*:\n",
        "\n",
        "* **Fidelity**$^\\textbf{+}$\n",
        "    $$\n",
        "        \\text{Fidelity}^+ = \\frac{1}{N}  \\sum_{i=1}^N (\\Phi(\\mathcal{G}_i) - \\Phi({\\mathcal{G}}_i^{1 - m}))\n",
        "    $$\n",
        "    with $N$ the number of instances, $\\Phi$ the STGNN, $\\mathcal{G}_i$ the $i^th$ input graph and $\\mathcal{G}_i^{1 - m}$ the complement of the important subgraph.\n",
        "* **Sparsity:**\n",
        "    $$\\text{Sparsity} = \\frac{1}{N} \\sum_{i=1}^N (1 -\\frac{m_i}{M_i})$$\n",
        "    with $N$ the number of instances, $m_i$ the number of nodes in the important subgraph used for the explanation at instance $i$ and $M_i$ the number of nodes of the input graph at instance $i$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3FYN5mEu8GAl"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import numpy as np\n",
        "\n",
        "EXPLAINED_DATA_DIR = os.path.join(BASE_DATA_DIR, 'explained')\n",
        "os.makedirs(EXPLAINED_DATA_DIR, exist_ok=True)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kdB__0mW8F-E",
        "outputId": "8b3df37c-7487-414b-9d8c-7a1ccd58ede7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Computing the explanations for the training set...\n",
            "[999/999] - 11818s - MAE: { severe_congestion 13.1 -congestion 2.62 -free_flow 1.95 - total: 5.87 } - RMSE: { severe_congestion 13.5 -congestion 3.09 -free_flow 2.12 - total: 6.21 } - MAPE: { severe_congestion 63.5% -congestion 5.15% -free_flow 2.88% - total: 23.8% } - Average time: 11.8s \n"
          ]
        }
      ],
      "source": [
        "from src.explanation.monte_carlo.explanation import get_all_explanations\n",
        "\n",
        "\n",
        "print('Computing the explanations for the training set...')\n",
        "x_train_explained, y_train_explained, train_scores = get_all_explanations(\n",
        "    x_train,\n",
        "    y_train,\n",
        "    distance_matrix,\n",
        "    spatial_temporal_gnn,\n",
        "    scaler,\n",
        "    n_rollouts=N_ROLLOUTS,\n",
        "    explanation_size_factor=EXPLANATION_SIZE_FACTOR,\n",
        "    cut_size_factor=CUT_SIZE_FACTOR,\n",
        "    exploration_weight=EXPLORATION_WEIGHT,\n",
        "    remove_value=REMOVE_VALUE,\n",
        "    divide_by_traffic_cluster_kind=True)\n",
        "\n",
        "# Save the explained data.\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'x_train.npy'), x_train_explained)\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'y_train.npy'), y_train_explained)\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'train_scores.npy'), train_scores)\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'x_train_time.npy'), x_train_time)\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'y_train_time.npy'), y_train_time)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WdIyprvn8F7d",
        "outputId": "bfd3b344-d90c-4eba-ede9-ceaa64df05b1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Computing the explanations for the validation set...\n",
            "[198/198] - 2467s - MAE: { severe_congestion 13.2 -congestion 2.79 -free_flow 1.81 - total: 5.82 } - RMSE: { severe_congestion 13.7 -congestion 3.3 -free_flow 2.01 - total: 6.2 } - MAPE: { severe_congestion 64.4% -congestion 5.61% -free_flow 2.69% - total: 23.8% } - Average time: 12.5s \n"
          ]
        }
      ],
      "source": [
        "from src.explanation.monte_carlo.explanation import get_all_explanations\n",
        "\n",
        "\n",
        "print('Computing the explanations for the validation set...')\n",
        "x_val_explained, y_val_explained, val_scores = get_all_explanations(\n",
        "    x_val,\n",
        "    y_val,\n",
        "    distance_matrix,\n",
        "    spatial_temporal_gnn,\n",
        "    scaler,\n",
        "    n_rollouts=N_ROLLOUTS,\n",
        "    explanation_size_factor=EXPLANATION_SIZE_FACTOR,\n",
        "    cut_size_factor=CUT_SIZE_FACTOR,\n",
        "    exploration_weight=EXPLORATION_WEIGHT,\n",
        "    remove_value=REMOVE_VALUE,\n",
        "    divide_by_traffic_cluster_kind=True)\n",
        "\n",
        "# Save the explained data.\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'x_val.npy'), x_val_explained)\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'y_val.npy'), y_val_explained)\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'val_scores.npy'), val_scores)\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'x_val_time.npy'), x_val_time)\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'y_val_time.npy'), y_val_time)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x-U5tti38F4V",
        "outputId": "247a5b9a-54ad-46d2-cd70-84a5bca7fe5a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Computing the explanations for the test set...\n",
            "[300/300] - 3489s - MAE: { severe_congestion 12.5 -congestion 2.65 -free_flow 1.97 - total: 5.67 } - RMSE: { severe_congestion 12.9 -congestion 3.13 -free_flow 2.15 - total: 6.02 } - MAPE: { severe_congestion 59.3% -congestion 5.24% -free_flow 2.93% - total: 22.4% } - Average time: 11.6s \n"
          ]
        }
      ],
      "source": [
        "from src.explanation.monte_carlo.explanation import get_all_explanations\n",
        "\n",
        "\n",
        "print('Computing the explanations for the test set...')\n",
        "x_test_explained, y_test_explained, test_scores = get_all_explanations(\n",
        "    x_test,\n",
        "    y_test,\n",
        "    distance_matrix,\n",
        "    spatial_temporal_gnn,\n",
        "    scaler,\n",
        "    n_rollouts=N_ROLLOUTS,\n",
        "    explanation_size_factor=EXPLANATION_SIZE_FACTOR,\n",
        "    cut_size_factor=CUT_SIZE_FACTOR,\n",
        "    exploration_weight=EXPLORATION_WEIGHT,\n",
        "    remove_value=REMOVE_VALUE,\n",
        "    divide_by_traffic_cluster_kind=True)\n",
        "\n",
        "# Save the explained data.\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'x_test.npy'), x_test_explained)\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'y_test.npy'), y_test_explained)\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'test_scores.npy'), test_scores)\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'x_test_time.npy'), x_test_time)\n",
        "np.save(os.path.join(EXPLAINED_DATA_DIR, 'y_test_time.npy'), y_test_time)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "Z5y3W1T2pYNc"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import numpy as np\n",
        "\n",
        "# Get the data and the values predicted by the STGNN.\n",
        "x_train_explained = np.load(os.path.join(BASE_DATA_DIR, 'explained', 'x_train.npy'))\n",
        "\n",
        "x_val_explained = np.load(os.path.join(BASE_DATA_DIR, 'explained', 'x_val.npy'))\n",
        "\n",
        "x_test_explained = np.load(os.path.join(BASE_DATA_DIR, 'explained', 'x_test.npy'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sjaJrmyBpYNc",
        "outputId": "794ad9d0-f067-4d66-e482-eda0dc5ec1ac"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Computing the fidelity+ for the training set...\n",
            "MAE+: { severe_congestion 16.4 -congestion 10.7 -free_flow 49.7 - total: 25.6% } -RMSE+: { severe_congestion 18.5 -congestion 12 -free_flow 53.3 - total: 27.9% } -MAPE+: { severe_congestion 71% -congestion 21.2% -free_flow 75.2% - total: 55.8% }\n",
            "\n",
            "Computing the fidelity+ for the validation set...\n",
            "MAE+: { severe_congestion 17.3 -congestion 11.3 -free_flow 52.4 - total: 27% } -RMSE+: { severe_congestion 19.5 -congestion 12.7 -free_flow 57.5 - total: 29.9% } -MAPE+: { severe_congestion 76.5% -congestion 23.1% -free_flow 79.8% - total: 59.8% }\n",
            "\n",
            "Computing the fidelity+ for the test set...\n",
            "MAE+: { severe_congestion 17.2 -congestion 9.6 -free_flow 41.8 - total: 22.8% } -RMSE+: { severe_congestion 19.1 -congestion 10.7 -free_flow 46 - total: 25.3% } -MAPE+: { severe_congestion 73.1% -congestion 19.2% -free_flow 63.2% - total: 51.8% }\n",
            "\n"
          ]
        }
      ],
      "source": [
        "from src.explanation.monte_carlo.explanation import print_all_fidelity_plus\n",
        "\n",
        "\n",
        "print('Computing the fidelity+ for the training set...')\n",
        "print_all_fidelity_plus(x_train, y_train, x_train_explained, spatial_temporal_gnn, scaler, remove_value=REMOVE_VALUE)\n",
        "print()\n",
        "\n",
        "print('Computing the fidelity+ for the validation set...')\n",
        "print_all_fidelity_plus(x_val, y_val, x_val_explained, spatial_temporal_gnn, scaler, remove_value=REMOVE_VALUE)\n",
        "print()\n",
        "\n",
        "print('Computing the fidelity+ for the test set...')\n",
        "print_all_fidelity_plus(x_test, y_test, x_test_explained, spatial_temporal_gnn, scaler, remove_value=REMOVE_VALUE)\n",
        "print()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Mx8iZu-JpYNc",
        "outputId": "c8b6269d-7b32-48ce-8f8c-9420ab47582c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Computing the sparsity for the training set...\n",
            "Sparsity: { severe_congestion 0.984 -congestion 0.988 -free_flow 0.984 - total: 0.985 }\n",
            "\n",
            "Computing the sparsity for the validation set...\n",
            "Sparsity: { severe_congestion 0.984 -congestion 0.986 -free_flow 0.981 - total: 0.984 }\n",
            "\n",
            "Computing the sparsity for the test set...\n",
            "Sparsity: { severe_congestion 0.985 -congestion 0.988 -free_flow 0.983 - total: 0.985 }\n",
            "\n"
          ]
        }
      ],
      "source": [
        "from src.explanation.monte_carlo.explanation import print_all_sparsity\n",
        "\n",
        "\n",
        "print('Computing the sparsity for the training set...')\n",
        "print_all_sparsity(x_train, x_train_explained)\n",
        "print()\n",
        "\n",
        "print('Computing the sparsity for the validation set...')\n",
        "print_all_sparsity(x_val, x_val_explained)\n",
        "print()\n",
        "\n",
        "print('Computing the sparsity for the test set...')\n",
        "print_all_sparsity(x_test, x_test_explained)\n",
        "print()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}