particles_container.hpp

#ifndef __PARTICLES_CONTAINER_HPP__
#define __PARTICLES_CONTAINER_HPP__
 
#include "Kokkos_Macros.hpp"
#include <Kokkos_Core.hpp>
#include <biocma_cst_config.hpp>
#include <cmath>
#include <common/common.hpp>
#include <common/env_var.hpp>
#include <common/has_serialize.hpp>
#include <cstdint>
#include <mc/alias.hpp>
#include <mc/prng/prng.hpp>
#include <mc/traits.hpp>
 
#include <Kokkos_Sort.hpp>
#include <sorting/Kokkos_BinSortPublicAPI.hpp>
#include <sorting/impl/Kokkos_CopyOpsForBinSortImpl.hpp>
#include <sorting/impl/Kokkos_SortByKeyImpl.hpp>
#include <stdexcept>
#include <utility>
namespace MC
{
 
  struct RuntimeParameters
  {
    uint64_t minimum_dead_particle_removal{};
    double buffer_ratio{};
    double allocation_factor = {};
    double shrink_ratio{};
    double dead_particle_ratio_threshold{};
 
    template <class Archive>
    void
    serialize(Archive& ar)
    {
      ar(minimum_dead_particle_removal,
         buffer_ratio,
         allocation_factor,
         shrink_ratio,
         dead_particle_ratio_threshold);
    }
  };

  template <ModelType Model> class ParticlesContainer
  {
  public:
    using UsedModel = Model;

    explicit ParticlesContainer(RuntimeParameters rt_param,
                                std::size_t n_particle,
                                std::size_t _n_samples);
    ParticlesContainer(); //=default;
    ParticlesContainer(const ParticlesContainer&) = default;
    ParticlesContainer(ParticlesContainer&&) = default;
    ParticlesContainer& operator=(const ParticlesContainer&) = default;
    ParticlesContainer& operator=(ParticlesContainer&&) = default;

    ~ParticlesContainer() = default;
 
    // NOLINTBEGIN(cppcoreguidelines-non-private-member-variables-in-classes)
    Model::SelfParticle model;
    Model::SelfContribs contribs;
 
    MC::ParticlePositions position;
    MC::ParticleStatus status;
    ParticleWeigths<typename Model::FloatType> weights;
    ParticleAges ages;
    // NOLINTEND(cppcoreguidelines-non-private-member-variables-in-classes)

    // template <typename CviewType>
    // KOKKOS_INLINE_FUNCTION void
    // get_contributions(const std::size_t idx,
    //                   const CviewType& contributions) const
    // {
    //   if (begin < end)
    //   {
    //     static_assert(ConstWeightModelType<Model>,
    //                   "ModelType: Constapply_weight()");
 
    //     const double weight = get_weight(idx);
    //     const auto pos = position(idx);
    //     auto access = contributions.access();
 
    //     for (int i = begin; i < end; ++i)
    //     {
    //       const int rel = i - begin;
    //       access(rel, pos) += weight * model(idx, i);
    //     }
    //   }
    // }

    [[nodiscard]] KOKKOS_INLINE_FUNCTION bool
    handle_division(const MC::pool_type& random_pool, std::size_t idx1) const;

    [[nodiscard]] KOKKOS_INLINE_FUNCTION Model::FloatType
    get_weight(std::size_t idx) const;
 
    // HOST

    void force_remove_dead();

    void merge_buffer();

    template <class Archive> void save(Archive& ar) const;

    template <class Archive> void load(Archive& ar);

    [[nodiscard]] double get_allocation_factor() const noexcept;

    [[nodiscard]] std::size_t capacity() const noexcept;

    [[nodiscard]] std::size_t get_inactive() const noexcept;

    void update_and_remove_inactive(std::size_t out, std::size_t dead);

    [[nodiscard]] KOKKOS_INLINE_FUNCTION std::size_t n_particles() const;
 
    void change_runtime(RuntimeParameters&& parameters) noexcept;
 
    void _sort(size_t n_c);

    void remove_inactive_particles(std::size_t to_remove);
 
// FIXME: used only in unit test
#ifndef NDEBUG
    [[maybe_unused]] [[nodiscard]] auto get_buffer_index() const;
#endif
 
  private:
    Model::SelfParticle buffer_model;
    ParticlePositions buffer_position;
    Kokkos::View<uint64_t, Kokkos::SharedSpace> buffer_index;
    std::size_t n_allocated_elements;
    uint64_t n_used_elements;
    std::size_t inactive_counter;
 
    void __allocate_buffer__();
    void _resize(std::size_t new_size, bool force = false);
    RuntimeParameters rt_params;
    // FIXME
  public:
    // int begin;
    // int end;
 
    // int begin;
    // int end;
  };
 
} // namespace MC
 
