DataIngestion_FeaturePreparation/1-DataIngestion.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Ingestion and Filtering - pipeline for the topology classifier with Apache Spark\n",
    "\n",
    "**1. Data Ingestion** is the first stage of the pipeline. Here we will read the ROOT file from HDFS into a Spark dataframe using [Spark-ROOT](https://github.com/diana-hep/spark-root) reader and then we will create the Low Level Features (LLF) and High Level Features datasets.\n",
    "\n",
    "To run this notebook we used the following configuration:\n",
    "* *Software stack*: Spark 2.4.3\n",
    "* *Platform*: CentOS 7, Python 3.6\n",
    "* *Spark cluster*: Analytix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install pyspark or use your favorite way to set Spark Home, here we use findspark\n",
    "import findspark\n",
    "findspark.init('/home/luca/Spark/spark-2.4.3-bin-hadoop2.7') #set path to SPARK_HOME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure according to your environment\n",
    "pyspark_python = \"<path to python>/bin/python\"\n",
    "spark_root_jar=\"https://github.com/diana-hep/spark-root/blob/master/jars/spark-root_2.11-0.1.17.jar?raw=true\"\n",
    "\n",
    "from pyspark.sql import SparkSession\n",
    "spark = SparkSession.builder \\\n",
    "        .appName(\"1-Data Ingestion\") \\\n",
    "        .master(\"yarn\") \\\n",
    "        .config(\"spark.driver.memory\",\"8g\") \\\n",
    "        .config(\"spark.executor.memory\",\"14g\") \\\n",
    "        .config(\"spark.executor.cores\",\"8\") \\\n",
    "        .config(\"spark.executor.instances\",\"50\") \\\n",
    "        .config(\"spark.dynamicAllocation.enabled\",\"false\") \\\n",
    "        .config(\"spark.jars\",spark_root_jar) \\\n",
    "        .config(\"spark.jars.packages\",\"org.diana-hep:root4j:0.1.6\") \\\n",
    "        .config(\"spark.pyspark.python\",pyspark_python) \\\n",
    "        .config(\"spark.eventLog.enabled\",\"false\") \\\n",
    "        .getOrCreate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "            <div>\n",
       "                <p><b>SparkSession - in-memory</b></p>\n",
       "                \n",
       "        <div>\n",
       "            <p><b>SparkContext</b></p>\n",
       "\n",
       "            <p><a href=\"http://pcitdbgpu1.dyndns.cern.ch:4040\">Spark UI</a></p>\n",
       "\n",
       "            <dl>\n",
       "              <dt>Version</dt>\n",
       "                <dd><code>v2.4.3</code></dd>\n",
       "              <dt>Master</dt>\n",
       "                <dd><code>yarn</code></dd>\n",
       "              <dt>AppName</dt>\n",
       "                <dd><code>1-Data Ingestion</code></dd>\n",
       "            </dl>\n",
       "        </div>\n",
       "        \n",
       "            </div>\n",
       "        "
      ],
      "text/plain": [
       "<pyspark.sql.session.SparkSession at 0x7fad28e67240>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check if Spark Session has been created correctly\n",
    "spark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a file containing functions that we will use later\n",
    "spark.sparkContext.addPyFile(\"utilFunctions.py\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read the data from HDFS\n",
    "<br>\n",
    "As first step we will read the samples into a Spark dataframe using Spark-Root. We will select only a subset of columns present in the original files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH = \"hdfs://analytix/Training/Spark/TopologyClassifier/lepFilter_rawData/\"\n",
    "\n",
    "samples = [\"qcd_lepFilter_13TeV, \"ttbar_lepFilter_13TeV\", \"Wlnu_lepFilter_13TeV\"]\n",
    "\n",
    "requiredColumns = [\n",
    "    \"EFlowTrack\",\n",
    "    \"EFlowNeutralHadron\",\n",
    "    \"EFlowPhoton\",\n",
    "    \"Electron\",\n",
    "    \"MuonTight\",\n",
    "    \"MuonTight_size\",\n",
    "    \"Electron_size\",\n",
    "    \"MissingET\",\n",
    "    \"Jet\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading Wlnu_lepFilter_13TeV sample...\n",
      "Loading qcd_lepFilter_13TeV sample...\n",
      "Loading ttbar_lepFilter_13TeV sample...\n"
     ]
    }
   ],
   "source": [
    "from pyspark.sql.functions import lit\n",
    "\n",
    "dfList = []\n",
    "\n",
    "for label,sample in enumerate(samples):\n",
    "    print(\"Loading {} sample...\".format(sample))\n",
    "    tmpDF = spark.read \\\n",
    "                .format(\"org.dianahep.sparkroot.experimental\") \\\n",
    "                .load(PATH + sample + \"/*.root\") \\\n",
    "                .select(requiredColumns) \\\n",
    "                .withColumn(\"label\", lit(label))\n",
    "    dfList.append(tmpDF)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Merge all samples into a single dataframe\n",
    "df = dfList[0]\n",
    "for tmpDF in dfList[1:]:\n",
    "    df = df.union(tmpDF)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's have a look at how many events there are for each class. Keep in mind that the labels are mapped as follow\n",
    "* $0=\\text{QCD}$\n",
    "* $1=\\text{t}\\bar{\\text{t}}$\n",
    "* $2=\\text{W}+\\text{jets}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----+--------+\n",
      "|label|   count|\n",
      "+-----+--------+\n",
      "|    2|26335315|\n",
      "|    1|13780026|\n",
      "|    0|14354796|\n",
      "| null|54470137|\n",
      "+-----+--------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Print the number of events per sample and the total (label=null)\n",
    "df.rollup(\"label\").count().orderBy(\"label\", ascending=False).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's print the schema of one of the required columns. This shows that the  \n",
    "**schema is complex and nested** (the full schema is even more complex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- EFlowTrack: array (nullable = true)\n",
      " |    |-- element: struct (containsNull = true)\n",
      " |    |    |-- fUniqueID: integer (nullable = true)\n",
      " |    |    |-- fBits: integer (nullable = true)\n",
      " |    |    |-- PID: integer (nullable = true)\n",
      " |    |    |-- Charge: integer (nullable = true)\n",
      " |    |    |-- PT: float (nullable = true)\n",
      " |    |    |-- Eta: float (nullable = true)\n",
      " |    |    |-- Phi: float (nullable = true)\n",
      " |    |    |-- EtaOuter: float (nullable = true)\n",
      " |    |    |-- PhiOuter: float (nullable = true)\n",
      " |    |    |-- X: float (nullable = true)\n",
      " |    |    |-- Y: float (nullable = true)\n",
      " |    |    |-- Z: float (nullable = true)\n",
      " |    |    |-- T: float (nullable = true)\n",
      " |    |    |-- XOuter: float (nullable = true)\n",
      " |    |    |-- YOuter: float (nullable = true)\n",
      " |    |    |-- ZOuter: float (nullable = true)\n",
      " |    |    |-- TOuter: float (nullable = true)\n",
      " |    |    |-- Dxy: float (nullable = true)\n",
      " |    |    |-- SDxy: float (nullable = true)\n",
      " |    |    |-- Xd: float (nullable = true)\n",
      " |    |    |-- Yd: float (nullable = true)\n",
      " |    |    |-- Zd: float (nullable = true)\n",
      " |    |    |-- EFlowTrack_Particle: struct (nullable = true)\n",
      " |    |    |    |-- TObject: struct (nullable = true)\n",
      " |    |    |    |    |-- fUniqueID: integer (nullable = true)\n",
      " |    |    |    |    |-- fBits: integer (nullable = true)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.select(\"EFlowTrack\").printSchema()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create derivate datasets\n",
    "\n",
    "Now we will create the LLF and HLF datasets. This is done by the function `convert` below which takes as input an event (i.e. the list of particles present in that event) and do the following steps:\n",
    "1. Select the events with at least one isolated electron/muon (implemented in `selection`)\n",
    "2. Create the list of 801 particles and the 19 low level features for each of them\n",
    "3. Compute the high level features "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import Lorentz Vector and other functions for pTmaps\n",
    "from utilFunctions import *\n",
    "\n",
    "def selection(event, TrkPtMap, NeuPtMap, PhotonPtMap, PTcut=23., ISOcut=0.45):\n",
    "    \"\"\"\n",
    "    This function simulates the trigger selection. \n",
    "    Foreach event the presence of one isolated muon or electron with pT >23 GeV is required\n",
    "    \"\"\"\n",
    "    if event.Electron_size == 0 and event.MuonTight_size == 0: \n",
    "        return False, False, False\n",
    "    \n",
    "    foundMuon = None \n",
    "    foundEle =  None \n",
    "    \n",
    "    l = LorentzVector()\n",
    "    \n",
    "    for ele in event.Electron:\n",
    "        if ele.PT <= PTcut: continue\n",
    "        l.SetPtEtaPhiM(ele.PT, ele.Eta, ele.Phi, 0.)\n",
    "        \n",
    "        pfisoCh = PFIso(l, 0.3, TrkPtMap, True)\n",
    "        pfisoNeu = PFIso(l, 0.3, NeuPtMap, False)\n",
    "        pfisoGamma = PFIso(l, 0.3, PhotonPtMap, False)\n",
    "        if foundEle == None and (pfisoCh+pfisoNeu+pfisoGamma)<ISOcut:\n",
    "            foundEle = [l.E(), l.Px(), l.Py(), l.Pz(), l.Pt(), l.Eta(), l.Phi(),\n",
    "                        0., 0., 0., pfisoCh, pfisoGamma, pfisoNeu,\n",
    "                        0., 0., 0., 1., 0., float(ele.Charge)]\n",
    "    \n",
    "    for muon in event.MuonTight:\n",
    "        if muon.PT <= PTcut: continue\n",
    "        l.SetPtEtaPhiM(muon.PT, muon.Eta, muon.Phi, 0.)\n",
    "        \n",
    "        pfisoCh = PFIso(l, 0.3, TrkPtMap, True)\n",
    "        pfisoNeu = PFIso(l, 0.3, NeuPtMap, False)\n",
    "        pfisoGamma = PFIso(l, 0.3, PhotonPtMap, False)\n",
    "        if foundMuon == None and (pfisoCh+pfisoNeu+pfisoGamma)<ISOcut:\n",
    "            foundMuon = [l.E(), l.Px(), l.Py(), l.Pz(), l.Pt(), l.Eta(), l.Phi(),\n",
    "                         0., 0., 0., pfisoCh, pfisoGamma, pfisoNeu,\n",
    "                         0., 0., 0., 0., 1., float(muon.Charge)]\n",
    "            \n",
    "    if foundEle != None and foundMuon != None:\n",
    "        if foundEle[5] > foundMuon[5]:\n",
    "            return True, foundEle, foundMuon\n",
    "        else:\n",
    "            return True, foundMuon, foundEle\n",
    "    if foundEle != None: return True, foundEle, foundMuon\n",
    "    if foundMuon != None: return True, foundMuon, foundEle\n",
    "    \n",
    "    return False, None, None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "from pyspark.ml.linalg import Vectors\n",
    "\n",
    "def convert(event):\n",
    "    \"\"\"\n",
    "    This function takes as input an event, applies trigger selection \n",
    "    and create LLF and HLF datasets\n",
    "    \"\"\"\n",
    "    q = LorentzVector()\n",
    "    particles = []\n",
    "    TrkPtMap = ChPtMapp(0.3, event)\n",
    "    NeuPtMap = NeuPtMapp(0.3, event)\n",
    "    PhotonPtMap = PhotonPtMapp(0.3, event)\n",
    "    if TrkPtMap.shape[0] == 0: return Row()\n",
    "    if NeuPtMap.shape[0] == 0: return Row()\n",
    "    if PhotonPtMap.shape[0] == 0: return Row()\n",
    "    \n",
    "    #\n",
    "    # Get leptons\n",
    "    #\n",
    "    selected, lep, otherlep = selection(event, TrkPtMap, NeuPtMap, PhotonPtMap)\n",
    "    if not selected: return Row()\n",
    "    particles.append(lep)\n",
    "    lepMomentum = LorentzVector(lep[1], lep[2], lep[3], lep[0])\n",
    "    \n",
    "    #\n",
    "    # Select Tracks\n",
    "    #\n",
    "    nTrk = 0\n",
    "    for h in event.EFlowTrack:\n",
    "        if nTrk>=450: continue\n",
    "        if h.PT<=0.5: continue\n",
    "        q.SetPtEtaPhiM(h.PT, h.Eta, h.Phi, 0.)\n",
    "        if lepMomentum.DeltaR(q) > 0.0001:\n",
    "            pfisoCh = PFIso(q, 0.3, TrkPtMap, True)\n",
    "            pfisoNeu = PFIso(q, 0.3, NeuPtMap, False)\n",
    "            pfisoGamma = PFIso(q, 0.3, PhotonPtMap, False)\n",
    "            particles.append([q.E(), q.Px(), q.Py(), q.Pz(),\n",
    "                              h.PT, h.Eta, h.Phi, h.X, h.Y, h.Z,\n",
    "                              pfisoCh, pfisoGamma, pfisoNeu,\n",
    "                              1., 0., 0., 0., 0., float(np.sign(h.PID))])\n",
    "            nTrk += 1\n",
    "    \n",
    "    #\n",
    "    # Select Photons\n",
    "    #\n",
    "    nPhoton = 0\n",
    "    for h in event.EFlowPhoton:\n",
    "        if nPhoton >= 150: continue\n",
    "        if h.ET <= 1.: continue\n",
    "        q.SetPtEtaPhiM(h.ET, h.Eta, h.Phi, 0.)\n",
    "        pfisoCh = PFIso(q, 0.3, TrkPtMap, True)\n",
    "        pfisoNeu = PFIso(q, 0.3, NeuPtMap, False)\n",
    "        pfisoGamma = PFIso(q, 0.3, PhotonPtMap, False)\n",
    "        particles.append([q.E(), q.Px(), q.Py(), q.Pz(),\n",
    "                          h.ET, h.Eta, h.Phi, 0., 0., 0.,\n",
    "                          pfisoCh, pfisoGamma, pfisoNeu,\n",
    "                          0., 0., 1., 0., 0., 0.])\n",
    "        nPhoton += 1\n",
    "    \n",
    "    #\n",
    "    # Select Neutrals\n",
    "    #\n",
    "    nNeu = 0\n",
    "    for h in event.EFlowNeutralHadron:\n",
    "        if nNeu >= 200: continue\n",
    "        if h.ET <= 1.: continue\n",
    "        q.SetPtEtaPhiM(h.ET, h.Eta, h.Phi, 0.)\n",
    "        pfisoCh = PFIso(q, 0.3, TrkPtMap, True)\n",
    "        pfisoNeu = PFIso(q, 0.3, NeuPtMap, False)\n",
    "        pfisoGamma = PFIso(q, 0.3, PhotonPtMap, False)\n",
    "        particles.append([q.E(), q.Px(), q.Py(), q.Pz(),\n",
    "                          h.ET, h.Eta, h.Phi, 0., 0., 0.,\n",
    "                          pfisoCh, pfisoGamma, pfisoNeu,\n",
    "                          0., 1., 0., 0., 0., 0.])\n",
    "        nNeu += 1\n",
    "        \n",
    "    for iTrk in range(nTrk, 450):\n",
    "        particles.append([0., 0., 0., 0., 0., 0., 0., 0.,0.,\n",
    "                          0.,0., 0., 0., 0., 0., 0., 0., 0., 0.])\n",
    "    for iPhoton in range(nPhoton, 150):\n",
    "        particles.append([0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
    "                          0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])\n",
    "    for iNeu in range(nNeu, 200):\n",
    "        particles.append([0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
    "                          0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])        \n",
    "    #\n",
    "    # High Level Features\n",
    "    #\n",
    "    myMET = event.MissingET[0]\n",
    "    MET = myMET.MET\n",
    "    phiMET = myMET.Phi\n",
    "    MT = 2.*MET*lepMomentum.Pt()*(1-math.cos(lepMomentum.Phi()-phiMET))\n",
    "    HT = 0.\n",
    "    nJets = 0.\n",
    "    nBjets = 0.\n",
    "    for jet in event.Jet:\n",
    "        if jet.PT > 30 and abs(jet.Eta)<2.6:\n",
    "            nJets += 1\n",
    "            HT += jet.PT\n",
    "            if jet.BTag>0: \n",
    "                nBjets += 1\n",
    "    LepPt = lep[4]\n",
    "    LepEta = lep[5]\n",
    "    LepPhi = lep[6]\n",
    "    LepIsoCh = lep[10]\n",
    "    LepIsoGamma = lep[11]\n",
    "    LepIsoNeu = lep[12]\n",
    "    LepCharge = lep[18]\n",
    "    LepIsEle = lep[16]\n",
    "    hlf = Vectors.dense([HT, MET, phiMET, MT, nJets, nBjets, LepPt, LepEta, LepPhi,\n",
    "           LepIsoCh, LepIsoGamma, LepIsoNeu, LepCharge, LepIsEle])     \n",
    "    #\n",
    "    # return the Row of low level features and high level features\n",
    "    #\n",
    "    return Row(lfeatures=particles, hfeatures=hlf, label=event.label)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally apply the function to all the events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = df.rdd \\\n",
    "            .map(convert) \\\n",
    "            .filter(lambda row: len(row) > 0) \\\n",
    "            .toDF()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- hfeatures: vector (nullable = true)\n",
      " |-- label: long (nullable = true)\n",
      " |-- lfeatures: array (nullable = true)\n",
      " |    |-- element: array (containsNull = true)\n",
      " |    |    |-- element: double (containsNull = true)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "features.printSchema()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save the datasets as Parquet files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.07 s, sys: 700 ms, total: 1.77 s\n",
      "Wall time: 3h 13min 36s\n"
     ]
    }
   ],
   "source": [
    "dataset_path = \"hdfs://analytix/Training/Spark/TopologyClassifier/dataIngestion_full_13TeV\"\n",
    "num_partitions = 3000 # used in DataFrame coalesce operation to limit number of output files\n",
    "\n",
    "%time features.coalesce(num_partitions).write.partitionBy(\"label\").parquet(dataset_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of events written to Parquet: 25468470\n"
     ]
    }
   ],
   "source": [
    "print(\"Number of events written to Parquet:\", spark.read.parquet(dataset_path).count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark.stop()"
   ]
  },
  {
   "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.8"
  },
  "sparkconnect": {
   "bundled_options": [],
   "list_of_options": [
    {
     "name": "spark.executor.cores",
     "value": "4"
    },
    {
     "name": "spark.executor.memory",
     "value": "8G"
    },
    {
     "name": "spark.executor.instances",
     "value": "10"
    },
    {
     "name": "spark.dynamicAllocation.enabled",
     "value": "false"
    },
    {
     "name": "spark.jars.packages",
     "value": "org.diana-hep:spark-root_2.11:0.1.16"
    },
    {
     "name": "spark.app.name",
     "value": "DataIngestion"
    }
   ]
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}