feat: Enhance phys/root/filter functionality

phys/root/filter/test/Snakefile

@@ -1,15 +1,52 @@
-rule define_columns:
+rule filter_str:
     input:
         "ntuple0.root",
     output:
-        "ntuple0_output.root",
+        "ntuple0_str_output.root",
     log:
-        "logs/filter/filter.log",
+        "logs/filter/filter_str.log",
+    params:
+        input_tree_name="TestTree",
+        output_tree_name="TestTree",
+        criteria="(pt > 1400) && (pz > 19000)",
+        branches_to_save=["pz", "pt", "p"],
+    threads: 2
+    wrapper:
+        "master/phys/root/filter"
+
+rule filter_list:
+    input:
+        "ntuple0.root",
+    output:
+        "ntuple0_list_output.root",
+    log:
+        "logs/filter/filter_list.log",
     params:
         input_tree_name="TestTree",
         output_tree_name="TestTree",
-        criteria="pt > 1400",
+        criteria=["pt > 1400", "pz > 19000"],
         branches_to_save=["pz", "pt", "p"],
+        verbose=True
     threads: 2
     wrapper:
         "master/phys/root/filter"
+
+rule filter_dict:
+    input:
+        "ntuple0.root",
+    output:
+        "ntuple0_dict_output.root",
+    log:
+        "logs/filter/filter.log",
+    params:
+        input_tree_name="TestTree",
+        output_tree_name="TestTree",
+        criteria={
+            "PT cut": "pt > 1400",
+            "PZ cut": "pz > 19000"
+        },
+        branches_to_save=["pz", "pt", "p"],
+        verbose=True
+    threads: 1
+    wrapper:
+        "master/phys/root/filter"

phys/root/filter/wrapper.py

@@ -1,18 +1,39 @@
-__author__ = "Anfeng Li"
+__author__ = "Anfeng Li, Jamie Gooding"
 __copyright__ = "Copyright 2024, Anfeng Li"
-__email__ = "[email protected]"
+__email__ = "[email protected], [email protected]"
 __license__ = "MIT"
 
-
+from typing import Dict, List
 import ROOT
 
-ROOT.EnableImplicitMT(snakemake.threads)
+if snakemake.threads > 1:
+    ROOT.EnableImplicitMT(snakemake.threads)
+else:
+    ROOT.DisableImplicitMT()
+
+# Parse criteria
+_smk_criteria = snakemake.params.get("criteria", "true")
+if isinstance(_smk_criteria, str):
+    criteria = [_smk_criteria]
+    labels = [_smk_criteria]
+elif isinstance(_smk_criteria, List):
+    criteria = _smk_criteria
+    labels = _smk_criteria
+elif isinstance(_smk_criteria, Dict):
+    criteria = _smk_criteria.values()
+    labels = _smk_criteria.keys()
+else:
+    raise TypeError("Parameter 'criteria' should be type of 'str', 'list' or 'dict'")
 
-criteria = snakemake.params.get("criteria", "true")
 branches_to_save = snakemake.params.get("branches_to_save", None)
 
 df = ROOT.RDataFrame(snakemake.params.input_tree_name, snakemake.input[0])
-df = df.Filter(criteria)
+for criterion, label in zip(criteria, labels):
+    df = df.Filter(criterion, label)
+
+if snakemake.params.get("verbose", False):
+    report = df.Report()
+
 if branches_to_save is not None:
     df.Snapshot(
         snakemake.params.output_tree_name,
@@ -21,3 +42,6 @@
     )
 else:
     df.Snapshot(snakemake.params.output_tree_name, snakemake.output[0])
+
+if snakemake.params.get("verbose", False):
+    report.Print()

test.py

@@ -6829,10 +6829,26 @@ def test_root_define_columns():
 
 
 @skip_if_not_modified
-def test_root_filter():
+def test_root_filter_str():
     run(
         "phys/root/filter",
-        ["snakemake", "--cores", "2", "--use-conda", "-F"],
+        ["snakemake", "--cores", "2", "--use-conda", "-F", "ntuple0_str_output.root"],
+    )
+
+
+@skip_if_not_modified
+def test_root_filter_list():
+    run(
+        "phys/root/filter",
+        ["snakemake", "--cores", "2", "--use-conda", "-F", "ntuple0_list_output.root"],
+    )
+
+
+@skip_if_not_modified
+def test_root_filter_dict():
+    run(
+        "phys/root/filter",
+        ["snakemake", "--cores", "2", "--use-conda", "-F", "ntuple0_dict_output.root"],
     )