[RF][HS3] Combined data sets and combined/conditional distributions (HS3 #119; depends on #24001)

roofit/histfactory/test/testHistFactory.cxx

@@ -32,6 +32,7 @@
 
 #include "../../roofitcore/test/gtest_wrapper.h"
 
+#include <algorithm>
 #include <regex>
 #include <set>
 
@@ -700,15 +701,10 @@ INSTANTIATE_TEST_SUITE_P(HistFactory, HFFixtureFit,
                                           testing::Values(ROOFIT_EVAL_BACKENDS_WITH_CODEGEN)),
                          getNameFromInfo);
 
-// Regression test for the HS3 importer's handling of "shapefactor" modifiers.
-// HistFactory's ShapeFactor is exported with type "shapesys" (the only
-// modifier type the writer ever emits for a ParamHistFunc), but valid HS3 JSON
-// may also use the dedicated "shapefactor" type for an unconstrained
-// ParamHistFunc - and the importer must accept it. This test starts from a
-// HistFactory model built with MakeModelAndMeasurementFast, rewrites the
-// modifier type to "shapefactor" in the JSON, and checks that the importer
-// recognises it. Without the fix in JSONFactories_HistFactory.cxx the import
-// throws "modifier ... of unknown type 'shapefactor'".
+// HS3 round trip of a free HistFactory ShapeFactor. It is written as a
+// "shapefactor" modifier without "constraints", and must survive the
+// JSON -> workspace -> JSON round trip with free gammas. The deprecated
+// "shapesys" alias is read as well.
 TEST(HistFactory, HS3ImportShapeFactorModifier)
 {
    using namespace RooStats::HistFactory;
@@ -745,12 +741,9 @@ TEST(HistFactory, HS3ImportShapeFactorModifier)
    chan.AddSample(sig);
 
    // ShapeFactor on the background: an unconstrained, bin-by-bin scaling.
-   // Make the gammas constant so that the workspace is well-defined for
-   // re-export (free shapefactor gammas have no constraints attached).
    Sample bkg("background", "background", inputFile);
    ShapeFactor sf;
    sf.SetName("bkgShape");
-   sf.SetConstant(true);
    bkg.AddShapeFactor(sf);
    chan.AddSample(bkg);
 
@@ -762,26 +755,30 @@ TEST(HistFactory, HS3ImportShapeFactorModifier)
 
    const std::string js = RooJSONFactoryWSTool{*ws}.exportJSONtoString();
 
-   // Rewrite the modifier type for "bkgShape" from "shapesys" to "shapefactor".
-   // The HistFactory exporter always writes "shapesys", but the importer should
-   // accept the more accurate "shapefactor" type as well.
-   const std::regex pattern{"\"name\":\"bkgShape\",\"parameters\":\\[([^\\]]*)\\],\"type\":\"shapesys\""};
-   const std::string jsShapeFactor =
-      std::regex_replace(js, pattern, "\"name\":\"bkgShape\",\"parameters\":[$1],\"type\":\"shapefactor\"");
-   ASSERT_NE(js, jsShapeFactor) << "Failed to substitute shapesys -> shapefactor in JSON";
+   // The free ShapeFactor is written as "shapefactor", without constraints.
+   const std::regex pattern{"\"name\":\"bkgShape\",\"parameters\":\\[[^\\]]*\\],\"type\":\"shapefactor\""};
+   EXPECT_TRUE(std::regex_search(js, pattern)) << js;
+   EXPECT_EQ(js.find("shapesys"), std::string::npos) << js;
 
    RooWorkspace wsFromJson{"new"};
-   ASSERT_NO_THROW(RooJSONFactoryWSTool{wsFromJson}.importJSONfromString(jsShapeFactor))
-      << "Importer rejected the 'shapefactor' modifier type";
+   ASSERT_NO_THROW(RooJSONFactoryWSTool{wsFromJson}.importJSONfromString(js));
 
-   // The imported workspace should expose the same ParamHistFunc gammas.
-   EXPECT_NE(wsFromJson.var("gamma_bkgShape_bin_0"), nullptr);
-   EXPECT_NE(wsFromJson.var("gamma_bkgShape_bin_1"), nullptr);
+   // The imported workspace should expose the same free ParamHistFunc gammas.
+   ASSERT_NE(wsFromJson.var("gamma_bkgShape_bin_0"), nullptr);
+   ASSERT_NE(wsFromJson.var("gamma_bkgShape_bin_1"), nullptr);
+   EXPECT_FALSE(wsFromJson.var("gamma_bkgShape_bin_0")->isConstant());
 
-   // Re-exporting should give back the original JSON, since the writer emits
-   // type "shapesys" in both cases.
+   // Re-exporting should give back the original JSON.
    const std::string js2 = RooJSONFactoryWSTool{wsFromJson}.exportJSONtoString();
    EXPECT_EQ(js, js2) << "JSON -> WS -> JSON roundtrip changed the JSON";
+
+   // The deprecated "shapesys" alias gives the same model.
+   const std::string jsShapeSys =
+      std::regex_replace(js, std::regex{"\"type\":\"shapefactor\""}, "\"type\":\"shapesys\"");
+   ASSERT_NE(js, jsShapeSys);
+   RooWorkspace wsFromAlias{"alias"};
+   ASSERT_NO_THROW(RooJSONFactoryWSTool{wsFromAlias}.importJSONfromString(jsShapeSys));
+   EXPECT_EQ(RooJSONFactoryWSTool{wsFromAlias}.exportJSONtoString(), js);
 }
 
 // Issue #20697: Sample::AddShapeFactor() now allows to set the initial value
