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