apply AI model for HB tracking

common-tools/clas-tracking/src/main/java/org/jlab/clas/tracking/kalmanfilter/zReference/KFitter.java

@@ -60,7 +60,7 @@ public class KFitter extends AKFitter {
     Matrix result = new Matrix();
     Matrix result_inv = new Matrix();
     Matrix adj = new Matrix();
-
+    
     public KFitter(boolean filter, int iterations, int dir, Swim swim, double Z[], Libr mo) {
         super(filter, iterations, dir, swim, mo);
         this.Z = Z;
@@ -861,7 +861,7 @@ private void calcFinalChisq(int sector) {
 
     // Since no vertex inforamtion, the starting point for path length is the final point at the last layer.
     // After vertex information is obtained, transition for the starting point from the final point to vertex will be taken.
-    private void calcFinalChisq(int sector, boolean nofilter) {
+    public void calcFinalChisq(int sector, boolean nofilter) {
         int k = svzLength - 1;
         this.chi2 = 0;
         double path = 0;
@@ -1157,6 +1157,14 @@ public StateVecs getStateVecs() {
     public double getNDFDAF(){
         return ndfDAF;
     }
+
+    public void setSvzLength(int svzlength){
+        this.svzLength = svzlength;
+    }
+
+    public int getSvzLength(){
+        return svzLength;
+    }
 
     public void printlnMeasVecs() {
         for (int i = 0; i < mv.measurements.size(); i++) {

reconstruction/ai/src/main/java/org/jlab/rec/ai/dcHBTrackState/HBTrackStateEstimator.java

@@ -0,0 +1,150 @@
+package org.jlab.rec.ai.dcHBTrackState;
+
+import ai.djl.MalformedModelException;
+import ai.djl.ModelException;
+import ai.djl.inference.Predictor;
+import ai.djl.ndarray.*;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.repository.zoo.*;
+import ai.djl.training.util.ProgressBar;
+import ai.djl.translate.*;
+
+import java.io.IOException;
+import java.nio.file.Paths;
+import java.util.concurrent.*;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+import org.jlab.clas.reco.ReconstructionEngine;
+import org.jlab.io.base.*;
+import org.jlab.utils.system.ClasUtilsFile;
+import org.jlab.service.ai.PredictorPool;
+
+
+public class HBTrackStateEstimator{
+    // ---------------- Configuration ----------------
+    private String modelFile;
+
+    ZooModel<float[][], float[]> model;
+    PredictorPool predictors;
+
+    // ---------------- Statistics for normalization of inputs and outputs of training samples ----------------
+    //// Note: Statistics of hits and track states depends on training samples, so need to be renewed when training samples change!!!
+    // Statistics of hits: doca, xm, xr, yr, z
+    private float[] HIT_MEAN;
+    private float[] HIT_STD;
+
+    // Statistics of track state: x, y, tx, ty, Q at z = 229 cm in the tilted sector frame
+    private float[] STATE_MEAN;
+    private float[] STATE_STD;
+
+    public HBTrackStateEstimator(String modelFile){
+        this.modelFile = modelFile;
+
+        if(modelFile.contains("inbending")){
+            HIT_MEAN = new float[]{0.52949071f, -45.771999f,  -45.744694f,  57.336819f, 373.046356f};
+            HIT_STD = new float[]{0.40272677f, 47.928203f, 48.379021f, 32.645191f, 111.54994f};
+            STATE_MEAN = new float[]{-33.564308f, 0.010787425f, -0.15567796f, 0.0017755219f, 0.317530721f};
+            STATE_STD = new float[]{28.667490f, 17.761129f, 0.11940812f, 0.074460238f, 0.74185127f};        
+        }
+        else if(modelFile.contains("outbending")){
+            HIT_MEAN = new float[]{0.53385729f, -59.236504f,  -59.200584f,  50.136387f, 372.057922f};
+            HIT_STD = new float[]{0.40085429f, 51.385536f, 51.840462f, 31.498201f, 111.50029f};
+            STATE_MEAN = new float[]{-39.446106f, 0.17583229f, -0.18047817f, 0.0014163271f, -0.082320645f};
+            STATE_STD = new float[]{33.733425f, 17.226780f, 0.14071095f, 0.072449364f, 0.72273886f};  
+        }
+        else{
+           Logger.getLogger(getClass().getName()).log(Level.SEVERE, "Name of model file does not include inbending or outbending"); 
+        }
+
+        System.setProperty("ai.djl.pytorch.num_interop_threads", "1");
+        System.setProperty("ai.djl.pytorch.num_threads", "1");
+        System.setProperty("ai.djl.pytorch.graph_optimizer", "false");                
+        try {
+            String modelPath = ClasUtilsFile.getResourceDir("CLAS12DIR", "etc/data/nnet/hbTSE/" + modelFile);
+
+            Criteria<float[][], float[]> criteria = Criteria.builder()
+                .setTypes(float[][].class, float[].class)
+                .optModelPath(Paths.get(modelPath))
+                .optEngine("PyTorch")
+                .optTranslator(getTranslator())
+                .optProgress(new ProgressBar())
+                .build();
+
+            model = criteria.loadModel();
+
+            int threads = 64;
+            predictors = new PredictorPool(threads, model);
+
+
+        } catch (IOException | ModelException e) {
+            Logger.getLogger(getClass().getName()).log(Level.SEVERE, null, e);
+        }                
+    }    
+
+    // ---------------- Translator ----------------
+    private Translator<float[][], float[]> getTranslator() {
+        return new Translator<float[][], float[]>() {
+
+            @Override
+            public NDList processInput(TranslatorContext ctx, float[][] hits) {
+                NDManager manager = ctx.getNDManager();
+                int n = hits.length;
+
+                float[][] norm = new float[n][5];
+                for (int i = 0; i < n; i++)
+                    for (int j = 0; j < 5; j++)
+                        norm[i][j] = (hits[i][j] - HIT_MEAN[j]) / HIT_STD[j];
+
+                NDArray x = manager.create(norm);    
+                x = x.reshape(1, n, 5); 
+                return new NDList(x);
+            }
+
+            @Override
+            public float[] processOutput(TranslatorContext ctx, NDList list) {
+                NDArray out = list.get(0); // [1,5]
+                float[] y = out.toFloatArray();
+
+                for (int i = 0; i < 5; i++)
+                    y[i] = y[i] * STATE_STD[i] + STATE_MEAN[i];
+
+                return y;
+            }
+
+            @Override
+            public Batchifier getBatchifier() {
+                return null;
+            }
+        };
+    }
+
+
+    public float[] predict(float[][] hits) {
+        if (hits == null) return null;
+
+        if (hits.length == 0) {
+            throw new IllegalArgumentException("HBInitialStateEstimator: empty hits");
+        }
+
+        for (int i = 0; i < hits.length; i++) {
+            if (hits[i].length != 5) {
+                throw new IllegalArgumentException(
+                    "Expect 5 features per hit, got " + hits[i].length
+                );
+            }
+        }        
+
+        try {
+            Predictor<float[][], float[]> predictor = predictors.take();
+            try {
+                return predictor.predict(hits);
+            } finally {
+                predictors.put(predictor);
+            }
+        } catch (TranslateException | InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+}
+

reconstruction/dc/pom.xml

@@ -41,6 +41,11 @@
       <artifactId>clas-tracking</artifactId>
       <version>13.5.2-SNAPSHOT</version>
     </dependency>
+    <dependency>
+      <groupId>org.jlab.clas12.detector</groupId>
+      <artifactId>clas12detector-ai</artifactId>
+      <version>13.5.2-SNAPSHOT</version>
+    </dependency>    
     <dependency>
       <groupId>org.jlab.clas</groupId>
       <artifactId>clas-jcsg</artifactId>

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

@@ -34,6 +34,8 @@
 import org.jlab.clas.tracking.utilities.MatrixOps.Libr;
 import org.jlab.clas.tracking.utilities.RungeKuttaDoca;
 
+import org.jlab.rec.ai.dcHBTrackState.HBTrackStateEstimator;
+
 /**
  * A class with a method implementing an algorithm that finds lists of track
  * candidates in the DC
@@ -204,6 +206,106 @@ private int calcInitTrkQ(double a, double TORSCALE) {
 
         return q;
     }
+
+    public List<Track> getTrackCandsAI(CrossList crossList, DCGeant4Factory DcDetector, Swim dcSwim, HBTrackStateEstimator hbTSEstimator){
+        List<Track> cands = new ArrayList();
+
+        for (List<Cross> aCrossList : crossList) {
+            Track cand = new Track();
+
+            if (aCrossList.size() == 3 && this.PassNSuperlayerTracking(aCrossList, cand)) {
+                cand.addAll(aCrossList);
+                cand.setSector(aCrossList.get(0).get_Sector());
+                int sector = cand.getSector();
+
+                List<Surface> measSurfaces = getMeasSurfaces(cand, DcDetector);
+
+                int numHits = 0;
+                for(Surface surf : measSurfaces){
+                    numHits += surf.nMeas;
+                }
+                float[][] hits = new float[numHits][5];
+                int indexHit = 0;
+                for(int i = 0; i < measSurfaces.size(); i++){
+                    hits[indexHit][0] = (float)measSurfaces.get(i).doca[0];
+                    hits[indexHit][1] = (float)measSurfaces.get(i).wireLine[0].origin().x();
+                    hits[indexHit][2] = (float)measSurfaces.get(i).wireLine[0].end().x();
+                    hits[indexHit][3] = (float)measSurfaces.get(i).wireLine[0].end().y();
+                    hits[indexHit][4] = (float)measSurfaces.get(i).wireLine[0].end().z();
+
+                    indexHit++;
+
+                    if(measSurfaces.get(i).nMeas == 2){                
+                        hits[indexHit][0] = (float)measSurfaces.get(i).doca[1];
+                        hits[indexHit][1] = (float)measSurfaces.get(i).wireLine[1].origin().x();
+                        hits[indexHit][2] = (float)measSurfaces.get(i).wireLine[1].end().x();
+                        hits[indexHit][3] = (float)measSurfaces.get(i).wireLine[1].end().y();
+                        hits[indexHit][4] = (float)measSurfaces.get(i).wireLine[1].end().z();
+
+                        indexHit++;
+                    }                                
+                } 
+
+                float[] estSV = hbTSEstimator.predict(hits);
+                StateVecs svs = new StateVecs();
+                org.jlab.clas.tracking.kalmanfilter.AStateVecs.StateVec initSV = svs.new StateVec(0);
+                initSV.x = estSV[0];
+                initSV.y = estSV[1];
+                initSV.z = 229.; // State vector at z = 229 for AI training samples
+                initSV.tx = estSV[2];
+                initSV.ty = estSV[3];
+                initSV.Q = estSV[4];
+
+                RungeKuttaDoca rk = new RungeKuttaDoca();
+                rk.SwimToZ(sector, initSV, dcSwim, measSurfaces.get(0).wireLine[0].end().z(), new float[3]); 
+
+                KFitter kFZRef = new KFitter(true, 1, 1, dcSwim, Constants.getInstance().Z, Libr.JNP);
+                Matrix initCMatrix = new Matrix();                       
+                initSV.CM = new Matrix();                 
+                kFZRef.init(measSurfaces, initSV);						
+
+                org.jlab.clas.tracking.kalmanfilter.AStateVecs.StateVec finalSV = svs.new StateVec(initSV);
+                rk.SwimToZ(sector, finalSV, dcSwim, measSurfaces.get(measSurfaces.size()-1).wireLine[0].end().z(), new float[3]);
+                kFZRef.getStateVecs().transported(true).put(measSurfaces.size()-1, finalSV);
+
+                kFZRef.setSvzLength(measSurfaces.size());
+                kFZRef.calcFinalChisq(sector, true);
+
+                StateVec stateVec = new StateVec(finalSV.x,
+                                finalSV.y, finalSV.tx, finalSV.ty);
+                int q = (int) Math.signum(finalSV.Q);
+                double p = 1. / Math.abs(finalSV.Q);
+                stateVec.setZ(finalSV.z);
+
+                //set the track parameters 
+                cand.set_P(p);
+                cand.set_Q(q);
+
+                // candidate parameters 
+                cand.set_FitChi2(kFZRef.chi2);
+                cand.set_FitNDF(kFZRef.NDF);
+
+                cand.setFinalStateVec(stateVec);
+                cand.set_Id(cands.size() + 1);
+                this.setTrackPars(cand, null,
+                        null, stateVec,
+                        stateVec.getZ(),
+                        DcDetector, dcSwim);
+
+                if(cand.get_Vtx0() != null){
+                    Point3D VTCS = cand.get(cand.size()-1).getCoordsInTiltedSector(cand.get_Vtx0().x(), cand.get_Vtx0().y(), cand.get_Vtx0().z());
+                    double deltaPathToVtx =  kFZRef.getDeltaPathToVtx(sector, VTCS.z());
+
+                    List<org.jlab.rec.dc.trajectory.StateVec> kfStateVecsAlongTrajectory = setKFStateVecsAlongTrajectory(kFZRef, deltaPathToVtx);
+                    cand.setStateVecs(kfStateVecsAlongTrajectory);                                  
+                }
+
+                if (kFZRef.chi2 < Constants.MAXCHI2) cands.add(cand);                                              
+            }
+        }
+
+        return cands;
+    }
 
     /**
      * @param crossList the input list of crosses

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

@@ -12,6 +12,7 @@
 import org.jlab.rec.dc.banks.Banks;
 import org.jlab.clas.tracking.kalmanfilter.zReference.KFitter;
 import org.jlab.clas.tracking.kalmanfilter.zReference.DAFilter;
+import org.jlab.rec.ai.dcHBTrackState.HBTrackStateEstimator;
 
 public class DCEngine extends ReconstructionEngine {
 
@@ -39,6 +40,9 @@ public class DCEngine extends ReconstructionEngine {
     private String   dafChi2Cut     = null;
     private String   dafAnnealingFactorsTB = null;
 
+    protected String hbTSEModelFileInbending = "transformer_32d_4h_3l_inbending.pt"; // AI model file for HB track state estimator for inbending runs  
+    protected String hbTSEModelFileOutbending = "transformer_32d_4h_3l_outbending.pt"; // AI model file for HB track state estimator for outbending runs  
+
     public static final Logger LOGGER = Logger.getLogger(ReconstructionEngine.class.getName());
 
 
@@ -127,6 +131,14 @@ else if(this.getEngineConfigString("dcT2DFunc").equalsIgnoreCase("Polynomial"))
             dafAnnealingFactorsTB=this.getEngineConfigString("dafAnnealingFactorsTB");
             KFitter.setDafAnnealingFactorsTB(dafAnnealingFactorsTB);
         }
+
+        if (getEngineConfigString("hbTSEModelFileInbending") != null){
+            hbTSEModelFileInbending = getEngineConfigString("hbTSEModelFileInbending");
+        }
+
+        if (getEngineConfigString("hbTSEModelFileOutbending") != null){
+            hbTSEModelFileOutbending = getEngineConfigString("hbTSEModelFileOutbending");
+        }
 
         // Set geometry shifts for alignment code
         if(this.getEngineConfigString("alignmentShifts")!=null) {