@@ -919,3 +916,77 @@ TEST(HistFactory, ShapeFactorValueAndRange)
       }
    }
 }
+
+/// The reference workspace of ROOT 6.16 is the standard case of a RooSimultaneous with combined data sets: it is
+/// written as a "combined_dist" and "combined" data sets (HS3 #119), read back into the same objects, and the
+/// likelihood is the same. The document is stable from the third export on: the 6.16 workspace has "binWidth_*"
+/// constants that the imported model doesn't use, and the HistFactory importer sets the default parameter errors
+/// (step widths), which the 6.16 workspace lacks. Both are unrelated to the combined objects.
+TEST(HistFactory, HS3RoundTripOfReferenceWorkspace)
+{
+   RooHelpers::LocalChangeMsgLevel changeMsgLvl(RooFit::WARNING);
+
+   std::unique_ptr<TFile> file(TFile::Open("./ref_6.16_example_UsingC_combined_meas_model.root"));
+   ASSERT_TRUE(file && file->IsOpen());
+   RooWorkspace *ws = nullptr;
+   file->GetObject("combined", ws);
+   ASSERT_NE(ws, nullptr);
+   auto *mc = dynamic_cast<RooStats::ModelConfig *>(ws->obj("ModelConfig"));
+   ASSERT_NE(mc, nullptr);
+
+   const std::string json = RooJSONFactoryWSTool{*ws}.exportJSONtoString();
+   EXPECT_NE(json.find(R"("type":"combined_dist")"), std::string::npos);
+   EXPECT_NE(json.find(R"("type":"combined")"), std::string::npos);
+   EXPECT_EQ(json.find("combined_datasets"), std::string::npos);
+   EXPECT_EQ(json.find("combined_distributions"), std::string::npos);
+
+   RooWorkspace wsFromJson("new");
+   RooJSONFactoryWSTool{wsFromJson}.importJSONfromString(json);
+
+   auto *newmc = dynamic_cast<RooStats::ModelConfig *>(wsFromJson.obj("ModelConfig"));
+   ASSERT_NE(newmc, nullptr);
+   auto *newpdf = dynamic_cast<RooSimultaneous *>(newmc->GetPdf());
+   ASSERT_NE(newpdf, nullptr);
+   EXPECT_STREQ(newpdf->indexCat().GetName(), "channelCat");
+   EXPECT_NE(newmc->GetObservables()->find("channelCat"), nullptr);
+   EXPECT_EQ(wsFromJson.allData().size(), ws->allData().size());
+   for (const char *dataName : {"obsData", "asimovData"}) {
+      RooAbsData *data = ws->data(dataName);
+      RooAbsData *newdata = wsFromJson.data(dataName);
+      ASSERT_NE(data, nullptr) << dataName;
+      ASSERT_NE(newdata, nullptr) << dataName;
+      EXPECT_EQ(newdata->numEntries(), data->numEntries()) << dataName;
+      EXPECT_DOUBLE_EQ(newdata->sumEntries(), data->sumEntries()) << dataName;
+      EXPECT_NE(newdata->get()->find("channelCat"), nullptr) << dataName;
+
+      // The likelihood is compared up to a constant: whether RooFit uses its binned likelihood for the channels
+      // (which differs from the event-wise one by a constant) is not part of the document.
+      using namespace RooFit;
+      std::unique_ptr<RooAbsReal> nll{mc->GetPdf()->createNLL(*data, GlobalObservables(*mc->GetGlobalObservables()))};
+      std::unique_ptr<RooAbsReal> newnll{
+         newpdf->createNLL(*newdata, GlobalObservables(*newmc->GetGlobalObservables()))};
+      const double nominal = nll->getVal();
+      const double newNominal = newnll->getVal();
+      const double alpha = ws->var("alpha_syst2")->getVal();
+      const double mu = ws->var("SigXsecOverSM")->getVal();
+      for (RooWorkspace *w : {ws, &wsFromJson}) {
+         w->var("alpha_syst2")->setVal(alpha + 0.5);
+         w->var("SigXsecOverSM")->setVal(mu + 0.5);
+      }
+      const double difference = nll->getVal() - nominal;
+      EXPECT_NE(difference, 0.0) << dataName;
+      EXPECT_NEAR(newnll->getVal() - newNominal, difference, 1e-8 * std::max(1.0, std::abs(difference))) << dataName;
+      for (RooWorkspace *w : {ws, &wsFromJson}) {
+         w->var("alpha_syst2")->setVal(alpha);
+         w->var("SigXsecOverSM")->setVal(mu);
+      }
+   }
+
+   const std::string json2 = RooJSONFactoryWSTool{wsFromJson}.exportJSONtoString();
+   RooWorkspace wsFromJson2("new2");
+   RooJSONFactoryWSTool{wsFromJson2}.importJSONfromString(json2);
+   const std::string json3 = RooJSONFactoryWSTool{wsFromJson2}.exportJSONtoString();
+   RooWorkspace wsFromJson3("new3");
+   RooJSONFactoryWSTool{wsFromJson3}.importJSONfromString(json3);
+   EXPECT_EQ(RooJSONFactoryWSTool{wsFromJson3}.exportJSONtoString(), json3);
+}

