[WIP] Training on AMD / ROCm

mlpf/data_cms/genjob_pu55to75.sh

@@ -75,11 +75,10 @@ ls -lrt
 echo "process.RandomNumberGeneratorService.generator.initialSeed = $SEED" >> step2_phase1_new.py
 cmsRun step2_phase1_new.py > /dev/null
 cmsRun step3_phase1_new.py > /dev/null
-#cmsRun $CMSSWDIR/src/Validation/RecoParticleFlow/test/pfanalysis_ntuple.py
 mv pfntuple.root pfntuple_${SEED}.root
-# python3 ${MLPF_PATH}/mlpf/data_cms/postprocessing2.py --input pfntuple_${SEED}.root --outpath ./
-# bzip2 -z pfntuple_${SEED}.pkl
-# cp *.pkl.bz2 $OUTDIR/$SAMPLE/raw/
+python3 ${MLPF_PATH}/mlpf/data_cms/postprocessing2.py --input pfntuple_${SEED}.root --outpath ./
+bzip2 -z pfntuple_${SEED}.pkl
+cp *.pkl.bz2 $OUTDIR/$SAMPLE/raw/
 
 #copy ROOT outputs
 #cp step2_phase1_new.root $OUTDIR/$SAMPLE/root/step2_${SEED}.root

mlpf/data_cms/postprocessing_jobs.py

@@ -33,14 +33,16 @@ def write_script(infiles, outfiles):
 
 
 samples = [
-    "/local/joosep/mlpf/cms/20240823_simcluster/pu55to75/TTbar_14TeV_TuneCUETP8M1_cfi",
+    "/local/joosep/mlpf/cms/20240823_simcluster/nopu/TTbar_14TeV_TuneCUETP8M1_cfi",
+    # "/local/joosep/mlpf/cms/20240823_simcluster/pu55to75/TTbar_14TeV_TuneCUETP8M1_cfi",
+    # "/local/joosep/mlpf/cms/20240823_simcluster/pu55to75/QCDForPF_14TeV_TuneCUETP8M1_cfi",
 ]
 
 ichunk = 1
 for sample in samples:
     infiles = list(glob.glob(f"{sample}/root/pfntuple*.root"))
     for infiles_chunk in chunks(infiles, 10):
-        outfiles_chunk = [inf.replace(".root", ".pkl.bz2").replace("/root/", "/raw2/") for inf in infiles_chunk]
+        outfiles_chunk = [inf.replace(".root", ".pkl.bz2").replace("/root/", "/raw/") for inf in infiles_chunk]
         os.makedirs(os.path.dirname(outfiles_chunk[0]), exist_ok=True)
         scr = write_script(infiles_chunk, outfiles_chunk)
         ofname = f"jobscripts/postproc_{ichunk}.sh"

mlpf/data_cms/prepare_args.py

