CalibrationAlgorithm.h

/**************************************************************************
 * BASF2 (Belle Analysis Framework 2)                                     *
 * Copyright(C) 2010 - Belle II Collaboration                             *
 *                                                                        *
 * Author: The Belle II Collaboration                                     *
 * Contributors:  Tadeas Bilka, David Dossett                             *
 *                                                                        *
 * This software is provided "as is" without any warranty.                *
 **************************************************************************/
 
#pragma once
#include <Python.h>
#include <memory>
#include <string>
#include <vector>
#include <map>
#include <any>
#include <utility>
#include <list>
#include <nlohmann/json.hpp>
#include <TClonesArray.h>
#include <TDirectory.h>
#include <TFile.h>
#include <TTree.h>
#include <framework/database/Database.h>
#include <framework/database/IntervalOfValidity.h>
#include <framework/logging/Logger.h>
#include <calibration/Utilities.h>
#include <calibration/dataobjects/RunRange.h>
 
namespace Belle2 {
  class CalibrationAlgorithm {
  public:
    enum EResult {
      c_OK,           
      c_Iterate,      
      c_NotEnoughData,
      c_Failure,      
      c_Undefined     
    };
 
    class ExecutionData {
    public:
      ExecutionData() {};
      ~ExecutionData() {};
      void reset()
      {
        B2DEBUG(100, "Resetting ExecutionData of algorithm");
        m_requestedRuns.clear();
        m_iteration = -1;
        m_result = c_Undefined;
        m_payloads.clear();
        m_iov = IntervalOfValidity();
        clearCalibrationData();
      }
      void clearCalibrationData()
      {
        m_mapCalibData.clear();
      }
      const std::vector<Calibration::ExpRun>& getRequestedRuns() const {return m_requestedRuns;}
      void setRequestedRuns(const std::vector<Calibration::ExpRun>& requestedRuns) {m_requestedRuns = requestedRuns;}
      int getIteration() const {return m_iteration;}
      void setIteration(int iteration)
      {
        B2DEBUG(29, "Setting Iteration of Algorithm to " << iteration);
        m_iteration = iteration;
      }
      EResult getResult() const {return m_result;}
      void setResult(EResult result) {m_result = result;}
      void setRequestedIov(const IntervalOfValidity& iov = IntervalOfValidity(0, 0, -1, -1)) {m_iov = iov;}
      const IntervalOfValidity& getRequestedIov() const {return m_iov;}
      std::list<Database::DBImportQuery>& getPayloads() {return m_payloads;}
      std::list<Database::DBImportQuery> getPayloadValues() {return m_payloads;}
      std::shared_ptr<TNamed> getCalibObj(const std::string& name, const RunRange& runRange) const
      {
        auto it = m_mapCalibData.find(std::make_pair(name, runRange));
        if (it == m_mapCalibData.end()) {
          return nullptr;
        }
        return it->second;
      }
      void setCalibObj(const std::string& name, const RunRange& runRange, const std::shared_ptr<TNamed>& objectPtr)
      {
        m_mapCalibData[std::make_pair(name, runRange)] = objectPtr;
      }
 
    private:
      std::vector<Calibration::ExpRun> m_requestedRuns{};
      int m_iteration{ -1};
      EResult m_result{c_Undefined};
      IntervalOfValidity m_iov;
      std::list<Database::DBImportQuery> m_payloads{};
      std::map<std::pair<std::string, RunRange>, std::shared_ptr<TNamed>> m_mapCalibData;
    };
 
    explicit CalibrationAlgorithm(const std::string& collectorModuleName) : m_prefix(collectorModuleName) {}
 
    virtual ~CalibrationAlgorithm() {}
 
    std::string getPrefix() const {return m_prefix;}
 
    bool checkPyExpRun(PyObject* pyObj);
 
    Calibration::ExpRun convertPyExpRun(PyObject* pyObj);
 
    std::string getCollectorName() const {return getPrefix();}
 
    void setPrefix(const std::string& prefix) {m_prefix = prefix;}
 
    void setInputFileNames(PyObject* inputFileNames);
 
    PyObject* getInputFileNames();
 
