IDAO.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load_ext blackcellmagic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://idao.world/wp-content/uploads/2018/10/16-%D0%BD%D0%B0-9-1600%D1%85900-768x432.png)\n",
    "\n",
    "This is the jupyter notebook used for model training and inference for team **Fantastic21** for the **International Data Analysis Olympiad 2019** (https://idao.world/) competition.\n",
    "\n",
    "Team: **Fantastic21**  \n",
    "Members:\n",
    "1. **Phung Cheng Shyong** (chengshyongphung@hotmail.com)\n",
    "2. **Kok Gin Xian** (xian_kgx@hotmail.com)\n",
    "3. **Alvin Ting Kee Ngoh**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installation of additional packages for running in Google Cloud Platform (GCP)\n",
    "!pip install imblearn joblib category_encoders lightgbm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "import math\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from imblearn.under_sampling import RandomUnderSampler\n",
    "from imblearn.pipeline import make_pipeline\n",
    "from joblib import dump, load\n",
    "\n",
    "# from sklearn.imputer import SimpleImputer\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.metrics import accuracy_score, roc_auc_score\n",
    "from sklearn.base import BaseEstimator, TransformerMixin\n",
    "import category_encoders\n",
    "\n",
    "import lightgbm as lgb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setting of environment variables for working with GCP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT = \"IDAO\"\n",
    "BUCKET = \"team469\"\n",
    "\n",
    "os.environ[\"PROJECT\"] = PROJECT\n",
    "os.environ[\"BUCKET\"] = BUCKET"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Data loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train_part_1 = pd.read_csv(\"data/train_part_1_v2.csv\")\n",
    "df_train_part_2 = pd.read_csv(\"data/train_part_2_v2.csv\")\n",
    "\n",
    "df_test = pd.read_csv(\"data/test_public_v2.csv\")\n",
    "\n",
    "df_submission = pd.read_csv(\"data/sample_submission.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train_part_1.set_index(\"id\", inplace=True)\n",
    "df_train_part_2.set_index(\"id\", inplace=True)\n",
    "\n",
    "df_test.set_index(\"id\", inplace=True)\n",
    "\n",
    "df_submission.set_index(\"id\", inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Concatenate train part 1 and train part 2 into a single dataset\n",
    "df_train = pd.concat([df_train_part_1, df_train_part_2]).reset_index(drop=True)\n",
    "\n",
    "del df_train_part_1, df_train_part_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Categorical features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "categorical_features = list(df_test.dtypes[df_test.dtypes == \"object\"].index)\n",
    "categorical_features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Excluded features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "excluded_features = [\n",
    "    \"label\",  # This is what we need to predict\n",
    "    \"particle_type\",  # Not in test set\n",
    "    \"weight\",  # Not in test set\n",
    "    \"sWeight\",  # Not in test set\n",
    "    \"kinWeight\",  # Not in test set\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Custom scikit-learn transformers and utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RadiusFromCoordinates(BaseEstimator, TransformerMixin):\n",
    "    \"\"\"\n",
    "    Custom scikit-learn transformer for calculating radius from coordinate axes.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        # The name of the new column/feature\n",
    "        feature_name,\n",
    "        # The coordinate axes for computation of radius\n",
    "        axis_cols=[],\n",
    "    ):\n",
    "        self.feature_name = feature_name\n",
    "        self.axis_cols = axis_cols\n",
    "\n",
    "    def fit(self, X, y=None):\n",
    "        return self\n",
    "\n",
    "    def transform(self, X, y=None):\n",
    "        accu = np.zeros((len(X),))\n",
    "\n",
    "        for axis in self.axis_cols:\n",
    "            accu += X[axis] ** 2\n",
    "\n",
    "        accu = accu ** 0.5\n",
    "        accu = pd.Series(accu, name=self.feature_name)\n",
    "\n",
    "        return pd.concat([X, accu], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AggregatedFeaturesFromArrayFeatures(BaseEstimator, TransformerMixin):\n",
    "    \"\"\"\n",
    "    Custom scikit-learn transformer for computing statistics (min, max, mean, median, size, sum, and variance)\n",
    "    from array features.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, array_cols=[]):\n",
    "        self.array_cols = array_cols\n",
    "        \n",
    "    def fit(self, X, y=None):\n",
    "        return self\n",
    "    \n",
    "    def transform(self, X, y=None):\n",
    "        new_cols = []\n",
    "        \n",
    "        for col in self.array_cols:\n",
    "            splitted = X[col].str.replace(\"[\", \"\").str.replace(\"]\", \"\").str.split()\n",
    "            splitted = list(splitted)\n",
    "            splitted = [pd.to_numeric(l) for l in splitted]\n",
    "            \n",
    "            name = \"{}__min\".format(col)\n",
    "            new_cols.append(pd.Series([np.min(l) for l in splitted], name=name, index=X.index))\n",
    "            name = \"{}__max\".format(col)\n",
    "            new_cols.append(pd.Series([np.max(l) for l in splitted], name=name, index=X.index))\n",
    "            name = \"{}__mean\".format(col)\n",
    "            new_cols.append(pd.Series([np.mean(l) for l in splitted], name=name, index=X.index))\n",
    "            name = \"{}__median\".format(col)\n",
    "            new_cols.append(pd.Series([np.median(l) for l in splitted], name=name, index=X.index))\n",
    "            name = \"{}__size\".format(col)\n",
    "            new_cols.append(pd.Series([np.size(l) for l in splitted], name=name, index=X.index))\n",
    "            name = \"{}__sum\".format(col)\n",
    "            new_cols.append(pd.Series([np.sum(l) for l in splitted], name=name, index=X.index))\n",
    "            name = \"{}__var\".format(col)\n",
    "            new_cols.append(pd.Series([np.var(l) for l in splitted], name=name, index=X.index))\n",
    "#             name = \"{}__mode\".format(col)\n",
    "#             new_cols.append(pd.Series([np.mode(l) for l in splitted], name=name, index=X.index))\n",
    "        \n",
    "        return pd.concat([X] + new_cols, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def enrich_df_with_foi_station_hit_count(df):\n",
    "    \"\"\"\n",
    "    Function for enriching dataframe with various FOI hit count features.\n",
    "    \"\"\"\n",
    "\n",
    "    df[\"FOI_hit_count_S[0]\"] = df[\"FOI_hits_S\"].str.count(\"0\")\n",
    "    df[\"FOI_hit_count_S[1]\"] = df[\"FOI_hits_S\"].str.count(\"1\")\n",
    "    df[\"FOI_hit_count_S[2]\"] = df[\"FOI_hits_S\"].str.count(\"2\")\n",
    "    df[\"FOI_hit_count_S[3]\"] = df[\"FOI_hits_S\"].str.count(\"3\")\n",
    "\n",
    "    df[\"FOI_hit_count_sum\"] = (\n",
    "        df[\"FOI_hit_count_S[0]\"]\n",
    "        + df[\"FOI_hit_count_S[1]\"]\n",
    "        + df[\"FOI_hit_count_S[2]\"]\n",
    "        + df[\"FOI_hit_count_S[3]\"]\n",
    "    )\n",
    "\n",
    "    df[\"FOI_hit_count_ratio_S[0]\"] = df[\"FOI_hit_count_S[0]\"] / df[\"FOI_hit_count_sum\"]\n",
    "    df[\"FOI_hit_count_ratio_S[1]\"] = df[\"FOI_hit_count_S[1]\"] / df[\"FOI_hit_count_sum\"]\n",
    "    df[\"FOI_hit_count_ratio_S[2]\"] = df[\"FOI_hit_count_S[2]\"] / df[\"FOI_hit_count_sum\"]\n",
    "    df[\"FOI_hit_count_ratio_S[3]\"] = df[\"FOI_hit_count_S[3]\"] / df[\"FOI_hit_count_sum\"]\n",
    "\n",
    "    df[\"FOI_hit_count_ratio_S[1][0]\"] = (\n",
    "        df[\"FOI_hit_count_ratio_S[1]\"] / df[\"FOI_hit_count_ratio_S[0]\"]\n",
    "    )\n",
    "    df[\"FOI_hit_count_ratio_S[2][0]\"] = (\n",
    "        df[\"FOI_hit_count_ratio_S[2]\"] / df[\"FOI_hit_count_ratio_S[0]\"]\n",
    "    )\n",
    "    df[\"FOI_hit_count_ratio_S[3][0]\"] = (\n",
    "        df[\"FOI_hit_count_ratio_S[3]\"] / df[\"FOI_hit_count_ratio_S[0]\"]\n",
    "    )\n",
    "    df[\"FOI_hit_count_ratio_S[2][1]\"] = (\n",
    "        df[\"FOI_hit_count_ratio_S[2]\"] / df[\"FOI_hit_count_ratio_S[1]\"]\n",
    "    )\n",
    "    df[\"FOI_hit_count_ratio_S[3][1]\"] = (\n",
    "        df[\"FOI_hit_count_ratio_S[3]\"] / df[\"FOI_hit_count_ratio_S[1]\"]\n",
    "    )\n",
    "    df[\"FOI_hit_count_ratio_S[3][2]\"] = (\n",
    "        df[\"FOI_hit_count_ratio_S[3]\"] / df[\"FOI_hit_count_ratio_S[2]\"]\n",
    "    )\n",
    "    \n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def enrich_df_with_Lextra_Matched_Hit_Vector(df):\n",
    "    \"\"\"\n",
    "    Function for enriching dataframe with vector components of MatchedHit to Lextra vector at various stations.\n",
    "    \"\"\"\n",
    "\n",
    "    df[\"Lextra_Matched_Hit_Vector_X[0]\"] = df[\"Lextra_X[0]\"] - df[\"MatchedHit_X[0]\"]\n",
    "    df[\"Lextra_Matched_Hit_Vector_X[1]\"] = df[\"Lextra_X[1]\"] - df[\"MatchedHit_X[1]\"]\n",
    "    df[\"Lextra_Matched_Hit_Vector_X[2]\"] = df[\"Lextra_X[2]\"] - df[\"MatchedHit_X[2]\"]\n",
    "    df[\"Lextra_Matched_Hit_Vector_X[3]\"] = df[\"Lextra_X[3]\"] - df[\"MatchedHit_X[3]\"]\n",
    "\n",
    "    df[\"Lextra_Matched_Hit_Vector_Y[0]\"] = df[\"Lextra_Y[0]\"] - df[\"MatchedHit_Y[0]\"]\n",
    "    df[\"Lextra_Matched_Hit_Vector_Y[1]\"] = df[\"Lextra_Y[1]\"] - df[\"MatchedHit_Y[1]\"]\n",
    "    df[\"Lextra_Matched_Hit_Vector_Y[2]\"] = df[\"Lextra_Y[2]\"] - df[\"MatchedHit_Y[2]\"]\n",
    "    df[\"Lextra_Matched_Hit_Vector_Y[3]\"] = df[\"Lextra_Y[3]\"] - df[\"MatchedHit_Y[3]\"]\n",
    "    \n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def two_vector_angle_and_magnitudes(\n",
    "    # Vector 1\n",
    "    x_dst_1,\n",
    "    y_dst_1,\n",
    "    z_dst_1,\n",
    "    x_src_1,\n",
    "    y_src_1,\n",
    "    z_src_1,\n",
    "    # Vector 2\n",
    "    x_dst_2,\n",
    "    y_dst_2,\n",
    "    z_dst_2,\n",
    "    x_src_2,\n",
    "    y_src_2,\n",
    "    z_src_2,\n",
    "):\n",
    "    \"\"\"\n",
    "    Function for computing the angle between two vectors and their magnitudes.\n",
    "    \"\"\"\n",
    "\n",
    "    dx_1 = x_dst_1 - x_src_1\n",
    "    dy_1 = y_dst_1 - y_src_1\n",
    "    dz_1 = z_dst_1 - z_src_1\n",
    "\n",
    "    dx_2 = x_dst_2 - x_src_2\n",
    "    dy_2 = y_dst_2 - y_src_2\n",
    "    dz_2 = z_dst_2 - z_src_2\n",
    "\n",
    "    v_1_dot_v_2 = np.array(\n",
    "        [\n",
    "            np.dot([a, b, c], [e, f, g])\n",
    "            for a, b, c, e, f, g in zip(dx_1, dy_1, dz_1, dx_2, dy_2, dz_2)\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    v_1_magnitude = (dx_1 ** 2 + dy_1 ** 2 + dz_1 ** 2) ** 0.5\n",
    "    v_2_magnitude = (dx_2 ** 2 + dy_2 ** 2 + dz_2 ** 2) ** 0.5\n",
    "    angle = np.arccos(v_1_dot_v_2 / (v_1_magnitude * v_2_magnitude))\n",
    "\n",
    "    return angle, v_1_magnitude, v_2_magnitude"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeaturesDropper(BaseEstimator, TransformerMixin):\n",
    "    \"\"\"\n",
    "    Custom scikit-learn transformer for dropping of features/columns from a dataframe.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, columns=[]):\n",
    "        self.columns = columns\n",
    "\n",
    "    def fit(self, X, y=None):\n",
    "        return self\n",
    "\n",
    "    def transform(self, X, y=None):\n",
    "        return X.drop(self.columns, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LeftJoin(BaseEstimator, TransformerMixin):\n",
    "    \"\"\"\n",
    "    Custom scikit-learn transformer for performing a left join between two dataframes.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, csv_path, on=None, left_on=None, right_on=None):\n",
    "        self.csv_path = csv_path\n",
    "\n",
    "        if on and (left_on or right_on):\n",
    "            raise Exception(\"Both on and left/right on set!\")\n",
    "        if (left_on and not right_on) or (not left_on and right_on):\n",
    "            raise Exception(\"Only left or right on set!\")\n",
    "        if not on and not left_on and not right_on:\n",
    "            raise Exception(\"Neither on nor left_on and right_on set!\")\n",
    "\n",
    "        self.on = on\n",
    "        self.left_on = left_on\n",
    "        self.right_on = right_on\n",
    "\n",
    "    def fit(self, X, y=None):\n",
    "        return self\n",
    "\n",
    "    def transform(self, X, y=None):\n",
    "        X_right = pd.read_csv(self.csv_path)\n",
    "\n",
    "        if self.on:\n",
    "            return X.merge(X_right, left_on=self.left_on, right_on=self.right_on)\n",
    "        elif self.left_on and self.right_on:\n",
    "            return X.merge(X_right, on=self.on)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ColumnRenamer(BaseEstimator, TransformerMixin):\n",
    "    \"\"\"\n",
    "    Custom scikit-learn transformer for renaming a column in a dataframe.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, col, col_new_name):\n",
    "        self.col = col\n",
    "        self.col_new_name = col_new_name\n",
    "    \n",
    "    def fit(self, X, y=None):\n",
    "        return self\n",
    "    \n",
    "    def transform(self, X, y=None):\n",
    "        return X.rename(columns={\n",
    "            self.col: self.col_new_name\n",
    "        })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DivideTwoFeatures(BaseEstimator, TransformerMixin):\n",
    "    \"\"\"\n",
    "    Custom scikit-learn transformer for creating a new column/feature obtained by dividing two columns/features in a dataframe.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, col1, col2, new_col_name):\n",
    "        self.col1 = col1\n",
    "        self.col2 = col2\n",
    "        self.new_col_name = new_col_name\n",
    "\n",
    "    def fit(self, X, y=None):\n",
    "        return self\n",
    "\n",
    "    def transform(self, X, y=None):\n",
    "        tmp = pd.Series(X[self.col1] / X[self.col2], name=self.new_col_name)\n",
    "        return pd.concat([X, tmp], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiplyFeatures(BaseEstimator, TransformerMixin):\n",
    "    \"\"\"\n",
    "    Custom scikit-learn transformer for creating a new column/feature that is the multiplication of a list of\n",
    "    columns/features in a dataframe.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, new_col_name, columns):\n",
    "        self.new_col_name = new_col_name\n",
    "        self.columns = columns\n",
    "\n",
    "    def fit(self, X, y=None):\n",
    "        return self\n",
    "\n",
    "    def transform(self, X, y=None):\n",
    "        accu = np.ones((len(X),))\n",
    "\n",
    "        for col in self.columns:\n",
    "            accu *= X[col]\n",
    "\n",
    "        accu = pd.Series(accu, name=self.new_col_name)\n",
    "\n",
    "        return pd.concat([X, accu], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiColumnSum(BaseEstimator, TransformerMixin):\n",
    "    \"\"\"\n",
    "    Custom scikit-learn transformer for creating a new column/feature that is the sum of a alist of columns/features\n",
    "    in a dataframe.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, new_col_name, columns):\n",
    "        self.columns = columns\n",
    "        self.new_col_name = new_col_name\n",
    "\n",
    "    def fit(self, X, y=None):\n",
    "        return self\n",
    "\n",
    "    def transform(self, X, y=None):\n",
    "        accu = np.zeros((len(X),))\n",
    "\n",
    "        for col in self.columns:\n",
    "            accu += X[col]\n",
    "\n",
    "        accu = pd.Series(accu, name=self.new_col_name)\n",
    "        return pd.concat([X, accu], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Angle and magnitude of coordinates\n",
    "def enrich_df_with_matchhit_angle_and_radius_at_station_3D(df):\n",
    "    for i in range(4):\n",
    "        df[\"MatchedHit_vector_angle_Z[{}]\".format(i)], df[\n",
    "            \"MatchHit_R_from_origin_S[{}]\".format(i)\n",
    "        ], _ = two_vector_angle_and_magnitudes(\n",
    "            df[\"MatchedHit_X[{}]\".format(i)].values,\n",
    "            df[\"MatchedHit_Y[{}]\".format(i)].values,\n",
    "            df[\"MatchedHit_Z[{}]\".format(i)].values,\n",
    "            np.zeros_like(df[\"MatchedHit_X[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Y[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Z[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_X[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Y[{}]\".format(i)]),\n",
    "            df[\"MatchedHit_Z[{}]\".format(i)].values,\n",
    "            np.zeros_like(df[\"MatchedHit_X[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Y[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Z[{}]\".format(i)]),\n",
    "        )\n",
    "\n",
    "        df[\n",
    "            \"MatchedHit_vector_angle_Y[{}]\".format(i)\n",
    "        ], _, _ = two_vector_angle_and_magnitudes(\n",
    "            df[\"MatchedHit_X[{}]\".format(i)].values,\n",
    "            df[\"MatchedHit_Y[{}]\".format(i)].values,\n",
    "            df[\"MatchedHit_Z[{}]\".format(i)].values,\n",
    "            np.zeros_like(df[\"MatchedHit_X[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Y[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Z[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_X[{}]\".format(i)]),\n",
    "            df[\"MatchedHit_Y[{}]\".format(i)].values,\n",
    "            np.zeros_like(df[\"MatchedHit_Z[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_X[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Y[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Z[{}]\".format(i)]),\n",
    "        )\n",
    "\n",
    "        df[\n",
    "            \"MatchedHit_vector_angle_X[{}]\".format(i)\n",
    "        ], _, _ = two_vector_angle_and_magnitudes(\n",
    "            df[\"MatchedHit_X[{}]\".format(i)].values,\n",
    "            df[\"MatchedHit_Y[{}]\".format(i)].values,\n",
    "            df[\"MatchedHit_Z[{}]\".format(i)].values,\n",
    "            np.zeros_like(df[\"MatchedHit_X[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Y[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Z[{}]\".format(i)]),\n",
    "            df[\"MatchedHit_X[{}]\".format(i)].values,\n",
    "            np.zeros_like(df[\"MatchedHit_Y[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Z[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_X[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Y[{}]\".format(i)]),\n",
    "            np.zeros_like(df[\"MatchedHit_Z[{}]\".format(i)]),\n",
    "        )\n",
    "\n",
    "    df[\"MatchedHit_vector_angle_Z_sum\"] = (\n",
    "        df[\"MatchedHit_vector_angle_Z[0]\"]\n",
    "        + df[\"MatchedHit_vector_angle_Z[1]\"]\n",
    "        + df[\"MatchedHit_vector_angle_Z[2]\"]\n",
    "        + df[\"MatchedHit_vector_angle_Z[3]\"]\n",
    "    )\n",
    "    df[\"MatchedHit_vector_angle_Y_sum\"] = (\n",
    "        df[\"MatchedHit_vector_angle_Y[0]\"]\n",
    "        + df[\"MatchedHit_vector_angle_Y[1]\"]\n",
    "        + df[\"MatchedHit_vector_angle_Y[2]\"]\n",
    "        + df[\"MatchedHit_vector_angle_Y[3]\"]\n",
    "    )\n",
    "    df[\"MatchedHit_vector_angle_X_sum\"] = (\n",
    "        df[\"MatchedHit_vector_angle_X[0]\"]\n",
    "        + df[\"MatchedHit_vector_angle_X[1]\"]\n",
    "        + df[\"MatchedHit_vector_angle_X[2]\"]\n",
    "        + df[\"MatchedHit_vector_angle_X[3]\"]\n",
    "    )\n",
    "    df[\"MatchHit_R_from_origin_S_sum\"] = (\n",
    "        df[\"MatchHit_R_from_origin_S[0]\"]\n",
    "        + df[\"MatchHit_R_from_origin_S[1]\"]\n",
    "        + df[\"MatchHit_R_from_origin_S[2]\"]\n",
    "        + df[\"MatchHit_R_from_origin_S[3]\"]\n",
    "    )\n",
    "    for i in range(4):\n",
    "        df[\"MatchedHit_vector_angle_Z_ratio[{}]\".format(i)] = (\n",
    "            df[\"MatchedHit_vector_angle_Z[{}]\".format(i)]\n",
    "            / df[\"MatchedHit_vector_angle_Z_sum\"]\n",
    "        )\n",
    "        df[\"MatchedHit_vector_angle_Y_ratio[{}]\".format(i)] = (\n",
    "            df[\"MatchedHit_vector_angle_Y[{}]\".format(i)]\n",
    "            / df[\"MatchedHit_vector_angle_Y_sum\"]\n",
    "        )\n",
    "        df[\"MatchedHit_vector_angle_X_ratio[{}]\".format(i)] = (\n",
    "            df[\"MatchedHit_vector_angle_X[{}]\".format(i)]\n",
    "            / df[\"MatchedHit_vector_angle_X_sum\"]\n",
    "        )\n",
    "\n",
    "    for s1, s2 in [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]:\n",
    "        df[\"MatchedHit_vector_angle_Z_diff_[{},{}]\".format(s2, s1)] = (\n",
    "            df[\"MatchedHit_vector_angle_Z[{}]\".format(s2)]\n",
    "            - df[\"MatchedHit_vector_angle_Z[{}]\".format(s1)]\n",
    "        )\n",
    "        df[\"MatchedHit_vector_angle_Z_ratio_[{},{}]\".format(s2, s1)] = (\n",
    "            df[\"MatchedHit_vector_angle_Z[{}]\".format(s2)]\n",
    "            / df[\"MatchedHit_vector_angle_Z[{}]\".format(s1)]\n",
    "        )\n",
    "        df[\"MatchedHit_vector_angle_Z[{}>{}]\".format(s2, s1)] = (\n",
    "            df[\"MatchedHit_vector_angle_Z[{}]\".format(s2)]\n",
    "            > df[\"MatchedHit_vector_angle_Z[{}]\".format(s1)]\n",
    "        ).astype(\"int\")\n",
    "        df[\"MatchedHit_vector_angle_Y_diff_[{},{}]\".format(s2, s1)] = (\n",
    "            df[\"MatchedHit_vector_angle_Y[{}]\".format(s2)]\n",
    "            - df[\"MatchedHit_vector_angle_Y[{}]\".format(s1)]\n",
    "        )\n",
    "        df[\"MatchedHit_vector_angle_Y_ratio_[{},{}]\".format(s2, s1)] = (\n",
    "            df[\"MatchedHit_vector_angle_Y[{}]\".format(s2)]\n",
    "            / df[\"MatchedHit_vector_angle_Y[{}]\".format(s1)]\n",
    "        )\n",
    "        df[\"MatchedHit_vector_angle_Y[{}>{}]\".format(s2, s1)] = (\n",
    "            df[\"MatchedHit_vector_angle_Y[{}]\".format(s2)]\n",
    "            > df[\"MatchedHit_vector_angle_Y[{}]\".format(s1)]\n",
    "        ).astype(\"int\")\n",
    "        df[\"MatchedHit_vector_angle_X_diff_[{},{}]\".format(s2, s1)] = (\n",
    "            df[\"MatchedHit_vector_angle_X[{}]\".format(s2)]\n",
    "            - df[\"MatchedHit_vector_angle_X[{}]\".format(s1)]\n",
    "        )\n",
    "        df[\"MatchedHit_vector_angle_X_ratio_[{},{}]\".format(s2, s1)] = (\n",
    "            df[\"MatchedHit_vector_angle_X[{}]\".format(s2)]\n",
    "            / df[\"MatchedHit_vector_angle_X[{}]\".format(s1)]\n",
    "        )\n",
    "        df[\"MatchedHit_vector_angle_X[{}>{}]\".format(s2, s1)] = (\n",
    "            df[\"MatchedHit_vector_angle_X[{}]\".format(s2)]\n",
    "            > df[\"MatchedHit_vector_angle_X[{}]\".format(s1)]\n",
    "        ).astype(\"int\")\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Angle and magnitude of coordinates between two stations\n",
    "def enrich_df_with_matchhit_angle_and_radius_between_two_stations(\n",
    "    df, src_station, dst_station\n",
    "):\n",
    "    df[\n",
    "        \"MatchHit_angle_between_stations_Z_[{},{}]\".format(src_station, dst_station)\n",
    "    ], df[\n",
    "        \"MatchHit_R_between_stations_[{},{}]\".format(src_station, dst_station)\n",
    "    ], _ = two_vector_angle_and_magnitudes(\n",
    "        df[\"MatchedHit_X[{}]\".format(dst_station)].values,\n",
    "        df[\"MatchedHit_Y[{}]\".format(dst_station)].values,\n",
    "        df[\"MatchedHit_Z[{}]\".format(dst_station)].values,\n",
    "        df[\"MatchedHit_X[{}]\".format(src_station)].values,\n",
    "        df[\"MatchedHit_Y[{}]\".format(src_station)].values,\n",
    "        df[\"MatchedHit_Z[{}]\".format(src_station)].values,\n",
    "        np.zeros_like(df[\"MatchedHit_X[{}]\".format(dst_station)]),\n",
    "        np.zeros_like(df[\"MatchedHit_Y[{}]\".format(dst_station)]),\n",
    "        df[\"MatchedHit_Z[{}]\".format(dst_station)].values,\n",
    "        np.zeros_like(df[\"MatchedHit_X[{}]\".format(src_station)]),\n",
    "        np.zeros_like(df[\"MatchedHit_Y[{}]\".format(src_station)]),\n",
    "        df[\"MatchedHit_Z[{}]\".format(src_station)].values,\n",
    "    )\n",
    "\n",
    "    df[\n",
    "        \"MatchHit_angle_between_stations_Y_[{},{}]\".format(src_station, dst_station)\n",
    "    ], _, _ = two_vector_angle_and_magnitudes(\n",
    "        df[\"MatchedHit_X[{}]\".format(dst_station)].values,\n",
    "        df[\"MatchedHit_Y[{}]\".format(dst_station)].values,\n",
    "        df[\"MatchedHit_Z[{}]\".format(dst_station)].values,\n",
    "        df[\"MatchedHit_X[{}]\".format(src_station)].values,\n",
    "        df[\"MatchedHit_Y[{}]\".format(src_station)].values,\n",
    "        df[\"MatchedHit_Z[{}]\".format(src_station)].values,\n",
    "        np.zeros_like(df[\"MatchedHit_X[{}]\".format(dst_station)]),\n",
    "        df[\"MatchedHit_Y[{}]\".format(dst_station)].values,\n",
    "        np.zeros_like(df[\"MatchedHit_Z[{}]\".format(dst_station)]),\n",
    "        np.zeros_like(df[\"MatchedHit_X[{}]\".format(src_station)]),\n",
    "        df[\"MatchedHit_Y[{}]\".format(src_station)].values,\n",
    "        np.zeros_like(df[\"MatchedHit_Z[{}]\".format(dst_station)]),\n",
    "    )\n",
    "\n",
    "    df[\n",
    "        \"MatchHit_angle_between_stations_X_[{},{}]\".format(src_station, dst_station)\n",
    "    ], _, _ = two_vector_angle_and_magnitudes(\n",
    "        df[\"MatchedHit_X[{}]\".format(dst_station)].values,\n",
    "        df[\"MatchedHit_Y[{}]\".format(dst_station)].values,\n",
    "        df[\"MatchedHit_Z[{}]\".format(dst_station)].values,\n",
    "        df[\"MatchedHit_X[{}]\".format(src_station)].values,\n",
    "        df[\"MatchedHit_Y[{}]\".format(src_station)].values,\n",
    "        df[\"MatchedHit_Z[{}]\".format(src_station)].values,\n",
    "        df[\"MatchedHit_X[{}]\".format(dst_station)].values,\n",
    "        np.zeros_like(df[\"MatchedHit_Y[{}]\".format(dst_station)]),\n",
    "        np.zeros_like(df[\"MatchedHit_Z[{}]\".format(dst_station)]),\n",
    "        df[\"MatchedHit_X[{}]\".format(src_station)].values,\n",
    "        np.zeros_like(df[\"MatchedHit_Y[{}]\".format(dst_station)]),\n",
    "        np.zeros_like(df[\"MatchedHit_Z[{}]\".format(dst_station)]),\n",
    "    )\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Polar coordinates conversion\n",
    "def enrich_df_with_matchedhit_radial_features(df):\n",
    "    for a1, a2 in [(\"Y\", \"X\"), (\"Y\", \"Z\"), (\"X\", \"Z\")]:\n",
    "        for i in range(4):\n",
    "            df[\"MatchedHit_angle_{}{}[{}]\".format(a1, a2, i)] = np.arctan2(\n",
    "                df[\"MatchedHit_{}[{}]\".format(a1, i)],\n",
    "                df[\"MatchedHit_{}[{}]\".format(a2, i)],\n",
    "            )\n",
    "            # df[\"MatchedHit_modulus_{}{}[{}]\".format(a1, a2, i)] = np.linalg.norm([df[\"MatchedHit_{}[{}]\".format(a1, i)], df[\"MatchedHit_{}[{}]\".format(a2, i)]])\n",
    "            # df[\"MatchedHit_sign_{}{}[{}]\".format(a1, a2, i)] =np.where((np.arctan2(df[\"MatchedHit_{}[{}]\".format(a1, i)], df[\"MatchedHit_{}[{}]\".format(a2, i)])>=0), 1, -1)\n",
    "            try:\n",
    "                df[\"Lextra_angle_{}{}[{}]\".format(a1, a2, i)] = np.arctan2(\n",
    "                    df[\"Lextra_{}[{}]\".format(a1, i)], df[\"Lextra_{}[{}]\".format(a2, i)]\n",
    "                )\n",
    "                # df[\"Lextra_modulus_{}{}[{}]\".format(a1, a2, i)] = np.linalg.norm([df[\"Lextra_{}[{}]\".format(a1, i)], df[\"Lextra_{}[{}]\".format(a2, i)]])\n",
    "                # df[\"Lextra_sign_{}{}[{}]\".format(a1, a2, i)] =np.where((np.arctan2(df[\"Lextra_{}[{}]\".format(a1, i)], df[\"Lextra_{}[{}]\".format(a2, i)])>=0), 1, -1)\n",
    "            except:\n",
    "                pass\n",
    "\n",
    "        df[\"MatchedHit_angle_sum_{}{}\".format(a1, a2)] = (\n",
    "            df[\"MatchedHit_angle_{}{}[0]\".format(a1, a2)]\n",
    "            + df[\"MatchedHit_angle_{}{}[1]\".format(a1, a2)]\n",
    "            + df[\"MatchedHit_angle_{}{}[2]\".format(a1, a2)]\n",
    "            + df[\"MatchedHit_angle_{}{}[3]\".format(a1, a2)]\n",
    "        )\n",
    "        # df[\"MatchedHit_modulus_sum_{}{}\".format(a1, a2)] = df[\"MatchedHit_modulus_{}{}[0]\".format(a1, a2)] + df[\"MatchedHit_modulus_{}{}[1]\".format(a1, a2)] + df[\"MatchedHit_modulus_{}{}[2]\".format(a1, a2)] + df[\"MatchedHit_modulus_{}{}[3]\".format(a1, a2)]\n",
    "        try:\n",
    "            df[\"Lextra_angle_sum_{}{}\".format(a1, a2)] = (\n",
    "                df[\"Lextra_angle_{}{}[0]\".format(a1, a2)]\n",
    "                + df[\"Lextra_angle_{}{}[1]\".format(a1, a2)]\n",
    "                + df[\"Lextra_angle_{}{}[2]\".format(a1, a2)]\n",
    "                + df[\"Lextra_angle_{}{}[3]\".format(a1, a2)]\n",
    "            )\n",
    "            # df[\"Lextra_modulus_sum_{}{}\".format(a1, a2)] = df[\"Lextra_modulus_{}{}[0]\".format(a1, a2)] + df[\"Lextra_modulus_{}{}[1]\".format(a1, a2)] + df[\"Lextra_modulus_{}{}[2]\".format(a1, a2)] + df[\"Lextra_modulus_{}{}[3]\".format(a1, a2)]\n",
    "        except:\n",
    "            pass\n",
    "\n",
    "        for i in range(4):\n",
    "            df[\"MatchedHit_angle_ratio_{}{}[{}]\".format(a1, a2, i)] = (\n",
    "                df[\"MatchedHit_angle_{}{}[{}]\".format(a1, a2, i)]\n",
    "                / df[\"MatchedHit_angle_sum_{}{}\".format(a1, a2)]\n",
    "            )\n",
    "            # df[\"MatchedHit_modulus_ratio_{}{}[{}]\".format(a1, a2, i)] = df[\"MatchedHit_modulus_{}{}[{}]\".format(a1, a2, i)] / df[\"MatchedHit_modulus_sum_{}{}\".format(a1, a2)]\n",
    "            try:\n",
    "                df[\"Lextra_angle_ratio_{}{}[{}]\".format(a1, a2, i)] = (\n",
    "                    df[\"Lextra_angle_{}{}[{}]\".format(a1, a2, i)]\n",
    "                    / df[\"Lextra_angle_sum_{}{}\".format(a1, a2)]\n",
    "                )\n",
    "                # df[\"Lextra_angle_modulus_{}{}[{}]\".format(a1, a2, i)] = df[\"Lextra_modulus_{}{}[{}]\".format(a1, a2, i)] / df[\"Lextra_modulus_sum_{}{}\".format(a1, a2)]\n",
    "            except:\n",
    "                pass\n",
    "\n",
    "        for s1, s2 in [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]:\n",
    "            df[\"MatchedHit_angle_diff_{}{}[{},{}]\".format(a1, a2, s2, s1)] = (\n",
    "                df[\"MatchedHit_angle_{}{}[{}]\".format(a1, a2, s2)]\n",
    "                - df[\"MatchedHit_angle_{}{}[{}]\".format(a1, a2, s1)]\n",
    "            )\n",
    "            df[\"MatchedHit_angle_ratio_{}{}[{},{}]\".format(a1, a2, s2, s1)] = (\n",
    "                df[\"MatchedHit_angle_{}{}[{}]\".format(a1, a2, s2)]\n",
    "                / df[\"MatchedHit_angle_{}{}[{}]\".format(a1, a2, s1)]\n",
    "            )\n",
    "            df[\"MatchedHit_angle_{}{}[{}>{}]\".format(a1, a2, s2, s1)] = (\n",
    "                df[\"MatchedHit_angle_{}{}[{}]\".format(a1, a2, s2)]\n",
    "                > df[\"MatchedHit_angle_{}{}[{}]\".format(a1, a2, s1)]\n",
    "            ).astype(\"int\")\n",
    "            # df[\"MatchedHit_modulus_diff_{}{}[{},{}]\".format(a1, a2, s2, s1)] = df[\"MatchedHit_modulus_{}{}[{}]\".format(a1, a2, s2)] -df[\"MatchedHit_modulus_{}{}[{}]\".format(a1, a2, s1)]\n",
    "            # df[\"MatchedHit_modulus_ratio_{}{}[{},{}]\".format(a1, a2, s2, s1)] = df[\"MatchedHit_modulus_{}{}[{}]\".format(a1, a2, s2)] /df[\"MatchedHit_modulus_{}{}[{}]\".format(a1, a2, s1)]\n",
    "            # df[\"MatchedHit_modulus_{}{}[{}>{}]\".format(a1, a2, s2, s1)] = (df[\"MatchedHit_modulus_{}{}[{}]\".format(a1, a2, s2)] > df[\"MatchedHit_modulus_{}{}[{}]\".format(a1, a2, s1)]).astype(\"int\")\n",
    "            # df[\"MatchedHit_sign_ratio_{}{}[{},{}]\".format(a1, a2, s2, s1)] = df[\"MatchedHit_sign_{}{}[{}]\".format(a1, a2, s2)] /df[\"MatchedHit_sign_{}{}[{}]\".format(a1, a2, s1)]\n",
    "\n",
    "            try:\n",
    "                df[\"Lextra_angle_diff_{}{}[{},{}]\".format(a1, a2, s2, s1)] = (\n",
    "                    df[\"Lextra_angle_{}{}[{}]\".format(a1, a2, s2)]\n",
    "                    - df[\"Lextra_angle_{}{}[{}]\".format(a1, a2, s1)]\n",
    "                )\n",
    "                df[\"Lextra_angle_ratio_{}{}[{},{}]\".format(a1, a2, s2, s1)] = (\n",
    "                    df[\"Lextra_angle_{}{}[{}]\".format(a1, a2, s2)]\n",
    "                    / df[\"Lextra_angle_{}{}[{}]\".format(a1, a2, s1)]\n",
    "                )\n",
    "                df[\"Lextra_angle_{}{}[{}>{}]\".format(a1, a2, s2, s1)] = (\n",
    "                    df[\"Lextra_angle_{}{}[{}]\".format(a1, a2, s2)]\n",
    "                    > df[\"Lextra_angle_{}{}[{}]\".format(a1, a2, s1)]\n",
    "                ).astype(\"int\")\n",
    "                # df[\"Lextra_modulus_diff_{}{}[{},{}]\".format(a1, a2, s2, s1)] = df[\"Lextra_modulus_{}{}[{}]\".format(a1, a2, s2)] - df[\"Lextra_modulus_{}{}[{}]\".format(a1, a2, s1)]\n",
    "                # df[\"Lextra_modulus_ratio_{}{}[{},{}]\".format(a1, a2, s2, s1)] = df[\"Lextra_modulus_{}{}[{}]\".format(a1, a2, s2)]/df[\"Lextra_modulus_{}{}[{}]\".format(a1, a2, s1)]\n",
    "                # df[\"Lextra_modulus_{}{}[{}>{}]\".format(a1, a2, s2, s1)] = (df[\"Lextra_modulus_{}{}[{}]\".format(a1, a2, s2)] > df[\"Lextra_modulus_{}{}[{}]\".format(a1, a2, s1)]).astype(\"int\")\n",
    "                # df[\"Lextra_sign_ratio_{}{}[{},{}]\".format(a1, a2, s2, s1)] = df[\"Lextra_sign_{}{}[{}]\".format(a1, a2, s2)]/df[\"Lextra_sign_{}{}[{}]\".format(a1, a2, s1)]\n",
    "\n",
    "            except:\n",
    "                pass\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Polar coordinates conversion of the lectra Matched Hit Vector\n",
    "def enrich_df_with_Lextra_matchedhit_radial_features(df):\n",
    "    for a1, a2 in [(\"Y\", \"X\")]:\n",
    "        for i in range(4):\n",
    "            df[\n",
    "                \"Lextra_Matched_Hit_Vector_angle_{}{}[{}]\".format(a1, a2, i)\n",
    "            ] = np.arctan2(\n",
    "                df[\"Lextra_Matched_Hit_Vector_{}[{}]\".format(a1, i)],\n",
    "                df[\"Lextra_Matched_Hit_Vector_{}[{}]\".format(a2, i)],\n",
    "            )\n",
    "\n",
    "        df[\"Lextra_Matched_Hit_Vector_angle_sum_{}{}\".format(a1, a2)] = (\n",
    "            df[\"Lextra_Matched_Hit_Vector_angle_{}{}[0]\".format(a1, a2)]\n",
    "            + df[\"Lextra_Matched_Hit_Vector_angle_{}{}[1]\".format(a1, a2)]\n",
    "            + df[\"MatchedHit_angle_{}{}[2]\".format(a1, a2)]\n",
    "            + df[\"MatchedHit_angle_{}{}[3]\".format(a1, a2)]\n",
    "        )\n",
    "\n",
    "        for i in range(4):\n",
    "            df[\"Lextra_Matched_Hit_Vector_ratio_{}{}[{}]\".format(a1, a2, i)] = (\n",
    "                df[\"Lextra_Matched_Hit_Vector_angle_{}{}[{}]\".format(a1, a2, i)]\n",
    "                / df[\"Lextra_Matched_Hit_Vector_angle_sum_{}{}\".format(a1, a2)]\n",
    "            )\n",
    "\n",
    "        for s1, s2 in [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]:\n",
    "            df[\n",
    "                \"Lextra_Matched_Hit_Vector_angle_diff_{}{}[{},{}]\".format(\n",
    "                    a1, a2, s2, s1\n",
    "                )\n",
    "            ] = (\n",
    "                df[\"Lextra_Matched_Hit_Vector_angle_{}{}[{}]\".format(a1, a2, s2)]\n",
    "                - df[\"Lextra_Matched_Hit_Vector_angle_{}{}[{}]\".format(a1, a2, s1)]\n",
    "            )\n",
    "            df[\n",
    "                \"Lextra_Matched_Hit_Vector_angle_ratio_{}{}[{},{}]\".format(\n",
    "                    a1, a2, s2, s1\n",
    "                )\n",
    "            ] = (\n",
    "                df[\"Lextra_Matched_Hit_Vector_angle_{}{}[{}]\".format(a1, a2, s2)]\n",
    "                / df[\"Lextra_Matched_Hit_Vector_angle_{}{}[{}]\".format(a1, a2, s1)]\n",
    "            )\n",
    "            df[\"Lextra_Matched_Hit_Vector_angle_{}{}[{}>{}]\".format(a1, a2, s2, s1)] = (\n",
    "                df[\"Lextra_Matched_Hit_Vector_angle_{}{}[{}]\".format(a1, a2, s2)]\n",
    "                > df[\"Lextra_Matched_Hit_Vector_angle_{}{}[{}]\".format(a1, a2, s1)]\n",
    "            ).astype(\"int\")\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Deriving information from time feature\n",
    "def enrich_df_with_matchedhit_time_features(df):\n",
    "    df[\"MatchedHit_T[sum]\"] = (\n",
    "        df[\"MatchedHit_T[0]\"]\n",
    "        + df[\"MatchedHit_T[1]\"]\n",
    "        + df[\"MatchedHit_T[2]\"]\n",
    "        + df[\"MatchedHit_T[3]\"]\n",
    "    )\n",
    "    for i in range(4):\n",
    "        df[\"MatchedHit_T_ratio[{}]\".format(i)] = (\n",
    "            df[\"MatchedHit_T[{}]\".format(i)] / df[\"MatchedHit_T[sum]\"]\n",
    "        )\n",
    "\n",
    "    for s1, s2 in [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]:\n",
    "        df[\"MatchedHit_T_diff[{},{}]\".format(s2, s1)] = (\n",
    "            df[\"MatchedHit_T[{}]\".format(s2)] - df[\"MatchedHit_T[{}]\".format(s1)]\n",
    "        )\n",
    "        df[\"MatchedHit_T_ratio[{},{}]\".format(s2, s1)] = (\n",
    "            df[\"MatchedHit_T[{}]\".format(s1)] / df[\"MatchedHit_T[{}]\".format(s2)]\n",
    "        )\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_df(df):\n",
    "    \"\"\"\n",
    "    Function for preparing dataframe prior to model training/inference.\n",
    "    \"\"\"\n",
    "    \n",
    "    df = AggregatedFeaturesFromArrayFeatures(categorical_features).fit_transform(df)\n",
    "    \n",
    "    enrich_df_with_foi_station_hit_count(df)\n",
    "    enrich_df_with_matchedhit_radial_features(df)\n",
    "    enrich_df_with_matchedhit_time_features(df)\n",
    "    enrich_df_with_matchhit_angle_and_radius_at_station_3D(df)\n",
    "    enrich_df_with_Lextra_Matched_Hit_Vector(df)\n",
    "    enrich_df_with_Lextra_matchedhit_radial_features(df)\n",
    "    \n",
    "    for s1, s2 in [\n",
    "        (0, 1),\n",
    "        (0, 2),\n",
    "        (0, 3),\n",
    "        (1, 2),\n",
    "        (1, 3),\n",
    "        (2, 3)\n",
    "    ]:\n",
    "        enrich_df_with_matchhit_angle_and_radius_between_two_stations(df, s1, s2)\n",
    "        \n",
    "    pipeline = make_pipeline(\n",
    "        MultiColumnSum(\"ncl[sum]\", [\"ncl[0]\", \"ncl[1]\", \"ncl[2]\", \"ncl[3]\"]),\n",
    "    \n",
    "        DivideTwoFeatures(\"ncl[0]\", \"ncl[sum]\", \"ncl_ratio[0]\"),\n",
    "        DivideTwoFeatures(\"ncl[1]\", \"ncl[sum]\", \"ncl_ratio[1]\"),\n",
    "        DivideTwoFeatures(\"ncl[2]\", \"ncl[sum]\", \"ncl_ratio[2]\"),\n",
    "        DivideTwoFeatures(\"ncl[3]\", \"ncl[sum]\", \"ncl_ratio[3]\"),\n",
    "        \n",
    "        MultiColumnSum(\"MatchedHit_X+DX[0]\", [\"MatchedHit_X[0]\", \"MatchedHit_DX[0]\"]),\n",
    "        MultiColumnSum(\"MatchedHit_X+DX[1]\", [\"MatchedHit_X[1]\", \"MatchedHit_DX[1]\"]),\n",
    "        MultiColumnSum(\"MatchedHit_X+DX[2]\", [\"MatchedHit_X[2]\", \"MatchedHit_DX[2]\"]),\n",
    "        MultiColumnSum(\"MatchedHit_X+DX[3]\", [\"MatchedHit_X[3]\", \"MatchedHit_DX[3]\"]),\n",
    "        \n",
    "        MultiColumnSum(\"MatchedHit_Y+DY[0]\", [\"MatchedHit_Y[0]\", \"MatchedHit_DY[0]\"]),\n",
    "        MultiColumnSum(\"MatchedHit_Y+DY[1]\", [\"MatchedHit_Y[1]\", \"MatchedHit_DY[1]\"]),\n",
    "        MultiColumnSum(\"MatchedHit_Y+DY[2]\", [\"MatchedHit_Y[2]\", \"MatchedHit_DY[2]\"]),\n",
    "        MultiColumnSum(\"MatchedHit_Y+DY[3]\", [\"MatchedHit_Y[3]\", \"MatchedHit_DY[3]\"]),\n",
    "        \n",
    "        MultiColumnSum(\"MatchedHit_Z+DZ[0]\", [\"MatchedHit_Z[0]\", \"MatchedHit_DZ[0]\"]),\n",
    "        MultiColumnSum(\"MatchedHit_Z+DZ[1]\", [\"MatchedHit_Z[1]\", \"MatchedHit_DZ[1]\"]),\n",
    "        MultiColumnSum(\"MatchedHit_Z+DZ[2]\", [\"MatchedHit_Z[2]\", \"MatchedHit_DZ[2]\"]),\n",
    "        MultiColumnSum(\"MatchedHit_Z+DZ[3]\", [\"MatchedHit_Z[3]\", \"MatchedHit_DZ[3]\"]),\n",
    "        \n",
    "        MultiColumnSum(\"MatchedHit_T+DT[0]\", [\"MatchedHit_T[0]\", \"MatchedHit_DT[0]\"]),\n",
    "        MultiColumnSum(\"MatchedHit_T+DT[1]\", [\"MatchedHit_T[1]\", \"MatchedHit_DT[1]\"]),\n",
    "        MultiColumnSum(\"MatchedHit_T+DT[2]\", [\"MatchedHit_T[2]\", \"MatchedHit_DT[2]\"]),\n",
    "        MultiColumnSum(\"MatchedHit_T+DT[3]\", [\"MatchedHit_T[3]\", \"MatchedHit_DT[3]\"]),\n",
    "        \n",
    "        DivideTwoFeatures(\"MatchedHit_X[0]\", \"MatchedHit_X[1]\", \"MatchedHit_X_ratio[0,1]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_X[0]\", \"MatchedHit_X[2]\", \"MatchedHit_X_ratio[0,2]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_X[0]\", \"MatchedHit_X[3]\", \"MatchedHit_X_ratio[0,3]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_X[1]\", \"MatchedHit_X[2]\", \"MatchedHit_X_ratio[1,2]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_X[1]\", \"MatchedHit_X[3]\", \"MatchedHit_X_ratio[1,3]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_X[2]\", \"MatchedHit_X[3]\", \"MatchedHit_X_ratio[2,3]\"),\n",
    "        \n",
    "        DivideTwoFeatures(\"MatchedHit_Y[0]\", \"MatchedHit_Y[1]\", \"MatchedHit_Y_ratio[0,1]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_Y[0]\", \"MatchedHit_Y[2]\", \"MatchedHit_Y_ratio[0,2]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_Y[0]\", \"MatchedHit_Y[3]\", \"MatchedHit_Y_ratio[0,3]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_Y[1]\", \"MatchedHit_Y[2]\", \"MatchedHit_Y_ratio[1,2]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_Y[1]\", \"MatchedHit_Y[3]\", \"MatchedHit_Y_ratio[1,3]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_Y[2]\", \"MatchedHit_Y[3]\", \"MatchedHit_Y_ratio[2,3]\"),\n",
    "        \n",
    "        DivideTwoFeatures(\"MatchedHit_Z[0]\", \"MatchedHit_Z[1]\", \"MatchedHit_Z_ratio[0,1]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_Z[0]\", \"MatchedHit_Z[2]\", \"MatchedHit_Z_ratio[0,2]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_Z[0]\", \"MatchedHit_Z[3]\", \"MatchedHit_Z_ratio[0,3]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_Z[1]\", \"MatchedHit_Z[2]\", \"MatchedHit_Z_ratio[1,2]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_Z[1]\", \"MatchedHit_Z[3]\", \"MatchedHit_Z_ratio[1,3]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_Z[2]\", \"MatchedHit_Z[3]\", \"MatchedHit_Z_ratio[2,3]\"),\n",
    "\n",
    "        MultiplyFeatures(\"count[0]\", [\"ncl[0]\", \"avg_cs[0]\"]),\n",
    "        MultiplyFeatures(\"count[1]\", [\"ncl[1]\", \"avg_cs[1]\"]),\n",
    "        MultiplyFeatures(\"count[2]\", [\"ncl[2]\", \"avg_cs[2]\"]),\n",
    "        MultiplyFeatures(\"count[3]\", [\"ncl[3]\", \"avg_cs[3]\"]),\n",
    "\n",
    "        MultiColumnSum(\"count[sum]\", [\"count[0]\", \"count[1]\", \"count[2]\", \"count[3]\"]),\n",
    "\n",
    "        DivideTwoFeatures(\"count[0]\", \"count[sum]\", \"count_ratio[0]\"),\n",
    "        DivideTwoFeatures(\"count[1]\", \"count[sum]\", \"count_ratio[1]\"),\n",
    "        DivideTwoFeatures(\"count[2]\", \"count[sum]\", \"count_ratio[2]\"),\n",
    "        DivideTwoFeatures(\"count[3]\", \"count[sum]\", \"count_ratio[3]\"),\n",
    "\n",
    "        DivideTwoFeatures(\"count[1]\", \"count[0]\", \"count_ratio[1,0]\"),\n",
    "        DivideTwoFeatures(\"count[2]\", \"count[0]\", \"count_ratio[2,0]\"),\n",
    "        DivideTwoFeatures(\"count[3]\", \"count[0]\", \"count_ratio[3,0]\"),\n",
    "        DivideTwoFeatures(\"count[2]\", \"count[1]\", \"count_ratio[2,1]\"),\n",
    "        DivideTwoFeatures(\"count[3]\", \"count[1]\", \"count_ratio[3,1]\"),\n",
    "        DivideTwoFeatures(\"count[3]\", \"count[2]\", \"count_ratio[3,2]\"),\n",
    "        \n",
    "        MultiColumnSum(\"avg_cs[sum]\", [\"avg_cs[0]\", \"avg_cs[1]\", \"avg_cs[2]\", \"avg_cs[3]\"]),\n",
    "\n",
    "        DivideTwoFeatures(\"avg_cs[0]\", \"avg_cs[sum]\", \"avg_cs_ratio[0]\"),\n",
    "        DivideTwoFeatures(\"avg_cs[1]\", \"avg_cs[sum]\", \"avg_cs_ratio[1]\"),\n",
    "        DivideTwoFeatures(\"avg_cs[2]\", \"avg_cs[sum]\", \"avg_cs_ratio[2]\"),\n",
    "        DivideTwoFeatures(\"avg_cs[3]\", \"avg_cs[sum]\", \"avg_cs_ratio[3]\"),\n",
    "        \n",
    "        DivideTwoFeatures(\"avg_cs[1]\", \"avg_cs[0]\", \"avg_cs_ratio[1,0]\"),\n",
    "        DivideTwoFeatures(\"avg_cs[2]\", \"avg_cs[0]\", \"avg_cs_ratio[2,0]\"),\n",
    "        DivideTwoFeatures(\"avg_cs[3]\", \"avg_cs[0]\", \"avg_cs_ratio[3,0]\"),\n",
    "        DivideTwoFeatures(\"avg_cs[2]\", \"avg_cs[1]\", \"avg_cs_ratio[2,1]\"),\n",
    "        DivideTwoFeatures(\"avg_cs[3]\", \"avg_cs[1]\", \"avg_cs_ratio[3,1]\"),\n",
    "        DivideTwoFeatures(\"avg_cs[3]\", \"avg_cs[2]\", \"avg_cs_ratio[3,2]\"),\n",
    "\n",
    "        RadiusFromCoordinates(\"Lextra_R[0]\", [\"Lextra_X[0]\", \"Lextra_Y[0]\"]),\n",
    "        RadiusFromCoordinates(\"Lextra_R[1]\", [\"Lextra_X[1]\", \"Lextra_Y[1]\"]),\n",
    "        RadiusFromCoordinates(\"Lextra_R[2]\", [\"Lextra_X[2]\", \"Lextra_Y[2]\"]),\n",
    "        RadiusFromCoordinates(\"Lextra_R[3]\", [\"Lextra_X[3]\", \"Lextra_Y[3]\"]),\n",
    "\n",
    "        DivideTwoFeatures(\"Lextra_R[0]\", \"Lextra_R[1]\", \"Lextra_R_ratio[0,1]\"),\n",
    "        DivideTwoFeatures(\"Lextra_R[0]\", \"Lextra_R[2]\", \"Lextra_R_ratio[0,2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_R[0]\", \"Lextra_R[3]\", \"Lextra_R_ratio[0,3]\"),\n",
    "        DivideTwoFeatures(\"Lextra_R[1]\", \"Lextra_R[2]\", \"Lextra_R_ratio[1,2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_R[1]\", \"Lextra_R[3]\", \"Lextra_R_ratio[1,3]\"),\n",
    "        DivideTwoFeatures(\"Lextra_R[2]\", \"Lextra_R[3]\", \"Lextra_R_ratio[2,3]\"),\n",
    "      \n",
    "        RadiusFromCoordinates(\"Lextra_Matched_Hit_Vector_R[0]\", [\"Lextra_Matched_Hit_Vector_X[0]\", \"Lextra_Matched_Hit_Vector_Y[0]\"]),\n",
    "        RadiusFromCoordinates(\"Lextra_Matched_Hit_Vector_R[1]\", [\"Lextra_Matched_Hit_Vector_X[1]\", \"Lextra_Matched_Hit_Vector_Y[1]\"]),\n",
    "        RadiusFromCoordinates(\"Lextra_Matched_Hit_Vector_R[2]\", [\"Lextra_Matched_Hit_Vector_X[2]\", \"Lextra_Matched_Hit_Vector_Y[2]\"]),\n",
    "        RadiusFromCoordinates(\"Lextra_Matched_Hit_Vector_R[3]\", [\"Lextra_Matched_Hit_Vector_X[3]\", \"Lextra_Matched_Hit_Vector_Y[3]\"]),\n",
    "\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_R[0]\", \"Lextra_Matched_Hit_Vector_R[1]\", \"Lextra_Matched_Hit_Vector_R_ratio[0,1]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_R[0]\", \"Lextra_Matched_Hit_Vector_R[2]\", \"Lextra_Matched_Hit_Vector_R_ratio[0,2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_R[0]\", \"Lextra_Matched_Hit_Vector_R[3]\", \"Lextra_Matched_Hit_Vector_R_ratio[0,3]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_R[1]\", \"Lextra_Matched_Hit_Vector_R[2]\", \"Lextra_Matched_Hit_Vector_R_ratio[1,2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_R[1]\", \"Lextra_Matched_Hit_Vector_R[3]\", \"Lextra_Matched_Hit_Vector_R_ratio[1,3]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_R[2]\", \"Lextra_Matched_Hit_Vector_R[3]\", \"Lextra_Matched_Hit_Vector_R_ratio[2,3]\"),\n",
    "      \n",
    "        DivideTwoFeatures(\"Lextra_X[0]\", \"MatchedHit_X[0]\", \"Lextra_to_MachedHit_ratio_X[0]\"),\n",
    "        DivideTwoFeatures(\"Lextra_X[1]\", \"MatchedHit_X[1]\", \"Lextra_to_MachedHit_ratio_X[1]\"),\n",
    "        DivideTwoFeatures(\"Lextra_X[2]\", \"MatchedHit_X[2]\", \"Lextra_to_MachedHit_ratio_X[2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_X[3]\", \"MatchedHit_X[3]\", \"Lextra_to_MachedHit_ratio_X[3]\"),\n",
    "        \n",
    "        DivideTwoFeatures(\"Lextra_Y[0]\", \"MatchedHit_Y[0]\", \"Lextra_to_MachedHit_ratio_Y[0]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Y[1]\", \"MatchedHit_Y[1]\", \"Lextra_to_MachedHit_ratio_Y[1]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Y[2]\", \"MatchedHit_Y[2]\", \"Lextra_to_MachedHit_ratio_Y[2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Y[3]\", \"MatchedHit_Y[3]\", \"Lextra_to_MachedHit_ratio_Y[3]\"),\n",
    "\n",
    "        RadiusFromCoordinates(\"MatchedHit_R[0]\", [\"MatchedHit_X[0]\", \"MatchedHit_Y[0]\", \"MatchedHit_Z[0]\"]),\n",
    "        RadiusFromCoordinates(\"MatchedHit_R[1]\", [\"MatchedHit_X[1]\", \"MatchedHit_Y[1]\", \"MatchedHit_Z[1]\"]),\n",
    "        RadiusFromCoordinates(\"MatchedHit_R[2]\", [\"MatchedHit_X[2]\", \"MatchedHit_Y[2]\", \"MatchedHit_Z[2]\"]),\n",
    "        RadiusFromCoordinates(\"MatchedHit_R[3]\", [\"MatchedHit_X[3]\", \"MatchedHit_Y[3]\", \"MatchedHit_Z[3]\"]),\n",
    "\n",
    "        DivideTwoFeatures(\"MatchedHit_R[0]\", \"MatchedHit_R[1]\", \"MatchedHit_R_ratio[0,1]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_R[0]\", \"MatchedHit_R[2]\", \"MatchedHit_R_ratio[0,2]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_R[0]\", \"MatchedHit_R[3]\", \"MatchedHit_R_ratio[0,3]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_R[1]\", \"MatchedHit_R[2]\", \"MatchedHit_R_ratio[1,2]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_R[1]\", \"MatchedHit_R[3]\", \"MatchedHit_R_ratio[1,3]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_R[2]\", \"MatchedHit_R[3]\", \"MatchedHit_R_ratio[2,3]\"),\n",
    "        \n",
    "        RadiusFromCoordinates(\"MatchedHit_R1[0]\", [\"MatchedHit_X[0]\", \"MatchedHit_Y[0]\"]),\n",
    "        RadiusFromCoordinates(\"MatchedHit_R1[1]\", [\"MatchedHit_X[1]\", \"MatchedHit_Y[1]\"]),\n",
    "        RadiusFromCoordinates(\"MatchedHit_R1[2]\", [\"MatchedHit_X[2]\", \"MatchedHit_Y[2]\"]),\n",
    "        RadiusFromCoordinates(\"MatchedHit_R1[3]\", [\"MatchedHit_X[3]\", \"MatchedHit_Y[3]\"]),\n",
    "\n",
    "        DivideTwoFeatures(\"MatchedHit_R1[0]\", \"MatchedHit_R1[1]\", \"MatchedHit_R1_ratio[0,1]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_R1[0]\", \"MatchedHit_R1[2]\", \"MatchedHit_R1_ratio[0,2]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_R1[0]\", \"MatchedHit_R1[3]\", \"MatchedHit_R1_ratio[0,3]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_R1[1]\", \"MatchedHit_R1[2]\", \"MatchedHit_R1_ratio[1,2]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_R1[1]\", \"MatchedHit_R1[3]\", \"MatchedHit_R1_ratio[1,3]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_R1[2]\", \"MatchedHit_R1[3]\", \"MatchedHit_R1_ratio[2,3]\"),\n",
    "      \n",
    "        DivideTwoFeatures(\"MatchedHit_T[0]\", \"MatchedHit_T[1]\", \"MatchedHit_T_ratio[0,1]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_T[0]\", \"MatchedHit_T[2]\", \"MatchedHit_T_ratio[0,2]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_T[0]\", \"MatchedHit_T[3]\", \"MatchedHit_T_ratio[0,3]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_T[1]\", \"MatchedHit_T[2]\", \"MatchedHit_T_ratio[1,2]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_T[1]\", \"MatchedHit_T[3]\", \"MatchedHit_T_ratio[1,3]\"),\n",
    "        DivideTwoFeatures(\"MatchedHit_T[2]\", \"MatchedHit_T[3]\", \"MatchedHit_T_ratio[2,3]\"),\n",
    "        \n",
    "        DivideTwoFeatures(\"Lextra_X[0]\", \"Lextra_X[1]\", \"Lextra_X_ratio[0,1]\"),\n",
    "        DivideTwoFeatures(\"Lextra_X[0]\", \"Lextra_X[2]\", \"Lextra_X_ratio[0,2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_X[0]\", \"Lextra_X[3]\", \"Lextra_X_ratio[0,3]\"),\n",
    "        DivideTwoFeatures(\"Lextra_X[1]\", \"Lextra_X[2]\", \"Lextra_X_ratio[1,2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_X[1]\", \"Lextra_X[3]\", \"Lextra_X_ratio[1,3]\"),\n",
    "        DivideTwoFeatures(\"Lextra_X[2]\", \"Lextra_X[3]\", \"Lextra_X_ratio[3,3]\"),\n",
    "      \n",
    "        DivideTwoFeatures(\"Lextra_Y[0]\", \"Lextra_Y[1]\", \"Lextra_Y_ratio[0,1]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Y[0]\", \"Lextra_Y[2]\", \"Lextra_Y_ratio[0,2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Y[0]\", \"Lextra_Y[3]\", \"Lextra_Y_ratio[0,3]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Y[1]\", \"Lextra_Y[2]\", \"Lextra_Y_ratio[1,2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Y[1]\", \"Lextra_Y[3]\", \"Lextra_Y_ratio[1,3]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Y[2]\", \"Lextra_Y[3]\", \"Lextra_Y_ratio[3,3]\"),\n",
    "      \n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_X[0]\", \"Lextra_Matched_Hit_Vector_X[1]\", \"Lextra_Matched_Hit_Vector_X_ratio[0,1]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_X[0]\", \"Lextra_Matched_Hit_Vector_X[2]\", \"Lextra_Matched_Hit_Vector_X_ratio[0,2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_X[0]\", \"Lextra_Matched_Hit_Vector_X[3]\", \"Lextra_Matched_Hit_Vector_X_ratio[0,3]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_X[1]\", \"Lextra_Matched_Hit_Vector_X[2]\", \"Lextra_Matched_Hit_Vector_X_ratio[1,2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_X[1]\", \"Lextra_Matched_Hit_Vector_X[3]\", \"Lextra_Matched_Hit_Vector_X_ratio[1,3]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_X[2]\", \"Lextra_Matched_Hit_Vector_X[3]\", \"Lextra_Matched_Hit_Vector_X_ratio[2,3]\"),\n",
    "  \n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_Y[0]\", \"Lextra_Matched_Hit_Vector_Y[1]\", \"Lextra_Matched_Hit_Vector_Y_ratio[0,1]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_Y[0]\", \"Lextra_Matched_Hit_Vector_Y[2]\", \"Lextra_Matched_Hit_Vector_Y_ratio[0,2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_Y[0]\", \"Lextra_Matched_Hit_Vector_Y[3]\", \"Lextra_Matched_Hit_Vector_Y_ratio[0,3]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_Y[1]\", \"Lextra_Matched_Hit_Vector_Y[2]\", \"Lextra_Matched_Hit_Vector_Y_ratio[1,2]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_Y[1]\", \"Lextra_Matched_Hit_Vector_Y[3]\", \"Lextra_Matched_Hit_Vector_Y_ratio[1,3]\"),\n",
    "        DivideTwoFeatures(\"Lextra_Matched_Hit_Vector_Y[2]\", \"Lextra_Matched_Hit_Vector_Y[3]\", \"Lextra_Matched_Hit_Vector_Y_ratio[2,3]\")\n",
    "    )\n",
    "    df = pipeline.fit_transform(df)\n",
    "   \n",
    "    df = category_encoders.OneHotEncoder(\n",
    "        cols=[\n",
    "            \"ndof\", \n",
    "            \"MatchedHit_TYPE[0]\", \n",
    "            \"MatchedHit_TYPE[1]\", \n",
    "            \"MatchedHit_TYPE[2]\", \n",
    "            \"MatchedHit_TYPE[3]\"\n",
    "        ]\n",
    "    ).fit_transform(df)\n",
    "    \n",
    "    drop_features = [\n",
    "        'FOI_hits_X', \n",
    "        'FOI_hits_Y', \n",
    "        'FOI_hits_Z', \n",
    "        'FOI_hits_DX', \n",
    "        'FOI_hits_DY', \n",
    "        'FOI_hits_DZ', \n",
    "        'FOI_hits_T', \n",
    "        'FOI_hits_DT',\n",
    "        'FOI_hits_S'\n",
    "    ]\n",
    "    df = FeaturesDropper(drop_features).fit_transform(df)\n",
    "    \n",
    "    # Replace characters in feature names that are not supported by xgboost\n",
    "    df.columns = [col.replace(\"[\", \"(\") for col in df.columns]\n",
    "    df.columns = [col.replace(\"]\", \")\") for col in df.columns]\n",
    "    df.columns = [col.replace(\"<\", \"_lt_\") for col in df.columns]\n",
    "    df.columns = [col.replace(\">\", \"_gt_\") for col in df.columns]\n",
    "    \n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train_test = prepare_df(df_test[0:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train_test.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Resample training dataset to make the label distribution balanced."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resampler = RandomUnderSampler(random_state=21)\n",
    "df_train_resampled, _ = resampler.fit_resample(df_train, df_train.label)\n",
    "df_train_resampled = pd.DataFrame(df_train_resampled, columns=df_train.columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for col in list(df_train_resampled.columns):\n",
    "    if col not in categorical_features:\n",
    "        df_train_resampled[col] = pd.to_numeric(df_train_resampled[col])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare train dataframe for training/inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train_resampled1 = prepare_df(df_train_resampled)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define and train regression models to predict **weight**, **sWeight**, and **kinWeight** from other features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# weight model\n",
    "model1 = make_pipeline(lgb.LGBMRegressor(n_estimators=1000, random_state=21))\n",
    "model1.fit(\n",
    "    df_train_resampled1[df_train_resampled1.columns.difference(excluded_features)],\n",
    "    df_train_resampled1.weight,\n",
    ")\n",
    "\n",
    "# sWeight model\n",
    "model3 = make_pipeline(lgb.LGBMRegressor(n_estimators=1000, random_state=21))\n",
    "model3.fit(\n",
    "    df_train_resampled1[df_train_resampled1.columns.difference(excluded_features)],\n",
    "    df_train_resampled1.sWeight,\n",
    ")\n",
    "\n",
    "# kinWeight model\n",
    "model4 = make_pipeline(lgb.LGBMRegressor(n_estimators=1000, random_state=21))\n",
    "model4.fit(\n",
    "    df_train_resampled1[df_train_resampled1.columns.difference(excluded_features)],\n",
    "    df_train_resampled1.kinWeight,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make predictions from the **weight**, **sWeight** and **kinWeight** models with the training data and then concatenate the predictions to the original training dataframe as new features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions_weight = model1.predict(\n",
    "    df_train_resampled1[df_train_resampled1.columns.difference(excluded_features)],\n",
    ")\n",
    "predictions_weight = pd.DataFrame(predictions_weight, columns=['prediction_weight'])\n",
    "\n",
    "predictions_sWeight = model3.predict(\n",
    "    df_train_resampled1[df_train_resampled1.columns.difference(excluded_features)],\n",
    ")\n",
    "predictions_sWeight = pd.DataFrame(predictions_sWeight, columns=['prediction_sWeight'])\n",
    "\n",
    "predictions_kinWeight = model4.predict(\n",
    "    df_train_resampled1[df_train_resampled1.columns.difference(excluded_features)],\n",
    ")\n",
    "predictions_kinWeight=pd.DataFrame(predictions_kinWeight, columns=['prediction_kinWeight'])\n",
    "\n",
    "\n",
    "df_train_resampled_all = pd.concat(\n",
    "    [\n",
    "        df_train_resampled1, \n",
    "        predictions_weight, \n",
    "        predictions_sWeight, \n",
    "        predictions_kinWeight\n",
    "    ],\n",
    "    axis=1\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define and train a classifier on the concatenated dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model2 = make_pipeline(lgb.LGBMClassifier(n_estimators=1000, random_state=21))\n",
    "model2.fit(\n",
    "    df_train_resampled_all[\n",
    "        df_train_resampled_all.columns.difference(excluded_features)\n",
    "    ],\n",
    "    df_train_resampled_all.label,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load private test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test_private = pd.read_csv(\"data/test_private_v2_track_1.csv\")\n",
    "df_test_private.set_index(\"id\", inplace=True)\n",
    "df_test_private.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare private test dataframe for training/inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test_private1 = prepare_df(df_test_private)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make predictions from private test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# weight predictions\n",
    "predictions_weight_private = model1.predict(\n",
    "    df_test_private1[df_test_private1.columns.difference(excluded_features)]\n",
    ")\n",
    "predictions_weight_private = pd.DataFrame(\n",
    "    predictions_weight_private, columns=[\"prediction_weight\"]\n",
    ")\n",
    "\n",
    "# sWeight predictions\n",
    "predictions_sWeight_private = model3.predict(\n",
    "    df_test_private1[df_test_private1.columns.difference(excluded_features)]\n",
    ")\n",
    "predictions_sWeight_private = pd.DataFrame(\n",
    "    predictions_sWeight_private, columns=[\"prediction_sWeight\"]\n",
    ")\n",
    "\n",
    "# kinWeight predictions\n",
    "predictions_kinWeight_private = model4.predict(\n",
    "    df_test_private1[df_test_private1.columns.difference(excluded_features)]\n",
    ")\n",
    "predictions_kinWeight_private = pd.DataFrame(\n",
    "    predictions_kinWeight_private, columns=[\"prediction_kinWight\"]\n",
    ")\n",
    "\n",
    "# Private dataset concatenated with weight, sWeight and kinWeight features\n",
    "df_test_private_all = pd.concat(\n",
    "    [\n",
    "        df_test_private1,\n",
    "        predictions_weight_private,\n",
    "        predictions_sWeight_private,\n",
    "        predictions_kinWeight_private,\n",
    "    ],\n",
    "    axis=1,\n",
    ")\n",
    "test_private_score = model2.predict_proba(\n",
    "    df_test_private_all[df_test_private_all.columns.difference(excluded_features)]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write submission csv file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame(\n",
    "    data={\"prediction\": test_private_score[:, 1]}, index=df_test_private_all.index\n",
    ").to_csv(\"submission_private_a.csv\", index_label=\"id\", header=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}