Fix lwtnn schema validation and add FastSim scale factor example (#343)

* Accept LWTNN nodes in schema validation

* Simplify tests / examples

---------

Co-authored-by: Nick Smith <nick.smith@cern.ch>

Author: sbein

src/correctionlib/schemav2.py

@@ -319,7 +319,7 @@ def validate_content(cls, content: list[CategoryItem]) -> list[CategoryItem]:
 def walk_content(content: Content, func: Callable[[Content], None]) -> None:
     """Visit all content nodes in a tree, applying func to each node."""
     func(content)
-    if isinstance(content, (float, Formula, FormulaRef, HashPRNG)):
+    if isinstance(content, (float, Formula, FormulaRef, HashPRNG, LWTNN)):
         pass
     elif isinstance(content, (Binning, MultiBinning)):
         for bin in content.content:
@@ -354,6 +354,12 @@ def _validate_input(allowed_names: set[str], node: Content) -> None:
             if inp not in allowed_names:
                 msg = f"{nodename} input {inp!r} not found in Correction inputs {allowed_names}"
                 raise ValueError(msg)
+    elif isinstance(node, LWTNN):
+        for inp in node.opaque.get("inputs", []):
+            name = inp.get("name")
+            if name not in allowed_names:
+                msg = f"{nodename} input {name!r} not found in Correction inputs {allowed_names}"
+                raise ValueError(msg)
     # FormulaRef has no direct input names
 
 

src/lwtnn_demo.cc

@@ -4,33 +4,25 @@
 
 #include "correction.h"
 
-static constexpr double PT_EPS  = 1e-4;
-static constexpr double ISO_EPS = 1e-6;
-
-static double safe_log10(double x, double eps) {
-  return std::log10(std::max(x, eps));
-}
-
-int main(int argc, char** argv) {
-  if (argc != 2) {
+int main(int argc, char **argv)
+{
+  if (argc != 2)
+  {
     std::cerr << "Usage: " << argv[0] << " lwtnn_correction.json\n";
     return 1;
   }
 
   const std::string json_path = argv[1];
   auto correction_set = correction::CorrectionSet::from_file(json_path);
-  auto correction = correction_set->at("electron_sf");
+  auto correction = correction_set->at("electron_fastsim_sf");
 
   // Mock "GEN-matched" electron values (replace with real NanoAOD lookup later)
-  const double gen_pt  = 15.0;
+  const double gen_pt = 15.0;
   const double gen_eta = 0.4;
   const double gen_phi = 2.1;
   const double gen_iso = 1e-3;
 
-  const double pt_log10  = safe_log10(gen_pt,  PT_EPS);
-  const double iso_log10 = safe_log10(gen_iso, ISO_EPS);
-
-  double sf = correction->evaluate({pt_log10, gen_eta, gen_phi, iso_log10});
+  double sf = correction->evaluate({gen_pt, gen_eta, gen_phi, gen_iso});
 
   std::cout << std::setprecision(17);
   std::cout << "sf_fullOverFast " << sf << "\n";