jetDsSpecSubs.cxx

// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.

// Jet substructure and spectrum task for D_s mesons
//
// This task is used to reconstruct and analyse jets containing charged D_s
// mesons
//
/// \author Monalisa Melo <monalisa.melo@cern.ch>, Universidade de São Paulo
//

#include "PWGHF/Core/DecayChannels.h"
#include "PWGJE/Core/JetDerivedDataUtilities.h"
#include "PWGJE/Core/JetUtilities.h"
#include "PWGJE/DataModel/Jet.h"
#include "PWGJE/DataModel/JetReducedData.h"
#include "PWGJE/DataModel/JetSubtraction.h"

#include "Common/Core/RecoDecay.h"
#include "Common/DataModel/TrackSelectionTables.h"

#include "Framework/ASoA.h"
#include "Framework/AnalysisDataModel.h"
#include "Framework/HistogramRegistry.h"
#include <Framework/AnalysisTask.h>
#include <Framework/ConfigContext.h>
#include <Framework/Configurable.h>
#include <Framework/DataProcessorSpec.h>
#include <Framework/HistogramSpec.h>
#include <Framework/InitContext.h>
#include <Framework/runDataProcessing.h>

#include "TVector3.h"

#include <cmath>
#include <string>
#include <vector>

using namespace o2;
using namespace o2::framework;
using namespace o2::framework::expressions;

namespace o2::aod
{
namespace jet_distance
{
DECLARE_SOA_COLUMN(JetHfDist, jetHfDist, float);
DECLARE_SOA_COLUMN(JetPt, jetPt, float);
DECLARE_SOA_COLUMN(JetEta, jetEta, float);
DECLARE_SOA_COLUMN(JetPhi, jetPhi, float);
DECLARE_SOA_COLUMN(JetNConst, jetNConst, int);
DECLARE_SOA_COLUMN(HfPt, hfPt, float);
DECLARE_SOA_COLUMN(HfEta, hfEta, float);
DECLARE_SOA_COLUMN(HfPhi, hfPhi, float);
DECLARE_SOA_COLUMN(HfMass, hfMass, float);
DECLARE_SOA_COLUMN(HfY, hfY, float);
DECLARE_SOA_COLUMN(HfMlScore0, hfMlScore0, float);
DECLARE_SOA_COLUMN(HfMlScore1, hfMlScore1, float);
DECLARE_SOA_COLUMN(HfMlScore2, hfMlScore2, float);

// extra
DECLARE_SOA_COLUMN(JetMass, jetMass, float);
DECLARE_SOA_COLUMN(JetGirth, jetGirth, float);
DECLARE_SOA_COLUMN(JetThrust, jetThrust, float);     // lambda_2^1
DECLARE_SOA_COLUMN(JetLambda11, jetLambda11, float); // lambda_1^1
} // namespace jet_distance

DECLARE_SOA_TABLE(JetDistanceTable, "AOD", "JETDISTTABLE",
                  jet_distance::JetHfDist,
                  jet_distance::JetPt,
                  jet_distance::JetEta,
                  jet_distance::JetPhi,
                  jet_distance::JetNConst,
                  jet_distance::HfPt,
                  jet_distance::HfEta,
                  jet_distance::HfPhi,
                  jet_distance::HfMass,
                  jet_distance::HfY,
                  jet_distance::HfMlScore0,
                  jet_distance::HfMlScore1,
                  jet_distance::HfMlScore2,
                  jet_distance::JetMass,
                  jet_distance::JetGirth,
                  jet_distance::JetThrust,
                  jet_distance::JetLambda11);
} // namespace o2::aod

