post_process.hpp

#ifndef __CORE_POST_PROCESS_PUBLIC_HPP__
#define __CORE_POST_PROCESS_PUBLIC_HPP__
#include "Kokkos_Macros.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>
#include <stdexcept>
 
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
  {
    std::optional<ParticlePropertyViewType<HostSpace>> particle_values;
    std::optional<ParticlePropertyViewType<HostSpace>> spatial_values;
    std::optional<ParticlePropertyViewType<HostSpace>> ages;
    std::optional<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);
 
        // Use kindices for partial to map with correct vector position
        // Warning UB  if len(kindex)!=len(vnames)
        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>
  ParticlePropertyViewType<ComputeSpace>
  get_particle_age_only(MC::ParticlesContainer<M>& container)
  {
    // USE list size not Kokkos View size.
    // bcause container allocates more
    // particles than needed
    const std::size_t n_p = container.n_particles();
    auto p_age = container.ages;
    ParticlePropertyViewType<ComputeSpace> ages_values("ages_values", 2, n_p);
 
    Kokkos::parallel_for(
        "get_age",
        Kokkos::RangePolicy<ComputeSpace>(0, n_p),
        KOKKOS_LAMBDA(const std::size_t i_particle) {
          ages_values(0, i_particle) = p_age(i_particle, 0);
          ages_values(1, i_particle) = p_age(i_particle, 1);
        });
    return ages_values;
  }
 
  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;
      properties.ages = std::nullopt;
      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;
        if (ar.size() != M::get_number().size())
        {
          throw std::invalid_argument("Partial export Model need to have same "
                                      "number of name and indices");
        }
      }
 
      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)
      {
        ParticlePropertyViewType<ComputeSpace> ages_values(
            "ages_values", 2, n_p);
 
        properties.ages = Kokkos::create_mirror_view_and_copy(
            Kokkos::HostSpace(), ages_values);
      }
 
      return properties;
    }
    else
    {
 
      if (with_age)
      {
        BonceBuffer properties;
        properties.ages = Kokkos::create_mirror_view_and_copy(
            Kokkos::HostSpace(), get_particle_age_only(container));
        return properties;
      }
      return std::nullopt;
    }
  }
 
} // namespace PostProcessing
#endif