root/RooFitHS3/RooJSONFactoryWSTool.h

0001 /*
0002  * Project: RooFit
0003  * Authors:
0004  *   Carsten D. Burgard, DESY/ATLAS, Dec 2021
0005  *
0006  * Copyright (c) 2022, CERN
0007  *
0008  * Redistribution and use in source and binary forms,
0009  * with or without modification, are permitted according to the terms
0010  * listed in LICENSE (http://roofit.sourceforge.net/license.txt)
0011  */
0012
0013 #ifndef RooFitHS3_RooJSONFactoryWSTool_h
0014 #define RooFitHS3_RooJSONFactoryWSTool_h
0015
0016 #include <RooFit/Detail/JSONInterface.h>
0017
0018 #include <RooArgList.h>
0019 #include <RooArgSet.h>
0020 #include <RooGlobalFunc.h>
0021 #include <RooWorkspace.h>
0022
0023 #include <map>
0024 #include <stdexcept>
0025 #include <set>
0026 #include <unordered_map>
0027
0028 namespace RooFit {
0029 namespace JSONIO {
0030 namespace Detail {
0031 class Domains;
0032 }
0033 } // namespace JSONIO
0034 } // namespace RooFit
0035 namespace RooStats {
0036 class ModelConfig;
0037 }
0038 class RooRealVar;
0039
0040 class RooJSONFactoryWSTool {
0041 public:
0042    static constexpr bool useListsInsteadOfDicts = true;
0043
0044    struct Config {
0045       bool allowExportInvalidNames = true;
0046       bool allowSanitizeNames = true;
0047       bool importNoDomainParametersAsRooConstVars = true;
0048    };
0049
0050    static Config &config();
0051
0052    static RooWorkspace sanitizeWS(const RooWorkspace &ws);
0053    static RooWorkspace cleanWS(const RooWorkspace &ws, bool onlyModelConfig = false);
0054
0055    struct CombinedData {
0056       std::string name;
0057       std::map<std::string, std::string> components;
0058    };
0059
0060    RooJSONFactoryWSTool(RooWorkspace &ws);
0061
0062    ~RooJSONFactoryWSTool();
0063
0064    static std::string name(const RooFit::Detail::JSONNode &n);
0065    static bool isValidName(const std::string &str);
0066    static bool testValidName(const std::string &str, bool forcError);
0067    static std::string sanitizeName(const std::string str);
0068    static void rebuildModelConfigInWorkspace(RooStats::ModelConfig *mc, RooWorkspace &ws);
0069
0070    static RooFit::Detail::JSONNode &appendNamedChild(RooFit::Detail::JSONNode &node, std::string const &name);
0071    static RooFit::Detail::JSONNode const *findNamedChild(RooFit::Detail::JSONNode const &node, std::string const &name);
0072
0073    static void fillSeq(RooFit::Detail::JSONNode &node, RooAbsCollection const &coll, size_t nMax = -1);
0074
0075    template <class T>
0076    T *request(const std::string &objname, const std::string &requestAuthor)
0077    {
0078       if (T *out = requestImpl<T>(objname)) {
0079          return out;
0080       }
0081       throw DependencyMissingError(requestAuthor, objname, T::Class()->GetName());
0082    }
0083
0084    template <class T>
0085    T *requestArg(const RooFit::Detail::JSONNode &node, const std::string &key)
0086    {
0087       std::string requestAuthor(RooJSONFactoryWSTool::name(node));
0088       if (!node.has_child(key)) {
0089          RooJSONFactoryWSTool::error("no \"" + key + "\" given in \"" + requestAuthor + "\"");
0090       }
0091       return request<T>(node[key].val(), requestAuthor);
0092    }
0093
0094    template <class T, class Coll_t>
0095    Coll_t requestCollection(const RooFit::Detail::JSONNode &node, const std::string &seqName)
0096    {
0097       std::string requestAuthor(RooJSONFactoryWSTool::name(node));
0098       if (!node.has_child(seqName)) {
0099          RooJSONFactoryWSTool::error("no \"" + seqName + "\" given in \"" + requestAuthor + "\"");
0100       }
0101       if (!node[seqName].is_seq()) {
0102          RooJSONFactoryWSTool::error("\"" + seqName + "\" in \"" + requestAuthor + "\" is not a sequence");
0103       }
0104
0105       Coll_t out;
0106       for (const auto &elem : node[seqName].children()) {
0107          out.add(*request<T>(elem.val(), requestAuthor));
0108       }
0109       return out;
0110    }
0111
0112    template <class T>
0113    RooArgSet requestArgSet(const RooFit::Detail::JSONNode &node, const std::string &seqName)
0114    {
0115       return requestCollection<T, RooArgSet>(node, seqName);
0116    }
0117
0118    template <class T>
0119    RooArgList requestArgList(const RooFit::Detail::JSONNode &node, const std::string &seqName)
0120    {
0121       return requestCollection<T, RooArgList>(node, seqName);
0122    }
0123
0124    RooWorkspace *workspace() { return &_workspace; }
0125
0126    template <class Obj_t>
0127    Obj_t &wsImport(Obj_t const &obj)
0128    {
0129       _workspace.import(obj, RooFit::RecycleConflictNodes(true), RooFit::Silence(true));
0130       return *static_cast<Obj_t *>(_workspace.obj(obj.GetName()));
0131    }
0132
0133    template <class Obj_t, typename... Args_t>
0134    Obj_t &wsEmplace(RooStringView name, Args_t &&...args)
0135    {
0136       return wsImport(Obj_t(name.c_str(), name.c_str(), std::forward<Args_t>(args)...));
0137    }
0138
0139    [[noreturn]] static void error(const char *s);
0140    [[noreturn]] inline static void error(const std::string &s) { error(s.c_str()); }
0141    static std::ostream &warning(const std::string &s);
0142
0143    static RooArgSet readAxes(const RooFit::Detail::JSONNode &node);
0144    static std::unique_ptr<RooDataHist>
0145    readBinnedData(const RooFit::Detail::JSONNode &n, const std::string &namecomp, RooArgSet const &vars);
0146
0147    bool importJSON(std::string const &filename);
0148    bool importJSON(std::istream &os);
0149    bool exportJSON(std::string const &fileName);
0150    bool exportJSON(std::ostream &os);
0151
0152    std::string exportJSONtoString();
0153    bool importJSONfromString(const std::string &s);
0154    void importJSONElement(const std::string &name, const std::string &jsonString);
0155    void importVariableElement(const RooFit::Detail::JSONNode &n);
0156
0157    void importFunction(const RooFit::Detail::JSONNode &p, bool importAllDependants);
0158    void importFunction(const std::string &jsonString, bool importAllDependants);
0159
0160    static std::unique_ptr<RooFit::Detail::JSONTree> createNewJSONTree();
0161
0162    static RooFit::Detail::JSONNode &makeVariablesNode(RooFit::Detail::JSONNode &rootNode);
0163
0164    // error handling helpers
0165    class DependencyMissingError : public std::exception {
0166       std::string _parent, _child, _class, _message;
0167
0168    public:
0169       DependencyMissingError(const std::string &p, const std::string &c, const std::string &classname)
0170          : _parent(p), _child(c), _class(classname)
0171       {
0172          _message = "object '" + _parent + "' is missing dependency '" + _child + "' of type '" + _class + "'";
0173       };
0174       const std::string &parent() const { return _parent; }
0175       const std::string &child() const { return _child; }
0176       const std::string &classname() const { return _class; }
0177       const char *what() const noexcept override { return _message.c_str(); }
0178    };
0179
0180    template <typename... Keys_t>
0181    static RooFit::Detail::JSONNode &getRooFitInternal(RooFit::Detail::JSONNode &node, Keys_t const &...keys)
0182    {
0183       return node.get("misc", "ROOT_internal", keys...);
0184    }
0185
0186    static void exportAxis(RooFit::Detail::JSONNode &obsNode, RooRealVar const &var);
0187
0188    static void
0189    exportHisto(RooArgSet const &vars, std::size_t n, double const *contents, RooFit::Detail::JSONNode &output);
0190
0191    static void exportArray(std::size_t n, double const *contents, RooFit::Detail::JSONNode &output);
0192
0193    void exportCategory(RooAbsCategory const &cat, RooFit::Detail::JSONNode &node);
0194
0195    void queueExport(RooAbsArg const &arg) { _serversToExport.push_back(&arg); }
0196    void queueExportTemporary(RooAbsArg *arg)
0197    {
0198       _serversToExport.push_back(arg);
0199       _serversToDelete.push_back(arg);
0200    }
0201
0202    std::string exportTransformed(const RooAbsReal *original, const std::string &suffix, const std::string &formula);
0203
0204    void setAttribute(const std::string &obj, const std::string &attrib);
0205    bool hasAttribute(const std::string &obj, const std::string &attrib);
0206    std::string getStringAttribute(const std::string &obj, const std::string &attrib);
0207    void setStringAttribute(const std::string &obj, const std::string &attrib, const std::string &value);
0208
0209 private:
0210    template <class T>
0211    T *requestImpl(const std::string &objname);
0212    void exportObject(RooAbsArg const &func, std::set<std::string> &exportedObjectNames);
0213
0214    // To export multiple objects sorted alphabetically
0215    template <class T>
0216    void exportObjects(T const &args, std::set<std::string> &exportedObjectNames)
0217    {
0218       RooArgSet argSet;
0219       for (RooAbsArg const *arg : args) {
0220          argSet.add(*arg);
0221       }
0222       argSet.sort();
0223       for (RooAbsArg *arg : argSet) {
0224          exportObject(*arg, exportedObjectNames);
0225       }
0226    }
0227
0228    void exportData(RooAbsData const &data);
0229    RooJSONFactoryWSTool::CombinedData exportCombinedData(RooAbsData const &data);
0230
0231    void importAllNodes(const RooFit::Detail::JSONNode &n);
0232
0233    void importVariable(const RooFit::Detail::JSONNode &p);
0234    void importDependants(const RooFit::Detail::JSONNode &n);
0235
0236    void exportVariable(const RooAbsArg *v, RooFit::Detail::JSONNode &n, bool storeConstant, bool storeBins);
0237    void exportVariables(const RooArgSet &allElems, RooFit::Detail::JSONNode &n, bool storeConstant, bool storeBins);
0238
0239    void exportAllObjects(RooFit::Detail::JSONNode &n);
0240
0241    void exportModelConfig(RooFit::Detail::JSONNode &rootnode, RooStats::ModelConfig const &mc,
0242                           const std::vector<RooJSONFactoryWSTool::CombinedData> &combined,
0243                           const std::vector<RooAbsData *> &single);
0244
0245    void exportSingleModelConfig(RooFit::Detail::JSONNode &rootnode, RooStats::ModelConfig const &mc,
0246                                 std::string const &analysisName,
0247                                 std::map<std::string, std::string> const *dataComponents);
0248
0249    // member variables
0250    const RooFit::Detail::JSONNode *_rootnodeInput = nullptr;
0251    const RooFit::Detail::JSONNode *_attributesNode = nullptr;
0252    RooFit::Detail::JSONNode *_rootnodeOutput = nullptr;
0253    RooFit::Detail::JSONNode *_varsNode = nullptr;
0254    RooWorkspace &_workspace;
0255
0256    // objects to represent intermediate information
0257    std::unique_ptr<RooFit::JSONIO::Detail::Domains> _domains;
0258    std::vector<RooAbsArg const *> _serversToExport;
0259    std::vector<RooAbsArg const *> _serversToDelete;
0260
0261    // Name-keyed indices over the top-level "functions" and "distributions"
0262    // sequences of the input JSON. Built once at the start of importAllNodes()
0263    // so that requestImpl() lookups become O(1) instead of an O(N) scan over
0264    // every sibling node.
0265    std::unordered_map<std::string, RooFit::Detail::JSONNode const *> _functionsByName;
0266    std::unordered_map<std::string, RooFit::Detail::JSONNode const *> _distributionsByName;
0267 };
0268 #endif