Nodal Vector Poisson solver with EB

.github/update_stub.sh

@@ -14,8 +14,8 @@ set -eu -o pipefail
 this_dir=$(cd $(dirname $0) && pwd)
 
 pybind11-stubgen --exit-code --enum-class-locations="GrowthStrategy:amrex.space1d" -o ${this_dir}/../src/ amrex.space1d
-pybind11-stubgen --exit-code --enum-class-locations="GrowthStrategy:amrex.space2d" -o ${this_dir}/../src/ amrex.space2d
-pybind11-stubgen --exit-code --enum-class-locations="GrowthStrategy:amrex.space3d" -o ${this_dir}/../src/ amrex.space3d
+pybind11-stubgen --exit-code --enum-class-locations="GrowthStrategy:amrex.space2d" --enum-class-locations="EBSupport:amrex.space2d" -o ${this_dir}/../src/ amrex.space2d
+pybind11-stubgen --exit-code --enum-class-locations="GrowthStrategy:amrex.space3d" --enum-class-locations="EBSupport:amrex.space3d" -o ${this_dir}/../src/ amrex.space3d
 
 # Fix circular default argumetn for
 #   strategy: GrowthStrategy = amrex.space3d.GrowthStrategy.Poisson
@@ -26,3 +26,11 @@ sed -i 's/amrex.space3d.GrowthStrategy/GrowthStrategy/g' src/amrex/space3d/amrex
 sed -i 's/= GrowthStrategy.Poisson/= "GrowthStrategy.Poisson"/g' src/amrex/space1d/amrex_1d_pybind/__init__.pyi
 sed -i 's/= GrowthStrategy.Poisson/= "GrowthStrategy.Poisson"/g' src/amrex/space2d/amrex_2d_pybind/__init__.pyi
 sed -i 's/= GrowthStrategy.Poisson/= "GrowthStrategy.Poisson"/g' src/amrex/space3d/amrex_3d_pybind/__init__.pyi
+
+# Fix circular default argument for
+#   support: EBSupport = amrex.space3d.EBSupport.full
+sed -i 's/amrex.space2d.EBSupport/EBSupport/g' src/amrex/space2d/amrex_2d_pybind/__init__.pyi
+sed -i 's/amrex.space3d.EBSupport/EBSupport/g' src/amrex/space3d/amrex_3d_pybind/__init__.pyi
+
+sed -i 's/= EBSupport.full/= "EBSupport.full"/g' src/amrex/space2d/amrex_2d_pybind/__init__.pyi
+sed -i 's/= EBSupport.full/= "EBSupport.full"/g' src/amrex/space3d/amrex_3d_pybind/__init__.pyi

.gitignore

@@ -157,3 +157,4 @@ docs/doxyxml/
 **/plt*
 ## Checkpoints
 **/chk*
+tests/fep_coils/

src/CMakeLists.txt

@@ -6,4 +6,5 @@ add_subdirectory(EB)
 #add_subdirectory(Extern)
 #add_subdirectory(LinearSolvers)
 add_subdirectory(Particle)
+add_subdirectory(VectorPoisson)
 #add_subdirectory(SDC)

src/EB/EBFabFactory.cpp

@@ -5,9 +5,14 @@
  */
 #include "pyAMReX.H"
 
+#include <AMReX_EBCellFlag.H>
 #include <AMReX_EBFabFactory.H>
+#include <AMReX_Gpu.H>
+#include <AMReX_GpuContainers.H>
 #include <AMReX_MultiFab.H>
 
+#include <stdexcept>
+
 
 void init_EBFabFactory (py::module& m)
 {
@@ -37,4 +42,137 @@ void init_EBFabFactory (py::module& m)
         py::arg("support"),
         "Make EBFArrayBoxFactory for given Geometry, BoxArray and DistributionMapping"
     );