    std::vector<Calibration::ExpRun> getRunListFromAllData() const;
 
    RunRange getRunRangeFromAllData() const;
 
    IntervalOfValidity getIovFromAllData() const;
 
    void fillRunToInputFilesMap();
 
    std::string getGranularity() const {return m_granularityOfData;};
 
    EResult execute(std::vector<Calibration::ExpRun> runs = {}, int iteration = 0, IntervalOfValidity iov = IntervalOfValidity());
 
    EResult execute(PyObject* runs, int iteration = 0, IntervalOfValidity iov = IntervalOfValidity());
 
    std::list<Database::DBImportQuery>& getPayloads() {return m_data.getPayloads();}
 
    std::list<Database::DBImportQuery> getPayloadValues() {return m_data.getPayloadValues();}
 
    bool commit();
 
    bool commit(std::list<Database::DBImportQuery> payloads);
 
    const std::string& getDescription() const {return m_description;}
 
    //  of JSON object creation. Failure probably means that your JSON string was badly formatted.
    bool loadInputJson(const std::string& jsonString);
 
    const std::string dumpOutputJson() const {return m_jsonExecutionOutput.dump();}
 
    //  areas of the CAF. Basically you search for features in the data that you want to find and make sure that the CAF knows
    //  there is a boundary where payloads should probably start/end. The output boundaries should be the starting ExpRun
    //  of the new boundary.
    const std::vector<Calibration::ExpRun> findPayloadBoundaries(std::vector<Calibration::ExpRun> runs, int iteration = 0);
 
  protected:
    // Developers implement this function ------------
 
    virtual EResult calibrate() = 0;
 
    //  Implementing this is optional because most people will never call findPayloadBoundaries in their CAF job.
    //  It returns false by default so that the boundaries vector is empty if you forgot to implement this.
    //
    //  We omit the names of arguments here so that we don't generate lots of compiler warnings in algorithms that don't
    //  implement this function.
    virtual bool isBoundaryRequired(const Calibration::ExpRun& /*currentRun*/)
    {
      B2ERROR("You didn't implement a isBoundaryRequired() member function in your CalibrationAlgorithm but you are calling it!");
      return false;
    }
 
    //  We omit the names of arguments here so that we don't generate lots of compiler warnings in algorithms that don't
    //  implement this function.
    virtual void boundaryFindingSetup(std::vector<Calibration::ExpRun> /*runs*/, int /*iteration = 0*/) {};
 
    //  This runs right after 'findPayloadBoundaries' and is supposed to  correct any changes you made in 'boundaryFindingSetup'
    //  or 'isBoundaryRequired'.
    virtual void boundaryFindingTearDown() {};
 
    // When using the boundaries functionality from isBoundaryRequired, this is used to store the boundaries. It is cleared when
    // calling findPayloadBoundaries, before boundaryFindingSetup is called.
    std::vector<Calibration::ExpRun> m_boundaries;
 
    // Helpers ---------------- Data retrieval -------
 
    const std::vector<Calibration::ExpRun>& getRunList() const {return m_data.getRequestedRuns();}
 
    int getIteration() const { return m_data.getIteration(); }
 
    void setInputFileNames(std::vector<std::string> inputFileNames);
 
    std::vector<std::string> getVecInputFileNames() const {return m_inputFileNames;}
 
    template<class T>
    std::shared_ptr<T> getObjectPtr(const std::string& name, const std::vector<Calibration::ExpRun>& requestedRuns);
 
    template<class T>
    std::shared_ptr<T> getObjectPtr(std::string name)
    {
      if (m_runsToInputFiles.size() == 0)
        fillRunToInputFilesMap();
      return getObjectPtr<T>(name, m_data.getRequestedRuns());
    }
 
    // Helpers ---------------- Database storage -----
 
    std::string getGranularityFromData() const;
 
    void saveCalibration(TClonesArray* data, const std::string& name);
 
    void saveCalibration(TClonesArray* data, const std::string& name, const IntervalOfValidity& iov);
 
    void saveCalibration(TObject* data);
 
    void saveCalibration(TObject* data, const IntervalOfValidity& iov);
 
