final fixes

muonDIS/add_muonresponse.py

@@ -1,9 +1,9 @@
-#!/usr/bin/env python3
-
+#!/usr/bin/env python
 """Script to add the incoming muon's MC hits in the SBT to the simulation file."""
 
 import argparse
 import logging
+import os
 
 import ROOT as r
 
@@ -22,67 +22,100 @@
 )
 args = parser.parse_args()
 
-# Open the existing ROOT file in update mode
-logging.warning(f"{args.inputfile} opened for modification")
-inputFile = r.TFile.Open(args.inputfile, "UPDATE")
-try:
-    input_tree = inputFile.cbmsim
-except Exception as e:
-    print(f"Error: {e}")
-    inputFile.Close()
-    exit(1)
-
-# Open the external file with additional vetoPoints
-muonFile = r.TFile.Open(args.muonfile, "READ")
-try:
-    muon_tree = muonFile.DIS
-except Exception as e:
-    print(f"Error: {e}")
-    muonFile.Close()
-    exit(1)
-
-inputFile.cd()
-output_tree = input_tree.CloneTree(
-    0
-)  # Clone the structure of the existing tree, but do not copy the entries
-
-# Access the vetoPoint branch in the original and external trees
-originalVetoPoints = r.TClonesArray("vetoPoint")
-muonVetoPoints = r.TClonesArray("vetoPoint")
-combinedVetoPoints = r.TClonesArray("vetoPoint")
-
-input_tree.SetBranchAddress("vetoPoint", originalVetoPoints)
-muon_tree.SetBranchAddress("muon_vetoPoints", muonVetoPoints)
-output_tree.SetBranchAddress("vetoPoint", combinedVetoPoints)
-
-
-for input_event, muon_event in zip(input_tree, muon_tree):
-    interaction_point = r.TVector3()
-    input_event.MCTrack[0].GetStartVertex(interaction_point)
 
-    combinedVetoPoints.Clear()
-
-    index = 0
+def read_file():
+    """Inspecting the produced file for successfully added muon veto points."""
+    inputFile = r.TFile.Open(args.inputfile, "READ")
+    input_tree = inputFile.cbmsim
+    input_tree.GetEntry(0)
 
-    for hit in originalVetoPoints:
-        if combinedVetoPoints.GetSize() == index:
-            combinedVetoPoints.Expand(index + 1)
-        combinedVetoPoints[index] = hit
-        index += 1
+    muons = False
+    for hit in input_tree.vetoPoint:
+        if hit.GetTrackID() == 0:
+            muons = True
+    inputFile.Close()
 
-    for hit in muonVetoPoints:
-        if hit.GetZ() < interaction_point.Z():
+    if muons:
+        print("File is already updated with incoming muon info. Nothing to do.")
+        return False
+    else:
+        print(
+            "Incoming muon's vetopoint info missing in file, proceeding with modification"
+        )
+        return True
+
+
+def modify_file():
+    """Modify the input file with muon info."""
+    if args.inputfile[0:4] == "/eos":
+        args.inputfile = os.environ["EOSSHIP"] + args.inputfile
+
+    logging.warning(f"{args.inputfile} opened for modification")
+
+    inputFile = r.TFile.Open(args.inputfile, "UPDATE")
+    try:
+        input_tree = inputFile.cbmsim
+    except Exception as e:
+        print(f"Error: {e}")
+        inputFile.Close()
+        exit(1)
+
+    # Open the external file with additional vetoPoints
+    muonFile = r.TFile.Open(args.muonfile, "READ")
+    try:
+        muon_tree = muonFile.DIS
+    except Exception as e:
+        print(f"Error: {e}")
+        muonFile.Close()
+        exit(1)
+
+    inputFile.cd()
+    output_tree = input_tree.CloneTree(
+        0
+    )  # Clone the structure of the existing tree, but do not copy the entries
+
+    # Access the vetoPoint branch in the original and external trees
+    originalVetoPoints = r.TClonesArray("vetoPoint")
+    muonVetoPoints = r.TClonesArray("vetoPoint")
+    combinedVetoPoints = r.TClonesArray("vetoPoint")
+
+    input_tree.SetBranchAddress("vetoPoint", originalVetoPoints)
+    muon_tree.SetBranchAddress("muon_vetoPoints", muonVetoPoints)
+    output_tree.SetBranchAddress("vetoPoint", combinedVetoPoints)
+
+    for input_event, muon_event in zip(input_tree, muon_tree):
+        interaction_point = r.TVector3()
+        input_event.MCTrack[0].GetStartVertex(interaction_point)
+
+        combinedVetoPoints.Clear()
+
+        index = 0
+
+        for hit in originalVetoPoints:
             if combinedVetoPoints.GetSize() == index:
                 combinedVetoPoints.Expand(index + 1)
             combinedVetoPoints[index] = hit
             index += 1
 
