SamplingContainer.H

Go to the documentation of this file.

#ifndef SAMPLINGCONTAINER_H
#define SAMPLINGCONTAINER_H
 
#include <memory>
#include <cstdint>
 
#include "AMReX_AmrParticles.H"
#include "src/utilities/DerivedQuantity.H"
#include "AMReX_REAL.H"
 
using namespace amrex::literals;
 
namespace kynema_sgf {
 
class Field;
 
namespace sampling {
 
class SamplerBase;
 
static constexpr int SNStructReal = 0;
static constexpr int SNStructInt = 3;
static constexpr int SNArrayReal = 0;
static constexpr int SNArrayInt = 0;

struct IIx
{
    enum Indices : std::uint8_t {
        uid = 0, 
        sid,     
        nid      
    };
};

class SamplingContainer
    : public amrex::AmrParticleContainer<
          SNStructReal,
          SNStructInt,
          SNArrayReal,
          SNArrayInt>
{
public:
    explicit SamplingContainer(amrex::AmrCore& mesh)
        : amrex::AmrParticleContainer<
              SNStructReal,
              SNStructInt,
              SNArrayReal,
              SNArrayInt>(&mesh)
        , m_mesh(mesh)
    {}

    void setup_container(int num_real_components, int num_int_components = 0);

    void initialize_particles(
        const amrex::Vector<std::unique_ptr<SamplerBase>>& /*samplers*/);

    void set_interpolation_order(const int order)
    {
        AMREX_ALWAYS_ASSERT(order == 0 || order == 1);
        m_interpolation_order = order;
    }

    template <typename FType>
    void interpolate_fields(const amrex::Vector<FType>& fields, const int scomp)
    {
        BL_PROFILE("kynema-sgf::SamplingContainer::interpolate_fields");
 
        const int nlevels = m_mesh.finestLevel() + 1;
 
        for (int lev = 0; lev < nlevels; ++lev) {
            for (ParIterType pti(*this, lev); pti.isValid(); ++pti) {
                int scomp_curr = scomp;
                for (const auto* fld : fields) {
                    const bool use_nearest = (m_interpolation_order == 0);
                    if (!use_nearest) {
                        AMREX_ALWAYS_ASSERT(
                            fld->num_grow() > amrex::IntVect{0});
                    }
                    const auto farr = (*fld)(lev).const_array(pti);
                    interpolate(
                        pti, farr, lev, fld->field_location(), fld->num_comp(),
                        scomp_curr, use_nearest);
 
                    scomp_curr += fld->num_comp();
                }
            }
        }
        amrex::Gpu::streamSynchronize();
    }

    void interpolate_derived_fields(
        const DerivedQtyMgr& derived_mgr, const FieldRepo& repo, int scomp);

    void populate_buffer(std::vector<amrex::Real>& buf);
 
    long num_sampling_particles() const { return m_total_particles; }
 
    long& num_sampling_particles() { return m_total_particles; }

    template <typename FType>
    void sample_field_nearest(
        const int np,
        const int ic,
        SamplingContainer::ParticleVector& pvec,
        SamplingContainer::RealVector& pavec,
        const FType& farr,
        const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& problo,
        const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& dxi,
        const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& dx,
        const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& offset)
    {
        BL_PROFILE("kynema-sgf::SamplingContainer::sample_nearest_impl");
 
        auto* pstruct = pvec.data();
        auto* parr = pavec.data();
 
        amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE(int ip) {
            auto& p = pstruct[ip];
            const amrex::Real x =
                (p.pos(0) - problo[0] - offset[0] * dx[0]) * dxi[0];
            const amrex::Real y =
                (p.pos(1) - problo[1] - offset[1] * dx[1]) * dxi[1];
            const amrex::Real z =
                (p.pos(2) - problo[2] - offset[2] * dx[2]) * dxi[2];
 
            // Round to the nearest data point index
            const int i = static_cast<int>(std::round(x));
            const int j = static_cast<int>(std::round(y));
            const int k = static_cast<int>(std::round(z));
 
            parr[ip] = farr(i, j, k, ic);
        });
    }
 
