Add flow: wrap option to handle circular axes (#276)

* New flow behavior `wrap`.

* Add test case for `wrap` flow behavior.

* Factor out flow behavior parsing to helper function.

* Correct arg type

---------

Co-authored-by: Nick Smith <[email protected]>

Author: dsavoiu

include/correction.h

@@ -175,7 +175,7 @@ class HashPRNG {
 };
 
 // common internal for Binning and MultiBinning
-enum class _FlowBehavior {value, clamp, error};
+enum class _FlowBehavior {value, clamp, error, wrap};
 
 using _NonUniformBins = std::vector<double>;
 

src/correction.cc

@@ -25,6 +25,22 @@
 
 using namespace correction;
 
+//! helper function for parsing flow behavior from string
+_FlowBehavior parse_flowbehavior(const rapidjson::Value& flowbehavior) {
+  if ( flowbehavior == "clamp" ) {
+    return _FlowBehavior::clamp;
+  }
+  else if ( flowbehavior == "error" ) {
+    return _FlowBehavior::error;
+  }
+  else if ( flowbehavior == "wrap" ) {
+    return _FlowBehavior::wrap;
+  }
+  else {
+    return _FlowBehavior::value;
+  }
+}
+
 class correction::JSONObject {
   public:
     JSONObject(rapidjson::Value::ConstObject&& json) : json_(json) { }
@@ -149,20 +165,33 @@ namespace {
             return bins->n; // the default value is stored at the end of the content array, after the last bin
           case _FlowBehavior::clamp:
             return value < bins->low ? 0 : bins->n - 1; // assuming we always have at least 1 bin
+          case _FlowBehavior::wrap:
+            break;
           case _FlowBehavior::error:
             const std::string belowOrAbove = value < bins->low ? "below" : "above";
             auto msg = "Index " + belowOrAbove + " bounds in " + name + " for input argument " + std::to_string(variableIdx) + " value: " + std::to_string(value);
             throw std::runtime_error(std::move(msg));
         }
       }
 
-      std::size_t binIdx = bins->n * ((value - bins->low) / (bins->high - bins->low));
+      double norm_value = ((value - bins->low) / (bins->high - bins->low));
+      if (flow == _FlowBehavior::wrap) {
+        norm_value -= std::floor(norm_value);
+      }
+      std::size_t binIdx = bins->n * norm_value;
       return binIdx;
     }
 
     // otherwise we have non-uniform binning
     using namespace std::string_literals;
     const auto bins = std::get<_NonUniformBins>(bins_);
+    if ( flow == _FlowBehavior::wrap ) {
+      double low = bins[0];
+      double high = bins[bins.size() - 1];
+      double norm_value = (value - low) / (high - low);
+      norm_value -= std::floor(norm_value);
+      value = low + norm_value * (high - low);
+    }
 
     auto it = std::upper_bound(std::begin(bins), std::end(bins), value);
     if ( it == std::begin(bins) ) { // underflow
@@ -172,6 +201,9 @@ namespace {
       else if ( flow == _FlowBehavior::error ) {
         throw std::runtime_error("Index below bounds in "s + name + " for input argument " + std::to_string(variableIdx) + " value: " + std::to_string(value));
       }
+      else if ( flow == _FlowBehavior::wrap ) {
+        throw std::logic_error("I should not have ever seen an underflow");
+      }
       else { // clamp
         it++;
       }
@@ -183,6 +215,9 @@ namespace {
       else if ( flow == _FlowBehavior::error ) {
         throw std::runtime_error("Index above bounds in "s + name + " for input argument " + std::to_string(variableIdx) + " value: " + std::to_string(value));
       }
+      else if ( flow == _FlowBehavior::wrap ) {
+        throw std::logic_error("I should not have ever seen an overflow");
+      }
       else { // clamp
         it--;
       }
@@ -461,15 +496,9 @@ Binning::Binning(const JSONObject& json, const Correction& context)
   }
   Content default_value{0.};
   const auto& flowbehavior = json.getRequiredValue("flow");
-  if ( flowbehavior == "clamp" ) {
-    flow_ = _FlowBehavior::clamp;
-  }
-  else if ( flowbehavior == "error" ) {
-    flow_ = _FlowBehavior::error;
-  }
-  else {
-    flow_ = _FlowBehavior::value;
-    default_value = resolve_content(flowbehavior, context);
+  flow_ = parse_flowbehavior(flowbehavior);
+  if (flow_ == _FlowBehavior::value) {
+      default_value = resolve_content(flowbehavior, context);
   }
 
   // set bin contents
@@ -540,16 +569,9 @@ MultiBinning::MultiBinning(const JSONObject& json, const Correction& context)
   }
 
   const auto& flowbehavior = json.getRequiredValue("flow");
-  if ( flowbehavior == "clamp" ) {
-    flow_ = _FlowBehavior::clamp;
-  }
-  else if ( flowbehavior == "error" ) {
-    flow_ = _FlowBehavior::error;
-  }
-  else {
-    flow_ = _FlowBehavior::value;
-    // store default value at end of content array
-    content_.push_back(resolve_content(flowbehavior, context));
+  flow_ = parse_flowbehavior(flowbehavior);
+  if (flow_ == _FlowBehavior::value) {
+      content_.push_back(resolve_content(flowbehavior, context));
   }
 }
 

src/correctionlib/schemav2.py

@@ -232,7 +232,7 @@ class Binning(Model):
         description="Edges of the binning, either as a list of monotonically increasing floats or as an instance of UniformBinning. edges[i] <= x < edges[i+1] => f(x, ...) = content[i](...)"
     )
     content: List[Content]
-    flow: Union[Content, Literal["clamp", "error"]] = Field(
+    flow: Union[Content, Literal["clamp", "error", "wrap"]] = Field(
         description="Overflow behavior for out-of-bounds values"
     )
 
@@ -287,7 +287,7 @@ class MultiBinning(Model):
         to the element at i0 in dimension 0, i1 in dimension 1, etc. and d0 = len(edges[0])-1, etc.
     """
     )
-    flow: Union[Content, Literal["clamp", "error"]] = Field(
+    flow: Union[Content, Literal["clamp", "error", "wrap"]] = Field(
         description="Overflow behavior for out-of-bounds values"
     )
 

tests/test_core.py

@@ -751,6 +751,17 @@ def binning(flow, uniform=True):
         assert corr.evaluate(2.9) == 2.0
         assert corr.evaluate(3.0) == 42.0
 
+        corr = binning(flow="wrap", uniform=use_uniform_binning)
+        assert corr.evaluate(-3.0) == 1.0
+        assert corr.evaluate(-2.9) == 1.0
+        assert corr.evaluate(-1.0) == 2.0
+        assert corr.evaluate(0.0) == 1.0
+        assert corr.evaluate(1.0) == 1.0 if use_uniform_binning else 2.0
+        assert corr.evaluate(2.9) == 2.0
+        assert corr.evaluate(3.0) == 1.0
+        assert corr.evaluate(4.6) == 2.0
+        assert corr.evaluate(6.1) == 1.0
+
     def multibinning(flow, uniform=True):
         if uniform:
             edges_x = schema.UniformBinning(n=2, low=0.0, high=3.0)