Zpeak_nano_multipledataset_v2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Once deleted, variables cannot be recovered. Proceed (y/[n])? y\n"
     ]
    }
   ],
   "source": [
    "%reset\n",
    "import pyspark.sql\n",
    "session = pyspark.sql.SparkSession.builder \\\n",
    "    .master('local') \\\n",
    "    .appName(\"Zpeak\") \\\n",
    "    .config('spark.jars.packages','org.diana-hep:spark-root_2.11:0.1.16') \\\n",
    "    .config('spark.driver.extraClassPath','/opt/hadoop/share/hadoop/common/lib/EOSfs.jar') \\\n",
    "    .getOrCreate()\n",
    "    \n",
    "#sc = session.sparkContext\n",
    "sqlContext = session"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading NanoAod root files using spark-root package\n",
    "\n",
    "Loading the nanoaod root file serving from cern public eos, and load the tree in dedicated TDirectory in the root file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TT\n",
      "WW\n",
      "SingleMuon\n",
      "ZZ\n",
      "DYJetsToLL\n",
      "WZ\n",
      "Number of element:\n",
      "DFList size = 6\n",
      "Number of columns:\n",
      "DFList[0][0] = 979\n",
      "DFList[0][0] = 964\n",
      "DFList[1][0] = 869\n",
      "DFList[2][0] = 964\n",
      "DFList[3][0] = 979\n",
      "DFList[4][0] = 964\n"
     ]
    }
   ],
   "source": [
    "from pyspark.sql.functions import lit\n",
    "from samples import *\n",
    "from __future__ import division\n",
    "\n",
    "DFList = [] # will contain six dataframe from six processes with added xsec, eff, kfactor, weight, samples\n",
    "\n",
    "for s in samples:\n",
    "    print s\n",
    "    #dsPath = \"hdfs://10.64.22.72:9000/\"+BASE+samples[s]['filename']\n",
    "    dsPath = \"root://eospublic.cern.ch//eos/opstest/cmspd-bigdata/\"+samples[s]['filename']\n",
    "    \n",
    "    tempDF = sqlContext.read \\\n",
    "                .format(\"org.dianahep.sparkroot\") \\\n",
    "                .option(\"tree\", \"Events\") \\\n",
    "                .load(dsPath)\\\n",
    "                .withColumn(\"proc_xsec\", lit(samples[s]['xsec']))\\\n",
    "                .withColumn(\"proc_eff\", lit(samples[s]['eff']))\\\n",
    "                .withColumn(\"proc_kfactor\", lit(samples[s]['kfactor']))\\\n",
    "                .withColumn(\"proc_weight\", lit(samples[s]['weight']))\\\n",
    "                .withColumn(\"proc_sample\", lit(s))\n",
    "    nev = tempDF.count()\n",
    "    if samples[s]['xsec'] == None: # Data       \n",
    "        DFList.append([tempDF.withColumn(\"proc_nevents\",lit(nev))])        \n",
    "    else: # MC\n",
    "        samples[s]['weight'] = ( samples[s]['xsec'] * samples[s]['eff'] * samples[s]['kfactor'] * (LUMI / nev))\n",
    "        DFList.append([tempDF.withColumn(\"proc_nevents\",lit(nev))])        \n",
    "\n",
    "print \"Number of element:\"\n",
    "print \"DFList size =\", len(DFList)\n",
    "print \"Number of columns:\"\n",
    "print \"DFList[0][0] =\", len(DFList[0][0].columns)\n",
    "print \"DFList[0][0] =\", len(DFList[1][0].columns)\n",
    "print \"DFList[1][0] =\", len(DFList[2][0].columns)\n",
    "print \"DFList[2][0] =\", len(DFList[3][0].columns)\n",
    "print \"DFList[3][0] =\", len(DFList[4][0].columns)\n",
    "print \"DFList[4][0] =\", len(DFList[5][0].columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reduces columns and Concatenates dataframes\n",
    "\n",
    "Select interested attributes (kinematics of a physics object), and concatenate all dataframes into one gaint dataframe for ease of interpretation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Partitions:  6\n"
     ]
    }
   ],
   "source": [
    "header = list(DFList[0][0].columns) #infer the column header from one of the dataframe\n",
    "globalAttr = [g for g in header if 'proc_' in g]\n",
    "muonAttr = [muon for muon in header if 'Muon' in muon and 'genPart' not in muon]\n",
    "jetAttr = [jet for jet in header if 'Jet' in jet and 'HLT' not in jet]\n",
    "electronAttr = [electron for electron in header if 'Electron' in electron]\n",
    "Attributes = muonAttr + globalAttr\n",
    "DF = None\n",
    "if DFList:\n",
    "    DF = DFList[0][0].select(Attributes)\n",
    "    for x in DFList[1:]:\n",
    "        DF = DF.union(x[0].select(Attributes))\n",
    "    print \"Partitions:  %d\" % DF.rdd.getNumPartitions()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inspect gaint dataframe\n",
      "+-----------+-----------+---------+------------+\n",
      "|proc_sample|proc_weight|proc_xsec|proc_nevents|\n",
      "+-----------+-----------+---------+------------+\n",
      "|         TT|        1.0|   831.76|      382077|\n",
      "|         TT|        1.0|   831.76|      382077|\n",
      "|         TT|        1.0|   831.76|      382077|\n",
      "+-----------+-----------+---------+------------+\n",
      "only showing top 3 rows\n",
      "\n",
      "+-----------+-----------+---------+------------+\n",
      "|proc_sample|proc_weight|proc_xsec|proc_nevents|\n",
      "+-----------+-----------+---------+------------+\n",
      "| DYJetsToLL|        1.0|   5765.4|      549790|\n",
      "| DYJetsToLL|        1.0|   5765.4|      549790|\n",
      "| DYJetsToLL|        1.0|   5765.4|      549790|\n",
      "+-----------+-----------+---------+------------+\n",
      "only showing top 3 rows\n",
      "\n",
      "+-----------+-----------+---------+------------+\n",
      "|proc_sample|proc_weight|proc_xsec|proc_nevents|\n",
      "+-----------+-----------+---------+------------+\n",
      "|         ZZ|        1.0|     16.6|      480144|\n",
      "|         ZZ|        1.0|     16.6|      480144|\n",
      "|         ZZ|        1.0|     16.6|      480144|\n",
      "+-----------+-----------+---------+------------+\n",
      "only showing top 3 rows\n",
      "\n",
      "+-----------+-----------+---------+------------+\n",
      "|proc_sample|proc_weight|proc_xsec|proc_nevents|\n",
      "+-----------+-----------+---------+------------+\n",
      "|         WZ|        1.0|     47.2|      424048|\n",
      "|         WZ|        1.0|     47.2|      424048|\n",
      "|         WZ|        1.0|     47.2|      424048|\n",
      "+-----------+-----------+---------+------------+\n",
      "only showing top 3 rows\n",
      "\n",
      "+-----------+-----------+---------+------------+\n",
      "|proc_sample|proc_weight|proc_xsec|proc_nevents|\n",
      "+-----------+-----------+---------+------------+\n",
      "|         WW|        1.0|    118.7|      348954|\n",
      "|         WW|        1.0|    118.7|      348954|\n",
      "|         WW|        1.0|    118.7|      348954|\n",
      "+-----------+-----------+---------+------------+\n",
      "only showing top 3 rows\n",
      "\n",
      "+-----------+-----------+---------+------------+\n",
      "|proc_sample|proc_weight|proc_xsec|proc_nevents|\n",
      "+-----------+-----------+---------+------------+\n",
      "| SingleMuon|        1.0|     null|      477186|\n",
      "| SingleMuon|        1.0|     null|      477186|\n",
      "| SingleMuon|        1.0|     null|      477186|\n",
      "+-----------+-----------+---------+------------+\n",
      "only showing top 3 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print \"Inspect gaint dataframe\"\n",
    "for i in [\"TT\",\"DYJetsToLL\",\"ZZ\",\"WZ\",\"WW\",\"SingleMuon\"]:\n",
    "    DF.select('proc_sample','proc_weight','proc_xsec','proc_nevents').filter(DF.proc_sample.like(\"%s\"%i)).show(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# User Defined Function Definition and helper function\n",
    "\n",
    "User Defined Function (UDF) is a function passes row by row to compute derived quantity such as invaraint mass of two physics objects involving multiple column. Helper function defined to facilitate event selection (using spark feature such as filter instead of python defined function (thanks to NANOAOD structure!))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.types import *\n",
    "\n",
    "schema = StructType([\n",
    "    StructField(\"pass\", BooleanType(), False), # True if filled / False if default(empty) \n",
    "    StructField(\"ss\", BooleanType(), False), # same sign?\n",
    "    StructField(\"mass\", FloatType(), False), # Z mass\n",
    "    StructField(\"pt\", FloatType(), False), # Z pt\n",
    "    StructField(\"eta\", FloatType(), False), # Z eta\n",
    "    StructField(\"phi\", FloatType(), False), # Z phi\n",
    "    StructField(\"dPhi\", FloatType(), False), # mu1,mu2 DeltaPhi\n",
    "    StructField(\"dR\", FloatType(), False), # mu1,mu2 DeltaR\n",
    "    StructField(\"dEta\", FloatType(), False), # mu1,mu2 DeltaEta\n",
    "    \n",
    "    StructField(\"mu1_pt\", FloatType(), False), # leading mu pT \n",
    "    StructField(\"mu2_pt\", FloatType(), False), # sub-leading mu pT \n",
    "    StructField(\"mu1_eta\", FloatType(), False), # leading mu eta\n",
    "    StructField(\"mu2_eta\", FloatType(), False), # sub-leading mu eta\n",
    "    StructField(\"mu1_phi\", FloatType(), False), # leading mu phi\n",
    "    StructField(\"mu2_phi\", FloatType(), False), # sub-leading mu phi\n",
    "    #StructField(\"mu1_charge\", FloatType(), False), # leading mu charge \n",
    "    #StructField(\"mu2_charge\", FloatType(), False), # sub-leading mu charge\n",
    "])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Jacopo's helper functions\n",
    "def deltaPhi(phi1,phi2):\n",
    "    ## Catch if being called with two objects\n",
    "    if type(phi1) != float and type(phi1) != int:\n",
    "        phi1 = phi1.phi\n",
    "    if type(phi2) != float and type(phi2) != int:\n",
    "        phi2 = phi2.phi\n",
    "    ## Otherwise\n",
    "    dphi = (phi1-phi2)\n",
    "    while dphi >  pi: dphi -= 2*pi\n",
    "    while dphi < -pi: dphi += 2*pi\n",
    "    return dphi\n",
    "\n",
    "def deltaR(eta1,phi1,eta2=None,phi2=None):\n",
    "    ## catch if called with objects\n",
    "    if eta2 == None:\n",
    "        return deltaR(eta1.eta,eta1.phi,phi1.eta,phi1.phi)\n",
    "    ## otherwise\n",
    "    return hypot(eta1-eta2, deltaPhi(phi1,phi2))\n",
    "\n",
    "def invMass(pt1,pt2, eta1, eta2, phi1, phi2, mass1, mass2):\n",
    "    \n",
    "        theta1 = 2.0*atan(exp(-eta1))\n",
    "        px1 = pt1 * cos(phi1)\n",
    "        py1 = pt1 * sin(phi1)\n",
    "        pz1 = pt1 / tan(theta1)\n",
    "        E1 = sqrt(px1**2 + py1**2 + pz1**2 + mass1**2)\n",
    "        \n",
    "        theta2 = 2.0*atan(exp(-eta2))\n",
    "        px2 = pt2 * cos(phi2)\n",
    "        py2 = pt2 * sin(phi2)\n",
    "        pz2 = pt2 / tan(theta2)\n",
    "        E2 = sqrt(px2**2 + py2**2 + pz2**2 + mass2**2)\n",
    "        \n",
    "        themass = sqrt((E1 + E2)**2 - (px1 + px2)**2 - (py1 + py2)**2 - (pz1 + pz2)**2)\n",
    "        thept = sqrt((px1 + px2)**2 + (py1 + py2)**2)\n",
    "        thetheta = atan( thept / (pz1 + pz2) )        \n",
    "        theeta = 0.5*log( (sqrt((px1 + px2)**2 + (py1 + py2)**2 + (pz1 + pz2)**2)+(pz1 + pz2))/(sqrt((px1 + px2)**2 + (py1 + py2)**2 + (pz1 + pz2)**2)-(pz1 + pz2)) )\n",
    "        thephi = asin((py1 + py2)/thept)\n",
    "        \n",
    "        delPhi = deltaPhi(phi1,phi2)\n",
    "        delR = deltaR(eta1,phi1,eta2,phi2)\n",
    "        delEta = eta1-eta2\n",
    "        \n",
    "        return (\n",
    "                themass, thept, theeta, thephi,\n",
    "                delPhi, delR, delEta\n",
    "               )\n",
    "#########\n",
    "\n",
    "#https://test-cms-nanoaod-integration.web.cern.ch/integration/master/mc80X_doc.html\n",
    "from math import *\n",
    "def invariantMass(pt,eta,phi,mass,charge,tightid):\n",
    "    MLL = ( False,False,0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 , 0.0, 0.0,  0.0, 0.0,  0.0, 0.0 )\n",
    "    \n",
    "    \"\"\"\n",
    "    Perform muon selection:\n",
    "      pT > 30\n",
    "      abs(eta) < 2.1\n",
    "      80GeV < DiMuonInvMass < 100GeV\n",
    "      tightId muon\n",
    "      N(mu) == 2\n",
    "    \"\"\"\n",
    "    \n",
    "    if len(pt) < 2:\n",
    "        return MLL\n",
    "    \n",
    "    #Identify muon candidate\n",
    "    leadingIdx = None\n",
    "    trailingIdx = None\n",
    " \n",
    "    for idx in range(len(pt)):\n",
    "        if leadingIdx == None:\n",
    "            if pt[idx] > 30 and abs(eta[idx]) < 2.1 and tightid[idx]:\n",
    "                leadingIdx = idx\n",
    "        elif trailingIdx == None:\n",
    "            if pt[idx] > 30 and abs(eta[idx]) < 2.1 and tightid[idx]:\n",
    "                trailingIdx = idx\n",
    "        else:\n",
    "            if pt[idx] > 30 and abs(eta[idx]) < 2.1 and tightid[idx]:\n",
    "                return MLL\n",
    "            \n",
    "    # Once we have the candidate, compute Invariant mass\n",
    "    if leadingIdx != None and trailingIdx != None:\n",
    "        #print \"Candidate found\"\n",
    "        sameSign= ((charge[leadingIdx] * charge[trailingIdx]) > 0)\n",
    "        dimuInvMass = (True, sameSign,) + \\\n",
    "        invMass(pt[leadingIdx], pt[trailingIdx],\n",
    "                                eta[leadingIdx], eta[trailingIdx],\n",
    "                                phi[leadingIdx], phi[trailingIdx],\n",
    "                                mass[leadingIdx], mass[trailingIdx])\n",
    "        #append candidate kinematics's\n",
    "        dimuInvMass = dimuInvMass+ ( pt[leadingIdx], pt[trailingIdx],\n",
    "                                     eta[leadingIdx], eta[trailingIdx],\n",
    "                                     phi[leadingIdx], phi[trailingIdx]\n",
    "                                   )\n",
    "        return dimuInvMass\n",
    "    else:\n",
    "        # if N(mu) != 2\n",
    "        return MLL\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test the UDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(68.12910614340414, 112.5281196536678, 0.981381526514867, -0.701345090742176, -1.2114648, 1.2430088934990942, -0.27825195)\n",
      "(False, False, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)\n"
     ]
    }
   ],
   "source": [
    "#test helper function\n",
    "#invMass(pt1,pt2, eta1, eta2, phi1, phi2, mass1, mass2)\n",
    "print invMass(95.00175, 35.536867,0.79174805, 1.07,-1.0014648, 0.21,0.4, 0.4)\n",
    "\n",
    "#invariantMass(pt,eta,phi,mass,charge,tightid)\n",
    "print invariantMass([95.00175],[0.79174805, 1.07],[-1.0014648, 0.21],[0.4, 0.4],[-1, 1],[\"true\", \"true\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computing the invaraince mass of two lepton \n",
    "\n",
    "Deploying the UDF to take kinemtaics variables of muons, constructing TLorenztVecotr on the fly and append a new column to reduced table (Muon object selection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----------+--------------------+-----------+\n",
      "|proc_sample|           Zparticle|proc_weight|\n",
      "+-----------+--------------------+-----------+\n",
      "|         TT|[true,false,144.8...|        1.0|\n",
      "|         TT|[true,false,71.81...|        1.0|\n",
      "|         TT|[true,true,108.34...|        1.0|\n",
      "+-----------+--------------------+-----------+\n",
      "only showing top 3 rows\n",
      "\n",
      "+-----------+--------------------+-----------+\n",
      "|proc_sample|           Zparticle|proc_weight|\n",
      "+-----------+--------------------+-----------+\n",
      "| DYJetsToLL|[true,false,92.16...|        1.0|\n",
      "| DYJetsToLL|[true,false,78.37...|        1.0|\n",
      "| DYJetsToLL|[true,false,88.75...|        1.0|\n",
      "+-----------+--------------------+-----------+\n",
      "only showing top 3 rows\n",
      "\n",
      "+-----------+--------------------+-----------+\n",
      "|proc_sample|           Zparticle|proc_weight|\n",
      "+-----------+--------------------+-----------+\n",
      "|         ZZ|[true,false,80.83...|        1.0|\n",
      "|         ZZ|[true,false,102.3...|        1.0|\n",
      "|         ZZ|[true,false,90.68...|        1.0|\n",
      "+-----------+--------------------+-----------+\n",
      "only showing top 3 rows\n",
      "\n",
      "+-----------+--------------------+-----------+\n",
      "|proc_sample|           Zparticle|proc_weight|\n",
      "+-----------+--------------------+-----------+\n",
      "|         WZ|[true,false,90.03...|        1.0|\n",
      "|         WZ|[true,false,88.58...|        1.0|\n",
      "|         WZ|[true,false,89.61...|        1.0|\n",
      "+-----------+--------------------+-----------+\n",
      "only showing top 3 rows\n",
      "\n",
      "+-----------+--------------------+-----------+\n",
      "|proc_sample|           Zparticle|proc_weight|\n",
      "+-----------+--------------------+-----------+\n",
      "|         WW|[true,false,251.5...|        1.0|\n",
      "|         WW|[true,false,152.2...|        1.0|\n",
      "|         WW|[true,false,105.0...|        1.0|\n",
      "+-----------+--------------------+-----------+\n",
      "only showing top 3 rows\n",
      "\n",
      "+-----------+--------------------+-----------+\n",
      "|proc_sample|           Zparticle|proc_weight|\n",
      "+-----------+--------------------+-----------+\n",
      "| SingleMuon|[true,false,89.35...|        1.0|\n",
      "| SingleMuon|[true,false,92.42...|        1.0|\n",
      "| SingleMuon|[true,false,82.21...|        1.0|\n",
      "+-----------+--------------------+-----------+\n",
      "only showing top 3 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from pyspark.sql.functions import udf\n",
    "MllUDF = udf(invariantMass,schema)\n",
    "\n",
    "DF = DF.withColumn('Zparticle',MllUDF(\"Muon_pt\",\"Muon_eta\",\"Muon_phi\",\"Muon_mass\",\"Muon_charge\",\"Muon_tightId\"))\n",
    "\n",
    "for i in [\"TT\",\"DYJetsToLL\",\"ZZ\",\"WZ\",\"WW\",\"SingleMuon\"]:\n",
    "    DF.select('proc_sample','Zparticle','proc_weight').filter(DF.proc_sample.like(\"%s\"%i)).filter(\"Zparticle.pass\").show(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----------+--------------------+-----------+\n",
      "|proc_sample|           Zparticle|proc_weight|\n",
      "+-----------+--------------------+-----------+\n",
      "|         TT|[true,false,144.8...|        1.0|\n",
      "|         TT|[true,false,71.81...|        1.0|\n",
      "|         TT|[true,true,108.34...|        1.0|\n",
      "+-----------+--------------------+-----------+\n",
      "only showing top 3 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "DF = DF.where(\"Zparticle.pass == True\")\n",
    "DF = DF.cache()\n",
    "\n",
    "for i in [\"TT\"]:\n",
    "    DF.select('proc_sample','Zparticle','proc_weight').filter(DF.proc_sample.like(\"%s\" %i)).show(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+---------+--------------------+\n",
      "|   mu1_pt|   mu2_pt|             Muon_pt|\n",
      "+---------+---------+--------------------+\n",
      "|135.60971| 49.35242|[135.60971, 49.35...|\n",
      "| 86.59694|30.364235|[86.59694, 30.364...|\n",
      "| 38.90641|33.944725|[38.90641, 33.944...|\n",
      "+---------+---------+--------------------+\n",
      "only showing top 3 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "DF.select('Zparticle.mu1_pt','Zparticle.mu2_pt','Muon_pt').filter(DF.proc_sample.like(\"TT\")).show(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting the Zpeak mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load libraries, amp up passMuDF with histogrammar\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import histogrammar as hg\n",
    "import histogrammar.sparksql\n",
    "hg.sparksql.addMethods(DF)\n",
    "\n",
    "# Monkey-patch histogrammar to fix a bug\n",
    "def x(self, jvm, converter):\n",
    "    return converter.Select(self.quantity.asSparkSQL(), self.cut._sparksql(jvm, converter))\n",
    "hg.Select._sparksql = x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "plots = hg.UntypedLabel(\n",
    "    # Count events (both raw and weighted)\n",
    "    nRawEvents = hg.Sum(lit(1.0)),\n",
    "    nEvents = hg.Sum(DF['proc_weight']),\n",
    "    # Make some regular \"histograms\" of muon pT\n",
    "    LeadPt = hg.Bin(50, 0, 150, DF['Zparticle.mu1_pt'], hg.Count(DF['proc_weight'])),\n",
    "    SecondPt = hg.Bin(50, 0, 150, DF['Zparticle.mu2_pt'], hg.Count(DF['proc_weight'])),\n",
    "    # Let's make an additional cut for only same-sign events\n",
    "    ssCut = hg.Select(\n",
    "        DF['Zparticle.ss'] == True,\n",
    "        # .. and count them\n",
    "        hg.Sum(DF['proc_weight']),\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make a set of histograms per-sample\n",
    "bulkHisto = hg.Categorize(quantity = DF[\"proc_sample\"], value = plots)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "'JavaPackage' object is not callable",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-14-8f64dc8a1f43>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# Fill from spark -- this is needed if you do a hg.Categorize (static) instead of a passMuDF.Categorize (the class method)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mbulkHisto\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfillsparksql\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdf\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mDF\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/usr/lib/python2.7/site-packages/histogrammar/defs.pyc\u001b[0m in \u001b[0;36mfillsparksql\u001b[0;34m(self, df)\u001b[0m\n\u001b[1;32m   1096\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1097\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mfillsparksql\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdf\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1098\u001b[0;31m         \u001b[0mconverter\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_sc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_jvm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0morg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdianahep\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhistogrammar\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msparksql\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpyspark\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mAggregatorConverter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1099\u001b[0m         \u001b[0magg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_sparksql\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_sc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_jvm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconverter\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1100\u001b[0m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconverter\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhistogrammar\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_jdf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0magg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: 'JavaPackage' object is not callable"
     ]
    }
   ],
   "source": [
    "# Fill from spark -- this is needed if you do a hg.Categorize (static) instead of a passMuDF.Categorize (the class method)\n",
    "bulkHisto.fillsparksql(df=DF)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}