ERF_PlaneAverage.H

Go to the documentation of this file.

#ifndef ERF_PlaneAverage_H
#define ERF_PlaneAverage_H
 
#include "AMReX_Gpu.H"
#include "AMReX_iMultiFab.H"
#include "AMReX_MultiFab.H"
#include "AMReX_GpuContainers.H"
#include "ERF_DirectionSelector.H"
 
/**
 * Basic averaging and interpolation operations
 */
 
class PlaneAverage {
public:
    /**
     * Construct a plane averaging object.
     * @param[in] field_in MultiFab to be averaged.
     * @param[in] geom_in Geometry of the domain.
     * @param[in] axis_in Axis along which to average.
     * @param[in] n_ghost_to_inc Ghost cells to include in the average.
     */
    AMREX_FORCE_INLINE
    explicit PlaneAverage (const amrex::MultiFab* field_in,
                           amrex::Geometry geom_in,
                           int axis_in,
                           amrex::IntVect n_ghost_to_inc=amrex::IntVect{0});
    PlaneAverage () = delete;
    ~PlaneAverage () = default;
 
    AMREX_FORCE_INLINE
    void operator()();
 
    /** evaluate line average at specific location for any average component */
    [[nodiscard]] AMREX_FORCE_INLINE
    amrex::Real line_average_interpolated (amrex::Real x, int comp) const;
 
    /** change precision of text file output */
    void set_precision (int p) { m_precision = p; }
 
    [[nodiscard]] amrex::Real dx () const { return m_dx; }
    [[nodiscard]] amrex::Real xlo () const { return m_xlo; }
 
    [[nodiscard]] int axis () const { return m_axis; }
    [[nodiscard]] int level () const { return m_level; }
    [[nodiscard]] int ncomp () const { return m_ncomp; }
    [[nodiscard]] int ncell_plane () const { return m_ncell_plane; }
    [[nodiscard]] int ncell_line () const { return m_ncell_line; }
 
    [[nodiscard]] const amrex::Vector<amrex::Real>& line_average () const
    {
        return m_line_average;
    }
 
    /**
     * Copy the line average of a specific component to a host vector.
     * @param[in] comp Component index.
     * @param[out] l_vec Host vector to store the line average.
     */
    AMREX_FORCE_INLINE
    void line_average (int comp, amrex::Gpu::HostVector<amrex::Real>& l_vec);
 
    [[nodiscard]] const amrex::Vector<amrex::Real>& line_centroids () const
    {
        return m_line_xcentroid;
    }
 
    [[nodiscard]] const amrex::MultiFab& field () const { return *m_field; }
 
protected:
    int m_ncomp; /** number of average components */
 
    /** line storage for the average velocity and tracer variables */
    amrex::Vector<amrex::Real> m_line_average;
 
    amrex::Vector<amrex::Real> m_line_xcentroid; /** line storage for centroids of each cell along a line*/
 
    amrex::Real m_dx;  /** grid spacing in axis direction*/
    amrex::Real m_xlo; /** bottom of domain in axis direction */
 
    int m_ncell_plane; /** number of cells in plane */
    int m_ncell_line;  /** number of cells along line */
 
    int m_precision = 4; /** precision for line plot text file */
    const int m_level = 0; /** level for plane averaging for now fixed at level=0 */
 
    const amrex::MultiFab* m_field;
    amrex::Geometry m_geom;
    const int m_axis;
    amrex::IntVect m_ghost_to_inc = amrex::IntVect{0};
    amrex::IntVect m_ixtype;
 
public:
    /** fill line storage with averages */
    template <typename IndexSelector>
    AMREX_FORCE_INLINE
    void compute_averages (const IndexSelector& idxOp, const amrex::MultiFab& mfab);
};
 
 
/**
 * Construct a plane averaging object.
 * @param[in] field_in MultiFab to be averaged.
 * @param[in] geom_in Geometry of the domain.
 * @param[in] axis_in Axis along which to average.
 * @param[in] n_ghost_to_inc Ghost cells to include in the average.
 */
