Dataset.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
 
#pragma once
#ifndef INCLUDE_GUARD_BELLE2_MVA_DATASET_HEADER
#define INCLUDE_GUARD_BELLE2_MVA_DATASET_HEADER
 
#include <mva/interface/Options.h>
 
#include <analysis/VariableManager/Manager.h>
 
#include <TFile.h>
#include <TChain.h>
 
#include <string>
 
namespace Belle2 {
  namespace MVA {
 
    class Dataset {
 
    public:
      explicit Dataset(const GeneralOptions& general_options);
 
      virtual ~Dataset() = default;
 
      Dataset(const Dataset&) = delete;
 
      Dataset& operator=(const Dataset&) = delete;
 
      virtual unsigned int getNumberOfFeatures() const = 0;
 
      virtual unsigned int getNumberOfSpectators() const = 0;
 
      virtual unsigned int getNumberOfEvents() const = 0;
 
      virtual void loadEvent(unsigned int iEvent) = 0;
 
      virtual float getSignalFraction();
 
      virtual unsigned int getFeatureIndex(const std::string& feature);
 
      virtual unsigned int getSpectatorIndex(const std::string& spectator);
 
      virtual std::vector<float> getFeature(unsigned int iFeature);
 
      virtual std::vector<float> getSpectator(unsigned int iSpectator);
 
      virtual std::vector<float> getWeights();
 
      virtual std::vector<float> getTargets();
 
      virtual std::vector<bool> getSignals();
 
      GeneralOptions m_general_options; 
      std::vector<float> m_input; 
      std::vector<float> m_spectators; 
      float m_weight; 
      float m_target; 
      bool m_isSignal; 
    };
 
 
    class SingleDataset : public Dataset {
 
    public:
      SingleDataset(const GeneralOptions& general_options, const std::vector<float>& input, float target = 1.0,
                    const std::vector<float>& spectators = std::vector<float>());
 
      virtual unsigned int getNumberOfFeatures() const override { return m_input.size(); }
 
      virtual unsigned int getNumberOfSpectators() const override { return m_spectators.size(); }
 
      virtual unsigned int getNumberOfEvents() const override { return 1; }
 
      virtual void loadEvent(unsigned int) override { };
 
      virtual std::vector<float> getFeature(unsigned int iFeature) override { return std::vector<float> {m_input[iFeature]}; }
 
      virtual std::vector<float> getSpectator(unsigned int iSpectator) override { return std::vector<float> {m_spectators[iSpectator]}; }
 
    };
 
    class MultiDataset : public Dataset {
 
    public:
      MultiDataset(const GeneralOptions& general_options, const std::vector<std::vector<float>>& input,
                   const std::vector<std::vector<float>>& spectators,
                   const std::vector<float>& targets = {}, const std::vector<float>& weights = {});
 
      virtual unsigned int getNumberOfFeatures() const override { return m_input.size(); }
 
      virtual unsigned int getNumberOfSpectators() const override { return m_spectators.size(); }
 
      virtual unsigned int getNumberOfEvents() const override { return m_matrix.size(); }
 
      virtual void loadEvent(unsigned int iEvent) override;
 
 
    private:
      std::vector<std::vector<float>> m_matrix; 
      std::vector<std::vector<float>> m_spectator_matrix; 
      std::vector<float> m_targets; 
      std::vector<float> m_weights; 
    };
 
    class SubDataset : public Dataset {
 
    public:
      SubDataset(const GeneralOptions& general_options, const std::vector<bool>& events, Dataset& dataset);
 
      virtual unsigned int getNumberOfFeatures() const override { return m_feature_indices.size(); }
 
      virtual unsigned int getNumberOfSpectators() const override { return m_spectator_indices.size(); }
 
      virtual unsigned int getNumberOfEvents() const override { return m_use_event_indices ? m_event_indices.size() : m_dataset.getNumberOfEvents(); }
 
      virtual void loadEvent(unsigned int iEvent) override;
 
