KfField.h

Go to the documentation of this file.

/* Copyright (C) 2024 GSI Helmholtzzentrum fuer Schwerionenforschung, Darmstadt
   SPDX-License-Identifier: GPL-3.0-only
   Authors: Sergei Zharko [committer] */

 
// FIXME: naming fFieldVariant - FieldVariant_t
 
#pragma once
 
#include "KfDefs.h"
#include "KfFieldRegion.h"
#include "KfFieldSlice.h"
#include "KfVisit.h"
 
#include <boost/serialization/access.hpp>
#include <boost/serialization/split_free.hpp>
 
#include <array>
#include <optional>
#include <set>
 
// TODO: Treat field at primary vertex as another field layer (field slice)
 
namespace cbm::algo::kf
{
  namespace detail
  {
    template<typename T, EFieldMode FldMode>
    class alignas(VcMemAlign) ConcreteField {
    };
 

    template<typename T>
    class alignas(VcMemAlign) ConcreteField<T, EFieldMode::Original> {
      template<typename, EFieldMode>
      friend class ConcreteField;
 
     public:
      ConcreteField() = default;

      ConcreteField(FieldFn_t fieldFn, std::vector<T>& fieldSliceZ, FieldRegion<T>& vtxField, EFieldType type)
        : fvFieldSliceZ(std::move(fieldSliceZ))
        , fFieldFn(fieldFn)
        , fPrimVertexField(std::move(vtxField))
        , fFieldType(type)
      {
      }

      template<typename I>
      ConcreteField(const ConcreteField<I, EFieldMode::Original>& other)
        : fFieldFn(other.fFieldFn)
        , fPrimVertexField(std::move(FieldRegion<T>::CopyConvert(other.fPrimVertexField)))
        , fFieldType(other.fFieldType)
      {
        fvFieldSliceZ.reserve(other.fvFieldSliceZ.size());
        for (const auto& slice : other.fvFieldSliceZ) {
          fvFieldSliceZ.emplace_back(utils::simd::Cast<I, T>(slice));
        }
      }

      ConcreteField(const ConcreteField&) = default;

      ConcreteField(ConcreteField&&) = default;

      ConcreteField& operator=(const ConcreteField&) = default;

      ConcreteField& operator=(ConcreteField&&) = default;

      const FieldFn_t& GetFieldFunction() const { return fFieldFn; }

      EFieldType GetFieldType() const { return fFieldType; }

      FieldValue<T> GetFieldValue(int sliceId, T x, T y) const
      {
        return GlobalField::GetFieldValue(fFieldFn, x, y, fvFieldSliceZ[sliceId]);
      }

      const FieldValue<T> GetFieldValueForLine(int sliceId, const TrackParam<T>& tp) const
      {
        T dz = fvFieldSliceZ[sliceId] - tp.GetZ();
        return GetFieldValue(sliceId, tp.GetX() + tp.GetTx() * dz, tp.GetY() + tp.GetTy() * dz);
      }

      const FieldRegion<T>& GetPrimVertexField() const { return fPrimVertexField; }

      int GetNofFieldSlices() const { return fvFieldSliceZ.size(); }

      void RemoveSlice(int iLayer) { fvFieldSliceZ.erase(fvFieldSliceZ.begin() + iLayer); }

      std::string ToString(int indentLevel, int verbose) const;

      static constexpr EFieldMode Mode() { return EFieldMode::Original; }
 
     private:
      std::vector<T> fvFieldSliceZ{};             
      FieldFn_t fFieldFn{defs::ZeroFieldFn};      
      FieldRegion<T> fPrimVertexField{};          
      EFieldType fFieldType{EFieldType::Normal};  
    };
 

    template<typename T>
    class alignas(VcMemAlign) ConcreteField<T, EFieldMode::Interpolated> {
      template<typename, EFieldMode>
      friend class ConcreteField;
 
     public:
      using SlicesContainer_t = std::vector<FieldSlice<T>>;

      ConcreteField() = default;

