Add: MultiFabRegister

Add a central location to store all `MultiFab`s in.

Source/WarpX.H

@@ -49,6 +49,7 @@
 #include "Utils/WarpXAlgorithmSelection.H"
 #include "Utils/export.H"
 
+#include <ablastr/fields/MultiFabRegister.H>
 #include <ablastr/utils/Enums.H>
 
 #include <AMReX.H>
@@ -1500,6 +1501,9 @@ private:
     amrex::Vector<amrex::Real> t_old;
     amrex::Vector<amrex::Real> dt;
 
+    // Field container
+    ablastr::fields::MultiFabRegister m_multifab;
+
     // Particle container
     std::unique_ptr<MultiParticleContainer> mypc;
     std::unique_ptr<MultiDiagnostics> multi_diags;

Source/ablastr/fields/CMakeLists.txt

@@ -1,5 +1,11 @@
 foreach(D IN LISTS WarpX_DIMS)
     warpx_set_suffix_dims(SD ${D})
+
+    target_sources(ablastr_${SD}
+      PRIVATE
+        MultiFabRegister.cpp
+    )
+
     if(ABLASTR_FFT AND D EQUAL 3)
         target_sources(ablastr_${SD}
           PRIVATE

Source/ablastr/fields/Make.package

@@ -1,4 +1,5 @@
 ifeq ($(USE_FFT),TRUE)
+    CEXE_sources += MultiFabRegister.cpp
     ifeq ($(DIM),3)
         CEXE_sources += IntegratedGreenFunctionSolver.cpp
     endif

Source/ablastr/fields/MultiFabRegister.H

@@ -0,0 +1,126 @@
+/* Copyright 2024 The ABLAST Community
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ * Authors: Axel Huebl, ...
+ */
+#ifndef ABLASTR_FIELDS_MF_REGISTER_H
+#define ABLASTR_FIELDS_MF_REGISTER_H
+
+
+#include <AMReX_MultiFab.H>
+
+#include <map>
+#include <memory>
+#include <optional>
+#include <string>
+#include <stdexcept>
+#include <vector>
+
+
+namespace ablastr::fields
+{
+    struct MultiFabOwner
+    {
+        // TODO: also add iMultiFab via std::variant
+
+        /** owned (i)MultiFab data */
+        amrex::MultiFab m_mf;
+
+        /** redistribute */
+        bool redistribute = true;
+    };
+
+    /** This is a register of fields aka amrex::(i)MultiFabs.
+     *
+     * This is owned by a simulation instance. All used fields should be registered here.
+     */
+    struct MultiFabRegister
+    {
+        /** title
+         *
+         * body body
+         * body
+         * body
+         *
+         * @param name ...
+         * @param ba ...
+         * @param dm ...
+         * @param ncomp ...
+         * @param ngrow ...
+         * @param level ...
+         * @param redistribute ...
+         * @param initial_value ...
+         * @return pointer to newly allocated MultiFab
+         */
+        amrex::MultiFab*
+        alloc_init (
+            std::string name,
+            const amrex::BoxArray& ba,
+            const amrex::DistributionMapping& dm,
+            const int ncomp,
+            const amrex::IntVect& ngrow,
+            const int level,
+            bool redistribute = true,
+            std::optional<const amrex::Real> initial_value = std::nullopt
+        );
+
+        /** title
+         *
+         * body body
+         * body
+         *
+         * @param name ...
+         */
+        void
+        alloc_like (
+            std::string other_key
+        );
+
+        /** title
+         *
+         * body body
+         * body
+         *
+         * @param name ...
+         * @return ...
+         */
+        amrex::MultiFab*
+        get (
+            std::string name
+        );
+
+        /** title
+         *
+         * body body
+         * body
+         *
+         * @return ...
+         */
+        std::vector<std::string>
+        list ();
+
+        /** title
+         *
+         * body body
+         * body
+         *
+         * @param name ...
+         * @return ...
+         */
+        void
+        erase (
+            std::string name
+        );
+
+    private:
+        std::map<
+            std::string,
+            MultiFabOwner
+        > m_mf_register;
+    };
+
+} // namespace ablastr::fields
+
+#endif  // ABLASTR_FIELDS_MF_REGISTER_H

Source/ablastr/fields/MultiFabRegister.cpp

@@ -0,0 +1,129 @@
+/* Copyright 2024 The ABLAST Community
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ * Authors: Axel Huebl, ...
+ */
+
+
+#include "MultiFabRegister.H"
+
+
+
+namespace ablastr::fields
+{
+
+    amrex::MultiFab*
+    MultiFabRegister::alloc_init (
+        std::string name,
+        const amrex::BoxArray& ba,
+        const amrex::DistributionMapping& dm,
+        const int ncomp,
+        const amrex::IntVect& ngrow,
+        const int level,
+        bool redistribute,
+        std::optional<const amrex::Real> initial_value
+    )
+    {
+        // create name
+        //     Add the suffix "[level=level]"
+        name.append("[level=").append(std::to_string(level)).append("]");
+
+        // Checks
+        // TODO: does the key already exist? error
+
+        // allocate
+        const auto tag = amrex::MFInfo().SetTag(name);
+        auto [it, success] = m_mf_register.emplace(
+            std::make_pair(
+                name,
+                MultiFabOwner{{ba, dm, ncomp, ngrow, tag}, redistribute}
+            )
+        );
+        if (!success) {
+            throw std::runtime_error("MultiFabRegister::alloc_init failed for " + name);
+        }
+
+        // a short-hand alias for the code below
+        amrex::MultiFab & mf = it->second.m_mf;
+
+        // initialize with value
+        if (initial_value) {
+            mf.setVal(*initial_value);
+        }
+
+        return &mf;
+    }
+
+    void
+    MultiFabRegister::alloc_like (
+        std::string /* other_key */
+    )
+    {
+        throw std::runtime_error("MultiFabRegister::alloc_like not yet implemented");
+
+        // Checks
+        // TODO: does the key already exist? error
+    }
+
+    /** title
+     *
+     * body body
+     * body
+     *
+     * @param name ...
+     * @return ...
+     */
+    amrex::MultiFab*
+    MultiFabRegister::get (
+        std::string name
+    )
+    {
+        if (m_mf_register.count(name) == 0) {
+            throw std::runtime_error("MultiFabRegister::get name does not exist in register: " + name);
+        }
+        amrex::MultiFab & mf = m_mf_register[name].m_mf;
+
+        return &mf;
+    }
+
+    /** title
+     *
+     * body body
+     * body
+     *
+     * @return ...
+     */
+    std::vector<std::string>
+    MultiFabRegister::list ()
+    {
+        std::vector<std::string> names;
+        names.reserve(m_mf_register.size());
+        for (auto const & str : m_mf_register) { names.push_back(str.first); }
+
+        return names;
+    }
+
+
+    /** title
+     *
+     * body body
+     * body
+     *
+     * @param name ...
+     * @return ...
+     */
+    void
+    MultiFabRegister::erase (
+        std::string name
+    )
+    {
+        if (m_mf_register.count(name) != 1) {
+            throw std::runtime_error("MultiFabRegister::remove name does not exist in register: " + name);
+        }
+        m_mf_register.erase(name);
+    }
+
+
+} // namespace ablastr::fields