MultiFabRegister.H

Go to the documentation of this file.

/* Copyright 2024 The ABLASTR Community
 *
 * This file is part of ABLASTR.
 *
 * License: BSD-3-Clause-LBNL
 * Authors: Axel Huebl
 */
#ifndef ABLASTR_FIELDS_MF_REGISTER_H
#define ABLASTR_FIELDS_MF_REGISTER_H
 
#include <AMReX_BaseFwd.H>
#include <AMReX_Enum.H>
#include <AMReX_Extension.H>
#include <AMReX_IntVect.H>
#include <AMReX_MultiFab.H>
#include <AMReX_REAL.H>
#include <AMReX_Vector.H>
 
#include <array>
#include <map>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
 
 
namespace
{
    // type trait helpers in lieu of an amrex::is_amrex_enum
    template <typename, typename = std::void_t<>>
    struct is_castable_to_string : std::false_type {};
 
    template <typename T>
    struct is_castable_to_string<T, std::void_t<decltype(static_cast<std::string>(std::declval<T>()))>> : std::true_type {};

    template<typename T>
    std::string getExtractedName (T name)
    {
        if constexpr(is_castable_to_string<T>())
        {
            // already a unique string key
            return std::string(std::move(name));
        } else
        {
            // user-defined AMREX_ENUM or compile error
            return amrex::getEnumNameString(name);
        }
    }
}
 
namespace ablastr::fields
{
    class Direction
    {
        int dir = 0;
 
    public:
        constexpr explicit Direction (int d) : dir(d) {}
 
#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER) || defined(WARPX_DIM_RSPHERE)
        static const Direction r;
        static const Direction theta;
 
        inline static const std::string r_string = "r";
        inline static const std::string theta_string = "theta";
#endif
#if defined(WARPX_DIM_RSPHERE)
        static const Direction phi;
 
        inline static const std::string phi_string = "phi";
#endif
#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z)
        static const Direction x;
        static const Direction y;
 
        inline static const std::string x_string = "x";
        inline static const std::string y_string = "y";
#endif
#if !defined(WARPX_DIM_RSPHERE)
        static const Direction z;
 
        inline static const std::string z_string = "z";
#endif
 
        bool operator<(const Direction& other) const
        {
            return other.dir < this->dir;
        }
 
        operator std::string() const // NOLINT(google-explicit-constructor)
        {
#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER) || defined(WARPX_DIM_RSPHERE)
            if (dir == r) { return r_string; }
            if (dir == theta) { return theta_string; }
#endif
#if defined(WARPX_DIM_RSPHERE)
            if (dir == phi) { return phi_string; }
#endif
#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z)
            if (dir == x) { return x_string; }
            if (dir == y) { return y_string; }
#endif
#if !defined(WARPX_DIM_RSPHERE)
            if (dir == z) { return z_string; }
#endif
            throw std::runtime_error("invalid direction: " + std::to_string(dir));
            return std::to_string(dir);
        }
 
        Direction (std::string const & s) : // NOLINT(google-explicit-constructor)
            dir{-1}
        {
#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER) || defined(WARPX_DIM_RSPHERE)
            if (s == r_string) { dir = r; }
            if (s == theta_string) { dir = theta; }
#endif
#if defined(WARPX_DIM_RSPHERE)
            if (s == phi_string) { dir = phi; }
#endif
#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z)
            if (s == x_string) { dir = x; }
            if (s == y_string) { dir = y; }
#endif
#if !defined(WARPX_DIM_RSPHERE)
            if (s == z_string) { dir = z; }
#endif
 
            if (dir == -1) {
                throw std::runtime_error("invalid direction: " + s);
            }
        }
 
        // e.g., "x"
        Direction (char const * c) : dir{Direction{c}} {} // NOLINT(google-explicit-constructor)
 
        // e.g., 'x'
        Direction (char const c) : dir{Direction{c}} {} // NOLINT(google-explicit-constructor)
 
        // rule of five: define default constructors and the destructor
        //   because we defined a special one above for char
        ~Direction () = default;
        Direction (const Direction&) = default;
        Direction& operator= (const Direction&) = default;
        Direction (Direction&&) = default;
        Direction& operator= (Direction&&) = default;
 
