src/concentrator/HGCalConcentratorAutoEncoderImpl.cc

0001 #include "L1Trigger/L1THGCal/interface/concentrator/HGCalConcentratorAutoEncoderImpl.h"
0002 #include "DataFormats/ForwardDetId/interface/HGCalTriggerDetId.h"
0003 #include <iomanip>
0004
0005 // Following example of implementing graphloading from here:
0006 // https://gitlab.cern.ch/mrieger/CMSSW-TensorFlowExamples/-/blob/master/GraphLoading/
0007
0008 HGCalConcentratorAutoEncoderImpl::HGCalConcentratorAutoEncoderImpl(const edm::ParameterSet& conf)
0009     : cellRemap_(conf.getParameter<std::vector<int>>("cellRemap")),
0010       cellRemapNoDuplicates_(conf.getParameter<std::vector<int>>("cellRemapNoDuplicates")),
0011       encoderShape_(conf.getParameter<std::vector<uint>>("encoderShape")),
0012       decoderShape_(conf.getParameter<std::vector<uint>>("decoderShape")),
0013       bitsPerInput_(conf.getParameter<int>("nBitsPerInput")),
0014       maxBitsPerOutput_(conf.getParameter<int>("maxBitsPerOutput")),
0015       outputBitsPerLink_(conf.getParameter<std::vector<int>>("bitsPerLink")),
0016       modelFilePaths_(conf.getParameter<std::vector<edm::ParameterSet>>("modelFiles")),
0017       linkToGraphMap_(conf.getParameter<std::vector<unsigned int>>("linkToGraphMap")),
0018       zeroSuppresionThreshold_(conf.getParameter<double>("zeroSuppresionThreshold")),
0019       bitShiftNormalization_(conf.getParameter<bool>("bitShiftNormalization")),
0020       saveEncodedValues_(conf.getParameter<bool>("saveEncodedValues")),
0021       preserveModuleSum_(conf.getParameter<bool>("preserveModuleSum")) {
0022   // find total size of the expected input shape
0023   // used for checking the maximum size used in cell Remap
0024   nInputs_ = 1;
0025   for (const auto& i : encoderShape_) {
0026     nInputs_ *= i;
0027   }
0028
0029   // check the size of the inputs shapes
0030   if (encoderShape_.size() != encoderTensorDims_) {
0031     throw cms::Exception("BadInitialization")
0032         << "Encoder input shapes are currently expected to be " << encoderTensorDims_ << " values";
0033   }
0034
0035   if (decoderShape_.size() != decoderTensorDims_) {
0036     throw cms::Exception("BadInitialization")
0037         << "Encoder input shapes are currently expected to be " << decoderTensorDims_ << " values long";
0038   }
0039
0040   if (cellRemap_.size() != nInputs_) {
0041     throw cms::Exception("BadInitialization")
0042         << "Size of cellRemap (" << cellRemap_.size()
0043         << ") does not agree with the total size specified for the encoder inputs based on the encoderShape variable ("
0044         << nInputs_ << ")";
0045   }
0046
0047   if (cellRemap_.size() != cellRemapNoDuplicates_.size()) {
0048     throw cms::Exception("BadInitialization")
0049         << "Size of cellRemap (" << cellRemap_.size() << ") does not agree with size of cellRemapNoDuplicates ("
0050         << cellRemapNoDuplicates_.size() << ")";
0051   }
0052
0053   for (unsigned i = 0; i < cellRemap_.size(); i++) {
0054     if (cellRemap_[i] > nTriggerCells_ - 1) {
0055       throw cms::Exception("BadInitialization")
0056           << "cellRemap value " << cellRemap_[i] << " is larger than the number of trigger cells " << nTriggerCells_;
0057     }
0058     if (cellRemapNoDuplicates_[i] > nTriggerCells_ - 1) {
0059       throw cms::Exception("BadInitialization") << "cellRemapNoDuplicates value " << cellRemapNoDuplicates_[i]
0060                                                 << " is larger than the number of trigger cells " << nTriggerCells_;
0061     }
0062   }
0063
0064   tensorflow::setLogging("0");
0065
0066   for (const auto& modelFilePset : modelFilePaths_) {
0067     std::string encoderPath = modelFilePset.getParameter<edm::FileInPath>("encoderModelFile").fullPath();
0068     std::string decoderPath = modelFilePset.getParameter<edm::FileInPath>("decoderModelFile").fullPath();
0069
0070     graphDef_encoder_ = std::unique_ptr<tensorflow::GraphDef>{tensorflow::loadGraphDef(encoderPath)};
0071
0072     // create a new session and add the graphDef
0073     session_encoder_.push_back(
0074         std::unique_ptr<tensorflow::Session>{tensorflow::createSession(graphDef_encoder_.get())});
0075
0076     graphDef_decoder_ = std::unique_ptr<tensorflow::GraphDef>{tensorflow::loadGraphDef(decoderPath)};
0077
0078     // create a new session and add the graphDef
0079     session_decoder_.push_back(
0080         std::unique_ptr<tensorflow::Session>{tensorflow::createSession(graphDef_decoder_.get())});
0081
0082     //extract encoder tenser names from first graph, check that rest of the names are consistent
0083     if (modelFilePset == modelFilePaths_.front()) {
0084       inputTensorName_encoder_ = graphDef_encoder_.get()->node(0).name();
0085       outputTensorName_encoder_ = graphDef_encoder_.get()->node(graphDef_encoder_.get()->node_size() - 1).name();
0086       inputTensorName_decoder_ = graphDef_decoder_.get()->node(0).name();
0087       outputTensorName_decoder_ = graphDef_decoder_.get()->node(graphDef_decoder_.get()->node_size() - 1).name();
0088     } else {
0089       if (inputTensorName_encoder_ != graphDef_encoder_.get()->node(0).name()) {
0090         throw cms::Exception("BadInitialization") << "provided list of encoder graphs have different input nodes";
0091       }
0092       if (outputTensorName_encoder_ != graphDef_encoder_.get()->node(graphDef_encoder_.get()->node_size() - 1).name()) {
0093         throw cms::Exception("BadInitialization") << "provided list of encoder graphs have different output nodes";
0094       }
0095       if (inputTensorName_decoder_ != graphDef_decoder_.get()->node(0).name()) {
0096         throw cms::Exception("BadInitialization") << "provided list of decoder graphs have different input nodes";
0097       }
0098       if (outputTensorName_decoder_ != graphDef_decoder_.get()->node(graphDef_decoder_.get()->node_size() - 1).name()) {
0099         throw cms::Exception("BadInitialization") << "provided list of decoder graphs have different output nodes";
0100       }
0101     }
0102   }
0103
0104   // check that the appropriate number of links have been specified
0105   if (linkToGraphMap_.size() <= maxNumberOfLinks_) {
0106     throw cms::Exception("BadInitialization")
0107         << "Autoencoder graph number must be specified for all link allocation possibilities. Only "
0108         << linkToGraphMap_.size() << " values specified while " << maxNumberOfLinks_ << "links are possible";
0109   }
0110
0111   // check that all graph indices specified exist in the model file lists
0112   for (const auto& graphNumber : linkToGraphMap_) {
0113     if (graphNumber >= modelFilePaths_.size()) {
0114       throw cms::Exception("BadInitialization")
0115           << "Autoencoder graph number  " << graphNumber << " is larger than the size of the provided list of graphs "
0116           << modelFilePaths_.size();
0117     }
0118   }
0119 }
0120
0121 void HGCalConcentratorAutoEncoderImpl::select(unsigned nLinks,
0122                                               const std::vector<l1t::HGCalTriggerCell>& trigCellVecInput,
0123                                               std::vector<l1t::HGCalTriggerCell>& trigCellVecOutput,
0124                                               std::vector<l1t::HGCalConcentratorData>& ae_encodedLayer_Output) {
0125   std::array<double, nTriggerCells_> mipPt;
0126   std::array<double, nTriggerCells_> uncompressedCharge;
0127   std::array<double, nTriggerCells_> compressedCharge;
0128   std::array<double, maxAEInputSize_> ae_inputArray;
0129   std::array<double, nTriggerCells_> ae_outputArray;
0130
0131   //reset inputs to 0 to account for zero suppressed trigger cells
0132   mipPt.fill(0);
0133   uncompressedCharge.fill(0);
0134   compressedCharge.fill(0);
0135   ae_inputArray.fill(0);
0136   ae_outputArray.fill(0);
0137
0138   double modSum = 0;
0139
0140   int bitsPerOutput = outputBitsPerLink_.at(nLinks);
0141
0142   // largest expected input and output values, used for bit truncation
0143   // values of -1 for the number of bits used to keep full precision, in which case the MaxIntSize variables are not used
0144   double inputMaxIntSize = 1;
0145   if (bitsPerInput_ > 0)
0146     inputMaxIntSize = 1 << bitsPerInput_;
0147   double outputMaxIntSize = 1;
0148   if (bitsPerOutput > 0)
0149     outputMaxIntSize = 1 << bitsPerOutput;
0150   double outputMaxIntSizeGlobal = 1;
0151   if (maxBitsPerOutput_ > 0)
0152     outputMaxIntSizeGlobal = 1 << maxBitsPerOutput_;
0153
0154   for (const auto& trigCell : trigCellVecInput) {
0155     if (triggerTools_.isScintillator(trigCell.detId()))
0156       return;  //currently, only silicon modules are setup to work (mapping of scinillators would be different, and needs to be investigated)
0157
0158     HGCalTriggerDetId id(trigCell.detId());
0159     uint cellu = id.triggerCellU();
0160     uint cellv = id.triggerCellV();
0161     int inputIndex = cellUVremap_[cellu][cellv];
0162     if (inputIndex < 0) {
0163       throw cms::Exception("BadInitialization")
0164           << "Invalid index provided for trigger cell u=" << cellu << " v=" << cellv << " in cellUVRemap[" << cellu
0165           << "][" << cellv << "]";
0166     }
0167
0168     mipPt[inputIndex] = trigCell.mipPt();
0169     uncompressedCharge[inputIndex] = trigCell.uncompressedCharge();
0170     compressedCharge[inputIndex] = trigCell.compressedCharge();
0171
0172     modSum += trigCell.mipPt();
0173   }
0174
0175   double normalization = modSum;
0176   if (modSum > 0) {
0177     //Use a bit shift normalization like will be implemented in ECON, rather than floating point sum
0178     //Normalizes to the MSB value of the module sum
0179     if (bitShiftNormalization_) {
0180       int msb = int(log2(modSum));
0181       normalization = pow(2, msb);
0182     }
0183
0184     //normalize inputs to module sum
0185     for (uint i = 0; i < nInputs_; i++) {
0186       int remapIndex = cellRemap_[i];
0187       if (remapIndex < 0)
0188         continue;
0189       ae_inputArray[i] = mipPt[remapIndex] / normalization;
0190       //round to precision of input, if bitsPerInput_ is -1 keep full precision
0191       if (bitsPerInput_ > 0) {
0192         ae_inputArray[i] = std::round(ae_inputArray[i] * inputMaxIntSize) / inputMaxIntSize;
0193       }
0194     }
0195   }
0196
0197   tensorflow::Tensor encoder_input(tensorflow::DT_FLOAT,
0198                                    {encoderShape_[0], encoderShape_[1], encoderShape_[2], encoderShape_[3]});
0199
0200   float* d = encoder_input.flat<float>().data();
0201   for (uint i = 0; i < nInputs_; i++, d++) {
0202     *d = ae_inputArray[i];
0203   }
0204
0205   int graphIndex = linkToGraphMap_.at(nLinks);
0206
0207   std::vector<tensorflow::Tensor> encoder_outputs;
0208   tensorflow::run(session_encoder_.at(graphIndex).get(),
0209                   {{inputTensorName_encoder_, encoder_input}},
0210                   {outputTensorName_encoder_},
0211                   &encoder_outputs);
0212
0213   if (encoder_outputs.empty()) {
0214     throw cms::Exception("BadInitialization") << "Autoencoder graph returning empty output vector";
0215   }
0216
0217   d = encoder_outputs[0].flat<float>().data();
0218   for (int i = 0; i < encoder_outputs[0].NumElements(); i++, d++) {
0219     ae_encodedLayer_[i] = *d;
0220     //truncate the encoded layer bits
0221     if (bitsPerOutput > 0 && maxBitsPerOutput_ > 0) {
0222       ae_encodedLayer_[i] = std::round(ae_encodedLayer_[i] * outputMaxIntSize) / outputMaxIntSize;
0223     }
0224   }
0225
0226   tensorflow::Tensor decoder_input(tensorflow::DT_FLOAT, {decoderShape_[0], decoderShape_[1]});
0227   d = decoder_input.flat<float>().data();
0228   for (int i = 0; i < nEncodedLayerNodes_; i++, d++) {
0229     *d = ae_encodedLayer_[i];
0230   }
0231
0232   std::vector<tensorflow::Tensor> decoder_outputs;
0233   tensorflow::run(session_decoder_.at(graphIndex).get(),
0234                   {{inputTensorName_decoder_, decoder_input}},
0235                   {outputTensorName_decoder_},
0236                   &decoder_outputs);
0237
0238   double outputSum = 0.;
0239
0240   d = decoder_outputs[0].flat<float>().data();
0241   for (uint i = 0; i < nInputs_; i++, d++) {
0242     int remapIndex = cellRemapNoDuplicates_[i];
0243     if (remapIndex < 0)
0244       continue;
0245     outputSum += *d * normalization;
0246     ae_outputArray[remapIndex] = *d;
0247   }
0248
0249   double renormalizationFactor = 1.;
0250   if (preserveModuleSum_) {
0251     renormalizationFactor = modSum / outputSum;
0252   }
0253
0254   // Add data back into trigger cells
0255   if (modSum > 0) {
0256     //get detID for everything but cell, take first entry detID and subtract off cellU and cellV contribution
0257     HGCalTriggerDetId id(trigCellVecInput.at(0).detId());
0258     int subdet = id.subdet();
0259     int zp = id.zside();
0260     int type = id.type();
0261     int layer = id.layer();
0262     int waferU = id.waferU();
0263     int waferV = id.waferV();
0264
0265     //use first TC to find mipPt conversions to Et and ADC
0266     float mipPtToEt_conv = trigCellVecInput[0].et() / trigCellVecInput[0].mipPt();
0267     float mipToADC_conv = trigCellVecInput[0].hwPt() / (trigCellVecInput[0].mipPt() * cosh(trigCellVecInput[0].eta()));
0268
0269     for (int i = 0; i < nTriggerCells_; i++) {
0270       if (ae_outputArray[i] > 0) {
0271         int cellU = ae_outputCellU_[i];
0272         int cellV = ae_outputCellV_[i];
0273
0274         HGCalTriggerDetId id(subdet, zp, type, layer, waferU, waferV, cellU, cellV);
0275
0276         if (!triggerTools_.getTriggerGeometry()->validTriggerCell(id))
0277           continue;
0278
0279         GlobalPoint point = triggerTools_.getTCPosition(id);
0280
0281         double mipPt = ae_outputArray[i] * normalization * renormalizationFactor;
0282         double adc = mipPt * cosh(point.eta()) * mipToADC_conv;
0283         double et = mipPt * mipPtToEt_conv;
0284
0285         if (mipPt < zeroSuppresionThreshold_)
0286           continue;
0287
0288         l1t::HGCalTriggerCell triggerCell(reco::LeafCandidate::LorentzVector(), adc, 0, 0, 0, id);
0289         //Keep the pre-autoencoder charge for this cell
0290         triggerCell.setUncompressedCharge(uncompressedCharge[i]);
0291         triggerCell.setCompressedCharge(compressedCharge[i]);
0292         triggerCell.setMipPt(mipPt);
0293
0294         math::PtEtaPhiMLorentzVector p4(et, point.eta(), point.phi(), 0.);
0295
0296         triggerCell.setP4(p4);
0297         triggerCell.setPosition(point);
0298
0299         trigCellVecOutput.push_back(triggerCell);
0300       }
0301     }
0302
0303     if (saveEncodedValues_) {
0304       id = HGCalTriggerDetId(subdet, zp, type, layer, waferU, waferV, 0, 0);
0305       for (int i = 0; i < nEncodedLayerNodes_; i++) {
0306         l1t::HGCalConcentratorData encodedLayerData(ae_encodedLayer_[i] * outputMaxIntSizeGlobal, i, id);
0307         ae_encodedLayer_Output.push_back(encodedLayerData);
0308       }
0309     }
0310   }
0311 }