struct JetDsSpecSubs {
  HistogramRegistry registry{
    "registry",
    {
      {"h_collisions", "event status;event status;entries", {HistType::kTH1F, {{4, 0.0, 4.0}}}},
      {"h_track_pt", "track pT;#it{p}_{T,track} (GeV/#it{c});entries", {HistType::kTH1F, {{200, 0., 200.}}}},
      {"h_track_eta", "track #eta;#eta_{track};entries", {HistType::kTH1F, {{100, -1.0, 1.0}}}},
      {"h_track_phi", "track #varphi;#varphi_{track};entries", {HistType::kTH1F, {{80, -1.0, 7.}}}},
      {"h_jet_pt", "jet pT;#it{p}_{T,jet} (GeV/#it{c});entries", {HistType::kTH1F, {{200, 0., 200.}}}},
      {"h_jet_eta", "jet #eta;#eta_{jet};entries", {HistType::kTH1F, {{100, -1.0, 1.0}}}},
      {"h_jet_phi", "jet #phi;#phi_{jet};entries", {HistType::kTH1F, {{80, -1.0, 7.}}}},
      {"h_collision_counter", "# of collisions;", {HistType::kTH1F, {{200, 0., 200.}}}},
      {"h_jet_counter", ";type;counts", {HistType::kTH1F, {{2, 0., 2.}}}},
      {"h_ds_jet_projection", ";z^{D_{S},jet}_{||};dN/dz^{D_{S},jet}_{||}", {HistType::kTH1F, {{1000, 0., 2.}}}},
      {"h_ds_jet_distance_vs_projection", ";#DeltaR_{D_{S},jet};z^{D_{S},jet}_{||}", {HistType::kTH2F, {{1000, 0., 1.}, {1000, 0., 2.}}}},
      {"h_ds_jet_distance", ";#DeltaR_{D_{S},jet};dN/d(#DeltaR)", {HistType::kTH1F, {{1000, 0., 1.}}}},
      {"h_ds_jet_pt", ";p_{T,D_{S} jet};dN/dp_{T,D_{S} jet}", {HistType::kTH1F, {{1000, 0., 100.}}}},
      {"h_ds_jet_eta", ";#eta_{D_{S} jet};entries", {HistType::kTH1F, {{250, -1., 1.}}}},
      {"h_ds_jet_phi", ";#phi_{D_{S} jet};entries", {HistType::kTH1F, {{250, -1., 7.}}}},
      {"hSparse_ds", ";m_{D_{S}};p_{T,D_{S}};z^{D_{S},jet}_{||};#DeltaR_{D_{S},jet}", {HistType::kTHnSparseF, {{60, 1.7, 2.1}, {60, 0., 100.}, {60, 0., 2.}, {60, 0., 1.0}}}},
      {"h_ds_mass", ";m_{D_{S}} (GeV/c^{2});entries", {HistType::kTH1F, {{1000, 0., 6.}}}},
      {"h_ds_eta", ";#eta_{D_{S}};entries", {HistType::kTH1F, {{250, -1., 1.}}}},
      {"h_ds_phi", ";#phi_{D_{S}};entries", {HistType::kTH1F, {{250, -1., 7.}}}},
      {"h_ds_jet_mass", ";m_{jet}^{ch} (GeV/#it{c}^{2});entries", {HistType::kTH1F, {{200, 0., 50.}}}},
      {"h_ds_jet_lambda11", ";#lambda_{1}^{1};entries", {HistType::kTH1F, {{200, 0., 1.0}}}},
      {"h_ds_jet_lambda12", ";#lambda_{2}^{1} (thrust);entries", {HistType::kTH1F, {{200, 0., 1.0}}}},
      {"h_ds_jet_girth", ";g (#equiv #lambda_{1}^{1}R);entries", {HistType::kTH1F, {{200, 0., 0.5}}}},
      {"h2_dsjet_pt_lambda11", ";#it{p}_{T,jet} (GeV/#it{c});#lambda_{1}^{1}", {HistType::kTH2F, {{100, 0., 100.}, {200, 0., 1.0}}}},
      {"h2_dsjet_pt_lambda12", ";#it{p}_{T,jet} (GeV/#it{c});#lambda_{2}^{1}", {HistType::kTH2F, {{100, 0., 100.}, {200, 0., 1.0}}}},
      {"h2_dsjet_pt_mass", ";#it{p}_{T,jet} (GeV/#it{c});m_{jet}^{ch} (GeV/#it{c}^{2})", {HistType::kTH2F, {{100, 0., 100.}, {200, 0., 50.0}}}},
      {"h2_dsjet_pt_girth", ";#it{p}_{T,jet} (GeV/#it{c});g", {HistType::kTH2F, {{100, 0., 100.}, {200, 0., 0.5}}}},
    }};

  Configurable<float> vertexZCut{"vertexZCut", 10.0f, "Accepted z-vertex range"};
  Configurable<float> jetPtMin{"jetPtMin", 5.0, "minimum jet pT cut"};
  Configurable<float> jetR{"jetR", 0.4, "jet resolution parameter"};

  Configurable<std::string> eventSelections{"eventSelections", "sel8", "choose event selection"};
  Configurable<std::string> trackSelections{"trackSelections", "globalTracks", "set track selections"};