/*
 * IMPLEMENTATION
 */
namespace MC
{
  namespace
  {
 
    using TeamPolicy = Kokkos::TeamPolicy<ComputeSpace>;
    using TeamMember = TeamPolicy::member_type;

    template <ModelType M> struct CompactParticlesFunctor
    {
 
      CompactParticlesFunctor(MC::ParticleStatus _status,
                              M::SelfParticle _model,
                              M::SelfContribs _contribs,
                              MC::ParticlePositions _position,
                              MC::ParticleAges _ages,
                              std::size_t _to_remove,
                              std::size_t _last_used_index)
          : status(std::move(_status)), model(std::move(_model)),
            contribs(std::move(_contribs)), position(std::move(_position)),
            ages(std::move(_ages)), offset("offset"), to_remove(_to_remove),
            last_used_index(_last_used_index)
      {
 
        Kokkos::deep_copy(offset, 0);
      }
 
      KOKKOS_INLINE_FUNCTION void
      operator()(const int i, std::size_t& update, const bool final) const
      {

 
        // Check if the particle is non-idle.
        const bool is_inactive = (status(i) != MC::Status::Idle);
 
        // Capture the current prefix sum of inactive particles.
        const std::size_t scan_index = update;
 
        // Increment the shared counter if the particle is inactive.
        update += is_inactive ? 1 : 0;
 
        // Final pass: Compact the container by replacing inactive particles.
        if (final && is_inactive && scan_index < to_remove)
        {
 
          const auto inactive_slot = i; // Index of the "gap" to fill
 
          // Atomically find a replacement particle from the end of the
          // container. The replacement must be alive (idle) and not the same as
          // the particle being removed.
          // Index of the particle that will "fill" the inactive slot
          auto replacement_index
              = last_used_index - Kokkos::atomic_fetch_add(&offset(), 1);
          while (status(replacement_index) != MC::Status::Idle
                 || replacement_index
                        == static_cast<std::size_t>(inactive_slot))
          {
            replacement_index
                = last_used_index - Kokkos::atomic_fetch_add(&offset(), 1);
          }
 
          // Mark the removed particle's position as idle.
          status(inactive_slot) = MC::Status::Idle;
 
          // Merge position EZ
          // No need atomic: final executed exactly once
          Kokkos::atomic_exchange(&position(inactive_slot),
                                  position(replacement_index));
 
          // TODO Use hierachical parallism here, thread range is likely to work
          for (std::size_t i_properties = 0; i_properties < M::n_var;
               ++i_properties)
          {
            model(inactive_slot, i_properties)
                = model(replacement_index, i_properties);
          }
 
          for (std::size_t i_c = 0; i_c < M::n_c; ++i_c)
          {
            contribs(inactive_slot, i_c) = contribs(replacement_index, i_c);
          }
 
          ages(inactive_slot, 0) = ages(replacement_index, 0);
          ages(inactive_slot, 1) = ages(replacement_index, 1);
        }
      }
 
      MC::ParticleStatus status;
      M::SelfParticle model;
      M::SelfContribs contribs;
      MC::ParticlePositions position;
      MC::ParticleAges ages;
      Kokkos::View<std::size_t, ComputeSpace> offset;
      std::size_t to_remove;
      std::size_t last_used_index;
    };