    template <typename FType>
    void sample_field(
        const int np,
        const int ic,
        SamplingContainer::ParticleVector& pvec,
        SamplingContainer::RealVector& pavec,
        const FType& farr,
        const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& problo,
        const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& dxi,
        const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& dx,
        const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& offset)
    {
        BL_PROFILE("kynema-sgf::SamplingContainer::sample_impl");
 
        auto* pstruct = pvec.data();
        auto* parr = pavec.data();
 
        amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE(int ip) {
            auto& p = pstruct[ip];
            // Determine offsets within the containing cell
            const amrex::Real x =
                (p.pos(0) - problo[0] - offset[0] * dx[0]) * dxi[0];
            const amrex::Real y =
                (p.pos(1) - problo[1] - offset[1] * dx[1]) * dxi[1];
            const amrex::Real z =
                (p.pos(2) - problo[2] - offset[2] * dx[2]) * dxi[2];
 
            // Index of the low corner
            const int i = static_cast<int>(std::floor(x));
            const int j = static_cast<int>(std::floor(y));
            const int k = static_cast<int>(std::floor(z));
 
            // Interpolation weights in each direction (linear basis)
            const amrex::Real wx_hi = (x - i);
            const amrex::Real wy_hi = (y - j);
            const amrex::Real wz_hi = (z - k);
 
            const amrex::Real wx_lo = 1.0_rt - wx_hi;
            const amrex::Real wy_lo = 1.0_rt - wy_hi;
            const amrex::Real wz_lo = 1.0_rt - wz_hi;
 
            parr[ip] = (wx_lo * wy_lo * wz_lo * farr(i, j, k, ic)) +
                       (wx_lo * wy_lo * wz_hi * farr(i, j, k + 1, ic)) +
                       (wx_lo * wy_hi * wz_lo * farr(i, j + 1, k, ic)) +
                       (wx_lo * wy_hi * wz_hi * farr(i, j + 1, k + 1, ic)) +
                       (wx_hi * wy_lo * wz_lo * farr(i + 1, j, k, ic)) +
                       (wx_hi * wy_lo * wz_hi * farr(i + 1, j, k + 1, ic)) +
                       (wx_hi * wy_hi * wz_lo * farr(i + 1, j + 1, k, ic)) +
                       (wx_hi * wy_hi * wz_hi * farr(i + 1, j + 1, k + 1, ic));
        });
    }
 
private:
    template <typename FType>
    void interpolate(
        const ParIterType& pti,
        const FType& farr,
        const int lev,
        const FieldLoc floc,
        const int ncomp,
        const int scomp,
        const bool use_nearest = false)
    {
        const auto& geom = m_mesh.Geom(lev);
        const auto dx = geom.CellSizeArray();
        const auto dxi = geom.InvCellSizeArray();
        const auto plo = geom.ProbLoArray();
        const int np = pti.numParticles();
        auto& pvec = pti.GetArrayOfStructs()();
        int fidx = scomp;
 
        // Determine the stagger offset for this field location
        amrex::GpuArray<amrex::Real, AMREX_SPACEDIM> offset{};
        switch (floc) {
        case FieldLoc::NODE:
            offset = {0.0_rt, 0.0_rt, 0.0_rt};
            break;
        case FieldLoc::CELL:
            offset = {0.5_rt, 0.5_rt, 0.5_rt};
            break;
        case FieldLoc::XFACE:
            offset = {0.0_rt, 0.5_rt, 0.5_rt};
            break;
        case FieldLoc::YFACE:
            offset = {0.5_rt, 0.0_rt, 0.5_rt};
            break;
        case FieldLoc::ZFACE:
            offset = {0.5_rt, 0.5_rt, 0.0_rt};
            break;
        }
 
        for (int ic = 0; ic < ncomp; ++ic) {
            auto& parr = pti.GetStructOfArrays().GetRealData(fidx++);
            if (use_nearest) {
                sample_field_nearest(
                    np, ic, pvec, parr, farr, plo, dxi, dx, offset);
            } else {
                sample_field(np, ic, pvec, parr, farr, plo, dxi, dx, offset);
            }
        }
    }
 
    const amrex::AmrCore& m_mesh;
 
    long m_total_particles{0};

    int m_interpolation_order{1};
};
 
} // namespace sampling
} // namespace kynema_sgf
 
#endif /* SAMPLINGCONTAINER_H */