      ConcreteField(SlicesContainer_t& fieldSlices, FieldRegion<T>& vtxField, EFieldType type)
        : fvFieldSlices(std::move(fieldSlices))
        , fPrimVertexField(std::move(vtxField))
        , fFieldType(type)
      {
      }

      template<typename I>
      ConcreteField(const ConcreteField<I, EFieldMode::Interpolated>& other)
        : fPrimVertexField(std::move(FieldRegion<T>::CopyConvert(other.fPrimVertexField)))
        , fFieldType(other.fFieldType)
      {
        fvFieldSlices.reserve(other.fvFieldSlices.size());
        for (const auto& slice : other.fvFieldSlices) {
          fvFieldSlices.emplace_back(FieldSlice<T>(slice));
        }
      }

      ConcreteField(const ConcreteField&) = default;

      ConcreteField(ConcreteField&&) = default;

      ConcreteField& operator=(const ConcreteField&) = default;

      ConcreteField& operator=(ConcreteField&&) = default;

      const auto& GetFieldSlice(int sliceID) const { return fvFieldSlices[sliceID]; }

      EFieldType GetFieldType() const { return fFieldType; }

      FieldValue<T> GetFieldValue(int sliceID, T x, T y) const { return fvFieldSlices[sliceID].GetFieldValue(x, y); }

      const FieldValue<T> GetFieldValueForLine(int sliceID, const TrackParam<T>& tp) const
      {
        T dz = fvFieldSlices[sliceID].GetZref() - tp.GetZ();
        return GetFieldValue(sliceID, tp.GetX() + tp.GetTx() * dz, tp.GetY() + tp.GetTy() * dz);
      }

      const FieldRegion<T>& GetPrimVertexField() const { return fPrimVertexField; }

      int GetNofFieldSlices() const { return fvFieldSlices.size(); }

      void RemoveSlice(int iLayer) { fvFieldSlices.erase(fvFieldSlices.begin() + iLayer); }

      std::string ToString(int indentLevel, int verbose) const;

      static constexpr EFieldMode Mode() { return EFieldMode::Interpolated; }
 
     private:
      SlicesContainer_t fvFieldSlices{};          
      FieldRegion<T> fPrimVertexField{};          
      EFieldType fFieldType{EFieldType::Normal};  

      friend class boost::serialization::access;
      template<class Archive>
      void serialize(Archive& ar, const unsigned int)
      {
        ar& fvFieldSlices;
        ar& fPrimVertexField;
        ar& fFieldType;
      }
    };
  }  // namespace detail
 

  template<typename T>
  class alignas(VcMemAlign) Field {
    template<typename>
    friend class Field;
    friend class boost::serialization::access;
 
   public:
    template<EFieldMode Mode>
    using ConcreteField_t = typename detail::ConcreteField<T, Mode>;
 
    using Interpolated_t    = ConcreteField_t<EFieldMode::Interpolated>;
    using Original_t        = ConcreteField_t<EFieldMode::Original>;
    using FieldVariant_t    = typename std::variant<Interpolated_t, Original_t>;
    using SlicesContainer_t = typename Interpolated_t::SlicesContainer_t;

    Field() = default;

    Field(const Field&) = default;

    Field(Field&&) = default;

    ~Field() = default;

    Field& operator=(const Field& other) = default;

    Field& operator=(Field&& other) = default;

    template<EFieldMode FieldMode, typename... Args>
    static constexpr Field Create(Args... args)
    {
      return Field(args..., Tag<FieldMode>{});
    }

    template<typename I>
    static Field CopyConvert(const Field<I>& other)
    {
      return kf::visit([&other](const auto& c) { return Field(other, Tag<std::decay_t<decltype(c)>::Mode()>{}); },
                       other.fFieldVariant);
    }

    template<EFieldMode FieldMode>
    const ConcreteField_t<FieldMode>* GetConcrete() const
    {
      return std::get_if<ConcreteField_t<FieldMode>>(&fFieldVariant);
    }

    EFieldMode GetFieldMode() const
    {
      return kf::visit([](const auto& c) { return std::decay_t<decltype(c)>::Mode(); }, fFieldVariant);
    }

