NuclearInteraction.hpp

Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file NuclearInteraction.hpp

// This file is part of the Acts project.
//
// Copyright (C) 2020-2021 CERN for the benefit of the Acts project
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#pragma once

#include "Acts/Definitions/Algebra.hpp"
#include "Acts/Definitions/Common.hpp"
#include "Acts/Definitions/PdgParticle.hpp"
#include "Acts/Material/MaterialSlab.hpp"
#include "Acts/Utilities/UnitVectors.hpp"
#include "Acts/Utilities/VectorHelpers.hpp"
#include "ActsFatras/EventData/Barcode.hpp"
#include "ActsFatras/EventData/Particle.hpp"
#include "ActsFatras/EventData/ProcessType.hpp"
#include "ActsFatras/Physics/NuclearInteraction/NuclearInteractionParameters.hpp"

#include <any>
#include <cmath>
#include <iterator>
#include <limits>
#include <optional>
#include <random>
#include <utility>
#include <vector>

namespace ActsFatras {

struct NuclearInteraction {
  using Scalar = Particle::Scalar;
  detail::MultiParticleNuclearInteractionParametrisation
      multiParticleParameterisation;
  //~ unsigned int nMatchingTrials = std::numeric_limits<unsigned int>::max();
  unsigned int nMatchingTrials = 100;
  unsigned int nMatchingTrialsTotal = 1000;

  template <typename generator_t>
  std::pair<Scalar, Scalar> generatePathLimits(generator_t& generator,
                                               const Particle& particle) const {
    // Fast exit: No paramtrization provided
    if (multiParticleParameterisation.empty()) {
      return std::make_pair(std::numeric_limits<Scalar>::infinity(),
                            std::numeric_limits<Scalar>::infinity());
    }
    // Find the parametrisation that corresponds to the particle type
    for (const auto& particleParametrisation : multiParticleParameterisation) {
      if (particleParametrisation.first == particle.pdg()) {
        std::uniform_real_distribution<double> uniformDistribution{0., 1.};

        // Get the parameters
        const detail::NuclearInteractionParametrisation&
            singleParticleParametrisation = particleParametrisation.second;
        const detail::NuclearInteractionParameters& parametrisation =
            findParameters(uniformDistribution(generator),
                           singleParticleParametrisation,
                           particle.absoluteMomentum());

        // Set the L0 limit if not done already
        const auto& distribution =
            parametrisation.nuclearInteractionProbability;
        auto limits =
            std::make_pair(std::numeric_limits<Scalar>::infinity(),
                           sampleContinuousValues(
                               uniformDistribution(generator), distribution));
        return limits;
      }
    }
    return std::make_pair(std::numeric_limits<Scalar>::infinity(),
                          std::numeric_limits<Scalar>::infinity());
  }

