BetheHeitlerApprox.hpp

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

// This file is part of the Acts project.
//
// Copyright (C) 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/TrackFitting/detail/GsfUtils.hpp"

#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <cstddef>
#include <fstream>
#include <mutex>
#include <random>
#include <stdexcept>
#include <string>
#include <tuple>

#include <boost/container/static_vector.hpp>

namespace Acts {

namespace detail {

struct GaussianComponent {
  ActsScalar weight = 0.0;
  ActsScalar mean = 0.0;
  ActsScalar var = 0.0;
};

inline void transformComponent(const GaussianComponent &cmp,
                               double &transformed_weight,
                               double &transformed_mean,
                               double &transformed_var) {
  const auto &[weight, mean, var] = cmp;

  transformed_weight = std::log(weight) - std::log(1 - weight);
  transformed_mean = std::log(mean) - std::log(1 - mean);
  transformed_var = std::log(var);
}

inline auto inverseTransformComponent(double transformed_weight,
                                      double transformed_mean,
                                      double transformed_var) {
  GaussianComponent cmp;
  cmp.weight = 1. / (1 + std::exp(-transformed_weight));
  cmp.mean = 1. / (1 + std::exp(-transformed_mean));
  cmp.var = std::exp(transformed_var);

  return cmp;
}

}  // namespace detail

struct BetheHeitlerApproxSingleCmp {
  constexpr auto numComponents() const { return 1; }

  constexpr bool validXOverX0(ActsScalar x) const {
    return x < 0.002;
    ;
  }

  static auto mixture(const ActsScalar x) {
    std::array<detail::GaussianComponent, 1> ret{};

    ret[0].weight = 1.0;

    const double c = x / std::log(2);
    ret[0].mean = std::pow(2, -c);
    ret[0].var = std::pow(3, -c) - std::pow(4, -c);

    return ret;
  }
};

template <int NComponents, int PolyDegree>
class AtlasBetheHeitlerApprox {
  static_assert(NComponents > 0);
  static_assert(PolyDegree > 0);

 public:
  struct PolyData {
    std::array<ActsScalar, PolyDegree + 1> weightCoeffs;
    std::array<ActsScalar, PolyDegree + 1> meanCoeffs;
    std::array<ActsScalar, PolyDegree + 1> varCoeffs;
  };

  using Data = std::array<PolyData, NComponents>;

  constexpr static double noChangeLimit = 0.0001;
  constexpr static double singleGaussianLimit = 0.002;
  constexpr static double lowerLimit = 0.10;
  constexpr static double higherLimit = 0.20;

 private:
  Data m_low_data;
  Data m_high_data;
  bool m_low_transform;
  bool m_high_transform;

 public:
  constexpr AtlasBetheHeitlerApprox(const Data &low_data, const Data &high_data,
                                    bool low_transform, bool high_transform)
      : m_low_data(low_data),
        m_high_data(high_data),
        m_low_transform(low_transform),
        m_high_transform(high_transform) {}

  constexpr auto numComponents() const { return NComponents; }

  constexpr bool validXOverX0(ActsScalar x) const { return x < higherLimit; }

  auto mixture(ActsScalar x) const {
    using Array =
        boost::container::static_vector<detail::GaussianComponent, NComponents>;
    // Build a polynom
    auto poly = [](ActsScalar xx,
                   const std::array<ActsScalar, PolyDegree + 1> &coeffs) {
      ActsScalar sum{0.};
      for (const auto c : coeffs) {
        sum = xx * sum + c;
      }
      assert((std::isfinite(sum) && "polynom result not finite"));
      return sum;
    };

    // Lambda which builds the components
    auto make_mixture = [&](const Data &data, double xx, bool transform) {
      // Value initialization should garanuee that all is initialized to zero
      Array ret(NComponents);
      ActsScalar weight_sum = 0;
      for (int i = 0; i < NComponents; ++i) {
        // These transformations must be applied to the data according to ATHENA
        // (TrkGaussianSumFilter/src/GsfCombinedMaterialEffects.cxx:79)
        if (transform) {
          ret[i] = detail::inverseTransformComponent(
              poly(xx, data[i].weightCoeffs), poly(xx, data[i].meanCoeffs),
              poly(xx, data[i].varCoeffs));
        } else {
          ret[i].weight = poly(xx, data[i].weightCoeffs);
          ret[i].mean = poly(xx, data[i].meanCoeffs);
          ret[i].var = poly(xx, data[i].varCoeffs);
        }

        weight_sum += ret[i].weight;
      }

      for (int i = 0; i < NComponents; ++i) {
        ret[i].weight /= weight_sum;
      }

      return ret;
    };

    // Return no change
    if (x < noChangeLimit) {
      Array ret(1);

      ret[0].weight = 1.0;
      ret[0].mean = 1.0;  // p_initial = p_final
      ret[0].var = 0.0;

      return ret;
    }
    // Return single gaussian approximation
    if (x < singleGaussianLimit) {
      Array ret(1);
      ret[0] = BetheHeitlerApproxSingleCmp::mixture(x)[0];
      return ret;
    }
    // Return a component representation for lower x0
    if (x < lowerLimit) {
      return make_mixture(m_low_data, x, m_low_transform);
    }
    // Return a component representation for higher x0
    // Cap the x because beyond the parameterization goes wild
    const auto high_x = std::min(higherLimit, x);
    return make_mixture(m_high_data, high_x, m_high_transform);
  }

  static auto loadFromFiles(const std::string &low_parameters_path,
                            const std::string &high_parameters_path) {
    auto read_file = [](const std::string &filepath) {
      std::ifstream file(filepath);

      if (!file) {
        throw std::invalid_argument("Could not open '" + filepath + "'");
      }

      std::size_t n_cmps = 0, degree = 0;
      bool transform_code = false;

      file >> n_cmps >> degree >> transform_code;

      if (NComponents != n_cmps) {
        throw std::invalid_argument("Wrong number of components in '" +
                                    filepath + "'");
      }

      if (PolyDegree != degree) {
        throw std::invalid_argument("Wrong polynom order in '" + filepath +
                                    "'");
      }

      Data data;

      for (auto &cmp : data) {
        for (auto &coeff : cmp.weightCoeffs) {
          file >> coeff;
        }
        for (auto &coeff : cmp.meanCoeffs) {
          file >> coeff;
        }
        for (auto &coeff : cmp.varCoeffs) {
          file >> coeff;
        }
      }

      return std::make_tuple(data, transform_code);
    };

    const auto [low_data, low_transform] = read_file(low_parameters_path);
    const auto [high_data, high_transform] = read_file(high_parameters_path);

    return AtlasBetheHeitlerApprox(low_data, high_data, low_transform,
                                   high_transform);
  }
};

AtlasBetheHeitlerApprox<6, 5> makeDefaultBetheHeitlerApprox();

}  // namespace Acts