Add Switch node to correctionlib for flexible logic

include/correction.h

@@ -41,7 +41,8 @@ class HashPRNG;
 class Binning;
 class MultiBinning;
 class Category;
-typedef std::variant<double, Formula, FormulaRef, Transform, HashPRNG, Binning, MultiBinning, Category> Content;
+class Switch;
+typedef std::variant<double, Formula, FormulaRef, Transform, HashPRNG, Binning, MultiBinning, Category, Switch> Content;
 class Correction;
 
 class FormulaAst {
@@ -234,6 +235,23 @@ class Category {
     size_t variableIdx_;
 };
 
+class Switch {
+  public:
+    Switch(const JSONObject& json, const Correction& context);
+    double evaluate(const std::vector<Variable::Type>& values) const;
+
+  private:
+    struct Selection {
+      size_t variableIdx;
+      std::string op;
+      double value;
+      std::unique_ptr<const Content> content;
+    };
+
+    std::vector<Selection> selections_;
+    std::unique_ptr<const Content> default_;
+};
+
 class Correction {
   public:
     typedef std::shared_ptr<const Correction> Ref;

src/correction.cc

@@ -130,6 +130,7 @@ namespace {
       else if ( type == "formularef" ) { return FormulaRef(obj, context); }
       else if ( type == "transform" ) { return Transform(obj, context); }
       else if ( type == "hashprng" ) { return HashPRNG(obj, context); }
+      else if ( type == "switch" ) { return Switch(obj, context); }
       else { throw std::runtime_error("Unrecognized Content object nodetype"); }
     }
     throw std::runtime_error("Invalid Content node type");
@@ -670,6 +671,62 @@ double Category::evaluate(const std::vector<Variable::Type>& values) const {
   return std::visit(node_evaluate{values}, *child);
 }
 
+Switch::Switch(const JSONObject& json, const Correction& context) {
+  json.getRequired<rapidjson::Value::ConstArray>("inputs");
+
+  for (const auto& item : json.getRequired<rapidjson::Value::ConstArray>("selections")) {
+    if ( ! item.IsObject() ) { throw std::runtime_error("Expected Comparison object in Switch selections"); }
+    JSONObject comp(item.GetObject());
+
+    Selection sel;
+    sel.variableIdx = find_input_index(comp.getRequired<const char*>("variable"), context.inputs());
+
+    if ( context.inputs()[sel.variableIdx].type() == Variable::VarType::string ) {
+      throw std::runtime_error("Switch node currently only supports numeric comparisons");
+    }
+
+    sel.op = comp.getRequired<std::string_view>("op");
+    sel.value = comp.getRequired<double>("value");
+    sel.content = std::make_unique<Content>(resolve_content(comp.getRequiredValue("content"), context));
+
+    selections_.push_back(std::move(sel));
+  }
+
+  default_ = std::make_unique<Content>(resolve_content(json.getRequiredValue("default"), context));
+}
+
+double Switch::evaluate(const std::vector<Variable::Type>& values) const {
+  for (const auto& sel : selections_) {
+    double val = 0.0;
+    const auto& input = values[sel.variableIdx];
+
+    if ( auto v = std::get_if<double>(&input) ) {
+      val = *v;
+    } else if ( auto v = std::get_if<int64_t>(&input) ) {
+      val = static_cast<double>(*v);
+    } else {
+      throw std::runtime_error("Switch node encountered non-numeric input value during evaluation");
+    }
+
+    bool pass = false;
+
+    if      (sel.op == "<")  pass = (val < sel.value);
+    else if (sel.op == ">")  pass = (val > sel.value);
+    else if (sel.op == "<=") pass = (val <= sel.value);
+    else if (sel.op == ">=") pass = (val >= sel.value);
+    else if (sel.op == "==") pass = (val == sel.value);
+    else if (sel.op == "!=") pass = (val != sel.value);
+    else {
+      throw std::runtime_error("Unknown operator in Switch node: " + sel.op);
+    }
+
+    if (pass) {
+      return std::visit(node_evaluate{values}, *sel.content);
+    }
+  }
+  return std::visit(node_evaluate{values}, *default_);
+}
+
 Correction::Correction(const JSONObject& json) :
   name_(json.getRequired<const char *>("name")),
   description_(json.getOptional<const char*>("description").value_or("")),

src/correctionlib/schemav2.py

@@ -61,6 +61,7 @@ def __rich__(self) -> str:
             "FormulaRef",
             "Transform",
             "HashPRNG",
+            "Switch",
         ],
         Field(discriminator="nodetype"),
     ],