  template <typename generator_t>
  bool run(generator_t& generator, Particle& particle,
           std::vector<Particle>& generated) const {
    // Fast exit: No paramtrization provided
    if (multiParticleParameterisation.empty()) {
      return false;
    }

    // Find the parametrisation that corresponds to the particle type
    for (const auto& particleParametrisation : multiParticleParameterisation) {
      if (particleParametrisation.first == particle.pdg()) {
        std::uniform_real_distribution<double> uniformDistribution{0., 1.};

        // Get the parameters
        const detail::NuclearInteractionParametrisation&
            singleParticleParametrisation = particleParametrisation.second;
        const detail::NuclearInteractionParameters& parametrisation =
            findParameters(uniformDistribution(generator),
                           singleParticleParametrisation,
                           particle.absoluteMomentum());

        std::normal_distribution<double> normalDistribution{0., 1.};
        // Dice the interaction type
        const bool interactSoft =
            softInteraction(normalDistribution(generator),
                            parametrisation.softInteractionProbability);

        // Get the final state multiplicity
        const unsigned int multiplicity = finalStateMultiplicity(
            uniformDistribution(generator),
            interactSoft ? parametrisation.softMultiplicity
                         : parametrisation.hardMultiplicity);

        // Get the parameters for the kinematics
        const std::vector<detail::NuclearInteractionParameters::
                              ParametersWithFixedMultiplicity>&
            parametrisationOfType =
                interactSoft ? parametrisation.softKinematicParameters
                             : parametrisation.hardKinematicParameters;
        const detail::NuclearInteractionParameters::
            ParametersWithFixedMultiplicity& parametrisationOfMultiplicity =
                parametrisationOfType[multiplicity];
        if (!parametrisationOfMultiplicity.validParametrisation) {
          return false;
        }

        // Get the kinematics
        const auto kinematics = sampleKinematics(
            generator, parametrisationOfMultiplicity, parametrisation.momentum);
        if (!kinematics.has_value()) {
          return run(generator, particle, generated);
        }

        // Get the particle types
        const std::vector<int> pdgIds =
            samplePdgIds(generator, parametrisation.pdgMap, multiplicity,
                         particle.pdg(), interactSoft);

        // Construct the particles
        const auto particles = convertParametersToParticles(
            generator, pdgIds, kinematics->first, kinematics->second, particle,
            parametrisation.momentum, interactSoft);

        // Kill the particle in a hard process
        if (!interactSoft) {
          particle.setAbsoluteMomentum(0);
        }

        generated.insert(generated.end(), particles.begin(), particles.end());
        return !interactSoft;
      }
    }
    // Fast exit if no parametrisation for the particle was provided
    return false;
  }

 private:
  const detail::NuclearInteractionParameters& findParameters(
      double rnd,
      const detail::NuclearInteractionParametrisation& parametrisation,
      float particleMomentum) const;

  inline bool softInteraction(double rnd, float probability) const {
    return rnd <= probability;
  }

  unsigned int finalStateMultiplicity(
      double rnd,
      const detail::NuclearInteractionParameters::CumulativeDistribution&
          distribution) const;

  template <typename generator_t>
  std::vector<int> samplePdgIds(
      generator_t& generator,
      const detail::NuclearInteractionParameters::PdgMap& pdgMap,
      unsigned int multiplicity, int particlePdg, bool soft) const;

  template <typename generator_t>
  Acts::ActsDynamicVector sampleInvariantMasses(
      generator_t& generator,
      const detail::NuclearInteractionParameters::
          ParametersWithFixedMultiplicity& parametrisation) const;

  template <typename generator_t>
  Acts::ActsDynamicVector sampleMomenta(
      generator_t& generator,
      const detail::NuclearInteractionParameters::
          ParametersWithFixedMultiplicity& parametrisation,
      float initialMomentum) const;

  bool match(const Acts::ActsDynamicVector& momenta,
             const Acts::ActsDynamicVector& invariantMasses,
             float parametrizedMomentum) const;

  template <typename generator_t>
  std::optional<std::pair<Acts::ActsDynamicVector, Acts::ActsDynamicVector>>
  sampleKinematics(generator_t& generator,
                   const detail::NuclearInteractionParameters::
                       ParametersWithFixedMultiplicity& parameters,
                   float momentum) const;

  std::pair<ActsFatras::Particle::Scalar, ActsFatras::Particle::Scalar>
  globalAngle(ActsFatras::Particle::Scalar phi1,
              ActsFatras::Particle::Scalar theta1, float phi2,
              float theta2) const;

  template <typename generator_t>
  std::vector<Particle> convertParametersToParticles(
      generator_t& generator, const std::vector<int>& pdgId,
      const Acts::ActsDynamicVector& momenta,
      const Acts::ActsDynamicVector& invariantMasses, Particle& initialParticle,
      float parametrizedMomentum, bool soft) const;

  unsigned int sampleDiscreteValues(
      double rnd,
      const detail::NuclearInteractionParameters::CumulativeDistribution&
          distribution) const;