    template <ModelType M> struct InsertFunctor
    {
      InsertFunctor(std::size_t _original_size,
                    M::SelfParticle _model,
                    MC::ParticlePositions _position,
                    MC::ParticleAges _ages,
                    M::SelfParticle _buffer_model,
                    MC::ParticlePositions _buffer_position)
          : original_size(_original_size), model(std::move(_model)),
            ages(std::move(_ages)), position(std::move(_position)),
            buffer_model(std::move(_buffer_model)),
            buffer_position(std::move(_buffer_position))
      {
      }
      KOKKOS_INLINE_FUNCTION
      void
      operator()(const TeamMember& team) const
      {
        auto range = M::n_var;
        const int i = team.league_rank();
 
        Kokkos::parallel_for(
            Kokkos::TeamVectorRange(team, range),
            [&](const int& j)
            { model(original_size + i, j) = buffer_model(i, j); });
        position(original_size + i) = buffer_position(i);
 
        // Actually needs buffer to store mother's hydraulic time
        // But set new hydraulic time to 0 to not create new buffer a save
        // memory usage
        ages(original_size + i, 0) = 0;
        ages(original_size + i, 1) = 0;
      }
 
      std::size_t original_size;
      M::SelfParticle model;
      MC::ParticleAges ages;
      MC::ParticlePositions position;
      M::SelfParticle buffer_model;
      MC::ParticlePositions buffer_position;
    };
 
  }; // namespace
 
  template <ModelType Model>
  [[nodiscard]] KOKKOS_INLINE_FUNCTION std::size_t
  ParticlesContainer<Model>::n_particles() const
  {
    return n_used_elements;
  };
 
  template <ModelType Model>
  [[nodiscard]] std::size_t
  ParticlesContainer<Model>::get_inactive() const noexcept
  {
    return inactive_counter;
  }
 
  template <ModelType Model>
  void
  ParticlesContainer<Model>::force_remove_dead()
  {
    this->remove_inactive_particles(this->inactive_counter);
  }
#ifndef NDEBUG
 
  template <ModelType Model>
  [[maybe_unused]] [[nodiscard]] auto
  ParticlesContainer<Model>::get_buffer_index() const
  {
    return buffer_index();
  }
#endif
 
  template <ModelType Model>
  [[nodiscard]] double
  ParticlesContainer<Model>::get_allocation_factor() const noexcept
  {
    return rt_params.allocation_factor;
  }
 
  template <ModelType Model>
  [[nodiscard]] std::size_t
  ParticlesContainer<Model>::capacity() const noexcept
  {
    return n_allocated_elements;
  }
 
  template <ModelType Model>
  template <class Archive>
  void
  ParticlesContainer<Model>::load(Archive& ar)
  {
    // Basically store everything that is usefull
    ar(n_allocated_elements, n_used_elements, rt_params);
 
    // Allocation is done in the deserialize
    deserialize_view(ar, weights);
    deserialize_view(ar, position);
    deserialize_view(ar, status);
    deserialize_view(ar, model);
    deserialize_view(ar, ages);
 
    if (Model::n_var != model.extent(1))
    {
      throw std::runtime_error(
          "Error when deserialze, model number of property mismatch");
    }
 
    Kokkos::resize(
        this->contribs,
        n_allocated_elements,
        Model::n_c); // Dont forget to allocate contribs which is not saved yet
#ifndef NDEBUG
    Kokkos::printf("ParticlesContainer::load: Check if load_tuning_constant "
                   "works with different value");
#endif
 
    __allocate_buffer__(); // Dont forget to allocate buffer
  }
 
  template <ModelType Model>
  template <class Archive>
  void
  ParticlesContainer<Model>::save(Archive& ar) const
  {
 
    // Basically store everything that is usefull
    ar(n_allocated_elements, n_used_elements, rt_params);
    serialize_view(ar, weights);
    serialize_view(ar, position);
    serialize_view(ar, status);
    serialize_view(ar, model);
    serialize_view(ar, ages);
  }
 
  template <ModelType Model>
  void
  ParticlesContainer<Model>::update_and_remove_inactive(const std::size_t out,
                                                        const std::size_t dead)
  {
    // Actually out particle and dead ones are just inactive
    inactive_counter += out;
    inactive_counter += dead;
 
    const auto _threshold = std::max(
        rt_params.minimum_dead_particle_removal,
        static_cast<uint64_t>(static_cast<double>(n_used_elements)
                              * rt_params.dead_particle_ratio_threshold));
 
    if (inactive_counter > _threshold)
    {
      this->remove_inactive_particles(inactive_counter);
    }
  }
 
  template <ModelType Model>
  KOKKOS_INLINE_FUNCTION bool
  ParticlesContainer<Model>::handle_division(const MC::pool_type& random_pool,
                                             std::size_t idx1) const
  {
    if (Kokkos::atomic_load(&buffer_index()) < buffer_model.extent(0))
    {
      const auto idx2 = Kokkos::atomic_fetch_add(&buffer_index(), 1);
      Model::division(random_pool, idx1, idx2, model, buffer_model);
      buffer_position(idx2) = position(idx1);
      ages(idx1, 1) = 0;
      return true;
    }
    return false;
  }
 
