[RF][HS3] category_conditional_dist and RooMultiPdf export (HS3 #147; depends on #24005)

roofit/hs3/inc/RooFitHS3/RooJSONFactoryWSTool.h

@@ -46,6 +46,10 @@ namespace RooStats {
 class ModelConfig;
 }
 class RooRealVar;
+class RooAbsCategory;
+class RooCategory;
+class RooMultiPdf;
+class RooSimultaneous;
 
 class RooJSONFactoryWSTool {
 public:
@@ -59,6 +63,19 @@ class RooJSONFactoryWSTool {
       /// before ROOT 6.26 by their structure, and don't write them, since `histfactory_dist` implies them. Turn this
       /// off if a workspace has such factors that are not meant as bin widths: they are then written as normfactors.
       bool identifyLegacyBinWidthFactors = true;
+      /// HS3 requires the components of a combined distribution to have distinct variates, e.g. "mass_SR" and
+      /// "mass_CR" instead of one "mass". If enabled, the observables of every combined data set are renamed on
+      /// export to "obs_<key>" (the labels of the key joined by "_"), or "obs_<key>_<observable>" if a component
+      /// has several observables; the components of the matching RooSimultaneous are written with the renamed
+      /// observables (functions depending on them get the suffix "_<key>"). This allows exporting workspaces whose
+      /// components share an observable. Off by default: without it, such workspaces can't be exported.
+      bool renameCombinedObservables = false;
+      /// HS3 requires the terms of a likelihood and the components of a combined distribution to have distinct
+      /// variates, e.g. "mass_SR" and "mass_CR" for the observables of two channels instead of one "mass". By
+      /// default, the export fails otherwise. If true, the document is written anyway, with a warning for every
+      /// shared variate. It is not valid HS3 and can't be read back before the variates are renamed, e.g. by hand in
+      /// the JSON.
+      bool allowExportSharedVariates = false;
    };
 
    static Config &config();
@@ -68,7 +85,7 @@ class RooJSONFactoryWSTool {
 
    struct CombinedData {
       std::string name;
-      std::map<std::string, std::string> components;
+      std::vector<std::string> columns;
    };
 
    RooJSONFactoryWSTool(RooWorkspace &ws);
@@ -199,14 +216,13 @@ class RooJSONFactoryWSTool {
    }
 
    static void exportAxis(RooFit::Detail::JSONNode &obsNode, RooRealVar const &var);
+   static void exportCategoricalAxis(RooFit::Detail::JSONNode &obsNode, RooCategory const &cat);
 
    static void
    exportHisto(RooArgSet const &vars, std::size_t n, double const *contents, RooFit::Detail::JSONNode &output);
 
    static void exportArray(std::size_t n, double const *contents, RooFit::Detail::JSONNode &output);
 
-   void exportCategory(RooAbsCategory const &cat, RooFit::Detail::JSONNode &node);
-
    void queueExport(RooAbsArg const &arg) { _serversToExport.push_back(&arg); }
    void queueExportTemporary(RooAbsArg *arg)
    {
@@ -243,6 +259,9 @@ class RooJSONFactoryWSTool {
       }
    }
 
+   void exportSimultaneous(RooSimultaneous const &simPdf, std::set<std::string> &exportedObjectNames);
+   void recordCategory(RooAbsCategory const &cat);
+   void exportMultiPdf(RooMultiPdf const &multiPdf, std::set<std::string> &exportedObjectNames);
    void exportData(RooAbsData const &data);
    RooJSONFactoryWSTool::CombinedData exportCombinedData(RooAbsData const &data);
 
@@ -263,7 +282,7 @@ class RooJSONFactoryWSTool {
                           const std::vector<RooAbsData *> &single);
 
    void exportSingleModelConfig(RooStats::ModelConfig const &mc, std::string const &analysisName,
-                                std::map<std::string, std::string> const *dataComponents);
+                                std::string const *dataName);
 
    friend void RooFit::HS3::setExportedObject(RooJSONFactoryWSTool &tool, RooFit::HS3::HistFactoryDist dist,
                                               std::vector<RooFit::HS3::DistributionRef> channelConstraints);
@@ -277,6 +296,13 @@ class RooJSONFactoryWSTool {
    std::unique_ptr<RooFit::HS3::ToolExportState> _export;
    std::vector<RooAbsArg const *> _serversToExport;
    std::vector<RooAbsArg const *> _serversToDelete;
+   // Renaming of the observables of combined data sets, see Config::renameCombinedObservables: for each data set and
+   // each selection label, the new name of each observable; for each RooSimultaneous and selection label, the
+   // component with renamed observables to export instead of the original one. The objects are owned here.
+   std::map<std::string, std::map<std::string, std::map<std::string, std::string>>> _combinedObservableRenames;
+   std::map<std::string, std::map<std::string, RooAbsPdf const *>> _renamedComponents;
+   std::vector<std::unique_ptr<RooAbsArg>> _renamedObjects;
+   void prepareCombinedObservableRenaming();
 
    // The objects that are currently being imported, innermost last, to detect circular dependencies.
    std::vector<std::string> _importStack;

roofit/hs3/src/HS3Document.h

@@ -188,9 +188,21 @@ struct BinnedData {
    std::optional<GaussianUncertainty> uncertainty;
 };
 
+/// The sub-data set of a combined data set for one key, which has one label per column.
+struct CombinedDataComponent {
+   std::vector<std::string> key;
+   DataRef data;
+};
+
+/// Data sets that are selected by a key (HS3 #119). The columns are categorical axes of the domains.
+struct CombinedData {
+   std::vector<std::string> columns;
+   std::vector<CombinedDataComponent> components;
+};
+
 struct Data {
    std::string name;
-   std::variant<PointData, UnbinnedData, BinnedData> body;
+   std::variant<PointData, UnbinnedData, BinnedData, CombinedData> body;
 };
 
 struct LikelihoodTerm {
@@ -456,6 +468,8 @@ struct ModelConfigInfo {
    std::optional<std::string> mcName;
    std::optional<std::string> pdfName;
    std::optional<std::string> combinedDataName;
+   /// The observables of the ModelConfig, written if no data set of the likelihood gives them.
+   std::optional<std::vector<std::string>> observables;
 };
 
 struct StepWidth {
@@ -473,6 +487,9 @@ struct RooFitInternal {
    /// The constraint terms of HistFactory distributions that no modifier names, by distribution. They are
    /// also listed in the "aux_distributions" of the likelihoods, which is what other implementations use.
    std::map<std::string, std::vector<DistributionRef>> channelConstraints;
+   /// Axes that the components of a combined data set share, but that are no variate of any distribution, e.g. an
+   /// event number, by combined data set. They have no statistical meaning, and may be shared.
+   std::map<std::string, std::vector<std::string>> bookkeepingColumns;
 };
 
 struct Diagnostic {

roofit/hs3/src/HS3Parser.cxx

@@ -574,6 +574,60 @@ void describe(Reader &ar, BinnedData &data)
    }
 }
 
+// A combined data set. Only the structure is checked here: the columns and keys are checked against the domains by the
+// validation, and the sub-data sets when the data set is built.
+void readCombinedData(Reader &ar, std::string const &name, CombinedData &data)
+{
+   std::string const context = "the combined data set '" + name + "'";
+   Cursor const &c = ar.cursor();
+   if (ar.has("axes")) {
+      c.error(context + " must not have \"axes\": its columns are declared by the domains");
+      return;
+   }
+   auto nonEmptyList = [&](const char *key) -> JSONNode const * {
+      if (!ar.has(key) || !c.node()[key].is_seq() || c.node()[key].num_children() == 0) {
+         return nullptr;
+      }
+      return &c.node()[key];
+   };
+   JSONNode const *columns = nonEmptyList("columns");
+   for (std::size_t i = 0; columns && i < columns->num_children(); ++i) {
+      JSONNode const &column = columns->child(i);
+      if (column.is_container() || column.val().empty()) {
+         columns = nullptr;
+         break;
+      }
+      data.columns.push_back(column.val());
+   }
+   if (!columns) {
+      c.error("\"columns\" of " + context + " must be a non-empty list of names of categorical axes");
+      return;
+   }
+   JSONNode const *components = nonEmptyList("components");
+   if (!components) {
+      c.error("\"components\" of " + context + " must be a non-empty list");
+      return;
+   }
+   for (JSONNode const &component : components->children()) {
+      if (!component.is_map() || !component.has_child("key") || !component.has_child("dataset")) {
+         c.error("every component of " + context + " must have a \"key\" and a \"dataset\"");
+         return;
+      }
+      JSONNode const &key = component["key"];
+      CombinedDataComponent &out = data.components.emplace_back();
+      if (key.is_seq() && key.num_children() == data.columns.size()) {
+         for (JSONNode const &label : key.children()) {
+            out.key.push_back(label.is_container() ? std::string{} : label.val());
+         }
+      } else {
+         c.error("every component of " + context + " must have a \"key\" with " + std::to_string(data.columns.size()) +
+                 " labels, one per column");
+         return;
+      }
+      out.data.name = component["dataset"].val();
+   }
+}
+
 void readValue(Cursor const &c, Data &out)
 {
    if (!c.node().is_map()) {
@@ -596,6 +650,10 @@ void readValue(Cursor const &c, Data &out)
       PointData body;
       describe(ar, body);
       out.body = std::move(body);
+   } else if (type == "combined") {
+      CombinedData body;
+      readCombinedData(ar, out.name, body);
+      out.body = std::move(body);
    } else if (ar.has("type")) {
       c.error("unknown type of dataset '" + type + "'");
    }
@@ -1099,6 +1157,7 @@ void describe(Reader &ar, ModelConfigInfo &info)
    ar("mcName", info.mcName);
    ar("pdfName", info.pdfName);
    ar("combined_data_name", info.combinedDataName);
+   ar("observables", info.observables);
 }
 
 namespace {
@@ -1168,6 +1227,7 @@ void readMisc(Cursor const &root, RooFitInternal &out)
    readMap(ar, "combined_datasets", out.combinedDatasets);
    readMap(ar, "ModelConfigs", out.modelConfigs);
    readMap(ar, "histfactory_constraints", out.channelConstraints);
+   readMap(ar, "bookkeeping_columns", out.bookkeepingColumns);
 }
 
 void readMetadata(Cursor const &root, Metadata &out, ParseOptions const &options)

roofit/hs3/src/HS3Validate.cxx

@@ -45,6 +45,7 @@ class Validator {
       checkLikelihoods();
       checkAnalyses();
       checkData();
+      checkCombinedData();
       for (std::size_t i = 0; i < _doc.distributions.size(); ++i) {
          auto const &dist = _doc.distributions[i];
          if (auto const *histFactory = std::get_if<HistFactoryDist>(&dist.body)) {
@@ -130,6 +131,7 @@ class Validator {
    void checkLikelihoods();
    void checkAnalyses();
    void checkData();
+   void checkCombinedData();
 
    HS3Document const &_doc;
    Diagnostics &_diagnostics;
@@ -329,6 +331,74 @@ void Validator::checkCategoricalAxes()
    }
 }
 
+// The columns of combined data sets are categorical axes of the domains, or categories that exist already, and the keys
+// consist of their categories. Every sub-data set is a data set of the document that belongs to one combined data set.
+void Validator::checkCombinedData()
+{
+   std::map<std::string, std::vector<std::string> const *> categorical;
+   for (Domain const &domain : _doc.domains) {
+      for (auto const &axis : domain.axes) {
+         if (auto const *cat = std::get_if<CategoricalAxis>(&axis)) {
+            categorical.emplace(cat->name, &cat->categories);
+         }
+      }
+   }
+   std::set<std::string> combinedNames;
+   for (auto const &data : _doc.data) {
+      if (std::holds_alternative<CombinedData>(data.body)) {
+         combinedNames.insert(data.name);
+      }
+   }
+   std::set<std::string> usedComponents;
+   for (std::size_t i = 0; i < _doc.data.size(); ++i) {
+      auto const *combined = std::get_if<CombinedData>(&_doc.data[i].body);
+      if (!combined) {
+         continue;
+      }
+      std::string const path = at("data", i, _doc.data[i].name);
+      std::string const context = "the combined data set '" + _doc.data[i].name + "'";
+      std::vector<std::vector<std::string> const *> columnCategories;
+      std::set<std::string> columns;
+      for (std::string const &column : combined->columns) {
+         auto found = categorical.find(column);
+         if (found == categorical.end() && !isExternal(column)) {
+            error(path,
+                  "the column '" + column + "' of " + context + " is not declared by a categorical axis of a domain");
+            return;
+         }
+         if (!columns.insert(column).second) {
+            error(path, "the column '" + column + "' of " + context + " is listed twice");
+            return;
+         }
+         columnCategories.push_back(found != categorical.end() ? found->second : nullptr);
+      }
+      std::set<std::vector<std::string>> keys;
+      for (CombinedDataComponent const &component : combined->components) {
+         for (std::size_t k = 0; k < component.key.size(); ++k) {
+            auto const *categories = columnCategories[k];
+            if (categories &&
+                std::find(categories->begin(), categories->end(), component.key[k]) == categories->end()) {
+               error(path, "the label '" + component.key[k] + "' in position " + std::to_string(k) + " of a key of " +
+                              context + " is not one of the categories of '" + combined->columns[k] + "'");
+            }
+         }
+         if (!keys.insert(component.key).second) {
+            std::string key = "[";
+            for (std::size_t k = 0; k < component.key.size(); ++k) {
+               key += (k == 0 ? "\"" : ", \"") + component.key[k] + "\"";
+            }
+            error(path, "the key " + key + "] is used by two components of " + context);
+         }
+         std::string const &name = component.data.name;
+         if (!isDefinedIn(name, "data") || combinedNames.count(name) || !usedComponents.insert(name).second) {
+            error(path, "the component '" + name + "' of " + context +
+                           " is not a data set of the document, or it is already a component of another combined "
+                           "data set");
+         }
+      }
+   }
+}
+
 void Validator::checkLikelihoods()
 {
    for (std::size_t i = 0; i < _doc.likelihoods.size(); ++i) {

roofit/hs3/src/HS3Writer.cxx

@@ -272,6 +272,16 @@ void writeData(JSONNode &dataArray, Data const &data)
                        node["weights"].fill_seq(*d.weights);
                     }
                  },
+                 [&](CombinedData const &d) {
+                    node["type"] << "combined";
+                    writeStrings(node["columns"], d.columns);
+                    JSONNode &components = node["components"].set_seq();
+                    for (CombinedDataComponent const &component : d.components) {
+                       JSONNode &out = components.append_child().set_map();
+                       writeStrings(out["key"], component.key);
+                       out["dataset"] << component.data.name;
+                    }
+                 },
                  [&](PointData const &d) {
                     node["type"] << "point";
                     node["value"] << d.value;
@@ -368,6 +378,9 @@ void writeRooFitInternal(JSONNode &root, RooFitInternal const &internal)
          if (item.second.combinedDataName) {
             node["combined_data_name"] << *item.second.combinedDataName;
          }
+         if (item.second.observables) {
+            node["observables"].fill_seq(*item.second.observables);
+         }
       }
    }
    if (!internal.channelConstraints.empty()) {
@@ -376,6 +389,12 @@ void writeRooFitInternal(JSONNode &root, RooFitInternal const &internal)
          writeRefs(constraints[item.first], item.second);
       }
    }
+   if (!internal.bookkeepingColumns.empty()) {
+      JSONNode &columns = rootInternal()["bookkeeping_columns"].set_map();
+      for (auto const &item : internal.bookkeepingColumns) {
+         writeStrings(columns[item.first], item.second);
+      }
+   }
    if (!internal.stepWidths.empty()) {
       JSONNode &stepWidths = root.get("misc", "minimization")["parameter_stepwidths"].set_seq();
       for (auto const &stepWidth : internal.stepWidths) {