[HS3] Preserve explicit binnings for RooGenericPdf and RooFormulaVa

roofit/hs3/src/JSONFactories_HistFactory.cxx

@@ -1230,8 +1230,8 @@ bool normSysSupportsMultiplicativeMerge(int interpolationCode)
 // single modifier. The shared metadata (constraint, parameter and interpolation code) must be identical across the
 // duplicates; the type-specific `combine` callable performs the actual merge and any additional validation.
 template <class Modifiers, class CombineFn>
-void mergeDuplicateModifiers(const Channel &channel, const Sample &sample, Modifiers &modifiers,
-                             std::string_view type, CombineFn combine)
+void mergeDuplicateModifiers(const Channel &channel, const Sample &sample, Modifiers &modifiers, std::string_view type,
+                             CombineFn combine)
 {
    Modifiers mergedModifiers;
    mergedModifiers.reserve(modifiers.size());
@@ -1269,38 +1269,34 @@ void mergeDuplicateModifiers(const Channel &channel, const Sample &sample, Modif
 
 void mergeDuplicateNormSys(const Channel &channel, Sample &sample)
 {
-   mergeDuplicateModifiers(channel, sample, sample.normsys, "normsys",
-                           [&](NormSys &merged, const NormSys &modifier) {
-                              if (!normSysSupportsMultiplicativeMerge(merged.interpolationCode)) {
-                                 duplicateModifierError(
-                                    channel, sample, "normsys", merged.name,
-                                    "multiplicative combination is only valid for log-space interpolation codes");
-                              }
-                              merged.low *= modifier.low;
-                              merged.high *= modifier.high;
-                           });
+   mergeDuplicateModifiers(channel, sample, sample.normsys, "normsys", [&](NormSys &merged, const NormSys &modifier) {
+      if (!normSysSupportsMultiplicativeMerge(merged.interpolationCode)) {
+         duplicateModifierError(channel, sample, "normsys", merged.name,
+                                "multiplicative combination is only valid for log-space interpolation codes");
+      }
+      merged.low *= modifier.low;
+      merged.high *= modifier.high;
+   });
 }
 
 void mergeDuplicateHistoSys(const Channel &channel, Sample &sample)
 {
    const std::size_t nBins = sample.hist.size();
-   mergeDuplicateModifiers(channel, sample, sample.histosys, "histosys",
-                           [&](HistoSys &merged, const HistoSys &modifier) {
-                              if (merged.interpolationCode != 4) {
-                                 duplicateModifierError(
-                                    channel, sample, "histosys", merged.name,
-                                    "non-default interpolation cannot currently be represented by the HS3 exporter");
-                              }
-                              if (merged.low.size() != nBins || merged.high.size() != nBins ||
-                                  modifier.low.size() != nBins || modifier.high.size() != nBins) {
-                                 duplicateModifierError(channel, sample, "histosys", merged.name,
-                                                        "histogram binning differs");
-                              }
-                              for (std::size_t bin = 0; bin < nBins; ++bin) {
-                                 merged.low[bin] += modifier.low[bin] - sample.hist[bin];
-                                 merged.high[bin] += modifier.high[bin] - sample.hist[bin];
-                              }
-                           });
+   mergeDuplicateModifiers(
+      channel, sample, sample.histosys, "histosys", [&](HistoSys &merged, const HistoSys &modifier) {
+         if (merged.interpolationCode != 4) {
+            duplicateModifierError(channel, sample, "histosys", merged.name,
+                                   "non-default interpolation cannot currently be represented by the HS3 exporter");
+         }
+         if (merged.low.size() != nBins || merged.high.size() != nBins || modifier.low.size() != nBins ||
+             modifier.high.size() != nBins) {
+            duplicateModifierError(channel, sample, "histosys", merged.name, "histogram binning differs");
+         }
+         for (std::size_t bin = 0; bin < nBins; ++bin) {
+            merged.low[bin] += modifier.low[bin] - sample.hist[bin];
+            merged.high[bin] += modifier.high[bin] - sample.hist[bin];
+         }
+      });
 }
 
 void ensureUniqueModifiers(const Channel &channel, const Sample &sample)

roofit/hs3/src/JSONFactories_RooFitCore.cxx

@@ -41,6 +41,7 @@
 #include <RooPoisson.h>
 #include <RooPolynomial.h>
 #include <RooPolyVar.h>
+#include <RooAbsRealLValue.h>
 #include <RooRealSumFunc.h>
 #include <RooRealSumPdf.h>
 #include <RooRealVar.h>
@@ -52,6 +53,7 @@
 #include <RooWorkspace.h>
 #include <RooRealIntegral.h>
 #include <RooSpline.h>
+#include <RooUniformBinning.h>
 #include <TSpline.h>
 
 #include <TF1.h>
@@ -62,9 +64,14 @@
 #include "static_execute.h"
 
 #include <algorithm>
+#include <charconv>
 #include <cctype>
+#include <cmath>
+#include <limits>
+#include <memory>
 #include <set>
 #include <string_view>
+#include <vector>
 
 using RooFit::Detail::JSONNode;
 
@@ -163,6 +170,101 @@ void translateImportedExpression(TString &expr)
    replaceIdentifier(expr, "EULER", "TMath::E()");
 }
 
+int readPositiveInteger(const JSONNode &node, const std::string &context)
+{
+   const std::string value = node.val();
+   int out = 0;
+   const auto result = std::from_chars(value.data(), value.data() + value.size(), out);
+   if (result.ec != std::errc{} || result.ptr != value.data() + value.size() || out <= 0) {
+      RooJSONFactoryWSTool::error("\"nbins\" in " + context + " must be a positive integer");
+   }
+   return out;
+}
+
+std::unique_ptr<RooAbsBinning>
+readFormulaAxisBinning(const JSONNode &axis, const std::string &axisName, const std::string &formulaName)
+{
+   const std::string context = "axis '" + axisName + "' of generic formula '" + formulaName + "'";
+   const bool hasEdges = axis.has_child("edges");
+   const bool hasMin = axis.has_child("min");
+   const bool hasMax = axis.has_child("max");
+   const bool hasNBins = axis.has_child("nbins");
+
+   if (hasEdges && (hasMin || hasMax || hasNBins)) {
+      RooJSONFactoryWSTool::error(context + " must use either \"edges\" or \"min\"/\"max\"/\"nbins\"");
+   }
+
+   if (hasEdges) {
+      const JSONNode &edgesNode = axis["edges"];
+      if (!edgesNode.is_seq()) {
+         RooJSONFactoryWSTool::error("\"edges\" in " + context + " must be a sequence");
+      }
+
+      std::vector<double> edges;
+      edges.reserve(edgesNode.num_children());
+      for (const JSONNode &edgeNode : edgesNode.children()) {
+         const double edge = edgeNode.val_double();
+         if (!std::isfinite(edge)) {
+            RooJSONFactoryWSTool::error("\"edges\" in " + context + " must contain only finite values");
+         }
+         if (!edges.empty() && edge <= edges.back()) {
+            RooJSONFactoryWSTool::error("\"edges\" in " + context + " must be strictly increasing");
+         }
+         edges.push_back(edge);
+      }
+      if (edges.size() < 2) {
+         RooJSONFactoryWSTool::error("\"edges\" in " + context + " must contain at least two values");
+      }
+      return std::make_unique<RooBinning>(static_cast<int>(edges.size() - 1), edges.data());
+   }
+
+   if (!hasMin || !hasMax || !hasNBins) {
+      RooJSONFactoryWSTool::error(context + " must define \"min\", \"max\", and \"nbins\"");
+   }
+
+   const double min = axis["min"].val_double();
+   const double max = axis["max"].val_double();
+   if (!std::isfinite(min) || !std::isfinite(max) || max <= min) {
+      RooJSONFactoryWSTool::error("\"min\" and \"max\" in " + context + " must be finite and increasing");
+   }
+   return std::make_unique<RooUniformBinning>(min, max, readPositiveInteger(axis["nbins"], context));
+}
+
+template <class RooArg_t>
+void importFormulaBinnings(RooArg_t &arg, const JSONNode &node)
+{
+   if (!node.has_child("axes")) {
+      return;
+   }
+
+   const JSONNode &axes = node["axes"];
+   if (!axes.is_seq()) {
+      RooJSONFactoryWSTool::error("\"axes\" in generic formula '" + std::string(arg.GetName()) +
+                                  "' must be a sequence");
+   }
+
+   std::set<std::string> axisNames;
+   for (const JSONNode &axis : axes.children()) {
+      if (!axis.is_map() || !axis.has_child("name")) {
+         RooJSONFactoryWSTool::error("each axis in generic formula '" + std::string(arg.GetName()) +
+                                     "' must be a map with a \"name\"");
+      }
+      const std::string axisName = axis["name"].val();
+      if (!axisNames.insert(axisName).second) {
+         RooJSONFactoryWSTool::error("duplicate axis '" + axisName + "' in generic formula '" + arg.GetName() + "'");
+      }
+
+      auto *observable = dynamic_cast<RooAbsRealLValue *>(arg.getParameter(axisName.c_str()));
+      if (!observable) {
+         RooJSONFactoryWSTool::error(
+            "axis '" + axisName + "' is not a real-valued formula variable of generic formula '" + arg.GetName() + "'");
+      }
+
+      std::unique_ptr<RooAbsBinning> binning = readFormulaAxisBinning(axis, axisName, arg.GetName());
+      arg.setBinning(*observable, *binning, /*checkFlatness=*/false);
+   }
+}
+
 template <class RooArg_t>
 bool importFormulaArg(RooJSONFactoryWSTool *tool, const JSONNode &p)
 {
@@ -176,7 +278,9 @@ bool importFormulaArg(RooJSONFactoryWSTool *tool, const JSONNode &p)
    for (const auto &d : extractArguments(formula.Data())) {
       dependents.add(*tool->request<RooAbsReal>(d, name));
    }
-   tool->wsImport(RooArg_t{name.c_str(), formula, dependents});
+   RooArg_t arg{name.c_str(), formula, dependents};
+   importFormulaBinnings(arg, p);
+   tool->wsImport(arg);
    return true;
 }
 
@@ -774,6 +878,25 @@ bool exportFormulaArg(RooJSONFactoryWSTool *, const RooAbsArg *func, JSONNode &e
       expression.ReplaceAll(("@" + std::to_string(idx)).c_str(), par->GetName());
    }
    elem["expression"] << expression.Data();
+
+   for (const RooAbsArg *dependent : pdf->dependents()) {
+      auto const *observable = dynamic_cast<const RooAbsRealLValue *>(dependent);
+      if (!observable) {
+         continue;
+      }
+      const RooAbsBinning *binning = pdf->getBinning(*observable);
+      if (!binning) {
+         continue;
+      }
+
+      auto &axes = elem["axes"];
+      if (!axes.is_seq()) {
+         axes.set_seq();
+      }
+      auto &axis = axes.append_child().set_map();
+      axis["name"] << observable->GetName();
+      writeAxisBinning(axis, *binning);
+   }
    return true;
 }
 

