[HS3] Fix HistFactory export and RooConstVar roundtripping

Author: Phmonski

roofit/hs3/src/Domains.cxx

@@ -162,7 +162,7 @@ void Domains::ProductDomain::populate(RooWorkspace &ws) const
 {
    for (auto const &item : _map) {
       const auto &name = item.first;
-      if (!ws.var(name)) {
+      if (!ws.arg(name)) {
          const auto &elem = item.second;
          const double vMin = elem.hasMin ? elem.min : -RooNumber::infinity();
          const double vMax = elem.hasMax ? elem.max : RooNumber::infinity();

roofit/hs3/src/JSONFactories_HistFactory.cxx

@@ -283,6 +283,36 @@ getOrCreateConstraint(RooJSONFactoryWSTool &tool, const JSONNode &mod, RooRealVa
                                   "'");
    }
 }
+double poissonTau(RooPoisson const &constraint, RooAbsArg const &gamma)
+{
+   auto const *mean = dynamic_cast<RooProduct const *>(&constraint.getMean());
+   if (!mean) {
+      RooJSONFactoryWSTool::error("Poisson gamma constraint mean is not a RooProduct: " +
+                                  std::string(constraint.GetName()));
+   }
+
+   for (RooAbsArg *arg : mean->servers()) {
+      if (arg == &gamma) {
+         continue;
+      }
+
+      if (auto const *tau = dynamic_cast<RooConstVar const *>(arg)) {
+         return tau->getVal();
+      }
+
+      // Imported workspaces can sometimes represent
+      // constants as constant RooRealVars.
+      if (auto const *real = dynamic_cast<RooAbsReal const *>(arg)) {
+         if (real->isConstant() || endsWith(std::string(real->GetName()), "_tau")) {
+            return real->getVal();
+         }
+      }
+   }
+
+   RooJSONFactoryWSTool::error("Could not find tau component in Poisson gamma constraint mean: " +
+                               std::string(constraint.GetName()));
+   return std::numeric_limits<double>::quiet_NaN();
+}
 
 bool importHistSample(RooJSONFactoryWSTool &tool, RooDataHist &dh, RooArgSet const &varlist,
                       RooAbsArg const *mcStatObject, const std::string &fprefix, const JSONNode &p,
@@ -334,6 +364,7 @@ bool importHistSample(RooJSONFactoryWSTool &tool, RooDataHist &dh, RooArgSet con
             // this is dealt with at a different place, ignore it for now
          } else if (modtype == "normfactor") {
             RooRealVar &constrParam = getOrCreate<RooRealVar>(ws, sysname, 1., -3, 5);
+            constrParam.setError(0.0);
             normElems.add(constrParam);
             if (mod.has_child("constraint_name") || mod.has_child("constraint_type")) {
                // for norm factors, constraints are optional
@@ -1060,7 +1091,7 @@ Channel readChannel(RooJSONFactoryWSTool *tool, const std::string &pdfname, cons
                if (constraint) {
                   sample.barlowBeestonLightConstraintType = constraint->IsA();
                   if (RooPoisson *constraint_p = dynamic_cast<RooPoisson *>(constraint)) {
-                     double erel = 1. / std::sqrt(constraint_p->getX().getVal());
+                     double erel = 1. / std::sqrt(poissonTau(*constraint_p, *g));
                      channel.rel_errors[idx] = erel;
                   } else if (RooGaussian *constraint_g = dynamic_cast<RooGaussian *>(constraint)) {
                      double erel = constraint_g->getSigma().getVal() / constraint_g->getMean().getVal();
@@ -1094,7 +1125,7 @@ Channel readChannel(RooJSONFactoryWSTool *tool, const std::string &pdfname, cons
                if (!constraint) {
                   sys.constraints.push_back(0.0);
                } else if (auto constraint_p = dynamic_cast<RooPoisson *>(constraint)) {
-                  sys.constraints.push_back(1. / std::sqrt(constraint_p->getX().getVal()));
+                  sys.constraints.push_back(1. / std::sqrt(poissonTau(*constraint_p, *g)));
                   if (!sys.constraint) {
                      sys.constraintType = RooPoisson::Class();
                   }

roofit/hs3/src/RooJSONFactoryWSTool.cxx

@@ -371,6 +371,41 @@ JSONNode const *findRooFitInternal(JSONNode const &node, Keys_t const &...keys)
    return node.find("misc", "ROOT_internal", keys...);
 }
 
+bool isMarkedConstVar(JSONNode const *attributesNode, std::string const &name)
+{
+   if (!attributesNode) {
+      return false;
+   }
+   if (auto *attrNode = attributesNode->find(name)) {
+      return attrNode->has_child("is_const_var") && (*attrNode)["is_const_var"].val_int() == 1;
+   }
+   return false;
+}
+
+void importMarkedConstVars(RooJSONFactoryWSTool &tool, JSONNode const &rootNode, JSONNode const *attributesNode)
+{
+   if (!attributesNode) {
+      return;
+   }
+
+   JSONNode const *varsNode = getVariablesNode(rootNode);
+   if (!varsNode) {
+      return;
+   }
+
+   RooWorkspace *workspace = tool.workspace();
+   for (const auto &varNode : varsNode->children()) {
+      std::string name = RooJSONFactoryWSTool::name(varNode);
+      if (!isMarkedConstVar(attributesNode, name) || workspace->arg(name)) {
+         continue;
+      }
+      if (!varNode.has_child("value")) {
+         RooJSONFactoryWSTool::error("cannot instantiate RooConstVar '" + name + "' without \"value\"!");
+      }
+      tool.wsEmplace<RooConstVar>(name, varNode["value"].val_double());
+   }
+}
+
 /**
  * @brief Check if a RooAbsArg is a literal constant variable.
  *
@@ -411,10 +446,9 @@ void exportAttributes(const RooAbsArg *arg, JSONNode &rootnode)
       node = &RooJSONFactoryWSTool::getRooFitInternal(rootnode, "attributes").set_map()[arg->GetName()].set_map();
    };
 
-   // RooConstVars are not a thing in HS3, and also for RooFit they are not
-   // that important: they are just constants. So we don't need to remember
-   // any information about them.
    if (dynamic_cast<RooConstVar const *>(arg)) {
+      initializeNode();
+      (*node)["is_const_var"] << 1;
       return;
    }
 
@@ -1688,22 +1722,20 @@ void RooJSONFactoryWSTool::importVariable(const JSONNode &p)
    std::string name(RooJSONFactoryWSTool::name(p));
    RooJSONFactoryWSTool::testValidName(name, true);
 
-   if (_workspace.var(name))
+   if (_workspace.var(name) || (isMarkedConstVar(_attributesNode, name) && _workspace.arg(name)))
       return;
    if (!p.is_map()) {
       std::stringstream ss;
       ss << "RooJSONFactoryWSTool() node '" << name << "' is not a map, skipping.";
       oocoutE(nullptr, InputArguments) << ss.str() << std::endl;
       return;
    }
-   if (_attributesNode) {
-      if (auto *attrNode = _attributesNode->find(name)) {
-         // We should not create RooRealVar objects for RooConstVars!
-         if (attrNode->has_child("is_const_var") && (*attrNode)["is_const_var"].val_int() == 1) {
-            wsEmplace<RooConstVar>(name, p["value"].val_double());
-            return;
-         }
+   if (isMarkedConstVar(_attributesNode, name)) {
+      if (!p.has_child("value")) {
+         RooJSONFactoryWSTool::error("cannot instantiate RooConstVar '" + name + "' without \"value\"!");
       }
+      wsEmplace<RooConstVar>(name, p["value"].val_double());
+      return;
    }
    configureVariable(*_domains, p, wsEmplace<RooRealVar>(name, 1.));
 }
@@ -2165,16 +2197,17 @@ void RooJSONFactoryWSTool::importAllNodes(const JSONNode &n)
       error(ss.str());
    }
 
+   _rootnodeInput = &n;
+
+   _attributesNode = findRooFitInternal(*_rootnodeInput, "attributes");
+
    _domains = std::make_unique<RooFit::JSONIO::Detail::Domains>();
    if (auto domains = n.find("domains")) {
       _domains->readJSON(*domains);
    }
+   importMarkedConstVars(*this, n, _attributesNode);
    _domains->populate(_workspace);
 
-   _rootnodeInput = &n;
-
-   _attributesNode = findRooFitInternal(*_rootnodeInput, "attributes");
-
    this->importDependants(n);
 
    if (auto paramPointsNode = n.find("parameter_points")) {

roofit/hs3/test/testRooFitHS3.cxx

@@ -259,6 +259,108 @@ TEST(RooFitHS3, RooGaussian)
    EXPECT_EQ(status, 0);
 }
 
+TEST(RooFitHS3, RooGaussianConstVarSigmaExport)
+{
+   RooRealVar x{"x", "x", 0.0, -10.0, 10.0};
+   RooRealVar mean{"mean", "mean", 0.0};
+   mean.setConstant(true);
+
+   RooConstVar sigmaConst{"sigma_const", "sigma_const", 1.0};
+   RooGaussian gaussConst{"gauss_const", "gauss_const", x, mean, sigmaConst};
+
+   RooGaussian gaussLiteral{"gauss_literal", "gauss_literal", x, mean, RooFit::RooConst(2.0)};
+
+   RooRealVar sigmaReal{"sigma_real", "sigma_real", 1.0, 0.1, 10.0};
+   sigmaReal.setConstant(true);
+   RooGaussian gaussReal{"gauss_real", "gauss_real", x, mean, sigmaReal};
+
+   RooWorkspace ws;
+   ws.import(gaussConst, RooFit::Silence());
+   ws.import(gaussLiteral, RooFit::RecycleConflictNodes(), RooFit::Silence());
+   ws.import(gaussReal, RooFit::RecycleConflictNodes(), RooFit::Silence());
+
+   const std::string json = RooJSONFactoryWSTool{ws}.exportJSONtoString();
+
+   EXPECT_NE(json.find("\"sigma\":\"sigma_const\""), std::string::npos);
+   EXPECT_NE(json.find("\"name\":\"sigma_const\""), std::string::npos);
+   EXPECT_NE(json.find("\"is_const_var\":1"), std::string::npos);
+   EXPECT_EQ(json.find("\"sigma\":1.0"), std::string::npos);
+   EXPECT_NE(json.find("\"sigma\":2.0"), std::string::npos);
+
+   EXPECT_NE(json.find("\"sigma\":\"sigma_real\""), std::string::npos);
+   EXPECT_NE(json.find("\"name\":\"sigma_real\""), std::string::npos);
+
+   RooWorkspace imported;
+   ASSERT_TRUE(RooJSONFactoryWSTool{imported}.importJSONfromString(json));
+   EXPECT_NE(dynamic_cast<RooConstVar *>(imported.obj("sigma_const")), nullptr);
+   EXPECT_EQ(imported.var("sigma_const"), nullptr);
+   EXPECT_NE(dynamic_cast<RooRealVar *>(imported.obj("sigma_real")), nullptr);
+
+   const std::string roundTripJson = RooJSONFactoryWSTool{imported}.exportJSONtoString();
+   EXPECT_NE(roundTripJson.find("\"sigma\":\"sigma_const\""), std::string::npos);
+   EXPECT_NE(roundTripJson.find("\"is_const_var\":1"), std::string::npos);
+
+   const std::string legacyJson = R"({
+      "metadata":{"hs3_version":"0.2"},
+      "parameter_points":[{"name":"default_values","parameters":[
+         {"name":"x","value":0.0},
+         {"name":"mean","value":0.0},
+         {"name":"sigma_const","value":1.0,"const":true}
+      ]}],
+      "distributions":[{"name":"gauss","type":"gaussian_dist","x":"x","mean":"mean","sigma":"sigma_const"}]
+   })";
+   RooWorkspace legacyImport;
+   ASSERT_TRUE(RooJSONFactoryWSTool{legacyImport}.importJSONfromString(legacyJson));
+   EXPECT_NE(dynamic_cast<RooRealVar *>(legacyImport.obj("sigma_const")), nullptr);
+   EXPECT_EQ(dynamic_cast<RooConstVar *>(legacyImport.obj("sigma_const")), nullptr);
+
+   const std::string markedConstWithDomainJson = R"({
+      "metadata":{"hs3_version":"0.2"},
+      "domains":[{"name":"default_domain","type":"product_domain","axes":[
+         {"name":"x","min":-10.0,"max":10.0},
+         {"name":"mean","min":-10.0,"max":10.0},
+         {"name":"sigma_const","min":0.0,"max":10.0}
+      ]}],
+      "parameter_points":[{"name":"default_values","parameters":[
+         {"name":"x","value":0.0},
+         {"name":"mean","value":0.0},
+         {"name":"sigma_const","value":1.0,"const":true}
+      ]}],
+      "distributions":[{"name":"gauss","type":"gaussian_dist","x":"x","mean":"mean","sigma":"sigma_const"}],
+      "misc":{"ROOT_internal":{"attributes":{"sigma_const":{"is_const_var":1}}}}
+   })";
+   RooWorkspace markedImport;
+   ASSERT_TRUE(RooJSONFactoryWSTool{markedImport}.importJSONfromString(markedConstWithDomainJson));
+   EXPECT_NE(dynamic_cast<RooConstVar *>(markedImport.obj("sigma_const")), nullptr);
+   EXPECT_EQ(markedImport.var("sigma_const"), nullptr);
+}
+
+TEST(RooFitHS3, RooConstVarCollectionProxyExport)
+{
+   RooRealVar x{"x", "x", 0.0, -10.0, 10.0};
+   RooRealVar mean1{"mean1", "mean1", -1.0};
+   RooRealVar mean2{"mean2", "mean2", 1.0};
+   RooRealVar sigma{"sigma", "sigma", 1.0, 0.1, 10.0};
+
+   RooGaussian g1{"g1", "g1", x, mean1, sigma};
+   RooGaussian g2{"g2", "g2", x, mean2, sigma};
+   RooConstVar frac{"frac_const", "frac_const", 0.25};
+   RooAddPdf model{"model", "model", RooArgList{g1, g2}, RooArgList{frac}};
+
+   RooWorkspace ws;
+   ws.import(model, RooFit::Silence());
+
+   const std::string json = RooJSONFactoryWSTool{ws}.exportJSONtoString();
+   EXPECT_NE(json.find("\"coefficients\":[\"frac_const\"]"), std::string::npos);
+   EXPECT_NE(json.find("\"name\":\"frac_const\""), std::string::npos);
+   EXPECT_NE(json.find("\"is_const_var\":1"), std::string::npos);
+
+   RooWorkspace imported;
+   ASSERT_TRUE(RooJSONFactoryWSTool{imported}.importJSONfromString(json));
+   EXPECT_NE(dynamic_cast<RooConstVar *>(imported.obj("frac_const")), nullptr);
+   EXPECT_EQ(imported.var("frac_const"), nullptr);
+}
+
 TEST(RooFitHS3, RooBernstein)
 {
    int status = validate({"RooBernstein::bernstein(x[0, 10], { a[1], 3, b[5, 0, 20] })"});