RooCompositeDataStore.cxx

Go to the documentation of this file.

/*****************************************************************************
 * Project: RooFit                                                           *
 * Package: RooFitCore                                                       *
 * @(#)root/roofitcore:$Id$
 * Authors:                                                                  *
 *   WV, Wouter Verkerke, UC Santa Barbara, verkerke@slac.stanford.edu       *
 *   DK, David Kirkby,    UC Irvine,         dkirkby@uci.edu                 *
 *                                                                           *
 * Copyright (c) 2000-2005, Regents of the University of California          *
 *                          and Stanford University. All rights reserved.    *
 *                                                                           *
 * 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)             *
 *****************************************************************************/
 
/**
\file RooCompositeDataStore.cxx
\class RooCompositeDataStore
\ingroup Roofitcore
 
Combines several disjunct datasets into one. This is useful for simultaneous PDFs
that do not depend on the same observable such as a PDF depending on `x` combined with another one depending
on `y`.
The composite storage will store two different datasets, `{x}` and `{y}`, but they can be passed as a single
dataset to RooFit operations. A category tag will define which dataset has to be passed to which likelihood.
 
When iterated from start to finish, datasets will be traversed in the order of the category index.
**/
 
#include "RooCompositeDataStore.h"
 
#include "RooMsgService.h"
#include "RooFormulaVar.h"
#include "RooRealVar.h"
#include "RooCategory.h"
 
#include <algorithm>
#include <iomanip>
#include <iostream>
 
using std::map, std::list, std::string;
 
 
 
////////////////////////////////////////////////////////////////////////////////
 
RooCompositeDataStore::RooCompositeDataStore()
{
}
 
 
 
////////////////////////////////////////////////////////////////////////////////
/// Convert map by label to map by index for more efficient internal use
 
RooCompositeDataStore::RooCompositeDataStore(
        RooStringView name, RooStringView title,
        const RooArgSet& vars, RooCategory& indexCat,map<std::string,RooAbsDataStore*> const& inputData) :
  RooAbsDataStore(name,title,RooArgSet(vars,indexCat)), _indexCat(&indexCat)
{
  for (const auto& iter : inputData) {
    const RooAbsCategory::value_type idx = indexCat.lookupIndex(iter.first);
    _dataMap[idx] = iter.second;
  }
}
 
 
////////////////////////////////////////////////////////////////////////////////
/// Convert map by label to map by index for more efficient internal use
 
RooCompositeDataStore::RooCompositeDataStore(const RooCompositeDataStore& other, const char* newname) :
  RooAbsDataStore(other,newname), _indexCat(other._indexCat), _curStore(other._curStore), _curIndex(other._curIndex), _ownComps(true)
{
  for (const auto& item : other._dataMap) {
    RooAbsDataStore* clonedata = item.second->clone() ;
    _dataMap[item.first] = clonedata ;
  }
}
 
 
////////////////////////////////////////////////////////////////////////////////
/// Update index category pointer, if it is contained in input argument vars
 
RooCompositeDataStore::RooCompositeDataStore(const RooCompositeDataStore& other, const RooArgSet& vars, const char* newname) :
  RooAbsDataStore(other,vars,newname), _indexCat(other._indexCat), _curStore(other._curStore), _curIndex(other._curIndex), _ownComps(true)
{
  RooCategory* newIdx = static_cast<RooCategory*>(vars.find(other._indexCat->GetName())) ;
  if (newIdx) {
    _indexCat = newIdx ;
  }
 
  // Convert map by label to map by index for more efficient internal use
  for (const auto& item : other._dataMap) {
    RooAbsDataStore* clonedata = item.second->clone(vars) ;
    _dataMap[item.first] = clonedata ;
  }
}
 
 
 
 
////////////////////////////////////////////////////////////////////////////////
/// Destructor
 
