{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "2664c5a5", "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": "fad54964", "metadata": {}, "outputs": [], "source": [ "import xgboost as xgb" ] }, { "cell_type": "code", "execution_count": 3, "id": "06140c21", "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": 4, "id": "a356a240", "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": 5, "id": "2ae41356", "metadata": {}, "outputs": [], "source": [ "@nb.njit\n", "def match_fjet_to_jet(fjets, jets, builder, FJET_DR = 0.8):\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": 6, "id": "d93ef2b6", "metadata": {}, "outputs": [], "source": [ "def to_np_array(ak_array, axis=-1, max_n=10, pad=0):\n", " return ak.fill_none(ak.pad_none(ak_array, max_n, clip=True, axis=axis), pad, axis=axis).to_numpy()" ] }, { "cell_type": "markdown", "id": "42e8de48", "metadata": {}, "source": [ "### BDT WP by background misidentification rate\n", "Tight: 0.3%\n", "\n", "Medium: 1%\n", "\n", "Loose: 2%" ] }, { "cell_type": "code", "execution_count": 7, "id": "f2ad1605", "metadata": {}, "outputs": [], "source": [ "WP_tight = 0.95626426\n", "WP_medium = 0.93498826\n", "WP_loose = 0.911348" ] }, { "cell_type": "code", "execution_count": 8, "id": "d7cceadc", "metadata": {}, "outputs": [], "source": [ "WP = WP_loose\n", "pred_file = \"//Users/billyli/UCSD/hhh/reports/bv2/pred_baseline_bdt_loose.h5\"" ] }, { "cell_type": "code", "execution_count": 9, "id": "6d2502b7", "metadata": {}, "outputs": [], "source": [ "test_file = \"//Users/billyli/UCSD/hhh/reports/bv2/hhh_test.h5\"" ] }, { "cell_type": "code", "execution_count": 10, "id": "ea6b8317", "metadata": {}, "outputs": [], "source": [ "in_file = h5py.File(test_file)" ] }, { "cell_type": "code", "execution_count": 11, "id": "401061e3", "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": "1632da1f", "metadata": {}, "outputs": [], "source": [ "# preliminary\n", "N_JETS = 10\n", "HIGGS_MASS = 125" ] }, { "cell_type": "markdown", "id": "92c92853", "metadata": {}, "source": [ "### Reconstruct boosted H" ] }, { "cell_type": "code", "execution_count": 13, "id": "1820e38a", "metadata": {}, "outputs": [], "source": [ "def get_test_XY(file):\n", " bh1 = file[\"TARGETS\"][\"bh1\"][\"bb\"][:]\n", " bh2 = file[\"TARGETS\"][\"bh2\"][\"bb\"][:]\n", " bh3 = file[\"TARGETS\"][\"bh3\"][\"bb\"][:]\n", "\n", " mask_fj1_bh1 = (bh1 == 0).astype(float)\n", " mask_fj1_bh2 = (bh2 == 0).astype(float)\n", " mask_fj1_bh3 = (bh3 == 0).astype(float)\n", " mask_fj1 = mask_fj1_bh1 + mask_fj1_bh2 + mask_fj1_bh3\n", "\n", " mask_fj2_bh1 = (bh1 == 1).astype(float)\n", " mask_fj2_bh2 = (bh2 == 1).astype(float)\n", " mask_fj2_bh3 = (bh3 == 1).astype(float)\n", " mask_fj2 = mask_fj2_bh1 + mask_fj2_bh2 + mask_fj2_bh3\n", "\n", " mask_fj3_bh1 = (bh1 == 2).astype(float)\n", " mask_fj3_bh2 = (bh2 == 2).astype(float)\n", " mask_fj3_bh3 = (bh3 == 2).astype(float)\n", " mask_fj3 = mask_fj3_bh1 + mask_fj3_bh2 + mask_fj3_bh3\n", "\n", " mask_signal = np.stack([mask_fj1, mask_fj2, mask_fj3], axis=1).flatten()\n", "\n", " # get zero mask\n", " mask_zero = file[\"INPUTS\"][\"BoostedJets\"][\"MASK\"][:].astype(float).flatten()\n", "\n", " feature_names = [\n", " \"fj_pt\",\n", " \"fj_eta\",\n", " # \"fj_phi\",\n", " \"fj_mass\",\n", " \"fj_sdmass\",\n", " # \"fj_charge\",\n", " \"fj_chargedenergyfrac\",\n", " \"fj_ncharged\",\n", " # \"fj_neutralenergyfrac\",\n", " \"fj_nneutral\",\n", " \"fj_tau21\",\n", " \"fj_tau32\",\n", " ]\n", " arrays = []\n", " for key in feature_names:\n", " feature = file[\"INPUTS\"][\"BoostedJets\"][key][:].astype(float).flatten()\n", " arrays.append(feature)\n", " data = np.stack(arrays, axis=1)\n", " labels = mask_signal.astype(bool)\n", " print(data.shape)\n", " print(labels.shape)\n", " return data, labels" ] }, { "cell_type": "code", "execution_count": 14, "id": "ad91ad5a", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "(128487, 9)\n", "(128487,)\n" ] } ], "source": [ "hhh_data, hhh_labels = get_test_XY(in_file)" ] }, { "cell_type": "code", "execution_count": 15, "id": "30133d49", "metadata": {}, "outputs": [], "source": [ "# reconstruct BDT dataset\n", "feature_names = [\n", " \"fj_pt\",\n", " \"fj_eta\",\n", " # \"fj_phi\",\n", " \"fj_mass\",\n", " \"fj_sdmass\",\n", " # \"fj_charge\",\n", " \"fj_chargedenergyfrac\",\n", " \"fj_ncharged\",\n", " # \"fj_neutralenergyfrac\",\n", " \"fj_nneutral\",\n", " \"fj_tau21\",\n", " \"fj_tau32\",\n", " ]\n", "test = xgb.DMatrix(data=hhh_data, label=hhh_labels, feature_names=feature_names)" ] }, { "cell_type": "code", "execution_count": 16, "id": "9133ec8b", "metadata": {}, "outputs": [], "source": [ "# load model\n", "param = {}\n", "\n", "param[\"seed\"] = 42 # set seed for reproducibility\n", "\n", "# Booster parameters\n", "param[\"eta\"] = 0.1 # learning rate\n", "param[\"max_depth\"] = 5 # maximum depth of a tree\n", "# param[\"subsample\"] = 0.8 # fraction of events to train tree on\n", "# param[\"colsample_bytree\"] = 0.8 # fraction of features to train tree on\n", "\n", "# Learning task parameters\n", "# param[\"scale_pos_weight\"] = scale_pos_weight\n", "param[\"objective\"] = \"binary:logistic\" # objective function\n", "param[\n", " \"eval_metric\"\n", "] = \"error\" # evaluation metric for cross validation, note: last one is used for early stopping\n", "param = list(param.items())\n", "\n", "num_trees = 150 # number of trees to make\n", "booster = xgb.Booster(param, model_file=\"bdt.json\")" ] }, { "cell_type": "code", "execution_count": 17, "id": "60993088", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "(42829, 3)\n" ] } ], "source": [ "# predict data\n", "pred_label = (booster.predict(test) > WP).reshape(-1, 3)\n", "print(pred_label.shape)" ] }, { "cell_type": "code", "execution_count": 18, "id": "118d59dc", "metadata": {}, "outputs": [], "source": [ "# load jets and fat jets from test h5 file\n", "js = load_jets(in_file)\n", "js_idx = ak.local_index(js)\n", "fjs = load_fjets(in_file)\n", "fj_idx = ak.local_index(fjs)\n", "\n", "# select real fjets based on pT and mass cut\n", "fj_mask = fjs['mask']\n", "fj_cond = pred_label & fj_mask\n", "fjs_selected = fjs[fj_cond]\n", "\n", "# save the qualified fjets indices\n", "# they will be bH candidates\n", "bh_fj_idx = fj_idx[fj_cond]\n", "bh_fj_idx = to_np_array(bh_fj_idx, max_n=3, pad=-1)\n", "\n", "# convert indices to AP and DP\n", "bhs_dp = np.zeros(shape=bh_fj_idx.shape)\n", "fjs_ap = np.zeros(shape=bh_fj_idx.shape)\n", "bhs_dp[bh_fj_idx!=-1] = 1\n", "fjs_ap[bh_fj_idx!=-1] = 1" ] }, { "cell_type": "markdown", "id": "9f12c161", "metadata": {}, "source": [ "### Select un_padded jets" ] }, { "cell_type": "code", "execution_count": 19, "id": "e3d30de8", "metadata": {}, "outputs": [], "source": [ "# find ak4jets that matched to selected ak8jets (dR check)\n", "not_padded = js['mask']\n", "j_cond = not_padded\n", "js_selected = js[j_cond]" ] }, { "cell_type": "markdown", "id": "cedbcf9e", "metadata": {}, "source": [ "### Reconstruct resolved H" ] }, { "cell_type": "code", "execution_count": 20, "id": "8eda8e8b", "metadata": {}, "outputs": [], "source": [ "# calculate how many resolved Higgs should be reconstructed\n", "# this number is limitied by how many jets you have after overlapping removal\n", "# and how many boosted Higgs that you have reconstructed\n", "N_jet = ak.num(js_selected, axis=-1).to_numpy(allow_missing=False)\n", "N_bH = ak.num(fjs_selected, axis=-1).to_numpy(allow_missing=False)\n", "# N_rH = np.minimum(np.floor(N_jet/2), 3-N_bH)\n", "N_rH = np.minimum(np.floor(N_jet/2), 3)" ] }, { "cell_type": "code", "execution_count": 21, "id": "b4f5933a", "metadata": {}, "outputs": [], "source": [ "# 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 = {}\n", "for nH in range(0, 1+3):\n", " JET_ASSIGNMENTS[nH] = {}\n", " for nj in range(0, nH*2):\n", " JET_ASSIGNMENTS[nH][nj] = []\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" ] }, { "cell_type": "code", "execution_count": 22, "id": "3b38686d", "metadata": {}, "outputs": [], "source": [ "# just consider top 2*N_rH jets\n", "N_rH_max = 3\n", "event_idx = ak.local_index(N_rH)\n", "\n", "rH_b1 = np.repeat(-1*np.ones(shape=N_rH.shape).reshape(1,-1), N_rH_max, axis=0)\n", "rH_b2 = np.repeat(-1*np.ones(shape=N_rH.shape).reshape(1,-1), N_rH_max, axis=0)\n", "\n", "rH_dp = np.repeat(-1*np.ones(shape=N_rH.shape).reshape(1,-1), N_rH_max, axis=0)\n", "rH_ap = np.repeat(-1*np.ones(shape=N_rH.shape).reshape(1,-1), N_rH_max, axis=0)\n", "\n", "for i in range(1, N_rH_max+1):\n", " nj = 2*i\n", " \n", " mask_i_rH = N_rH == i\n", " event_i_rH = event_idx[mask_i_rH]\n", "\n", " mjj = (js[event_i_rH][:, JET_ASSIGNMENTS[i][nj][:,:,0]]+js[event_i_rH][:, JET_ASSIGNMENTS[i][nj][:,:,1]]).mass\n", " chi2 = ak.sum(np.square(mjj - HIGGS_MASS), axis=-1)\n", " chi2_argmin = ak.argmin(chi2, axis=-1)\n", " \n", " for j in range(0, i):\n", " rH_b1[j][event_i_rH] = JET_ASSIGNMENTS[i][nj][chi2_argmin][:, j, 0]\n", " rH_b2[j][event_i_rH] = JET_ASSIGNMENTS[i][nj][chi2_argmin][:, j, 1]\n", " rH_dp[j][event_i_rH] = 1\n", " rH_ap[j][event_i_rH] = 1\n", " \n", " for k in range(i, N_rH_max):\n", " rH_dp[k][event_i_rH] = 0\n", " rH_ap[k][event_i_rH] = 0" ] }, { "cell_type": "code", "execution_count": 23, "id": "1662a585", "metadata": {}, "outputs": [], "source": [ "# save all assignment to the h5file\n", "# boosted \n", "datasets = {}\n", "datasets[\"TARGETS/bh1/bb\"] = bh_fj_idx[:,0]+10\n", "datasets[\"TARGETS/bh2/bb\"] = bh_fj_idx[:,1]+10\n", "datasets[\"TARGETS/bh3/bb\"] = bh_fj_idx[:,2]+10\n", "\n", "datasets[\"TARGETS/bh1/detection_probability\"] = bhs_dp[:,0]\n", "datasets[\"TARGETS/bh2/detection_probability\"] = bhs_dp[:,1]\n", "datasets[\"TARGETS/bh3/detection_probability\"] = bhs_dp[:,2]\n", "\n", "datasets[\"TARGETS/bh1/assignment_probability\"] = bhs_dp[:,0]\n", "datasets[\"TARGETS/bh2/assignment_probability\"] = bhs_dp[:,1]\n", "datasets[\"TARGETS/bh3/assignment_probability\"] = bhs_dp[:,2]\n", "\n", "# resolved\n", "for i in range(1, N_rH_max+1):\n", " datasets[f\"TARGETS/h{i}/b1\"] = rH_b1[i-1]\n", " datasets[f\"TARGETS/h{i}/b2\"] = rH_b2[i-1]\n", "\n", " datasets[f\"TARGETS/h{i}/detection_probability\"] = rH_dp[i-1]\n", " datasets[f\"TARGETS/h{i}/assignment_probability\"] = rH_ap[i-1]" ] }, { "cell_type": "code", "execution_count": 24, "id": "b6b04fe7", "metadata": {}, "outputs": [], "source": [ "all_datasets = {}\n", "for dataset_name, data in datasets.items():\n", " if dataset_name not in all_datasets:\n", " all_datasets[dataset_name] = []\n", " all_datasets[dataset_name].append(data)" ] }, { "cell_type": "code", "execution_count": 25, "id": "d7607a1b", "metadata": {}, "outputs": [], "source": [ "with h5py.File(pred_file, \"w\") as output:\n", " for jet_type_name, jet_type in in_file[\"INPUTS\"].items():\n", " for feature_name, feature in jet_type.items():\n", " dataset_name = f\"INPUTS/{jet_type_name}/{feature_name}\"\n", " data = np.array(feature)\n", " output.create_dataset(dataset_name, data=data)\n", " for dataset_name, all_data in all_datasets.items():\n", " concat_data = np.concatenate(all_data, axis=0)\n", " output.create_dataset(dataset_name, data=concat_data)" ] }, { "cell_type": "code", "execution_count": 26, "id": "f620bdfc", "metadata": {}, "outputs": [], "source": [ "pred_h5 = h5py.File(pred_file)" ] }, { "cell_type": "code", "execution_count": 27, "id": "d91e4dc0", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 27, "metadata": {}, "output_type": "execute_result" } ], "source": [ "pred_h5['TARGETS']['h3'].keys()" ] }, { "cell_type": "code", "execution_count": 28, "id": "17aa80d9", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 28, "metadata": {}, "output_type": "execute_result" } ], "source": [ "pred_h5.keys()" ] }, { "cell_type": "code", "execution_count": null, "id": "a880420d", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "6e9ac5bc", "metadata": {}, "outputs": [], "source": [] } ], "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.9.10" } }, "nbformat": 4, "nbformat_minor": 5 }