Merge pull request #1 from jpata/master

Merge new commits from jpata:master

README_clic.md

@@ -0,0 +1,10 @@
+Download the .slcio files from https://atlaswww.hep.anl.gov/hepsim/list.php?find=rfull201 -> `gev380ee_pythia6_ttbar` to
+
+```
+data/clic/gev380ee_pythia6_ttbar_rfull201/root/
+```
+
+Run
+```bash
+ls data/clic/gev380ee_pythia6_ttbar_rfull201/root/*.slcio | parallel -j4 singularity exec delphes/centos7hepsim.sif ./clic/process_data.sh {}
+```

clic/dumper.py

@@ -20,25 +20,37 @@
 import math
 import json
 import bz2
+import sys
 
-save_calohits = True
+perfile = 10
 
 def genParticleToDict(par):
     mom = par.getMomentum()
-    parent_pdgid = 0
+    parent_pdgid0 = 0
+    parent_idx0 = -1
+    parent_pdgid1 = 0
+    parent_idx1 = -1
+
     if len(par.getParents()) > 0:
-        parent_pdgid = par.getParents()[0].getPDG()
-
+        parent_pdgid0 = par.getParents()[0].getPDG()
+        parent_idx0 = genparticle_dict[par.getParents()[0]]
+    if len(par.getParents()) > 1:
+        parent_pdgid1 = par.getParents()[1].getPDG()
+        parent_idx1 = genparticle_dict[par.getParents()[1]]
+
     vec = {
         "pdgid": par.getPDG(),
         "status": par.getGeneratorStatus(),
         "mass": par.getMass(),
         "charge": par.getCharge(),
-        "pdgid_parent0": parent_pdgid,
         "px": mom[0],
         "py": mom[1],
         "pz": mom[2],
-        "energy": par.getEnergy()
+        "energy": par.getEnergy(),
+        "pdgid_parent0": parent_pdgid0,
+        "idx_parent0": parent_idx0,
+        "pdgid_parent1": parent_pdgid1,
+        "idx_parent1": parent_idx1
     }
     return vec
 
@@ -146,14 +158,15 @@ def caloHitToDict(par, calohit_to_cluster, genparticle_dict, calohit_recotosim):
     return vec
 
 if __name__ == "__main__":
-    infile = "/home/joosep/Downloads/pythia6_ttbar_0001_pandora.slcio"
+    infile = sys.argv[1]
 
     factory = LCFactory.getInstance()
     reader = factory.createLCReader()
     reader.open(infile)
     event_data = []
 
     nEvent = 0
+    ioutfile = 0
     while True:
         evt = reader.readNextEvent()
         if (evt == None):
@@ -214,9 +227,12 @@ def caloHitToDict(par, calohit_to_cluster, genparticle_dict, calohit_recotosim):
 
         genparticles = []
         genparticle_dict = {}
-        for i in range(nMc): # loop over all particles 
+        for i in range(nMc):
             par=col.getElementAt(i)
             genparticle_dict[par] = i
+
+        for i in range(nMc):
+            par=col.getElementAt(i)
             vec = genParticleToDict(par)
             genparticles.append(vec)
 
@@ -311,11 +327,14 @@ def caloHitToDict(par, calohit_to_cluster, genparticle_dict, calohit_recotosim):
         }
 
         event_data.append(event)
+        if len(event_data) >= perfile:
+            ofi = bz2.BZ2File(infile.replace(".slcio", "_%d.json.bz2"%ioutfile), "w")
+            json.dump(event_data, ofi, indent=2, sort_keys=True)
+            ofi.close()
+            event_data = []
+            ioutfile += 1
         nEvent += 1
 
     #save the event data to a file 
-    ofi = bz2.BZ2File(infile.replace(".slcio", ".json.bz2"), "w")
-    json.dump(event_data, ofi, indent=2, sort_keys=True)
-    ofi.close()
 
     reader.close() # close the file

mlpf/data/postprocessing2.py

