femtoUtils.h

// Copyright 2019-2022 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.

/// \file femtoUtils.h
/// \brief Collision selection
/// \author Anton Riedel, TU München, anton.riedel@cern.ch

#ifndef PWGCF_FEMTO_CORE_FEMTOUTILS_H_
#define PWGCF_FEMTO_CORE_FEMTOUTILS_H_

#include <Common/Core/TableHelper.h>

#include <CommonConstants/PhysicsConstants.h>
#include <Framework/InitContext.h>
#include <Framework/Logger.h>

#include <TPDGCode.h>

#include <cmath>
#include <concepts>
#include <cstdint>
#include <optional>
#include <unordered_map>

namespace o2::analysis::femto
{
namespace utils
{

template <typename T1, typename T2>
inline std::optional<T2> getIndex(const T1& index, const std::unordered_map<T1, T2>& map)
{
  auto it = map.find(index);
  if (it != map.end()) {
    return it->second;
  }
  return std::nullopt;
}

template <typename T>
float itsSignal(T const& track)
{
  uint32_t clsizeflag = track.itsClusterSizes();
  auto clSizeLayer0 = (clsizeflag >> (0 * 4)) & 0xf;
  auto clSizeLayer1 = (clsizeflag >> (1 * 4)) & 0xf;
  auto clSizeLayer2 = (clsizeflag >> (2 * 4)) & 0xf;
  auto clSizeLayer3 = (clsizeflag >> (3 * 4)) & 0xf;
  auto clSizeLayer4 = (clsizeflag >> (4 * 4)) & 0xf;
  auto clSizeLayer5 = (clsizeflag >> (5 * 4)) & 0xf;
  auto clSizeLayer6 = (clsizeflag >> (6 * 4)) & 0xf;
  int numLayers = 7;
  int sumClusterSizes = clSizeLayer0 + clSizeLayer1 + clSizeLayer2 + clSizeLayer3 + clSizeLayer4 + clSizeLayer5 + clSizeLayer6;
  double cosLamnda = 1. / std::cosh(track.eta());
  double signal = (static_cast<double>(sumClusterSizes) / numLayers) * cosLamnda;
  return static_cast<float>(signal);
};

inline double getPdgMass(int pdgCode)
{
  // use this function instead of TDatabasePDG to return masses defined in the PhysicsConstants.h header
  // this approach saves a lot of memory and important partilces like deuteron are missing in TDatabasePDG anyway
  double mass = 0.f;
  // add new particles if necessary here
  switch (std::abs(pdgCode)) {
    case kPiPlus:
      mass = o2::constants::physics::MassPiPlus;
      break;
    case kKPlus:
      mass = o2::constants::physics::MassKPlus;
      break;
    case kProton:
      mass = o2::constants::physics::MassProton;
      break;
    case kLambda0:
      mass = o2::constants::physics::MassLambda;
      break;
    case o2::constants::physics::Pdg::kPhi:
      mass = o2::constants::physics::MassPhi;
      break;
    case kRho770_0:
      mass = 775.26; // not defined in O2?
      break;
    case kRho770Plus:
      mass = 775.11; // not defined in O2?
      break;
    case o2::constants::physics::Pdg::kK0Star892:
      mass = o2::constants::physics::MassK0Star892;
      break;
    case o2::constants::physics::Pdg::kLambdaCPlus:
      mass = o2::constants::physics::MassLambdaCPlus;
      break;
    case o2::constants::physics::Pdg::kDeuteron:
      mass = o2::constants::physics::MassDeuteron;
      break;
    case o2::constants::physics::Pdg::kTriton:
      mass = o2::constants::physics::MassTriton;
      break;
    case o2::constants::physics::Pdg::kHelium3:
      mass = o2::constants::physics::MassHelium3;
      break;
    case kSigmaMinus:
      mass = o2::constants::physics::MassSigmaMinus;
      break;
    case kSigmaPlus:
      mass = o2::constants::physics::MassSigmaPlus;
      break;
    case kXiMinus:
      mass = o2::constants::physics::MassXiMinus;
      break;
    case kOmegaMinus:
      mass = o2::constants::physics::MassOmegaMinus;
      break;
    default:
      LOG(warn) << "PDG code is not suppored. Return 0...";
  }
  return mass;
}

template <typename T>
float qn(T const& col)
{
  float qn = std::sqrt(col.qvecFT0CReVec()[0] * col.qvecFT0CReVec()[0] + col.qvecFT0CImVec()[0] * col.qvecFT0CImVec()[0]) * std::sqrt(col.sumAmplFT0C());
  return qn;
}

/// Recalculate pT for Kinks (Sigmas) using kinematic constraints
inline float calcPtnew(float pxMother, float pyMother, float pzMother, float pxDaughter, float pyDaughter, float pzDaughter)
{
  float almost0 = 1e-6f;
  // Particle masses in GeV/c^2
  auto massPion = o2::constants::physics::MassPionCharged;
  auto massNeutron = o2::constants::physics::MassNeutron;
  auto massSigmaMinus = o2::constants::physics::MassSigmaMinus;

  // Calculate mother momentum and direction versor
  float pMother = std::sqrt(pxMother * pxMother + pyMother * pyMother + pzMother * pzMother);
  if (pMother < almost0)
    return -999.f;

  float versorX = pxMother / pMother;
  float versorY = pyMother / pMother;
  float versorZ = pzMother / pMother;

  // Calculate daughter energy
  float ePi = std::sqrt(massPion * massPion + pxDaughter * pxDaughter + pyDaughter * pyDaughter + pzDaughter * pzDaughter);

  // Scalar product of versor with daughter momentum
  float scalarProduct = versorX * pxDaughter + versorY * pyDaughter + versorZ * pzDaughter;

  // Solve quadratic equation for momentum magnitude
  float k = massSigmaMinus * massSigmaMinus + massPion * massPion - massNeutron * massNeutron;
  float a = 4.f * (ePi * ePi - scalarProduct * scalarProduct);
  float b = -4.f * scalarProduct * k;
  float c = 4.f * ePi * ePi * massSigmaMinus * massSigmaMinus - k * k;

  if (std::abs(a) < almost0)
    return -999.f;

  float d = b * b - 4.f * a * c;
  if (d < 0.f)
    return -999.f;

  float sqrtD = std::sqrt(d);
  float p1 = (-b + sqrtD) / (2.f * a);
  float p2 = (-b - sqrtD) / (2.f * a);

  // Pick physical solution: prefer P2 if positive, otherwise P1
  if (p2 < 0.f && p1 < 0.f)
    return -999.f;
  if (p2 < 0.f)
    return p1;

  // Choose solution closest to original momentum
  float p1Diff = std::abs(p1 - pMother);
  float p2Diff = std::abs(p2 - pMother);
  float p = (p1Diff < p2Diff) ? p1 : p2;

  // Calculate pT from recalibrated momentum
  float pxS = versorX * p;
  float pyS = versorY * p;
  return std::sqrt(pxS * pxS + pyS * pyS);
}

inline bool enableTable(const char* tableName, int userSetting, o2::framework::InitContext& initContext)
{
  if (userSetting == 1) {
    LOG(info) << "Enabled femto table (forced on): " << tableName;
    return true;
  }
  if (userSetting == 0) {
    LOG(info) << "Disabled femto table (forced off): " << tableName;
    return false;
  }
  bool required = o2::common::core::isTableRequiredInWorkflow(initContext, tableName);
  if (required) {
    LOG(info) << "Enabled femto table (auto): " << tableName;
  }
  return required;
}

// template <typename T>
// using HasMass = decltype(std::declval<T&>().mass());
//
// template <typename T>
// using HasSign = decltype(std::declval<T&>().sign());

template <typename T>
concept HasMass = requires(T t) {
  { t.mass() } -> std::convertible_to<float>; // or double, whatever mass() returns
};

template <typename T>
inline int signum(T x)
{
  return (T(0) < x) - (x < T(0));
}

}; // namespace utils
}; // namespace o2::analysis::femto
#endif // PWGCF_FEMTO_CORE_FEMTOUTILS_H_