      virtual std::vector<float> getFeature(unsigned int iFeature) override;
 
      virtual std::vector<float> getSpectator(unsigned int iSpectator) override;
 
    private:
      bool m_use_event_indices = false; 
      std::vector<unsigned int>
      m_feature_indices; 
      std::vector<unsigned int>
      m_spectator_indices; 
      std::vector<unsigned int>
      m_event_indices; 
      Dataset& m_dataset; 
    };
 
    class CombinedDataset : public Dataset {
 
    public:
      CombinedDataset(const GeneralOptions& general_options, Dataset& signal_dataset, Dataset& background_dataset);
 
      virtual unsigned int getNumberOfFeatures() const override { return m_signal_dataset.getNumberOfFeatures(); }
 
      virtual unsigned int getNumberOfSpectators() const override { return m_signal_dataset.getNumberOfSpectators(); }
 
      virtual unsigned int getNumberOfEvents() const override { return m_signal_dataset.getNumberOfEvents() + m_background_dataset.getNumberOfEvents(); }
 
      virtual void loadEvent(unsigned int iEvent) override;
 
      virtual std::vector<float> getFeature(unsigned int iFeature) override;
 
      virtual std::vector<float> getSpectator(unsigned int iSpectator) override;
 
    private:
      Dataset& m_signal_dataset; 
      Dataset& m_background_dataset; 
    };
 
 
    class ROOTDataset : public Dataset {
 
    public:
      explicit ROOTDataset(const GeneralOptions& _general_options);
 
      virtual unsigned int getNumberOfFeatures() const override { return m_input.size(); }
 
      virtual unsigned int getNumberOfSpectators() const override { return m_spectators.size(); }
 
      virtual unsigned int getNumberOfEvents() const override
      {
        return (m_general_options.m_max_events == 0) ? m_tree->GetEntries() : m_general_options.m_max_events;
      }
 
      virtual void loadEvent(unsigned int event) override;
 
      virtual std::vector<float> getFeature(unsigned int iFeature) override;
 
      virtual std::vector<float> getWeights() override;
 
      virtual std::vector<float> getSpectator(unsigned int iSpectator) override;
 
      virtual ~ROOTDataset();
 
    protected:
      typedef std::variant<double, float, int, bool> RootDatasetVarVariant;
 
      TChain* m_tree = nullptr; 
      std::vector<RootDatasetVarVariant> m_input_variant; 
      std::vector<RootDatasetVarVariant>
      m_spectators_variant; 
      RootDatasetVarVariant m_weight_variant; 
      RootDatasetVarVariant m_target_variant; 
    private:
 
      template<class T>
      std::vector<float> getVectorFromTTree(const std::string& variableType, const std::string& branchName, T& memberVariableTarget);
 
      std::vector<float> getVectorFromTTreeVariant(const std::string& variableType, const std::string& branchName,
                                                   RootDatasetVarVariant& memberVariableTarget);
 
      void setRootInputType();
 
      template<class T>
      void setScalarVariableAddress(const std::string& variableType, const std::string& variableName, T& variableTarget);
 
      void setScalarVariableAddressVariant(const std::string& variableType, const std::string& variableName,
                                           RootDatasetVarVariant& variableTarget);
 
      template<class T>
      void setVectorVariableAddress(const std::string& variableType, const std::vector<std::string>& variableName,
                                    T& variableTargets);
 
      void setVectorVariableAddressVariant(const std::string& variableType, const std::vector<std::string>& variableName,
                                           std::vector<RootDatasetVarVariant>& varVariantTargets);
 
      void setTargetRootInputType();
 
 
      void setBranchAddresses();
 
      bool checkForBranch(TTree*, const std::string&) const;
 
      float castVarVariantToFloat(RootDatasetVarVariant&) const;
 
      void initialiseVarVariantType(const std::string, RootDatasetVarVariant&);
 
      void initialiseVarVariantForBranch(const std::string, RootDatasetVarVariant&);
    };
  }
}
#endif