Merge branch 'master-Dichroic' into 'master'

Cherry-picking dichroic JSS

See merge request !30281

Reconstruction/Jet/JetSubStructureMomentTools/JetSubStructureMomentTools/EnergyCorrelatorGeneralizedRatiosTool.h

@@ -21,8 +21,11 @@ class EnergyCorrelatorGeneralizedRatiosTool :
       int modifyJet(xAOD::Jet &jet) const override;
 
     private:
+      bool m_doN3;
+      bool m_doLSeries;
       std::vector<float> m_betaVals;
       std::vector<float> m_cleaned_betaVals; // Local vector for cleaned up inputs
+      bool m_doDichroic;
 
 };
 

Reconstruction/Jet/JetSubStructureMomentTools/JetSubStructureMomentTools/EnergyCorrelatorGeneralizedTool.h

@@ -26,6 +26,7 @@ class EnergyCorrelatorGeneralizedTool :
       bool m_doLSeries;
       std::vector<float> m_betaVals;
       std::vector<float> m_cleaned_betaVals; // Local vector for cleaned up inputs
+      bool m_doDichroic;
 
 };
 

Reconstruction/Jet/JetSubStructureMomentTools/JetSubStructureMomentTools/EnergyCorrelatorRatiosTool.h

@@ -21,10 +21,11 @@ class EnergyCorrelatorRatiosTool :
       int modifyJet(xAOD::Jet &jet) const override;
 
     private:
-      bool m_includeECF4;
+      bool m_doC3;
       std::vector<float> m_betaVals;
       std::vector<float> m_cleaned_betaVals; // Local vector for cleaned up inputs
-            
+      bool m_doDichroic;
+
 };
 
 #endif

Reconstruction/Jet/JetSubStructureMomentTools/JetSubStructureMomentTools/EnergyCorrelatorTool.h

@@ -22,9 +22,10 @@ class EnergyCorrelatorTool :
 
     private:
       float m_Beta;
-      bool m_includeECF4;
+      bool m_doC3;
       std::vector<float> m_betaVals;
       std::vector<float> m_cleaned_betaVals; // Local vector for cleaned up inputs
+      bool m_doDichroic;
 
 };
 

Reconstruction/Jet/JetSubStructureMomentTools/JetSubStructureMomentTools/NSubjettinessRatiosTool.h

 /*
-  Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration
+  Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration
 */
 
 #ifndef jetsubstructuremomenttools_nsubjetinessratiostool_header
@@ -16,7 +16,10 @@ class NSubjettinessRatiosTool :
       NSubjettinessRatiosTool(std::string name);
 
       int modifyJet(xAOD::Jet &jet) const;
-};
 
+    private:
+      bool m_doDichroic;
+
+};
 
 #endif

Reconstruction/Jet/JetSubStructureMomentTools/JetSubStructureMomentTools/NSubjettinessTool.h

 /*
-  Copyright (C) 2002-2019 CERN for the benefit of the ATLAS collaboration
+  Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration
 */
 
 #ifndef jetsubstructuremomenttools_nsubjetinesstool_header
@@ -19,6 +19,7 @@ class NSubjettinessTool :
 
     private:
       float m_Alpha;
+      bool m_doDichroic;
       
 };
 

Reconstruction/Jet/JetSubStructureMomentTools/Root/EnergyCorrelatorGeneralizedRatiosTool.cxx