@@ -6,15 +6,15 @@
 outdir = "/local/joosep/mlpf/cms/20240823_simcluster"
 
 samples = [
-#    ("TTbar_14TeV_TuneCUETP8M1_cfi",                           105000, 110010, "genjob_pu55to75.sh", outdir + "/pu55to75"),
+    ("TTbar_14TeV_TuneCUETP8M1_cfi",                           100000, 120010, "genjob_pu55to75.sh", outdir + "/pu55to75"),
 #    ("ZTT_All_hadronic_14TeV_TuneCUETP8M1_cfi",                200000, 220010, "genjob_pu55to75.sh", outdir + "/pu55to75"),
-    ("QCDForPF_14TeV_TuneCUETP8M1_cfi",                        300000, 305000, "genjob_pu55to75.sh", outdir + "/pu55to75"),
+    ("QCDForPF_14TeV_TuneCUETP8M1_cfi",                        300000, 320010, "genjob_pu55to75.sh", outdir + "/pu55to75"),
 #    ("SMS-T1tttt_mGl-1500_mLSP-100_TuneCP5_14TeV_pythia8_cfi", 500000, 520010, "genjob_pu55to75.sh", outdir + "/pu55to75"),
 #    ("ZpTT_1500_14TeV_TuneCP5_cfi",                            600000, 620010, "genjob_pu55to75.sh", outdir + "/pu55to75"),
-#    ("VBF_TuneCP5_14TeV_pythia8_cfi",                          700000, 720010, "genjob_pu55to75.sh", outdir + "/pu55to75"),
+#    ("VBF_TuneCP5_14TeV_pythia8_cfi",                          700000, 705010, "genjob_pu55to75.sh", outdir + "/pu55to75"),
 
 #    ("TTbar_14TeV_TuneCUETP8M1_cfi",                           702000, 705000, "genjob_nopu.sh", outdir + "/nopu"),
-#    ("MultiParticlePFGun50_cfi",                               800000, 820000, "genjob_nopu.sh", outdir + "/nopu"),
+#    ("MultiParticlePFGun50_cfi",                               800000, 805000, "genjob_nopu.sh", outdir + "/nopu"),
 #    ("VBF_TuneCP5_14TeV_pythia8_cfi",                         900000, 920010, "genjob_nopu.sh", outdir + "/nopu"),
 #    ("QCDForPF_14TeV_TuneCUETP8M1_cfi",                      1000000,1020010, "genjob_nopu.sh", outdir + "/nopu"),
 
@@ -36,6 +36,6 @@
         os.makedirs(this_outdir + "/" + samp + "/root", exist_ok=True)
 
         for seed in range(seed0, seed1):
-            p = this_outdir + "/" + samp + "/raw2/pfntuple_{}.pkl.bz2".format(seed)
+            p = this_outdir + "/" + samp + "/root/pfntuple_{}.root".format(seed)
             if not os.path.isfile(p):
                 print(f"sbatch --mem-per-cpu 8G --partition main --time 20:00:00 --cpus-per-task 1 scripts/tallinn/cmssw-el8.sh mlpf/data_cms/{script} {samp} {seed}")

mlpf/heptfds/cms_pf/qcd.py

@@ -21,7 +21,7 @@
 class CmsPfQcd(tfds.core.GeneratorBasedBuilder):
     """DatasetBuilder for cms_pf_qcd dataset."""
 
-    VERSION = tfds.core.Version("2.1.0")
+    VERSION = tfds.core.Version("2.3.0")
     RELEASE_NOTES = {
         "1.3.0": "12_2_0_pre2 generation with updated caloparticle/trackingparticle",
         "1.3.1": "Remove PS again",
@@ -33,6 +33,7 @@ class CmsPfQcd(tfds.core.GeneratorBasedBuilder):
         "1.7.1": "Increase stats to 400k events",
         "2.0.0": "New truth def based primarily on CaloParticles",
         "2.1.0": "Additional stats",
+        "2.3.0": "Split CaloParticles along tracks",
     }
     MANUAL_DOWNLOAD_INSTRUCTIONS = """
     rsync -r --progress lxplus.cern.ch:/eos/user/j/jpata/mlpf/tensorflow_datasets/cms/cms_pf_qcd ~/tensorflow_datasets/

mlpf/heptfds/cms_pf/ttbar.py

@@ -21,7 +21,7 @@
 class CmsPfTtbar(tfds.core.GeneratorBasedBuilder):
     """DatasetBuilder for cms_pf dataset."""
 
-    VERSION = tfds.core.Version("2.2.0")
+    VERSION = tfds.core.Version("2.3.0")
     RELEASE_NOTES = {
         "1.0.0": "Initial release.",
         "1.1.0": "Add muon type, fix electron GSF association",
@@ -38,6 +38,7 @@ class CmsPfTtbar(tfds.core.GeneratorBasedBuilder):
         "2.0.0": "New truth def based primarily on CaloParticles",
         "2.1.0": "Additional stats",
         "2.2.0": "Split CaloParticles along tracks",
+        "2.3.0": "Increase stats",
     }
     MANUAL_DOWNLOAD_INSTRUCTIONS = """
     rsync -r --progress lxplus.cern.ch:/eos/user/j/jpata/mlpf/tensorflow_datasets/cms/cms_pf_ttbar ~/tensorflow_datasets/

mlpf/jet_utils.py

@@ -1,24 +1,24 @@
 import numpy as np
 
-import numba
+# import numba
 import awkward
 import vector
 
 
-@numba.njit
+# @numba.njit
 def deltaphi(phi1, phi2):
     diff = phi1 - phi2
     return np.arctan2(np.sin(diff), np.cos(diff))
 
 
-@numba.njit
+# @numba.njit
 def deltar(eta1, phi1, eta2, phi2):
     deta = eta1 - eta2
     dphi = deltaphi(phi1, phi2)
     return np.sqrt(deta**2 + dphi**2)
 
 
-@numba.njit
+# @numba.njit
 def match_jets(jets1, jets2, deltaR_cut):
     iev = len(jets1)
     jet_inds_1_ev = []

mlpf/plotting/plot_utils.py

@@ -530,6 +530,16 @@ def plot_jets(yvals, epoch=None, cp_dir=None, comet_experiment=None, title=None,
 
     plt.figure()
     b = np.linspace(0, 1000, 100)
+
+    pt = awkward.to_numpy(awkward.flatten(yvals["jets_target_pt"]))
+    plt.hist(
+        pt,
+        bins=b,
+        histtype="step",
+        lw=2,
+        label="Target",
+    )
+
     pt = awkward.to_numpy(awkward.flatten(yvals["jets_cand_pt"]))
     plt.hist(
         pt,
@@ -580,6 +590,67 @@ def plot_jets(yvals, epoch=None, cp_dir=None, comet_experiment=None, title=None,
     )
 
 
+
+    plt.figure()
+    b = np.linspace(-5, 5, 100)
+    eta = awkward.to_numpy(awkward.flatten(yvals["jets_target_eta"]))
+    plt.hist(
+        eta,
+        bins=b,
+        histtype="step",
+        lw=2,
+        label="Target",
+    )
+
+    eta = awkward.to_numpy(awkward.flatten(yvals["jets_cand_eta"]))
+    plt.hist(
+        eta,
+        bins=b,
+        histtype="step",
+        lw=2,
+        label="PF",
+    )
+
+    eta = awkward.to_numpy(awkward.flatten(yvals["jets_pred_eta"]))
+    plt.hist(
+        eta,
+        bins=b,
+        histtype="step",
+        lw=2,
+        label="MLPF",
+    )
+
+    eta = awkward.to_numpy(awkward.flatten(yvals["jets_gen_eta"]))
+    plt.hist(
+        eta,
+        bins=b,
+        histtype="step",
+        lw=2,
+        label="Truth",
+    )
+
+    plt.xlabel("jet $\eta$")
+    plt.ylabel("Jets / bin")
+    plt.yscale("log")
+    plt.legend(loc="best")
+    if title:
+        plt.title(title)
+    ax = plt.gca()
+    ylim = ax.get_ylim()
+    ax.set_ylim(ylim[0], 10 * ylim[1])
+
+    if dataset:
+        EXPERIMENT_LABELS[dataset](ax)
+    if sample:
+        sample_label(ax, sample)
+
+    save_img(
+        "jet_eta.png",
+        epoch,
+        cp_dir=cp_dir,
+        comet_experiment=comet_experiment,
+    )
+
 def plot_jet_ratio(
     yvals,
     epoch=None,

mlpf/pyg/PFDataset.py

@@ -70,6 +70,23 @@ def __getitem__(self, item):
             ret["ygen"][:, 0][(ret["X"][:, 0] == 10) & (ret["ygen"][:, 0] == 7)] = 2
             ret["ygen"][:, 0][(ret["X"][:, 0] == 11) & (ret["ygen"][:, 0] == 7)] = 2
 
+            # set pt for HO which would otherwise be 0
+            msk_ho = ret["X"][:, 0] == 10
+            eta = ret["X"][:, 2][msk_ho]
+            e = ret["X"][:, 5][msk_ho]
+            ret["X"][:, 1][msk_ho] = np.sqrt(e**2 - (np.tanh(eta) * e) ** 2)
+
+        # transform pt -> log(pt / elem pt), same for energy
+        ret["ygen"][:, 6] = np.log(ret["ygen"][:, 6] / ret["X"][:, 5])
+        ret["ygen"][:, 6][np.isnan(ret["ygen"][:, 6])] = 0.0
+        ret["ygen"][:, 6][np.isinf(ret["ygen"][:, 6])] = 0.0
+        ret["ygen"][:, 6][ret["ygen"][:, 0] == 0] = 0
+
+        ret["ygen"][:, 2] = np.log(ret["ygen"][:, 2] / ret["X"][:, 1])
+        ret["ygen"][:, 2][np.isnan(ret["ygen"][:, 2])] = 0.0
+        ret["ygen"][:, 2][np.isinf(ret["ygen"][:, 2])] = 0.0
+        ret["ygen"][:, 2][ret["ygen"][:, 0] == 0] = 0
+
         return ret
 
     def __len__(self):
@@ -214,10 +231,14 @@ def get_interleaved_dataloaders(world_size, rank, config, use_cuda, use_ray):
                 dataset.append(ds)
             dataset = torch.utils.data.ConcatDataset(dataset)
 
+            shuffle = split == "train"
             if world_size > 1:
-                sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+                sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=shuffle)
             else:
-                sampler = torch.utils.data.SequentialSampler(dataset)
+                if shuffle:
+                    sampler = torch.utils.data.RandomSampler(dataset)
+                else:
+                    sampler = torch.utils.data.SequentialSampler(dataset)
 
             # build dataloaders
             batch_size = config[f"{split}_dataset"][config["dataset"]][type_]["batch_size"] * config["gpu_batch_multiplier"]

mlpf/pyg/inference.py

@@ -42,13 +42,26 @@ def predict_one_batch(conv_type, model, i, batch, rank, jetdef, jet_ptcut, jet_m
     batch = batch.to(rank)
     ypred = model(batch.X, batch.mask)
 
+    # transform log (pt/elempt) -> pt
+    pred_cls = torch.argmax(ypred[0], axis=-1)
+    ypred[2][..., 0] = torch.exp(ypred[2][..., 0]) * batch.X[..., 1]
+    batch.ygen[..., 2] = torch.exp(batch.ygen[..., 2]) * batch.X[..., 1]
+
+    # transform log (E/elemE) -> E
+    ypred[2][..., 4] = torch.exp(ypred[2][..., 4]) * batch.X[..., 5]
+    batch.ygen[..., 6] = torch.exp(batch.ygen[..., 6]) * batch.X[..., 5]
+
+    ypred[2][..., 0][pred_cls == 0] = 0
+    ypred[2][..., 4][pred_cls == 0] = 0
+    batch.ygen[..., 2][batch.ygen[..., 0] == 0] = 0
+    batch.ygen[..., 6][batch.ygen[..., 0] == 0] = 0
+
     # convert all outputs to float32 in case running in float16 or bfloat16
     ypred = tuple([y.to(torch.float32) for y in ypred])
 
-    ygen = unpack_target(batch.ygen.to(torch.float32))
-    ycand = unpack_target(batch.ycand.to(torch.float32))
+    ygen = unpack_target(batch.ygen.to(torch.float32), model)
+    ycand = unpack_target(batch.ycand.to(torch.float32), model)
     ypred = unpack_predictions(ypred)
-
     genjets_msk = batch.genjets[:, :, 0].cpu() != 0
     genjets = awkward.unflatten(batch.genjets.cpu().to(torch.float64)[genjets_msk], torch.sum(genjets_msk, axis=1))
     genjets = vector.awk(
@@ -79,15 +92,18 @@ def predict_one_batch(conv_type, model, i, batch, rank, jetdef, jet_ptcut, jet_m
     jets_coll = {}
     for typ, ydata in zip(["cand", "target"], [ycand, ygen]):
         clsid = awkward.unflatten(ydata["cls_id"], counts)
+        pt = awkward.unflatten(ydata["pt"], counts)
+        eta = awkward.unflatten(ydata["eta"], counts)
+        phi = awkward.unflatten(ydata["phi"], counts)
+        e = awkward.unflatten(ydata["energy"], counts)
         msk = clsid != 0
-        p4 = awkward.unflatten(ydata["p4"], counts)
         vec = vector.awk(
             awkward.zip(
                 {
-                    "pt": p4[msk][:, :, 0],
-                    "eta": p4[msk][:, :, 1],
-                    "phi": p4[msk][:, :, 2],
-                    "e": p4[msk][:, :, 3],
+                    "pt": pt[msk],
+                    "eta": eta[msk],
+                    "phi": phi[msk],
+                    "e": e[msk],
                 }
             )
         )