From 7f081e1a27251b59f7db035403499a051ddc1126 Mon Sep 17 00:00:00 2001 From: skuditha Date: Mon, 18 May 2026 07:54:00 -0400 Subject: [PATCH 1/9] fix: changed atof_clusterid to atof_hit_id ito match the code --- etc/bankdefs/hipo4/alert.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/etc/bankdefs/hipo4/alert.json b/etc/bankdefs/hipo4/alert.json index 64d4faeec1..64c96f1de7 100644 --- a/etc/bankdefs/hipo4/alert.json +++ b/etc/bankdefs/hipo4/alert.json @@ -73,7 +73,7 @@ "info": "ALERT AI-assisted PrePID to be used for Kalman Filter", "entries": [ {"name":"trackid", "type":"I", "info":"AHDC trackid"}, - {"name":"clusterid", "type":"I", "info":"ATOF cluster id"}, + {"name":"atof_hit_id", "type":"I", "info":"matched ATOF hit id, -1 if no hit was matched"}, {"name":"prepid", "type":"I", "info":"argmax PID"}, {"name":"p2212", "type":"F", "info":"P(pid=2212)"}, {"name":"p45", "type":"F", "info":"P(pid=45)"}, From aaed859a8c156f8db117e19b84a15ff64478c2b6 Mon Sep 17 00:00:00 2001 From: skuditha Date: Mon, 18 May 2026 07:56:46 -0400 Subject: [PATCH 2/9] fix: updated the order of PID to match the trained model --- .../org/jlab/rec/alert/AIPID/ModelPrePID.java | 57 +++++++++++-------- 1 file changed, 33 insertions(+), 24 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..51d219293c 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 @@ -21,45 +21,49 @@ import java.util.logging.Logger; 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 static final int[] CLASS_IDS = new int[]{2212, 45, 46, 49, 47}; private final ZooModel model; public ModelPrePID() { - Translator my_translator = new Translator<>() { + Translator translator = new Translator<>() { @Override public NDList processInput(TranslatorContext ctx, float[] floats) { NDManager manager = ctx.getNDManager(); - // IMPORTANT: model expects (batch, 23). Provide (1, 23). - NDArray x = manager.create(floats, new Shape(1, 23)); + NDArray x = manager.create(floats, new Shape(1, PrePIDFeatureBuilder.INPUT_SIZE)); return new NDList(x); } @Override public float[] processOutput(TranslatorContext ctx, NDList ndList) { - NDArray logits = ndList.get(0); // (1,5) - NDArray probs = logits.softmax(1); // (1,5) + NDArray logits = ndList.get(0); // (1,5), model class order + NDArray probs = logits.softmax(1); // convert logits to probabilities - float[] p = probs.toFloatArray(); // length 5 (row-major) + float[] p = probs.toFloatArray(); // model order: p, d, t, he3, he4 - // argmax int bestIdx = 0; float best = p[0]; - for (int k = 1; k < 5; k++) { - if (p[k] > best) { best = p[k]; bestIdx = k; } + for (int k = 1; k < p.length; k++) { + if (p[k] > best) { + best = p[k]; + bestIdx = k; + } } int prepid = CLASS_IDS[bestIdx]; - // Return: prepid + probabilities in fixed class order return new float[]{ (float) prepid, - p[0], p[1], p[2], p[3], p[4] + p[0], // p2212: proton + p[1], // p45: deuteron + p[2], // p46: triton + p[4], // p47: helium4 + p[3] // p49: helium3 }; } }; @@ -73,8 +77,9 @@ public float[] processOutput(TranslatorContext ctx, NDList ndList) { Criteria criteria = Criteria.builder() .setTypes(float[].class, float[].class) .optModelPath(Paths.get(path)) + .optModelName("model_PrePID") .optEngine("PyTorch") - .optTranslator(my_translator) + .optTranslator(translator) .optProgress(new ProgressBar()) .build(); @@ -89,16 +94,20 @@ public ZooModel getModel() { return model; } - /** 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]"); + /** + * Returns float[]{prepid, p2212, p45, p46, p47, p49}. + * + * @param features61 raw 61-feature vector in the improved PrePID order + * @return PID prediction and probabilities in ALERT::ai:prepid bank order + * @throws ai.djl.translate.TranslateException if DJL inference fails + */ + public float[] prediction(float[] features61) throws TranslateException { + if (features61 == null || features61.length != PrePIDFeatureBuilder.INPUT_SIZE) { + LOGGER.warning("PrePID input must be float[" + PrePIDFeatureBuilder.INPUT_SIZE + "]"); return null; } - Predictor predictor = model.newPredictor(); - return predictor.predict(features23); + try (Predictor predictor = model.newPredictor()) { + return predictor.predict(features61); + } } } From 81161316cfd040f36abd03ab343a75b29ce48774 Mon Sep 17 00:00:00 2001 From: skuditha Date: Mon, 18 May 2026 07:57:38 -0400 Subject: [PATCH 3/9] fix: changed the code to reflect the change from cluster to hit ID --- .../main/java/org/jlab/rec/alert/AIPID/PrePIDResult.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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..9d0017982c 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 @@ -2,13 +2,13 @@ public class PrePIDResult { public final int trackid; - public final int clusterid; + public final int atofHitId; 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 atofHitId, int prepid, float p2212, float p45, float p46, float p47, float p49) { this.trackid = trackid; - this.clusterid = clusterid; + this.atofHitId = atofHitId; this.prepid = prepid; this.p2212 = p2212; this.p45 = p45; From 085cd90ccd238f853019e7f5d24e156d59689344 Mon Sep 17 00:00:00 2001 From: skuditha Date: Mon, 18 May 2026 07:59:27 -0400 Subject: [PATCH 4/9] feat: added a feature builder to facilitate feature engineering --- .../jlab/rec/alert/PrePIDFeatureBuilder.java | 373 ++++++++++++++++++ 1 file changed, 373 insertions(+) create mode 100644 reconstruction/alert/src/main/java/org/jlab/rec/alert/PrePIDFeatureBuilder.java diff --git a/reconstruction/alert/src/main/java/org/jlab/rec/alert/PrePIDFeatureBuilder.java b/reconstruction/alert/src/main/java/org/jlab/rec/alert/PrePIDFeatureBuilder.java new file mode 100644 index 0000000000..e3f2570ba5 --- /dev/null +++ b/reconstruction/alert/src/main/java/org/jlab/rec/alert/PrePIDFeatureBuilder.java @@ -0,0 +1,373 @@ +package org.jlab.rec.alert.AIPID; + +import java.util.HashSet; +import java.util.Set; +import org.jlab.io.base.DataBank; + +/** + * Builds the ALERT PrePID model input vector in the exact feature order used by + * the improved alert_prepid training pipeline. + * + *

The deployed TorchScript model consumes 61 raw features. Standardization is + * embedded inside the exported model, so this class must not standardize the + * values before inference.

+ */ +public final class PrePIDFeatureBuilder { + + public static final int INPUT_SIZE = 61; + + private static final double C_MM_PER_NS = 299.792458; + private static final double ATOF_TIME_OFFSET_NS = 124.0; + + private static final double MASS_PROTON_GEV = 0.9382720813; + private static final double MASS_DEUTERON_GEV = 1.8756129426; + private static final double MASS_TRITON_GEV = 2.80892113298; + private static final double MASS_HE3_GEV = 2.80839160743; + private static final double MASS_HE4_GEV = 3.7273794066; + + private PrePIDFeatureBuilder() { + // Utility class. + } + + public static MatchResolution resolveTrackMatch(DataBank projections, DataBank hits, int trackid) { + MatchResolution result = new MatchResolution(); + + if (projections == null || hits == null) { + return result; + } + + for (int ip = 0; ip < projections.rows(); ip++) { + if (projections.getInt("trackid", ip) != trackid) { + continue; + } + + final int hitId = projections.getInt("matched_atof_hit_id", ip); + if (hitId < 0) { + continue; + } + + final int hitRow = findRowByInt(hits, "id", hitId); + if (hitRow < 0) { + result.invalidHitRefCount++; + continue; + } + + result.validMatchCount++; + if (!result.hasValidMatch) { + result.hasValidMatch = true; + result.hitId = hitId; + result.hitRow = hitRow; + } + } + + return result; + } + + public static int findRowByInt(DataBank bank, String column, int value) { + if (bank == null) { + return -1; + } + for (int row = 0; row < bank.rows(); row++) { + if (bank.getInt(column, row) == value) { + return row; + } + } + return -1; + } + + /** + * Build the 61-feature PrePID input row for one AHDC track. + * + * @param tracks AHDC::track bank + * @param trackRow row in AHDC::track + * @param hits ATOF::hits bank + * @param match result of resolving ALERT::ai:projections for this track + * @return raw 61-feature vector in model order + */ + public static float[] buildFeatures(DataBank tracks, int trackRow, DataBank hits, MatchResolution match) { + final float[] x = new float[INPUT_SIZE]; + + final double tx = tracks.getFloat("x", trackRow); + final double ty = tracks.getFloat("y", trackRow); + final double tz = tracks.getFloat("z", trackRow); + final double px = tracks.getFloat("px", trackRow); + final double py = tracks.getFloat("py", trackRow); + final double pz = tracks.getFloat("pz", trackRow); + final int nHits = tracks.getInt("n_hits", trackRow); + final int sumAdc = tracks.getInt("sum_adc", trackRow); + final double dedx = tracks.getFloat("dEdx", trackRow); + final double chi2 = tracks.getFloat("chi2", trackRow); + + final double pT = Math.sqrt(px * px + py * py); + final double p = Math.sqrt(px * px + py * py + pz * pz); + final double theta = Math.atan2(pT, pz); + final double phi = Math.atan2(py, px); + // Defaults used when there is no valid ATOF match. + double atofTime = 0.0; + double hx = 0.0; + double hy = 0.0; + double hz = 0.0; + double hitEnergy = 0.0; + double flightPathMm = 0.0; + int sector = 0; + int layer = 0; + int component = 0; + int clusterNHits = 0; + int clusterUniqueLayers = 0; + double clusterTotalEnergy = 0.0; + double clusterMeanTime = 0.0; + double clusterTimeSpread = 0.0; + double clusterEnergyWeightedTime = 0.0; + double clusterMeanX = 0.0; + double clusterMeanY = 0.0; + double clusterMeanZ = 0.0; + double clusterEnergyWeightedX = 0.0; + double clusterEnergyWeightedY = 0.0; + double clusterEnergyWeightedZ = 0.0; + int hasAtofMatch = 0; + + if (match != null && match.hasValidMatch && hits != null && match.hitRow >= 0) { + hasAtofMatch = 1; + + sector = hits.getInt("sector", match.hitRow); + layer = hits.getInt("layer", match.hitRow); + component = hits.getInt("component", match.hitRow); + final int clusterId = hits.getInt("clusterid", match.hitRow); + final double rawTime = hits.getFloat("time", match.hitRow); + + atofTime = rawTime - ATOF_TIME_OFFSET_NS; + hx = hits.getFloat("x", match.hitRow); + hy = hits.getFloat("y", match.hitRow); + hz = hits.getFloat("z", match.hitRow); + hitEnergy = hits.getFloat("energy", match.hitRow); + + final double dx = hx - tx; + final double dy = hy - ty; + final double dz = hz - tz; + flightPathMm = Math.sqrt(dx * dx + dy * dy + dz * dz); + + ClusterSummary cluster = summarizeCluster(hits, clusterId); + clusterNHits = cluster.nHits; + clusterUniqueLayers = cluster.uniqueLayers; + clusterTotalEnergy = cluster.totalEnergy; + clusterMeanTime = cluster.meanTime; + clusterTimeSpread = cluster.timeSpread; + clusterEnergyWeightedTime = cluster.energyWeightedTime; + clusterMeanX = cluster.meanX; + clusterMeanY = cluster.meanY; + clusterMeanZ = cluster.meanZ; + clusterEnergyWeightedX = cluster.energyWeightedX; + clusterEnergyWeightedY = cluster.energyWeightedY; + clusterEnergyWeightedZ = cluster.energyWeightedZ; + } + + final double beta = (Double.isFinite(flightPathMm) && Double.isFinite(atofTime) && Math.abs(atofTime) > 1.0e-12) + ? flightPathMm / (C_MM_PER_NS * atofTime) + : 0.0; + final boolean betaValid = Double.isFinite(beta) && Math.abs(beta) > 1.0e-12; + final double invBeta = betaValid ? 1.0 / beta : 0.0; + final double invBeta2 = betaValid ? invBeta * invBeta : 0.0; + + final double rigidityGev = p / 1000.0; + final double mass2Term = betaValid ? invBeta2 - 1.0 : 0.0; + final double mass2Q1 = rigidityGev * rigidityGev * mass2Term; + final double mass2Q2 = (2.0 * rigidityGev) * (2.0 * rigidityGev) * mass2Term; + + final float bProton = expectedBeta(rigidityGev, 1, MASS_PROTON_GEV); + final float bDeuteron = expectedBeta(rigidityGev, 1, MASS_DEUTERON_GEV); + final float bTriton = expectedBeta(rigidityGev, 1, MASS_TRITON_GEV); + final float bHe3 = expectedBeta(rigidityGev, 2, MASS_HE3_GEV); + final float bHe4 = expectedBeta(rigidityGev, 2, MASS_HE4_GEV); + + final double logDedx = safeLog1p(dedx); + + x[0] = finiteFloat(tx); // track_x + x[1] = finiteFloat(ty); // track_y + x[2] = finiteFloat(tz); // track_z + x[3] = finiteFloat(px); // track_px + x[4] = finiteFloat(py); // track_py + x[5] = finiteFloat(pz); // track_pz + x[6] = nHits; // n_hits + x[7] = sumAdc; // sum_adc + x[8] = finiteFloat(flightPathMm); // path + x[9] = finiteFloat(dedx); // dEdx + x[10] = finiteFloat(chi2); // chi2 + x[11] = finiteFloat(atofTime); // atof_time + x[12] = finiteFloat(hx); // hit_x + x[13] = finiteFloat(hy); // hit_y + x[14] = finiteFloat(hz); // hit_z + x[15] = finiteFloat(hitEnergy); // hit_energy + x[16] = finiteFloat(p); // p_over_q + x[17] = finiteFloat(safeLog1p(p)); // log1p_p_over_q + x[18] = finiteFloat(pT); // pT + x[19] = finiteFloat(theta); // theta + x[20] = finiteFloat(phi); // phi + x[21] = finiteFloat(safeRatio(sumAdc, nHits)); // sum_adc_per_hit + x[22] = finiteFloat(safeRatio(chi2, nHits)); // chi2_per_hit + x[23] = finiteFloat(safeLog1p(sumAdc)); // log1p_sum_adc + x[24] = finiteFloat(logDedx); // log1p_dEdx + x[25] = finiteFloat(safeLog1p(chi2)); // log1p_chi2 + x[26] = finiteFloat(betaValid ? beta : 0.0); // beta + x[27] = finiteFloat(invBeta); // inv_beta + x[28] = finiteFloat(invBeta2); // inv_beta2 + x[29] = finiteFloat(mass2Q1); // mass2_q1 + x[30] = signedLog1pAbs(mass2Q1); // signed_log1p_abs_mass2_q1 + x[31] = finiteFloat(mass2Q2); // mass2_q2 + x[32] = signedLog1pAbs(mass2Q2); // signed_log1p_abs_mass2_q2 + x[33] = bProton > 0.0f ? finiteFloat(invBeta - 1.0 / bProton) : 0.0f; + x[34] = bDeuteron > 0.0f ? finiteFloat(invBeta - 1.0 / bDeuteron) : 0.0f; + x[35] = bTriton > 0.0f ? finiteFloat(invBeta - 1.0 / bTriton) : 0.0f; + x[36] = bHe3 > 0.0f ? finiteFloat(invBeta - 1.0 / bHe3) : 0.0f; + x[37] = bHe4 > 0.0f ? finiteFloat(invBeta - 1.0 / bHe4) : 0.0f; + x[38] = finiteFloat(dedx * beta * beta); // dedx_times_beta2 + x[39] = betaValid ? finiteFloat(dedx / (beta * beta)) : 0.0f; // dedx_over_beta2 + x[40] = finiteFloat(safeLog1p(hitEnergy)); // log1p_hit_energy + x[41] = finiteFloat(hitEnergy * beta * beta); // hit_energy_times_beta2 + x[42] = finiteFloat(p * beta); // p_times_beta + x[43] = betaValid ? finiteFloat(p / beta) : 0.0f; // p_over_beta + x[44] = finiteFloat(p * logDedx); // p_over_q_times_log1p_dEdx + x[45] = finiteFloat(clusterTotalEnergy); // cluster_total_energy + x[46] = finiteFloat(clusterMeanTime); // cluster_mean_time + x[47] = finiteFloat(clusterTimeSpread); // cluster_time_spread + x[48] = finiteFloat(clusterEnergyWeightedTime); // cluster_energy_weighted_time + x[49] = finiteFloat(clusterMeanX); // cluster_mean_x + x[50] = finiteFloat(clusterMeanY); // cluster_mean_y + x[51] = finiteFloat(clusterMeanZ); // cluster_mean_z + x[52] = finiteFloat(clusterEnergyWeightedX); // cluster_energy_weighted_x + x[53] = finiteFloat(clusterEnergyWeightedY); // cluster_energy_weighted_y + x[54] = finiteFloat(clusterEnergyWeightedZ); // cluster_energy_weighted_z + x[55] = sector; // atof_sector + x[56] = layer; // atof_layer + x[57] = component; // atof_component + x[58] = clusterNHits; // cluster_n_hits + x[59] = clusterUniqueLayers; // cluster_unique_layers + x[60] = hasAtofMatch; // has_atof_match + + return x; + } + + private static ClusterSummary summarizeCluster(DataBank hits, int clusterId) { + ClusterSummary summary = new ClusterSummary(); + if (hits == null || clusterId < 0) { + return summary; + } + + double sumTimeRaw = 0.0; + double minTimeRaw = Double.POSITIVE_INFINITY; + double maxTimeRaw = Double.NEGATIVE_INFINITY; + double sumX = 0.0; + double sumY = 0.0; + double sumZ = 0.0; + double ewTimeRaw = 0.0; + double ewX = 0.0; + double ewY = 0.0; + double ewZ = 0.0; + Set uniqueLayers = new HashSet<>(); + + for (int row = 0; row < hits.rows(); row++) { + if (hits.getInt("clusterid", row) != clusterId) { + continue; + } + + final double energy = hits.getFloat("energy", row); + final double timeRaw = hits.getFloat("time", row); + final double x = hits.getFloat("x", row); + final double y = hits.getFloat("y", row); + final double z = hits.getFloat("z", row); + + summary.nHits++; + summary.totalEnergy += energy; + sumTimeRaw += timeRaw; + minTimeRaw = Math.min(minTimeRaw, timeRaw); + maxTimeRaw = Math.max(maxTimeRaw, timeRaw); + sumX += x; + sumY += y; + sumZ += z; + ewTimeRaw += energy * timeRaw; + ewX += energy * x; + ewY += energy * y; + ewZ += energy * z; + uniqueLayers.add(hits.getInt("layer", row)); + } + + if (summary.nHits <= 0) { + return summary; + } + + final double meanTimeRaw = sumTimeRaw / summary.nHits; + final double ewTimeRawNorm = summary.totalEnergy > 0.0 ? ewTimeRaw / summary.totalEnergy : meanTimeRaw; + + summary.uniqueLayers = uniqueLayers.size(); + summary.meanTime = meanTimeRaw - ATOF_TIME_OFFSET_NS; + summary.timeSpread = maxTimeRaw - minTimeRaw; + summary.energyWeightedTime = ewTimeRawNorm - ATOF_TIME_OFFSET_NS; + summary.meanX = sumX / summary.nHits; + summary.meanY = sumY / summary.nHits; + summary.meanZ = sumZ / summary.nHits; + summary.energyWeightedX = summary.totalEnergy > 0.0 ? ewX / summary.totalEnergy : summary.meanX; + summary.energyWeightedY = summary.totalEnergy > 0.0 ? ewY / summary.totalEnergy : summary.meanY; + summary.energyWeightedZ = summary.totalEnergy > 0.0 ? ewZ / summary.totalEnergy : summary.meanZ; + + return summary; + } + + private static float expectedBeta(double rigidityGev, int absCharge, double massGev) { + final double pGev = Math.abs(absCharge) * rigidityGev; + if (!Double.isFinite(pGev) || pGev <= 0.0 || massGev <= 0.0) { + return 0.0f; + } + return finiteFloat(pGev / Math.sqrt(pGev * pGev + massGev * massGev)); + } + + private static float safeLog1p(double value) { + if (!Double.isFinite(value) || value <= -1.0) { + return 0.0f; + } + return finiteFloat(Math.log1p(value)); + } + + private static float signedLog1pAbs(double value) { + if (!Double.isFinite(value)) { + return 0.0f; + } + final double sign = value < 0.0 ? -1.0 : 1.0; + return finiteFloat(sign * Math.log1p(Math.abs(value))); + } + + private static float safeRatio(double numerator, double denominator) { + if (!Double.isFinite(numerator) || !Double.isFinite(denominator) || Math.abs(denominator) < 1.0e-12) { + return 0.0f; + } + return finiteFloat(numerator / denominator); + } + + private static float finiteFloat(double value) { + if (!Double.isFinite(value)) { + return 0.0f; + } + return (float) value; + } + + public static final class MatchResolution { + public boolean hasValidMatch = false; + public int hitId = -1; + public int hitRow = -1; + public int validMatchCount = 0; + public int invalidHitRefCount = 0; + } + + private static final class ClusterSummary { + int nHits = 0; + int uniqueLayers = 0; + double totalEnergy = 0.0; + double meanTime = 0.0; + double timeSpread = 0.0; + double energyWeightedTime = 0.0; + double meanX = 0.0; + double meanY = 0.0; + double meanZ = 0.0; + double energyWeightedX = 0.0; + double energyWeightedY = 0.0; + double energyWeightedZ = 0.0; + } +} From 03ae2e7e37cd51047ae123348f0ff21e4d199758 Mon Sep 17 00:00:00 2001 From: skuditha Date: Mon, 18 May 2026 08:00:43 -0400 Subject: [PATCH 5/9] fix: updated the bank writer with the corrected atof hit ID instead of cluster ID --- .../src/main/java/org/jlab/rec/alert/banks/RecoBankWriter.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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..cd5bf72089 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 @@ -103,7 +103,7 @@ public int appendPrePIDBank(DataEvent event, ArrayList Date: Mon, 18 May 2026 08:02:58 -0400 Subject: [PATCH 6/9] idiot fix: moved the prepid feature builder to correct location --- .../java/org/jlab/rec/alert/{ => AIPID}/PrePIDFeatureBuilder.java | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename reconstruction/alert/src/main/java/org/jlab/rec/alert/{ => AIPID}/PrePIDFeatureBuilder.java (100%) diff --git a/reconstruction/alert/src/main/java/org/jlab/rec/alert/PrePIDFeatureBuilder.java b/reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/PrePIDFeatureBuilder.java similarity index 100% rename from reconstruction/alert/src/main/java/org/jlab/rec/alert/PrePIDFeatureBuilder.java rename to reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/PrePIDFeatureBuilder.java From 0ac5637e8b9da6e3b2343c814cb75947e4ef711a Mon Sep 17 00:00:00 2001 From: skuditha Date: Mon, 18 May 2026 08:03:36 -0400 Subject: [PATCH 7/9] update: improved pre-PID model called by the ALERT Engine --- .../org/jlab/service/alert/ALERTEngine.java | 90 ++++++------------- 1 file changed, 28 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 a899e20c8a..bd7f1a1437 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 @@ -21,6 +21,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.PrePIDFeatureBuilder; import org.jlab.rec.alert.banks.RecoBankWriter; import org.jlab.rec.alert.projections.TrackProjector; import org.jlab.rec.atof.hit.ATOFHit; @@ -244,7 +245,17 @@ public boolean processDataEvent(DataEvent event) { rbc.appendTrackMatchingAIBank(event, matched_ATOF_hit_id); // --------------------------------------------------------------------------------------- - // PrePID using AI (AHDC::track + ATOF::clusters matched via ALERT::ai:projections) + // PrePID using AI with the improved ALERT PrePID feature contract. + // + // Match resolution follows the shared alert_prepid implementation: + // AHDC::track.trackid + // -> ALERT::ai:projections.trackid + // -> first valid ALERT::ai:projections.matched_atof_hit_id != -1 + // -> ATOF::hits.id + // -> all ATOF::hits rows sharing that hit's clusterid for cluster features. + // + // The model consumes one 61-feature row. Standardization is embedded in + // the TorchScript model file. // --------------------------------------------------------------------------------------- if (event.hasBank("ALERT::ai:projections") && event.hasBank("AHDC::track") && event.hasBank("ATOF::hits")) { @@ -254,70 +265,25 @@ public boolean processDataEvent(DataEvent event) { ArrayList prepid_results = new ArrayList<>(); - 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 - - // 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; + for (int trkRow = 0; trkRow < bankTrk.rows(); trkRow++) { + int trackid = bankTrk.getInt("trackid", trkRow); + PrePIDFeatureBuilder.MatchResolution match = PrePIDFeatureBuilder.resolveTrackMatch(bankProj, bankHit, trackid); + float[] features = PrePIDFeatureBuilder.buildFeatures(bankTrk, trkRow, bankHit, match); try { - float[] pred = modelPrePID.prediction(x); + float[] pred = modelPrePID.prediction(features); + if (pred == null || pred.length < 6) { + continue; + } + int prepid = (int) pred[0]; - prepid_results.add(new PrePIDResult(trackid, hitid, prepid, pred[1], pred[2], pred[3], pred[4], pred[5])); + int atofHitId = match.hasValidMatch ? match.hitId : -1; + prepid_results.add(new PrePIDResult( + trackid, + atofHitId, + prepid, + pred[1], pred[2], pred[3], pred[4], pred[5] + )); } catch (TranslateException ex) { LOGGER.warning(() -> "Exception in ALERTEngine PrePID: " + ex); } From 0af959f5adc0f76bfb0f7e8b74eb6419800dad92 Mon Sep 17 00:00:00 2001 From: skuditha Date: Mon, 18 May 2026 08:29:58 -0400 Subject: [PATCH 8/9] fix: overrode the Batchifier to fix the added dimension issue --- .../java/org/jlab/rec/alert/AIPID/ModelPrePID.java | 12 ++++++++++++ 1 file changed, 12 insertions(+) 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 51d219293c..f12c9daf54 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 @@ -11,6 +11,7 @@ import ai.djl.repository.zoo.ZooModel; import ai.djl.training.util.ProgressBar; import ai.djl.translate.TranslateException; +import ai.djl.translate.Batchifier; import ai.djl.translate.Translator; import ai.djl.translate.TranslatorContext; @@ -24,6 +25,8 @@ public class ModelPrePID { static final Logger LOGGER = Logger.getLogger(ModelPrePID.class.getName()); + // Update to match the improved training class order: + // proton, deuteron, triton, helium3, helium4. private static final int[] CLASS_IDS = new int[]{2212, 45, 46, 49, 47}; private final ZooModel model; @@ -36,10 +39,17 @@ public ModelPrePID() { public NDList processInput(TranslatorContext ctx, float[] floats) { NDManager manager = ctx.getNDManager(); + // The improved TorchScript PrePID model expects one raw 61-feature row. + // Standardization is embedded in the exported model. NDArray x = manager.create(floats, new Shape(1, PrePIDFeatureBuilder.INPUT_SIZE)); return new NDList(x); } + @Override + public Batchifier getBatchifier() { + return null; + } + @Override public float[] processOutput(TranslatorContext ctx, NDList ndList) { NDArray logits = ndList.get(0); // (1,5), model class order @@ -57,6 +67,8 @@ public float[] processOutput(TranslatorContext ctx, NDList ndList) { } int prepid = CLASS_IDS[bestIdx]; + // Return the bank order expected by ALERT::ai:prepid: + // prepid, p2212, p45, p46, p47, p49. return new float[]{ (float) prepid, p[0], // p2212: proton From f19a955f11566af5c7c3ce48c34ee79ba75f1cc0 Mon Sep 17 00:00:00 2001 From: skuditha Date: Mon, 18 May 2026 08:46:30 -0400 Subject: [PATCH 9/9] change: atof hit ID to short instead of int. --- etc/bankdefs/hipo4/alert.json | 2 +- .../src/main/java/org/jlab/rec/alert/banks/RecoBankWriter.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/etc/bankdefs/hipo4/alert.json b/etc/bankdefs/hipo4/alert.json index 64c96f1de7..3d4833a324 100644 --- a/etc/bankdefs/hipo4/alert.json +++ b/etc/bankdefs/hipo4/alert.json @@ -73,7 +73,7 @@ "info": "ALERT AI-assisted PrePID to be used for Kalman Filter", "entries": [ {"name":"trackid", "type":"I", "info":"AHDC trackid"}, - {"name":"atof_hit_id", "type":"I", "info":"matched ATOF hit id, -1 if no hit was matched"}, + {"name":"atof_hit_id", "type":"S", "info":"matched ATOF hit id, -1 if no hit was matched"}, {"name":"prepid", "type":"I", "info":"argmax PID"}, {"name":"p2212", "type":"F", "info":"P(pid=2212)"}, {"name":"p45", "type":"F", "info":"P(pid=45)"}, 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 cd5bf72089..bb39dffb90 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 @@ -103,7 +103,7 @@ public int appendPrePIDBank(DataEvent event, ArrayList