Add conventional tracking into ai-assisted tracking as complementary (#543)

* As supplementary, take conventional tracking with remaining cluters after AI-assisted tracking

* add one more bit into track status to tell if track is from AI-asssisted tracking

Author: tongtongcao

reconstruction/dc/src/main/java/org/jlab/rec/dc/track/Track.java

@@ -66,6 +66,8 @@ public void setFinalStateVec(StateVec finalStateVec) {
     private Segment _singleSuperlayer ;
     private int _fitConvergenceStatus;
     private StateVec finalStateVec ;
+    
+    private boolean isAITrack = false;
 
     public Track() {
     }
@@ -98,6 +100,14 @@ public Segment getSingleSuperlayer() {
     public void setSingleSuperlayer(Segment _singleSuperlayer) {
         this._singleSuperlayer = _singleSuperlayer;
     }
+    
+    public void setIsAITrack(boolean isAITrack){
+        this.isAITrack = isAITrack;
+    }
+    
+    public boolean getIsAITrack(){
+        return isAITrack;
+    }
 
     public int getBitStatus() { 
         int status = 0;
@@ -119,6 +129,9 @@ public int getBitStatus() {
         for(int isl = 0; isl <6; isl++) {
             status |= segmentStatus[isl] << isl*2;
         }
+        
+        status |= (this.isAITrack ? 1 : 0) << 12; // The 13th bit tells if track is from AI-assisted trcking; 1: yes; 0: no
+        
         return status;
     }
 

reconstruction/dc/src/main/java/org/jlab/service/dc/DCHBPostClusterAI.java

@@ -3,6 +3,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.logging.Level;
 import org.jlab.clas.swimtools.Swim;
 import org.jlab.clas.swimtools.Swimmer;
@@ -20,6 +21,12 @@
 import org.jlab.rec.dc.segment.Segment;
 import org.jlab.rec.dc.track.Track;
 import org.jlab.rec.dc.track.TrackCandListFinder;
+import org.jlab.rec.dc.cluster.ClusterFinder;
+import org.jlab.rec.dc.cluster.ClusterFitter;
+import org.jlab.rec.dc.segment.SegmentFinder;
+import org.jlab.rec.dc.cross.CrossMaker;
+import org.jlab.rec.dc.trajectory.Road;
+import org.jlab.rec.dc.trajectory.RoadFinder;
 
 /**
  *
@@ -52,18 +59,15 @@ public boolean processDataEvent(DataEvent event) {
             return true;
         }
 
+        ////// AI-assisted tracking
         /* IO */
         HitReader reader      = new HitReader(this.getBanks(), Constants.getInstance().dcDetector);
         reader.initialize(event);
         RecoBankWriter writer = new RecoBankWriter(this.getBanks());
         // get Field
         Swim dcSwim = new Swim();
-        /* 2 */
-        
-        /* 5 */
         LOGGER.log(Level.FINE, "HB AI process event");
-        /* 7 */
-        /* 8 */
+
         //AI
         List<Track> trkcands = null;
         List<Cross> crosses = null;
@@ -77,7 +81,7 @@ public boolean processDataEvent(DataEvent event) {
         List<Hit> hits = reader.get_DCHits();
         fhits = new ArrayList<>();
         //II) process the hits
-        //1) exit if hit list is empty
+        // Exit if hit list is empty
         if (hits.isEmpty()) {
             return true;
         }
@@ -128,35 +132,182 @@ public boolean processDataEvent(DataEvent event) {
             // remove overlaps
             trkcandFinder.removeOverlappingTracks(trkcands);
             for (Track trk : trkcands) {
+                trk.setIsAITrack(true);
+                
                 // reset the id
                 trk.set_Id(trkId);
                 trkcandFinder.matchHits(trk.getStateVecs(),
                         trk,
                         Constants.getInstance().dcDetector,
                         dcSwim);
-                for (Cross c : trk) {
-                    c.set_CrossDirIntersSegWires();
-                    clusters.add(c.get_Segment1().get_fittedCluster());
-                    clusters.add(c.get_Segment2().get_fittedCluster());
-                    trkcandFinder.setHitDoubletsInfo(c.get_Segment1());
-                    trkcandFinder.setHitDoubletsInfo(c.get_Segment2());
-                    for (FittedHit h1 : c.get_Segment1()) {
-                        h1.set_AssociatedHBTrackID(trkId);
-                        //if(h1.get_AssociatedHBTrackID()>0) 
-                        fhits.add(h1); 
+                trkId++;
+            }
+        }
+        
+        ////// Find tracks by rest of clusters using conventional tracking
+        List<FittedCluster> clustersConv = null;
+        List<Segment> segmentsConv = null;
+        List<Cross> crossesConv = null;
+        List<Track> trkcandsConv = null;
+        
+        //1) read hits from the banks
+        Map<Integer, ArrayList<FittedHit>> hitsConv = reader.read_Hits(event);
+        
+        //2) find clusters from these hits
+        ClusterFinder clusFinder = new ClusterFinder();
+        ClusterFitter cf = new ClusterFitter();
+        clustersConv = clusFinder.RecomposeClusters(hitsConv, Constants.getInstance().dcDetector, cf);
+
+        //3) remove clusters which are on tracks
+        List<FittedCluster> removedClustersConv = new ArrayList();
+        for(FittedCluster cls : clustersConv){
+            boolean flag = false;
+            for(Track trk : trkcands){
+                if(flag) break;
+                for(Cross crs : trk){
+                    if(cls.get_Id() == crs.get_Segment1().get_Id() || cls.get_Id() == crs.get_Segment2().get_Id()) {
+                        removedClustersConv.add(cls);
+                        flag = true;
+                        break;
+                    }
+                }
+            }
+        }
+        clustersConv.removeAll(removedClustersConv);
+        clusters.addAll(clustersConv);
+        
+        //4) find segments from clusters
+        SegmentFinder segFinder = new SegmentFinder();
+        segmentsConv = segFinder.get_Segments(clustersConv,
+                event,
+                Constants.getInstance().dcDetector, false);
+        List<Segment> rmSegsConv = new ArrayList<>();
+        // clean up hit-based segments
+        double trkDocOverCellSize;
+        for (Segment se : segmentsConv) {
+            trkDocOverCellSize = 0;
+            for (FittedHit fh : se.get_fittedCluster()) {
+                trkDocOverCellSize += fh.get_ClusFitDoca() / fh.get_CellSize();
+            }
+            if (trkDocOverCellSize / se.size() > 1.1) {
+                rmSegsConv.add(se);
+            }
+        }
+        segmentsConv.removeAll(rmSegsConv);
+        segments.addAll(segmentsConv);
+        
+        //5) find crosses from segments
+        CrossMaker crossMake = new CrossMaker();
+        crossesConv = crossMake.find_Crosses(segmentsConv, Constants.getInstance().dcDetector);
+        crosses.addAll(crossesConv);
+        
+        //6) find cross lists from crosses
+        CrossList crosslistConv = crossLister.candCrossLists(event, crossesConv,
+                false,
+                null,
+                Constants.getInstance().dcDetector,
+                null,
+                dcSwim, false);       
+        
+        //7) find track candidates with 5 or 6 clusters
+        // track candidates with 6 clusters 
+        trkcandsConv = trkcandFinder.getTrackCands(crosslistConv,
+                Constants.getInstance().dcDetector,
+                Swimmer.getTorScale(),
+                dcSwim, false); 
+        
+        // track candidates with 5 clusters
+        RoadFinder rf = new RoadFinder();
+        List<Road> allRoadsConv = rf.findRoads(segmentsConv, Constants.getInstance().dcDetector);
+        List<Segment> Segs2RoadConv = new ArrayList<>();
+        List<Segment> psegmentsConv = new ArrayList<>();
+        for (Road r : allRoadsConv) { 
+            Segs2RoadConv.clear();
+            int missingSL = -1;
+            for (int ri = 0; ri < 3; ri++) {
+                if (r.get(ri).associatedCrossId == -1) {
+                    if (r.get(ri).get_Superlayer() % 2 == 1) {
+                        missingSL = r.get(ri).get_Superlayer() + 1;
+                    } else {
+                        missingSL = r.get(ri).get_Superlayer() - 1;
                     }
-                    for (FittedHit h2 : c.get_Segment2()) {
-                        h2.set_AssociatedHBTrackID(trkId);
-                        //if(h2.get_AssociatedHBTrackID()>0) 
-                        fhits.add(h2); 
+                }
+            } 
+            if(missingSL==-1) 
+                continue;
+            for (int ri = 0; ri < 3; ri++) {
+                for (Segment s : segmentsConv) {
+                    if (s.get_Sector() == r.get(ri).get_Sector() &&
+                            s.get_Region() == r.get(ri).get_Region() &&
+                            s.associatedCrossId == r.get(ri).associatedCrossId &&
+                            r.get(ri).associatedCrossId != -1) {
+                        if (s.get_Superlayer() % 2 == missingSL % 2)
+                            Segs2RoadConv.add(s); 
                     }
                 }
-                trk.calcTrajectory(trk.getId(), dcSwim, trk.get_Vtx0(), trk.get_pAtOrig(), trk.get_Q());
-                trkId++;
+            }
+            if (Segs2RoadConv.size() == 2) {
+                Segment pSegmentConv = rf.findRoadMissingSegment(Segs2RoadConv,
+                        Constants.getInstance().dcDetector,
+                        r.a);
+                if (pSegmentConv != null)
+                    psegmentsConv.add(pSegmentConv);
             }
         }
+
+        segmentsConv.addAll(psegmentsConv);
+        List<Cross> pcrossesConv = crossMake.find_Crosses(segmentsConv, Constants.getInstance().dcDetector);
+        CrossList pcrosslistConv = crossLister.candCrossLists(event, pcrossesConv,
+                false,
+                null,
+                Constants.getInstance().dcDetector,
+                null,
+                dcSwim, true);
+        List<Track> mistrkcandsConv = trkcandFinder.getTrackCands(pcrosslistConv,
+                Constants.getInstance().dcDetector,
+                Swimmer.getTorScale(),
+                dcSwim, false);
 
-                // no candidate found, stop here and save the hits,
+        //8) Select overlapping tracks from all track candidates with 5 or 6 clusters, and update hits in tracks 
+        trkcandsConv.addAll(mistrkcandsConv);
+        if (!trkcandsConv.isEmpty()) {
+            // remove overlaps
+            trkcandFinder.removeOverlappingTracks(trkcandsConv);
+            for (Track trk : trkcandsConv) {
+                // reset the id
+                trk.set_Id(trkId);
+                trkcandFinder.matchHits(trk.getStateVecs(),
+                        trk,
+                        Constants.getInstance().dcDetector,
+                        dcSwim);
+                trkId++;
+            }
+        }
+                
+        //////gather all the hits for pointer bank creation
+        trkcands.addAll(trkcandsConv);
+        trkId=1;
+        for (Track trk : trkcands) {            
+            trk.calcTrajectory(trk.getId(), dcSwim, trk.get_Vtx0(), trk.get_pAtOrig(), trk.get_Q());
+            for (Cross c : trk) {
+                c.set_CrossDirIntersSegWires();
+                trkcandFinder.setHitDoubletsInfo(c.get_Segment1());
+                trkcandFinder.setHitDoubletsInfo(c.get_Segment2());
+                for (FittedHit h1 : c.get_Segment1()) {          
+                    h1.set_AssociatedHBTrackID(trkId);
+                    //if(h1.get_AssociatedHBTrackID()>0) 
+                    fhits.add(h1); 
+                }
+                for (FittedHit h2 : c.get_Segment2()) {
+                    h2.set_AssociatedHBTrackID(trkId);
+                    //if(h2.get_AssociatedHBTrackID()>0) 
+                    fhits.add(h2); 
+                }
+            }
+            trkId++;
+        }
+           
+        // no candidate found, stop here and save the hits,
         // the clusters, the segments, the crosses
         if (trkcands.isEmpty()) {
             event.appendBanks(

reconstruction/dc/src/main/java/org/jlab/service/dc/DCTBEngine.java

@@ -189,6 +189,11 @@ public boolean processDataEvent(DataEvent event) {
                     trkbank.getFloat("tx", i), trkbank.getFloat("ty", i));
             HBFinalSV.setZ(trkbank.getFloat("z", i));
             HBtrk.setFinalStateVec(HBFinalSV);
+            
+            int status = trkbank.getShort("status", i);
+            int isAITrack = (status >> 12) & 1;
+            HBtrk.setIsAITrack((isAITrack == 1));
+                        
             TrackArray[HBtrk.get_Id()-1] = HBtrk; 
 //            TrackArray[HBtrk.get_Id()-1].set_Status(0);
         }