move_kernel.hpp

#ifndef __SIMULATION_MOVE_KERNEL_HPP__
#define __SIMULATION_MOVE_KERNEL_HPP__
 
#include "Kokkos_Assert.hpp"
#include <Kokkos_Core.hpp>
#include <Kokkos_Random.hpp>
#include <biocma_cst_config.hpp>
#include <cassert>
#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")
    {
    }
  };
 
  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
  }
 
  template <bool enable_leave, bool enable_move> struct MoveFunctor
  {
    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)
        : 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)) {};
 
    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]]
      {
        // Kokkos::printf("Skip %ld", idx);
        return;
      }
 
      if constexpr (enable_move)
      {
        handle_move(idx);
      }
 
      if constexpr (enable_leave)
      {
        handle_exit(idx, dead_count);
      }
    }
 
    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);
 
      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
    {
 
      for (size_t i = 0LU; i < move.index_leaving_flow.size(); ++i)
      {
        const auto position = positions(idx);
        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;
          status(idx) = MC::Status::Exit;
          events.wrap_incr<MC::EventType::Exit>();
          if constexpr (AutoGenerated::FlagCompileTime::use_probe)
          {
            // Ignore ret value, if probe is full were gonna miss events which is not really
            // important
            // auto _ = probe_at_t.set(properties.hydraulic_time);
          }
        }
      }
    }
 
    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;
  };
} // namespace Simulation::KernelInline
 
#endif