[RF] Also read other types than int into RooAbsCategory from TTree

For `RooAbsReal`, this was already implemented in [1] with the
TreeReadBuffer mechanism. This commit uses the same approach for the
RooAbsCategory as well, such that one can read TTree branches of all
fundamental types to RooFit categories.

Fixes root-project#10278.

[1] a18675a, "Add capability to read ULong64_t + more into RooDataSets."

Author: guitargeek

roofit/roofitcore/inc/RooAbsCategory.h

@@ -33,14 +33,15 @@ class TTree;
 class RooVectorDataStore;
 class Roo1DTable;
 class TIterator;
+struct TreeReadBuffer; /// A space to attach TBranches
 
 class RooAbsCategory : public RooAbsArg {
 public:
   /// The type used to denote a specific category state.
   using value_type = int;
 
   // Constructors, assignment etc.
-  RooAbsCategory() { };
+  RooAbsCategory();
   RooAbsCategory(const char *name, const char *title);
   RooAbsCategory(const RooAbsCategory& other, const char* name=0) ;
   ~RooAbsCategory() override;
@@ -215,13 +216,14 @@ class RooAbsCategory : public RooAbsArg {
   mutable value_type _currentIndex{std::numeric_limits<int>::min()}; ///< Current category state
   std::map<std::string, value_type> _stateNames;                     ///< Map state names to index numbers. Make sure state names are updated in recomputeShape().
   std::vector<std::string> _insertionOrder;                          ///< Keeps track in which order state numbers have been inserted. Make sure this is updated in recomputeShape().
-  mutable UChar_t _byteValue{0};                                     ///<! Transient cache for byte values from tree branches
   mutable std::map<value_type, std::unique_ptr<RooCatType, std::function<void(RooCatType*)>> > _legacyStates; ///<! Map holding pointers to RooCatType instances. Only for legacy interface. Don't use if possible.
-  bool _treeVar{false}; ///< Is this category attached to a tree?
 
   static const decltype(_stateNames)::value_type& invalidCategory();
 
-  ClassDefOverride(RooAbsCategory, 3) // Abstract discrete variable
+private:
+  std::unique_ptr<TreeReadBuffer> _treeReadBuffer; //! A buffer for reading values from trees
+
+  ClassDefOverride(RooAbsCategory, 4) // Abstract discrete variable
 };
 
 #endif

roofit/roofitcore/src/RooAbsCategory.cxx

@@ -52,13 +52,15 @@ the following replacements should be used:
 #include "RooMsgService.h"
 #include "RooVectorDataStore.h"
 #include "RooFitLegacy/RooAbsCategoryLegacyIterator.h"
+#include "TreeReadBuffer.h"
 
 #include "Compression.h"
 #include "TString.h"
 #include "TTree.h"
 #include "TLeaf.h"
 #include "TBranch.h"
 
+#include <functional>
 #include <memory>
 
 using namespace std;
@@ -72,6 +74,10 @@ const decltype(RooAbsCategory::_stateNames)::value_type& RooAbsCategory::invalid
   return invalid;
 }
 
+
+RooAbsCategory::RooAbsCategory() {}
+
+
 ////////////////////////////////////////////////////////////////////////////////
 /// Constructor
 
@@ -90,8 +96,7 @@ RooAbsCategory::RooAbsCategory(const char *name, const char *title) :
 RooAbsCategory::RooAbsCategory(const RooAbsCategory& other,const char* name) :
   RooAbsArg(other,name),  _currentIndex(other._currentIndex),
   _stateNames(other._stateNames),