@@ -283,7 +283,8 @@ def prepare_normalized_table(g, genparticle_energy_threshold=0.2):
     elem_branches = [
         "typ", "pt", "eta", "phi", "e",
         "layer", "depth", "charge", "trajpoint", 
-        "eta_ecal", "phi_ecal", "eta_hcal", "phi_hcal", "muon_dt_hits", "muon_csc_hits"
+        "eta_ecal", "phi_ecal", "eta_hcal", "phi_hcal", "muon_dt_hits", "muon_csc_hits",
+        "px", "py", "pz", "deltap", "sigmadeltap"
     ]
     target_branches = ["typ", "charge", "pt", "eta", "sin_phi", "cos_phi", "e"]
 

mlpf/launcher.py

@@ -12,19 +12,45 @@ def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--model-spec", type=str, default="parameters/delphes-gnn-skipconn.yaml", help="the model specification")
     parser.add_argument("--action", type=str, choices=["data", "train", "eval", "time"], help="Run training, validation or timing", default="train")
+    parser.add_argument("--modifier", type=str, choices=["retrain_energy", None], help="Apply a modification on the standard training", default=None)
     parser.add_argument("--weights", type=str, help="weight file to load", default=None)
     parser.add_argument("--ntrain", type=int, help="override the number of training events", default=None)
     parser.add_argument("--ntest", type=int, help="override the number of testing events", default=None)
     parser.add_argument("--recreate", action="store_true", help="recreate a new output dir", default=None)
+    parser.add_argument("--raw-path", type=str, help="Override the dataset raw files path", default=None)
+    parser.add_argument("--processed-path", type=str, help="Override the dataset processed files path", default=None)
     args = parser.parse_args()
     return args
 
+def apply_modifier_retrain_energy(config):
+    assert(config["parameters"]["model"] == "gnn_dense")
+    config["setup"]["trainable"] = "ffn_momentum4"
+    for loss in [
+        "classification_loss_coef",
+        "charge_loss_coef",
+        "pt_loss_coef",
+        "eta_loss_coef",
+        "sin_phi_loss_coef",
+        "cos_phi_loss_coef"]:
+        config["dataset"][loss] = 0.0
+
+    config["dataset"]["energy_loss_coef"] = 1.0
+    config["setup"]["batch_size"] = 20
+    return config
+
+modifiers = {
+    "retrain_energy": apply_modifier_retrain_energy
+}
+
 if __name__ == "__main__":
     args = parse_args()
     yaml_path = args.model_spec
 
     config = load_config(yaml_path)
 
+    if args.modifier:
+        config = modifiers[args.modifier](config)
+
     if config["backend"] == "tensorflow":
         tfmodel.model_setup.main(args, yaml_path, config)
     elif config["backend"] == "pytorch":

mlpf/tallinn/cms-gnn-dense-focal-retrain.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+#SBATCH -p gpu
+#SBATCH --gpus 5
+#SBATCH --mem-per-gpu=8G
+
+IMG=/home/software/singularity/base.simg:latest
+cd ~/particleflow
+
+#TF training
+singularity exec --nv $IMG python3 mlpf/launcher.py --model-spec parameters/cms-gnn-dense-focal.yaml --action train --modifier retrain_energy --recreate --weights experiments/cms-gnn-dense-focal-285ae825.gpu0.local/weights-300-1.175282.hdf5

mlpf/tallinn/cms-gnn-dense.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 #SBATCH -p gpu
-#SBATCH --gpus 1
+#SBATCH --gpus 5
 #SBATCH --mem-per-gpu=8G
 
 IMG=/home/software/singularity/base.simg:latest

mlpf/tfmodel/data.py

@@ -81,7 +81,7 @@ def prepare_data_cms(self, fn):
         if fn.endswith(".pkl"):
             data = pickle.load(open(fn, "rb"), encoding='iso-8859-1')
         elif fn.endswith(".pkl.bz2"):
-            data = pickle.load(bz2.BZ2File(fn, "rb"))
+            data = pickle.load(bz2.BZ2File(fn, "rb"), encoding='iso-8859-1')
 
         for event in data:
             Xelem = event["Xelem"]