    void saveCalibration(TObject* data, const std::string& name);
 
    void saveCalibration(TObject* data, const std::string& name, const IntervalOfValidity& iov);
 
    void updateDBObjPtrs(const unsigned int event, const int run, const int experiment);
 
    // -----------------------------------------------
 
    void setDescription(const std::string& description) {m_description = description;}
 
    void clearCalibrationData() {m_data.clearCalibrationData();}
 
    Calibration::ExpRun getAllGranularityExpRun() const {return m_allExpRun;}
 
    void resetInputJson() {m_jsonExecutionInput.clear();}
 
    void resetOutputJson() {m_jsonExecutionOutput.clear();}
 
    template<class T>
    void setOutputJsonValue(const std::string& key, const T& value) {m_jsonExecutionOutput[key] = value;}
 
    //  No attempt to catch exceptions is made here.
    template<class T>
    const T getOutputJsonValue(const std::string& key) const
    {
      return m_jsonExecutionOutput.at(key);
    }
 
    //  No attempt to catch them is made here.
    template<class T>
    const T getInputJsonValue(const std::string& key) const
    {
      return m_jsonExecutionInput.at(key);
    }
 
    //  up on weird errors where someone snuck an array into the member variable.
    const nlohmann::json& getInputJsonObject() const {return m_jsonExecutionInput;}
 
    bool inputJsonKeyExists(const std::string& key) const {return m_jsonExecutionInput.count(key);}
 
  private:
 
    static const Calibration::ExpRun m_allExpRun;
 
    std::string getExpRunString(Calibration::ExpRun& expRun) const;
 
    std::string getFullObjectPath(std::string name, Calibration::ExpRun expRun) const;
 
    std::vector<std::string> m_inputFileNames;
 
    std::map<Calibration::ExpRun, std::vector<std::string>> m_runsToInputFiles;
 
    std::string m_granularityOfData;
 
    ExecutionData m_data;
 
    std::string m_description{""};
 
    std::string m_prefix{""};
 
    //  We initialise to "{}" rather than allowing a JSON array/value as the top level type. This forces the user
    //  to use "key":value for storing data in this object. You should test for empty(), is_null() will always
    //  return false due to the empty top level object.
    //  Functionally similar to simple member variables for configuration of the algorithm.
    //  However these input values are easier to use from Python code without needing to know the details of the algorithm.
    //  These values are intended to be used for a single execution, not reused
    nlohmann::json m_jsonExecutionInput = nlohmann::json::object();
 