        /* TODO: just temporary int compatibility */
        operator int() const { return dir; } // NOLINT(google-explicit-constructor)
 
    };
 
#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER) || defined(WARPX_DIM_RSPHERE)
    inline constexpr Direction Direction::r = Direction{0};
    inline constexpr Direction Direction::theta = Direction{1};
#endif
#if defined(WARPX_DIM_RSPHERE)
    inline constexpr Direction Direction::phi = Direction{2};
#endif
 
#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z)
    inline constexpr Direction Direction::x = Direction{0};
    inline constexpr Direction Direction::y = Direction{1};
#endif
 
#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z) || defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER)
    inline constexpr Direction Direction::z = Direction{2};
#endif

    using ScalarField = amrex::MultiFab*;

    using ConstScalarField = amrex::MultiFab const *;

    //using VectorField = ablastr::utils::ConstMap<Direction, amrex::MultiFab *>;
    using VectorField = std::array<amrex::MultiFab *, 3>;

    //using VectorField = ablastr::utils::ConstMap<Direction, amrex::MultiFab const *>;
    using ConstVectorField = std::array<amrex::MultiFab const *, 3>;

    using MultiLevelScalarField = amrex::Vector<ScalarField>;

    using ConstMultiLevelScalarField = amrex::Vector<ConstScalarField>;

    using MultiLevelVectorField = amrex::Vector<VectorField>;

    using ConstMultiLevelVectorField = amrex::Vector<ConstVectorField>;

    struct MultiFabOwner
    {
        // TODO: also add iMultiFab via std::variant

        amrex::MultiFab m_mf;

        std::optional<Direction> m_dir = std::nullopt;

        int m_level = 0;

        bool m_remake = true;

        bool m_redistribute_on_remake = true;

        bool m_checkpoint_restart = false;

        std::string m_owner;

        AMREX_INLINE
        bool
        is_vector () const { return m_dir.has_value(); }

        AMREX_INLINE
        bool
        is_alias () const { return !m_owner.empty(); }
    };

    struct MultiFabRegister
    {
        // Avoid accidental copies when passing to member functions
        MultiFabRegister() = default;
        MultiFabRegister(MultiFabRegister const &) = delete;
        MultiFabRegister(MultiFabRegister&&) = delete;
        MultiFabRegister& operator=(MultiFabRegister const &) = delete;
        MultiFabRegister& operator=(MultiFabRegister&&) = delete;
        ~MultiFabRegister() = default;

        template<typename T>
        amrex::MultiFab*
        alloc_init (
            T name,
            int level,
            amrex::BoxArray const & ba,
            amrex::DistributionMapping const & dm,
            int ncomp,
            amrex::IntVect const & ngrow,
            std::optional<amrex::Real const> initial_value = std::nullopt,
            bool remake = true,
            bool redistribute_on_remake = true,
            bool checkpoint_restart = false
        )
        {
            return internal_alloc_init(
                getExtractedName(std::move(name)),
                level,
                ba,
                dm,
                ncomp,
                ngrow,
                initial_value,
                remake,
                redistribute_on_remake,
                checkpoint_restart
            );
        }

        template<typename T>
        amrex::MultiFab*
        alloc_init (
            T name,
            Direction dir,
            int level,
            amrex::BoxArray const & ba,
            amrex::DistributionMapping const & dm,
            int ncomp,
            amrex::IntVect const & ngrow,
            std::optional<amrex::Real const> initial_value = std::nullopt,
            bool remake = true,
            bool redistribute_on_remake = true,
            bool checkpoint_restart = false
        )
        {
            return internal_alloc_init(
                getExtractedName(std::move(name)),
                dir,
                level,
                ba,
                dm,
                ncomp,
                ngrow,
                initial_value,
                remake,
                redistribute_on_remake,
                checkpoint_restart
            );
        }

        template<typename N, typename A>
        amrex::MultiFab*
        alias_init (
            N new_name,
            A alias_name,
            int level,
            std::optional<amrex::Real const> initial_value = std::nullopt
        )
        {
            return internal_alias_init(
                getExtractedName(new_name),
                getExtractedName(alias_name),
                level,
                initial_value
            );
        }