RooCompositeDataStore::~RooCompositeDataStore()
{
  if (_ownComps) {
    for (auto& item : _dataMap) {
      delete item.second;
    }
  }
}
 
 
std::unique_ptr<RooAbsDataStore> RooCompositeDataStore::reduce(
        RooStringView name, RooStringView title, const RooArgSet& vars, const RooFormulaVar* cutVar,
        const char* cutRange, std::size_t nStart, std::size_t nStop)
{
  // for the components, we need the set of variables without the index category
  RooArgSet varsNoIndex{vars};
  if (RooAbsArg * indexCat = varsNoIndex.find(*_indexCat)) {
    varsNoIndex.remove(*indexCat,true) ;
  }
 
  // create an empty RooCompositeDataStore
  auto out = std::make_unique<RooCompositeDataStore>(name, title, varsNoIndex, *_indexCat, std::map<std::string,RooAbsDataStore*>{});
 
  RooCategory* indexCatForCut = nullptr;
  RooAbsCategory::value_type initialIndex = 0;
  if (cutVar) {
    // find _indexCat in the formula recursive servers, if present
    RooArgSet formulaObservables;
    cutVar->getObservables(&_vars, formulaObservables);
    indexCatForCut = dynamic_cast<RooCategory*>(formulaObservables.find(*_indexCat));
    if (indexCatForCut) {
      // save initial index to restore it later
      initialIndex = indexCatForCut->getCurrentIndex();
    }
  }
 
  // fill it with reduced versions of components
  for (const auto& item : _dataMap) {
    if (indexCatForCut) indexCatForCut->setIndex(item.first);
    out->_dataMap[item.first] = item.second->reduce(name, title, varsNoIndex, cutVar, cutRange, nStart, nStop).release();
  }
  if (indexCatForCut) {
    // restore initial index
    indexCatForCut->setIndex(initialIndex);
  }
 
  // indiceate component ownership and return
  out->_ownComps = true;
  return out;
}
 
 
////////////////////////////////////////////////////////////////////////////////
/// Forward fill request to appropriate subset
 
Int_t RooCompositeDataStore::fill()
{
  RooAbsDataStore* subset = _dataMap[_indexCat->getCurrentIndex()] ;
  const_cast<RooArgSet*>((subset->get()))->assignValueOnly(_vars) ;
  return subset->fill() ;
}
 
 
 
////////////////////////////////////////////////////////////////////////////////
/// Forward fill request to appropriate subset
 
double RooCompositeDataStore::sumEntries() const
{
  double sum(0) ;
 
  for (auto const& item : _dataMap) {
    sum+= item.second->sumEntries() ;
  }
  return sum ;
}
 
 
 
////////////////////////////////////////////////////////////////////////////////
/// Load the n-th data point (n='idx') in memory
/// and return a pointer to the internal RooArgSet
/// holding its coordinates.
 
const RooArgSet* RooCompositeDataStore::get(Int_t idx) const
{
  Int_t offset(0) ;
  for (auto const& item : _dataMap) {
    if (idx>=(offset+item.second->numEntries())) {
      offset += item.second->numEntries() ;
      continue ;
    }
    _vars.assign(*item.second->get(idx-offset)) ;
 
    _indexCat->setIndex(item.first) ;
    _curStore = item.second ;
    _curIndex = idx-offset ;
 
    return &_vars ;
  }
  return nullptr ;
}
 
 
 
////////////////////////////////////////////////////////////////////////////////
 
double RooCompositeDataStore::weight() const
{
  if (!_curStore) get(0) ;
  // coverity[FORWARD_NULL]
  return _curStore->weight(_curIndex) ;
}
 
 
////////////////////////////////////////////////////////////////////////////////
 
double RooCompositeDataStore::weightError(RooAbsData::ErrorType etype) const
{
  if (!_curStore) get(0) ;
  // coverity[FORWARD_NULL]
  return _curStore->weightError(etype) ;
}
 
 
 
 
////////////////////////////////////////////////////////////////////////////////
 
void RooCompositeDataStore::weightError(double& lo, double& hi, RooAbsData::ErrorType etype) const
{
  if (!_curStore) get(0) ;
  // coverity[FORWARD_NULL]
  return _curStore->weightError(lo,hi,etype) ;
}
 
 
 
 
////////////////////////////////////////////////////////////////////////////////
 
