From 227c4819b04873fbd5aca5f8b3cb1181eefd62f8 Mon Sep 17 00:00:00 2001 From: tongtongcao Date: Mon, 2 Nov 2026 21:28:55 -0500 Subject: [PATCH 1/2] optimize loading and application of AI model for foward HB tracking --- .../dcHBTrackState/HBTrackStateEstimator.java | 116 ++++++++++++++---- .../rec/dc/track/TrackCandListFinder.java | 6 +- .../java/org/jlab/service/dc/DCEngine.java | 26 ++-- .../org/jlab/service/dc/DCHBTrackingAI.java | 26 ++-- 4 files changed, 123 insertions(+), 51 deletions(-) diff --git a/reconstruction/ai/src/main/java/org/jlab/rec/ai/dcHBTrackState/HBTrackStateEstimator.java b/reconstruction/ai/src/main/java/org/jlab/rec/ai/dcHBTrackState/HBTrackStateEstimator.java index f0818b9c3b..a09115fb8c 100644 --- a/reconstruction/ai/src/main/java/org/jlab/rec/ai/dcHBTrackState/HBTrackStateEstimator.java +++ b/reconstruction/ai/src/main/java/org/jlab/rec/ai/dcHBTrackState/HBTrackStateEstimator.java @@ -14,6 +14,8 @@ import java.util.concurrent.*; import java.util.logging.Level; import java.util.logging.Logger; +import java.util.regex.Matcher; +import java.util.regex.Pattern; import org.jlab.clas.reco.ReconstructionEngine; import org.jlab.io.base.*; @@ -38,31 +40,22 @@ public class HBTrackStateEstimator{ private float[] STATE_MEAN; private float[] STATE_STD; - public HBTrackStateEstimator(String modelFile){ + public HBTrackStateEstimator(String modelFile, String hitStatsFile, String stateStatsFile){ 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 { + System.setProperty("ai.djl.pytorch.graph_optimizer", "false"); + + try { + String hitStatsFilePath = ClasUtilsFile.getResourceDir("CLAS12DIR", "etc/data/nnet/hbTSE/" + hitStatsFile); + String stateStatsFilePath = ClasUtilsFile.getResourceDir("CLAS12DIR", "etc/data/nnet/hbTSE/" + stateStatsFile); + + loadHitStats(hitStatsFilePath); + loadStateStats(stateStatsFilePath); + String modelPath = ClasUtilsFile.getResourceDir("CLAS12DIR", "etc/data/nnet/hbTSE/" + modelFile); - + Criteria criteria = Criteria.builder() .setTypes(float[][].class, float[].class) .optModelPath(Paths.get(modelPath)) @@ -76,10 +69,87 @@ else if(modelFile.contains("outbending")){ int threads = 64; predictors = new PredictorPool(threads, model); - } catch (IOException | ModelException e) { Logger.getLogger(getClass().getName()).log(Level.SEVERE, null, e); } + } + + private void loadHitStats(String filePath) throws IOException { + String json = java.nio.file.Files.readString(Paths.get(filePath)); + + String[] keys = {"doca", "xm", "xr", "yr", "z"}; + + HIT_MEAN = new float[5]; + HIT_STD = new float[5]; + + for (int i = 0; i < keys.length; i++) { + HIT_MEAN[i] = getJsonValue(json, keys[i] + "_mean"); + HIT_STD[i] = getJsonValue(json, keys[i] + "_std"); + + if (!Float.isFinite(HIT_MEAN[i]) || !Float.isFinite(HIT_STD[i]) || HIT_STD[i] <= 0) { + throw new IOException("Invalid hit statistics: " + keys[i]); + } + } + } + + private float getJsonValue(String json, String key) throws IOException { + Pattern pattern = Pattern.compile("\"" + Pattern.quote(key) + "\"\\s*:\\s*(-?\\d+(?:\\.\\d+)?(?:[eE][+-]?\\d+)?)"); + + Matcher matcher = pattern.matcher(json); + + if (!matcher.find()) { + throw new IOException("Missing JSON key: " + key); + } + + try { + return Float.parseFloat(matcher.group(1)); + } catch (NumberFormatException e) { + throw new IOException("Invalid JSON value: " + key, e); + } + } + + private void loadStateStats(String filePath) throws IOException { + String json = java.nio.file.Files.readString(Paths.get(filePath)); + + String[] keys = {"x", "y", "tx", "ty", "Q"}; + + STATE_MEAN = new float[5]; + STATE_STD = new float[5]; + + for (int i = 0; i < keys.length; i++) { + float[] values = getJsonArray(json, keys[i]); + + if (values.length != 2 || !Float.isFinite(values[0]) || !Float.isFinite(values[1]) || values[1] <= 0) { + throw new IOException("Invalid state statistics: " + keys[i]); + } + + STATE_MEAN[i] = values[0]; + STATE_STD[i] = values[1]; + } + } + + private float[] getJsonArray(String json, String key) throws IOException{ + Pattern pattern = Pattern.compile("\"" + Pattern.quote(key) + "\"\\s*:\\s*\\[([^\\]]+)\\]"); + + Matcher matcher = pattern.matcher(json); + + if (!matcher.find()) { + throw new IOException("Missing JSON key: " + key); + } + + String[] values = matcher.group(1).split(","); + + float[] result = new float[values.length]; + + try { + for (int i = 0; i < values.length; i++) { + result[i] = Float.parseFloat(values[i].trim()); + } + } catch (NumberFormatException e) { + throw new IOException("Invalid JSON values for key: " + key, e); + } + + return result; } // ---------------- Translator ---------------- @@ -145,6 +215,6 @@ public float[] predict(float[][] hits) { } catch (TranslateException | InterruptedException e) { throw new RuntimeException(e); } - } + } } diff --git a/reconstruction/dc/src/main/java/org/jlab/rec/dc/track/TrackCandListFinder.java b/reconstruction/dc/src/main/java/org/jlab/rec/dc/track/TrackCandListFinder.java index 952d3d80ea..ca447ec341 100644 --- a/reconstruction/dc/src/main/java/org/jlab/rec/dc/track/TrackCandListFinder.java +++ b/reconstruction/dc/src/main/java/org/jlab/rec/dc/track/TrackCandListFinder.java @@ -207,7 +207,7 @@ private int calcInitTrkQ(double a, double TORSCALE) { return q; } - public List getTrackCandsAI(CrossList crossList, DCGeant4Factory DcDetector, Swim dcSwim, HBTrackStateEstimator hbTSEstimator){ + public List getTrackCandsAI(CrossList crossList, DCGeant4Factory DcDetector, Swim dcSwim, HBTrackStateEstimator hbTSEstimator, boolean isInbendingHBTSEModel){ List cands = new ArrayList<>(); for (List aCrossList : crossList) { @@ -254,10 +254,10 @@ public List getTrackCandsAI(CrossList crossList, DCGeant4Factory DcDetect initSV.z = 229.; // State vector at z = 229 for AI training samples initSV.tx = estSV[2]; initSV.ty = estSV[3]; - initSV.Q = estSV[4]; + initSV.Q = estSV[4] * Math.signum(Swimmer.getTorScale()) * (isInbendingHBTSEModel ? -1 : 1); double absTorScale = Math.abs(Swimmer.getTorScale()); - initSV.Q = absTorScale>0.001 ? estSV[4]/absTorScale : estSV[4]; // Momentum magnitude is scaled by the scale factor of magnetic field + initSV.Q = absTorScale>0.001 ? initSV.Q/absTorScale : initSV.Q; // Momentum magnitude is scaled by the scale factor of magnetic field RungeKuttaDoca rk = new RungeKuttaDoca(); rk.SwimToZ(sector, initSV, dcSwim, measSurfaces.get(0).wireLine[0].end().z(), new float[3]); diff --git a/reconstruction/dc/src/main/java/org/jlab/service/dc/DCEngine.java b/reconstruction/dc/src/main/java/org/jlab/service/dc/DCEngine.java index 9e00490f75..359142de6f 100644 --- a/reconstruction/dc/src/main/java/org/jlab/service/dc/DCEngine.java +++ b/reconstruction/dc/src/main/java/org/jlab/service/dc/DCEngine.java @@ -39,8 +39,10 @@ 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 + protected String hbTSEModelFile = "transformer_32d_4h_3l_inbending.pt"; // AI model file for HB track state estimator + protected String hitStatsHBTSEModel = "hit_stats_inbending.json"; + protected String stateStatsHBTSEModel = "state_stats_inbending.json"; + protected boolean isInbendingHBTSEModel = true; // If the model was trained by inbending run public static final Logger LOGGER = Logger.getLogger(ReconstructionEngine.class.getName()); @@ -130,13 +132,21 @@ else if(this.getEngineConfigString("dcT2DFunc").equalsIgnoreCase("Polynomial")) KFitter.setDafAnnealingFactorsTB(dafAnnealingFactorsTB); } - if (getEngineConfigString("hbTSEModelFileInbending") != null){ - hbTSEModelFileInbending = getEngineConfigString("hbTSEModelFileInbending"); - } + if (getEngineConfigString("hbTSEModelFile") != null){ + hbTSEModelFile = getEngineConfigString("hbTSEModelFile"); + } + + if (getEngineConfigString("hitStatsHBTSEModel") != null){ + hitStatsHBTSEModel = getEngineConfigString("hitStatsHBTSEModel"); + } - if (getEngineConfigString("hbTSEModelFileOutbending") != null){ - hbTSEModelFileOutbending = getEngineConfigString("hbTSEModelFileOutbending"); - } + if (getEngineConfigString("stateStatsHBTSEModel") != null){ + stateStatsHBTSEModel = getEngineConfigString("stateStatsHBTSEModel"); + } + + if (getEngineConfigString("isInbendingHBTSEModel") != null){ + isInbendingHBTSEModel = Boolean.valueOf(getEngineConfigString("isInbendingHBTSEModel")); + } // Set geometry shifts for alignment code if(this.getEngineConfigString("alignmentShifts")!=null) { diff --git a/reconstruction/dc/src/main/java/org/jlab/service/dc/DCHBTrackingAI.java b/reconstruction/dc/src/main/java/org/jlab/service/dc/DCHBTrackingAI.java index 38f4db4dfe..0f51515f0d 100644 --- a/reconstruction/dc/src/main/java/org/jlab/service/dc/DCHBTrackingAI.java +++ b/reconstruction/dc/src/main/java/org/jlab/service/dc/DCHBTrackingAI.java @@ -36,21 +36,18 @@ * @author tongtong */ public class DCHBTrackingAI extends DCEngine { - private HBTrackStateEstimator hbTSEstimatorInbending = null; - private HBTrackStateEstimator hbTSEstimatorOutbending = null; + private HBTrackStateEstimator hbTSEstimator = null; public DCHBTrackingAI() { super("DCHTAI"); this.getBanks().init("HitBasedTrkg", "", "AI"); - hbTSEstimatorInbending = new HBTrackStateEstimator(hbTSEModelFileInbending); - hbTSEstimatorOutbending = new HBTrackStateEstimator(hbTSEModelFileOutbending); + hbTSEstimator = new HBTrackStateEstimator(hbTSEModelFile, hitStatsHBTSEModel, stateStatsHBTSEModel); } public DCHBTrackingAI(String outputBankPrefix) { super("DCHTAI"); this.getBanks().init("HitBasedTrkg", "", outputBankPrefix); - hbTSEstimatorInbending = new HBTrackStateEstimator(hbTSEModelFileInbending); - hbTSEstimatorOutbending = new HBTrackStateEstimator(hbTSEModelFileOutbending); + hbTSEstimator = new HBTrackStateEstimator(hbTSEModelFile, hitStatsHBTSEModel, stateStatsHBTSEModel); } @Override @@ -74,14 +71,9 @@ public boolean processDataEventUser(DataEvent event) { LOGGER.log(Level.INFO, "RUN=0: Skipping event"); return true; } - - HBTrackStateEstimator hbTSEstimator = null; - if(Swimmer.getTorScale() < 0) - hbTSEstimator = hbTSEstimatorInbending; - else if(Swimmer.getTorScale() > 0) - hbTSEstimator = hbTSEstimatorOutbending; - else{ - Logger.getLogger(getClass().getName()).log(Level.SEVERE, "Torus scale = 0"); + + if(hbTSEstimator == null){ + Logger.getLogger(getClass().getName()).log(Level.SEVERE, "Model is not found!"); return true; } @@ -147,7 +139,7 @@ else if(Swimmer.getTorScale() > 0) TrackCandListFinder trkcandFinder = new TrackCandListFinder(Constants.HITBASE); trkcands = trkcandFinder.getTrackCandsAI(crosslist, Constants.getInstance().dcDetector, - dcSwim, hbTSEstimator); + dcSwim, hbTSEstimator, isInbendingHBTSEModel); // track found int trkId = 1; @@ -241,7 +233,7 @@ else if(Swimmer.getTorScale() > 0) // track candidates with 6 clusters trkcandsConv = trkcandFinder.getTrackCandsAI(crosslistConv, Constants.getInstance().dcDetector, - dcSwim, hbTSEstimator); + dcSwim, hbTSEstimator, isInbendingHBTSEModel); // track candidates with 5 clusters RoadFinder rf = new RoadFinder(); @@ -299,7 +291,7 @@ else if(Swimmer.getTorScale() > 0) dcSwim, true); List mistrkcandsConv = trkcandFinder.getTrackCandsAI(pcrosslistConv, Constants.getInstance().dcDetector, - dcSwim, hbTSEstimator); + dcSwim, hbTSEstimator, isInbendingHBTSEModel); //8) Select overlapping tracks from all track candidates with 5 or 6 clusters, and update hits in tracks trkcandsConv.addAll(mistrkcandsConv); From 1b3d860df23c42a8e0a99accacdb58ff56dd25bc Mon Sep 17 00:00:00 2001 From: tongtongcao Date: Tue, 3 Nov 2026 14:57:54 -0500 Subject: [PATCH 2/2] define estimator in init() instead of constructor --- .../java/org/jlab/service/dc/DCHBTrackingAI.java | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/reconstruction/dc/src/main/java/org/jlab/service/dc/DCHBTrackingAI.java b/reconstruction/dc/src/main/java/org/jlab/service/dc/DCHBTrackingAI.java index 0f51515f0d..98f9231800 100644 --- a/reconstruction/dc/src/main/java/org/jlab/service/dc/DCHBTrackingAI.java +++ b/reconstruction/dc/src/main/java/org/jlab/service/dc/DCHBTrackingAI.java @@ -41,13 +41,20 @@ public class DCHBTrackingAI extends DCEngine { public DCHBTrackingAI() { super("DCHTAI"); this.getBanks().init("HitBasedTrkg", "", "AI"); - hbTSEstimator = new HBTrackStateEstimator(hbTSEModelFile, hitStatsHBTSEModel, stateStatsHBTSEModel); } public DCHBTrackingAI(String outputBankPrefix) { super("DCHTAI"); this.getBanks().init("HitBasedTrkg", "", outputBankPrefix); - hbTSEstimator = new HBTrackStateEstimator(hbTSEModelFile, hitStatsHBTSEModel, stateStatsHBTSEModel); + + } + + @Override + public boolean init() { + super.init(); + hbTSEstimator = new HBTrackStateEstimator(hbTSEModelFile, hitStatsHBTSEModel, stateStatsHBTSEModel); + + return true; } @Override @@ -71,7 +78,7 @@ public boolean processDataEventUser(DataEvent event) { LOGGER.log(Level.INFO, "RUN=0: Skipping event"); return true; } - + if(hbTSEstimator == null){ Logger.getLogger(getClass().getName()).log(Level.SEVERE, "Model is not found!"); return true;