Improvements for MC-DATA anchoring

MC/bin/o2dpg_sim_workflow.py

@@ -126,6 +126,8 @@
 #  help='Treat smaller sensors in a single digitization')
 parser.add_argument('--pregenCollContext', action='store_true', help=argparse.SUPPRESS) # Now the default, giving this option or not makes not difference. We keep it for backward compatibility
 parser.add_argument('--data-anchoring', type=str, default='', help="Take collision contexts (from data) stored in this path")
+parser.add_argument('--aod-output-folder', type=str, default='', help="Force this AOD folder in the AOD producer")
+parser.add_argument('--aod-parent-file', type=str, default='', help="Link this as parent file in the AOD")
 parser.add_argument('--no-combine-smaller-digi', action='store_true', help=argparse.SUPPRESS)
 parser.add_argument('--no-combine-dpl-devices', action='store_true', help=argparse.SUPPRESS)
 parser.add_argument('--no-mc-labels', action='store_true', default=False, help=argparse.SUPPRESS)
@@ -1794,6 +1796,10 @@ def getDigiTaskName(det):
    if created_by_option != '':
       created_by_option += ' ' + aod_creator
 
+   aod_timeframe_id = f"${{ALIEN_PROC_ID}}{aod_df_id}" if not args.run_anchored else ""
+   if len(args.aod_output_folder) > 0:
+       aod_timeframe_id = args.aod_output_folder
+
    AODtask = createTask(name='aod_'+str(tf), needs=aodneeds, tf=tf, cwd=timeframeworkdir, lab=["AOD"], mem='4000', cpu='1')
    AODtask['cmd'] = ('','ln -nfs ../bkg_Kine.root . ;')[doembedding]
    AODtask['cmd'] += '[ -f AO2D.root ] && rm AO2D.root; '
@@ -1810,12 +1816,13 @@ def getDigiTaskName(det):
       "--anchor-pass ${ALIEN_JDL_LPMANCHORPASSNAME:-unknown}",
       "--anchor-prod ${ALIEN_JDL_LPMANCHORPRODUCTION:-unknown}",
       "--reco-pass ${ALIEN_JDL_LPMPASSNAME:-unknown}",
+      f"--aod-parent {args.aod_parent_file}" if len(args.aod_parent_file) > 0 else "",
       created_by_option,
       "--combine-source-devices" if not args.no_combine_dpl_devices else "",
       "--disable-mc" if args.no_mc_labels else "",
       "--enable-truncation 0" if environ.get("O2DPG_AOD_NOTRUNCATE") or environ.get("ALIEN_JDL_O2DPG_AOD_NOTRUNCATE") else "",
       "--disable-strangeness-tracker" if args.no_strangeness_tracking else "",
-      f"--aod-timeframe-id ${{ALIEN_PROC_ID}}{aod_df_id}" if not args.run_anchored else "",
+      f"--aod-timeframe-id {aod_timeframe_id}" if len(aod_timeframe_id) > 0 else ""
    ])
    # Consider in future: AODtask['disable_alternative_reco_software'] = True # do not apply reco software here (we prefer latest aod converter)
    workflow['stages'].append(AODtask)

MC/run/ANCHOR/anchorMC_DataEmbedding.sh

@@ -329,7 +329,10 @@ for external_context in collission_context_*.root; do
   # extract timeframe from name
   anchoring_tf="${external_context#collission_context_}" # remove prefix 'collision_context_'
   anchoring_tf="${anchoring_tf%.root}" # remove suffix '.root'
-  echo "Treating timeframe ${anchoring_tf}"
+  # now we have a string with DF_FOLDERNAME:TIMEFRAMEID
+  df_folder="${anchoring_tf%%:*}"
+  anchoring_tf="${anchoring_tf##*:}"
+  echo "Treating timeframe ${anchoring_tf} coming from folder ${df_folder}"
 
   # we do it in a separate workspace
   workspace="TF_${anchoring_tf}"
@@ -366,7 +369,7 @@ for external_context in collission_context_*.root; do
   remainingargs="${remainingargs} ${ALIEN_JDL_CCDB_CONDITION_NOT_AFTER:+--condition-not-after ${ALIEN_JDL_CCDB_CONDITION_NOT_AFTER}}"
   # add external collision context injection
   if [ "${ALIEN_JDL_MC_DATA_EMBEDDING_AO2D}" ]; then
-    remainingargs="${remainingargs} --data-anchoring ${PWD}/../${external_context}"
+    remainingargs="${remainingargs} --data-anchoring ${PWD}/../${external_context} --aod-output-folder ${df_folder##*_} --aod-parent-file ${ALIEN_JDL_MC_DATA_EMBEDDING_AO2D}"
   fi
 
   echo_info "baseargs passed to o2dpg_sim_workflow_anchored.py: ${baseargs}"

MC/utils/o2dpg_data_embedding_utils.py

@@ -1,6 +1,13 @@
 # Set of python modules/util functions for the MC-to-DATA embedding
 # Mostly concerning extraction of MC collision context from existing data AO2D.root
 
+import warnings
+warnings.filterwarnings(
+    "ignore",
+    message="pandas.Int64Index is deprecated",
+    category=FutureWarning,
+)
+
 import ROOT
 import uproot
 import pandas as pd
@@ -240,7 +247,7 @@ def fetch_bccoll_to_localFile(alien_file, local_filename):
   return True
 
 
-def convert_to_digicontext(aod_timeframe=None, timeframeID=-1):
+def convert_to_digicontext(aod_timeframe, df_folder, timeframeID=-1):
     """
     converts AOD collision information from AO2D to collision context
     which can be used for MC
@@ -282,13 +289,16 @@ def convert_to_digicontext(aod_timeframe=None, timeframeID=-1):
 
     # set the bunch filling ---> NEED to fetch it from CCDB
     # digicontext.setBunchFilling(bunchFillings[0]);
+
+    # TODO: set the interaction rate (for TPC loopers)
+    # digicontext.mDigitizerInteractionRate = ...
 
     prefixes = ROOT.std.vector("std::string")();
     prefixes.push_back("sgn")
 
     digicontext.setSimPrefixes(prefixes);
     digicontext.printCollisionSummary();
-    digicontext.saveToFile(f"collission_context_{timeframeID}.root")
+    digicontext.saveToFile(f"collission_context_{df_folder}:{timeframeID}.root")
 
 
 def process_data_AO2D(file_name, run_number, upper_limit = -1):
@@ -315,7 +325,8 @@ def process_data_AO2D(file_name, run_number, upper_limit = -1):
                     break
                 tf = row['timeframeID']
                 cols = row['position_vectors']
-                convert_to_digicontext(cols, tf)
+                df = key.split('@')[0] # this is the DF folder name
+                convert_to_digicontext(cols, df, tf)
                 counter = counter + 1
 
 
@@ -324,7 +335,7 @@ def main():
 
     parser.add_argument("--run-number", type=int, help="Run number to anchor to", required=True)
     parser.add_argument("--aod-file", type=str, help="Data AO2D file (can be on AliEn)", required=True)
-    parser.add_argument("--limit", type=int, default=-1, help="Upper limit of timeframes to be extracted")
+    parser.add_argument("--limit", type=int, default=-1, help="Upper limit of timeframes to be extracted (-1 is no limit)")
     args = parser.parse_args()
 
     process_data_AO2D(args.aod_file, args.run_number, args.limit)