From ac52fe23c164a6a30d7e48fee82d21f4393be175 Mon Sep 17 00:00:00 2001 From: Uditha Weerasinghe Date: Thu, 30 Jul 2026 08:36:21 -0400 Subject: [PATCH 1/5] feat: added implementation for post-KF PID model --- .../jlab/rec/alert/AIPID/ModelPostPID.java | 81 +++++++++++++++++++ .../org/jlab/rec/alert/AIPID/PIDResult.java | 19 +++++ 2 files changed, 100 insertions(+) create mode 100644 reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/ModelPostPID.java create mode 100644 reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/PIDResult.java diff --git a/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/ModelPostPID.java b/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/ModelPostPID.java new file mode 100644 index 0000000000..5d8913b65f --- /dev/null +++ b/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/ModelPostPID.java @@ -0,0 +1,81 @@ +package org.jlab.rec.alert.AIPID; + +import ai.djl.MalformedModelException; +import ai.djl.inference.Predictor; +import ai.djl.ndarray.NDList; +import ai.djl.ndarray.types.Shape; +import ai.djl.repository.zoo.Criteria; +import ai.djl.repository.zoo.ModelNotFoundException; +import ai.djl.repository.zoo.ZooModel; +import ai.djl.training.util.ProgressBar; +import ai.djl.translate.TranslateException; +import ai.djl.translate.Translator; +import ai.djl.translate.TranslatorContext; +import java.io.IOException; +import java.nio.file.Paths; +import java.util.logging.Logger; +import org.jlab.utils.CLASResources; + +public class ModelPostPID { + + private static final Logger LOGGER = Logger.getLogger(ModelPostPID.class.getName()); + private static final int[] CLASS_IDS = {2212, 45, 46, 49, 47}; + private static final int INPUT_SIZE = 18; + + private final ZooModel model; + + public ModelPostPID() { + 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"); + + String path = CLASResources.getResourcePath("etc/data/nnet/rg-l/model_PID/"); + Criteria criteria = Criteria.builder() + .setTypes(float[].class, float[].class) + .optModelPath(Paths.get(path)) + .optEngine("PyTorch") + .optTranslator(translator()) + .optProgress(new ProgressBar()) + .build(); + try { + model = criteria.loadModel(); + } catch (IOException | ModelNotFoundException | MalformedModelException e) { + throw new RuntimeException(e); + } + } + + public float[] prediction(float[] features) throws TranslateException { + if (features == null || features.length != INPUT_SIZE) { + LOGGER.warning("PostPID input must be float[18]"); + return null; + } + try (Predictor predictor = model.newPredictor()) { + return predictor.predict(features); + } + } + + private static Translator translator() { + return new Translator<>() { + @Override + public NDList processInput(TranslatorContext ctx, float[] features) { + return new NDList(ctx.getNDManager().create(features, new Shape(1, INPUT_SIZE))); + } + + @Override + public float[] processOutput(TranslatorContext ctx, NDList output) { + float[] probabilities = output.get(0).toFloatArray(); + int bestIndex = 0; + for (int i = 1; i < probabilities.length; i++) { + if (probabilities[i] > probabilities[bestIndex]) { + bestIndex = i; + } + } + return new float[]{ + CLASS_IDS[bestIndex], + probabilities[0], probabilities[1], probabilities[2], + probabilities[3], probabilities[4] + }; + } + }; + } +} diff --git a/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/PIDResult.java b/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/PIDResult.java new file mode 100644 index 0000000000..5ae1730901 --- /dev/null +++ b/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/PIDResult.java @@ -0,0 +1,19 @@ +package org.jlab.rec.alert.AIPID; + +public class PIDResult { + public final int trackid; + public final int clusterid; + public final int pid; + public final float p2212, p45, p46, p47, p49; + + public PIDResult(int trackid, int clusterid, float[] prediction) { + this.trackid = trackid; + this.clusterid = clusterid; + this.pid = (int) prediction[0]; + this.p2212 = prediction[1]; + this.p45 = prediction[2]; + this.p46 = prediction[3]; + this.p49 = prediction[4]; + this.p47 = prediction[5]; + } +} From 23f0b70278623bdb78856f951399f5a04e593672 Mon Sep 17 00:00:00 2001 From: Uditha Weerasinghe Date: Thu, 30 Jul 2026 08:36:52 -0400 Subject: [PATCH 2/5] fix: updated PrePID model to give a prediction even for tracks without ATOF matches --- .../org/jlab/rec/alert/AIPID/ModelPrePID.java | 124 +++++++++--------- .../jlab/rec/alert/AIPID/PrePIDResult.java | 9 +- 2 files changed, 71 insertions(+), 62 deletions(-) diff --git a/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/ModelPrePID.java b/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/ModelPrePID.java index 4f91f4deed..431a618944 100644 --- a/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/ModelPrePID.java +++ b/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/ModelPrePID.java @@ -2,9 +2,7 @@ import ai.djl.MalformedModelException; import ai.djl.inference.Predictor; -import ai.djl.ndarray.NDArray; import ai.djl.ndarray.NDList; -import ai.djl.ndarray.NDManager; import ai.djl.ndarray.types.Shape; import ai.djl.repository.zoo.Criteria; import ai.djl.repository.zoo.ModelNotFoundException; @@ -13,92 +11,96 @@ import ai.djl.translate.TranslateException; import ai.djl.translate.Translator; import ai.djl.translate.TranslatorContext; - -import org.jlab.utils.CLASResources; - import java.io.IOException; import java.nio.file.Paths; import java.util.logging.Logger; +import org.jlab.utils.CLASResources; public class ModelPrePID { - - static final Logger LOGGER = Logger.getLogger(ModelPrePID.class.getName()); - // Must match training class order - private static final int[] CLASS_IDS = new int[]{2212, 45, 46, 47, 49}; - - private final ZooModel model; - - public ModelPrePID() { - Translator my_translator = new Translator<>() { + private static final Logger LOGGER = Logger.getLogger(ModelPrePID.class.getName()); + private static final int[] CLASS_IDS = {2212, 45, 46, 49, 47}; - @Override - public NDList processInput(TranslatorContext ctx, float[] floats) { - NDManager manager = ctx.getNDManager(); + private final ZooModel ahdcModel; + private final ZooModel atofModel; - // IMPORTANT: model expects (batch, 23). Provide (1, 23). - NDArray x = manager.create(floats, new Shape(1, 23)); - return new NDList(x); - } + public ModelPrePID() { + 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"); - @Override - public float[] processOutput(TranslatorContext ctx, NDList ndList) { - NDArray logits = ndList.get(0); // (1,5) - NDArray probs = logits.softmax(1); // (1,5) + ahdcModel = loadModel("model_prePID_AHDC", 11); + atofModel = loadModel("model_prePID_ATOF", 16); + } - float[] p = probs.toFloatArray(); // length 5 (row-major) + public ZooModel getModel() { + return ahdcModel; + } - // argmax - int bestIdx = 0; - float best = p[0]; - for (int k = 1; k < 5; k++) { - if (p[k] > best) { best = p[k]; bestIdx = k; } - } - int prepid = CLASS_IDS[bestIdx]; + public float[] prediction(float[] features) throws TranslateException { + if (features != null && features.length == 16) { + return predictionATOF(features); + } + return predictionAHDC(features); + } - // Return: prepid + probabilities in fixed class order - return new float[]{ - (float) prepid, - p[0], p[1], p[2], p[3], p[4] - }; - } - }; + public float[] predictionAHDC(float[] features) throws TranslateException { + return predict(ahdcModel, features, 11); + } - 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"); + public float[] predictionATOF(float[] features) throws TranslateException { + return predict(atofModel, features, 16); + } - String path = CLASResources.getResourcePath("etc/data/nnet/rg-l/model_PrePID/"); + private static float[] predict(ZooModel model, float[] features, + int expectedSize) throws TranslateException { + if (features == null || features.length != expectedSize) { + LOGGER.warning("PrePID input must be float[" + expectedSize + "]"); + return null; + } + try (Predictor predictor = model.newPredictor()) { + return predictor.predict(features); + } + } + private static ZooModel loadModel(String directory, int inputSize) { + String path = CLASResources.getResourcePath("etc/data/nnet/rg-l/" + directory + "/"); Criteria criteria = Criteria.builder() .setTypes(float[].class, float[].class) .optModelPath(Paths.get(path)) .optEngine("PyTorch") - .optTranslator(my_translator) + .optTranslator(translator(inputSize)) .optProgress(new ProgressBar()) .build(); - try { - model = criteria.loadModel(); + return criteria.loadModel(); } catch (IOException | ModelNotFoundException | MalformedModelException e) { throw new RuntimeException(e); } } - public ZooModel getModel() { - return model; - } + private static Translator translator(int inputSize) { + return new Translator<>() { + @Override + public NDList processInput(TranslatorContext ctx, float[] features) { + return new NDList(ctx.getNDManager().create(features, new Shape(1, inputSize))); + } - /** Returns float[]{prepid} where prepid in {2212,45,46,47,49}. - * @param features23 - * @return - * @throws ai.djl.translate.TranslateException */ - public float[] prediction(float[] features23) throws TranslateException { - if (features23 == null || features23.length != 23) { - LOGGER.warning("PrePID input must be float[23]"); - return null; - } - Predictor predictor = model.newPredictor(); - return predictor.predict(features23); + @Override + public float[] processOutput(TranslatorContext ctx, NDList output) { + float[] probabilities = output.get(0).toFloatArray(); + int bestIndex = 0; + for (int i = 1; i < probabilities.length; i++) { + if (probabilities[i] > probabilities[bestIndex]) { + bestIndex = i; + } + } + return new float[]{ + CLASS_IDS[bestIndex], + probabilities[0], probabilities[1], probabilities[2], + probabilities[3], probabilities[4] + }; + } + }; } } diff --git a/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/PrePIDResult.java b/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/PrePIDResult.java index cf50e4d665..58a37e0a61 100644 --- a/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/PrePIDResult.java +++ b/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/PrePIDResult.java @@ -6,7 +6,8 @@ public class PrePIDResult { public final int prepid; public final float p2212, p45, p46, p47, p49; - public PrePIDResult(int trackid, int clusterid, int prepid, float p2212, float p45, float p46, float p47, float p49) { + public PrePIDResult(int trackid, int clusterid, int prepid, + float p2212, float p45, float p46, float p47, float p49) { this.trackid = trackid; this.clusterid = clusterid; this.prepid = prepid; @@ -16,4 +17,10 @@ public PrePIDResult(int trackid, int clusterid, int prepid, float p2212, float p this.p47 = p47; this.p49 = p49; } + + public PrePIDResult(int trackid, int clusterid, float[] prediction) { + this(trackid, clusterid, (int) prediction[0], + prediction[1], prediction[2], prediction[3], + prediction[5], prediction[4]); + } } From 414053d11ac317f8dd8d45a9ea4234645e313b8f Mon Sep 17 00:00:00 2001 From: Uditha Weerasinghe Date: Thu, 30 Jul 2026 08:38:27 -0400 Subject: [PATCH 3/5] feat: updated RecoBankWriter to write the ALERT::ai:pid bank --- .../jlab/rec/alert/banks/RecoBankWriter.java | 36 ++++++++++++++----- 1 file changed, 28 insertions(+), 8 deletions(-) diff --git a/reconstruction/alert/src/main/java/org/jlab/rec/alert/banks/RecoBankWriter.java b/reconstruction/alert/src/main/java/org/jlab/rec/alert/banks/RecoBankWriter.java index a6c8c756b5..5f973eb9fa 100644 --- a/reconstruction/alert/src/main/java/org/jlab/rec/alert/banks/RecoBankWriter.java +++ b/reconstruction/alert/src/main/java/org/jlab/rec/alert/banks/RecoBankWriter.java @@ -4,8 +4,9 @@ import java.util.List; import org.jlab.io.base.DataBank; import org.jlab.io.base.DataEvent; +import org.jlab.rec.alert.AIPID.PIDResult; +import org.jlab.rec.alert.AIPID.PrePIDResult; import org.jlab.rec.alert.projections.TrackProjection; -//import org.jlab.rec.alert.AIpid.PIDResult; import ai.djl.util.Pair; @@ -17,7 +18,7 @@ * @author Whit Armstrong */ public class RecoBankWriter { - + /** * Writes the bank of track projections. * @@ -54,12 +55,12 @@ public static DataBank fillProjectionsBank(DataEvent event, ArrayList results) { + + public int appendPrePIDBank(DataEvent event, ArrayList results) { DataBank bank = event.createBank("ALERT::ai:prepid", results.size()); if (bank == null) { System.err.println("COULD NOT CREATE A ALERT::ai:prepid BANK!!!!!!"); return 1; } - for (int i = 0; i < results.size(); i++) { - org.jlab.rec.alert.AIPID.PrePIDResult r = results.get(i); + PrePIDResult r = results.get(i); bank.setInt("trackid", i, r.trackid); bank.setInt("clusterid", i, r.clusterid); bank.setInt("prepid", i, r.prepid); @@ -111,7 +111,27 @@ public int appendPrePIDBank(DataEvent event, ArrayList results) { + DataBank bank = event.createBank("ALERT::ai:pid", results.size()); + if (bank == null) { + System.err.println("COULD NOT CREATE A ALERT::ai:pid BANK!!!!!!"); + return 1; + } + for (int i = 0; i < results.size(); i++) { + PIDResult r = results.get(i); + bank.setInt("trackid", i, r.trackid); + bank.setInt("clusterid", i, r.clusterid); + bank.setInt("pid", i, r.pid); + bank.setFloat("prob_2212", i, r.p2212); + bank.setFloat("prob_45", i, r.p45); + bank.setFloat("prob_46", i, r.p46); + bank.setFloat("prob_47", i, r.p47); + bank.setFloat("prob_49", i, r.p49); + } event.appendBank(bank); return 0; } From 4cb2d3816bcc08f2d03b67b4d013dfcecd798909 Mon Sep 17 00:00:00 2001 From: Uditha Weerasinghe Date: Thu, 30 Jul 2026 08:39:12 -0400 Subject: [PATCH 4/5] fix/feat: updated the feature extraction for the prepid model and added the feature extraction for the postpid model --- .../org/jlab/service/alert/ALERTEngine.java | 213 +++++++++++++----- 1 file changed, 151 insertions(+), 62 deletions(-) diff --git a/reconstruction/alert/src/main/java/org/jlab/service/alert/ALERTEngine.java b/reconstruction/alert/src/main/java/org/jlab/service/alert/ALERTEngine.java index 2088e45ce6..ab82801fd8 100644 --- a/reconstruction/alert/src/main/java/org/jlab/service/alert/ALERTEngine.java +++ b/reconstruction/alert/src/main/java/org/jlab/service/alert/ALERTEngine.java @@ -24,6 +24,7 @@ import org.jlab.io.hipo.HipoDataSync; import org.jlab.rec.alert.TrackMatchingAI.ModelTrackMatching; import org.jlab.rec.alert.AIPID.ModelPrePID; +import org.jlab.rec.alert.AIPID.ModelPostPID; import org.jlab.rec.alert.banks.RecoBankWriter; import org.jlab.rec.alert.projections.TrackProjector; import org.jlab.rec.atof.hit.ATOFHit; @@ -60,6 +61,7 @@ import ai.djl.util.Pair; import org.jlab.rec.alert.AIPID.PrePIDResult; +import org.jlab.rec.alert.AIPID.PIDResult; /** @@ -115,6 +117,7 @@ public class ALERTEngine extends ReconstructionEngine { private ModelTrackMatching modelTrackMatching; private ModelPrePID modelPrePID; + private ModelPostPID modelPostPID; // AHDC track-finding strategy (driven by ALERT.Mode YAML key) private TrackFinder trackFinder; @@ -158,6 +161,7 @@ public boolean init() { modelTrackMatching = new ModelTrackMatching(); modelPrePID = new ModelPrePID(); + modelPostPID = new ModelPostPID(); Map tableMap = new HashMap<>(); tableMap.put("/calibration/alert/ahdc/gains", 3); @@ -178,7 +182,8 @@ public boolean init() { "AHDC::preclusters", "AHDC::clusters", "AHDC::track", "AHDC::interclusters", "AHDC::docaclusters", "AHDC::ai:prediction", "AHDC::mc", "AHDC::kftrack", - "ALERT::projections", "ALERT::ai:projections", "ALERT::prePID"); + "ALERT::projections", "ALERT::ai:projections", + "ALERT::ai:prePID", "ALERT::ai:postPID"); return true; } @@ -430,9 +435,12 @@ public boolean processDataEventUser(DataEvent event) { matched_ATOF_hit_id.add(new Pair<>(track_id, matchHitId)); } rbc.appendTrackMatchingAIBank(event, matched_ATOF_hit_id); + + HashMap prePidByTrack = new HashMap<>(); // --------------------------------------------------------------------------------------- - // PrePID using AI (AHDC::track + ATOF::clusters matched via ALERT::ai:projections) + // PrePID using AI. Use the ATOF model when a matched wedge and start time exist; + // otherwise use the AHDC-only model. // --------------------------------------------------------------------------------------- if (event.hasBank("ALERT::ai:projections") && event.hasBank("AHDC::track") && event.hasBank("ATOF::hits")) { @@ -441,71 +449,58 @@ public boolean processDataEventUser(DataEvent event) { DataBank bankHit = event.getBank("ATOF::hits"); ArrayList prepid_results = new ArrayList<>(); + double startTime = getEventStartTime(event); for (int i = 0; i < bankProj.rows(); i++) { int trackid = bankProj.getInt("trackid", i); - int hitid = bankProj.getInt("matched_atof_hit_id", i); // TODO: Fix to hit_id instead of clusterid + int hitid = bankProj.getInt("matched_atof_hit_id", i); + int trkRow = findRow(bankTrk, "trackid", trackid); + if (trkRow < 0) continue; + + float px = bankTrk.getFloat("px", trkRow); + float py = bankTrk.getFloat("py", trkRow); + float pz = bankTrk.getFloat("pz", trkRow); + double pMag = Math.sqrt(px * px + py * py + pz * pz); + double pt = Math.sqrt(px * px + py * py); + float[] ahdcFeatures = { + bankTrk.getFloat("x", trkRow), + bankTrk.getFloat("y", trkRow), + bankTrk.getFloat("z", trkRow), + (float) Math.log1p(pMag), + (float) Math.log1p(pt), + (float) Math.atan2(py, px), + pMag > 0 ? (float) Math.acos(pz / pMag) : Float.NaN, + bankTrk.getInt("n_hits", trkRow), + (float) Math.log1p(bankTrk.getFloat("path", trkRow)), + (float) Math.log1p(bankTrk.getFloat("dEdx", trkRow)), + (float) Math.log1p(bankTrk.getFloat("chi2", trkRow)) + }; + if (!allFinite(ahdcFeatures)) continue; - // TODO: refactor this to replace this with single line - int trkRow = -1; - for (int r = 0; r < bankTrk.rows(); r++) { - if (bankTrk.getInt("trackid", r) == trackid) { trkRow = r; break; } - } - if (trkRow < 0) continue; - - int hitRow = -1; - for (int r = 0; r < bankHit.rows(); r++) { - if (bankHit.getInt("id", r) == hitid) { hitRow = r; break; } - } - if (hitRow < 0) continue; - - // Build feature vector float[23] in the exact training order - float[] x = new float[23]; - - // AHDC::track (13) - x[0] = bankTrk.getFloat("x", trkRow); - x[1] = bankTrk.getFloat("y", trkRow); - x[2] = bankTrk.getFloat("z", trkRow); - x[3] = bankTrk.getFloat("px", trkRow); - x[4] = bankTrk.getFloat("py", trkRow); - x[5] = bankTrk.getFloat("pz", trkRow); - x[6] = bankTrk.getInt("n_hits", trkRow); - x[7] = bankTrk.getInt("sum_adc", trkRow); - x[8] = bankTrk.getFloat("path", trkRow); - x[9] = bankTrk.getFloat("dEdx", trkRow); - x[10] = bankTrk.getFloat("p_drift", trkRow); - x[11] = bankTrk.getFloat("chi2", trkRow); - x[12] = bankTrk.getFloat("sum_residuals", trkRow); - - /*// ATOF::clusters (10) - x[13] = bankClu.getInt("n_bar", cluRow); - x[14] = bankClu.getInt("n_wedge", cluRow); - x[15] = bankClu.getFloat("time", cluRow); - x[16] = bankClu.getFloat("x", cluRow); - x[17] = bankClu.getFloat("y", cluRow); - x[18] = bankClu.getFloat("z", cluRow); - x[19] = bankClu.getFloat("energy", cluRow); - x[20] = bankClu.getFloat("pathlength", cluRow); - x[21] = bankClu.getFloat("inpathlength", cluRow); - x[22] = bankClu.getInt("projID", cluRow);*/ - - // ATOF::Hits (Temporarily updating to the same 10 slots as ATOF Clusters would have if it worked) - x[13] = 0f; - x[14] = 0f; - x[15] = bankHit.getFloat("time", hitRow); - x[16] = bankHit.getFloat("x", hitRow); - x[17] = bankHit.getFloat("y", hitRow); - x[18] = bankHit.getFloat("z", hitRow); - x[19] = bankHit.getFloat("energy", hitRow); - x[20] = 0f; - x[21] = 0f; - x[22] = 0f; - try { - float[] pred = modelPrePID.prediction(x); - int prepid = (int) pred[0]; - prepid_results.add(new PrePIDResult(trackid, hitid, prepid, pred[1], pred[2], pred[3], pred[4], pred[5])); + float[] prediction = null; + int hitRow = findRow(bankHit, "id", hitid); + int clusterid = hitRow >= 0 ? bankHit.getInt("clusterid", hitRow) : -1; + if (hitRow >= 0 && Double.isFinite(startTime)) { + float[] atofFeatures = new float[16]; + System.arraycopy(ahdcFeatures, 0, atofFeatures, 0, ahdcFeatures.length); + atofFeatures[11] = bankHit.getFloat("x", hitRow); + atofFeatures[12] = bankHit.getFloat("y", hitRow); + atofFeatures[13] = bankHit.getFloat("z", hitRow); + atofFeatures[14] = (float) Math.log1p(bankHit.getFloat("energy", hitRow)); + atofFeatures[15] = (float) (bankHit.getFloat("time", hitRow) - startTime); + if (allFinite(atofFeatures)) { + prediction = modelPrePID.predictionATOF(atofFeatures); + } + } + if (prediction == null) { + prediction = modelPrePID.predictionAHDC(ahdcFeatures); + } + if (prediction != null) { + prepid_results.add(new PrePIDResult(trackid, clusterid, prediction)); + prePidByTrack.put(trackid, (int) prediction[0]); + } } catch (TranslateException ex) { LOGGER.warning(() -> "Exception in ALERTEngine PrePID: " + ex); } @@ -766,11 +761,105 @@ public boolean processDataEventUser(DataEvent event) { } DataBank recoKFHitsBank = ahdc_writer.fillAHDCHitsBank(event, AHDC_hits); event.appendBank(recoKFHitsBank); // remark: only hits assocuated to a track are saved - + + // Post-KF PID: score every valid ALERT::ai:projections pair independently. + if (event.hasBank("ALERT::ai:projections") && event.hasBank("AHDC::kftrack") + && event.hasBank("ATOF::hits") && event.hasBank("ATOF::clusters")) { + DataBank bankProj = event.getBank("ALERT::ai:projections"); + DataBank bankKF = event.getBank("AHDC::kftrack"); + DataBank bankHit = event.getBank("ATOF::hits"); + DataBank bankCluster = event.getBank("ATOF::clusters"); + ArrayList pidResults = new ArrayList<>(); + double startTime = getEventStartTime(event); + + if (Double.isFinite(startTime)) { + for (int i = 0; i < bankProj.rows(); i++) { + int trackid = bankProj.getInt("trackid", i); + int hitid = bankProj.getInt("matched_atof_hit_id", i); + int trackRow = findRow(bankKF, "trackid", trackid); + int hitRow = findRow(bankHit, "id", hitid); + if (trackRow < 0 || hitRow < 0) continue; + + int clusterid = bankHit.getInt("clusterid", hitRow); + int clusterRow = findRow(bankCluster, "id", clusterid); + if (clusterid < 0 || clusterRow < 0) continue; + + float px = bankKF.getFloat("px", trackRow); + float py = bankKF.getFloat("py", trackRow); + float pz = bankKF.getFloat("pz", trackRow); + double pMag = Math.sqrt(px * px + py * py + pz * pz); + double pt = Math.sqrt(px * px + py * py); + float[] features = { + bankKF.getFloat("x", trackRow), + bankKF.getFloat("y", trackRow), + bankKF.getFloat("z", trackRow), + (float) Math.log1p(pMag), + (float) Math.log1p(pt), + (float) Math.atan2(py, px), + pMag > 0 ? (float) Math.acos(pz / pMag) : Float.NaN, + bankKF.getInt("n_hits", trackRow), + bankCluster.getInt("n_bar", clusterRow), + bankCluster.getInt("n_wedge", clusterRow), + (float) Math.log1p(bankKF.getFloat("dEdx", trackRow)), + bankKF.getFloat("sum_residuals", trackRow), + (float) Math.log1p(bankCluster.getFloat("energy", clusterRow)), + bankCluster.getFloat("x", clusterRow), + bankCluster.getFloat("y", clusterRow), + bankCluster.getFloat("z", clusterRow), + (float) (bankCluster.getFloat("time", clusterRow) - startTime), + bankCluster.getFloat("pathlength", clusterRow) + }; + if (!allFinite(features)) continue; + + try { + float[] prediction = modelPostPID.prediction(features); + if (prediction != null) { + pidResults.add(new PIDResult(trackid, clusterid, prediction)); + } + } catch (TranslateException ex) { + LOGGER.warning(() -> "Exception in ALERTEngine PostPID: " + ex); + } + } + } + rbc.appendPIDBank(event, pidResults); + } + + return true; + } + + private static int findRow(DataBank bank, String field, int value) { + for (int row = 0; row < bank.rows(); row++) { + if (bank.getInt(field, row) == value) return row; + } + return -1; + } + private static boolean allFinite(float[] values) { + for (float value : values) { + if (!Float.isFinite(value)) return false; + } return true; } + private static double getEventStartTime(DataEvent event) { + if (event.hasBank("MC::Particle")) { + DataBank bank = event.getBank("MC::Particle"); + if (bank.rows() > 0) return bank.getFloat("vt", 0); + } + if (event.hasBank("REC::Event")) { + DataBank bank = event.getBank("REC::Event"); + if (bank.rows() > 0) { + double startTime = bank.getFloat("startTime", 0); + if (startTime > 0) { + return 0.0; // REC::Event startTime is already accounted for in ATOF timing; return 0.0 to avoid double-counting. + } else { + return Double.NaN; + } + } + } + return Double.NaN; + } + /** Extract a deduplicated list of ATOF hits from {@code ATOF::hits} for the * GNN graph builder. Dedup key is {@code (sector, layer, component)} — * inference-time variant of the Python dedup which also keys on track id From 700fd3c9cb2d3fa6b5516e720b8c42f8a80cb2fd Mon Sep 17 00:00:00 2001 From: Uditha Weerasinghe <92590276+skuditha@users.noreply.github.com> Date: Fri, 31 Jul 2026 09:29:08 -0400 Subject: [PATCH 5/5] fixes based on Mathieu's comments --- .../src/main/java/org/jlab/service/alert/ALERTEngine.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/reconstruction/alert/src/main/java/org/jlab/service/alert/ALERTEngine.java b/reconstruction/alert/src/main/java/org/jlab/service/alert/ALERTEngine.java index ab82801fd8..49745bfb4f 100644 --- a/reconstruction/alert/src/main/java/org/jlab/service/alert/ALERTEngine.java +++ b/reconstruction/alert/src/main/java/org/jlab/service/alert/ALERTEngine.java @@ -183,7 +183,7 @@ public boolean init() { "AHDC::interclusters", "AHDC::docaclusters", "AHDC::ai:prediction", "AHDC::mc", "AHDC::kftrack", "ALERT::projections", "ALERT::ai:projections", - "ALERT::ai:prePID", "ALERT::ai:postPID"); + "ALERT::ai:prepid", "ALERT::ai:pid"); return true; } @@ -820,8 +820,9 @@ public boolean processDataEventUser(DataEvent event) { LOGGER.warning(() -> "Exception in ALERTEngine PostPID: " + ex); } } + rbc.appendPIDBank(event, pidResults); } - rbc.appendPIDBank(event, pidResults); + } return true;