@@ -20,6 +20,9 @@ EnergyCorrelatorGeneralizedRatiosTool::EnergyCorrelatorGeneralizedRatiosTool(std
   JetSubStructureMomentToolsBase(name)
 {
   declareProperty("BetaList", m_betaVals = {});
+  declareProperty("DoN3", m_doN3 = false);
+  declareProperty("DoLSeries", m_doLSeries = false);
+  declareProperty("DoDichroic", m_doDichroic = false);
 }
 
 StatusCode EnergyCorrelatorGeneralizedRatiosTool::initialize() {
@@ -53,67 +56,125 @@ int EnergyCorrelatorGeneralizedRatiosTool::modifyJet(xAOD::Jet &jet) const {
       return 1;
     }
 
-    if (!jet.isAvailable<float>(m_prefix+"ECFG_3_1"+suffix)) {
-      ATH_MSG_WARNING("Energy correlation function " << m_prefix << "ECFG_3_1" << suffix << "' is not available. Exiting..");
-      return 1;
-    }
-
     if (!jet.isAvailable<float>(m_prefix+"ECFG_3_2"+suffix)) {
       ATH_MSG_WARNING("Energy correlation function " << m_prefix << "ECFG_3_2" << suffix << "' is not available. Exiting..");
       return 1;
     }
 
-    if (!jet.isAvailable<float>(m_prefix+"ECFG_4_2"+suffix)) {
-      ATH_MSG_WARNING("Energy correlation function " << m_prefix << "ECFG_4_2" << suffix << "' is not available. Exiting..");
-      return 1;
+    if(m_doN3) {
+      if (!jet.isAvailable<float>(m_prefix+"ECFG_3_1"+suffix)) {
+        ATH_MSG_WARNING("Energy correlation function " << m_prefix << "ECFG_3_1" << suffix << "' is not available. Exiting..");
+        return 1;
+      }
+
+      if (!jet.isAvailable<float>(m_prefix+"ECFG_4_2"+suffix)) {
+        ATH_MSG_WARNING("Energy correlation function " << m_prefix << "ECFG_4_2" << suffix << "' is not available. Exiting..");
+        return 1;
+      }
+    }
+
+    if (m_doDichroic) {
+      if (!jet.isAvailable<float>(m_prefix+"ECFG_2_1_ungroomed"+suffix)) {
+        ATH_MSG_WARNING("Energy correlation function " << m_prefix << "ECFG_2_1_ungroomed" << suffix << "' is not available. Exiting..");
+        return 1;
+      }
+
+      if (!jet.isAvailable<float>(m_prefix+"ECFG_3_1_ungroomed"+suffix)) {
+        ATH_MSG_WARNING("Energy correlation function " << m_prefix << "ECFG_3_1_ungroomed" << suffix << "' is not available. Exiting..");
+        return 1;
+      }
+
+      if(m_doN3) {
+        if (!jet.isAvailable<float>(m_prefix+"ECFG_3_2_ungroomed"+suffix)) {
+          ATH_MSG_WARNING("Energy correlation function " << m_prefix << "ECFG_3_2_ungroomed" << suffix << "' is not available. Exiting..");
+          return 1;
+        }
+      }
     }
 
     float ecfg_2_1 = jet.getAttribute<float>(m_prefix+"ECFG_2_1"+suffix);
+    float ecfg_3_1 = jet.getAttribute<float>(m_prefix+"ECFG_3_1"+suffix);
     float ecfg_3_2 = jet.getAttribute<float>(m_prefix+"ECFG_3_2"+suffix);
-
     float ecfg_4_2 = jet.getAttribute<float>(m_prefix+"ECFG_4_2"+suffix);
-    float ecfg_3_1 = jet.getAttribute<float>(m_prefix+"ECFG_3_1"+suffix);
+
+    float ecfg_2_1_ungroomed = -999.0;
+    float ecfg_3_1_ungroomed = -999.0;
+    float ecfg_3_2_ungroomed = -999.0;
+
+    if (m_doDichroic) {
+      ecfg_2_1_ungroomed = jet.getAttribute<float>(m_prefix+"ECFG_2_1_ungroomed"+suffix);
+      ecfg_3_1_ungroomed = jet.getAttribute<float>(m_prefix+"ECFG_3_1_ungroomed"+suffix);
+      ecfg_3_2_ungroomed = jet.getAttribute<float>(m_prefix+"ECFG_3_2_ungroomed"+suffix);
+    }
 
     // N2
 
-    if(ecfg_2_1 > 1e-8) // Prevent div-0
-      jet.setAttribute(m_prefix+"N2"+suffix, ecfg_3_2  / (pow(ecfg_2_1, 2.0)));
-    else
+    if(ecfg_2_1 > 1e-8) { // Prevent div-0
+      jet.setAttribute(m_prefix+"N2"+suffix, ecfg_3_2 / (pow(ecfg_2_1, 2.0)));
+
+      if(ecfg_3_2_ungroomed > 1e-8 && ecfg_2_1_ungroomed > 1e-8)
+        jet.setAttribute(m_prefix+"N2_dichroic"+suffix, ecfg_3_2_ungroomed / (ecfg_2_1_ungroomed * ecfg_2_1));
+      else
+        jet.setAttribute(m_prefix+"N2_dichroic"+suffix, -999.0);
+    }
+    else {
       jet.setAttribute(m_prefix+"N2"+suffix, -999.0);
+      jet.setAttribute(m_prefix+"N2_dichroic"+suffix, -999.0);
+    }
 
     // N3
     if(ecfg_3_1 > 1e-8) // Prevent div-0
-      jet.setAttribute(m_prefix+"N3"+suffix, ecfg_4_2  / (pow(ecfg_3_1, 2.0)));
+      jet.setAttribute(m_prefix+"N3"+suffix, ecfg_4_2 / (pow(ecfg_3_1, 2.0)));
     else
       jet.setAttribute(m_prefix+"N3"+suffix, -999.0);
 
     // M2
-    if(ecfg_2_1 > 1e-8) // Prevent div-0
-      jet.setAttribute(m_prefix+"M2"+suffix, ecfg_3_2  / ecfg_2_1);
-    else
+    if(ecfg_2_1 > 1e-8) { // Prevent div-0
+      jet.setAttribute(m_prefix+"M2"+suffix, ecfg_3_1 / ecfg_2_1);
+
+      if(ecfg_3_1_ungroomed > 1e-8)
+        jet.setAttribute(m_prefix+"M2_dichroic"+suffix, ecfg_3_1_ungroomed / ecfg_2_1);
+      else
+        jet.setAttribute(m_prefix+"M2_dichroic"+suffix, -999.0);
+    }
+    else {
       jet.setAttribute(m_prefix+"M2"+suffix, -999.0);
+      jet.setAttribute(m_prefix+"M2_dichroic"+suffix, -999.0);
+    }
 
   }
 
-  // These are for t/H discrimination
-  if (!jet.isAvailable<float>(m_prefix+"ECFG_2_1_2") ||
-      !jet.isAvailable<float>(m_prefix+"ECFG_3_1_1") ||
-      !jet.isAvailable<float>(m_prefix+"ECFG_3_2_1") ||
-      !jet.isAvailable<float>(m_prefix+"ECFG_3_2_2") ||
-      !jet.isAvailable<float>(m_prefix+"ECFG_4_2_2") ||
-      !jet.isAvailable<float>(m_prefix+"ECFG_3_3_1") ||
-      !jet.isAvailable<float>(m_prefix+"ECFG_4_4_1")) {
-    ATH_MSG_WARNING("L series energy correlation functions with prefix '" << m_prefix << "' are not all available. Exiting..");
-    return 1;
+  if (m_doLSeries) {
+    // These are for t/H discrimination
+    if (!jet.isAvailable<float>(m_prefix+"ECFG_2_1_2") ||
+        !jet.isAvailable<float>(m_prefix+"ECFG_3_1_1") ||
+        !jet.isAvailable<float>(m_prefix+"ECFG_3_2_1") ||
+        !jet.isAvailable<float>(m_prefix+"ECFG_3_2_2") ||
+        !jet.isAvailable<float>(m_prefix+"ECFG_4_2_2") ||
+        !jet.isAvailable<float>(m_prefix+"ECFG_3_3_1") ||
+        !jet.isAvailable<float>(m_prefix+"ECFG_4_4_1")) {
+      ATH_MSG_WARNING("L series energy correlation functions with prefix '" << m_prefix << "' are not all available. Exiting..");
+      return 1;
+    }
   }
 
-  float ecfg_2_1_2 = jet.getAttribute<float>(m_prefix+"ECFG_2_1_2");
-  float ecfg_3_1_1 = jet.getAttribute<float>(m_prefix+"ECFG_3_1_1");
-  float ecfg_3_2_1 = jet.getAttribute<float>(m_prefix+"ECFG_3_2_1");  
-  float ecfg_3_2_2 = jet.getAttribute<float>(m_prefix+"ECFG_3_2_2");
-  float ecfg_4_2_2 = jet.getAttribute<float>(m_prefix+"ECFG_4_2_2");
-  float ecfg_3_3_1 = jet.getAttribute<float>(m_prefix+"ECFG_3_3_1");
-  float ecfg_4_4_1 = jet.getAttribute<float>(m_prefix+"ECFG_4_4_1");
+  float ecfg_2_1_2 = -999.0;
+  float ecfg_3_1_1 = -999.0;
+  float ecfg_3_2_1 = -999.0;
+  float ecfg_3_2_2 = -999.0;
+  float ecfg_4_2_2 = -999.0;
+  float ecfg_3_3_1 = -999.0;
+  float ecfg_4_4_1 = -999.0;
+
+  if (m_doLSeries) {
+    ecfg_2_1_2 = jet.getAttribute<float>(m_prefix+"ECFG_2_1_2");
+    ecfg_3_1_1 = jet.getAttribute<float>(m_prefix+"ECFG_3_1_1");
+    ecfg_3_2_1 = jet.getAttribute<float>(m_prefix+"ECFG_3_2_1");  
+    ecfg_3_2_2 = jet.getAttribute<float>(m_prefix+"ECFG_3_2_2");
+    ecfg_4_2_2 = jet.getAttribute<float>(m_prefix+"ECFG_4_2_2");
+    ecfg_3_3_1 = jet.getAttribute<float>(m_prefix+"ECFG_3_3_1");
+    ecfg_4_4_1 = jet.getAttribute<float>(m_prefix+"ECFG_4_4_1");
+  }
 
   // L-series variables
   // (experimental for ttH t/H discrimination)

Reconstruction/Jet/JetSubStructureMomentTools/Root/EnergyCorrelatorGeneralizedTool.cxx

@@ -23,8 +23,9 @@ EnergyCorrelatorGeneralizedTool::EnergyCorrelatorGeneralizedTool(std::string nam
   ATH_MSG_DEBUG("Initializing EnergyCorrelatorGeneralized tool");
   declareProperty("Beta", m_Beta = 1.0);
   declareProperty("BetaList", m_betaVals = {});
-  declareProperty("DoN3", m_doN3 = false);  
-  declareProperty("DoLSeries", m_doLSeries = false);  
+  declareProperty("DoN3", m_doN3 = false);
+  declareProperty("DoLSeries", m_doLSeries = false);
+  declareProperty("DoDichroic", m_doDichroic = false);
 }
 
 StatusCode EnergyCorrelatorGeneralizedTool::initialize() {
@@ -54,7 +55,23 @@ StatusCode EnergyCorrelatorGeneralizedTool::initialize() {
 int EnergyCorrelatorGeneralizedTool::modifyJet(xAOD::Jet &injet) const {
   
   fastjet::PseudoJet jet;
+  fastjet::PseudoJet jet_ungroomed;
+  
   bool decorate = SetupDecoration(jet,injet);
+  bool decorate_ungroomed = false;
+  if(m_doDichroic) {
+    // Get parent jet here and replace injet
+    ElementLink<xAOD::JetContainer> parentLink = injet.auxdata<ElementLink<xAOD::JetContainer> >("Parent");
+
+    // Return error is parent element link is broken
+    if(!parentLink.isValid()) {
+      ATH_MSG_ERROR("Parent element link is not valid. Aborting");
+      return 1;
+    }
+
+    const xAOD::Jet* parentJet = *(parentLink);
+    decorate_ungroomed = SetupDecoration(jet_ungroomed,*parentJet);
+  }
 
   for(float beta : m_cleaned_betaVals) {
 
@@ -62,9 +79,11 @@ int EnergyCorrelatorGeneralizedTool::modifyJet(xAOD::Jet &injet) const {
 
     // These are used for N2 and M2
     float ECFG_2_1_value = -999, ECFG_3_2_value = -999;
+    float ECFG_2_1_value_ungroomed = -999, ECFG_3_2_value_ungroomed = -999;
 
     // These are used for N3
     float ECFG_3_1_value = -999, ECFG_4_2_value = -999;
+    float ECFG_3_1_value_ungroomed = -999;
 
     if (decorate) {
 
@@ -75,6 +94,11 @@ int EnergyCorrelatorGeneralizedTool::modifyJet(xAOD::Jet &injet) const {
       ECFG_2_1_value = ECFG_2_1.result(jet);
       ECFG_3_2_value = ECFG_3_2.result(jet);
 
+      if (decorate_ungroomed) {
+        ECFG_2_1_value_ungroomed = ECFG_2_1.result(jet_ungroomed);
+        ECFG_3_2_value_ungroomed = ECFG_3_2.result(jet_ungroomed);
+      }
+
       // These are used for N3
 
       if(m_doN3) {
@@ -82,6 +106,11 @@ int EnergyCorrelatorGeneralizedTool::modifyJet(xAOD::Jet &injet) const {
         JetSubStructureUtils::EnergyCorrelatorGeneralized ECFG_3_1(1, 3, beta, JetSubStructureUtils::EnergyCorrelator::pt_R);
         ECFG_3_1_value = ECFG_3_1.result(jet);
         ECFG_4_2_value = ECFG_4_2.result(jet);
+
+        if (decorate_ungroomed) {
+          ECFG_3_1_value_ungroomed = ECFG_3_1.result(jet_ungroomed);
+        }
+
       }
 
     }
@@ -91,6 +120,10 @@ int EnergyCorrelatorGeneralizedTool::modifyJet(xAOD::Jet &injet) const {
     injet.setAttribute(m_prefix+"ECFG_3_1"+suffix, ECFG_3_1_value);
     injet.setAttribute(m_prefix+"ECFG_4_2"+suffix, ECFG_4_2_value);
 
+    injet.setAttribute(m_prefix+"ECFG_2_1_ungroomed"+suffix, ECFG_2_1_value_ungroomed);
+    injet.setAttribute(m_prefix+"ECFG_3_2_ungroomed"+suffix, ECFG_3_2_value_ungroomed);
+    injet.setAttribute(m_prefix+"ECFG_3_1_ungroomed"+suffix, ECFG_3_1_value_ungroomed);
+
   }
 
   // These are for t/H discrimination

Reconstruction/Jet/JetSubStructureMomentTools/Root/EnergyCorrelatorRatiosTool.cxx

@@ -20,7 +20,8 @@ EnergyCorrelatorRatiosTool::EnergyCorrelatorRatiosTool(std::string name) :
   JetSubStructureMomentToolsBase(name)
 {
   declareProperty("BetaList", m_betaVals = {});
-  declareProperty("IncludeECF4",  m_includeECF4  = false);
+  declareProperty("DoC3",  m_doC3 = false);
+  declareProperty("DoDichroic", m_doDichroic = false);
 }
 
 StatusCode EnergyCorrelatorRatiosTool::initialize() {
@@ -64,20 +65,61 @@ int EnergyCorrelatorRatiosTool::modifyJet(xAOD::Jet &jet) const {
       return 1;
     }
 
-    if (m_includeECF4 && !jet.isAvailable<float>(m_prefix+"ECF4"+suffix)) {
+    if (m_doC3 && !jet.isAvailable<float>(m_prefix+"ECF4"+suffix)) {
       ATH_MSG_WARNING("Energy correlation function " << m_prefix << "ECF4" << suffix << " is not available. Exiting..");
       return 1;
     }
 
+    if(m_doDichroic) {
+      if (!jet.isAvailable<float>(m_prefix+"ECF1_ungroomed"+suffix)) {
+        ATH_MSG_WARNING("Energy correlation function " << m_prefix << "ECF1_ungroomed" << suffix << " is not available. Exiting..");
+        return 1;
+      }
+
+      if (!jet.isAvailable<float>(m_prefix+"ECF2_ungroomed"+suffix)) {
+        ATH_MSG_WARNING("Energy correlation function " << m_prefix << "ECF2_ungroomed" << suffix << " is not available. Exiting..");
+        return 1;
+      }
+
+      if (!jet.isAvailable<float>(m_prefix+"ECF3_ungroomed"+suffix)) {
+        ATH_MSG_WARNING("Energy correlation function " << m_prefix << "ECF3_ungroomed" << suffix << " is not available. Exiting..");
+        return 1;
+      }
+    }
+
     float ecf1 = jet.getAttribute<float>(m_prefix+"ECF1"+suffix);
     float ecf2 = jet.getAttribute<float>(m_prefix+"ECF2"+suffix);
     float ecf3 = jet.getAttribute<float>(m_prefix+"ECF3"+suffix);
 
+    float ecf4 = -999.0;
+    if(m_doC3) {
+      ecf4 = jet.getAttribute<float>(m_prefix+"ECF4"+suffix);
+    }
+
+    float ecf1_ungroomed = -999.0;
+    float ecf2_ungroomed = -999.0;
+    float ecf3_ungroomed = -999.0;
+
+    if(m_doDichroic) {
+      ecf1_ungroomed = jet.getAttribute<float>(m_prefix+"ECF1_ungroomed"+suffix);
+      ecf2_ungroomed = jet.getAttribute<float>(m_prefix+"ECF2_ungroomed"+suffix);
+      ecf3_ungroomed = jet.getAttribute<float>(m_prefix+"ECF3_ungroomed"+suffix);
+    }
+
     // D2
-    if(ecf2 > 1e-8) // Prevent div-0
+    if(ecf2 > 1e-8) { // Prevent div-0
       jet.setAttribute(m_prefix+"D2"+suffix, ecf3 * pow(ecf1, 3.0) / pow(ecf2, 3.0));
-    else
+
+      if(ecf2_ungroomed > 1e-8 && ecf3_ungroomed > 1e-8)
+        jet.setAttribute(m_prefix+"D2_dichroic"+suffix, ecf3_ungroomed * ecf1_ungroomed * pow(ecf1, 2.0) / ( pow(ecf2_ungroomed, 2.0) * ecf2 ));
+      else
+        jet.setAttribute(m_prefix+"D2_dichroic"+suffix, -999.0);
+
+    }
+    else {
       jet.setAttribute(m_prefix+"D2"+suffix, -999.0);
+      jet.setAttribute(m_prefix+"D2_dichroic"+suffix, -999.0);
+    }
 
     // C1
     if(ecf1 > 1e-8) // Prevent div-0
@@ -91,14 +133,11 @@ int EnergyCorrelatorRatiosTool::modifyJet(xAOD::Jet &jet) const {
     else
       jet.setAttribute(m_prefix+"C2"+suffix, -999.0);
 
-    if(m_includeECF4) {
-      float ecf4 = jet.getAttribute<float>(m_prefix+"ECF4"+suffix);
-      // C3
-      if(ecf3 > 1e-8) // Prevent div-0
-        jet.setAttribute(m_prefix+"C3"+suffix, ecf4 * ecf2 / pow(ecf3, 2.0));
-      else
-        jet.setAttribute(m_prefix+"C3"+suffix, -999.0);
-    }
+    // C3
+    if(ecf3 > 1e-8) // Prevent div-0
+      jet.setAttribute(m_prefix+"C3"+suffix, ecf4 * ecf2 / pow(ecf3, 2.0));
+    else
+      jet.setAttribute(m_prefix+"C3"+suffix, -999.0);
   }
 
   return 0;

Reconstruction/Jet/JetSubStructureMomentTools/Root/EnergyCorrelatorTool.cxx

@@ -22,7 +22,8 @@ EnergyCorrelatorTool::EnergyCorrelatorTool(std::string name) :
   ATH_MSG_DEBUG("Initializing EnergyCorrelator tool.");
   declareProperty("Beta", m_Beta = 1.0);
   declareProperty("BetaList", m_betaVals = {});
-  declareProperty("IncludeECF4", m_includeECF4 = false);
+  declareProperty("DoC3", m_doC3 = false);
+  declareProperty("DoDichroic", m_doDichroic = false);
 }
 
 StatusCode EnergyCorrelatorTool::initialize() {
@@ -52,7 +53,23 @@ StatusCode EnergyCorrelatorTool::initialize() {
 int EnergyCorrelatorTool::modifyJet(xAOD::Jet &injet) const {
   
   PseudoJet jet;
+  PseudoJet jet_ungroomed;
+
   bool decorate = SetupDecoration(jet,injet);
+  bool decorate_ungroomed = false;
+  if(m_doDichroic) {
+    // Get parent jet here and replace injet
+    ElementLink<xAOD::JetContainer> parentLink = injet.auxdata<ElementLink<xAOD::JetContainer> >("Parent");
+
+    // Return error is parent element link is broken
+    if(!parentLink.isValid()) {
+      ATH_MSG_ERROR("Parent element link is not valid. Aborting");
+      return 1;
+    }
+
+    const xAOD::Jet* parentJet = *(parentLink);
+    decorate_ungroomed = SetupDecoration(jet_ungroomed,*parentJet);
+  }
 
   for(float beta : m_cleaned_betaVals) {
 
@@ -68,6 +85,10 @@ int EnergyCorrelatorTool::modifyJet(xAOD::Jet &injet) const {
     float result_ECF3 = -999;
     float result_ECF4 = -999;
 
+    float result_ECF1_ungroomed = -999;
+    float result_ECF2_ungroomed = -999;
+    float result_ECF3_ungroomed = -999;
+
     if(decorate) {
 
       JetSubStructureUtils::EnergyCorrelator ECF1(1, beta, JetSubStructureUtils::EnergyCorrelator::pt_R);
@@ -78,18 +99,27 @@ int EnergyCorrelatorTool::modifyJet(xAOD::Jet &injet) const {
       result_ECF2 = ECF2.result(jet);
       result_ECF3 = ECF3.result(jet);
 
-      if(m_includeECF4) {
+      if(m_doC3) {
         JetSubStructureUtils::EnergyCorrelator ECF4(4, beta, JetSubStructureUtils::EnergyCorrelator::pt_R);
         result_ECF4 = ECF4.result(jet);
       }
+    
+      if(decorate_ungroomed) {
+        result_ECF1_ungroomed = ECF1.result(jet_ungroomed);
+        result_ECF2_ungroomed = ECF2.result(jet_ungroomed);
+        result_ECF3_ungroomed = ECF3.result(jet_ungroomed);
+      }
+
     }
 
     injet.setAttribute(m_prefix+"ECF1"+suffix, result_ECF1);
     injet.setAttribute(m_prefix+"ECF2"+suffix, result_ECF2);
     injet.setAttribute(m_prefix+"ECF3"+suffix, result_ECF3);
-
-    if(m_includeECF4)
-      injet.setAttribute(m_prefix+"ECF4"+suffix, result_ECF4);
+    injet.setAttribute(m_prefix+"ECF4"+suffix, result_ECF4);
+    
+    injet.setAttribute(m_prefix+"ECF1_ungroomed"+suffix, result_ECF1_ungroomed);
+    injet.setAttribute(m_prefix+"ECF2_ungroomed"+suffix, result_ECF2_ungroomed);
+    injet.setAttribute(m_prefix+"ECF3_ungroomed"+suffix, result_ECF3_ungroomed);
 
   }
 

Reconstruction/Jet/JetSubStructureMomentTools/Root/NSubjettinessRatiosTool.cxx

 /*
-  Copyright (C) 2002-2019 CERN for the benefit of the ATLAS collaboration
+  Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration
 */
 
 #include "JetSubStructureMomentTools/NSubjettinessRatiosTool.h"
@@ -7,6 +7,7 @@
 NSubjettinessRatiosTool::NSubjettinessRatiosTool(std::string name) : 
   JetSubStructureMomentToolsBase(name)
 {
+  declareProperty("DoDichroic", m_doDichroic = false);
 }
 
 int NSubjettinessRatiosTool::modifyJet(xAOD::Jet &jet) const {
@@ -20,51 +21,126 @@ int NSubjettinessRatiosTool::modifyJet(xAOD::Jet &jet) const {
       !jet.isAvailable<float>(m_prefix+"Tau3_wta") ||
       !jet.isAvailable<float>(m_prefix+"Tau4_wta")) {
     
-    ATH_MSG_WARNING("Tau decorations for "<<m_prefix<<"are not all available. Exiting..");
+    ATH_MSG_ERROR("Tau decorations for " << m_prefix << " are not all available. Exiting..");
     return 1;
   }
-  
+
+  if (m_doDichroic) { 
+    if (!jet.isAvailable<float>(m_prefix+"Tau2_ungroomed") ||
+        !jet.isAvailable<float>(m_prefix+"Tau3_ungroomed") ||
+        !jet.isAvailable<float>(m_prefix+"Tau4_ungroomed") ||
+        !jet.isAvailable<float>(m_prefix+"Tau2_wta_ungroomed") ||
+        !jet.isAvailable<float>(m_prefix+"Tau3_wta_ungroomed") ||
+        !jet.isAvailable<float>(m_prefix+"Tau4_wta_ungroomed")) {
+
+      ATH_MSG_ERROR("Ungroomed tau decorations for " << m_prefix << " are not all available. Exiting..");
+      return 1;
+    }
+  }
+
   // Regular
   float tau1 = jet.getAttribute<float>(m_prefix+"Tau1");
   float tau2 = jet.getAttribute<float>(m_prefix+"Tau2");
   float tau3 = jet.getAttribute<float>(m_prefix+"Tau3");
   float tau4 = jet.getAttribute<float>(m_prefix+"Tau4");
 
+  float tau2_ungroomed = -999.0;
+  float tau3_ungroomed = -999.0;
+  float tau4_ungroomed = -999.0;
+
+  if (m_doDichroic) {
+    tau2_ungroomed = jet.getAttribute<float>(m_prefix+"Tau2_ungroomed");
+    tau3_ungroomed = jet.getAttribute<float>(m_prefix+"Tau3_ungroomed");
+    tau4_ungroomed = jet.getAttribute<float>(m_prefix+"Tau4_ungroomed");
+  }
+
   //Prevent div-0 and check against default value (-999) of the decoration
-  if(tau1 > 1e-8)
+  if(tau1 > 1e-8) {
     jet.setAttribute(m_prefix+"Tau21", tau2/tau1);