Skip to content

Commit d24aa6f

Browse files
skudithaUditha Weerasinghebaltzell
authored
feat: ALERT pre- and post-KF AI PID inference (#1368)
* feat: added implementation for post-KF PID model * fix: updated PrePID model to give a prediction even for tracks without ATOF matches * feat: updated RecoBankWriter to write the ALERT::ai:pid bank * fix/feat: updated the feature extraction for the prepid model and added the feature extraction for the postpid model * fixes based on Mathieu's comments * added guard for appendPIDBank --------- Co-authored-by: Uditha Weerasinghe <skuditha@jlab.org> Co-authored-by: Nathan Baltzell <baltzell@gmx.com>
1 parent d48a81c commit d24aa6f

6 files changed

Lines changed: 351 additions & 132 deletions

File tree

‎reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/ModelPostPID.java‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
package org.jlab.rec.alert.AIPID;
2+
3+
import ai.djl.MalformedModelException;
4+
import ai.djl.inference.Predictor;
5+
import ai.djl.ndarray.NDList;
6+
import ai.djl.ndarray.types.Shape;
7+
import ai.djl.repository.zoo.Criteria;
8+
import ai.djl.repository.zoo.ModelNotFoundException;
9+
import ai.djl.repository.zoo.ZooModel;
10+
import ai.djl.training.util.ProgressBar;
11+
import ai.djl.translate.TranslateException;
12+
import ai.djl.translate.Translator;
13+
import ai.djl.translate.TranslatorContext;
14+
import java.io.IOException;
15+
import java.nio.file.Paths;
16+
import java.util.logging.Logger;
17+
import org.jlab.utils.CLASResources;
18+
19+
public class ModelPostPID {
20+
21+
private static final Logger LOGGER = Logger.getLogger(ModelPostPID.class.getName());
22+
private static final int[] CLASS_IDS = {2212, 45, 46, 49, 47};
23+
private static final int INPUT_SIZE = 18;
24+
25+
private final ZooModel<float[], float[]> model;
26+
27+
public ModelPostPID() {
28+
System.setProperty("ai.djl.pytorch.num_interop_threads", "1");
29+
System.setProperty("ai.djl.pytorch.num_threads", "1");
30+
System.setProperty("ai.djl.pytorch.graph_optimizer", "false");
31+
32+
String path = CLASResources.getResourcePath("etc/data/nnet/rg-l/model_PID/");
33+
Criteria<float[], float[]> criteria = Criteria.builder()
34+
.setTypes(float[].class, float[].class)
35+
.optModelPath(Paths.get(path))
36+
.optEngine("PyTorch")
37+
.optTranslator(translator())
38+
.optProgress(new ProgressBar())
39+
.build();
40+
try {
41+
model = criteria.loadModel();
42+
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
43+
throw new RuntimeException(e);
44+
}
45+
}
46+
47+
public float[] prediction(float[] features) throws TranslateException {
48+
if (features == null || features.length != INPUT_SIZE) {
49+
LOGGER.warning("PostPID input must be float[18]");
50+
return null;
51+
}
52+
try (Predictor<float[], float[]> predictor = model.newPredictor()) {
53+
return predictor.predict(features);
54+
}
55+
}
56+
57+
private static Translator<float[], float[]> translator() {
58+
return new Translator<>() {
59+
@Override
60+
public NDList processInput(TranslatorContext ctx, float[] features) {
61+
return new NDList(ctx.getNDManager().create(features, new Shape(1, INPUT_SIZE)));
62+
}
63+
64+
@Override
65+
public float[] processOutput(TranslatorContext ctx, NDList output) {
66+
float[] probabilities = output.get(0).toFloatArray();
67+
int bestIndex = 0;
68+
for (int i = 1; i < probabilities.length; i++) {
69+
if (probabilities[i] > probabilities[bestIndex]) {
70+
bestIndex = i;
71+
}
72+
}
73+
return new float[]{
74+
CLASS_IDS[bestIndex],
75+
probabilities[0], probabilities[1], probabilities[2],
76+
probabilities[3], probabilities[4]
77+
};
78+
}
79+
};
80+
}
81+
}

‎reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/ModelPrePID.java‎

Lines changed: 63 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,7 @@
22

33
import ai.djl.MalformedModelException;
44
import ai.djl.inference.Predictor;
5-
import ai.djl.ndarray.NDArray;
65
import ai.djl.ndarray.NDList;
7-
import ai.djl.ndarray.NDManager;
86
import ai.djl.ndarray.types.Shape;
97
import ai.djl.repository.zoo.Criteria;
108
import ai.djl.repository.zoo.ModelNotFoundException;
@@ -13,92 +11,96 @@
1311
import ai.djl.translate.TranslateException;
1412
import ai.djl.translate.Translator;
1513
import ai.djl.translate.TranslatorContext;
16-
17-
import org.jlab.utils.CLASResources;
18-
1914
import java.io.IOException;
2015
import java.nio.file.Paths;
2116
import java.util.logging.Logger;
17+
import org.jlab.utils.CLASResources;
2218

2319
public class ModelPrePID {
24-
25-
static final Logger LOGGER = Logger.getLogger(ModelPrePID.class.getName());
26-
// Must match training class order
27-
private static final int[] CLASS_IDS = new int[]{2212, 45, 46, 47, 49};
28-
29-
private final ZooModel<float[], float[]> model;
30-
31-
public ModelPrePID() {
3220

33-
Translator<float[], float[]> my_translator = new Translator<>() {
21+
private static final Logger LOGGER = Logger.getLogger(ModelPrePID.class.getName());
22+
private static final int[] CLASS_IDS = {2212, 45, 46, 49, 47};
3423

35-
@Override
36-
public NDList processInput(TranslatorContext ctx, float[] floats) {
37-
NDManager manager = ctx.getNDManager();
24+
private final ZooModel<float[], float[]> ahdcModel;
25+
private final ZooModel<float[], float[]> atofModel;
3826

39-
// IMPORTANT: model expects (batch, 23). Provide (1, 23).
40-
NDArray x = manager.create(floats, new Shape(1, 23));
41-
return new NDList(x);
42-
}
27+
public ModelPrePID() {
28+
System.setProperty("ai.djl.pytorch.num_interop_threads", "1");
29+
System.setProperty("ai.djl.pytorch.num_threads", "1");
30+
System.setProperty("ai.djl.pytorch.graph_optimizer", "false");
4331

44-
@Override
45-
public float[] processOutput(TranslatorContext ctx, NDList ndList) {
46-
NDArray logits = ndList.get(0); // (1,5)
47-
NDArray probs = logits.softmax(1); // (1,5)
32+
ahdcModel = loadModel("model_prePID_AHDC", 11);
33+
atofModel = loadModel("model_prePID_ATOF", 16);
34+
}
4835

49-
float[] p = probs.toFloatArray(); // length 5 (row-major)
36+
public ZooModel<float[], float[]> getModel() {
37+
return ahdcModel;
38+
}
5039

51-
// argmax
52-
int bestIdx = 0;
53-
float best = p[0];
54-
for (int k = 1; k < 5; k++) {
55-
if (p[k] > best) { best = p[k]; bestIdx = k; }
56-
}
57-
int prepid = CLASS_IDS[bestIdx];
40+
public float[] prediction(float[] features) throws TranslateException {
41+
if (features != null && features.length == 16) {
42+
return predictionATOF(features);
43+
}
44+
return predictionAHDC(features);
45+
}
5846

59-
// Return: prepid + probabilities in fixed class order
60-
return new float[]{
61-
(float) prepid,
62-
p[0], p[1], p[2], p[3], p[4]
63-
};
64-
}
65-
};
47+
public float[] predictionAHDC(float[] features) throws TranslateException {
48+
return predict(ahdcModel, features, 11);
49+
}
6650

67-
System.setProperty("ai.djl.pytorch.num_interop_threads", "1");
68-
System.setProperty("ai.djl.pytorch.num_threads", "1");
69-
System.setProperty("ai.djl.pytorch.graph_optimizer", "false");
51+
public float[] predictionATOF(float[] features) throws TranslateException {
52+
return predict(atofModel, features, 16);
53+
}
7054

71-
String path = CLASResources.getResourcePath("etc/data/nnet/rg-l/model_PrePID/");
55+
private static float[] predict(ZooModel<float[], float[]> model, float[] features,
56+
int expectedSize) throws TranslateException {
57+
if (features == null || features.length != expectedSize) {
58+
LOGGER.warning("PrePID input must be float[" + expectedSize + "]");
59+
return null;
60+
}
61+
try (Predictor<float[], float[]> predictor = model.newPredictor()) {
62+
return predictor.predict(features);
63+
}
64+
}
7265

66+
private static ZooModel<float[], float[]> loadModel(String directory, int inputSize) {
67+
String path = CLASResources.getResourcePath("etc/data/nnet/rg-l/" + directory + "/");
7368
Criteria<float[], float[]> criteria = Criteria.builder()
7469
.setTypes(float[].class, float[].class)
7570
.optModelPath(Paths.get(path))
7671
.optEngine("PyTorch")
77-
.optTranslator(my_translator)
72+
.optTranslator(translator(inputSize))
7873
.optProgress(new ProgressBar())
7974
.build();
80-
8175
try {
82-
model = criteria.loadModel();
76+
return criteria.loadModel();
8377
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
8478
throw new RuntimeException(e);
8579
}
8680
}
8781

88-
public ZooModel<float[], float[]> getModel() {
89-
return model;
90-
}
82+
private static Translator<float[], float[]> translator(int inputSize) {
83+
return new Translator<>() {
84+
@Override
85+
public NDList processInput(TranslatorContext ctx, float[] features) {
86+
return new NDList(ctx.getNDManager().create(features, new Shape(1, inputSize)));
87+
}
9188

92-
/** Returns float[]{prepid} where prepid in {2212,45,46,47,49}.
93-
* @param features23
94-
* @return
95-
* @throws ai.djl.translate.TranslateException */
96-
public float[] prediction(float[] features23) throws TranslateException {
97-
if (features23 == null || features23.length != 23) {
98-
LOGGER.warning("PrePID input must be float[23]");
99-
return null;
100-
}
101-
Predictor<float[], float[]> predictor = model.newPredictor();
102-
return predictor.predict(features23);
89+
@Override
90+
public float[] processOutput(TranslatorContext ctx, NDList output) {
91+
float[] probabilities = output.get(0).toFloatArray();
92+
int bestIndex = 0;
93+
for (int i = 1; i < probabilities.length; i++) {
94+
if (probabilities[i] > probabilities[bestIndex]) {
95+
bestIndex = i;
96+
}
97+
}
98+
return new float[]{
99+
CLASS_IDS[bestIndex],
100+
probabilities[0], probabilities[1], probabilities[2],
101+
probabilities[3], probabilities[4]
102+
};
103+
}
104+
};
103105
}
104106
}

‎reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/PIDResult.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
package org.jlab.rec.alert.AIPID;
2+
3+
public class PIDResult {
4+
public final int trackid;
5+
public final int clusterid;
6+
public final int pid;
7+
public final float p2212, p45, p46, p47, p49;
8+
9+
public PIDResult(int trackid, int clusterid, float[] prediction) {
10+
this.trackid = trackid;
11+
this.clusterid = clusterid;
12+
this.pid = (int) prediction[0];
13+
this.p2212 = prediction[1];
14+
this.p45 = prediction[2];
15+
this.p46 = prediction[3];
16+
this.p49 = prediction[4];
17+
this.p47 = prediction[5];
18+
}
19+
}

‎reconstruction/alert/src/main/java/org/jlab/rec/alert/AIPID/PrePIDResult.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@ public class PrePIDResult {
66
public final int prepid;
77
public final float p2212, p45, p46, p47, p49;
88

9-
public PrePIDResult(int trackid, int clusterid, int prepid, float p2212, float p45, float p46, float p47, float p49) {
9+
public PrePIDResult(int trackid, int clusterid, int prepid,
10+
float p2212, float p45, float p46, float p47, float p49) {
1011
this.trackid = trackid;
1112
this.clusterid = clusterid;
1213
this.prepid = prepid;
@@ -16,4 +17,10 @@ public PrePIDResult(int trackid, int clusterid, int prepid, float p2212, float p
1617
this.p47 = p47;
1718
this.p49 = p49;
1819
}
20+
21+
public PrePIDResult(int trackid, int clusterid, float[] prediction) {
22+
this(trackid, clusterid, (int) prediction[0],
23+
prediction[1], prediction[2], prediction[3],
24+
prediction[5], prediction[4]);
25+
}
1926
}

‎reconstruction/alert/src/main/java/org/jlab/rec/alert/banks/RecoBankWriter.java‎

Lines changed: 28 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,9 @@
44
import java.util.List;
55
import org.jlab.io.base.DataBank;
66
import org.jlab.io.base.DataEvent;
7+
import org.jlab.rec.alert.AIPID.PIDResult;
8+
import org.jlab.rec.alert.AIPID.PrePIDResult;
79
import org.jlab.rec.alert.projections.TrackProjection;
8-
//import org.jlab.rec.alert.AIpid.PIDResult;
910

1011
import ai.djl.util.Pair;
1112

@@ -17,7 +18,7 @@
1718
* @author Whit Armstrong
1819
*/
1920
public class RecoBankWriter {
20-
21+
2122
/**
2223
* Writes the bank of track projections.
2324
*
@@ -54,12 +55,12 @@ public static DataBank fillProjectionsBank(DataEvent event, ArrayList<TrackProje
5455
}
5556
return bank;
5657
}
57-
58+
5859
/**
5960
* Appends the alert match banks to an event.
6061
*
6162
* @param event the {@link DataEvent} in which to append the banks
62-
* @param projections the {@link ArrayList} of {@link TrackProjection} containing the
63+
* @param projections the {@link ArrayList} of {@link TrackProjection} containing the
6364
* track projections info to be added
6465
*
6566
* @return 0 if it worked, 1 if it failed
@@ -91,17 +92,16 @@ public int appendTrackMatchingAIBank(DataEvent event, ArrayList<Pair<Integer, In
9192

9293
return 0;
9394
}
94-
95-
public int appendPrePIDBank(DataEvent event, ArrayList<org.jlab.rec.alert.AIPID.PrePIDResult> results) {
95+
96+
public int appendPrePIDBank(DataEvent event, ArrayList<PrePIDResult> results) {
9697

9798
DataBank bank = event.createBank("ALERT::ai:prepid", results.size());
9899
if (bank == null) {
99100
System.err.println("COULD NOT CREATE A ALERT::ai:prepid BANK!!!!!!");
100101
return 1;
101102
}
102-
103103
for (int i = 0; i < results.size(); i++) {
104-
org.jlab.rec.alert.AIPID.PrePIDResult r = results.get(i);
104+
PrePIDResult r = results.get(i);
105105
bank.setInt("trackid", i, r.trackid);
106106
bank.setInt("clusterid", i, r.clusterid);
107107
bank.setInt("prepid", i, r.prepid);
@@ -111,7 +111,27 @@ public int appendPrePIDBank(DataEvent event, ArrayList<org.jlab.rec.alert.AIPID.
111111
bank.setFloat("p47", i, r.p47);
112112
bank.setFloat("p49", i, r.p49);
113113
}
114+
event.appendBank(bank);
115+
return 0;
116+
}
114117

118+
public int appendPIDBank(DataEvent event, ArrayList<PIDResult> results) {
119+
DataBank bank = event.createBank("ALERT::ai:pid", results.size());
120+
if (bank == null) {
121+
System.err.println("COULD NOT CREATE A ALERT::ai:pid BANK!!!!!!");
122+
return 1;
123+
}
124+
for (int i = 0; i < results.size(); i++) {
125+
PIDResult r = results.get(i);
126+
bank.setInt("trackid", i, r.trackid);
127+
bank.setInt("clusterid", i, r.clusterid);
128+
bank.setInt("pid", i, r.pid);
129+
bank.setFloat("prob_2212", i, r.p2212);
130+
bank.setFloat("prob_45", i, r.p45);
131+
bank.setFloat("prob_46", i, r.p46);
132+
bank.setFloat("prob_47", i, r.p47);
133+
bank.setFloat("prob_49", i, r.p49);
134+
}
115135
event.appendBank(bank);
116136
return 0;
117137
}

0 commit comments

Comments
 (0)