  Scalar sampleContinuousValues(
      double rnd,
      const detail::NuclearInteractionParameters::CumulativeDistribution&
          distribution,
      bool interpolate = false) const;
};

template <typename generator_t>
std::vector<int> NuclearInteraction::samplePdgIds(
    generator_t& generator,
    const detail::NuclearInteractionParameters::PdgMap& pdgMap,
    unsigned int multiplicity, int particlePdg, bool soft) const {
  // Fast exit in case of no final state particles
  if (multiplicity == 0) {
    return {};
  }

  // The final state PDG IDs
  std::vector<int> pdgIds;
  pdgIds.reserve(multiplicity);

  std::uniform_real_distribution<float> uniformDistribution{0., 1.};

  // Find the producers probability distribution
  auto citProducer = pdgMap.cbegin();
  while (citProducer->first != particlePdg && citProducer != pdgMap.end()) {
    citProducer++;
  }

  const std::vector<std::pair<int, float>>& mapInitial = citProducer->second;
  // Set the first particle depending on the interaction type
  if (soft) {
    // Store the initial particle if the interaction is soft
    pdgIds.push_back(particlePdg);
  } else {
    // Otherwise dice the particle
    const float rndInitial = uniformDistribution(generator);

    pdgIds.push_back(
        std::lower_bound(mapInitial.begin(), mapInitial.end(), rndInitial,
                         [](const std::pair<int, float>& element,
                            float random) { return element.second < random; })
            ->first);
  }

  // Set the remaining particles
  for (unsigned int i = 1; i < multiplicity; i++) {
    // Find the producers probability distribution from the last produced
    // particle
    citProducer = pdgMap.cbegin();
    while (citProducer->first != pdgIds[i - 1] && citProducer != pdgMap.end()) {
      citProducer++;
    }

    // Set the next particle
    const std::vector<std::pair<int, float>>& map = citProducer->second;
    const float rnd = uniformDistribution(generator);
    pdgIds.push_back(
        std::lower_bound(map.begin(), map.end(), rnd,
                         [](const std::pair<int, float>& element,
                            float random) { return element.second < random; })
            ->first);
  }
  return pdgIds;
}

template <typename generator_t>
Acts::ActsDynamicVector NuclearInteraction::sampleInvariantMasses(
    generator_t& generator,
    const detail::NuclearInteractionParameters::ParametersWithFixedMultiplicity&
        parametrisation) const {
  // The resulting vector
  Acts::ActsDynamicVector parameters;
  const unsigned int size = parametrisation.eigenvaluesInvariantMass.size();
  parameters.resize(size);

  // Sample in the eigenspace
  for (unsigned int i = 0; i < size; i++) {
    float variance = parametrisation.eigenvaluesInvariantMass[i];
    std::normal_distribution<Acts::ActsScalar> dist{
        parametrisation.meanInvariantMass[i], sqrtf(variance)};
    parameters[i] = dist(generator);
  }
  // Transform to multivariate normal distribution
  parameters = parametrisation.eigenvectorsInvariantMass * parameters;

  // Perform the inverse sampling from the distributions
  for (unsigned int i = 0; i < size; i++) {
    const double cdf = (std::erff(parameters[i]) + 1) * 0.5;
    parameters[i] = sampleContinuousValues(
        cdf, parametrisation.invariantMassDistributions[i]);
  }
  return parameters;
}

template <typename generator_t>
Acts::ActsDynamicVector NuclearInteraction::sampleMomenta(
    generator_t& generator,
    const detail::NuclearInteractionParameters::ParametersWithFixedMultiplicity&
        parametrisation,
    float initialMomentum) const {
  // The resulting vector
  Acts::ActsDynamicVector parameters;
  const unsigned int size = parametrisation.eigenvaluesMomentum.size();
  parameters.resize(size);

  // Sample in the eigenspace
  for (unsigned int i = 0; i < size; i++) {
    float variance = parametrisation.eigenvaluesMomentum[i];
    std::normal_distribution<Acts::ActsScalar> dist{
        parametrisation.meanMomentum[i], sqrtf(variance)};
    parameters[i] = dist(generator);
  }

  // Transform to multivariate normal distribution
  parameters = parametrisation.eigenvectorsMomentum * parameters;

  // Perform the inverse sampling from the distributions
  for (unsigned int i = 0; i < size; i++) {
    const float cdf = (std::erff(parameters[i]) + 1) * 0.5;
    parameters[i] =
        sampleContinuousValues(cdf, parametrisation.momentumDistributions[i]);
  }

  // Scale the momenta
  Acts::ActsDynamicVector momenta = parameters.head(size - 1);
  const float sum = momenta.sum();
  const float scale = parameters.template tail<1>()(0, 0) / sum;
  momenta *= scale * initialMomentum;
  return momenta;
}

template <typename generator_t>
std::optional<std::pair<Acts::ActsDynamicVector, Acts::ActsDynamicVector>>
NuclearInteraction::sampleKinematics(
    generator_t& generator,
    const detail::NuclearInteractionParameters::ParametersWithFixedMultiplicity&
        parameters,
    float momentum) const {
  unsigned int trials = 0;
  Acts::ActsDynamicVector invariantMasses =
      sampleInvariantMasses(generator, parameters);
  Acts::ActsDynamicVector momenta =
      sampleMomenta(generator, parameters, momentum);
  // Repeat momentum evaluation until the parameters match
  while (!match(momenta, invariantMasses, momentum)) {
    if (trials == nMatchingTrialsTotal) {
      return std::nullopt;
    }
    // Re-sampole invariant masses if no fitting momenta were found
    if (trials++ % nMatchingTrials == 0) {
      invariantMasses = sampleInvariantMasses(generator, parameters);
    } else {
      momenta = sampleMomenta(generator, parameters, momentum);
    }
  }
  return std::make_pair(momenta, invariantMasses);
}

template <typename generator_t>
std::vector<Particle> NuclearInteraction::convertParametersToParticles(
    generator_t& generator, const std::vector<int>& pdgId,
    const Acts::ActsDynamicVector& momenta,
    const Acts::ActsDynamicVector& invariantMasses, Particle& initialParticle,
    float parametrizedMomentum, bool soft) const {
  std::uniform_real_distribution<double> uniformDistribution{0., 1.};
  const auto& initialDirection = initialParticle.direction();
  const double phi = Acts::VectorHelpers::phi(initialDirection);
  const double theta = Acts::VectorHelpers::theta(initialDirection);
  const unsigned int size = momenta.size();

  std::vector<Particle> result;
  result.reserve(size);

  // Build the particles
  for (unsigned int i = 0; i < size; i++) {
    const float momentum = momenta[i];
    const float invariantMass = invariantMasses[i];
    const float p1p2 = 2. * momentum * parametrizedMomentum;
    const float costheta = 1. - invariantMass * invariantMass / p1p2;

    const auto phiTheta =
        globalAngle(phi, theta, uniformDistribution(generator) * 2. * M_PI,
                    std::acos(costheta));
    const auto direction =
        Acts::makeDirectionFromPhiTheta(phiTheta.first, phiTheta.second);

    Particle p = Particle(initialParticle.particleId().makeDescendant(i),
                          static_cast<Acts::PdgParticle>(pdgId[i]));
    p.setProcess(ProcessType::eNuclearInteraction)
        .setPosition4(initialParticle.fourPosition())
        .setAbsoluteMomentum(momentum)
        .setDirection(direction)
        .setReferenceSurface(initialParticle.referenceSurface());

    // Store the particle
    if (i == 0 && soft) {
      initialParticle = p;
    } else {
      result.push_back(std::move(p));
    }
  }

  return result;
}
}  // namespace ActsFatras