  std::vector<int> eventSelectionBits;
  int trackSelection = -1;

  Produces<aod::JetDistanceTable> distJetTable;

  template <typename JET, typename TRACKS>
  float computeLambda(JET const& jet, TRACKS const& tracks, float a, float k)
  {
    if (jet.pt() <= 0.f) {
      return -1.f;
    }
    float sum = 0.f;
    for (auto const& trk : tracks) {
      const float dr = jetutilities::deltaR(jet, trk);
      sum += std::pow(trk.pt(), k) * std::pow(dr, a);
    }
    const float R = jet.r() / 100.f;
    const float denom = std::pow(jet.pt(), k) * std::pow(R, a);
    if (denom <= 0.f) {
      return -1.f;
    }
    return sum / denom;
  }

  template <typename TRACKS>
  float computeJetMassFromTracksMass(TRACKS const& tracks)
  {
    double sumPx = 0.0, sumPy = 0.0, sumPz = 0.0, sumE = 0.0;

    for (auto const& trk : tracks) {
      const double pt = trk.pt();
      const double phi = trk.phi();
      const double eta = trk.eta();

      const double px = pt * std::cos(phi);
      const double py = pt * std::sin(phi);
      const double pz = pt * std::sinh(eta);
      const double p = std::sqrt(px * px + py * py + pz * pz);

      sumPx += px;
      sumPy += py;
      sumPz += pz;
      sumE += p; // massless
    }

    const double m2 = sumE * sumE - (sumPx * sumPx + sumPy * sumPy + sumPz * sumPz);
    return (m2 > 0.0) ? static_cast<float>(std::sqrt(m2)) : 0.f;
  }

  void init(o2::framework::InitContext&)
  {
    eventSelectionBits = jetderiveddatautilities::initialiseEventSelectionBits(static_cast<std::string>(eventSelections));
    trackSelection = jetderiveddatautilities::initialiseTrackSelection(static_cast<std::string>(trackSelections));

    auto h = registry.get<TH1>(HIST("h_jet_counter"));
    h->GetXaxis()->SetBinLabel(1, "All jets");
    h->GetXaxis()->SetBinLabel(2, "Ds-tagged jets");
  }

  Filter jetCuts = aod::jet::pt > jetPtMin&& aod::jet::r == nround(jetR.node() * 100.0f);
  Filter collisionFilter = nabs(aod::jcollision::posZ) < vertexZCut;

  void processCollisions(aod::JetCollision const& collision, aod::JetTracks const& tracks)
  {
    registry.fill(HIST("h_collisions"), 0.5);
    if (!jetderiveddatautilities::selectCollision(collision, eventSelectionBits)) {
      return;
    }

    registry.fill(HIST("h_collisions"), 1.5);

    for (auto const& track : tracks) {
      if (!jetderiveddatautilities::selectTrack(track, trackSelection)) {
        continue;
      }
      registry.fill(HIST("h_track_pt"), track.pt());
      registry.fill(HIST("h_track_eta"), track.eta());
      registry.fill(HIST("h_track_phi"), track.phi());
    }
  }
  PROCESS_SWITCH(JetDsSpecSubs, processCollisions, "process JE collisions", false);

  void processDataCharged(soa::Filtered<aod::JetCollisions>::iterator const& collision,
                          soa::Filtered<aod::ChargedJets> const& jets)
  {
    if (!jetderiveddatautilities::selectCollision(collision, eventSelectionBits)) {
      return;
    }
    for (const auto& jet : jets) {
      registry.fill(HIST("h_jet_pt"), jet.pt());
      registry.fill(HIST("h_jet_eta"), jet.eta());
      registry.fill(HIST("h_jet_phi"), jet.phi());
    }
  }
  PROCESS_SWITCH(JetDsSpecSubs, processDataCharged, "charged jets in data", false);