+
+    m.def(
+        "makeStaircaseEBFabFactoryFromCupy",
+        [] (Geometry const& geom, BoxArray const& ba, DistributionMapping const& dm,
+            py::object mask_obj, Vector<int> const& ngrow, EBSupport support)
+        {
+            if (mask_obj.is_none()) {
+                throw py::type_error(
+                    "makeStaircaseEBFabFactoryFromCupy() requires a CuPy/NumPy mask object"
+                );
+            }
+
+            py::object mask_host_obj;
+            if (py::hasattr(mask_obj, "__cuda_array_interface__")) {
+                mask_host_obj = py::module_::import("cupy").attr("asnumpy")(mask_obj);
+            } else {
+                mask_host_obj = py::module_::import("numpy").attr("asarray")(mask_obj);
+            }
+
+            py::array_t<unsigned char, py::array::c_style | py::array::forcecast> mask_host(
+                mask_host_obj
+            );
+
+            if (mask_host.ndim() != 2) {
+                throw py::value_error(
+                    "makeStaircaseEBFabFactoryFromCupy() expects a 2D nodal mask"
+                );
+            }
+
+            const int nr_nodal = geom.Domain().length(0) + 1;
+            const int nz_nodal = geom.Domain().length(1) + 1;
+
+            const int in_nr = static_cast<int>(mask_host.shape(0));
+            const int in_nz = static_cast<int>(mask_host.shape(1));
+            if (in_nr < nr_nodal || in_nz < nz_nodal) {
+                throw py::value_error(
+                    "makeStaircaseEBFabFactoryFromCupy() mask is smaller than domain nodal shape"
+                );
+            }
+
+            // Accept masks with symmetric ghost padding and crop to the valid
+            // nodal domain centered window.
+            const int off_r = (in_nr - nr_nodal) / 2;
+            const int off_z = (in_nz - nz_nodal) / 2;
+
+            Vector<int> hmask(static_cast<std::size_t>(nr_nodal) * static_cast<std::size_t>(nz_nodal), 0);
+            auto const mh = mask_host.unchecked<2>();
+            for (int i = 0; i < nr_nodal; ++i) {
+                for (int j = 0; j < nz_nodal; ++j) {
+                    hmask[static_cast<std::size_t>(i) * static_cast<std::size_t>(nz_nodal)
+                        + static_cast<std::size_t>(j)] = (mh(off_r + i, off_z + j) != 0) ? 1 : 0;
+                }
+            }
+
+            auto eb_factory = makeEBFabFactory(geom, ba, dm, ngrow, support);
+
+            auto& flags = const_cast<FabArray<EBCellFlagFab>&>(eb_factory->getMultiEBCellFlagFab());
+            auto& vfrac = const_cast<MultiFab&>(eb_factory->getVolFrac());
+            auto& levset = const_cast<MultiFab&>(eb_factory->getLevelSet());
+
+            const IndexType levset_type = levset.boxArray().ixType();
+            const bool levset_is_nodal = levset_type.nodeCentered(0)
+#if (AMREX_SPACEDIM >= 2)
+                && levset_type.nodeCentered(1)
+#endif
+                ;
+            if (!levset_is_nodal) {
+                throw std::runtime_error(
+                    "makeStaircaseEBFabFactoryFromCupy() expected nodal EB levelset"
+                );
+            }
+
+            Gpu::DeviceVector<int> dmask(hmask.size());
+            Gpu::copy(Gpu::hostToDevice, hmask.begin(), hmask.end(), dmask.begin());
+            int const* nodal_mask = dmask.data();
+
+            const int nr_cc = geom.Domain().length(0);
+            const int nz_cc = geom.Domain().length(1);
+
+            for (MFIter mfi(vfrac, TilingIfNotGPU()); mfi.isValid(); ++mfi) {
+                const Box& bx = mfi.growntilebox(vfrac.nGrowVect());
+                auto const& flag_arr = flags.array(mfi);
+                auto const& vfrac_arr = vfrac.array(mfi);
+
+                ParallelFor(bx,
+                [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept {
+                    if (i >= 0 && i < nr_cc && j >= 0 && j < nz_cc) {
+                        const bool n00 = (nodal_mask[i * nz_nodal + j] != 0);
+                        const bool n10 = (nodal_mask[(i + 1) * nz_nodal + j] != 0);
+                        const bool n01 = (nodal_mask[i * nz_nodal + (j + 1)] != 0);
+                        const bool n11 = (nodal_mask[(i + 1) * nz_nodal + (j + 1)] != 0);
+                        const bool cell_covered = n00 && n10 && n01 && n11;
+
+                        auto cell = flag_arr(i, j, k);
+                        if (cell_covered) {
+                            cell.setCovered();
+                            vfrac_arr(i, j, k) = Real(0.0);
+                        } else {
+                            cell.setRegular();
+                            vfrac_arr(i, j, k) = Real(1.0);
+                        }
+                        flag_arr(i, j, k) = cell;
+                    }
+                });
+            }
+
+            for (MFIter mfi(levset, TilingIfNotGPU()); mfi.isValid(); ++mfi) {
+                const Box& bx = mfi.validbox();
+                auto const& levset_arr = levset.array(mfi);
+
+                ParallelFor(bx,
+                [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept {
+                    if (i >= 0 && i < nr_nodal && j >= 0 && j < nz_nodal) {
+                        const int idx = i * nz_nodal + j;
+                        levset_arr(i, j, k) = (nodal_mask[idx] != 0) ? Real(1.0) : Real(-1.0);
+                    }
+                });
+            }
+
+            Gpu::streamSynchronize();
+
+            const Periodicity period = eb_factory->Geom().periodicity();
+            vfrac.FillBoundary(period);
+            levset.FillBoundary(period);
+
+            return eb_factory;
+        },
+        py::arg("geom"), py::arg("ba"), py::arg("dm"), py::arg("mask"),
+        py::arg("ngrow"), py::arg("support") = EBSupport::full,
+        "Make staircase EBFArrayBoxFactory from a CuPy/NumPy nodal boolean mask "
+        "(True/1 means embedded boundary / covered)."
+    );
+
 }

src/VectorPoisson/CMakeLists.txt

@@ -0,0 +1,10 @@
+# VectorPoisson module for pyAMReX
+
+foreach(D IN LISTS AMReX_SPACEDIM)
+    target_sources(pyAMReX_${D}d
+      PRIVATE
+        VectorPoissonSolver.cpp
+        VectorPoissonSolverNodal.cpp
+        PyVectorPoisson.cpp
+    )
+endforeach()

src/VectorPoisson/PyVectorPoisson.H

@@ -0,0 +1,7 @@
+#pragma once
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/numpy.h>
+
+namespace py = pybind11;