-    output_tree.Fill()
+        for hit in muonVetoPoints:
+            if hit.GetZ() < interaction_point.Z():
+                if combinedVetoPoints.GetSize() == index:
+                    combinedVetoPoints.Expand(index + 1)
+                combinedVetoPoints[index] = hit
+                index += 1
+
+        output_tree.Fill()
+
+    # Write the updated tree back to the same file
+    inputFile.cd()
+    output_tree.Write("cbmsim", r.TObject.kOverwrite)
+    inputFile.Close()
+    muonFile.Close()
+
+    print(f"File updated with incoming muon info as {args.inputfile}")
+
 
-# Write the updated tree back to the same file
-inputFile.cd()  # Ensure we are in the correct directory of the ROOT file
-output_tree.Write("cbmsim", r.TObject.kOverwrite)
-inputFile.Close()
-muonFile.Close()
+add_muons = read_file()
 
-print(f"Updated file saved as {args.inputfile}")
+if add_muons:
+    modify_file()

muonDIS/makeMuonDIS.py

@@ -94,7 +94,7 @@ def makeMuonDIS():
     mutype = {-13: "gamma/mu+", 13: "gamma/mu-"}
 
     myPythia.SetMSTU(11, 11)
-    logging.info(f"Processing events from {nStart} to {nEnd}...")
+    logging.info(f"Processing events from {nStart} to {nEnd-1}...")
 
     nMade = 0
 
@@ -121,7 +121,7 @@ def makeMuonDIS():
         phi = r.TMath.ATan2(py, px)
 
         logging.debug(
-            f"Muon: PID={pid}, px={px}, py={py}, pz={pz}, E={E}, x={x}, y={y}, z={z}, w={w}"
+            f"\n\tMuon index:{k} \n\tPID = {pid}, weight = {w} \n\tpx = {px}, py = {py}, pz = {pz}, E = {E},\n\tx = {x}, y = {y}, z = {z}\n"
         )
 
         isProton = 1
@@ -136,15 +136,17 @@ def makeMuonDIS():
             if a == args.nDIS // 2:
                 myPythia.Initialize("FIXT", mutype[pid], "n0", p)  # target = "n0"
                 isProton = 0
-                logging.info("Switching to neutron interaction")
+                logging.debug("Switching to neutron interaction")
 
             dPartDIS.Clear()
             iMuon.Clear()
             muPart[9] = isProton
             iMuon[0] = muPart
             myPythia.GenerateEvent()
             myPythia.Pyedit(1)
-            logging.info(f"Event {a} generated, number of particles: {myPythia.GetN()}")
+            logging.debug(
+                f"DIS Event {a} generated, number of particles: {myPythia.GetN()}"
+            )
 
             for itrk in range(1, myPythia.GetN() + 1):
                 xsec = myPythia.GetPARI(1)
@@ -218,7 +220,7 @@ def makeMuonDIS():
     output_tree.Write()
     myPythia.SetMSTU(11, 6)
     logging.info(
-        f"Created DIS for  saved in muonDis_{args.nJob}.root, muon event index from {nStart} - {nEnd} , nDISPerMuon {args.nDIS}"
+        f"DIS generated for muons (index {nStart} - {nEnd-1}) , output saved in muonDis_{args.nJob}.root, nDISPerMuon = {args.nDIS}"
     )