diff --git a/src/correction.cc b/src/correction.cc index 267b1cd..2f9ade2 100644 --- a/src/correction.cc +++ b/src/correction.cc @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include // std::abort @@ -72,6 +73,44 @@ namespace { const std::vector& values; }; + // Per-thread scratch storage for Transform::evaluate. + // Depth indexing keeps nested Transform evaluations re-entrant-safe. + class TransformScratch { + public: + explicit TransformScratch(const std::vector& values): + slot_(acquire_slot()) + { + slot_ = values; + } + + ~TransformScratch() { + release_slot(); + } + + std::vector& values() { + return slot_; + } + + private: + static std::vector& acquire_slot() { + if (depth_ == slots_.size()) { + slots_.emplace_back(); + } + return slots_[depth_++]; + } + + static void release_slot() { + depth_--; + } + + std::vector& slot_; + static thread_local std::deque> slots_; + static thread_local std::size_t depth_; + }; + + thread_local std::deque> TransformScratch::slots_; + thread_local std::size_t TransformScratch::depth_ = 0; + std::size_t find_bin_idx(Variable::Type value_variant, const detail::EdgesType &bins_, const detail::FlowBehavior &flow, @@ -110,7 +149,7 @@ namespace { // otherwise we have non-uniform binning using namespace std::string_literals; - const auto bins = std::get(bins_); + const auto& bins = std::get(bins_); if ( flow == detail::FlowBehavior::wrap ) { double low = bins[0]; double high = bins[bins.size() - 1]; @@ -309,7 +348,9 @@ Transform::Transform(const JSONObject& json, const Correction& context) { } double Transform::evaluate(const std::vector& values) const { - std::vector new_values(values); + TransformScratch scratch(values); + auto& new_values = scratch.values(); + double vnew = std::visit(node_evaluate{values}, *rule_); auto& v = new_values[variableIdx_]; if ( std::holds_alternative(v) ) { @@ -558,27 +599,25 @@ Category::Category(const JSONObject& json, const Correction& context) double Category::evaluate(const std::vector& values) const { const Content* child = nullptr; if ( auto pval = std::get_if(&values[variableIdx_]) ) { - try { - child = &std::get(map_).at(*pval); - } catch (std::out_of_range& ex) { - if ( default_ ) { - child = default_.get(); - } - else { - throw std::out_of_range("Index not available in Category for input argument " + std::to_string(variableIdx_) + " val: " + *pval); - } + const auto& m = std::get(map_); + auto it = m.find(*pval); + if ( it != m.end() ) { + child = &it->second; + } else if ( default_ ) { + child = default_.get(); + } else { + throw std::out_of_range("Index not available in Category for input argument " + std::to_string(variableIdx_) + " val: " + *pval); } } else if ( auto pval = std::get_if(&values[variableIdx_]) ) { - try { - child = &std::get(map_).at(*pval); - } catch (std::out_of_range& ex) { - if ( default_ ) { - child = default_.get(); - } - else { - throw std::out_of_range("Index not available in Category for input argument " + std::to_string(variableIdx_) + " val: " + std::to_string(*pval)); - } + const auto& m = std::get(map_); + auto it = m.find(*pval); + if ( it != m.end() ) { + child = &it->second; + } else if ( default_ ) { + child = default_.get(); + } else { + throw std::out_of_range("Index not available in Category for input argument " + std::to_string(variableIdx_) + " val: " + std::to_string(*pval)); } } else { throw std::runtime_error("Invalid variable type"); @@ -682,6 +721,11 @@ size_t CompoundCorrection::input_index(const std::string_view name) const { } double CompoundCorrection::evaluate(const std::vector& values) const { + // Per-thread scratch storage. This call site is not re-entrant so we + // can use a simpler implementation than for TransformScratch + static thread_local std::vector ivalues; + static thread_local std::vector cvalues; + if ( values.size() != inputs_.size() ) { throw std::invalid_argument("Incorrect number of inputs (got " + std::to_string(values.size()) + ", expected " + std::to_string(inputs_.size()) + ")"); @@ -689,9 +733,9 @@ double CompoundCorrection::evaluate(const std::vector& values) c for (size_t i=0; i < inputs_.size(); ++i) { inputs_[i].validate(values[i]); } - std::vector ivalues(values); - std::vector cvalues; + ivalues = values; cvalues.reserve(values.size()); + double out = 0.; double sf = 0.; bool start{true}; diff --git a/tests/test_transform_nested.py b/tests/test_transform_nested.py new file mode 100644 index 0000000..a581584 --- /dev/null +++ b/tests/test_transform_nested.py @@ -0,0 +1,70 @@ +import correctionlib._core as core +from correctionlib import schemav2 as schema + + +def wrap(*corrs): + cset = schema.CorrectionSet( + schema_version=schema.VERSION, + corrections=list(corrs), + ) + return core.CorrectionSet.from_string(cset.model_dump_json()) + + +def test_transform_nested_rule_and_content(): + # Build a nested transform tree with transforms in both the rule and + # content paths to exercise recursive evaluation and scratch-buffer reuse. + cset = wrap( + schema.Correction( + name="nested_transform", + version=2, + inputs=[ + schema.Variable(name="x", type="real"), + schema.Variable(name="y", type="real"), + ], + output=schema.Variable(name="out", type="real"), + data=schema.Transform( + nodetype="transform", + input="x", + rule=schema.Transform( + nodetype="transform", + input="y", + rule=1.0, + content=schema.Formula( + nodetype="formula", + expression="x + y", + parser="TFormula", + variables=["x", "y"], + ), + ), + content=schema.Transform( + nodetype="transform", + input="y", + rule=schema.Transform( + nodetype="transform", + input="x", + rule=2.0, + content=schema.Formula( + nodetype="formula", + expression="x + y", + parser="TFormula", + variables=["x", "y"], + ), + ), + content=schema.Formula( + nodetype="formula", + expression="x + y", + parser="TFormula", + variables=["x", "y"], + ), + ), + ), + ) + ) + + corr = cset["nested_transform"] + # x=3, y=4 + # outer rule: transform y->1 then x+y => 3+1 = 4, so x becomes 4 + # outer content: transform y with rule: + # inner rule transform x->2 then x+y => 2+4 = 6, so y becomes 6 + # final content: x+y = 4+6 = 10 + assert corr.evaluate(3.0, 4.0) == 10.0