post_process.hpp

#ifndef __CORE_POST_PROCESS_PUBLIC_HPP__
#define __CORE_POST_PROCESS_PUBLIC_HPP__
#include <Kokkos_Core.hpp>
#include <Kokkos_Core_fwd.hpp>
#include <Kokkos_ScatterView.hpp>
#include <common/common.hpp>
#include <mc/alias.hpp>
#include <mc/particles_container.hpp>
#include <mc/traits.hpp>
 
template <typename MemorySpace>
using ParticlePropertyViewType =
    Kokkos::View<double**, Kokkos::LayoutRight, MemorySpace>;
 
using SubViewtype = Kokkos::Subview<ParticlePropertyViewType<ComputeSpace>,
                                    decltype(Kokkos::ALL),
                                    std::size_t>;
namespace PostProcessing
{
  struct BonceBuffer
  {
    ParticlePropertyViewType<HostSpace> particle_values;
    ParticlePropertyViewType<HostSpace> spatial_values;
    std::optional<ParticlePropertyViewType<HostSpace>> ages;
    std::vector<std::string> vnames;
  };
 
  namespace
  {
    template <typename Model, typename ExecutionSpace>
    struct GetPropertiesFunctor
    {
      using ComputeSpace = typename ExecutionSpace::memory_space;
 
      GetPropertiesFunctor(
          std::size_t n_p,
          MC::ParticlePositions position,
          typename Model::SelfParticle model,
          MC::ParticleAges ages,
          ParticlePropertyViewType<ComputeSpace> particle_values,
          ParticlePropertyViewType<ComputeSpace> spatial_values,
          ParticlePropertyViewType<ComputeSpace> ages_value,
          MC::ParticleStatus status)
          : n_p(n_p), position(std::move(position)), model(std::move(model)),
            ages(std::move(ages)), particle_values(std::move(particle_values)),
            spatial_values(std::move(spatial_values)),
            ages_value(std::move(ages_value)), status(std::move(status))
      {
      }
 
      KOKKOS_INLINE_FUNCTION
      void operator()(const int i_particle) const
        requires(HasExportPropertiesFull<Model>)
      {
        if (status(i_particle) != MC::Status::Idle)
        {
          return;
        }
 
        auto access = scatter_spatial_values.access();
        for (size_t k_var = 0; k_var < Model::n_var; ++k_var)
        {
          const auto current = model(i_particle, k_var);
          access(k_var, position(i_particle)) += current;
          particle_values(k_var, i_particle) = current;
        }
 
        const auto mass = Model::mass(i_particle, model);
        access(Model::n_var, position(i_particle)) += mass;
        particle_values(Model::n_var, i_particle) = mass;
        ages_value(0, i_particle) = ages(i_particle, 0);
        ages_value(1, i_particle) = ages(i_particle, 1);
      }
 
      KOKKOS_INLINE_FUNCTION
      void operator()(const int i_particle) const
        requires(HasExportPropertiesPartial<Model>)
      {
        if (status(i_particle) != MC::Status::Idle)
        {
          return;
        }
 
        auto access = scatter_spatial_values.access();
        for (size_t k_var = 0; k_var < kindices.extent(0); ++k_var)
        {
          const auto current = model(i_particle, kindices(k_var));
          access(k_var, position(i_particle)) += current;
          particle_values(k_var, i_particle) = current;
        }
 
        const auto mass = Model::mass(i_particle, model);
        access(kindices.extent(0), position(i_particle)) += mass;
        particle_values(kindices.extent(0), i_particle) = mass;
        ages_value(0, i_particle) = ages(i_particle, 0);
        ages_value(1, i_particle) = ages(i_particle, 1);
      }
 
      void run()
      {
 
        scatter_spatial_values =
            Kokkos::Experimental::create_scatter_view(spatial_values);
 
        // Pour PartialExport, on initialise kindices
        if constexpr (HasExportPropertiesPartial<Model>)
        {
          static const auto indices = Model::get_number();
          Kokkos::View<size_t*, HostSpace> host_index("host_index",
                                                      indices.size());
          for (size_t i = 0; i < indices.size(); ++i)
          {
            host_index(i) = indices[i];
          }
          kindices =
              Kokkos::create_mirror_view_and_copy(ComputeSpace(), host_index);
        }
 
        Kokkos::parallel_for("get_properties",
                             Kokkos::RangePolicy<ExecutionSpace>(0, n_p),
                             *this);
 
        Kokkos::fence();
        Kokkos::Experimental::contribute(spatial_values,
                                         scatter_spatial_values);
      }
 
      std::size_t n_p;
      MC::ParticlePositions position;
      typename Model::SelfParticle model;
      MC::ParticleAges ages;
      ParticlePropertyViewType<ComputeSpace> particle_values;
      ParticlePropertyViewType<ComputeSpace> spatial_values;
      ParticlePropertyViewType<ComputeSpace> ages_value;
      MC::ParticleStatus status;
      Kokkos::View<const size_t*, ComputeSpace> kindices;
      Kokkos::Experimental::
          ScatterView<double**, Kokkos::LayoutRight, ComputeSpace>
              scatter_spatial_values;
    };
  } // namespace
 
  template <ModelType M>
  std::optional<PostProcessing::BonceBuffer>
  get_properties(MC::ParticlesContainer<M>& container,
                 const std::size_t n_compartment,
                 const bool with_age)
  {
    if constexpr (HasExportProperties<M>)
    {
      container.force_remove_dead();
      BonceBuffer properties;
      const std::size_t n_p =
          container.n_particles(); // USE list size not Kokkos View size.
                                   // bcause container allocates more particles
                                   // than needed
      auto ar = M::names();
      properties.vnames = std::vector<std::string>(ar.begin(), ar.end());
      properties.vnames.emplace_back("mass");
 
      std::size_t n_var = M::n_var + 1;
      if constexpr (HasExportPropertiesPartial<M>)
      {
        n_var = ar.size() + 1;
      }
 
      ParticlePropertyViewType<ComputeSpace> spatial_values(
          "property_spatial", n_var, n_compartment);
      ParticlePropertyViewType<ComputeSpace> particle_values(
          "property_values", n_var, n_p);
 
      // TODO Currently view is allocated and filled even if with_age flag is
      // turned of, with_age controlls export but not postprocessing
      ParticlePropertyViewType<ComputeSpace> ages_values("ages_values", 2, n_p);
 
      GetPropertiesFunctor<M, Kokkos::DefaultExecutionSpace>(n_p,
                                                             container.position,
                                                             container.model,
                                                             container.ages,
                                                             particle_values,
                                                             spatial_values,
                                                             ages_values,
                                                             container.status)
          .run();
 
      properties.particle_values = Kokkos::create_mirror_view_and_copy(
          Kokkos::HostSpace(), particle_values);
      properties.spatial_values = Kokkos::create_mirror_view_and_copy(
          Kokkos::HostSpace(), spatial_values);
 
      if (with_age)
      {
 
        properties.ages = Kokkos::create_mirror_view_and_copy(
            Kokkos::HostSpace(), ages_values);
      }
      else
      {
        properties.ages = std::nullopt;
      }
 
      return properties;
    }
    else
    {
      return std::nullopt;
    }
  }
 
} // namespace PostProcessing
#endif