DeepTauBase.cc

Go to the documentation of this file.

/*
 * \class DeepTauBase
 *
 * Implementation of the base class for tau identification using Deep NN.
 *
 * \author Konstantin Androsov, INFN Pisa
 * \author Maria Rosaria Di Domenico, University of Siena & INFN Pisa
 */


#include "RecoTauTag/RecoTau/interface/DeepTauBase.h"

namespace deep_tau {

TauWPThreshold::TauWPThreshold(const std::string& cut_str)
{
    bool simple_value = false;
    try {
        size_t pos = 0;
        value_ = std::stod(cut_str, &pos);
        simple_value = (pos == cut_str.size());
    } catch(std::invalid_argument&) {
    } catch(std::out_of_range&) {
    }
    if(!simple_value) {
        static const std::string prefix = "[&](double *x, double *p) { const int decayMode = p[0];"
                                          "const double pt = p[1]; const double eta = p[2];";
        static const int n_params = 3;
        static const auto handler = [](int, Bool_t, const char*, const char*) -> void {};

        const std::string fn_str = prefix + cut_str + "}";
        auto old_handler = SetErrorHandler(handler);
        fn_ = std::make_unique<TF1>("fn_", fn_str.c_str(), 0, 1, n_params);
        SetErrorHandler(old_handler);
        if(!fn_->IsValid())
            throw cms::Exception("TauWPThreshold: invalid formula") << "Invalid WP cut formula = '" << cut_str << "'.";
    }
}

double TauWPThreshold::operator()(const pat::Tau& tau) const
{
    if(!fn_)
        return value_;
    fn_->SetParameter(0, tau.decayMode());
    fn_->SetParameter(1, tau.pt());
    fn_->SetParameter(2, tau.eta());
    return fn_->Eval(0);
}

DeepTauBase::Output::ResultMap DeepTauBase::Output::get_value(const edm::Handle<TauCollection>& taus,
                                                              const tensorflow::Tensor& pred,
                                                              const WPMap& working_points) const
{
    ResultMap output;
    output[""] = std::make_unique<TauDiscriminator>(TauRefProd(taus));
    for(const auto& wp : working_points)
        output[wp.first] = std::make_unique<TauDiscriminator>(TauRefProd(taus));

    for(size_t tau_index = 0; tau_index < taus->size(); ++tau_index) {
        float x = 0;
        for(size_t num_elem : num_)
            x += pred.matrix<float>()(tau_index, num_elem);
        if(x != 0 && !den_.empty()) {
            float den_val = 0;
            for(size_t den_elem : den_)
                den_val += pred.matrix<float>()(tau_index, den_elem);
            x = den_val != 0 ? x / den_val : std::numeric_limits<float>::max();
        }
        output[""]->setValue(tau_index, x);
        for(const auto& wp : working_points) {
            const auto& tau = taus->at(tau_index);
            const bool pass = x > (*wp.second)(tau);
            output[wp.first]->setValue(tau_index, pass);
        }
    }
    return output;
}

DeepTauBase::DeepTauBase(const edm::ParameterSet& cfg, const OutputCollection& outputCollection,
                         const DeepTauCache* cache) :
    tausToken_(consumes<TauCollection>(cfg.getParameter<edm::InputTag>("taus"))),
    pfcandToken_(consumes<pat::PackedCandidateCollection>(cfg.getParameter<edm::InputTag>("pfcands"))),
    vtxToken_(consumes<reco::VertexCollection>(cfg.getParameter<edm::InputTag>("vertices"))),
    outputs_(outputCollection),
    cache_(cache)
{
    for(const auto& output_desc : outputs_) {
        produces<TauDiscriminator>(output_desc.first);
        const auto& cut_pset = cfg.getParameter<edm::ParameterSet>(output_desc.first + "WP");
        for(const std::string& wp_name : cut_pset.getParameterNames()) {
            const auto& cut_str = cut_pset.getParameter<std::string>(wp_name);
            workingPoints_[output_desc.first][wp_name] = std::make_unique<Cutter>(cut_str);
            produces<TauDiscriminator>(output_desc.first + wp_name);
        }
    }
}

void DeepTauBase::produce(edm::Event& event, const edm::EventSetup& es)
{
    edm::Handle<TauCollection> taus;
    event.getByToken(tausToken_, taus);

    const tensorflow::Tensor& pred = getPredictions(event, es, taus);
    createOutputs(event, pred, taus);
}

void DeepTauBase::createOutputs(edm::Event& event, const tensorflow::Tensor& pred, edm::Handle<TauCollection> taus)
{
    for(const auto& output_desc : outputs_) {
        auto result_map = output_desc.second.get_value(taus, pred, workingPoints_.at(output_desc.first));
        for(auto& result : result_map)
            event.put(std::move(result.second), output_desc.first + result.first);
    }
}

std::unique_ptr<DeepTauCache> DeepTauBase::initializeGlobalCache(const edm::ParameterSet& cfg )
{
    const auto graph_name_vector = cfg.getParameter<std::vector<std::string>>("graph_file");
    std::map<std::string, std::string> graph_names;
    for(const auto& entry : graph_name_vector) {
        const size_t sep_pos = entry.find(':');
        std::string entry_name, graph_file;
        if(sep_pos != std::string::npos) {
            entry_name = entry.substr(0, sep_pos);
            graph_file = entry.substr(sep_pos + 1);
        } else {
            entry_name = "";
            graph_file = entry;
        }
        graph_file = edm::FileInPath(graph_file).fullPath();
        if(graph_names.count(entry_name))
            throw cms::Exception("DeepTauCache") << "Duplicated graph entries";
        graph_names[entry_name] = graph_file;
    }
    bool mem_mapped = cfg.getParameter<bool>("mem_mapped");
    return std::make_unique<DeepTauCache>(graph_names, mem_mapped);
}

DeepTauCache::DeepTauCache(const std::map<std::string, std::string>& graph_names, bool mem_mapped)
{
    for(const auto& graph_entry : graph_names) {
        tensorflow::SessionOptions options;
        tensorflow::setThreading(options, 1, "no_threads");

        const std::string& entry_name = graph_entry.first;
        const std::string& graph_file = graph_entry.second;
        if(mem_mapped) {
            memmappedEnv_[entry_name] = std::make_unique<tensorflow::MemmappedEnv>(tensorflow::Env::Default());
            const tensorflow::Status mmap_status = memmappedEnv_.at(entry_name)->InitializeFromFile(graph_file);
            if(!mmap_status.ok()) {
                throw cms::Exception("DeepTauCache: unable to initalize memmapped environment for ")
                    << graph_file << ". \n" << mmap_status.ToString();
            }

            graphs_[entry_name] = std::make_unique<tensorflow::GraphDef>();
            const tensorflow::Status load_graph_status = ReadBinaryProto(memmappedEnv_.at(entry_name).get(),
                                                                         tensorflow::MemmappedFileSystem::kMemmappedPackageDefaultGraphDef,
                                                                         graphs_.at(entry_name).get());
            if(!load_graph_status.ok())
                throw cms::Exception("DeepTauCache: unable to load graph from ") << graph_file << ". \n"
                                                                                  << mmap_status.ToString();
            options.config.mutable_graph_options()->mutable_optimizer_options()->set_opt_level(::tensorflow::OptimizerOptions::L0);
            options.env = memmappedEnv_.at(entry_name).get();

            sessions_[entry_name] = tensorflow::createSession(graphs_.at(entry_name).get(), options);

        } else {
            graphs_[entry_name].reset(tensorflow::loadGraphDef(graph_file));
            sessions_[entry_name] = tensorflow::createSession(graphs_.at(entry_name).get(), options);
        }
    }
}

DeepTauCache::~DeepTauCache()
{
    for(auto& session_entry : sessions_)
        tensorflow::closeSession(session_entry.second);
}

} // namespace deep_tau