SpadesParticleContainer.H

Go to the documentation of this file.

#ifndef SPADESPARTICLECONTAINER_H
#define SPADESPARTICLECONTAINER_H
#include <AMReX.H>
#include <AMReX_AmrCore.H>
#include <AMReX_AmrParGDB.H>
#include <AMReX_NeighborParticles.H>
#include <AMReX_Random.H>
#include "ParticleData.H"
#include "ParticleOps.H"
#include "Utilities.H"
 
#if defined(AMREX_USE_CUDA) || defined(AMREX_USE_HIP)
#include <thrust/sort.h>
#include <thrust/execution_policy.h>
#endif
 
namespace spades::particles {
 
template <
    int NType,
    int NStructReal,
    int NStructInt,
    int NArrayReal,
    int NArrayInt>
class SpadesParticleContainer
    : public amrex::NeighborParticleContainer<
          NStructReal,
          NStructInt,
          NArrayReal,
          NArrayInt>
{
public:
    using ParticleType = typename amrex::NeighborParticleContainer<
        NStructReal,
        NStructInt,
        NArrayReal,
        NArrayInt>::ParticleType;
    using ParticleTileType = typename amrex::NeighborParticleContainer<
        NStructReal,
        NStructInt,
        NArrayReal,
        NArrayInt>::ParticleTileType;
    using IntVector = typename amrex::NeighborParticleContainer<
        NStructReal,
        NStructInt,
        NArrayReal,
        NArrayInt>::IntVector;
    using RealVector = typename amrex::NeighborParticleContainer<
        NStructReal,
        NStructInt,
        NArrayReal,
        NArrayInt>::RealVector;
 
    static std::string identifier() { return "spades_particles"; }
 
    explicit SpadesParticleContainer(amrex::AmrParGDB* par_gdb, int ngrow = 0);
 
    explicit SpadesParticleContainer(
        const amrex::Vector<amrex::Geometry>& geom,
        const amrex::Vector<amrex::DistributionMapping>& dmap,
        const amrex::Vector<amrex::BoxArray>& ba,
        int ngrow = 0);
 
    void initialize_state();
 
    void clear_state();
 
    void update_counts();
 
    void count_particles();
 
    void count_offsets();
 
    const amrex::iMultiFab& counts() const { return m_counts; };
 
    const amrex::iMultiFab& offsets() const { return m_offsets; };
 
    amrex::Long total_count(const int typ) const
    {
        BL_PROFILE("spades::SpadesParticleContainer::total_count()");
        return m_counts.sum(typ);
    }
 
    amrex::Long min_count(const int typ) const
    {
        BL_PROFILE("spades::SpadesParticleContainer::min_count()");
        return m_counts.min(typ);
    }
 
    amrex::Long max_count(const int typ) const
    {
        BL_PROFILE("spades::SpadesParticleContainer::max_count()");
        return m_counts.max(typ);
    }
 
    void check_sort(const amrex::MFIter& mfi);
 
    virtual void sort() = 0;
 
    template <typename CompareFunctor>
    void sort_impl(const CompareFunctor& compare);
 
    template <typename CompareFunctor>
    void nonencoded_sort_impl(const CompareFunctor& compare);
 
    void encoded_sort_impl();
 
    void print_messages(const std::string& header);
 
    void reposition_particles();
 
    virtual void write_plot_file(const std::string& plt_filename) = 0;
 
    int ngrow() const { return m_ngrow; }
 
    static constexpr int LEV{0};
 
    virtual void initialize_variable_names() = 0;
 
    void write_plot_file_impl(
        const std::string& plt_filename, const std::string& name);
 
    virtual void read_parameters()
    {
        {
            amrex::ParmParse pp("spades");
            pp.query("sort_type", m_sort_type);
            check_sort_type(m_sort_type);
        }
    }
 
    void check_sort_type(const std::string& sort_type)
    {
        const amrex::Vector<std::string> valid_types = {
            "nonencoded", "encoded"};
        if (std::find(valid_types.cbegin(), valid_types.cend(), sort_type) ==
            valid_types.cend()) {
            amrex::Abort("Invalid sort type. Must be nonencoded or encoded");
        }
    }
 
    ParticleArrays<NArrayReal, NArrayInt, ParticleType, RealVector, IntVector>
    particle_arrays(ParticleTileType& pti) const
    {
        auto& aos = pti.GetArrayOfStructs();
        auto& soa = pti.GetStructOfArrays();
        return ParticleArrays<
            NArrayReal, NArrayInt, ParticleType, RealVector, IntVector>(
            aos.dataPtr(), soa.GetRealData(), soa.GetIntData());
    }
 
protected:
    int m_ngrow;
 
    amrex::Vector<int> m_writeflags_real;
 
    amrex::Vector<int> m_writeflags_int;
 
    amrex::Vector<std::string> m_real_data_names;
 
    amrex::Vector<std::string> m_int_data_names;
 
    std::string m_sort_type{"nonencoded"};
 
    amrex::iMultiFab m_counts;
 
    amrex::iMultiFab m_offsets;
 
    amrex::MultiFab m_min_timestamp;
 
    amrex::MultiFab m_max_timestamp;
};
 
#include "SpadesParticleContainerImpl.H"
 
} // namespace spades::particles
#endif