roofit/hs3/inc/RooFitHS3/RooJSONFactoryWSTool.h

@@ -21,6 +21,7 @@
 #include <RooWorkspace.h>
 
 #include <map>
+#include <memory>
 #include <stdexcept>
 #include <set>
 #include <unordered_map>
@@ -36,6 +37,9 @@ namespace RooStats {
 class ModelConfig;
 }
 class RooRealVar;
+class RooAbsPdf;
+class RooCategory;
+class RooSimultaneous;
 
 class RooJSONFactoryWSTool {
 public:
@@ -45,6 +49,13 @@ class RooJSONFactoryWSTool {
       bool allowExportInvalidNames = true;
       bool allowSanitizeNames = true;
       bool importNoDomainParametersAsRooConstVars = 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;
    };
 
    static Config &config();
@@ -54,7 +65,7 @@ class RooJSONFactoryWSTool {
 
    struct CombinedData {
       std::string name;
-      std::map<std::string, std::string> components;
+      std::vector<std::string> columns;
    };
 
    RooJSONFactoryWSTool(RooWorkspace &ws);
@@ -184,14 +195,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)
    {
@@ -228,6 +238,7 @@ class RooJSONFactoryWSTool {
       }
    }
 
+   void exportSimultaneous(RooSimultaneous const &simPdf, std::set<std::string> &exportedObjectNames);
    void exportData(RooAbsData const &data);
    RooJSONFactoryWSTool::CombinedData exportCombinedData(RooAbsData const &data);
 
@@ -246,8 +257,7 @@ class RooJSONFactoryWSTool {
                           const std::vector<RooAbsData *> &single);
 
    void exportSingleModelConfig(RooFit::Detail::JSONNode &rootnode, RooStats::ModelConfig const &mc,
-                                std::string const &analysisName,
-                                std::map<std::string, std::string> const *dataComponents);
+                                std::string const &analysisName, std::string const *dataName);
 
    // member variables
    const RooFit::Detail::JSONNode *_rootnodeInput = nullptr;
@@ -261,6 +271,15 @@ class RooJSONFactoryWSTool {
    std::unique_ptr<RooFit::JSONIO::Detail::Domains> _domains;
    std::vector<RooAbsArg const *> _serversToExport;
    std::vector<RooAbsArg const *> _serversToDelete;
+   // The columns of the exported combined distributions: they are not parameters, and don't appear in parameter points.
+   std::set<std::string> _columnCategories;
+   // 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();
 
    // Name-keyed indices over the top-level "functions" and "distributions"
    // sequences of the input JSON. Built once at the start of importAllNodes()