  template <ModelType Model>
  void
  ParticlesContainer<Model>::merge_buffer()
  {
    PROFILE_SECTION("ParticlesContainer::merge_buffer")
    const auto original_size = n_used_elements;
    const auto n_add_item = buffer_index();
    if (n_add_item == 0)
    {
      return;
    }
    _resize(original_size + n_add_item);
    Kokkos::parallel_for("insert_merge",
                         TeamPolicy(n_add_item, Kokkos::AUTO, Model::n_var),
                         InsertFunctor<Model>(original_size,
                                              model,
                                              position,
                                              ages,
                                              buffer_model,
                                              buffer_position));
 
    buffer_index() = 0;
    n_used_elements += n_add_item;
    __allocate_buffer__();
  }
 
  template <ModelType Model>
  void
  ParticlesContainer<Model>::_resize(std::size_t new_size, bool force)
  {
    PROFILE_SECTION("ParticlesContainer::_resize")
 
    // Ensure new_size is greater than zero
    if (new_size > 0)
    {
      // Determine if resizing is necessary based on the condition
      if (new_size > n_allocated_elements || force)
      {
        // Calculate the new allocated size
        const auto new_allocated_size = static_cast<std::size_t>(std::ceil(
            static_cast<double>(new_size) * rt_params.allocation_factor));
 
        // Update the allocated size
        n_allocated_elements = new_allocated_size;
 
        // Perform the resizing on all relevant data containers
        Kokkos::resize(position, n_allocated_elements);
        Kokkos::resize(model,
                       n_allocated_elements,
                       Model::n_var); // use 2nd dim resize if dynamic
        Kokkos::resize(contribs,
                       n_allocated_elements,
                       Model::n_c); // use 2nd dim resize if dynamic
        Kokkos::resize(status, n_allocated_elements);
        Kokkos::resize(ages, n_allocated_elements);
 
        // Handle resizing for weights based on model type
        if constexpr (ConstWeightModelType<Model>)
        {
          Kokkos::resize(weights,
                         1); // Fixed size for ConstWeightModelType
        }
        else
        {
          Kokkos::resize(weights, n_allocated_elements);
        }
      }
    }
  }
 
  // template <ModelType Model>
  // void
  // ParticlesContainer<Model>::__allocate_buffer__()
  // {
  //   PROFILE_SECTION("ParticlesContainer::__allocate_buffer__")
 
  //   std::size_t buffer_size = buffer_position.extent(0);
 
  //   const double tmp = rt_params.buffer_ratio * n_allocated_elements;
 
  //   const std::size_t buffer_threshold = static_cast<std::size_t>(tmp);
 
  //   if (buffer_size < buffer_threshold)
  //   {
  //     buffer_size = static_cast<std::size_t>(
  //         std::ceil(rt_params.buffer_ratio * n_allocated_elements));
 
  //     // Realloc because not needed to keep buffer as it has been copied
  //     Kokkos::realloc(buffer_position, buffer_size);
  //     Kokkos::realloc(buffer_model, buffer_size, Model::n_var);
  //     buffer_index() = 0;
  //   }
  // }
 
  template <ModelType Model>
  void
  ParticlesContainer<Model>::__allocate_buffer__()
  {
    PROFILE_SECTION("ParticlesContainer::__allocate_buffer__")
 
    const auto required_buffer_size = static_cast<std::size_t>(
        std::ceil(rt_params.buffer_ratio * n_allocated_elements));
 
    if (buffer_position.extent(0) < required_buffer_size)
    {
      // Realloc because not needed to keep buffer as it has been copied
      Kokkos::realloc(buffer_position, required_buffer_size);
      Kokkos::realloc(buffer_model, required_buffer_size, Model::n_var);
      buffer_index() = 0;
    }
  }
 
// NOLINTBEGIN
#define alloc_without_init(name)                                               \
  Kokkos::view_alloc(Kokkos::WithoutInitializing, name)
  // NOLINTEND
 
  template <ModelType M>
  ParticlesContainer<M>::ParticlesContainer(
      RuntimeParameters rt_param,
      std::size_t n_particle,
      [[maybe_unused]] std::size_t _n_samples)
      : model(alloc_without_init("particle_model"), 0, 0),
 