    FieldValue<T> GetFieldValue(int sliceId, T x, T y) const
    {
      return kf::visit([sliceId, x, y](const auto& c) -> FieldValue<T> { return c.GetFieldValue(sliceId, x, y); },
                       fFieldVariant);
    }

    const FieldValue<T> GetFieldValueForLine(int sliceId, const TrackParam<T>& tp) const
    {
      return kf::visit([sliceId, &tp](const auto& c) -> FieldValue<T> { return c.GetFieldValueForLine(sliceId, tp); },
                       fFieldVariant);
    }

    EFieldType GetFieldType() const
    {
      return kf::visit([](const auto& c) -> EFieldType { return c.GetFieldType(); }, fFieldVariant);
    }

    //FieldRegion<T> MakeFieldRegion(const Node& n0, const Node& n1, const Node& n2)
    //{
    //}

    // TODO:  rename GetFieldRegion -> MakeFieldRegion
    // FIXME: use kf::visit
    FieldRegion<T> GetFieldRegion(const FieldValue<T>& b0, const FieldValue<T>& b1, const FieldValue<T>& b2) const
    {
      if (auto* held = std::get_if<Original_t>(&fFieldVariant)) {
        // FIXME: FieldRegion constructor
        return FieldRegion<T>(held->GetFieldType(), held->GetFieldFunction());
      }
      else {
        return FieldRegion<T>(b0, b1, b2);
      }
    }

    const FieldRegion<T>& GetPrimVertexField() const
    {
      return kf::visit([](const auto& c) -> const FieldRegion<T>& { return c.GetPrimVertexField(); }, fFieldVariant);
    }
 

    void RemoveSlice(int iLayer)
    {
      kf::visit([iLayer](auto&& c) -> void { c.RemoveSlice(iLayer); }, fFieldVariant);
    }

    std::string ToString(int indentLevel, int verbose) const;
 
   private:
    Field(SlicesContainer_t& fieldSlices, FieldRegion<T>& vtxField, EFieldType fieldType, Tag<EFieldMode::Interpolated>)
      : fFieldVariant(FieldVariant_t(std::in_place_type<Interpolated_t>, fieldSlices, vtxField, fieldType))
    {
    }

    Field(FieldFn_t fieldFn, std::vector<T>& fieldSliceZ, FieldRegion<T>& vtxField, EFieldType fieldType,
          Tag<EFieldMode::Original>)
      : fFieldVariant(FieldVariant_t(std::in_place_type<Original_t>, fieldFn, fieldSliceZ, vtxField, fieldType))
    {
    }

    template<typename I, EFieldMode FieldMode>
    Field(const Field<I>& other, Tag<FieldMode>)
      : fFieldVariant(
        FieldVariant_t(std::in_place_type<ConcreteField_t<FieldMode>>,
                       ConcreteField_t<FieldMode>(std::get<detail::ConcreteField<I, FieldMode>>(other.fFieldVariant))))
    {
    }
 
    // TODO: Fully re-organize serialization

    template<class Archive>
    void load(Archive& ar, const unsigned int /*version*/)
    {
      auto field = Interpolated_t{};
      ar >> field;
      fFieldVariant.template emplace<Interpolated_t>(std::move(field));
    }

    // FIXME: use kf::variant
    template<class Archive>
    void save(Archive& ar, const unsigned int /*version*/) const
    {
      // FIXME (?) Which type should be serialized? Only double? Or others can be serialized as well?
      if (std::holds_alternative<Interpolated_t>(fFieldVariant)) {
        ar << std::get<Interpolated_t>(fFieldVariant);
      }
      else {
        throw std::logic_error("Attempt to serialize a kf::Field object with fFieldMode == EFieldMode::Original");
      }
    }

    template<class Archive>
    void serialize(Archive& ar, const unsigned int version)
    {
      boost::serialization::split_member(ar, *this, version);
    }
 
    FieldVariant_t fFieldVariant;  
  };
 