@@ -296,11 +297,33 @@ def validate_content(cls, content: list[CategoryItem]) -> list[CategoryItem]:
         return content
 
 
+class Comparison(Model):
+    variable: str = Field(description="The name of the input variable")
+    op: Literal[">", "<", ">=", "<=", "==", "!="] = Field(
+        description="Comparison operator"
+    )
+    value: float = Field(description="Value to compare against")
+    content: Content = Field(description="Content to return if comparison is true")
+
+
+class Switch(Model):
+    nodetype: Literal["switch"]
+    inputs: list[str] = Field(
+        description="The names of the input variables used in the selections"
+    )
+    selections: list[Comparison] = Field(
+        description="List of checks to perform. First one to evaluate true returns its content."
+    )
+    default: Content = Field(description="Default content if no selection matches")
+
+
 Transform.model_rebuild()
 Binning.model_rebuild()
 MultiBinning.model_rebuild()
 CategoryItem.model_rebuild()
 Category.model_rebuild()
+Comparison.model_rebuild()
+Switch.model_rebuild()
 
 
 def walk_content(content: Content, func: Callable[[Content], None]) -> None:
@@ -321,6 +344,10 @@ def walk_content(content: Content, func: Callable[[Content], None]) -> None:
     elif isinstance(content, Transform):
         walk_content(content.rule, func)
         walk_content(content.content, func)
+    elif isinstance(content, Switch):
+        for selection in content.selections:
+            walk_content(selection.content, func)
+        walk_content(content.default, func)
     else:
         raise RuntimeError(f"Unknown content node type: {type(content)}")
 
@@ -331,7 +358,7 @@ def _validate_input(allowed_names: set[str], node: Content) -> None:
         if node.input not in allowed_names:
             msg = f"{nodename} input {node.input!r} not found in Correction inputs {allowed_names}"
             raise ValueError(msg)
-    elif isinstance(node, (MultiBinning, HashPRNG)):
+    elif isinstance(node, (MultiBinning, HashPRNG, Switch)):
         for inp in node.inputs:
             if inp not in allowed_names:
                 msg = f"{nodename} input {inp!r} not found in Correction inputs {allowed_names}"

tests/test_switch.py

@@ -0,0 +1,48 @@
+import json
+
+import correctionlib
+import correctionlib.schemav2 as cs
+
+corr = cs.Correction(
+    name="boundary_test",
+    version=1,
+    inputs=[cs.Variable(name="eta", type="real")],
+    output=cs.Variable(name="weight", type="real"),
+    data=cs.Switch(
+        nodetype="switch",
+        inputs=["eta"],
+        selections=[
+            cs.Comparison(
+                variable="eta",
+                op="<=",
+                value=3.0,
+                content=1.0,
+            )
+        ],
+        default=0.0,
+    ),
+)
+
+json_str = corr.model_dump_json(exclude_unset=True)
+print(f"Generated JSON with Switch node:\n{json_str}\n")
+
+cset_json = json.dumps({"schema_version": 2, "corrections": [json.loads(json_str)]})
+cset = correctionlib.CorrectionSet.from_string(cset_json)
+evaluator = cset["boundary_test"]
+
+print("-" * 40)
+print("TESTING C++ EVALUATOR")
+print("-" * 40)
+
+val_inclusive = 3.0
+res_inclusive = evaluator.evaluate([val_inclusive])
+print(f"Input: {val_inclusive:<10} | Expected: 1.0 | Got: {res_inclusive}")
+
+val_exclusive = 3.00001
+res_exclusive = evaluator.evaluate([val_exclusive])
+print(f"Input: {val_exclusive:<10} | Expected: 0.0 | Got: {res_exclusive}")
+
+if res_inclusive == 1.0 and res_exclusive == 0.0:
+    print("\n[SUCCESS] The Switch node correctly handles inclusive boundaries!")
+else:
+    print("\n[FAIL] Logic error in Switch node implementation.")