Add different encoding options

Author: romainschotter

PWGLF/Tasks/Strangeness/strangenessderivedbinnedinfo.cxx

@@ -124,6 +124,17 @@ struct strangenessderivedbinnedinfo {
     Configurable<float> maxIR{"maxIR", -1, "maximum IR collisions"};
   } eventSelections;
 
+  static constexpr float defaultSqrtScalingParameters[1][4] = {{0.1, 0.1, 0, 128}};
+
+  // preselection options
+  struct : ConfigurableGroup {
+    std::string prefix = "encodingOpts";
+    Configurable<bool> useSqrtEncodingForOccupancy{"useSqrtEncodingForOccupancy", false, "Store sqrt(occupancy) instead of occupancy"};
+    Configurable<bool> useSqrtEncodingForRadius{"useSqrtEncodingForRadius", false, "Store sqrt(radius) instead of radius"};
+    Configurable<bool> useSqrtScalingForEncodingPt{"useSqrtScalingForEncodingPt", false, "Store sqrt scaling(pT) instead of pT"};
+    Configurable<LabeledArray<float>> sqrtScalingParameters{"sqrtScalingParameters", {defaultSqrtScalingParameters[0], 4, {"sigma0", "sigma1", "clampMin", "clampMax"}}, "Sqrt scaling parameters"};
+  } encodingOpts;
+
   struct : ConfigurableGroup {
     Configurable<int> v0TypeSelection{"v0Selections.v0TypeSelection", 1, "select on a certain V0 type (leave negative if no selection desired)"};
 
@@ -246,6 +257,20 @@ struct strangenessderivedbinnedinfo {
   // PDG database
   Service<o2::framework::O2DatabasePDG> pdgDB;
 
+  // Sqrt scaling function
+  // Author: Marian Ivanov
+  int codeSqrtScaling(float val, float sigma0, float sigma1, int clampMin, int clampMax) {
+    float code_f = std::asinh((sigma1 * val) / sigma0) / sigma0;
+    return std::clamp(static_cast<int>(std::round(code_f)), clampMin, clampMax);
+  }
+
+  // Function to decode the sqrt scaling 
+  // Author: Marian Ivanov
+  float decodeSqrtScaling(int code, float sigma0, float sigma1, int clampMin, int clampMax) {
+      float code_f = static_cast<float>(code);
+      return (sigma0 / sigma1) * std::sinh(sigma0 * code_f);
+  }
+
   void init(InitContext const&)
   {
     if (analyseK0Short + analyseLambda + analyseAntiLambda + analyseXi + analyseOmega > 1) {
@@ -481,6 +506,8 @@ struct strangenessderivedbinnedinfo {
     } else {
       centrality = collision.centFT0C();
       occupancy = eventSelections.useFT0CbasedOccupancy ? collision.ft0cOccupancyInTimeRange() : collision.trackOccupancyInTimeRange();
+      if(encodingOpts.useSqrtEncodingForOccupancy)
+        occupancy = std::sqrt(occupancy);
     }
     histos.fill(HIST("hEventCentrality"), centrality);
     histos.fill(HIST("hEventOccupancy"), occupancy);
@@ -772,14 +799,17 @@ struct strangenessderivedbinnedinfo {
         if (v0.v0Type() != v0Selections.v0TypeSelection && v0Selections.v0TypeSelection > -1)
           continue; // skip V0s that are not standard
 
+        float pT = encodingOpts.useSqrtScalingForEncodingPt ? codeSqrtScaling(v0.pt(), encodingOpts.sqrtScalingParameters->get("sigma0"), encodingOpts.sqrtScalingParameters->get("sigma1"), encodingOpts.sqrtScalingParameters->get("clampMin"), encodingOpts.sqrtScalingParameters->get("clampMax")) : v0.pt();
+        float decayRadius = encodingOpts.useSqrtEncodingForRadius ? decayRadius = std::sqrt(v0.v0radius()) : v0.v0radius();
+
         if (analyseK0Short && isV0Selected(v0, collision, v0.yK0Short())) {
-          histos.fill(HIST("h9dMassPtPhiEtaPtArmV0AlphaV0RadiusCentOcc"), v0.mK0Short(), v0.pt(), v0.phi(), v0.eta(), v0.qtarm(), v0.alpha(), v0.v0radius(), centrality, occupancy);
+          histos.fill(HIST("h9dMassPtPhiEtaPtArmV0AlphaV0RadiusCentOcc"), v0.mK0Short(), pT, v0.phi(), v0.eta(), v0.qtarm(), v0.alpha(), decayRadius, centrality, occupancy);
         }
         if (analyseLambda && isV0Selected(v0, collision, v0.yLambda())) {
-          histos.fill(HIST("h9dMassPtPhiEtaPtArmV0AlphaV0RadiusCentOcc"), v0.mLambda(), v0.pt(), v0.phi(), v0.eta(), v0.qtarm(), v0.alpha(), v0.v0radius(), centrality, occupancy);
+          histos.fill(HIST("h9dMassPtPhiEtaPtArmV0AlphaV0RadiusCentOcc"), v0.mLambda(), pT, v0.phi(), v0.eta(), v0.qtarm(), v0.alpha(), decayRadius, centrality, occupancy);
         }
         if (analyseAntiLambda && isV0Selected(v0, collision, v0.yLambda())) {
-          histos.fill(HIST("h9dMassPtPhiEtaPtArmV0AlphaV0RadiusCentOcc"), v0.mAntiLambda(), v0.pt(), v0.phi(), v0.eta(), v0.qtarm(), v0.alpha(), v0.v0radius(), centrality, occupancy);
+          histos.fill(HIST("h9dMassPtPhiEtaPtArmV0AlphaV0RadiusCentOcc"), v0.mAntiLambda(), pT, v0.phi(), v0.eta(), v0.qtarm(), v0.alpha(), decayRadius, centrality, occupancy);
         }
       } // end v0 loop
     }
@@ -791,11 +821,14 @@ struct strangenessderivedbinnedinfo {
             std::abs(cascade.bacheloreta()) > cascSelections.daughterEtaCut)
           continue; // remove acceptance that's badly reproduced by MC / superfluous in future
 
+        float pT = encodingOpts.useSqrtScalingForEncodingPt ? codeSqrtScaling(cascade.pt(), encodingOpts.sqrtScalingParameters->get("sigma0"), encodingOpts.sqrtScalingParameters->get("sigma1"), encodingOpts.sqrtScalingParameters->get("clampMin"), encodingOpts.sqrtScalingParameters->get("clampMax")) : cascade.pt();
+        float decayRadius = encodingOpts.useSqrtEncodingForRadius ? decayRadius = std::sqrt(cascade.cascradius()) : cascade.cascradius();
+
         if (analyseXi && isCascadeSelected(cascade, collision, cascade.yXi())) {
-          histos.fill(HIST("h9dMassPtPhiEtaPtArmV0AlphaV0RadiusCentOcc"), cascade.m(1), cascade.pt(), cascade.phi(), cascade.eta(), 0., 0., cascade.cascradius(), centrality, occupancy);
+          histos.fill(HIST("h9dMassPtPhiEtaPtArmV0AlphaV0RadiusCentOcc"), cascade.m(1), pT, cascade.phi(), cascade.eta(), 0., 0., decayRadius, centrality, occupancy);
         }
         if (analyseOmega && isCascadeSelected(cascade, collision, cascade.yOmega())) {
-          histos.fill(HIST("h9dMassPtPhiEtaPtArmV0AlphaV0RadiusCentOcc"), cascade.pt(), cascade.m(2), cascade.pt(), cascade.phi(), cascade.eta(), 0., 0., cascade.cascradius(), centrality, occupancy);
+          histos.fill(HIST("h9dMassPtPhiEtaPtArmV0AlphaV0RadiusCentOcc"), cascade.m(2), pT, cascade.phi(), cascade.eta(), 0., 0., decayRadius, centrality, occupancy);
         }
       } // end cascade loop
     }