-  _insertionOrder(other._insertionOrder),
-  _treeVar(other._treeVar)
+  _insertionOrder(other._insertionOrder)
 {
   setValueDirty() ;
   setShapeDirty() ;
@@ -439,41 +444,65 @@ void RooAbsCategory::attachToVStore(RooVectorDataStore& vstore)
 /// Attach the category index and label as branches to the given
 /// TTree. The index field will be attached as integer with name
 /// `<name>_idx`. If a branch `<name>` exists, it attaches to this branch.
-void RooAbsCategory::attachToTree(TTree& t, Int_t bufSize)
+void RooAbsCategory::attachToTree(TTree& tree, Int_t bufSize)
 {
   // First check if there is an integer branch matching the category name
   TString cleanName(cleanBranchName()) ;
-  TBranch* branch = t.GetBranch(cleanName) ;
+  TBranch* branch = tree.GetBranch(cleanName) ;
   if (!branch) {
     cleanName += "_idx";
-    branch = t.GetBranch(cleanName);
+    branch = tree.GetBranch(cleanName);
   }
 
   if (branch) {
-    TString typeName(((TLeaf*)branch->GetListOfLeaves()->At(0))->GetTypeName()) ;
-    if (!typeName.CompareTo("Int_t")) {
-      // Imported TTree: attach only index field as branch
-
-      coutI(DataHandling) << "RooAbsCategory::attachToTree(" << GetName() << ") TTree branch " << GetName()
-           << " will be interpreted as category index" << endl ;
-
-      t.SetBranchAddress(cleanName, &_currentIndex) ;
-      setAttribute("INTIDXONLY_TREE_BRANCH",kTRUE) ;
-      _treeVar = true;
-      return ;
-    } else if (!typeName.CompareTo("UChar_t")) {
-      coutI(DataHandling) << "RooAbsReal::attachToTree(" << GetName() << ") TTree UChar_t branch " << GetName()
-           << " will be interpreted as category index" << endl ;
-      t.SetBranchAddress(cleanName,&_byteValue) ;
-      setAttribute("UCHARIDXONLY_TREE_BRANCH",kTRUE) ;
-      _treeVar = true;
+    TLeaf* leaf = (TLeaf*)branch->GetListOfLeaves()->At(0) ;
+
+    // Check that leaf is _not_ an array
+    Int_t dummy ;
+    TLeaf* counterLeaf = leaf->GetLeafCounter(dummy) ;
+    if (counterLeaf) {
+      coutE(Eval) << "RooAbsCategory::attachToTree(" << GetName() << ") ERROR: TTree branch " << GetName()
+        << " is an array and cannot be attached to a RooAbsCategory" << endl ;
       return ;
     }
+
+    TString typeName(leaf->GetTypeName()) ;
+
+
+    // For different type names, store a function to attach
+    std::map<std::string, std::function<std::unique_ptr<TreeReadBuffer>()>> typeMap {
+      {"Float_t",   [&](){ return createTreeReadBuffer<Float_t  >(cleanName, tree); }},
+      {"Double_t",  [&](){ return createTreeReadBuffer<Double_t >(cleanName, tree); }},
+      {"UChar_t",   [&](){ return createTreeReadBuffer<UChar_t  >(cleanName, tree); }},
+      {"Bool_t",    [&](){ return createTreeReadBuffer<Bool_t   >(cleanName, tree); }},
+      {"Char_t",    [&](){ return createTreeReadBuffer<Char_t   >(cleanName, tree); }},
+      {"UInt_t",    [&](){ return createTreeReadBuffer<UInt_t   >(cleanName, tree); }},
+      {"Long64_t",  [&](){ return createTreeReadBuffer<Long64_t >(cleanName, tree); }},
+      {"ULong64_t", [&](){ return createTreeReadBuffer<ULong64_t>(cleanName, tree); }},
+      {"Short_t",   [&](){ return createTreeReadBuffer<Short_t  >(cleanName, tree); }},
+      {"UShort_t",  [&](){ return createTreeReadBuffer<UShort_t >(cleanName, tree); }},
+    };
+
+    auto typeDetails = typeMap.find(typeName.Data());
+    if (typeDetails != typeMap.end()) {
+      coutI(DataHandling) << "RooAbsCategory::attachToTree(" << GetName() << ") TTree " << typeName << " branch \"" << cleanName
+                  << "\" will be converted to int." << endl ;
+      _treeReadBuffer = typeDetails->second();
+    } else {
+      _treeReadBuffer = nullptr;
+
+      if (!typeName.CompareTo("Int_t")) {
+        tree.SetBranchAddress(cleanName, &_currentIndex);
+      }
+      else {
+        coutE(InputArguments) << "RooAbsCategory::attachToTree(" << GetName() << ") data type " << typeName << " is not supported." << endl ;
+      }
+    }
   } else {
     TString format(cleanName);
     format.Append("/I");
     void* ptr = &_currentIndex;
-    t.Branch(cleanName, ptr, (const Text_t*)format, bufSize);
+    tree.Branch(cleanName, ptr, (const Text_t*)format, bufSize);
   }
 }
 
@@ -533,33 +562,17 @@ void RooAbsCategory::copyCache(const RooAbsArg *source, Bool_t /*valueOnly*/, Bo
    auto other = static_cast<const RooAbsCategory*>(source);
    assert(dynamic_cast<const RooAbsCategory*>(source));
 
-   _currentIndex = other->_currentIndex;
-
-   if (setValDirty) {
-     setValueDirty();
-   }
-
-   if (!_treeVar)
-     return;
-
-   if (source->getAttribute("INTIDXONLY_TREE_BRANCH")) {
-     // Lookup cat state from other-index because label is missing
-     if (hasIndex(other->_currentIndex)) {
-       _currentIndex = other->_currentIndex;
-     } else {
-       coutE(DataHandling) << "RooAbsCategory::copyCache(" << GetName() << ") ERROR: index of source arg "
-           << source->GetName() << " is invalid (" << other->_currentIndex
-           << "), value not updated" << endl;
-     }
-   } else if (source->getAttribute("UCHARIDXONLY_TREE_BRANCH")) {
-     // Lookup cat state from other-index because label is missing
-     Int_t tmp = static_cast<int>(other->_byteValue);
-     if (hasIndex(tmp)) {
-       _currentIndex = tmp;
-     } else {
-       coutE(DataHandling) << "RooAbsCategory::copyCache(" << GetName() << ") ERROR: index of source arg "
-           << source->GetName() << " is invalid (" << tmp << "), value not updated" << endl;
+   value_type tmp = other->_treeReadBuffer ? *other->_treeReadBuffer : other->_currentIndex;
+   // Lookup cat state from other-index because label is missing
+   if (hasIndex(tmp)) {
+     _currentIndex = tmp;
+     if (setValDirty) {
+       setValueDirty();
      }
+   } else {
+     coutE(DataHandling) << "RooAbsCategory::copyCache(" << GetName() << ") ERROR: index of source arg "
+         << source->GetName() << " is invalid (" << other->_currentIndex
+         << "), value not updated" << endl;
    }
 }
 

roofit/roofitcore/src/RooAbsReal.cxx

@@ -84,6 +84,7 @@
 #include "RooCachedReal.h"
 #include "RooHelpers.h"
 #include "RunContext.h"
+#include "TreeReadBuffer.h"
 #include "ValueChecking.h"
 
 #include "ROOT/StringUtils.hxx"
@@ -3205,12 +3206,6 @@ RooAbsFunc *RooAbsReal::bindVars(const RooArgSet &vars, const RooArgSet* nset, B
 
 
 
-struct TreeReadBuffer {
-  virtual ~TreeReadBuffer() = default;
-  virtual operator double() = 0;
-};
-
-
 ////////////////////////////////////////////////////////////////////////////////
 /// Copy the cached value of another RooAbsArg to our cache.
 /// Warning: This function just copies the cached values of source,
@@ -3238,30 +3233,6 @@ void RooAbsReal::attachToVStore(RooVectorDataStore& vstore)
 }
 
 
-namespace {
-/// Helper for reading branches with various types from a TTree, and convert all to double.
-template<typename T>
-struct TypedTreeReadBuffer final : public TreeReadBuffer {
-  operator double() override {
-    return _value;
-  }
-  T _value;
-};
-
-/// Create a TreeReadBuffer to hold the specified type, and attach to the branch passed as argument.
-/// \tparam T Type of branch to be read.
-/// \param[in] branchName Attach to this branch.
-/// \param[in] tree Tree to attach to.
-template<typename T>
-std::unique_ptr<TreeReadBuffer> createTreeReadBuffer(const TString& branchName, TTree& tree) {
-  auto buf = new TypedTreeReadBuffer<T>();
-  tree.SetBranchAddress(branchName.Data(), &buf->_value);
-  return std::unique_ptr<TreeReadBuffer>(buf);
-}
-
-}
-
-
 ////////////////////////////////////////////////////////////////////////////////
 /// Attach object to a branch of given TTree. By default it will
 /// register the internal value cache RooAbsReal::_value as branch

roofit/roofitcore/src/TreeReadBuffer.h

@@ -0,0 +1,44 @@
+/*
+ * Project: RooFit
+ * Authors:
+ *   Stephan Hageboeck, CERN 2020
+ *
+ * Copyright (c) 2022, CERN
+ *
+ * Redistribution and use in source and binary forms,
+ * with or without modification, are permitted according to the terms
+ * listed in LICENSE (http://roofit.sourceforge.net/license.txt)
+ */
+
+#ifndef RooFit_TreeReadBuffer_h
+#define RooFit_TreeReadBuffer_h
+
+#include <TTree.h>
+
+struct TreeReadBuffer {
+   virtual ~TreeReadBuffer() = default;
+   virtual operator double() = 0;
+   virtual operator int() = 0;
+};
+
+/// Helper for reading branches with various types from a TTree, and convert all to double.
+template <typename T>
+struct TypedTreeReadBuffer final : public TreeReadBuffer {
+   operator double() override { return _value; }
+   operator int() override { return _value; }
+   T _value;
+};
+
+/// Create a TreeReadBuffer to hold the specified type, and attach to the branch passed as argument.
+/// \tparam T Type of branch to be read.
+/// \param[in] branchName Attach to this branch.
+/// \param[in] tree Tree to attach to.
+template <typename T>
+std::unique_ptr<TreeReadBuffer> createTreeReadBuffer(const TString &branchName, TTree &tree)
+{
+   auto buf = new TypedTreeReadBuffer<T>();
+   tree.SetBranchAddress(branchName.Data(), &buf->_value);
+   return std::unique_ptr<TreeReadBuffer>(buf);
+}
+
+#endif

roofit/roofitcore/test/CMakeLists.txt

@@ -9,6 +9,7 @@
 
 ROOT_ADD_GTEST(simple simple.cxx LIBRARIES RooFitCore)
 ROOT_ADD_GTEST(testRooCacheManager testRooCacheManager.cxx LIBRARIES RooFitCore)
+ROOT_ADD_GTEST(testRooCategory testRooCategory.cxx LIBRARIES RooFitCore)
 ROOT_ADD_GTEST(testWorkspace testWorkspace.cxx LIBRARIES RooFitCore RooStats)
 if(NOT MSVC OR win_broken_tests)
   ROOT_ADD_GTEST(testRooDataHist testRooDataHist.cxx LIBRARIES RooFitCore

roofit/roofitcore/test/testRooCategory.cxx

@@ -0,0 +1,26 @@
+// Tests for the RooCategory
+// Author: Jonas Rembser, CERN  04/2021
+
+#include <RooCategory.h>
+#include <RooDataSet.h>
+#include <RooGlobalFunc.h>
+
+#include <TTree.h>
+
+#include <gtest/gtest.h>
+
+// GitHub issue 10278: RooDataSet incorrectly loads RooCategory values from TTree branch of type Short_t
+TEST(RooCategory, CategoryDefineMultiState)
+{
+   TTree tree("test_tree", "Test tree");
+   Short_t cat_in;
+   tree.Branch("cat", &cat_in);
+
+   cat_in = 2; // category B
+   tree.Fill();
+
+   RooCategory cat("cat", "Category", {{"B_cat", 2}, {"A_cat", 3}});
+   RooDataSet data("data", "RooDataSet", RooArgSet(cat), RooFit::Import(tree));
+
+   EXPECT_EQ(static_cast<RooCategory &>((*data.get(0))["cat"]).getCurrentIndex(), 2);
+}