diff --git a/reconstruction/ai/src/main/java/org/jlab/service/ai/DCClsComboEngine.java b/reconstruction/ai/src/main/java/org/jlab/service/ai/DCClsComboEngine.java index dc08ac1ee1..e1b3060ce3 100644 --- a/reconstruction/ai/src/main/java/org/jlab/service/ai/DCClsComboEngine.java +++ b/reconstruction/ai/src/main/java/org/jlab/service/ai/DCClsComboEngine.java @@ -32,12 +32,16 @@ import org.jlab.rec.ai.dcCluster.DCClusterCombo; public class DCClsComboEngine extends ReconstructionEngine { - String inputBank = "HitBasedTrkg::Clusters"; + + final static String CONF_INPUT_BANK_PREFIX = "inputBankPrefix"; + String inputBank = "HitBasedTrkg::Clusters"; + + final static String CONF_OUTPUT_BANK_PREFIX = "outputBankPrefix"; String outputBank = "ai::tracks"; final static String CONF_THREADS = "threads"; - - final static String LIMIT_CLSCOMBOS_PERSECTOR = "limitClsCombosPerSector"; + + final static String CONF_LIMIT_CLSCOMBOS_PERSECTOR = "limitClsCombosPerSector"; int limitClsCombosPerSector = (int)1e5; final static String CONF_MODEL_FILE_6CLS = "modelFile6Cls"; @@ -75,10 +79,20 @@ public boolean init() { System.setProperty("ai.djl.pytorch.num_threads", "1"); System.setProperty("ai.djl.pytorch.graph_optimizer", "false"); + if (getEngineConfigString(CONF_INPUT_BANK_PREFIX) != null){ + String inputBankPrefix = getEngineConfigString(CONF_INPUT_BANK_PREFIX); + if(inputBankPrefix.equals("DN")) inputBank = "HitBasedTrkg::DNClusters"; + } + + if (getEngineConfigString(CONF_OUTPUT_BANK_PREFIX) != null){ + String outBankPrefix = getEngineConfigString(CONF_OUTPUT_BANK_PREFIX); + if(outBankPrefix.equals("aidn")) inputBank = "aidn::tracks"; + } + int threads = Integer.parseInt(getEngineConfigString(CONF_THREADS,"64")); - if (getEngineConfigString(LIMIT_CLSCOMBOS_PERSECTOR) != null) - limitClsCombosPerSector = Integer.parseInt(getEngineConfigString(LIMIT_CLSCOMBOS_PERSECTOR)); + if (getEngineConfigString(CONF_LIMIT_CLSCOMBOS_PERSECTOR) != null) + limitClsCombosPerSector = Integer.parseInt(getEngineConfigString(CONF_LIMIT_CLSCOMBOS_PERSECTOR)); if (getEngineConfigString(CONF_THRESHOLD_6CLS) != null) threshold6Cls = Float.parseFloat(getEngineConfigString(CONF_THRESHOLD_6CLS));