PlaneAverage::PlaneAverage (const amrex::MultiFab* field_in,
                            amrex::Geometry geom_in,
                            int axis_in,
                            amrex::IntVect n_ghost_to_inc)
    : m_field(field_in), m_geom(geom_in), m_axis(axis_in), m_ghost_to_inc(n_ghost_to_inc)
{
    AMREX_ALWAYS_ASSERT(m_axis >= 0 && m_axis < AMREX_SPACEDIM);
 
    m_xlo    = m_geom.ProbLo(m_axis);
    m_dx     = m_geom.CellSize(m_axis);
    m_ncomp  = m_field->nComp();
    m_ixtype = m_field->boxArray().ixType().toIntVect();
 
    amrex::Box domain = m_geom.Domain();
    m_ghost_to_inc = std::min(field_in->nGrowVect(),m_ghost_to_inc);
    domain.grow(axis_in, m_ghost_to_inc[axis_in]);
 
    amrex::IntVect dom_lo(domain.loVect());
    amrex::IntVect dom_hi(domain.hiVect());
 
    // Domain guarantees enough space and we usually average in z-dir (spans domain)
    m_ncell_line = dom_hi[m_axis] - dom_lo[m_axis] + 1 + m_ixtype[m_axis];
 
    // First estimate is with domain (updated with MFiter later)
    m_ncell_plane = 1;
    auto period = m_geom.periodicity();
    for (int i = 0; i < AMREX_SPACEDIM; ++i) {
      int p_fac = (!period.isPeriodic(i)) ? 1 : 0;
       if (i != m_axis) m_ncell_plane *= (dom_hi[i] - dom_lo[i] + 1 + p_fac*m_ixtype[i]);
    }
 
    m_line_average.resize(static_cast<size_t>(m_ncell_line) * m_ncomp, zero);
    m_line_xcentroid.resize(m_ncell_line);
 
    for (int i = 0; i < m_ncell_line; ++i) {
       m_line_xcentroid[i] = m_xlo + (i + myhalf) * m_dx;
    }
}
 
/**
 * Evaluate line average at specific location for any average component.
 * @param[in] x Physical coordinate along the averaging axis.
 * @param[in] comp Component index.
 * @return Interpolated value of the line average.
 */
amrex::Real
PlaneAverage::line_average_interpolated (amrex::Real x, int comp) const
{
    AMREX_ALWAYS_ASSERT(comp >= 0 && comp < m_ncomp);
 
    amrex::Real c = zero;
    int ind = 0;
 
    if (x > m_xlo + myhalf * m_dx) {
        ind = static_cast<int>(floor((x - m_xlo) / m_dx - myhalf));
        const amrex::Real x1 = m_xlo + (ind + myhalf) * m_dx;
        c = (x - x1) / m_dx;
    }
 
    if (ind + 1 >= m_ncell_line) {
        ind = m_ncell_line - 2;
        c = one;
    }
 
    AMREX_ALWAYS_ASSERT(ind >= 0 && ind + 1 < m_ncell_line);
 
    return m_line_average[m_ncomp * ind + comp] * (one - c) +
           m_line_average[m_ncomp * (ind + 1) + comp] * c;
}
 
/**
 * Copy the line average of a specific component to a host vector.
 * @param[in] comp Component index.
 * @param[out] l_vec Host vector to store the line average.
 */
void
PlaneAverage::line_average (int comp, amrex::Gpu::HostVector<amrex::Real>& l_vec)
{
    AMREX_ALWAYS_ASSERT(comp >= 0 && comp < m_ncomp);
 
    for (int i = 0; i < m_ncell_line; i++) {
        l_vec[i] = m_line_average[m_ncomp * i + comp];
    }
}
 
/**
 * Compute the plane averages for the stored field.
 */
void
PlaneAverage::operator()()
{
    std::fill(m_line_average.begin(), m_line_average.end(), zero);
    switch (m_axis) {
    case 0:
        compute_averages(XDir(), *m_field);
        break;
    case 1:
        compute_averages(YDir(), *m_field);
        break;
    case 2:
        compute_averages(ZDir(), *m_field);
        break;
    default:
        amrex::Abort("axis must be equal to 0, 1, or 2");
        break;
    }
}
 
/**
 * Compute averages across a plane.
 * @tparam IndexSelector Type of index selector functor.
 * @param[in] idxOp Index selector to map cell indices to the line index.
 * @param[in] mfab MultiFab containing the data to average.
 */
