move_kernel.hpp

#ifndef __SIMULATION_MOVE_KERNEL_HPP__
#define __SIMULATION_MOVE_KERNEL_HPP__
 
#include <Kokkos_Assert.hpp>
#include <Kokkos_Core.hpp>
#include <Kokkos_Printf.hpp>
#include <Kokkos_Random.hpp>
#include <biocma_cst_config.hpp>
#include <cassert>
#include <mc/alias.hpp>
#include <mc/domain.hpp>
#include <mc/events.hpp>
#include <mc/prng/prng.hpp>
#include <mc/traits.hpp>
#include <simulation/alias.hpp>
#include <simulation/probability_leaving.hpp>
#include <simulation/probe.hpp>
#include <utility>
 
namespace Simulation::KernelInline
{
  constexpr bool enable_leave = true;
  constexpr bool disable_leave = false;
  constexpr bool disable_move = false;
  constexpr bool enable_move = true;
 
  template <typename ExecSpace> struct MoveInfo
  {
    ConstNeighborsView<ExecSpace> neighbors;
    DiagonalView<ExecSpace> diag_transition;
    CumulativeProbabilityView<ExecSpace> cumulative_probability;
    LeavingFlowType leaving_flow;
    LeavingFlowIndexType index_leaving_flow;
    Kokkos::View<double*, ExecSpace> liquid_volume;
 
    MoveInfo()
        : diag_transition("diag_transition", 0),
          leaving_flow("leaving_flow", 0),
          index_leaving_flow("index_leaving_flow", 0),
          liquid_volume("liquid_volume", 0)
    {
    }
 
    MoveInfo(const std::size_t n_compartments, const std::size_t n_flows)
        : diag_transition("diag_transition", n_compartments),
          leaving_flow("leaving_flow", n_flows),
          index_leaving_flow("index_leaving_flow", n_flows),
          liquid_volume("liquid_volume", n_compartments)
 
    {
    }
  };
 
  KOKKOS_INLINE_FUNCTION std::size_t __find_next_compartment(
      const ConstNeighborsView<ComputeSpace>& neighbors,
      const CumulativeProbabilityView<ComputeSpace>& cumulative_probability,
      const std::size_t i_compartment,
      const double random_number)
  {
    const int max_neighbor = static_cast<int>(neighbors.extent(1));
 
    std::size_t next =
        neighbors(i_compartment, 0); // Default to the first neighbor
 
    // Iterate through the neighbors to find the appropriate next compartment
    for (int k_neighbor = 0; k_neighbor < max_neighbor - 1; ++k_neighbor)
    {
 
      // Get the cumulative probability range for the current neighbor
      const auto pi = cumulative_probability(i_compartment, k_neighbor);
      const auto pn = cumulative_probability(i_compartment, k_neighbor + 1);
 
      // Use of a Condition mask to avoid branching.
      next = (random_number <= pn && pi <= random_number)
                 ? neighbors(i_compartment, k_neighbor + 1)
                 : next;
    }
 
    return next; // Return the index of the chosen next compartment
  }
 
  struct MoveFunctor
  {
    MoveFunctor(MC::ParticlePositions p,
                MC::ParticleStatus _status,
                MoveInfo<ComputeSpace> m,
                MC::KPRNG::pool_type _random_pool,
                MC::EventContainer _events,
                ProbeAutogeneratedBuffer _probes,
                MC::ParticleAges _ages)
        : d_t(0.), positions(std::move(p)), n_particles(0), move(std::move(m)),
          random_pool(_random_pool), status(std::move(_status)),
          events(std::move(_events)), probes(std::move(_probes)),
          ages(std::move(_ages)), enable_move(false), enable_leave(false) {};
 
