Merge pull request #115 from FAST-HEP/BK_repeatable_SystWeight_stage

Improve SystematicWeights stage

CHANGELOG.md

@@ -5,6 +5,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+## [0.17.3] - 2020-02-27
+### Changed
+- Allow SystematicWeights stage to be used twice, PR #115 [@BenKrikler](https://github.com/benkrikler/)
 
 ## [0.17.2] - 2020-02-25
 ### Changed

fast_carpenter/define/systematics.py

@@ -31,6 +31,11 @@ class SystematicWeights():
           variable to use for the "nominal" variation, or a dictionary containing
           any of the keys, ``nominal``, ``up``, or ``down``.  Each of these should
           then have a value providing the expression to use for that variation/
+      out_format (str): The format string to use to build the name of the
+          output variations.  Defaults to "weight_{}".  Should contain a pair
+          of empty braces which will be replaced with the name for the current
+          variation, e.g. "nominal" or "PileUp_up".
+      extra_variations (list[str]): A list of additional variations to allow
 
     Other Parameters:
       name (str):  The name of this stage (handled automatically by fast-flow)
@@ -53,11 +58,11 @@ class SystematicWeights():
         weight_energy_scale_down =  WeightEnergyScaleDown * TriggerEfficiency * ReconEfficiency
         weight_recon_up =  WeightEnergyScale * TriggerEfficiency * ReconEfficiency_up
     """
-    def __init__(self, name, out_dir, weights):
+    def __init__(self, name, out_dir, weights, out_format="weight_{}", extra_variations=[]):
         self.name = name
         self.out_dir = out_dir
-        weights = _normalize_weights(name, weights)
-        variations = _build_variations(name, weights)
+        weights = _normalize_weights(name, weights, tuple(extra_variations))
+        variations = _build_variations(name, weights, out_fmt=out_format)
         self.variable_maker = Define(name + "_builder", out_dir, variations)
 
     def event(self, chunk):
@@ -66,33 +71,41 @@ def event(self, chunk):
         return self.variable_maker.event(chunk)
 
 
-def _normalize_weights(stage_name, variable_list):
+def _normalize_weights(stage_name, variable_list, valid_vars):
     if not isinstance(variable_list, dict):
         msg = "{}: Didn't receive a list of variables"
         raise BadSystematicWeightsConfig(msg.format(stage_name))
-    return {name: _normalize_one_variation(stage_name, cfg) for name, cfg in variable_list.items()}
+    return {name: _normalize_one_variation(stage_name, cfg, name, valid_vars=valid_vars)
+            for name, cfg in variable_list.items()}
 
 
-def _build_variations(stage_name, weights, out_name="weight_{}"):
+def _build_variations(stage_name, weights, out_fmt="weight_{}"):
+    def _combine_weights(w):
+        return "(" + ")*(".join(w) + ")"
+
     nominal_weights = {n: w["nominal"] for n, w in weights.items()}
-    variations = [{out_name.format("nominal"): "*".join(nominal_weights.values())}]
-    weights_to_vary = {(n, var): w[var] for n, w in weights.items() for var in ("up", "down") if var in w}
+    variations = [{out_fmt.format("nominal"): _combine_weights(nominal_weights.values())}]
+    weights_to_vary = {(n, var): w[var] for n, w in weights.items() for var in w if var != "nominal"}
     for (name, direction), variable in weights_to_vary.items():
         combination = nominal_weights.copy()
         combination[name] = variable
-        combination = "*".join(combination.values())
-        variations.append({out_name.format(name + "_" + direction): combination})
+        combination = _combine_weights(combination.values())
+        variations.append({out_fmt.format(name + "_" + direction): combination})
     return variations
 
 
-def _normalize_one_variation(stage_name, cfg):
+def _normalize_one_variation(stage_name, cfg, name, valid_vars):
     if isinstance(cfg, six.string_types):
         return dict(nominal=cfg)
     if not isinstance(cfg, dict):
         msg = "{}: Each systematic weight should be either a dict or just a string"
         raise BadSystematicWeightsConfig(msg.format(stage_name))
-    bad_keys = [key for key in cfg if key not in ("nominal", "up", "down")]
+    if "nominal" not in cfg:
+        msg = "{}: No nominal weight provided for '{}'"
+        raise BadSystematicWeightsConfig(msg.format(stage_name, name))
+
+    bad_keys = [key for key in cfg if key not in ("nominal", "up", "down") + valid_vars]
     if bad_keys:
-        msg = "{}: Received unknown keys '{}'"
-        raise BadSystematicWeightsConfig(msg.format(stage_name, bad_keys))
+        msg = "{}: Received unknown keys,'{}', for '{}'"
+        raise BadSystematicWeightsConfig(msg.format(stage_name, bad_keys, name))
     return cfg

fast_carpenter/version.py

@@ -12,5 +12,5 @@ def split_version(version):
     return tuple(result)
 
 
-__version__ = '0.17.2'
+__version__ = '0.17.3'
 version_info = split_version(__version__) # noqa

setup.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.17.2
+current_version = 0.17.3
 commit = True
 tag = False
 

tests/define/test_systematics.py

@@ -41,3 +41,52 @@ def test_systematic_variations_1_mc(systematic_variations_1, fake_file):
     assert all(chunk.tree.weight_nominal == [0, 4, 8])
     assert all(chunk.tree.weight_varied_up == [0, 6., 12.])
     assert all(chunk.tree.weight_varied_down == [0, 3., 6.])
+
+
+def test_normalize_one_variation():
+    cfg = dict(nominal="one", up="two")
+    stage_name = "test_normalize_one_variation"
+    out = fast_syst._normalize_one_variation(stage_name, cfg, "test1", tuple())
+    assert len(out) == 2
+    assert out["nominal"] == "one"
+    assert out["up"] == "two"
+
+    with pytest.raises(fast_syst.BadSystematicWeightsConfig) as e:
+        fast_syst._normalize_one_variation(stage_name, dict(a="bad"), "test2", tuple())
+    assert "nominal" in str(e)
+
+    with pytest.raises(fast_syst.BadSystematicWeightsConfig) as e:
+        fast_syst._normalize_one_variation(stage_name, dict(nominal="one", a="bad"), "test3", tuple())
+    assert "unknown key" in str(e)
+
+    out = fast_syst._normalize_one_variation(stage_name, dict(nominal="one", a="bad"), "test3", tuple("a"))
+    assert len(out) == 2
+    assert out["nominal"] == "one"
+    assert out["a"] == "bad"
+
+    out = fast_syst._normalize_one_variation(stage_name, "just_a_string", "test4", tuple())
+    assert len(out) == 1
+    assert out["nominal"] == "just_a_string"
+
+
+def test_build_variations():
+    pileup = dict(nominal="PILEUP", up="PILEUP_UP", down="PILEUP_DOWN")
+    isolation = dict(nominal="Iso", up="IsoUp")
+    another = dict(nominal="Blahblah", left="BlahblahLeft")
+
+    all_vars = dict(pileup=pileup, isolation=isolation, another=another)
+    stage_name = "test_build_variations"
+    formulae = fast_syst._build_variations(stage_name, all_vars, out_fmt="Weight_{}_test")
+    assert len(formulae) == 5
+    assert isinstance(formulae, list)
+    assert all(map(lambda x: isinstance(x, dict) and len(x) == 1, formulae))
+    formulae = {list(i.keys())[0]: list(i.values())[0] for i in formulae}
+
+    def prep(string):
+        return set(string.split("*"))
+
+    assert prep(formulae["Weight_nominal_test"]) == prep("(PILEUP)*(Iso)*(Blahblah)")
+    assert prep(formulae["Weight_pileup_up_test"]) == prep("(PILEUP_UP)*(Iso)*(Blahblah)")
+    assert prep(formulae["Weight_pileup_down_test"]) == prep("(PILEUP_DOWN)*(Iso)*(Blahblah)")
+    assert prep(formulae["Weight_isolation_up_test"]) == prep("(PILEUP)*(IsoUp)*(Blahblah)")
+    assert prep(formulae["Weight_another_left_test"]) == prep("(PILEUP)*(Iso)*(BlahblahLeft)")