bool RooCompositeDataStore::isWeighted() const
{
  for (auto const& item : _dataMap) {
    if (item.second->isWeighted()) return true ;
  }
  return false ;
}
 
 
////////////////////////////////////////////////////////////////////////////////
 
void RooCompositeDataStore::loadValues(const RooAbsDataStore*, const RooFormulaVar*, const char*, std::size_t, std::size_t)
{
  throw(std::runtime_error("RooCompositeDataSore::loadValues() NOT IMPLEMENTED")) ;
}
 
 
 
////////////////////////////////////////////////////////////////////////////////
/// Change name of internal observable named 'from' into 'to'
 
bool RooCompositeDataStore::changeObservableName(const char* from, const char* to)
{
 
  // Find observable to be changed
  RooAbsArg* var = _vars.find(from) ;
 
  // Check that we found it
  if (!var) {
    coutE(InputArguments) << "RooCompositeDataStore::changeObservableName(" << GetName() << " no observable " << from << " in this dataset" << std::endl ;
    return true ;
  }
 
  // Process name change
  var->SetName(to) ;
 
  // Forward name change request to component datasets
  bool ret(false) ;
  for (auto const& item : _dataMap) {
    ret |= item.second->changeObservableName(from,to) ;
  }
 
  return ret ;
}
 
 
 
////////////////////////////////////////////////////////////////////////////////
/// WVE ownership issue here!! Caller (a RooAbsData) should take ownership of all
/// arguments, but only does for the first one here...
 
RooAbsArg* RooCompositeDataStore::addColumn(RooAbsArg& newVar, bool adjustRange)
{
  RooAbsArg* ret(nullptr) ;
  for (auto const& item : _dataMap) {
    ret = item.second->addColumn(newVar,adjustRange) ;
  }
  if (ret) {
    _vars.add(*ret) ;
  }
  return ret ;
}
 
 
////////////////////////////////////////////////////////////////////////////////
 
RooAbsDataStore* RooCompositeDataStore::merge(const RooArgSet& /*allVars*/, list<RooAbsDataStore*> /*dstoreList*/)
{
  throw string("RooCompositeDataStore::merge() is not implemented yet") ;
}
 
 
 
 
 
////////////////////////////////////////////////////////////////////////////////
 
void RooCompositeDataStore::append(RooAbsDataStore& other)
{
  Int_t nevt = other.numEntries() ;
  for (int i=0 ; i<nevt ; i++) {
    _vars.assign(*other.get(i)) ;
    fill() ;
  }
}
 
 
 
////////////////////////////////////////////////////////////////////////////////
 
Int_t RooCompositeDataStore::numEntries() const
{
  Int_t n(0) ;
  for (auto const& item : _dataMap) {
    n += item.second->numEntries() ;
  }
  return n ;
}
 
 
 
 
////////////////////////////////////////////////////////////////////////////////
 
void RooCompositeDataStore::reset()
{
  for (auto const& item : _dataMap) {
    item.second->reset() ;
  }
}
 
 
 
////////////////////////////////////////////////////////////////////////////////
 
void RooCompositeDataStore::attachBuffers(const RooArgSet& extObs)
{
  for (auto const& item : _dataMap) {
    item.second->attachBuffers(extObs);
  }
  return ;
}
 
 
 
////////////////////////////////////////////////////////////////////////////////
 
void RooCompositeDataStore::resetBuffers()
{
  for (auto const& item : _dataMap) {
    item.second->resetBuffers();
  }
  return ;
}
 
 
////////////////////////////////////////////////////////////////////////////////
 