  class FieldBuilder {
    struct SliceRef {
      double fHalfSizeX{defs::Undef<double>};  
      double fHalfSizeY{defs::Undef<double>};  
      double fRefZ{defs::Undef<double>};       

      SliceRef(double halfX, double halfY, double refZ) : fHalfSizeX(halfX), fHalfSizeY(halfY), fRefZ(refZ) {}

      bool operator<(const SliceRef& r) const { return fRefZ < r.fRefZ; }
    };
 
   public:
    FieldBuilder() = default;

    FieldBuilder(const FieldBuilder&) = default;

    ~FieldBuilder() = default;

    FieldBuilder& operator=(const FieldBuilder&) = default;

    void AddSliceReference(double halfSizeX, double halfSizeY, double refZ);

    EFieldType GetFieldType() const { return fFieldType; }

    const FieldFn_t& GetFieldFunction() const { return fFieldFn; }

    template<typename T>
    Field<T> MakeField(EFieldMode fldMode) const;

    void Reset() { *this = FieldBuilder(); }

    void ResetSliceReferences() { fSliceReferences.clear(); }

    void SetFieldFunction(const FieldFn_t& fieldFn, EFieldType fldType)
    {
      fFieldFn   = fieldFn;
      fFieldType = fldType;
    }

    void SetStep(double step = 2.5) { fTargetStep = step; }

    void SetTarget(double x, double y, double z) { fTarget = {x, y, z}; }
 
   private:
    std::set<SliceRef> fSliceReferences;      
    FieldFn_t fFieldFn{defs::ZeroFieldFn};    
    double fTargetStep{2.5};                  
    EFieldType fFieldType{EFieldType::Null};  

    std::array<double, 3> fTarget{{defs::Undef<double>, defs::Undef<double>, defs::Undef<double>}};
  };
 
  // -------------------------------------------------------------------------------------------------------------------
  //
  template<typename T>
  Field<T> FieldBuilder::MakeField(EFieldMode fldMode) const
  {
    // Check initialization
    if (std::any_of(fTarget.begin(), fTarget.end(), [](double x) { return !utils::IsFinite(x); })) {
      throw std::logic_error("FieldBuilder::MakeField: target is undefined");
    }
    if (!fSliceReferences.size()) {  // TODO: Remove requirement of slice references
      throw std::logic_error("FieldBuilder::MakeField: no slice references were provided");
    }
    if (!fFieldFn) {
      throw std::logic_error("FieldBuilder::CreateField: no field function is provided");
    }
 
    // Initialize the Field object
    switch (fldMode) {
      case EFieldMode::Original: {
        std::vector<T> fieldSlices;
        fieldSlices.reserve(fSliceReferences.size());
        for (const auto& ref : fSliceReferences) {
          fieldSlices.emplace_back(ref.fRefZ);
        }
        // FIXME: rewrite signature (Field::CreateOriginal, Field::CreateInterpolated)
        return Field<T>::template Create<EFieldMode::Original>(fFieldFn, std::move(fieldSlices),
                                                               FieldRegion<T>(fFieldType, fFieldFn), fFieldType);
      } break;
      case EFieldMode::Interpolated:
      default: {
        typename Field<T>::SlicesContainer_t fieldSlices;
        fieldSlices.reserve(fSliceReferences.size());
        for (const auto& ref : fSliceReferences) {
          fieldSlices.emplace_back(fFieldFn, ref.fHalfSizeX, ref.fHalfSizeY, ref.fRefZ);
        }
 
        std::array<FieldValue<T>, 3> fld;
        {
          // PV field initialization
          double z = fTarget[2];
          for (auto& fldNode : fld) {
            fldNode = GlobalField::GetFieldValue<double>(fFieldFn, fTarget[0], fTarget[1], z);
            z += fTargetStep;
          }
        }
        return Field<T>::template Create<EFieldMode::Interpolated>(std::move(fieldSlices),
                                                                   FieldRegion<T>(fld[0], fld[1], fld[2]), fFieldType);
      } break;
    }
  }
}  // namespace cbm::algo::kf