template <typename IndexSelector>
void
PlaneAverage::compute_averages (const IndexSelector& idxOp, const amrex::MultiFab& mfab)
{
    amrex::Vector<amrex::Real> cnt_h(m_ncell_line,zero);
    amrex::AsyncArray<amrex::Real> cnt_d(cnt_h.data(), cnt_h.size());
    amrex::AsyncArray<amrex::Real> lavg(m_line_average.data(), m_line_average.size());
    amrex::Real* cnt_avg  = cnt_d.data();
    amrex::Real* line_avg = lavg.data();
    const int ncomp = m_ncomp;
 
    amrex::Box domain = amrex::convert(m_geom.Domain(),m_ixtype);
 
    amrex::IntVect ng = amrex::IntVect(0);
    int offset = m_ghost_to_inc[m_axis];
    ng[m_axis] = offset;
 
    std::unique_ptr<amrex::iMultiFab> mask = OwnerMask(*m_field, m_geom.periodicity(), ng);
 
#ifdef _OPENMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
    for (amrex::MFIter mfi(mfab, amrex::TilingIfNotGPU()); mfi.isValid(); ++mfi) {
        amrex::Box tbx  = mfi.tilebox();
        if (tbx.smallEnd(m_axis) == domain.smallEnd(m_axis)) tbx.growLo(m_axis,offset);
        if (tbx.bigEnd  (m_axis) == domain.bigEnd  (m_axis)) tbx.growHi(m_axis,offset);
        amrex::Box pbx = PerpendicularBox<IndexSelector>(tbx, amrex::IntVect(0));
 
        const amrex::Array4<const amrex::Real>& fab_arr  = mfab.const_array(mfi);
        const amrex::Array4<const int        >& mask_arr = mask->const_array(mfi);
 
        amrex::ParallelFor(amrex::Gpu::KernelInfo().setReduction(true), pbx, [=]
                    AMREX_GPU_DEVICE( int p_i, int p_j, int p_k,
                                      amrex::Gpu::Handler const& handler) noexcept
        {
            // Loop over the direction perpendicular to the plane.
            // This reduces the atomic pressure on the destination arrays.
            amrex::Box lbx = ParallelBox<IndexSelector>(tbx, amrex::IntVect{p_i, p_j, p_k});
 
            for (int k = lbx.smallEnd(2); k <= lbx.bigEnd(2); ++k) {
                for (int j = lbx.smallEnd(1); j <= lbx.bigEnd(1); ++j) {
                    for (int i = lbx.smallEnd(0); i <= lbx.bigEnd(0); ++i) {
                        int ind = idxOp.getIndx(i, j, k) + offset;
                        // NOTE: This factor is to avoid an if statement that will break
                        //       the devicereducesum since all threads won't participate.
                        //       This more performant than Gpu::Atomic::Add.
                        amrex::Real fac = (mask_arr(i,j,k)) ? one : zero;
                        amrex::Gpu::deviceReduceSum(&cnt_avg[ind], fac, handler);
                        for (int n = 0; n < ncomp; ++n) {
                            amrex::Gpu::deviceReduceSum(&line_avg[ncomp * ind + n],
                                                        fab_arr(i, j, k, n) * fac, handler);
                        }
                    }
                }
            }
        });
    }
 
    cnt_d.copyToHost(cnt_h.data(), cnt_h.size());
    lavg.copyToHost(m_line_average.data(), m_line_average.size());
    amrex::ParallelDescriptor::ReduceRealSum(cnt_h.data(), cnt_h.size());
    amrex::ParallelDescriptor::ReduceRealSum(m_line_average.data(), m_line_average.size());
    //
    // With AMR, a level need not cover every plane in the direction we average over: a
    // fine level typically covers only the bottom of the domain.  A plane holding no
    // cells has nothing to divide by, and leaving it at the accumulated sum -- exactly
    // zero, since m_line_average was zeroed before the reduction -- is indistinguishable
    // from a legitimate average of zero.  A consumer that divides one profile by another,
    // as the subsidence forcing in MakeSources does, then gets 0/0 = NaN far from the
    // cause.  So fill each empty plane from the closest plane that does hold cells,
    // preferring the one below it, which extends the profile as a constant beyond the
    // range of the level.
    //
    amrex::Vector<int> plane_has_cells(m_ncell_line);
    int ncell_line_empty = 0;
    for (int ind(0); ind<m_ncell_line; ++ind) {
        plane_has_cells[ind] = (cnt_h[ind] > zero);
        if (plane_has_cells[ind]) {
            for (int n(0); n<ncomp; ++n) {
                m_line_average[ncomp*ind + n] /= cnt_h[ind];
            }
        } else {
            ++ncell_line_empty;
        }
    }
 
    if (ncell_line_empty == m_ncell_line) {
        amrex::Abort("PlaneAverage: no cells were found in any plane");
    }
 
    if (ncell_line_empty > 0) {
        // Sweeping up first and then down leaves every empty plane holding the value of
        // the nearest non-empty plane below it, or above it if there is none below.
        for (int ind(1); ind<m_ncell_line; ++ind) {
            if (!plane_has_cells[ind] && plane_has_cells[ind-1]) {
                for (int n(0); n<ncomp; ++n) {
                    m_line_average[ncomp*ind + n] = m_line_average[ncomp*(ind-1) + n];
                }
                plane_has_cells[ind] = 1;
            }
        }
        for (int ind(m_ncell_line-2); ind>=0; --ind) {
            if (!plane_has_cells[ind] && plane_has_cells[ind+1]) {
                for (int n(0); n<ncomp; ++n) {
                    m_line_average[ncomp*ind + n] = m_line_average[ncomp*(ind+1) + n];
                }
                plane_has_cells[ind] = 1;
            }
        }
    }
 
    // Take the size of a plane from one that actually holds cells: plane 0 need not be
    // one of them on a level whose grids do not reach the bottom of the domain.
    amrex::Real cnt_max = zero;
    for (int ind(0); ind<m_ncell_line; ++ind) { cnt_max = amrex::max(cnt_max,cnt_h[ind]); }
    m_ncell_plane = static_cast<int>(cnt_max);
}
#endif /* ERF_PlaneAverage.H */