Add a test for issue 100 in DaVinciTests

Closes #100 (closed).

The test is a shorter version of the original script focusing only on what was giving trouble.

DaVinciTests/tests/qmtest/mc.qms/test_davinci-issue-100_multiple_bkgcat_mc-truth.qmt

@@ -32,7 +32,6 @@ validateWithReference(preproc = preprocessor, counter_preproc = counter_preproce
 countErrorLines({"FATAL": 0, "ERROR": 0}, stdout=stdout)
 
 import os
-from DaVinciTests.QMTest.check_helpers import get_pandas_dataframe, list_fields_with_nan
 
 filename = "./davinci-issue-100_multiple_bkgcat_mc-truth_ntuple.root"
 
@@ -49,7 +48,9 @@ df_shapes = dict(zip(samples, [(113, 44), (133, 44), (87, 44), (72, 44)]))
 l_branches_with_nans =\
     dict(zip(samples, [["J_psi_1S_TRUEM", "J_psi_1S_TRUEPT", "muplus_TRUEM", "muplus_TRUEPT"]] + 3*[["J_psi_1S_TRUEM", "J_psi_1S_TRUEPT"]]))
 
-def run_basic_checks(sample):
+for sample in samples:
+    from DaVinciTests.QMTest.check_helpers import get_pandas_dataframe, list_fields_with_nan
+
     df = get_pandas_dataframe(filename, ntuples[sample])
 
     # Check ntuple structure
@@ -63,26 +64,81 @@ def run_basic_checks(sample):
     if sorted(l_test) != sorted(l_branches_with_nans[sample]):
         causes.append(f"Ntuple {ntuples[sample]}: unexpected list of branches with NaN values")
 
-for sample in samples:
-    run_basic_checks(sample)
-
-# ==> Extra checks on sample "Tuple_Hlt2TrackEff_DiMuon_SeedMuon_mup_Tag"
+# ==> Extra checks on
+sample = "Tuple_Hlt2TrackEff_DiMuon_SeedMuon_mup_Tag"
+df = get_pandas_dataframe(filename, ntuples[sample])
 # Checks PIDs are correctly assigned
 if not ( ( (df["J_psi_1S_TRUEID"]==443).value_counts()[False] == 4 )
          and ( (df["J_psi_1S_MC_MOTHER_ID"].abs()==531).value_counts()[False] == 4 )
          and ( (df["muplus_TRUEID"].abs()==13).value_counts()[False] == 1 )  # all true mu+ except 1
-         and ( (df["muminus_TRUEID"].abs() == 13).all().all() )  # all true mu-
+         and ( (df["muminus_TRUEID"].abs() == 13).all() )  # all true mu- 
          and ( (df["muplus_MC_MOTHER_ID"]==443).sum() == 110 )
          and ( (df["muminus_MC_MOTHER_ID"]==443).sum() == 112 )
        ):
-    causes.append("Ntuple contains unexpected MC ID values")
-
+    causes.append(f"Ntuple {ntuples[sample]} contains unexpected MC ID values")
 # Check background categories
 if not ( ( (df["J_psi_1S_BKGCAT"]==0).sum() == 98 )
+         and ( (df["muplus_BKGCAT"]==-1).all() )  # all entries are -1 for final-state particles
+         and ( (df["muminus_BKGCAT"]==-1).all() )  # all entries are -1 for final-state particles
+       ):
+    causes.append(f"Ntuple {ntuples[sample]} contains unexpected BKGCAT values")
+
+# ==> Extra checks on
+sample = "Tuple_Hlt2TrackEff_DiMuon_SeedMuon_mum_Tag"
+df = get_pandas_dataframe(filename, ntuples[sample])
+# Checks PIDs are correctly assigned
+if not ( ( (df["J_psi_1S_TRUEID"]==443).value_counts()[False] == 9 )
+         and ( (df["J_psi_1S_MC_MOTHER_ID"].abs()==531).value_counts()[False] == 9 )
+         and ( (df["muplus_TRUEID"].abs()==13).all() )  # all true mu+
+         and ( (df["muminus_TRUEID"].abs() == 13).value_counts()[False] == 6 )
+         and ( (df["muplus_MC_MOTHER_ID"]==443).sum() == 131 )
+         and ( (df["muminus_MC_MOTHER_ID"]==443).sum() == 125 )
+       ):
+    causes.append(f"Ntuple {ntuples[sample]} contains unexpected MC ID values")
+# Check background categories
+if not ( ( (df["J_psi_1S_BKGCAT"]==0).sum() == 120 )
+         and ( (df["muplus_BKGCAT"]==-1).all() )  # all entries are -1 for final-state particles
+         and ( (df["muminus_BKGCAT"]==-1).all() )  # all entries are -1 for final-state particles
+       ):
+    causes.append(f"Ntuple {ntuples[sample]} contains unexpected BKGCAT values")
+
+# ==> Extra checks on
+sample = "Tuple_Hlt2TrackEff_DiMuon_VeloMuon_mup_Tag"
+df = get_pandas_dataframe(filename, ntuples[sample])
+# Checks PIDs are correctly assigned
+if not ( ( (df["J_psi_1S_TRUEID"]==443).value_counts()[False] == 5 )
+         and ( (df["J_psi_1S_MC_MOTHER_ID"].abs()==531).value_counts()[False] == 5 )
+         and ( (df["muplus_TRUEID"].abs()==13).value_counts()[False] == 2 )  # all true mu+ except 2
+         and ( (df["muminus_TRUEID"].abs() == 13).value_counts()[False] == 2 )  # all true mu- except 2
+         and ( (df["muplus_MC_MOTHER_ID"]==443).sum() == 84 )
+         and ( (df["muminus_MC_MOTHER_ID"]==443).sum() == 84 )
+       ):
+    causes.append(f"Ntuple {ntuples[sample]} contains unexpected MC ID values")
+# Check background categories
+if not ( ( (df["J_psi_1S_BKGCAT"]==0).sum() == 79 )
+         and ( (df["muplus_BKGCAT"]==-1).all() )  # all entries are -1 for final-state particles
+         and ( (df["muminus_BKGCAT"]==-1).all() )  # all entries are -1 for final-state particles
+       ):
+    causes.append(f"Ntuple {ntuples[sample]} contains unexpected BKGCAT values")
+
+# ==> Extra checks on
+sample = "Tuple_Hlt2TrackEff_DiMuon_VeloMuon_mum_Tag"
+df = get_pandas_dataframe(filename, ntuples[sample])
+# Checks PIDs are correctly assigned
+if not ( ( (df["J_psi_1S_TRUEID"]==443).value_counts()[False] == 1 )
+         and ( (df["J_psi_1S_MC_MOTHER_ID"].abs()==531).value_counts()[False] == 1 )
+         and ( (df["muplus_TRUEID"].abs()==13).all() )  # all true mu+
+         and ( (df["muminus_TRUEID"].abs() == 13).value_counts()[False] )  # all true mu- except 1
+         and ( (df["muplus_MC_MOTHER_ID"]==443).all() )
+         and ( (df["muminus_MC_MOTHER_ID"]==443).sum() == 71 )
+       ):
+    causes.append(f"Ntuple {ntuples[sample]} contains unexpected MC ID values")
+# Check background categories
+if not ( ( (df["J_psi_1S_BKGCAT"]==0).sum() == 66 )
          and ( (df["muplus_BKGCAT"]==-1).all().all() )  # all entries are -1 for final-state particles
          and ( (df["muminus_BKGCAT"]==-1).all().all() )  # all entries are -1 for final-state particles
        ):
-    causes.append("Ntuple contains unexpected BKGCAT values")
+    causes.append(f"Ntuple {ntuples[sample]} contains unexpected BKGCAT values")
 
 print('Test successfully completed!')
 os.system(f"rm {filename}")