        template<typename N, typename A>
        amrex::MultiFab*
        alias_init (
            N new_name,
            A alias_name,
            Direction dir,
            int level,
            std::optional<amrex::Real const> initial_value = std::nullopt
        )
        {
            return internal_alias_init(
                getExtractedName(new_name),
                getExtractedName(alias_name),
                dir,
                level,
                initial_value
            );
        }

        template<typename T>
        [[nodiscard]] bool
        has (
            T name,
            int level
        ) const
        {
            return internal_has(
                getExtractedName(std::move(name)),
                level
            );
        }

        template<typename T>
        [[nodiscard]] bool
        has (
            T name,
            Direction dir,
            int level
        ) const
        {
            return internal_has(
                getExtractedName(std::move(name)),
                dir,
                level
            );
        }

        template<typename T>
        [[nodiscard]] bool
        has_vector (
            T name,
            int level
        ) const
        {
            return internal_has_vector(
                getExtractedName(std::move(name)),
                level
            );
        }

        template<typename T>
        [[nodiscard]] amrex::MultiFab*
        get (
            T name,
            int level
        )
        {
            return internal_get(
                getExtractedName(std::move(name)),
                level
            );
        }

        template<typename T>
        [[nodiscard]] amrex::MultiFab*
        get (
            T name,
            Direction dir,
            int level
        )
        {
            return internal_get(
                getExtractedName(std::move(name)),
                dir,
                level
            );
        }

        template<typename T>
        [[nodiscard]] amrex::MultiFab const *
        get (
            T name,
            int level
        ) const
        {
            return internal_get(
                getExtractedName(std::move(name)),
                level
            );
        }

        template<typename T>
        [[nodiscard]] amrex::MultiFab const *
        get (
            T name,
            Direction dir,
            int level
        ) const
        {
            return internal_get(
                getExtractedName(std::move(name)),
                dir,
                level
            );
        }

        template<typename T>
        [[nodiscard]] MultiLevelScalarField
        get_mr_levels (
            T name,
            int finest_level,
            bool skip_level_0=false
        )
        {
            return internal_get_mr_levels(
                getExtractedName(std::move(name)),
                finest_level,
                skip_level_0
            );
        }
        template<typename T>
        [[nodiscard]] ConstMultiLevelScalarField
        get_mr_levels (
            T name,
            int finest_level,
            bool skip_level_0=false
        ) const
        {
            return internal_get_mr_levels(
                getExtractedName(std::move(name)),
                finest_level,
                skip_level_0
            );
        }


        template<typename T>
        [[nodiscard]] VectorField
        get_alldirs  (
            T name,
            int level
        )
        {
            return internal_get_alldirs(
                getExtractedName(std::move(name)),
                level
            );
        }
        template<typename T>
        [[nodiscard]] ConstVectorField
        get_alldirs  (
            T name,
            int level
        ) const
        {
            return internal_get_alldirs(
                getExtractedName(std::move(name)),
                level
            );
        }


        template<typename T>
        [[nodiscard]] MultiLevelVectorField
        get_mr_levels_alldirs (
            T name,
            int finest_level,
            bool skip_level_0=false
        )
        {
            return internal_get_mr_levels_alldirs(
                getExtractedName(std::move(name)),
                finest_level,
                skip_level_0
            );
        }
        template<typename T>
        [[nodiscard]] ConstMultiLevelVectorField
        get_mr_levels_alldirs (
            T name,
            int finest_level,
            bool skip_level_0=false
        ) const
        {
            return internal_get_mr_levels_alldirs(
                getExtractedName(std::move(name)),
                finest_level,
                skip_level_0
            );
        }


        [[nodiscard]] std::vector<std::string>
        list () const;

        template<typename T>
        void
        erase (
            T name,
            int level
        )
        {
            internal_erase(getExtractedName(std::move(name)), level);
        }

        template<typename T>
        void
        erase (
            T name,
            Direction dir,
            int level
        )
        {
            internal_erase(getExtractedName(name), dir, level);
        }

        void
        clear_level (
            int level
        );

        void
        remake_level (
            int other_level,
            amrex::DistributionMapping const & new_dm
        );

        void
        write_checkpoints (
            int level,
            std::string const & dir
        );

        void
        read_restarts (
            int level,
            std::string const & dir
        );

        [[nodiscard]] std::string
        mf_name (
            std::string name,
            int level
        ) const;

