MVAUtils/src/TMVAZipReader.cc

0001 #include "CommonTools/MVAUtils/interface/TMVAZipReader.h"
0002 #include "FWCore/Utilities/interface/Exception.h"
0003
0004 #include <cstdio>
0005 #include <cstdlib>
0006 #include <zlib.h>
0007
0008 using namespace std;
0009
0010 // From http://stackoverflow.com/questions/874134/find-if-string-endswith-another-string-in-c
0011 bool reco::details::hasEnding(std::string const& fullString, std::string const& ending) {
0012   if (fullString.length() >= ending.length()) {
0013     return (0 == fullString.compare(fullString.length() - ending.length(), ending.length(), ending));
0014   } else {
0015     return false;
0016   }
0017 }
0018
0019 char* reco::details::readGzipFile(const std::string& weightFile) {
0020   FILE* f = fopen(weightFile.c_str(), "r");
0021   if (f == nullptr) {
0022     throw cms::Exception("InvalidFileState") << "Failed to open MVA file = " << weightFile << " !!\n";
0023   }
0024   int magic;
0025   int size;
0026   fread(&magic, 4, 1, f);
0027   fseek(f, -4, SEEK_END);
0028   fread(&size, 4, 1, f);
0029   fclose(f);
0030   // printf("%x, %i\n", magic, size);
0031
0032   gzFile file = gzopen(weightFile.c_str(), "r");
0033
0034   int bytes_read;
0035   char* buffer = (char*)malloc(size);
0036   bytes_read = gzread(file, buffer, size - 1);
0037   buffer[bytes_read] = '\0';
0038   if (!gzeof(file)) {
0039     int err;
0040     const char* error_string;
0041     error_string = gzerror(file, &err);
0042     if (err) {
0043       free(buffer);
0044       throw cms::Exception("InvalidFileState") << "Error while reading gzipped file = " << weightFile << " !!\n"
0045                                                << error_string;
0046     }
0047   }
0048   gzclose(file);
0049   return buffer;
0050 }
0051
0052 TMVA::IMethod* reco::details::loadTMVAWeights(TMVA::Reader* reader,
0053                                               const std::string& method,
0054                                               const std::string& weightFile,
0055                                               bool verbose) {
0056   TMVA::IMethod* ptr = nullptr;
0057
0058   verbose = false;
0059   if (verbose)
0060     std::cout << "Booking TMVA Reader with " << method << " and weight file: " << weightFile << std::endl;
0061
0062   if (reco::details::hasEnding(weightFile, ".xml")) {
0063     if (verbose)
0064       std::cout << "Weight file is pure xml." << std::endl;
0065     // Let TMVA read the file
0066     ptr = reader->BookMVA(method, weightFile);
0067   } else if (reco::details::hasEnding(weightFile, ".gz") || reco::details::hasEnding(weightFile, ".gzip")) {
0068     if (verbose)
0069       std::cout << "Unzipping file." << std::endl;
0070     char* c = readGzipFile(weightFile);
0071
0072     // We can't use tmpnam, gcc emits a warning about security.
0073     // This is also technically insecure in the same way, since we append
0074     // a suffix and then open another file.
0075     char tmpFilename[] = "/tmp/tmva.XXXXXX";
0076     int fdToUselessFile = mkstemp(tmpFilename);
0077     std::string weight_file_name(tmpFilename);
0078     weight_file_name += ".xml";
0079     FILE* theActualFile = fopen(weight_file_name.c_str(), "w");
0080     if (theActualFile != nullptr) {
0081       // write xml
0082       fputs(c, theActualFile);
0083       fputs("\n", theActualFile);
0084       fclose(theActualFile);
0085       close(fdToUselessFile);
0086     } else {
0087       throw cms::Exception("CannotWriteFile") << "Error while writing file = " << weight_file_name << " !!\n";
0088     }
0089     if (verbose)
0090       std::cout << "Booking MvA" << std::endl;
0091     ptr = reader->BookMVA(method, weight_file_name);
0092     if (verbose)
0093       std::cout << "Cleaning up" << std::endl;
0094     remove(weight_file_name.c_str());
0095     remove(tmpFilename);
0096
0097     // Someday this will work.
0098     // reader->BookMVA(TMVA::Types::Instance().GetMethodType(TString(method)), c);
0099     if (verbose) {
0100       std::cout << "Reader booked" << std::endl;
0101     }
0102     free(c);
0103   } else {
0104     throw cms::Exception("BadTMVAWeightFilename") << "I don't understand the extension on the filename: " << weightFile
0105                                                   << ", it should be .xml, .gz, or .gzip" << std::endl;
0106   }
0107
0108   return ptr;
0109 }