Merge branch 'master' into changept

SegmentLinking · May 29, 2024 · f68c1e7 · f68c1e7
2 parents 8b75100 + f93339a
commit f68c1e7
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 7 deletions.
diff --git a/SDL/PixelTriplet.h b/SDL/PixelTriplet.h
@@ -2557,9 +2557,15 @@ namespace SDL {
                     mdsInGPU.anchorRt[fourthMDIndex],
                     mdsInGPU.anchorRt[fifthMDIndex]};
 
-    rzChiSquared = computePT5RZChiSquared(acc, modulesInGPU, lowerModuleIndices, rtPix, zPix, rts, zs);
+    rzChiSquared = 0;
 
-    if (/*pixelRadius*/ 0 < 5.0f * kR1GeVf) {  // FIXME: pixelRadius is not defined yet
+    //get the appropriate radii and centers
+    centerX = segmentsInGPU.circleCenterX[pixelSegmentArrayIndex];
+    centerY = segmentsInGPU.circleCenterY[pixelSegmentArrayIndex];
+    pixelRadius = segmentsInGPU.circleRadius[pixelSegmentArrayIndex];
+
+    if (pixelRadius < 5.0f * kR1GeVf) {  //only apply r-z chi2 cuts for <5GeV tracks
+      rzChiSquared = computePT5RZChiSquared(acc, modulesInGPU, lowerModuleIndices, rtPix, zPix, rts, zs);
       pass = pass and passPT5RZChiSquaredCuts(modulesInGPU,
                                               lowerModuleIndex1,
                                               lowerModuleIndex2,
@@ -2583,11 +2589,6 @@ namespace SDL {
                    mdsInGPU.anchorY[fourthMDIndex],
                    mdsInGPU.anchorY[fifthMDIndex]};
 
-    //get the appropriate radii and centers
-    centerX = segmentsInGPU.circleCenterX[pixelSegmentArrayIndex];
-    centerY = segmentsInGPU.circleCenterY[pixelSegmentArrayIndex];
-    pixelRadius = segmentsInGPU.circleRadius[pixelSegmentArrayIndex];
-
     float T5CenterX = quintupletsInGPU.regressionG[quintupletIndex];
     float T5CenterY = quintupletsInGPU.regressionF[quintupletIndex];
     quintupletRadius = quintupletsInGPU.regressionRadius[quintupletIndex];

diff --git a/code/core/write_sdl_ntuple.cc b/code/core/write_sdl_ntuple.cc
@@ -136,6 +136,23 @@ void createOptionalOutputBranches()
     ana.tx->createBranch<vector<float>>("t5_rzChiSquared");
     ana.tx->createBranch<vector<float>>("t5_nonAnchorChiSquared");
 
+    // Occupancy branches
+    ana.tx->createBranch<vector<int>>("module_layers");
+    ana.tx->createBranch<vector<int>>("module_subdets");
+    ana.tx->createBranch<vector<int>>("module_rings");
+    ana.tx->createBranch<vector<int>>("module_rods");
+    ana.tx->createBranch<vector<int>>("module_modules");
+    ana.tx->createBranch<vector<bool>>("module_isTilted");
+    ana.tx->createBranch<vector<float>>("module_eta");
+    ana.tx->createBranch<vector<float>>("module_r");
+    ana.tx->createBranch<vector<int>>("md_occupancies");
+    ana.tx->createBranch<vector<int>>("sg_occupancies");
+    ana.tx->createBranch<vector<int>>("t3_occupancies");
+    ana.tx->createBranch<int>("tc_occupancies");
+    ana.tx->createBranch<vector<int>>("t5_occupancies");
+    ana.tx->createBranch<int>("pT3_occupancies");
+    ana.tx->createBranch<int>("pT5_occupancies");
+
 #endif
 }
 
@@ -297,10 +314,76 @@ void setOptionalOutputBranches(SDL::Event<SDL::Acc>* event)
     setPixelQuintupletOutputBranches(event);
     setQuintupletOutputBranches(event);
     setPixelTripletOutputBranches(event);
+    setOccupancyBranches(event);
 
 #endif
 }
 
+//________________________________________________________________________________________________________________________________
+void setOccupancyBranches(SDL::Event<SDL::Acc>* event)
+{
+    SDL::modules& modulesInGPU = (*event->getModules());
+    SDL::miniDoublets& mdsInGPU = (*event->getMiniDoublets());
+    SDL::segments& segmentsInGPU = (*event->getSegments());
+    SDL::triplets& tripletsInGPU = (*event->getTriplets());
+    SDL::quintuplets&  quintupletsInGPU = (*event->getQuintuplets());
+    SDL::pixelQuintuplets& pixelQuintupletsInGPU = (*event->getPixelQuintuplets());
+    SDL::pixelTriplets& pixelTripletsInGPU = (*event->getPixelTriplets());
+    SDL::trackCandidates& trackCandidatesInGPU = (*event->getTrackCandidates());
+
+    std::vector<int> moduleLayer;
+    std::vector<int> moduleSubdet;
+    std::vector<int> moduleRing;
+    std::vector<int> moduleRod;
+    std::vector<int> moduleModule;
+    std::vector<float> moduleEta;
+    std::vector<float> moduleR;
+    std::vector<bool> moduleIsTilted;
+    std::vector<int> trackCandidateOccupancy;
+    std::vector<int> tripletOccupancy;
+    std::vector<int> segmentOccupancy;
+    std::vector<int> mdOccupancy;
+    std::vector<int> quintupletOccupancy;
+
+    for(unsigned int lowerIdx = 0; lowerIdx <= *(modulesInGPU.nLowerModules); lowerIdx++)
+    {
+        //layer = 0, subdet = 0 => pixel module
+        moduleLayer.push_back(modulesInGPU.layers[lowerIdx]);
+        moduleSubdet.push_back(modulesInGPU.subdets[lowerIdx]);
+        moduleRing.push_back(modulesInGPU.rings[lowerIdx]);
+        moduleRod.push_back(modulesInGPU.rods[lowerIdx]);
+        moduleEta.push_back(modulesInGPU.eta[lowerIdx]);
+        moduleR.push_back(modulesInGPU.r[lowerIdx]);
+        bool isTilted = (modulesInGPU.subdets[lowerIdx] == 5 and modulesInGPU.sides[lowerIdx] != 3);
+        moduleIsTilted.push_back(isTilted);
+        moduleModule.push_back(modulesInGPU.modules[lowerIdx]);
+        segmentOccupancy.push_back(segmentsInGPU.totOccupancySegments[lowerIdx]);
+        mdOccupancy.push_back(mdsInGPU.totOccupancyMDs[lowerIdx]);
+
+        if(lowerIdx < *(modulesInGPU.nLowerModules))
+        {
+            quintupletOccupancy.push_back(quintupletsInGPU.totOccupancyQuintuplets[lowerIdx]);
+            tripletOccupancy.push_back(tripletsInGPU.totOccupancyTriplets[lowerIdx]);
+        }
+    }
+
+    ana.tx->setBranch<vector<int>>("module_layers", moduleLayer);
+    ana.tx->setBranch<vector<int>>("module_subdets", moduleSubdet);
+    ana.tx->setBranch<vector<int>>("module_rings", moduleRing);
+    ana.tx->setBranch<vector<int>>("module_rods", moduleRod);
+    ana.tx->setBranch<vector<int>>("module_modules", moduleModule);
+    ana.tx->setBranch<vector<bool>>("module_isTilted", moduleIsTilted);
+    ana.tx->setBranch<vector<float>>("module_eta", moduleEta);
+    ana.tx->setBranch<vector<float>>("module_r", moduleR);
+    ana.tx->setBranch<vector<int>>("md_occupancies", mdOccupancy);
+    ana.tx->setBranch<vector<int>>("sg_occupancies", segmentOccupancy);
+    ana.tx->setBranch<vector<int>>("t3_occupancies", tripletOccupancy);
+    ana.tx->setBranch<int>("tc_occupancies", *(trackCandidatesInGPU.nTrackCandidates));
+    ana.tx->setBranch<int>("pT3_occupancies", *(pixelTripletsInGPU.totOccupancyPixelTriplets));
+    ana.tx->setBranch<vector<int>>("t5_occupancies", quintupletOccupancy);
+    ana.tx->setBranch<int>("pT5_occupancies", *(pixelQuintupletsInGPU.totOccupancyPixelQuintuplets));
+}
+
 //________________________________________________________________________________________________________________________________
 void setPixelQuintupletOutputBranches(SDL::Event<SDL::Acc>* event)
 {

diff --git a/code/core/write_sdl_ntuple.h b/code/core/write_sdl_ntuple.h
@@ -20,6 +20,7 @@ void createGnnNtupleBranches();
 void fillOutputBranches(SDL::Event<SDL::Acc>* event);
 void setOutputBranches(SDL::Event<SDL::Acc>* event);
 void setOptionalOutputBranches(SDL::Event<SDL::Acc>* event);
+void setOccupancyBranches(SDL::Event<SDL::Acc>* event);
 void setPixelQuintupletOutputBranches(SDL::Event<SDL::Acc>* event);
 void setQuintupletOutputBranches(SDL::Event<SDL::Acc>* event);
 void setPixelTripletOutputBranches(SDL::Event<SDL::Acc> *event);