    void update(double _d_t,
                std::size_t n_p,
                MoveInfo<ComputeSpace> move_i,
                MC::ParticlePositions _positions,
                MC::ParticleStatus _status,
                MC::ParticleAges _ages,
                bool b_move,
                bool b_leave)
    {
      this->d_t = _d_t;
      this->n_particles = n_p;
      this->enable_leave = b_leave;
      this->enable_move = b_move;
      this->move = std::move(move_i);
 
      this->positions = std::move(_positions);
      this->status = std::move(_status);
      this->ages = std::move(_ages);
    }
    MoveFunctor(double _d_t,
                MC::ParticlePositions p,
                MC::ParticleStatus _status,
                std::size_t n_p,
                MoveInfo<ComputeSpace> m,
                MC::KPRNG::pool_type _random_pool,
                MC::EventContainer _events,
                ProbeAutogeneratedBuffer _probes,
                MC::ParticleAges _ages,
                bool b_move,
                bool b_leave)
        : d_t(_d_t), positions(std::move(p)), n_particles(n_p),
          move(std::move(m)), random_pool(_random_pool),
          status(std::move(_status)), events(std::move(_events)),
          probes(std::move(_probes)), ages(std::move(_ages)),
          enable_move(b_move), enable_leave(b_leave) {};
 
    KOKKOS_INLINE_FUNCTION void
    operator()(const Kokkos::TeamPolicy<ComputeSpace>::member_type& team_handle,
               std::size_t& dead_count) const
    {
 
      GET_INDEX(n_particles);
      if (status(idx) != MC::Status::Idle) [[unlikely]]
      {
        return;
      }
 
      ages(idx, 0) += d_t;
      if (enable_move)
      {
        handle_move(idx);
      }
      if (enable_leave)
      {
        handle_exit(idx, dead_count);
      }
    }
 
    [[nodiscard]] bool need_launch() const
    {
      return enable_leave || enable_move;
    }
 
    KOKKOS_FUNCTION void handle_move(const std::size_t idx) const
    {
      auto generator = random_pool.get_state();
      const float rng1 = generator.frand(0., 1.);
      const double rng2 = generator.drand(0., 1.);
      random_pool.free_state(generator);
 
      const std::size_t i_current_compartment = positions(idx);
 
      KOKKOS_ASSERT(i_current_compartment < move.liquid_volume.extent(0));
 
      const bool mask_next =
          probability_leaving(rng1,
                              move.liquid_volume(i_current_compartment),
                              move.diag_transition(i_current_compartment),
                              d_t);
 
      positions(idx) =
          (mask_next) ? __find_next_compartment(move.neighbors,
                                                move.cumulative_probability,
                                                i_current_compartment,
                                                rng2)
                      : i_current_compartment;
 
      KOKKOS_ASSERT(positions(idx) < move.liquid_volume.extent(0));
 
      if constexpr (AutoGenerated::FlagCompileTime::enable_event_counter)
      {
        if (mask_next)
        {
          events.wrap_incr<MC::EventType::Move>();
        }
      }
    }
 
    KOKKOS_FUNCTION void handle_exit(std::size_t idx,
                                     std::size_t& dead_count) const
    {
      const auto position = positions(idx);
      for (size_t i = 0LU; i < move.index_leaving_flow.size(); ++i)
      {
 
        auto generator = random_pool.get_state();
        const float random_number = generator.frand(0., 1.);
        random_pool.free_state(generator);
 
        const auto& index = move.index_leaving_flow(i);
        const auto& flow = move.leaving_flow(i);
        if (position != index)
        {
          return;
        }
        if (probability_leaving(
                random_number, move.liquid_volume(position), flow, d_t))
        {
          dead_count += 1;
          if constexpr (AutoGenerated::FlagCompileTime::use_probe)
          {
            // Ignore ret value, if probe is full we´re gonna miss events which
            // is not really important
            if (!probes.set(ages(idx, 0)))
            {
              Kokkos::printf("[Kernel]: PROBES OVERFLOW\r\n");
            };
          }
          ages(idx, 0) = 0;
          status(idx) = MC::Status::Exit;
          events.wrap_incr<MC::EventType::Exit>();
        }
      }
    }
 
    double d_t;
    MC::ParticlePositions positions;
    std::size_t n_particles;
    MoveInfo<ComputeSpace> move;
    MC::KPRNG::pool_type random_pool;
    MC::ParticleStatus status;
    MC::EventContainer events;
    ProbeAutogeneratedBuffer probes;
    MC::ParticleAges ages;
 
    bool enable_move;
    bool enable_leave;
  };
} // namespace Simulation::KernelInline
 
#endif