{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "3f953076", "metadata": {}, "outputs": [], "source": [ "import itertools\n", "import logging\n", "from pathlib import Path\n", "import numba as nb\n", "\n", "import awkward as ak\n", "import click\n", "import h5py\n", "import numpy as np\n", "import vector\n", "\n", "from src.data.cms.convert_to_h5 import MIN_JETS, N_JETS, N_FJETS\n", "\n", "vector.register_awkward()\n", "\n", "logging.basicConfig(level=logging.INFO)" ] }, { "cell_type": "code", "execution_count": 2, "id": "9bfffafb", "metadata": {}, "outputs": [], "source": [ "pred_file = \"\"\n", "test_file = \"//home/billyli/UCSD/hhh/reports/bv2/hhh_test.h5\"" ] }, { "cell_type": "code", "execution_count": 3, "id": "8cd15fc1", "metadata": {}, "outputs": [], "source": [ "HIGGS_MASS = 125.0\n", "# precompute possible jet assignments lookup table\n", "MIN_JETS=6\n", "\n", "JET_ASSIGNMENTS = {}\n", "for nH in range(0, 1+3):\n", " JET_ASSIGNMENTS[nH] = {}\n", " for nj in range(nH*2, N_JETS + 1):\n", " a = list(itertools.combinations(range(nj), 2))\n", " b = np.array([ assignment for assignment in itertools.combinations(a, nH) if len(np.unique(assignment)) == nH*2])\n", " JET_ASSIGNMENTS[nH][nj] = b\n", "\n", "FJET_ASSIGNMENTS = {}" ] }, { "cell_type": "code", "execution_count": 4, "id": "3f2e1c98", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([[[0, 1],\n", " [2, 3],\n", " [4, 5]],\n", "\n", " [[0, 1],\n", " [2, 3],\n", " [4, 6]],\n", "\n", " [[0, 1],\n", " [2, 3],\n", " [4, 7]],\n", "\n", " ...,\n", "\n", " [[2, 7],\n", " [3, 4],\n", " [5, 6]],\n", "\n", " [[2, 7],\n", " [3, 5],\n", " [4, 6]],\n", "\n", " [[2, 7],\n", " [3, 6],\n", " [4, 5]]])" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "JET_ASSIGNMENTS[3][8]" ] }, { "cell_type": "code", "execution_count": 5, "id": "a81ba0d6", "metadata": {}, "outputs": [], "source": [ "a = list(itertools.combinations(range(3), 2))" ] }, { "cell_type": "code", "execution_count": 6, "id": "2fe0ab88", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[((0, 1), (0, 2), (1, 2))]" ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "list(itertools.combinations(a, 3))" ] }, { "cell_type": "code", "execution_count": 7, "id": "a98a8c90", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "3" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "len(set((0, 1)+(0, 2)+(1, 2)))" ] }, { "cell_type": "code", "execution_count": 8, "id": "31b33688", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "(0, 1, 0, 2, 1, 2)" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "(0, 1)+(0, 2)+(1, 2)" ] }, { "cell_type": "code", "execution_count": 9, "id": "845dde96", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "{0, 1, 2}" ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "set((0, 1)+(0, 2)+(1, 2))" ] }, { "cell_type": "code", "execution_count": 10, "id": "cd13e524", "metadata": {}, "outputs": [], "source": [ "in_file = h5py.File(test_file)" ] }, { "cell_type": "code", "execution_count": 11, "id": "167a6d65", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 11, "metadata": {}, "output_type": "execute_result" } ], "source": [ "in_file['INPUTS']['BoostedJets'].keys()" ] }, { "cell_type": "code", "execution_count": 12, "id": "6f12174f", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 12, "metadata": {}, "output_type": "execute_result" } ], "source": [ "in_file['INPUTS']['Jets'].keys()" ] }, { "cell_type": "code", "execution_count": 13, "id": "4ffe3d07", "metadata": {}, "outputs": [], "source": [ "# a function that loads jets from hhh_test.h5\n", "def load_jets(in_file):\n", " # load jets from the h5\n", " pt = ak.Array(in_file[\"INPUTS\"][\"Jets\"][\"pt\"])\n", " eta = ak.Array(in_file[\"INPUTS\"][\"Jets\"][\"eta\"])\n", " phi = ak.Array(in_file[\"INPUTS\"][\"Jets\"][\"phi\"])\n", " btag = ak.Array(in_file[\"INPUTS\"][\"Jets\"][\"btag\"])\n", " mass = ak.Array(in_file[\"INPUTS\"][\"Jets\"][\"mass\"])\n", " mask = ak.Array(in_file[\"INPUTS\"][\"Jets\"][\"MASK\"])\n", "\n", " jets = ak.zip(\n", " {\n", " \"pt\": pt,\n", " \"eta\": eta,\n", " \"phi\": phi,\n", " \"btag\": btag,\n", " \"mass\": mass,\n", " \"mask\": mask\n", " },\n", " with_name=\"Momentum4D\",\n", " )\n", " \n", " return jets" ] }, { "cell_type": "code", "execution_count": 14, "id": "718e93a4", "metadata": {}, "outputs": [], "source": [ "# a function that loads fat jets from hhh_test.h5\n", "def load_fjets(in_file):\n", " # load fatjets from h5\n", " fj_pt = ak.Array(in_file[\"INPUTS\"][\"BoostedJets\"][\"fj_pt\"])\n", " fj_eta = ak.Array(in_file[\"INPUTS\"][\"BoostedJets\"][\"fj_eta\"])\n", " fj_phi = ak.Array(in_file[\"INPUTS\"][\"BoostedJets\"][\"fj_phi\"])\n", " fj_mass = ak.Array(in_file[\"INPUTS\"][\"BoostedJets\"][\"fj_mass\"])\n", " fj_mask = ak.Array(in_file[\"INPUTS\"][\"BoostedJets\"][\"MASK\"])\n", "\n", " fjets = ak.zip(\n", " {\n", " \"pt\": fj_pt,\n", " \"eta\": fj_eta,\n", " \"phi\": fj_phi,\n", " 'mass': fj_mass,\n", " 'mask': fj_mask\n", " },\n", " with_name=\"Momentum4D\"\n", " )\n", " \n", " return fjets" ] }, { "cell_type": "code", "execution_count": 15, "id": "d36d4dc8", "metadata": {}, "outputs": [], "source": [ "fjs = load_fjets(in_file)\n", "fj_mask = fjs['mask']\n", "fjmass_cond = (fjs['mass']>110) & (fjs['mass']<140)\n", "fjpt_cond = fjs['pt']>300\n", "fj_cond = fjmass_cond & fjpt_cond & fj_mask\n", "fjs_selected = fjs[fj_cond]\n", "\n", "fj_idx = ak.local_index(fjs)\n", "bh_fj_idx = fj_idx[fj_cond]" ] }, { "cell_type": "code", "execution_count": 16, "id": "ed1d868d", "metadata": {}, "outputs": [], "source": [ "js = load_jets(in_file)\n", "js_idx = ak.local_index(js)" ] }, { "cell_type": "code", "execution_count": 17, "id": "cc406254", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
{pt: 99.1,\n",
       " eta: -0.14,\n",
       " phi: 3.02,\n",
       " btag: 1,\n",
       " mass: 27.2,\n",
       " mask: True}\n",
       "------------------\n",
       "type: Momentum4D[\n",
       "    pt: float64,\n",
       "    eta: float64,\n",
       "    phi: float64,\n",
       "    btag: float64,\n",
       "    mass: float64,\n",
       "    mask: bool\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 17, "metadata": {}, "output_type": "execute_result" } ], "source": [ "js[0][0]" ] }, { "cell_type": "code", "execution_count": 18, "id": "9bbf8f4a", "metadata": { "scrolled": false }, "outputs": [ { "data": { "text/html": [ "
[1.99]\n",
       "-----------------\n",
       "type: 1 * float64
" ], "text/plain": [ "" ] }, "execution_count": 18, "metadata": {}, "output_type": "execute_result" } ], "source": [ "fjs_selected[8].deltaR(js[0][0])" ] }, { "cell_type": "code", "execution_count": 19, "id": "e1a758c9", "metadata": {}, "outputs": [], "source": [ "FJET_DR = 0.8\n", "\n", "@nb.njit\n", "def match_fjet_to_jet(fjets, jets, builder):\n", " for fjets_event, jets_event in zip(fjets, jets):\n", " builder.begin_list()\n", " for i, jet in enumerate(jets_event):\n", " match_idx = -1\n", " for j, fjet in enumerate(fjets_event):\n", " if jet.deltaR(fjet) < FJET_DR:\n", " match_idx = j\n", " builder.append(match_idx)\n", " builder.end_list()\n", "\n", " return builder" ] }, { "cell_type": "code", "execution_count": 76, "id": "da003aac", "metadata": {}, "outputs": [], "source": [ "matched_fj_idx = match_fjet_to_jet(fjs_selected, js, ak.ArrayBuilder()).snapshot()\n", "unoverlapped = matched_fj_idx==-1\n", "unmasked = js['mask']\n", "j_cond = unoverlapped & unmasked\n", "js_selected = js[j_cond]\n", "idx_js_unoverlapped = js_idx[j_cond]" ] }, { "cell_type": "code", "execution_count": 77, "id": "0015dd31", "metadata": {}, "outputs": [], "source": [ "js_unoverlapped\n", "jcounts = ak.count(js_unoverlapped, axis=-1)" ] }, { "cell_type": "code", "execution_count": null, "id": "c187b9f0", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 141, "id": "58584575", "metadata": {}, "outputs": [], "source": [ "n_bhs_matched = ak.count(fjs_selected, axis=-1)\n", "NrHs = (3-n_bhs_matched)" ] }, { "cell_type": "code", "execution_count": 23, "id": "404e5bf3", "metadata": {}, "outputs": [], "source": [ "# jet_assignments = [JET_ASSIGNMENTS[NrH_event][jcount_event] for jcount_event, NrH_event in zip(jcounts, NrHs)]" ] }, { "cell_type": "code", "execution_count": 24, "id": "1d0d5f99", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([[0, 1],\n", " [2, 3],\n", " [4, 6]])" ] }, "execution_count": 24, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# jet_assignments[0][1]" ] }, { "cell_type": "code", "execution_count": null, "id": "9645d636", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 25, "id": "6dcbefd5", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([[[0, 1],\n", " [2, 3]],\n", "\n", " [[0, 2],\n", " [1, 3]],\n", "\n", " [[0, 3],\n", " [1, 2]]])" ] }, "execution_count": 25, "metadata": {}, "output_type": "execute_result" } ], "source": [ "a = [1, 2, 3, 4]\n", "JET_ASSIGNMENTS[2][4]" ] }, { "cell_type": "code", "execution_count": 33, "id": "399396e4", "metadata": {}, "outputs": [], "source": [ "def ja_list_2_ak(JET_ASSIGNMENTS, NrH, jcounts, builder):\n", " builder.begin_list()\n", " for NrH_e, jcounts_e in zip(NrH, jcounts):\n", " builder.begin_list()\n", " ja = JET_ASSIGNMENTS[NrH_e][jcounts_e]\n", " for comb in ja:\n", " builder.begin_list()\n", " for b1_b2 in comb:\n", " builder.begin_list()\n", " for bx in b1_b2:\n", " builder.append(bx)\n", " builder.end_list()\n", " builder.end_list()\n", " builder.end_list()\n", " \n", " builder.end_list()\n", " return builder" ] }, { "cell_type": "code", "execution_count": null, "id": "5c6a2ab1", "metadata": {}, "outputs": [], "source": [ "from numba.core import types\n", "from numba.typed import Dict\n", "\n", "JET_ASSIGNMENTS_nb = Dict.empty(\n", " key_type=types.unicode_type,\n", " value_type=types.float64[:,:,:],\n", ")\n", "\n", "for nH in range(0, 1+3):\n", " JET_ASSIGNMENTS_nb[nH] = {}\n", " for nj in range(nH*2, N_JETS + 1):\n", " a = list(itertools.combinations(range(nj), 2))\n", " b = np.array([ assignment for assignment in itertools.combinations(a, nH) if len(np.unique(assignment)) == nH*2])\n", " JET_ASSIGNMENTS[nH][nj] = b\n", "\n" ] }, { "cell_type": "code", "execution_count": 116, "id": "b7e6eb10", "metadata": {}, "outputs": [], "source": [ "JET_ASSIGNMENTS_ak = []\n", "\n", "for nH in range(0, 1+3):\n", " JET_ASSIGNMENTS_ak.append([])\n", " for nj in range(0, nH*2):\n", " JET_ASSIGNMENTS_ak[nH].append([])\n", " for nj in range(nH*2, N_JETS + 1):\n", " JET_ASSIGNMENTS_ak[nH].append([])\n", " a = list(itertools.combinations(range(nj), 2))\n", " b = np.array([ assignment for assignment in itertools.combinations(a, nH) if len(np.unique(assignment)) == nH*2])\n", " JET_ASSIGNMENTS_ak[nH][nj] = b\n", " \n", "JET_ASSIGNMENTS_ak = ak.Array(JET_ASSIGNMENTS_ak)" ] }, { "cell_type": "code", "execution_count": 111, "id": "7761d7a1", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[[[]], [[]], [[]], [[]], [[]], [[]], [[]], [[]], [[]], [[]], [[]]],\n",
       " [[], ..., [[[0, 1]], [[0, 2]], [[0, 3]], ..., [[7, ...]], [[7, 9]], [[8, 9]]]],\n",
       " [[], [], ..., [[[0, 1], [2, 3]], [[0, 1], [2, ...]], ..., [[6, 9], [7, 8]]]],\n",
       " [[], ..., [[[0, 1], [2, 3], [4, 5]], [...], ..., [[4, 9], [5, 8], [6, 7]]]]]\n",
       "--------------------------------------------------------------------------------\n",
       "type: 4 * var * var * var * var * int64
" ], "text/plain": [ "" ] }, "execution_count": 111, "metadata": {}, "output_type": "execute_result" } ], "source": [ "JET_ASSIGNMENTS_ak" ] }, { "cell_type": "code", "execution_count": 112, "id": "eb7667c1", "metadata": {}, "outputs": [], "source": [ "@nb.njit\n", "def ja_list_2_ak_2(JET_ASSIGNMENTS, NrH, jcounts, builder):\n", " builder.begin_list()\n", " for NrH_e, jcounts_e in zip(NrH, jcounts):\n", " builder.begin_list()\n", " ja = JET_ASSIGNMENTS[NrH_e][jcounts_e]\n", " for comb in ja:\n", " builder.begin_list()\n", " for b1_b2 in comb:\n", " builder.begin_list()\n", " for bx in b1_b2:\n", " builder.append(bx)\n", " builder.end_list()\n", " builder.end_list()\n", " builder.end_list()\n", " \n", " builder.end_list()\n", " return builder" ] }, { "cell_type": "code", "execution_count": 99, "id": "1e239a9c", "metadata": {}, "outputs": [], "source": [ "sufficient_jets = jcounts >= NrHs*2" ] }, { "cell_type": "code", "execution_count": 100, "id": "4178504a", "metadata": {}, "outputs": [], "source": [ "jcounts = jcounts[sufficient_jets]" ] }, { "cell_type": "code", "execution_count": 142, "id": "a4c907e1", "metadata": {}, "outputs": [], "source": [ "NrHs = NrHs[sufficient_jets]" ] }, { "cell_type": "code", "execution_count": 144, "id": "63a48516", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "1 * var * var * var * var * int64\n", "42815 * var * var * var * int64\n" ] } ], "source": [ "jet_assignments_ak = ja_list_2_ak_2(JET_ASSIGNMENTS_ak, NrHs, jcounts, ak.ArrayBuilder()).snapshot()\n", "print(jet_assignments_ak.type)\n", "jet_assignments_ak = ak.flatten(jet_assignments, axis=0)\n", "print(jet_assignments_ak.type)" ] }, { "cell_type": "code", "execution_count": 143, "id": "b0ce8c32", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "42815 * int64\n" ] } ], "source": [ "print(NrHs.type)" ] }, { "cell_type": "code", "execution_count": 65, "id": "a2dc7598", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "0\n" ] } ], "source": [ "print(jet_assignments[0][0][0][0])" ] }, { "cell_type": "code", "execution_count": 127, "id": "23c393d9", "metadata": {}, "outputs": [], "source": [ "@nb.njit\n", "def chi2_matching3(js, jcounts, NrHs, jet_assignments, builder):\n", " HIGGS_MASS = 125\n", " builder.begin_list()\n", " for js_e, jcount_e, ja_e, NrH_e in zip(js, jcounts, jet_assignments, NrHs):\n", " builder.begin_list()\n", " if NrH_e == 0:\n", " continue\n", "\n", " chi2_argmin = 0\n", " chi2_min = 99999\n", " for i, comb in enumerate(ja_e):\n", " chi2 = 0\n", " for b1_b2 in comb:\n", " j_b1_idx = b1_b2[0]\n", " j_b2_idx = b1_b2[1]\n", " \n", " j_b1 = js_e[j_b1_idx]\n", " j_b2 = js_e[j_b2_idx]\n", " mjj = (j_b1 + j_b2).mass\n", " \n", " chi2 += np.square(mjj - HIGGS_MASS)\n", " if chi2[[[[0, 1], [2, 3], [4, 5]], [[0, 1], ...], ..., [[0, 1], [2, 3], [4, 5]]]]\n", "--------------------------------------------------------------------------\n", "type: 1 * var * var * var * int64" ], "text/plain": [ "" ] }, "execution_count": 130, "metadata": {}, "output_type": "execute_result" } ], "source": [ "chi2_matching3(js_unoverlapped[0:10], jcounts[0:10], NrHs[0:10], jet_assignments[0:10], ak.ArrayBuilder()).snapshot()" ] }, { "cell_type": "code", "execution_count": null, "id": "d317e051", "metadata": {}, "outputs": [], "source": [ "def chi2_matching(js, jcounts, NrHs, jet_assignments):\n", " HIGGS_MASS = 125\n", " predictions = []\n", " for js_e, jcount_e, ja_e, NrH_e in zip(js, jcounts, jet_assignments, NrHs):\n", " if NrH_e == 0:\n", " continue\n", " \n", " j_b1 = js_e[ja_e[:,:,0]]\n", " j_b2 = js_e[ja_e[:,:,1]]\n", " mjj = (j_b1 + j_b2).mass\n", " chi2 = ak.sum(np.square(mjj - HIGGS_MASS), axis=-1)\n", " chi2_argmin = ak.argmin(chi2, axis=-1)\n", " \n", " prediction_e = [ja_e[i,:] for i in range(NrH_e)]\n", " predictions.append(prediction_e)\n", "\n", " return predictions" ] }, { "cell_type": "code", "execution_count": null, "id": "4833659e", "metadata": { "scrolled": true }, "outputs": [], "source": [ "chi2_matching(js_unoverlapped, jcounts, NrHs, jet_assignments)" ] }, { "cell_type": "code", "execution_count": null, "id": "11362fba", "metadata": {}, "outputs": [], "source": [ "def chi2_matching2(js, jcounts, NrH, JET_ASSIGNMENTS):\n", " HIGGS_MASS = 125\n", " j_b1 = []" ] }, { "cell_type": "code", "execution_count": null, "id": "66716acb", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "6a378eed", "metadata": {}, "outputs": [], "source": [ "NrHs_unique = np.unique(NrHs)\n", "jet_assignments = {}\n", "\n", "# loop through different number of resolved Higgs candidates\n", "for NrH in NrHs_unique:\n", " # filter out events that have same number of resolved Higgs candidate\n", " NrH_filter = NrHs==NrH\n", " js_NrH = js[NrH_filter]\n", " jcounts_NrH = jcounts[NrH_filter]\n", " jet_assignments_NrH = [JET_ASSIGNMENTS[NrH][jcount_NrH] for jcount_NrH in jcounts_NrH]\n", " \n", "\n", "for \n", "jet_assignments = \n", "jet_assignments = [JET_ASSIGNMENTS[NrH_event][jcount_event] for jcount_event, NrH_event in zip(jcounts, NrHs)]" ] }, { "cell_type": "code", "execution_count": null, "id": "683200fe", "metadata": {}, "outputs": [], "source": [ "#js_b1 = js[]\n", "# jet_assignments = ak.Array(jet_assignments)\n", "jet_assignments[0][1]" ] }, { "cell_type": "code", "execution_count": null, "id": "10c1e89b", "metadata": {}, "outputs": [], "source": [ "idx_js_b1 = [ pair[0] for ev in jet_assignments for comb in ev for pair in comb ]\n", "idx_js_b2 = [ pair[1] for ev in jet_assignments for comb in ev for pair in comb ]" ] }, { "cell_type": "code", "execution_count": 149, "id": "4c5390eb", "metadata": {}, "outputs": [], "source": [ "# load jets and fat jets from test h5 file\n", "in_file = h5py.File(test_file)\n", "js = load_jets(in_file)\n", "fjs = load_fjets(in_file)\n", "\n", "# select real fjets based on pT and mass cut\n", "fj_mask = fjs['mask']\n", "fjmass_cond = (fjs['mass']>110) & (fjs['mass']<140)\n", "fjpt_cond = fjs['pt']>300\n", "fj_cond = fjmass_cond & fjpt_cond & fj_mask\n", "fjs_selected = fjs[fj_cond]\n", "\n", "# save the qualified fjets indices\n", "# they will be bH candidates\n", "fj_idx = ak.local_index(fjs)\n", "bh_fj_idx = fj_idx[fj_cond]\n", "\n", "# find ak4jets that matched to selected ak8jets (dR check)\n", "matched_fj_idx = match_fjet_to_jet(fjs_selected, js, ak.ArrayBuilder()).snapshot()\n", "\n", "# remove overlapped ak4jets and padded jets\n", "unoverlapped = matched_fj_idx==-1\n", "unmasked = js['mask']\n", "j_cond = unoverlapped & unmasked\n", "js_selected = js[j_cond]\n", "idx_js_selected = js_idx[j_cond]\n", "\n", "# get the auxiliary information for chi2_matching \n", "jcounts = ak.count(js_unoverlapped, axis=-1)\n", "n_bhs_matched = ak.count(fjs_selected, axis=-1)\n", "NrHs = (3-n_bhs_matched).to_numpy()\n", "\n", "# chi2 \n", "\n", "# mask events that don't have enough resolved jets to match to 3H (X bH + Y rH)\n", "# those events will be extracted later\n", "suff = jcounts >= NrHs*2\n", "js_suff = js_selected[suff]\n", "jcounts_suff = jcounts[suff]\n", "NrHs_suff = NrHs[suff]\n", "\n", "# construct jet assignment look-up array that has \n", "# all combinations of input jets\n", "# for different numbers of resolved higgs and jets\n", "JET_ASSIGNMENTS_ak = []\n", "for nH in range(0, 1+3):\n", " JET_ASSIGNMENTS_ak.append([])\n", " for nj in range(0, nH*2):\n", " JET_ASSIGNMENTS_ak[nH].append([])\n", " for nj in range(nH*2, N_JETS + 1):\n", " JET_ASSIGNMENTS_ak[nH].append([])\n", " a = list(itertools.combinations(range(nj), 2))\n", " b = np.array([ assignment for assignment in itertools.combinations(a, nH) if len(np.unique(assignment)) == nH*2])\n", " JET_ASSIGNMENTS_ak[nH][nj] = b\n", "\n", "JET_ASSIGNMENTS_ak = ak.Array(JET_ASSIGNMENTS_ak)\n", "\n", "# find the jet assignment combinations for each event\n", "jet_assignments_ak = ja_list_2_ak_2(JET_ASSIGNMENTS_ak, NrHs, jcounts, ak.ArrayBuilder()).snapshot()\n", "jet_assignments_ak = ak.flatten(jet_assignments, axis=0)\n", "\n", "# assign ak4 jets to the resolved higgs by chi2\n", "rh_rj_idx_suff = chi2_matching3(js_suff[0:1000], jcounts_suff[0:1000], NrHs_suff[0:1000], jet_assignments_ak[0:1000], ak.ArrayBuilder()).snapshot()\n", "\n", "# For events that don't have enough jets\n", "# try reconstruct one less higgs\n", "\n", "\n", "# save all assignment to the h5file\n" ] }, { "cell_type": "code", "execution_count": 148, "id": "246017c2", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[[[[0, 1], [2, 3], [4, 5]], [[0, 1], ...], ..., [[0, 1], [2, 3], [4, 5]]]]\n" ] } ], "source": [ "print(rh_rj_idx_suff)" ] }, { "cell_type": "code", "execution_count": 122, "id": "b0efea04", "metadata": {}, "outputs": [], "source": [ "def main(test_file):\n", " # load jets and fat jets from test h5 file\n", " in_file = h5py.File(test_file)\n", " js = load_jets(in_file)\n", " fjs = load_fjets(in_file)\n", " \n", " # select real fjets based on pT and mass cut\n", " fj_mask = fjs['mask']\n", " fjmass_cond = (fjs['mass']>110) & (fjs['mass']<140)\n", " fjpt_cond = fjs['pt']>300\n", " fj_cond = fjmass_cond & fjpt_cond & fj_mask\n", " fjs_selected = fjs[fj_cond]\n", "\n", " # save the qualified fjets indices\n", " # they will be bH candidates\n", " fj_idx = ak.local_index(fjs)\n", " bh_fj_idx = fj_idx[fj_cond]\n", " \n", " # find ak4jets that matched to selected ak8jets (dR check)\n", " matched_fj_idx = match_fjet_to_jet(fjs_selected, js, ak.ArrayBuilder()).snapshot()\n", " \n", " # remove overlapped ak4jets and padded jets\n", " unoverlapped = matched_fj_idx==-1\n", " unmasked = js['mask']\n", " j_cond = unoverlapped & unmasked\n", " js_selected = js[j_cond]\n", " idx_js_selected = js_idx[j_cond]\n", " \n", " # get the auxiliary information for chi2_matching \n", " jcounts = ak.count(js_unoverlapped, axis=-1)\n", " n_bhs_matched = ak.count(fjs_selected, axis=-1)\n", " NrHs = (3-n_bhs_matched).to_numpy()\n", " \n", " # chi2 \n", " \n", " # mask events that don't have enough resolved jets to match to 3H (X bH + Y rH)\n", " # those events will be extracted later\n", " suff = jcounts >= NrHs*2\n", " js_suff = js_selected[suff]\n", " jcounts_suff = jcounts[suff]\n", " NrHs_suff = NrHs[suff]\n", " \n", " # construct jet assignment look-up array that has \n", " # all combinations of input jets\n", " # for different numbers of resolved higgs and jets\n", " JET_ASSIGNMENTS_ak = []\n", " for nH in range(0, 1+3):\n", " JET_ASSIGNMENTS_ak.append([])\n", " for nj in range(0, nH*2):\n", " JET_ASSIGNMENTS_ak[nH].append([])\n", " for nj in range(nH*2, N_JETS + 1):\n", " JET_ASSIGNMENTS_ak[nH].append([])\n", " a = list(itertools.combinations(range(nj), 2))\n", " b = np.array([ assignment for assignment in itertools.combinations(a, nH) if len(np.unique(assignment)) == nH*2])\n", " JET_ASSIGNMENTS_ak[nH][nj] = b\n", "\n", " JET_ASSIGNMENTS_ak = ak.Array(JET_ASSIGNMENTS_ak)\n", " \n", " # find the jet assignment combinations for each event\n", " jet_assignments_ak = ja_list_2_ak_2(JET_ASSIGNMENTS_ak, NrHs, jcounts, ak.ArrayBuilder()).snapshot()\n", " jet_assignments_ak = ak.flatten(jet_assignments, axis=1)\n", " \n", " # assign ak4 jets to the resolved higgs by chi2\n", " rh_rj_idx_suff = chi2_matching3(js_suff, jcounts_suff, NrHs_suff, jet_assignments, ak.ArrayBuilder()).snapshot()\n", " \n", " # For events that don't have enough jets\n", " # try reconstruct one less higgs\n", " \n", " \n", " # save all assignment to the h5file\n", " \n", " return\n", " \n", " \n", "\n", " \n", "\n", "\n", "main(test_file)\n" ] }, { "cell_type": "code", "execution_count": 123, "id": "b8606811", "metadata": {}, "outputs": [ { "ename": "NameError", "evalue": "name 'jet_assignments_ak' is not defined", "output_type": "error", "traceback": [ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", "Cell \u001b[0;32mIn[123], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mjet_assignments_ak\u001b[49m[\u001b[38;5;241m3\u001b[39m][\u001b[38;5;241m2\u001b[39m]\n", "\u001b[0;31mNameError\u001b[0m: name 'jet_assignments_ak' is not defined" ] } ], "source": [ "jet_assignments_ak[3][2]" ] }, { "cell_type": "markdown", "id": "6bb99e9f", "metadata": {}, "source": [ "## code segments waiting" ] }, { "cell_type": "code", "execution_count": null, "id": "29b39ed1", "metadata": {}, "outputs": [], "source": [ "\n", " \n", " # just consider top-6 jets\n", " nj = 6\n", " mjj = (jets[:, JET_ASSIGNMENTS[nj][:, :, 0]] + jets[:, JET_ASSIGNMENTS[nj][:, :, 1]]).mass\n", " chi2 = ak.sum(np.square(mjj - HIGGS_MASS), axis=-1)\n", " chi2_argmin = ak.argmin(chi2, axis=-1)\n", "\n", " h1_bs = np.concatenate(\n", " (\n", " np.array(in_file[\"TARGETS\"][\"h1\"][\"b1\"])[:, np.newaxis],\n", " np.array(in_file[\"TARGETS\"][\"h1\"][\"b2\"])[:, np.newaxis],\n", " ),\n", " axis=-1,\n", " )\n", " h2_bs = np.concatenate(\n", " (\n", " np.array(in_file[\"TARGETS\"][\"h2\"][\"b1\"])[:, np.newaxis],\n", " np.array(in_file[\"TARGETS\"][\"h2\"][\"b2\"])[:, np.newaxis],\n", " ),\n", " axis=-1,\n", " )\n", " h3_bs = np.concatenate(\n", " (\n", " np.array(in_file[\"TARGETS\"][\"h3\"][\"b1\"])[:, np.newaxis],\n", " np.array(in_file[\"TARGETS\"][\"h3\"][\"b2\"])[:, np.newaxis],\n", " ),\n", " axis=-1,\n", " )\n", "\n", " # chi2 on fjets to find Higgs\n", " \n", " num_events = len(fj_pt) \n", " bh1_b_pred = np.ones(shape=(num_events, 1), dtype=int)\n", " bh2_b_pred = np.ones(shape=(num_events, 1), dtype=int)*2\n", " bh3_b_pred = np.ones(shape=(num_events, 1), dtype=int)*3\n", "\n", " bh1_b = np.array(in_file[\"TARGETS\"][\"bh1\"][\"bb\"])\n", " bh2_b = np.array(in_file[\"TARGETS\"][\"bh2\"][\"bb\"])\n", " bh3_b = np.array(in_file[\"TARGETS\"][\"bh3\"][\"bb\"])\n", "\n", " targets = [h1_bs, h2_bs, h3_bs, bh1_b, bh2_b, bh3_b]\n", "\n", " masks = np.concatenate(\n", " (\n", " np.array(in_file[\"TARGETS\"][\"h1\"][\"mask\"])[np.newaxis, :],\n", " np.array(in_file[\"TARGETS\"][\"h2\"][\"mask\"])[np.newaxis, :],\n", " np.array(in_file[\"TARGETS\"][\"h3\"][\"mask\"])[np.newaxis, :],\n", " np.array(in_file[\"TARGETS\"][\"bh1\"][\"mask\"])[np.newaxis, :],\n", " np.array(in_file[\"TARGETS\"][\"bh2\"][\"mask\"])[np.newaxis, :],\n", " np.array(in_file[\"TARGETS\"][\"bh3\"][\"mask\"])[np.newaxis, :]\n", " ),\n", " axis=0,\n", " )\n", "\n", " predictions = [\n", " JET_ASSIGNMENTS[nj][chi2_argmin][:, 0, :],\n", " JET_ASSIGNMENTS[nj][chi2_argmin][:, 1, :],\n", " JET_ASSIGNMENTS[nj][chi2_argmin][:, 2, :],\n", " bh1_b_pred,\n", " bh2_b_pred,\n", " bh3_b_pred,\n", " ]\n", "\n", " num_vectors = np.sum(mask, axis=-1).to_numpy() # number of unique objects in every event\n", " lines = 2 # how many lines are generated in the table\n", " results, jet_limits, clusters = evaluate_predictions(predictions, num_vectors, targets, masks, event_file, lines)\n", " display_table(results, jet_limits, clusters)\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "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.7.12" } }, "nbformat": 4, "nbformat_minor": 5 }