        contribs(alloc_without_init("particle_contribs"), 0),
        position(alloc_without_init("particle_position"), 0),
        status(alloc_without_init("particle_status"), 0),
        weights(alloc_without_init("particle_weigth"), 0),
        ages(alloc_without_init("particle_age"), 0),
        buffer_model("buffer_particle_model", 0),
        buffer_position("buffer_particle_position", 0),
        buffer_index("buffer_index"), n_allocated_elements(0),
        n_used_elements(n_particle), inactive_counter(0), rt_params(rt_param)
  {
 
    // load_tuning_constant();
 
    if (n_particle != 0)
    {
      _resize(n_particle);
      __allocate_buffer__();
    }
 
    // const auto bounds = M::get_bounds();
 
    // if (begin > end)
    // {
    //   throw std::invalid_argument("Model begin should be > end");
    // }
 
    // begin = bounds.begin;
    // end = bounds.end;
  }
 
  template <ModelType M>
  ParticlesContainer<M>::ParticlesContainer()
      : ParticlesContainer(RuntimeParameters{}, 0, 0)
  {
  }
 
  template <ModelType M>
  void
  ParticlesContainer<M>::remove_inactive_particles(std::size_t to_remove)
  {
 
    PROFILE_SECTION("ParticlesContainer::remove_inactive_particles")
    if (to_remove == 0)
    {
      return;
    }
 
    if (to_remove == n_used_elements)
    {
      _resize(0, true);
      n_used_elements = 0;
      inactive_counter = 0;
    }
    else if (to_remove > n_used_elements)
    {
      throw std::runtime_error(
          "remove_inactive_particles: Error in kernel cannot remove more "
          "element than existing");
    }
    else
    {
 
      const auto new_used_item = n_used_elements - to_remove;
 
      const auto last_used_index = n_used_elements - 1;
      Kokkos::parallel_scan(
          "find_and_fill_gap",
          Kokkos::RangePolicy<ComputeSpace>(0, n_used_elements),
          CompactParticlesFunctor<M>(status,
                                     model,
                                     contribs,
                                     position,
                                     ages,
                                     to_remove,
                                     last_used_index));
 
      Kokkos::fence();
      KOKKOS_ASSERT(this->position.extent(0) == n_allocated_elements);
      KOKKOS_ASSERT(this->contribs.extent(0) == n_allocated_elements);
      KOKKOS_ASSERT(this->model.extent(0) == n_allocated_elements);
      KOKKOS_ASSERT(this->status.extent(0) == n_allocated_elements);
 
      n_used_elements = new_used_item;
      const bool do_shrink = n_used_elements <= static_cast<std::size_t>(
                                 rt_params.shrink_ratio * n_allocated_elements);
 
      if (do_shrink)
      {
        _resize(n_used_elements * rt_params.allocation_factor, true);
      }
      if (do_shrink)
      {
        // force to true if we want to shrink
        _resize(n_used_elements * rt_params.allocation_factor, true);
      }
      inactive_counter = inactive_counter - to_remove;
    };
  }
 
  template <ModelType M>
  [[nodiscard]] KOKKOS_INLINE_FUNCTION M::FloatType
  ParticlesContainer<M>::get_weight(const std::size_t idx) const
  {
    if constexpr (ConstWeightModelType<M>)
    {
      return weights(0);
    }
    else
    {
      return weights(idx);
    }
  }
 
  template <ModelType M>
  void
  ParticlesContainer<M>::change_runtime(RuntimeParameters&& parameters) noexcept
  {
    this->rt_params = parameters;
  }
 
  template <ModelType M>
  void
  ParticlesContainer<M>::_sort(std::size_t n_c)
  {
    (void)n_c;
 
    // PROFILE_SECTION("SORT")
    // // const auto N = n_used_elements;
    const int bin_size = 2048;
 
    const auto sn
        = Kokkos::subview(position, std::pair<int, int>(0, n_used_elements));
 
    using ExecSpace = Kokkos::DefaultExecutionSpace;
    using view_type = decltype(position);
    auto binop = Kokkos::BinOp1D<view_type>(bin_size, 0, n_c);
 
    auto sorter = Kokkos::BinSort<view_type, decltype(binop)>(
        ExecSpace(), sn, binop, true);
  }
 
} // namespace MC
 
#endif