  void processDataChargedSubstructure(aod::JetCollision const& collision,
                                      soa::Join<aod::DsChargedJets, aod::DsChargedJetConstituents> const& jets,
                                      aod::CandidatesDsData const&,
                                      aod::JetTracks const&)
  {
    registry.fill(HIST("h_collision_counter"), 2.0);
    if (!jetderiveddatautilities::selectCollision(collision, eventSelectionBits) ||
        !(std::abs(collision.posZ()) < vertexZCut)) {
      return;
    }
    registry.fill(HIST("h_collision_counter"), 3.0);

    for (const auto& jet : jets) {

      registry.fill(HIST("h_jet_counter"), 0.5);

      bool hasDs = false;

      TVector3 jetVector(jet.px(), jet.py(), jet.pz());

      // Compute jet-level quantities once (independent of Ds candidates)
      auto jetTracks = jet.tracks_as<aod::JetTracks>();

      const float lambda11 = computeLambda(jet, jetTracks, 1.f, 1.f);
      const float lambda12 = computeLambda(jet, jetTracks, 2.f, 1.f);
      const float mjet = computeJetMassFromTracksMass(jetTracks);

      const float R = jet.r() / 100.f;
      const float girth = (lambda11 >= 0.f) ? (lambda11 * R) : -1.f;

      // Loop over Ds candidates (particle level)
      for (const auto& dsCandidate : jet.candidates_as<aod::CandidatesDsData>()) {

        hasDs = true;

        TVector3 dsVector(dsCandidate.px(), dsCandidate.py(), dsCandidate.pz());

        // zParallel defined as longitudinal momentum fraction along the jet axis
        const double zParallel = (jetVector * dsVector) / (jetVector * jetVector);
        const double axisDistance = jetutilities::deltaR(jet, dsCandidate);

        // --- Ds-level observables ---
        registry.fill(HIST("h_ds_jet_projection"), zParallel);
        registry.fill(HIST("h_ds_jet_distance_vs_projection"), axisDistance, zParallel);
        registry.fill(HIST("h_ds_jet_distance"), axisDistance);

        registry.fill(HIST("h_ds_mass"), dsCandidate.m());
        registry.fill(HIST("h_ds_eta"), dsCandidate.eta());
        registry.fill(HIST("h_ds_phi"), dsCandidate.phi());

        const float mass = dsCandidate.m();
        const float pt = dsCandidate.pt();
        const float z = zParallel;
        const float dR = axisDistance;

        // Main THnSparse: invariant mass, pT, z, and ΔR
        registry.fill(HIST("hSparse_ds"), mass, pt, z, dR);

        // --- output table ---
        auto scores = dsCandidate.mlScores();
        constexpr int kScore0 = 0;
        constexpr int kScore1 = 1;
        constexpr int kScore2 = 2;

        const float s0 = (scores.size() > kScore0) ? scores[kScore0] : -999.f;
        const float s1 = (scores.size() > kScore1) ? scores[kScore1] : -999.f;
        const float s2 = (scores.size() > kScore2) ? scores[kScore2] : -999.f;

        distJetTable(static_cast<float>(axisDistance),
                     jet.pt(), jet.eta(), jet.phi(),
                     static_cast<int>(jetTracks.size()),
                     dsCandidate.pt(), dsCandidate.eta(), dsCandidate.phi(),
                     dsCandidate.m(), dsCandidate.y(),
                     s0, s1, s2,
                     mjet, girth, lambda12, lambda11);
      }

      // Jet-level quantities (filled once per jet containing at least one Ds)
      if (hasDs) {

        registry.fill(HIST("h_jet_counter"), 1.5);

        // Jet properties
        registry.fill(HIST("h_ds_jet_pt"), jet.pt());
        registry.fill(HIST("h_ds_jet_eta"), jet.eta());
        registry.fill(HIST("h_ds_jet_phi"), jet.phi());

        // Jet substructure observables
        if (lambda11 >= 0.f) {
          registry.fill(HIST("h_ds_jet_lambda11"), lambda11);
          registry.fill(HIST("h2_dsjet_pt_lambda11"), jet.pt(), lambda11);
        }

        if (lambda12 >= 0.f) {
          registry.fill(HIST("h_ds_jet_lambda12"), lambda12);
          registry.fill(HIST("h2_dsjet_pt_lambda12"), jet.pt(), lambda12);
        }

        registry.fill(HIST("h_ds_jet_mass"), mjet);
        registry.fill(HIST("h2_dsjet_pt_mass"), jet.pt(), mjet);

        if (girth >= 0.f) {
          registry.fill(HIST("h_ds_jet_girth"), girth);
          registry.fill(HIST("h2_dsjet_pt_girth"), jet.pt(), girth);
        }
      }
    }
  }
  PROCESS_SWITCH(JetDsSpecSubs, processDataChargedSubstructure, "charged HF jet substructure", false);
};
WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
{
  return WorkflowSpec{adaptAnalysisTask<JetDsSpecSubs>(cfgc)};
}