void RooCompositeDataStore::dump()
{
  std::cout << "RooCompositeDataStore::dump()" << std::endl ;
  for (auto const& item : _dataMap) {
    std::cout << "state number " << item.first << " has store " << item.second->ClassName() << " with variables " << *item.second->get() ;
    if (item.second->isWeighted()) std::cout << " and is weighted " ;
    std::cout << std::endl ;
  }
}
 
 
////////////////////////////////////////////////////////////////////////////////
/// Get the weights of the events in the range [first, first+len).
/// This implementation will fill a vector with every event retrieved one by one
/// (even if the weight is constant). Then, it returns a span.
std::span<const double> RooCompositeDataStore::getWeightBatch(std::size_t first, std::size_t len) const {
  if (!_weightBuffer) {
    _weightBuffer = std::make_unique<std::vector<double>>();
    _weightBuffer->reserve(len);
 
    for (std::size_t i = 0; i < static_cast<std::size_t>(numEntries()); ++i) {
      _weightBuffer->push_back(weight(i));
    }
  }
 
  return {_weightBuffer->data() + first, len};
}
 
////////////////////////////////////////////////////////////////////////////////
/// Fill the internal per-column buffers for batch access by loading the
/// composite rows one by one, like getWeightBatch() does for the weights.
/// This also synthesizes a column for the index category, which is not
/// stored in any of the component datasets.
void RooCompositeDataStore::fillBatchBuffers() const
{
   const auto n = static_cast<std::size_t>(numEntries());
 
   if (!_realBatchBuffers.empty() || !_catBatchBuffers.empty()) {
      // Refill from scratch if entries were added or removed since the last
      // fill. Value mutations that keep the number of entries are not
      // detected, like for the weight buffer above.
      const std::size_t nFilled = !_realBatchBuffers.empty() ? _realBatchBuffers.begin()->second.size()
                                                             : _catBatchBuffers.begin()->second.size();
      if (nFilled == n)
         return;
      _realBatchBuffers.clear();
      _catBatchBuffers.clear();
   }
 
   std::vector<std::pair<RooAbsReal const *, std::vector<double> *>> realCols;
   std::vector<std::pair<RooAbsCategory const *, std::vector<RooAbsCategory::value_type> *>> catCols;
 
   for (RooAbsArg const *arg : _vars) {
      if (auto cat = dynamic_cast<RooAbsCategory const *>(arg)) {
         auto &buf = _catBatchBuffers[arg];
         buf.reserve(n);
         catCols.emplace_back(cat, &buf);
      } else if (auto real = dynamic_cast<RooAbsReal const *>(arg)) {
         auto &buf = _realBatchBuffers[arg];
         buf.reserve(n);
         realCols.emplace_back(real, &buf);
      }
   }
 
   for (std::size_t i = 0; i < n; ++i) {
      get(i);
      for (auto &col : realCols)
         col.second->push_back(col.first->getVal());
      for (auto &col : catCols)
         col.second->push_back(col.first->getCurrentIndex());
   }
}
 
////////////////////////////////////////////////////////////////////////////////
/// Get the batches of the real-valued columns in the range [first, first+len).
/// The columns of the component datasets are lazily concatenated into
/// internal buffers in composite row order.
RooAbsData::RealSpans RooCompositeDataStore::getBatches(std::size_t first, std::size_t len) const
{
   fillBatchBuffers();
 
   // Clamp against the actual number of entries, because the default
   // arguments of RooAbsData::getBatches() ask for the maximum length.
   first = std::min(first, static_cast<std::size_t>(numEntries()));
   len = std::min(len, static_cast<std::size_t>(numEntries()) - first);
 
   RooAbsData::RealSpans out;
   for (auto const &item : _realBatchBuffers) {
      out.emplace(item.first, std::span<const double>{item.second.data() + first, len});
   }
   return out;
}
 
////////////////////////////////////////////////////////////////////////////////
/// Get the batches of the category columns in the range [first, first+len),
/// including the index category. See getBatches().
RooAbsData::CategorySpans RooCompositeDataStore::getCategoryBatches(std::size_t first, std::size_t len) const
{
   fillBatchBuffers();
 
   // See getBatches() for the clamping rationale.
   first = std::min(first, static_cast<std::size_t>(numEntries()));
   len = std::min(len, static_cast<std::size_t>(numEntries()) - first);
 
   RooAbsData::CategorySpans out;
   for (auto const &item : _catBatchBuffers) {
      out.emplace(item.first, std::span<const RooAbsCategory::value_type>{item.second.data() + first, len});
   }
   return out;
}