        [[nodiscard]] std::string
        mf_name (
            std::string name,
            Direction dir,
            int level
        ) const;

        [[nodiscard]] bool
        internal_has (
            std::string const & internal_name
        );
        [[nodiscard]] amrex::MultiFab *
        internal_get (
            std::string const & internal_name
        );
 
    private:
 
        [[nodiscard]] amrex::MultiFab const *
        internal_get (
            std::string const & internal_name
        ) const;
 
        amrex::MultiFab*
        internal_alloc_init (
            std::string const & name,
            int level,
            amrex::BoxArray const & ba,
            amrex::DistributionMapping const & dm,
            int ncomp,
            amrex::IntVect const & ngrow,
            std::optional<amrex::Real const> initial_value = std::nullopt,
            bool remake = true,
            bool redistribute_on_remake = true,
            bool checkpoint_restart = false
        );
        amrex::MultiFab*
        internal_alloc_init (
            std::string const & name,
            Direction dir,
            int level,
            amrex::BoxArray const & ba,
            amrex::DistributionMapping const & dm,
            int ncomp,
            amrex::IntVect const & ngrow,
            std::optional<amrex::Real const> initial_value = std::nullopt,
            bool remake = true,
            bool redistribute_on_remake = true,
            bool checkpoint_restart = false
        );
 
        amrex::MultiFab*
        internal_alias_init (
            std::string const & new_name,
            std::string const & alias_name,
            int level,
            std::optional<amrex::Real const> initial_value = std::nullopt
        );
        amrex::MultiFab*
        internal_alias_init (
            std::string const & new_name,
            std::string const & alias_name,
            Direction dir,
            int level,
            std::optional<amrex::Real const> initial_value = std::nullopt
        );
 
        [[nodiscard]] bool
        internal_has (
            std::string const & name,
            int level
        ) const;
        [[nodiscard]] bool
        internal_has (
            std::string const & name,
            Direction dir,
            int level
        ) const;
        [[nodiscard]] bool
        internal_has_vector (
            std::string const & name,
            int level
        ) const;
 
        [[nodiscard]] amrex::MultiFab *
        internal_get (
            std::string const & name,
            int level
        );
        [[nodiscard]] amrex::MultiFab const *
        internal_get (
            std::string const & name,
            int level
        ) const;
        [[nodiscard]] amrex::MultiFab *
        internal_get (
            std::string const & name,
            Direction dir,
            int level
        );
        [[nodiscard]] amrex::MultiFab const *
        internal_get (
            std::string const & name,
            Direction dir,
            int level
        ) const;
        [[nodiscard]] MultiLevelScalarField
        internal_get_mr_levels (
            std::string const & name,
            int finest_level,
            bool skip_level_0
        );
        [[nodiscard]] ConstMultiLevelScalarField
        internal_get_mr_levels (
            std::string const & name,
            int finest_level,
            bool skip_level_0
        ) const;
        [[nodiscard]] VectorField
        internal_get_alldirs (
            std::string const & name,
            int level
        );
        [[nodiscard]] ConstVectorField
        internal_get_alldirs (
            std::string const & name,
            int level
        ) const;
        [[nodiscard]] MultiLevelVectorField
        internal_get_mr_levels_alldirs (
            std::string const & name,
            int finest_level,
            bool skip_level_0
        );
        [[nodiscard]] ConstMultiLevelVectorField
        internal_get_mr_levels_alldirs (
            std::string const & name,
            int finest_level,
            bool skip_level_0
        ) const;
 
        void
        internal_erase (
            std::string const & name,
            int level
        );
        void
        internal_erase (
            std::string const & name,
            Direction dir,
            int level
        );

        std::map<
            std::string,
            MultiFabOwner
        > m_mf_register;
 
      public:
        static inline std::vector<Direction> m_all_dirs =
#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z)
            {Direction::x, Direction::y, Direction::z};
#elif defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER)
            {Direction::r, Direction::theta, Direction::z};
#elif defined(WARPX_DIM_RSPHERE)
            {Direction::r, Direction::theta, Direction::phi};
#endif
    };

    VectorField
    a2m (
        std::array< std::unique_ptr<amrex::MultiFab>, 3 > const & old_vectorfield
    );
 
} // namespace ablastr::fields
 
#endif  // ABLASTR_FIELDS_MF_REGISTER_H