roofit/hs3/src/JSONIOUtils.cxx

@@ -1,5 +1,7 @@
 #include "JSONIOUtils.h"
 
+#include <RooAbsBinning.h>
+
 #include <string>
 
 using RooFit::Detail::JSONNode;
@@ -82,3 +84,19 @@ std::unique_ptr<RooFit::Detail::JSONTree> varJSONString(const JSONNode &treeRoot
 
    return jsonDict;
 }
+
+void writeAxisBinning(JSONNode &node, const RooAbsBinning &binning)
+{
+   if (binning.isUniform()) {
+      node["min"] << binning.lowBound();
+      node["max"] << binning.highBound();
+      node["nbins"] << binning.numBins();
+      return;
+   }
+
+   auto &edges = node["edges"].set_seq();
+   edges.append_child() << binning.binLow(0);
+   for (int i = 0; i < binning.numBins(); ++i) {
+      edges.append_child() << binning.binHigh(i);
+   }
+}

roofit/hs3/src/JSONIOUtils.h

@@ -4,9 +4,12 @@
 #include <string_view>
 #include <RooFit/Detail/JSONInterface.h>
 
+class RooAbsBinning;
+
 bool startsWith(std::string_view str, std::string_view prefix);
 bool endsWith(std::string_view str, std::string_view suffix);
 std::string removeSuffix(std::string_view str, std::string_view suffix);
 std::unique_ptr<RooFit::Detail::JSONTree> varJSONString(const RooFit::Detail::JSONNode &treeRoot);
+void writeAxisBinning(RooFit::Detail::JSONNode &node, const RooAbsBinning &binning);
 
 #endif

roofit/hs3/src/RooJSONFactoryWSTool.cxx

@@ -1232,18 +1232,7 @@ void RooJSONFactoryWSTool::exportAxis(JSONNode &obsNode, RooRealVar const &var)
    RooJSONFactoryWSTool::testValidName(name, false);
    obsNode["name"] << name;
 
-   auto const &binning = var.getBinning();
-   if (binning.isUniform()) {
-      obsNode["min"] << var.getMin();
-      obsNode["max"] << var.getMax();
-      obsNode["nbins"] << var.getBins();
-   } else {
-      auto &edges = obsNode["edges"].set_seq();
-      edges.append_child() << binning.binLow(0);
-      for (int i = 0; i < binning.numBins(); ++i) {
-         edges.append_child() << binning.binHigh(i);
-      }
-   }
+   writeAxisBinning(obsNode, var.getBinning());
 }
 
 /**