    //  As for input we initialise to an empty "{}" JSON object. Testing for empty() returns true, but is_null() does not.
    //  Nothing is done with these by default, however a calling process may decide to capture these values and use them
    //  as input to a following execution.
    nlohmann::json m_jsonExecutionOutput = nlohmann::json::object();
 
  };  // End of CalibrationAlgorithm definition
 
 
  /**************************************
   *                                    *
   * Implementation of larger templates *
   *                                    *
   **************************************/
  template<class T>
  std::shared_ptr<T> CalibrationAlgorithm::getObjectPtr(const std::string& name,
                                                        const std::vector<Calibration::ExpRun>& requestedRuns)
  {
    // Check if this object already exists
    RunRange runRangeRequested(requestedRuns);
    std::shared_ptr<T> objOutputPtr = std::dynamic_pointer_cast<T>(m_data.getCalibObj(name, runRangeRequested));
    if (objOutputPtr)
      return objOutputPtr;
 
    std::shared_ptr<T> mergedObjPtr(nullptr);
    bool mergedEmpty = true;
    TDirectory* dir = gDirectory;
 
    // Technically we could grab all the objects from all files, add to list and then merge at the end.
    // But I prefer the (maybe) more memory efficient way of merging with all objects
    // in a file before moving on to the next one, just in case TDirectory stuff screws us.
    TList list;
    list.SetOwner(false);
 
    // Construct the TDirectory names where we expect our objects to be
    std::string runRangeObjName(getPrefix() + "/" + Calibration::RUN_RANGE_OBJ_NAME);
 
    if (strcmp(getGranularity().c_str(), "run") == 0) {
      // Loop over our runs requested for the right files
      for (auto expRunRequested : requestedRuns) {
        // Find the relevant files for this ExpRun
        auto searchFiles = m_runsToInputFiles.find(expRunRequested);
        if (searchFiles == m_runsToInputFiles.end()) {
          B2WARNING("No input file found with data collected from run "
                    "(" << expRunRequested.first << "," << expRunRequested.second << ")");
          continue;
        } else {
          auto files = searchFiles->second;
          for (auto fileName : files) {
            RunRange* runRangeData;
            //Open TFile to get the objects
            std::unique_ptr<TFile> f;
            f.reset(TFile::Open(fileName.c_str(), "READ"));
            runRangeData = dynamic_cast<RunRange*>(f->Get(runRangeObjName.c_str()));
            // Check that nothing went wrong in the mapping and that this file definitely contains this run's data
            auto runSet = runRangeData->getExpRunSet();
            if (runSet.find(expRunRequested) == runSet.end()) {
              B2WARNING("Something went wrong with the mapping of ExpRun -> Input Files. "
                        "(" << expRunRequested.first << "," << expRunRequested.second << ") not in " << fileName);
            }
            // Get the path/directory of the Exp,Run TDirectory that holds the object(s)
            std::string objDirName = getFullObjectPath(name, expRunRequested);
            TDirectory* objDir = f->GetDirectory(objDirName.c_str());
            if (!objDir) {
              B2ERROR("Directory for requested object " << name << " not found: " << objDirName);
              return nullptr;
            }
            // Find all the objects inside, there may be more than one
            for (auto key : * (objDir->GetListOfKeys())) {
              std::string keyName = key->GetName();
              B2DEBUG(100, "Adding found object " << keyName << " in the directory " << objDir->GetPath());
              T* objOther = (T*)objDir->Get(keyName.c_str());
              if (objOther) {
                if (mergedEmpty) {
                  mergedObjPtr = std::shared_ptr<T>(dynamic_cast<T*>(objOther->Clone(name.c_str())));
                  mergedObjPtr->SetDirectory(0);
                  mergedEmpty = false;
                } else {
                  list.Add(objOther);
                }
              }
            }
            if (!mergedEmpty)
              mergedObjPtr->Merge(&list);
            list.Clear();
          }
        }
      }
    } else {
      for (auto fileName : m_inputFileNames) {
        //Open TFile to get the objects
        std::unique_ptr<TFile> f;
        f.reset(TFile::Open(fileName.c_str(), "READ"));
        Calibration::ExpRun allGranExpRun = getAllGranularityExpRun();
        std::string objDirName = getFullObjectPath(name, allGranExpRun);
        std::string objPath = objDirName + "/" + name + "_1";
        T* objOther = (T*)f->Get(objPath.c_str()); // Only one index for granularity == all
        B2DEBUG(100, "Adding " << objPath);
        if (objOther) {
          if (mergedEmpty) {
            mergedObjPtr = std::shared_ptr<T>(dynamic_cast<T*>(objOther->Clone(name.c_str())));
            mergedObjPtr->SetDirectory(0);
            mergedEmpty = false;
          } else {
            list.Add(objOther);
          }
        }
        if (!mergedEmpty)
          mergedObjPtr->Merge(&list);
        list.Clear();
      }
    }
    dir->cd();
    objOutputPtr = mergedObjPtr;
    if (!objOutputPtr) {
      B2ERROR("No data found for object " << name);
      return nullptr;
    }
    objOutputPtr->SetDirectory(0);
    // make a TNamed version to input to the map of previous calib objects
    std::shared_ptr<TNamed> storedObjPtr = std::static_pointer_cast<TNamed>(objOutputPtr);
    m_data.setCalibObj(name, runRangeRequested, storedObjPtr);
    B2DEBUG(100, "Passing back merged data " << name);
    return objOutputPtr;
  }
 
  template<> std::shared_ptr<TTree>
  CalibrationAlgorithm::getObjectPtr(
    const std::string& name,
    const std::vector<Calibration::ExpRun>& requestedRuns);
 
} // namespace Belle2