Skip to content

optimize loading and application of AI model for foward HB tracking #1563

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tongtongcao wants to merge 2 commits into
base: development
Choose a base branch
from
Open

optimize loading and application of AI model for foward HB tracking #1563

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
227c481
optimize loading and application of AI model for foward HB tracking
tongtongcao Nov 3, 2026
1b3d860
define estimator in init() instead of constructor
tongtongcao Nov 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
116 changes: 93 additions & 23 deletions reconstruction/ai/src/main/java/org/jlab/rec/ai/dcHBTrackState/HBTrackStateEstimator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.*;
Expand All @@ -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<float[][], float[]> criteria = Criteria.builder()
.setTypes(float[][].class, float[].class)
.optModelPath(Paths.get(modelPath))
Expand All @@ -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 ----------------
Expand Down Expand Up @@ -145,6 +215,6 @@ public float[] predict(float[][] hits) {
} catch (TranslateException | InterruptedException e) {
throw new RuntimeException(e);
}
}
}
}

6 changes: 3 additions & 3 deletions reconstruction/dc/src/main/java/org/jlab/rec/dc/track/TrackCandListFinder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,7 +207,7 @@ private int calcInitTrkQ(double a, double TORSCALE) {
return q;
}

public List<Track> getTrackCandsAI(CrossList crossList, DCGeant4Factory DcDetector, Swim dcSwim, HBTrackStateEstimator hbTSEstimator){
public List<Track> getTrackCandsAI(CrossList crossList, DCGeant4Factory DcDetector, Swim dcSwim, HBTrackStateEstimator hbTSEstimator, boolean isInbendingHBTSEModel){
List<Track> cands = new ArrayList<>();

for (List<Cross> aCrossList : crossList) {
Expand Down Expand Up @@ -254,10 +254,10 @@ public List<Track> 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]);
Expand Down
26 changes: 18 additions & 8 deletions reconstruction/dc/src/main/java/org/jlab/service/dc/DCEngine.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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());

Expand Down Expand Up @@ -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) {
Expand Down
31 changes: 15 additions & 16 deletions reconstruction/dc/src/main/java/org/jlab/service/dc/DCHBTrackingAI.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,21 +36,25 @@
* @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);
}

public DCHBTrackingAI(String outputBankPrefix) {
super("DCHTAI");
this.getBanks().init("HitBasedTrkg", "", outputBankPrefix);
hbTSEstimatorInbending = new HBTrackStateEstimator(hbTSEModelFileInbending);
hbTSEstimatorOutbending = new HBTrackStateEstimator(hbTSEModelFileOutbending);

}

@Override
public boolean init() {
super.init();
hbTSEstimator = new HBTrackStateEstimator(hbTSEModelFile, hitStatsHBTSEModel, stateStatsHBTSEModel);

return true;
}

@Override
Expand All @@ -75,13 +79,8 @@ public boolean processDataEventUser(DataEvent 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;
}

Expand Down Expand Up @@ -147,7 +146,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;
Expand Down Expand Up @@ -241,7 +240,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();
Expand Down Expand Up @@ -299,7 +298,7 @@ else if(Swimmer.getTorScale() > 0)
dcSwim, true);
List<Track> 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);
Expand Down
Loading