diff --git a/src/Base/CMakeLists.txt b/src/Base/CMakeLists.txt index e4d774e3..40a67d2f 100644 --- a/src/Base/CMakeLists.txt +++ b/src/Base/CMakeLists.txt @@ -20,12 +20,15 @@ foreach(D IN LISTS AMReX_SPACEDIM) CoordSys.cpp Dim3.cpp DistributionMapping.cpp + FabArray.cpp FArrayBox.cpp Geometry.cpp + iMultiFab.cpp IndexType.cpp IntVect.cpp RealVect.cpp SmallMatrix.cpp + MFInfo.cpp MultiFab.cpp ParallelDescriptor.cpp ParmParse.cpp diff --git a/src/Base/FabArray.cpp b/src/Base/FabArray.cpp new file mode 100644 index 00000000..2ca463bf --- /dev/null +++ b/src/Base/FabArray.cpp @@ -0,0 +1,301 @@ +/* Copyright 2021-2022 The AMReX Community + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#include "pyAMReX.H" + +#include +#include +#include +#include +#include +#include + +#include +#include + + +namespace +{ + template + void make_FabArray_T(py::module &m, std::string const &name) + { + using namespace amrex; + + using FAT = FabArray; + using value_type = typename FAT::value_type; + std::string const full_name = "FabArray_" + name; + py::class_ py_FAT(m, full_name.c_str()); + py_FAT + // define + .def("clear", &FAT::clear) + .def("ok", &FAT::ok) + + .def_property_readonly("arena", &FAT::arena, + "Provides access to the Arena this FabArray was build with.") + .def_property_readonly("has_EB_fab_factory", &FAT::hasEBFabFactory) + .def_property_readonly("factory", &FAT::Factory) + + //.def("array", py::overload_cast< const MFIter& >(&FAT::array)) + //.def("const_array", &FAT::const_array) + .def("array", [](FAT & fa, MFIter const & mfi) + { return fa.array(mfi); }, + // as long as the return value (argument 0) exists, keep the fa (argument 1) alive + py::keep_alive<0, 1>() + ) + .def("const_array", [](FAT & fa, MFIter const & mfi) + { return fa.const_array(mfi); }, + // as long as the return value (argument 0) exists, keep the fa (argument 1) alive + py::keep_alive<0, 1>() + ) + + /* setters */ + .def("set_val", + py::overload_cast< value_type >(&FAT::template setVal), + py::arg("val"), + "Set all components in the entire region of each FAB to val." + ) + .def("set_val", + py::overload_cast< value_type, int, int, int >(&FAT::template setVal), + py::arg("val"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, + "Set the value of num_comp components in the valid region of\n" + "each FAB in the FabArray, starting at component comp to val.\n" + "Also set the value of nghost boundary cells." + ) + .def("set_val", + py::overload_cast< value_type, int, int, IntVect const & >(&FAT::template setVal), + py::arg("val"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost"), + "Set the value of num_comp components in the valid region of\n" + "each FAB in the FabArray, starting at component comp to val.\n" + "Also set the value of nghost boundary cells." + ) + .def("set_val", + py::overload_cast< value_type, Box const &, int, int, int >(&FAT::template setVal), + py::arg("val"), py::arg("region"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, + "Set the value of num_comp components in the valid region of\n" + "each FAB in the FabArray, starting at component comp, as well\n" + "as nghost boundary cells, to val, provided they also intersect\n" + "with the Box region." + ) + .def("set_val", + py::overload_cast< value_type, Box const &, int, int, IntVect const & >(&FAT::template setVal), + py::arg("val"), py::arg("region"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost"), + "Set the value of num_comp components in the valid region of\n" + "each FAB in the FabArray, starting at component comp, as well\n" + "as nghost boundary cells, to val, provided they also intersect\n" + "with the Box region." + ) + + .def("abs", py::overload_cast< int, int, int >(&FAT::template abs), + py::arg("comp"), py::arg("ncomp"), py::arg("nghost")=0 + ) + .def("abs", py::overload_cast< int, int, IntVect const & >(&FAT::template abs), + py::arg("comp"), py::arg("ncomp"), py::arg("nghost") + ) + + .def("saxpy", + [](FAT & dst, value_type a, FAT const & x, int x_comp, int comp, int ncomp, IntVect const & nghost) + { + FAT::Saxpy(dst, a, x, x_comp, comp, ncomp, nghost); + }, + py::arg("a"), py::arg("x"), py::arg("x_comp"), py::arg("comp"), py::arg("ncomp"), py::arg("nghost"), + "self += a * x\n\n" + "Parameters\n" + "----------\n" + "a : scalar a\n" + "x : FabArray x\n" + "x_comp : starting component of x\n" + "comp : starting component of self\n" + "ncomp : number of components\n" + "nghost : number of ghost cells" + ) + .def("xpay", + [](FAT & self, value_type a, FAT const & x, int x_comp, int comp, int ncomp, IntVect const & nghost) + { + FAT::Xpay(self, a, x, x_comp, comp, ncomp, nghost); + }, + py::arg("a"), py::arg("x"), py::arg("xcomp"), py::arg("comp"), py::arg("ncomp"), py::arg("nghost"), + "self = x + a * self\n\n" + "Parameters\n" + "----------\n" + "a : scalar a\n" + "x : FabArray x\n" + "x_comp : starting component of x\n" + "comp : starting component of self\n" + "ncomp : number of components\n" + "nghost : number of ghost cells" + ) + .def("lin_comb", + []( + FAT & dst, + value_type a, FAT const & x, int x_comp, + value_type b, FAT const & y, int y_comp, + int comp, int ncomp, IntVect const & nghost) + { + FAT::LinComb(dst, a, x, x_comp, b, y, y_comp, comp, ncomp, nghost); + }, + py::arg("a"), py::arg("x"), py::arg("xcomp"), + py::arg("b"), py::arg("y"), py::arg("ycomp"), + py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "self = a * x + b * y\n\n" + "Parameters\n" + "----------\n" + "a : float\n" + " scalar a\n" + "x : FabArray\n" + "xcomp : int\n" + " starting component of x\n" + "b : float\n" + " scalar b\n" + "y : FabArray\n" + "ycomp : int\n" + " starting component of y\n" + "comp : int\n" + " starting component of self\n" + "numcomp : int\n" + " number of components\n" + "nghost : int\n" + " number of ghost cells" + ) + + .def("sum", + py::overload_cast< int, IntVect const&, bool >(&FAT::template sum, py::const_), + py::arg("comp"), py::arg("nghost"), py::arg("local"), + "Returns the sum of component \"comp\"" + ) + .def("sum_boundary", + py::overload_cast< Periodicity const & >(&FAT::SumBoundary), + py::arg("period"), + "Sum values in overlapped cells. The destination is limited to valid cells." + ) + .def("sum_boundary", py::overload_cast< int, int, Periodicity const & >(&FAT::SumBoundary), + py::arg("scomp"), py::arg("ncomp"), py::arg("period"), + "Sum values in overlapped cells. The destination is limited to valid cells." + ) + .def("sum_boundary", py::overload_cast< int, int, IntVect const&, Periodicity const & >(&FAT::SumBoundary), + py::arg("scomp"), py::arg("ncomp"), py::arg("nghost"), py::arg("period"), + "Sum values in overlapped cells. The destination is limited to valid cells." + ) + .def("sum_boundary", py::overload_cast< int, int, IntVect const&, IntVect const&, Periodicity const & >(&FAT::SumBoundary), + py::arg("scomp"), py::arg("ncomp"), py::arg("nghost"), py::arg("dst_nghost"), py::arg("period"), + "Sum values in overlapped cells. The destination is limited to valid cells." + ) + ; + + constexpr auto doc_fabarray_osync = R"(Synchronize nodal data. + + The synchronization will override valid regions by the intersecting valid regions with a higher precedence. + The smaller the global box index is, the higher precedence the box has. + With periodic boundaries, for cells in the same box, those near the lower corner have higher precedence than those near the upper corner. + + Parameters + ---------- + scomp : + starting component + ncomp : + number of components + period : + periodic length if it's non-zero)"; + + py_FAT + .def("override_sync", + py::overload_cast< Periodicity const & >(&FAT::OverrideSync), + py::arg("period"), + doc_fabarray_osync + ) + .def("override_sync", + py::overload_cast< int, int, Periodicity const & >(&FAT::OverrideSync), + py::arg("scomp"), py::arg("ncomp"), py::arg("period"), + doc_fabarray_osync + ) + ; + + constexpr auto doc_fabarray_fillb = R"(Copy on intersection within a FabArray. + + Data is copied from valid regions to intersecting regions of definition. + The purpose is to fill in the boundary regions of each FAB in the FabArray. + If cross=true, corner cells are not filled. If the length of periodic is provided, + periodic boundaries are also filled. + + If scomp is provided, this only copies ncomp components starting at scomp. + + Note that FabArray itself does not contains any periodicity information. + FillBoundary expects that its cell-centered version of its BoxArray is non-overlapping.)"; + + py_FAT + .def("fill_boundary", + py::overload_cast< bool >(&FAT::template FillBoundary), + py::arg("cross")=false, + doc_fabarray_fillb + ) + .def("fill_boundary", + py::overload_cast< Periodicity const &, bool >(&FAT::template FillBoundary), + py::arg("period"), + py::arg("cross")=false, + doc_fabarray_fillb + ) + .def("fill_boundary", + py::overload_cast< IntVect const &, Periodicity const &, bool >(&FAT::template FillBoundary), + py::arg("nghost"), + py::arg("period"), + py::arg("cross")=false, + doc_fabarray_fillb + ) + .def("fill_boundary", + py::overload_cast< int, int, bool >(&FAT::template FillBoundary), + py::arg("scomp"), + py::arg("ncomp"), + py::arg("cross")=false, + doc_fabarray_fillb + ) + .def("fill_boundary", + py::overload_cast< int, int, Periodicity const &, bool >(&FAT::template FillBoundary), + py::arg("scomp"), + py::arg("ncomp"), + py::arg("period"), + py::arg("cross")=false, + doc_fabarray_fillb + ) + .def("fill_boundary", + py::overload_cast< int, int, IntVect const &, Periodicity const &, bool >(&FAT::template FillBoundary), + py::arg("scomp"), + py::arg("ncomp"), + py::arg("nghost"), + py::arg("period"), + py::arg("cross")=false, + doc_fabarray_fillb + ) + ; + } +} + +void +init_FabArray(py::module &m) +{ + using namespace amrex; + + py::class_< FabArrayBase >(m, "FabArrayBase") + .def_property_readonly("is_all_cell_centered", &FabArrayBase::is_cell_centered) + .def_property_readonly("is_all_nodal", + py::overload_cast< >(&FabArrayBase::is_nodal, py::const_)) + .def("is_nodal", + py::overload_cast< int >(&FabArrayBase::is_nodal, py::const_)) + + .def_property_readonly("nComp", &FabArrayBase::nComp, + "Return number of variables (aka components) associated with each point.") + .def_property_readonly("num_comp", &FabArrayBase::nComp, + "Return number of variables (aka components) associated with each point.") + .def_property_readonly("size", &FabArrayBase::size, + "Return the number of FABs in the FabArray.") + + .def_property_readonly("n_grow_vect", &FabArrayBase::nGrowVect, + "Return the grow factor (per direction) that defines the region of definition.") + ; + + py::class_< FabFactory >(m, "FabFactory_IArrayBox"); + py::class_< FabFactory >(m, "FabFactory_FArrayBox"); + + make_FabArray_T(m, "IArrayBox"); + make_FabArray_T(m, "FArrayBox"); +} diff --git a/src/Base/MFInfo.cpp b/src/Base/MFInfo.cpp new file mode 100644 index 00000000..7582e197 --- /dev/null +++ b/src/Base/MFInfo.cpp @@ -0,0 +1,46 @@ +/* Copyright 2021-2022 The AMReX Community + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#include "pyAMReX.H" + +#include + + +void init_MFInfo(py::module &m) +{ + using namespace amrex; + + py::class_(m, "MFInfo") + .def_readwrite("alloc", &MFInfo::alloc) + .def_readwrite("arena", &MFInfo::arena) + .def_readwrite("tags", &MFInfo::tags) + + .def(py::init<>()) + + .def("set_alloc", &MFInfo::SetAlloc) + .def("set_arena", &MFInfo::SetArena) + //.def("set_tag", py::overload_cast< std::string >(&MFInfo::SetTag)) + .def("set_tag", [](MFInfo &info, std::string tag) { info.SetTag(std::move(tag)); }); + + py::class_(m, "MFItInfo") + .def_readwrite("do_tiling", &MFItInfo::do_tiling) + .def_readwrite("dynamic", &MFItInfo::dynamic) + .def_readwrite("device_sync", &MFItInfo::device_sync) + .def_readwrite("num_streams", &MFItInfo::num_streams) + .def_readwrite("tilesize", &MFItInfo::tilesize) + + .def(py::init<>()) + + .def("enable_tiling", &MFItInfo::EnableTiling, + py::arg("ts") /*=FabArrayBase::mfiter_tile_size*/ ) + .def("set_dynamic", &MFItInfo::SetDynamic, + py::arg("f")) + .def("disable_device_sync", &MFItInfo::DisableDeviceSync) + .def("set_device_sync", &MFItInfo::SetDeviceSync, + py::arg("f")) + .def("set_num_streams", &MFItInfo::SetNumStreams, + py::arg("n")) + .def("use_default_stream", &MFItInfo::UseDefaultStream); +} diff --git a/src/Base/MultiFab.H b/src/Base/MultiFab.H new file mode 100644 index 00000000..1cd55255 --- /dev/null +++ b/src/Base/MultiFab.H @@ -0,0 +1,684 @@ +/* Copyright 2021-2025 The AMReX Community + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#pragma once + +#include "pyAMReX.H" + +#include +#include +#include +#include +#include + +#include +#include + + +namespace +{ + template + void check_comp (T const & mf, std::string const type, int comp, std::string const name) + { + if (comp < 0 || comp >= mf.nComp()) + throw py::index_error(type + "::" + name + " comp out of bounds"); + } + + template + void check_nghost (T const & mf, std::string const type, int nghost, std::string const name) + { + if (nghost < 0 || nghost > mf.nGrowVect().min()) + throw py::index_error(type + "::" + name + " nghost out of bounds"); + } +} + +/** Helper to create MultiFab and iMultiFab + * + * @tparam PY_T py::class_ from pybind11 + * @param py_MultiFab the pybind11 class + * @param name either "MultiFab" or "iMultiFab" + */ +template +void make_MultiFab(PY_T & py_MultiFab, std::string const name) +{ + using namespace amrex; + + using T = typename PY_T::type; // MultiFab or iMultiFab + using fab_type = typename T::fab_type; // FArrayBox or IArrayBox + using value_type = typename T::value_type; // Real or int + + py_MultiFab + .def("__repr__", + [name](T const & mf) { + return ""; + } + ) + + /* Constructors */ + .def(py::init< >(), + R"(Constructs an empty (i)MultiFab. + + Data can be defined at a later time using the define member functions + inherited from FabArray.)" + ) + .def(py::init< Arena* >(), + py::arg("a"), + R"(Constructs an empty (i)MultiFab. + + Data can be defined at a later time using the define member functions. + If ``define`` is called later with a nullptr as MFInfo's arena, the + default Arena ``a`` will be used. If the arena in MFInfo is not a + nullptr, the MFInfo's arena will be used.)" + ) + ; + + constexpr auto const doc_mf_init = R"(Constructs an (i)MultiFab. + +The size of the FArrayBox is given by the Box grown by \p ngrow, and +the number of components is given by \p ncomp. If \p info is set to +not allocating memory, then no FArrayBoxes are allocated at +this time but can be defined later. + +Parameters +---------- +bxs : + a valid region +dm : + a DistributionMapping +ncomp : + number of components +ngrow : + number of cells the region grows +info : + (i)MultiFab info, including allocation Arena +factory : + FArrayBoxFactory for embedded boundaries)"; + + py_MultiFab + .def(py::init< const BoxArray&, const DistributionMapping&, int, int, + MFInfo const &, FabFactory const & >(), + py::arg("bxs"), py::arg("dm"), py::arg("ncomp"), py::arg("ngrow"), + py::arg("info"), py::arg("factory"), + doc_mf_init + ) + .def(py::init< const BoxArray&, const DistributionMapping&, int, int, + MFInfo const & >(), + py::arg("bxs"), py::arg("dm"), py::arg("ncomp"), py::arg("ngrow"), + py::arg("info"), + doc_mf_init + ) + .def(py::init< const BoxArray&, const DistributionMapping&, int, int>(), + py::arg("bxs"), py::arg("dm"), py::arg("ncomp"), py::arg("ngrow"), + doc_mf_init + ) + + .def(py::init< const BoxArray&, const DistributionMapping&, int, IntVect const&, + MFInfo const& >(), + py::arg("bxs"), py::arg("dm"), py::arg("ncomp"), py::arg("ngrow"), + py::arg("info"), + doc_mf_init + ) + .def(py::init< const BoxArray&, const DistributionMapping&, int, IntVect const&, + MFInfo const&, FabFactory const & >(), + py::arg("bxs"), py::arg("dm"), py::arg("ncomp"), py::arg("ngrow"), + py::arg("info"), py::arg("factory"), + doc_mf_init + ) + .def(py::init< const BoxArray&, const DistributionMapping&, int, IntVect const&>(), + py::arg("bxs"), py::arg("dm"), py::arg("ncomp"), py::arg("ngrow"), + doc_mf_init + ) + + //.def(py::init< T const&, MakeType, int, int >()) + + /* delayed defines */ + //.def("define", + // py::overload_cast< const BoxArray&, const DistributionMapping&, int, int, + // MFInfo const &, FabFactory const & + //>(&T::define)) + //.def("define", + // py::overload_cast< const BoxArray&, const DistributionMapping&, int, + // IntVect const&, MFInfo const &, FabFactory const & + //>(&T::define)) + + /* sizes, etc. */ + .def("min", + [name](T const & mf, int comp, int nghost, bool local) + { + check_comp(mf, name, comp, "min"); + check_nghost(mf, name, nghost, "min"); + return mf.min(comp, nghost, local); + }, + py::arg("comp") = 0, py::arg("nghost") = 0, py::arg("local") = false, + "Returns the minimum value of the specified component of the (i)MultiFab." + ) + .def("min", + [name](T const & mf, Box const & region, int comp, int nghost, bool local) { + check_comp(mf, name, comp, "min"); + check_nghost(mf, name, nghost, "min"); + return mf.min(region, comp, nghost, local); }, + py::arg("region"), py::arg("comp") = 0, py::arg("nghost") = 0, py::arg("local") = false, + "Returns the minimum value of the specified component of the (i)MultiFab over the region." + ) + + .def("max", + [name](T const & mf, int comp, int nghost, bool local) { + check_comp(mf, name, comp, "max"); + check_nghost(mf, name, nghost, "max"); + return mf.max(comp, nghost, local); }, + py::arg("comp") = 0, py::arg("nghost") = 0, py::arg("local") = false, + "Returns the maximum value of the specified component of the (i)MultiFab." + ) + .def("max", + [name](T const & mf, Box const & region, int comp, int nghost, bool local) { + check_comp(mf, name, comp, "max"); + check_nghost(mf, name, nghost, "max"); + return mf.max(region, comp, nghost, local); }, + py::arg("region"), py::arg("comp") = 0, py::arg("nghost") = 0, py::arg("local") = false, + "Returns the maximum value of the specified component of the (i)MultiFab over the region." + ) + + .def("minIndex", &T::minIndex) + .def("maxIndex", &T::maxIndex) + ; + + /* norms */ + if constexpr (std::is_same_v) { + py_MultiFab + .def("norm0", py::overload_cast< int, int, bool, bool >(&T::norm0, py::const_)) + .def("norm0", py::overload_cast< iMultiFab const &, int, int, bool >(&T::norm0, py::const_)) + + .def("norminf", + //py::overload_cast< int, int, bool, bool >(&T::norminf, py::const_) + [](T const & mf, int comp, int nghost, bool local, bool ignore_covered) { + return mf.norminf(comp, nghost, local, ignore_covered); + } + ) + //.def("norminf", py::overload_cast< iMultiFab const &, int, int, bool >(&T::norminf, py::const_)) + + .def("norm1", py::overload_cast< int, Periodicity const&, bool >(&T::norm1, py::const_)) + .def("norm1", py::overload_cast< int, int, bool >(&T::norm1, py::const_)) + .def("norm1", py::overload_cast< Vector const &, int, bool >(&T::norm1, py::const_)) + + .def("norm2", py::overload_cast< int >(&T::norm2, py::const_)) + .def("norm2", py::overload_cast< int, Periodicity const& >(&T::norm2, py::const_)) + .def("norm2", py::overload_cast< Vector const & >(&T::norm2, py::const_)) + ; + } + + py_MultiFab + /* simple math */ + .def("sum", + // py::overload_cast< int, bool >(&T::sum, py::const_), + [](T const & mf, int comp , bool local) { return mf.sum(comp, local); }, + py::arg("comp") = 0, py::arg("local") = false, + "Returns the sum of component 'comp' over the (i)MultiFab -- no ghost cells are included." + ) + .def("sum", + // py::overload_cast< Box const &, int, bool >(&T::sum, py::const_), + [](T const & mf, Box const & region, int comp , bool local) { return mf.sum(region, comp, local); }, + py::arg("region"), py::arg("comp") = 0, py::arg("local") = false, + "Returns the sum of component 'comp' in the given 'region'. -- no ghost cells are included." + ) + ; + + // TODO: Missing in AMReX iMultiFab as of v25.01 + if constexpr (std::is_same_v) { + py_MultiFab + .def("sum_unique", + py::overload_cast< int, bool, Periodicity const& >(&T::sum_unique, py::const_), + py::arg("comp") = 0, + py::arg("local") = false, + py::arg_v("period", Periodicity::NonPeriodic(), "Periodicity.non_periodic()"), + "Same as sum with local=false, but for non-cell-centered data, this" + "skips non-unique points that are owned by multiple boxes." + ) + .def("sum_unique", + py::overload_cast< Box const&, int, bool >(&T::sum_unique, py::const_), + py::arg("region"), + py::arg("comp") = 0, + py::arg("local") = false, + "Returns the unique sum of component `comp` in the given " + "region. Non-unique points owned by multiple boxes in the MultiFab are" + "only added once. No ghost cells are included. This function does not take" + "periodicity into account in the determination of uniqueness of points." + ) + ; + } + + py_MultiFab + .def("plus", + py::overload_cast< value_type, int >(&T::plus), + py::arg("val"), py::arg("nghost")=0, + "Adds the scalar value val to the value of each cell in the\n" + "valid region of each component of the MultiFab. The value\n" + "of nghost specifies the number of cells in the boundary\n" + "region that should be modified." + ) + .def("plus", + py::overload_cast< value_type, int, int, int >(&T::plus), + py::arg("val"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, + "Adds the scalar value \\p val to the value of each cell in the\n" + "specified subregion of the MultiFab.\n\n" + "The subregion consists of the \\p num_comp components starting at component \\p comp.\n" + "The value of nghost specifies the number of cells in the\n" + "boundary region of each FArrayBox in the subregion that should\n" + "be modified." + ) + .def("plus", + py::overload_cast< value_type, const Box&, int >(&T::plus), + py::arg("val"), py::arg("region"), py::arg("nghost")=0, + "Adds the scalar value val to the value of each cell in the\n" + "valid region of each component of the MultiFab, that also\n" + "intersects the Box region. The value of nghost specifies the\n" + "number of cells in the boundary region of each FArrayBox in\n" + "the subregion that should be modified." + ) + .def("plus", + py::overload_cast< value_type, const Box&, int, int, int >(&T::plus), + py::arg("val"), py::arg("region"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, + "Identical to the previous version of plus(), with the\n" + "restriction that the subregion is further constrained to\n" + "the intersection with Box region." + ) + .def("plus", + py::overload_cast< T const &, int, int, int >(&T::plus), + py::arg("mf"), py::arg("strt_comp"), py::arg("num_comp"), py::arg("nghost")=0, + "This function adds the values of the cells in mf to the corresponding\n" + "cells of this MultiFab. mf is required to have the same BoxArray or\n" + "\"valid region\" as this MultiFab. The addition is done only to num_comp\n" + "components, starting with component number strt_comp. The parameter\n" + "nghost specifies the number of boundary cells that will be modified.\n" + "If nghost == 0, only the valid region of each FArrayBox will be\n" + "modified." + ) + + .def("minus", + py::overload_cast< T const &, int, int, int >(&T::minus), + py::arg("mf"), py::arg("strt_comp"), py::arg("num_comp"), py::arg("nghost")=0, + "This function subtracts the values of the cells in mf from the\n" + "corresponding cells of this MultiFab. mf is required to have the\n" + "same BoxArray or \"valid region\" as this MultiFab. The subtraction is\n" + "done only to num_comp components, starting with component number\n" + "strt_comp. The parameter nghost specifies the number of boundary\n" + "cells that will be modified. If nghost == 0, only the valid region of\n" + "each FArrayBox will be modified." + ) + + // renamed: ImportError: overloading a method with both static and instance methods is not supported + .def("divi", + py::overload_cast< T const &, int, int, int >(&T::divide), + py::arg("mf"), py::arg("strt_comp"), py::arg("num_comp"), py::arg("nghost")=0, + "This function divides the values of the cells in mf from the\n" + "corresponding cells of this MultiFab. mf is required to have the\n" + "same BoxArray or \"valid region\" as this MultiFab. The division is\n" + "done only to num_comp components, starting with component number\n" + "strt_comp. The parameter nghost specifies the number of boundary\n" + "cells that will be modified. If nghost == 0, only the valid region of\n" + "each FArrayBox will be modified. Note, nothing is done to protect\n" + "against divide by zero." + ) + + .def("mult", + py::overload_cast< value_type, int >(&T::mult), + py::arg("val"), py::arg("nghost")=0, + "Scales the value of each cell in the valid region of each\n" + "component of the MultiFab by the scalar val (a[i] <- a[i]*val).\n" + "The value of nghost specifies the number of cells in the\n" + "boundary region that should be modified." + ) + .def("mult", + py::overload_cast< value_type, int, int, int >(&T::mult), + py::arg("val"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, + "Scales the value of each cell in the specified subregion of the\n" + "MultiFab by the scalar val (a[i] <- a[i]*val). The subregion\n" + "consists of the num_comp components starting at component comp.\n" + "The value of nghost specifies the number of cells in the\n" + "boundary region of each FArrayBox in the subregion that should\n" + "be modified." + ) + .def("mult", + py::overload_cast< value_type, Box const &, int, int, int >(&T::mult), + py::arg("val"), py::arg("region"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, + "Identical to the previous version of mult(), with the\n" + "restriction that the subregion is further constrained to the\n" + "intersection with Box region. The value of nghost specifies the\n" + "number of cells in the boundary region of each FArrayBox in\n" + "the subregion that should be modified." + ) + .def("mult", + py::overload_cast< value_type, Box const &, int >(&T::mult), + py::arg("val"), py::arg("region"), py::arg("nghost")=0, + "Scales the value of each cell in the valid region of each\n" + "component of the MultiFab by the scalar val (a[i] <- a[i]*val),\n" + "that also intersects the Box region. The value of nghost\n" + "specifies the number of cells in the boundary region of each\n" + "FArrayBox in the subregion that should be modified." + ) + ; + + if constexpr (std::is_same_v) { + py_MultiFab + .def("invert", + py::overload_cast< value_type, int >(&T::invert), + py::arg("numerator"), py::arg("nghost"), + "Replaces the value of each cell in the specified subregion of\n" + "the MultiFab with its reciprocal multiplied by the value of\n" + "numerator. The value of nghost specifies the number of cells\n" + "in the boundary region that should be modified." + ) + .def("invert", + py::overload_cast< value_type, int, int, int >(&T::invert), + py::arg("numerator"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, + "Replaces the value of each cell in the specified subregion of\n" + "the MultiFab with its reciprocal multiplied by the value of\n" + "numerator. The subregion consists of the num_comp components\n" + "starting at component comp. The value of nghost specifies the\n" + "number of cells in the boundary region of each FArrayBox in the\n" + "subregion that should be modified." + ) + .def("invert", + py::overload_cast< value_type, Box const &, int >(&T::invert), + py::arg("numerator"), py::arg("region"), py::arg("nghost"), + "Scales the value of each cell in the valid region of each\n" + "component of the MultiFab by the scalar val (a[i] <- a[i]*val),\n" + "that also intersects the Box region. The value of nghost\n" + "specifies the number of cells in the boundary region of each\n" + "FArrayBox in the subregion that should be modified." + ) + .def("invert", + py::overload_cast< value_type, Box const &, int, int, int >(&T::invert), + py::arg("numerator"), py::arg("region"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, + "Identical to the previous version of invert(), with the\n" + "restriction that the subregion is further constrained to the\n" + "intersection with Box region. The value of nghost specifies the\n" + "number of cells in the boundary region of each FArrayBox in the\n" + "subregion that should be modified." + ) + ; + } + + py_MultiFab + .def("negate", + py::overload_cast< int >(&T::negate), + py::arg("nghost")=0, + "Negates the value of each cell in the valid region of\n" + "the MultiFab. The value of nghost specifies the number of\n" + "cells in the boundary region that should be modified." + ) + .def("negate", + py::overload_cast< int, int, int >(&T::negate), + py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, + "Negates the value of each cell in the specified subregion of\n" + "the MultiFab. The subregion consists of the num_comp\n" + "components starting at component comp. The value of nghost\n" + "specifies the number of cells in the boundary region of each\n" + "FArrayBox in the subregion that should be modified." + ) + .def("negate", + py::overload_cast< Box const &, int >(&T::negate), + py::arg("region"), py::arg("nghost")=0, + "Negates the value of each cell in the valid region of\n" + "the MultiFab that also intersects the Box region. The value\n" + "of nghost specifies the number of cells in the boundary region\n" + "that should be modified." + ) + .def("negate", + py::overload_cast< Box const &, int, int, int >(&T::negate), + py::arg("region"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, + "Identical to the previous version of negate(), with the\n" + "restriction that the subregion is further constrained to\n" + "the intersection with Box region." + ) + ; + + /* static (standalone) simple math functions */ + if constexpr (std::is_same_v) { + py_MultiFab + .def("dot", + [](T const &self, int comp, T const &y, int y_comp, int numcomp, int nghost, bool local) { + return T::Dot(self, comp, y, y_comp, numcomp, nghost, local); + }, + py::arg("comp"), + py::arg("y"), py::arg("y_comp"), + py::arg("numcomp"), py::arg("nghost"), py::arg("local") = false, + "Returns the dot product of self with another MultiFab." + ) + .def("dot", + [](T const &self, int comp, int numcomp, int nghost, bool local) { + return T::Dot(self, comp, numcomp, nghost, local); + }, + py::arg("comp"), + py::arg("numcomp"), py::arg("nghost"), py::arg("local") = false, + "Returns the dot product with itself." + ) + .def("dot", + [](T const &self, const iMultiFab &mask, int comp, MultiFab const &y, int y_comp, int numcomp, + int nghost, bool local) { + return T::Dot(mask, self, comp, y, y_comp, numcomp, nghost, local); + }, + py::arg("mask"), py::arg("comp"), py::arg("y"), py::arg("y_comp"), + py::arg("numcomp"), py::arg("nghost"), py::arg("local") = false, + "Returns the dot product of self with another MultiFab where the mask is valid." + ) + ; + } + + py_MultiFab + .def("add", + [](T & self, T const & src, int srccomp, int comp, int numcomp, int nghost) { + T::Add(self, src, srccomp, comp, numcomp, nghost); + }, + py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "Add src to self including nghost ghost cells.\n" + "The two MultiFabs MUST have the same underlying BoxArray." + ) + ; + + // TODO: Missing in iMultiFab + if constexpr (std::is_same_v) { + py_MultiFab + .def("add", + [](T &self, T const &src, int srccomp, int comp, int numcomp, IntVect const &nghost) { + T::Add(self, src, srccomp, comp, numcomp, nghost); + }, + py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "Add src to self including nghost ghost cells.\n" + "The two MultiFabs MUST have the same underlying BoxArray." + ) + ; + } + + py_MultiFab + .def("subtract", + [](T & self, T const & src, int srccomp, int comp, int numcomp, int nghost) { + T::Subtract(self, src, srccomp, comp, numcomp, nghost); + }, + py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "Subtract src from self including nghost ghost cells.\n" + "The two MultiFabs MUST have the same underlying BoxArray." + ) + ; + + // TODO: Missing in iMultiFab + if constexpr (std::is_same_v) { + py_MultiFab + .def("subtract", + [](T & self, T const & src, int srccomp, int comp, int numcomp, IntVect const & nghost) { + T::Subtract(self, src, srccomp, comp, numcomp, nghost); + }, + py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "Subtract src from self including nghost ghost cells.\n" + "The two MultiFabs MUST have the same underlying BoxArray." + ) + ; + } + + py_MultiFab + .def("multiply", + [](T & self, T const & src, int srccomp, int comp, int numcomp, int nghost) { + T::Multiply(self, src, srccomp, comp, numcomp, nghost); + }, + py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "Multiply self by src including nghost ghost cells.\n" + "The two MultiFabs MUST have the same underlying BoxArray." + ) + ; + + // TODO: Missing in iMultiFab + if constexpr (std::is_same_v) { + py_MultiFab + .def("multiply", + [](T &self, T const &src, int srccomp, int comp, int numcomp, IntVect const &nghost) { + T::Multiply(self, src, srccomp, comp, numcomp, nghost); + }, + py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "Multiply self by src including nghost ghost cells.\n" + "The two MultiFabs MUST have the same underlying BoxArray." + ); + } + + py_MultiFab + .def("divide", + [](T & self, T const & src, int srccomp, int comp, int numcomp, int nghost) { + T::Divide(self, src, srccomp, comp, numcomp, nghost); + }, + py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "Divide self by src including nghost ghost cells.\n" + "The two MultiFabs MUST have the same underlying BoxArray." + ) + ; + + if constexpr (std::is_same_v) { + py_MultiFab + .def("divide", /* TODO: Missing in iMultiFab */ + [](T &self, T const &src, int srccomp, int comp, int numcomp, IntVect const &nghost) { + T::Divide(self, src, srccomp, comp, numcomp, nghost); + }, + py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "Divide self by src including nghost ghost cells.\n" + "The two MultiFabs MUST have the same underlying BoxArray." + ) + + .def("swap", /* TODO: Missing in iMultiFab */ + [](T &self, T &src, int srccomp, int comp, int numcomp, int nghost) { + T::Swap(self, src, srccomp, comp, numcomp, nghost); + }, + py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "Swap from src to self including nghost ghost cells.\n" + "The two MultiFabs MUST have the same underlying BoxArray.\n" + "The swap is local." + ) + .def("swap", /* TODO: Missing in iMultiFab */ + [](T &self, T &src, int srccomp, int comp, int numcomp, IntVect const &nghost) { + T::Swap(self, src, srccomp, comp, numcomp, nghost); + }, + py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "Swap from src to self including nghost ghost cells.\n" + "The two MultiFabs MUST have the same underlying BoxArray.\n" + "The swap is local." + ) + + .def("saxpy", + [](T & self, value_type a, T const & src, int srccomp, int comp, int numcomp, int nghost) { + T::Saxpy(self, a, src, srccomp, comp, numcomp, nghost); + }, + py::arg("a"), py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "self += a * src" + ) + + .def("xpay", + [](T & self, value_type a, T const & src, int srccomp, int comp, int numcomp, int nghost) { + T::Xpay(self, a, src, srccomp, comp, numcomp, nghost); + }, + py::arg("a"), py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "self = src + a * self" + ) + + .def("lin_comb", + [](T & self, value_type a, T const & x, int x_comp, value_type b, T const & y, int y_comp, int comp, int numcomp, int nghost) { + T::LinComb(self, a, x, x_comp, b, y, y_comp, comp, numcomp, nghost); + }, + py::arg("a"), py::arg("x"), py::arg("x_comp"), + py::arg("b"), py::arg("y"), py::arg("y_comp"), + py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "self = a * x + b * y" + ) + + .def("add_product", + [](T & self, MultiFab const & src1, int comp1, MultiFab const & src2, int comp2, int comp, int numcomp, int nghost) { + T::AddProduct(self, src1, comp1, src2, comp2, comp, numcomp, nghost); + }, + py::arg("src1"), py::arg("comp1"), + py::arg("src2"), py::arg("comp2"), + py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "self += src1 * src2" + ) + .def("add_product", + [](T & self, MultiFab const & src1, int comp1, MultiFab const & src2, int comp2, int comp, int numcomp, IntVect const & nghost) { + T::AddProduct(self, src1, comp1, src2, comp2, comp, numcomp, nghost); + }, + py::arg("src1"), py::arg("comp1"), + py::arg("src2"), py::arg("comp2"), + py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), + "self += src1 * src2" + ) + + /* simple data validity checks */ + .def("contains_nan", + py::overload_cast< bool >(&T::contains_nan, py::const_), + py::arg("local")=false + ) + .def("contains_nan", + py::overload_cast< int, int, int, bool >(&T::contains_nan, py::const_), + py::arg("scomp"), py::arg("ncomp"), py::arg("ngrow")=0, py::arg("local")=false + ) + .def("contains_nan", + py::overload_cast< int, int, IntVect const &, bool >(&T::contains_nan, py::const_), + py::arg("scomp"), py::arg("ncomp"), py::arg("ngrow"), py::arg("local")=false + ) + + .def("contains_inf", + py::overload_cast< bool >(&T::contains_inf, py::const_), + py::arg("local")=false + ) + .def("contains_inf", + py::overload_cast< int, int, int, bool >(&T::contains_inf, py::const_), + py::arg("scomp"), py::arg("ncomp"), py::arg("ngrow")=0, py::arg("local")=false + ) + .def("contains_inf", + py::overload_cast< int, int, IntVect const &, bool >(&T::contains_inf, py::const_), + py::arg("scomp"), py::arg("ncomp"), py::arg("ngrow"), py::arg("local")=false + ) + ; + } + + py_MultiFab + .def("box_array", &T::boxArray) + .def("dm", &T::DistributionMap) + .def_property_readonly("n_comp", &T::nComp) + .def_property_readonly("n_grow_vect", &T::nGrowVect) + + /* masks & ownership */ + // TODO: + // - OverlapMask -> std::unique_ptr + // - OwnerMask -> std::unique_ptr + ; + + /* Syncs */ + if constexpr (std::is_same_v) { + py_MultiFab + .def("average_sync", &T::AverageSync) + .def("weighted_sync", &T::WeightedSync) + .def("override_sync", py::overload_cast(&T::OverrideSync)); + } + + py_MultiFab + /* Init & Finalize */ + .def_static("initialize", &T::Initialize) + .def_static("finalize", &T::Finalize) + ; +} diff --git a/src/Base/MultiFab.cpp b/src/Base/MultiFab.cpp index 4766557a..e95f081b 100644 --- a/src/Base/MultiFab.cpp +++ b/src/Base/MultiFab.cpp @@ -4,77 +4,25 @@ * License: BSD-3-Clause-LBNL */ #include "pyAMReX.H" +#include "MultiFab.H" #include -#include #include #include #include -#include #include #include -#include #include -namespace { - void check_comp(amrex::MultiFab const & mf, const int comp, std::string const name) - { - if (comp < 0 || comp >= mf.nComp()) - throw py::index_error("MultiFab::" + name + " comp out of bounds"); - } - void check_nghost(amrex::MultiFab const & mf, const int nghost, std::string const name) - { - if (nghost < 0 || nghost > mf.nGrowVect().min()) - throw py::index_error("MultiFab::" + name + " nghost out of bounds"); - } -} -void init_MultiFab(py::module &m) +void init_MultiFab(py::module &m, py::class_< amrex::MFIter > & py_MFIter) { using namespace amrex; - py::class_< FabArrayBase > py_FabArrayBase(m, "FabArrayBase"); - py::class_< FabArray, FabArrayBase > py_FabArray_FArrayBox(m, "FabArray_FArrayBox"); py::class_< MultiFab, FabArray > py_MultiFab(m, "MultiFab"); - py::class_< FabFactory >(m, "FabFactory_FArrayBox"); - - py::class_< MFInfo >(m, "MFInfo") - .def_readwrite("alloc", &MFInfo::alloc) - .def_readwrite("arena", &MFInfo::arena) - .def_readwrite("tags", &MFInfo::tags) - - .def(py::init< >()) - - .def("set_alloc", &MFInfo::SetAlloc) - .def("set_arena", &MFInfo::SetArena) - //.def("set_tag", py::overload_cast< std::string >(&MFInfo::SetTag)) - .def("set_tag", [](MFInfo & info, std::string tag) { info.SetTag(std::move(tag)); }) - ; - - py::class_< MFItInfo >(m, "MFItInfo") - .def_readwrite("do_tiling", &MFItInfo::do_tiling) - .def_readwrite("dynamic", &MFItInfo::dynamic) - .def_readwrite("device_sync", &MFItInfo::device_sync) - .def_readwrite("num_streams", &MFItInfo::num_streams) - .def_readwrite("tilesize", &MFItInfo::tilesize) - - .def(py::init< >()) - - .def("enable_tiling", &MFItInfo::EnableTiling, - py::arg("ts") /*=FabArrayBase::mfiter_tile_size*/ ) - .def("set_dynamic", &MFItInfo::SetDynamic, - py::arg("f")) - .def("disable_device_sync", &MFItInfo::DisableDeviceSync) - .def("set_device_sync", &MFItInfo::SetDeviceSync, - py::arg("f")) - .def("set_num_streams", &MFItInfo::SetNumStreams, - py::arg("n")) - .def("use_default_stream", &MFItInfo::UseDefaultStream) - ; - - py::class_< MFIter >(m, "MFIter", py::dynamic_attr()) + py_MFIter .def("__repr__", [](MFIter const & mfi) { std::string r = "()) - //.def(py::init< iMultiFab const & >()) - //.def(py::init< iMultiFab const &, MFItInfo const & >()) + .def(py::init< iMultiFab const & >()) + .def(py::init< iMultiFab const &, MFItInfo const & >()) // helpers for iteration __next__ .def("_incr", &MFIter::operator++) @@ -131,265 +79,6 @@ void init_MultiFab(py::module &m) .def_property_readonly("length", &MFIter::length) ; - py_FabArrayBase - .def_property_readonly("is_all_cell_centered", &FabArrayBase::is_cell_centered) - .def_property_readonly("is_all_nodal", - py::overload_cast< >(&FabArrayBase::is_nodal, py::const_)) - .def("is_nodal", - py::overload_cast< int >(&FabArrayBase::is_nodal, py::const_)) - - .def_property_readonly("nComp", &FabArrayBase::nComp, - "Return number of variables (aka components) associated with each point.") - .def_property_readonly("num_comp", &FabArrayBase::nComp, - "Return number of variables (aka components) associated with each point.") - .def_property_readonly("size", &FabArrayBase::size, - "Return the number of FABs in the FabArray.") - - .def_property_readonly("n_grow_vect", &FabArrayBase::nGrowVect, - "Return the grow factor (per direction) that defines the region of definition.") - ; - - py_FabArray_FArrayBox - // define - .def("clear", &FabArray::clear) - .def("ok", &FabArray::ok) - - .def_property_readonly("arena", &FabArray::arena, - "Provides access to the Arena this FabArray was build with.") - .def_property_readonly("has_EB_fab_factory", &FabArray::hasEBFabFactory) - .def_property_readonly("factory", &FabArray::Factory) - - //.def("array", py::overload_cast< const MFIter& >(&FabArray::array)) - //.def("const_array", &FabArray::const_array) - .def("array", [](FabArray & fa, MFIter const & mfi) - { return fa.array(mfi); }, - // as long as the return value (argument 0) exists, keep the fa (argument 1) alive - py::keep_alive<0, 1>() - ) - .def("const_array", [](FabArray & fa, MFIter const & mfi) - { return fa.const_array(mfi); }, - // as long as the return value (argument 0) exists, keep the fa (argument 1) alive - py::keep_alive<0, 1>() - ) - - /* setters */ - .def("set_val", - py::overload_cast< amrex::Real >(&FabArray::setVal), - py::arg("val"), - "Set all components in the entire region of each FAB to val." - ) - .def("set_val", - py::overload_cast< amrex::Real, int, int, int >(&FabArray::setVal), - py::arg("val"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, - "Set the value of num_comp components in the valid region of\n" - "each FAB in the FabArray, starting at component comp to val.\n" - "Also set the value of nghost boundary cells." - ) - .def("set_val", - py::overload_cast< amrex::Real, int, int, IntVect const & >(&FabArray::setVal), - py::arg("val"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost"), - "Set the value of num_comp components in the valid region of\n" - "each FAB in the FabArray, starting at component comp to val.\n" - "Also set the value of nghost boundary cells." - ) - .def("set_val", - py::overload_cast< amrex::Real, Box const &, int, int, int >(&FabArray::setVal), - py::arg("val"), py::arg("region"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, - "Set the value of num_comp components in the valid region of\n" - "each FAB in the FabArray, starting at component comp, as well\n" - "as nghost boundary cells, to val, provided they also intersect\n" - "with the Box region." - ) - .def("set_val", - py::overload_cast< amrex::Real, Box const &, int, int, IntVect const & >(&FabArray::setVal), - py::arg("val"), py::arg("region"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost"), - "Set the value of num_comp components in the valid region of\n" - "each FAB in the FabArray, starting at component comp, as well\n" - "as nghost boundary cells, to val, provided they also intersect\n" - "with the Box region." - ) - - .def("abs", py::overload_cast< int, int, int >(&FabArray::abs), - py::arg("comp"), py::arg("ncomp"), py::arg("nghost")=0 - ) - .def("abs", py::overload_cast< int, int, IntVect const & >(&FabArray::abs), - py::arg("comp"), py::arg("ncomp"), py::arg("nghost") - ) - - .def("saxpy", - [](FabArray & dst, Real a, FabArray const & x, int x_comp, int comp, int ncomp, IntVect const & nghost) - { - FabArray::Saxpy(dst, a, x, x_comp, comp, ncomp, nghost); - }, - py::arg("a"), py::arg("x"), py::arg("x_comp"), py::arg("comp"), py::arg("ncomp"), py::arg("nghost"), - "self += a * x\n\n" - "Parameters\n" - "----------\n" - "a : scalar a\n" - "x : FabArray x\n" - "x_comp : starting component of x\n" - "comp : starting component of self\n" - "ncomp : number of components\n" - "nghost : number of ghost cells" - ) - .def("xpay", - [](FabArray & self, Real a, FabArray const & x, int x_comp, int comp, int ncomp, IntVect const & nghost) - { - FabArray::Xpay(self, a, x, x_comp, comp, ncomp, nghost); - }, - py::arg("a"), py::arg("x"), py::arg("xcomp"), py::arg("comp"), py::arg("ncomp"), py::arg("nghost"), - "self = x + a * self\n\n" - "Parameters\n" - "----------\n" - "a : scalar a\n" - "x : FabArray x\n" - "x_comp : starting component of x\n" - "comp : starting component of self\n" - "ncomp : number of components\n" - "nghost : number of ghost cells" - ) - .def("lin_comb", - []( - FabArray & dst, - Real a, FabArray const & x, int x_comp, - Real b, FabArray const & y, int y_comp, - int comp, int ncomp, IntVect const & nghost) - { - FabArray::LinComb(dst, a, x, x_comp, b, y, y_comp, comp, ncomp, nghost); - }, - py::arg("a"), py::arg("x"), py::arg("xcomp"), - py::arg("b"), py::arg("y"), py::arg("ycomp"), - py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "self = a * x + b * y\n\n" - "Parameters\n" - "----------\n" - "a : float\n" - " scalar a\n" - "x : FabArray\n" - "xcomp : int\n" - " starting component of x\n" - "b : float\n" - " scalar b\n" - "y : FabArray\n" - "ycomp : int\n" - " starting component of y\n" - "comp : int\n" - " starting component of self\n" - "numcomp : int\n" - " number of components\n" - "nghost : int\n" - " number of ghost cells" - ) - - .def("sum", - py::overload_cast< int, IntVect const&, bool >(&FabArray::template sum, py::const_), - py::arg("comp"), py::arg("nghost"), py::arg("local"), - "Returns the sum of component \"comp\"" - ) - .def("sum_boundary", - py::overload_cast< Periodicity const & >(&FabArray::SumBoundary), - py::arg("period"), - "Sum values in overlapped cells. The destination is limited to valid cells." - ) - .def("sum_boundary", py::overload_cast< int, int, Periodicity const & >(&FabArray::SumBoundary), - py::arg("scomp"), py::arg("ncomp"), py::arg("period"), - "Sum values in overlapped cells. The destination is limited to valid cells." - ) - .def("sum_boundary", py::overload_cast< int, int, IntVect const&, Periodicity const & >(&FabArray::SumBoundary), - py::arg("scomp"), py::arg("ncomp"), py::arg("nghost"), py::arg("period"), - "Sum values in overlapped cells. The destination is limited to valid cells." - ) - .def("sum_boundary", py::overload_cast< int, int, IntVect const&, IntVect const&, Periodicity const & >(&FabArray::SumBoundary), - py::arg("scomp"), py::arg("ncomp"), py::arg("nghost"), py::arg("dst_nghost"), py::arg("period"), - "Sum values in overlapped cells. The destination is limited to valid cells." - ) - ; - - constexpr auto doc_fabarray_fillb = R"(Copy on intersection within a FabArray. - - Data is copied from valid regions to intersecting regions of definition. - The purpose is to fill in the boundary regions of each FAB in the FabArray. - If cross=true, corner cells are not filled. If the length of periodic is provided, - periodic boundaries are also filled. - - If scomp is provided, this only copies ncomp components starting at scomp. - - Note that FabArray itself does not contains any periodicity information. - FillBoundary expects that its cell-centered version of its BoxArray is non-overlapping.)"; - - py_FabArray_FArrayBox - .def("fill_boundary", - py::overload_cast< bool >(&FabArray::template FillBoundary), - py::arg("cross")=false, - doc_fabarray_fillb - ) - .def("fill_boundary", - py::overload_cast< Periodicity const &, bool >(&FabArray::template FillBoundary), - py::arg("period"), - py::arg("cross")=false, - doc_fabarray_fillb - ) - .def("fill_boundary", - py::overload_cast< IntVect const &, Periodicity const &, bool >(&FabArray::template FillBoundary), - py::arg("nghost"), - py::arg("period"), - py::arg("cross")=false, - doc_fabarray_fillb - ) - .def("fill_boundary", - py::overload_cast< int, int, bool >(&FabArray::template FillBoundary), - py::arg("scomp"), - py::arg("ncomp"), - py::arg("cross")=false, - doc_fabarray_fillb - ) - .def("fill_boundary", - py::overload_cast< int, int, Periodicity const &, bool >(&FabArray::template FillBoundary), - py::arg("scomp"), - py::arg("ncomp"), - py::arg("period"), - py::arg("cross")=false, - doc_fabarray_fillb - ) - .def("fill_boundary", - py::overload_cast< int, int, IntVect const &, Periodicity const &, bool >(&FabArray::template FillBoundary), - py::arg("scomp"), - py::arg("ncomp"), - py::arg("nghost"), - py::arg("period"), - py::arg("cross")=false, - doc_fabarray_fillb - ) - ; - - constexpr auto doc_fabarray_osync = R"(Synchronize nodal data. - - The synchronization will override valid regions by the intersecting valid regions with a higher precedence. - The smaller the global box index is, the higher precedence the box has. - With periodic boundaries, for cells in the same box, those near the lower corner have higher precedence than those near the upper corner. - - Parameters - ---------- - scomp : - starting component - ncomp : - number of components - period : - periodic length if it's non-zero)"; - - py_FabArray_FArrayBox - .def("override_sync", - py::overload_cast< Periodicity const & >(&FabArray::OverrideSync), - py::arg("period"), - doc_fabarray_osync - ) - .def("override_sync", - py::overload_cast< int, int, Periodicity const & >(&FabArray::OverrideSync), - py::arg("scomp"), py::arg("ncomp"), py::arg("period"), - doc_fabarray_osync - ) - ; - m.def("htod_memcpy", py::overload_cast< FabArray &, FabArray const & >(&htod_memcpy), py::arg("dest"), py::arg("src"), @@ -412,578 +101,7 @@ void init_MultiFab(py::module &m) "Copy from a device to host FabArray for a specific (number of) component(s)." ); - py_MultiFab - .def("__repr__", - [](MultiFab const & mf) { - return ""; - } - ) - - /* Constructors */ - .def(py::init< >(), - R"(Constructs an empty MultiFab. - - Data can be defined at a later time using the define member functions - inherited from FabArray.)" - ) - .def(py::init< Arena* >(), - py::arg("a"), - R"(Constructs an empty MultiFab. - - Data can be defined at a later time using the define member functions. - If ``define`` is called later with a nullptr as MFInfo's arena, the - default Arena ``a`` will be used. If the arena in MFInfo is not a - nullptr, the MFInfo's arena will be used.)" - ) - ; - - constexpr auto doc_mf_init = R"(Constructs a MultiFab. - - The size of the FArrayBox is given by the Box grown by \p ngrow, and - the number of components is given by \p ncomp. If \p info is set to - not allocating memory, then no FArrayBoxes are allocated at - this time but can be defined later. - - Parameters - ---------- - bxs : - a valid region - dm : - a DistribuionMapping - ncomp : - number of components - ngrow : - number of cells the region grows - info : - MultiFab info, including allocation Arena - factory : - FArrayBoxFactory for embedded boundaries)"; - - py_MultiFab - .def(py::init< const BoxArray&, const DistributionMapping&, int, int, - MFInfo const &, FabFactory const & >(), - py::arg("bxs"), py::arg("dm"), py::arg("ncomp"), py::arg("ngrow"), - py::arg("info"), py::arg("factory"), - doc_mf_init - ) - .def(py::init< const BoxArray&, const DistributionMapping&, int, int, - MFInfo const & >(), - py::arg("bxs"), py::arg("dm"), py::arg("ncomp"), py::arg("ngrow"), - py::arg("info"), - doc_mf_init - ) - .def(py::init< const BoxArray&, const DistributionMapping&, int, int>(), - py::arg("bxs"), py::arg("dm"), py::arg("ncomp"), py::arg("ngrow"), - doc_mf_init - ) - - .def(py::init< const BoxArray&, const DistributionMapping&, int, IntVect const&, - MFInfo const& >(), - py::arg("bxs"), py::arg("dm"), py::arg("ncomp"), py::arg("ngrow"), - py::arg("info"), - doc_mf_init - ) - .def(py::init< const BoxArray&, const DistributionMapping&, int, IntVect const&, - MFInfo const&, FabFactory const & >(), - py::arg("bxs"), py::arg("dm"), py::arg("ncomp"), py::arg("ngrow"), - py::arg("info"), py::arg("factory"), - doc_mf_init - ) - .def(py::init< const BoxArray&, const DistributionMapping&, int, IntVect const&>(), - py::arg("bxs"), py::arg("dm"), py::arg("ncomp"), py::arg("ngrow"), - doc_mf_init - ) - - //.def(py::init< MultiFab const&, MakeType, int, int >()) - - /* delayed defines */ - //.def("define", - // py::overload_cast< const BoxArray&, const DistributionMapping&, int, int, - // MFInfo const &, FabFactory const & - //>(&MultiFab::define)) - //.def("define", - // py::overload_cast< const BoxArray&, const DistributionMapping&, int, - // IntVect const&, MFInfo const &, FabFactory const & - //>(&MultiFab::define)) - - /* sizes, etc. */ - .def("min", - [](MultiFab const & mf, int comp, int nghost, bool local) { - check_comp(mf, comp, "min"); - check_nghost(mf, nghost, "min"); - return mf.min(comp, nghost, local); }, - py::arg("comp") = 0, py::arg("nghost") = 0, py::arg("local") = false, - "Returns the minimum value of the specfied component of the MultiFab." - ) - .def("min", - [](MultiFab const & mf, Box const & region, int comp, int nghost, bool local) { - check_comp(mf, comp, "min"); - check_nghost(mf, nghost, "min"); - return mf.min(region, comp, nghost, local); }, - py::arg("region"), py::arg("comp") = 0, py::arg("nghost") = 0, py::arg("local") = false, - "Returns the minimum value of the specfied component of the MultiFab over the region." - ) - - .def("max", - [](MultiFab const & mf, int comp, int nghost, bool local) { - check_comp(mf, comp, "max"); - check_nghost(mf, nghost, "max"); - return mf.max(comp, nghost, local); }, - py::arg("comp") = 0, py::arg("nghost") = 0, py::arg("local") = false, - "Returns the maximum value of the specfied component of the MultiFab." - ) - .def("max", - [](MultiFab const & mf, Box const & region, int comp, int nghost, bool local) { - check_comp(mf, comp, "max"); - check_nghost(mf, nghost, "max"); - return mf.max(region, comp, nghost, local); }, - py::arg("region"), py::arg("comp") = 0, py::arg("nghost") = 0, py::arg("local") = false, - "Returns the maximum value of the specfied component of the MultiFab over the region." - ) - - .def("minIndex", &MultiFab::minIndex) - .def("maxIndex", &MultiFab::maxIndex) - - /* norms */ - .def("norm0", py::overload_cast< int, int, bool, bool >(&MultiFab::norm0, py::const_)) - //.def("norm0", py::overload_cast< iMultiFab const &, int, int, bool >(&MultiFab::norm0, py::const_)) - - .def("norminf", - //py::overload_cast< int, int, bool, bool >(&MultiFab::norminf, py::const_) - [](MultiFab const & mf, int comp, int nghost, bool local, bool ignore_covered) { - return mf.norminf(comp, nghost, local, ignore_covered); - } - ) - //.def("norminf", py::overload_cast< iMultiFab const &, int, int, bool >(&MultiFab::norminf, py::const_)) - - .def("norm1", py::overload_cast< int, Periodicity const&, bool >(&MultiFab::norm1, py::const_)) - .def("norm1", py::overload_cast< int, int, bool >(&MultiFab::norm1, py::const_)) - .def("norm1", py::overload_cast< Vector const &, int, bool >(&MultiFab::norm1, py::const_)) - - .def("norm2", py::overload_cast< int >(&MultiFab::norm2, py::const_)) - .def("norm2", py::overload_cast< int, Periodicity const& >(&MultiFab::norm2, py::const_)) - .def("norm2", py::overload_cast< Vector const & >(&MultiFab::norm2, py::const_)) - - /* simple math */ - - .def("sum", - // py::overload_cast< int, bool >(&MultiFab::sum, py::const_), - [](MultiFab const & mf, int comp , bool local) { return mf.sum(comp, local); }, - py::arg("comp") = 0, py::arg("local") = false, - "Returns the sum of component 'comp' over the MultiFab -- no ghost cells are included." - ) - .def("sum", - // py::overload_cast< Box const &, int, bool >(&MultiFab::sum, py::const_), - [](MultiFab const & mf, Box const & region, int comp , bool local) { return mf.sum(region, comp, local); }, - py::arg("region"), py::arg("comp") = 0, py::arg("local") = false, - "Returns the sum of component 'comp' in the given 'region'. -- no ghost cells are included." - ) - .def("sum_unique", - py::overload_cast< int, bool, Periodicity const& >(&MultiFab::sum_unique, py::const_), - py::arg("comp") = 0, - py::arg("local") = false, - py::arg_v("period", Periodicity::NonPeriodic(), "Periodicity.non_periodic()"), - "Same as sum with local=false, but for non-cell-centered data, this" - "skips non-unique points that are owned by multiple boxes." - ) - .def("sum_unique", - py::overload_cast< Box const&, int, bool >(&MultiFab::sum_unique, py::const_), - py::arg("region"), - py::arg("comp") = 0, - py::arg("local") = false, - "Returns the unique sum of component `comp` in the given " - "region. Non-unique points owned by multiple boxes in the MultiFab are" - "only added once. No ghost cells are included. This function does not take" - "periodicity into account in the determination of uniqueness of points." - ) - - .def("plus", - py::overload_cast< Real, int >(&MultiFab::plus), - py::arg("val"), py::arg("nghost")=0, - "Adds the scalar value val to the value of each cell in the\n" - "valid region of each component of the MultiFab. The value\n" - "of nghost specifies the number of cells in the boundary\n" - "region that should be modified." - ) - .def("plus", - py::overload_cast< Real, int, int, int >(&MultiFab::plus), - py::arg("val"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, - "Adds the scalar value \\p val to the value of each cell in the\n" - "specified subregion of the MultiFab.\n\n" - "The subregion consists of the \\p num_comp components starting at component \\p comp.\n" - "The value of nghost specifies the number of cells in the\n" - "boundary region of each FArrayBox in the subregion that should\n" - "be modified." - ) - .def("plus", - py::overload_cast< Real, const Box&, int >(&MultiFab::plus), - py::arg("val"), py::arg("region"), py::arg("nghost")=0, - "Adds the scalar value val to the value of each cell in the\n" - "valid region of each component of the MultiFab, that also\n" - "intersects the Box region. The value of nghost specifies the\n" - "number of cells in the boundary region of each FArrayBox in\n" - "the subregion that should be modified." - ) - .def("plus", - py::overload_cast< Real, const Box&, int, int, int >(&MultiFab::plus), - py::arg("val"), py::arg("region"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, - "Identical to the previous version of plus(), with the\n" - "restriction that the subregion is further constrained to\n" - "the intersection with Box region." - ) - .def("plus", - py::overload_cast< MultiFab const &, int, int, int >(&MultiFab::plus), - py::arg("mf"), py::arg("strt_comp"), py::arg("num_comp"), py::arg("nghost")=0, - "This function adds the values of the cells in mf to the corresponding\n" - "cells of this MultiFab. mf is required to have the same BoxArray or\n" - "\"valid region\" as this MultiFab. The addition is done only to num_comp\n" - "components, starting with component number strt_comp. The parameter\n" - "nghost specifies the number of boundary cells that will be modified.\n" - "If nghost == 0, only the valid region of each FArrayBox will be\n" - "modified." - ) - - .def("minus", - py::overload_cast< MultiFab const &, int, int, int >(&MultiFab::minus), - py::arg("mf"), py::arg("strt_comp"), py::arg("num_comp"), py::arg("nghost")=0, - "This function subtracts the values of the cells in mf from the\n" - "corresponding cells of this MultiFab. mf is required to have the\n" - "same BoxArray or \"valid region\" as this MultiFab. The subtraction is\n" - "done only to num_comp components, starting with component number\n" - "strt_comp. The parameter nghost specifies the number of boundary\n" - "cells that will be modified. If nghost == 0, only the valid region of\n" - "each FArrayBox will be modified." - ) - - // renamed: ImportError: overloading a method with both static and instance methods is not supported - .def("divi", - py::overload_cast< MultiFab const &, int, int, int >(&MultiFab::divide), - py::arg("mf"), py::arg("strt_comp"), py::arg("num_comp"), py::arg("nghost")=0, - "This function divides the values of the cells in mf from the\n" - "corresponding cells of this MultiFab. mf is required to have the\n" - "same BoxArray or \"valid region\" as this MultiFab. The division is\n" - "done only to num_comp components, starting with component number\n" - "strt_comp. The parameter nghost specifies the number of boundary\n" - "cells that will be modified. If nghost == 0, only the valid region of\n" - "each FArrayBox will be modified. Note, nothing is done to protect\n" - "against divide by zero." - ) - - .def("mult", - py::overload_cast< Real, int >(&MultiFab::mult), - py::arg("val"), py::arg("nghost")=0, - "Scales the value of each cell in the valid region of each\n" - "component of the MultiFab by the scalar val (a[i] <- a[i]*val).\n" - "The value of nghost specifies the number of cells in the\n" - "boundary region that should be modified." - ) - .def("mult", - py::overload_cast< Real, int, int, int >(&MultiFab::mult), - py::arg("val"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, - "Scales the value of each cell in the specified subregion of the\n" - "MultiFab by the scalar val (a[i] <- a[i]*val). The subregion\n" - "consists of the num_comp components starting at component comp.\n" - "The value of nghost specifies the number of cells in the\n" - "boundary region of each FArrayBox in the subregion that should\n" - "be modified." - ) - .def("mult", - py::overload_cast< Real, Box const &, int, int, int >(&MultiFab::mult), - py::arg("val"), py::arg("region"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, - "Identical to the previous version of mult(), with the\n" - "restriction that the subregion is further constrained to the\n" - "intersection with Box region. The value of nghost specifies the\n" - "number of cells in the boundary region of each FArrayBox in\n" - "the subregion that should be modified." - ) - .def("mult", - py::overload_cast< Real, Box const &, int >(&MultiFab::mult), - py::arg("val"), py::arg("region"), py::arg("nghost")=0, - "Scales the value of each cell in the valid region of each\n" - "component of the MultiFab by the scalar val (a[i] <- a[i]*val),\n" - "that also intersects the Box region. The value of nghost\n" - "specifies the number of cells in the boundary region of each\n" - "FArrayBox in the subregion that should be modified." - ) - - .def("invert", - py::overload_cast< Real, int >(&MultiFab::invert), - py::arg("numerator"), py::arg("nghost"), - "Replaces the value of each cell in the specified subregion of\n" - "the MultiFab with its reciprocal multiplied by the value of\n" - "numerator. The value of nghost specifies the number of cells\n" - "in the boundary region that should be modified." - ) - .def("invert", - py::overload_cast< Real, int, int, int >(&MultiFab::invert), - py::arg("numerator"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, - "Replaces the value of each cell in the specified subregion of\n" - "the MultiFab with its reciprocal multiplied by the value of\n" - "numerator. The subregion consists of the num_comp components\n" - "starting at component comp. The value of nghost specifies the\n" - "number of cells in the boundary region of each FArrayBox in the\n" - "subregion that should be modified." - ) - .def("invert", - py::overload_cast< Real, Box const &, int >(&MultiFab::invert), - py::arg("numerator"), py::arg("region"), py::arg("nghost"), - "Scales the value of each cell in the valid region of each\n" - "component of the MultiFab by the scalar val (a[i] <- a[i]*val),\n" - "that also intersects the Box region. The value of nghost\n" - "specifies the number of cells in the boundary region of each\n" - "FArrayBox in the subregion that should be modified." - ) - .def("invert", - py::overload_cast< Real, Box const &, int, int, int >(&MultiFab::invert), - py::arg("numerator"), py::arg("region"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, - "Identical to the previous version of invert(), with the\n" - "restriction that the subregion is further constrained to the\n" - "intersection with Box region. The value of nghost specifies the\n" - "number of cells in the boundary region of each FArrayBox in the\n" - "subregion that should be modified." - ) - - .def("negate", - py::overload_cast< int >(&MultiFab::negate), - py::arg("nghost")=0, - "Negates the value of each cell in the valid region of\n" - "the MultiFab. The value of nghost specifies the number of\n" - "cells in the boundary region that should be modified." - ) - .def("negate", - py::overload_cast< int, int, int >(&MultiFab::negate), - py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, - "Negates the value of each cell in the specified subregion of\n" - "the MultiFab. The subregion consists of the num_comp\n" - "components starting at component comp. The value of nghost\n" - "specifies the number of cells in the boundary region of each\n" - "FArrayBox in the subregion that should be modified." - ) - .def("negate", - py::overload_cast< Box const &, int >(&MultiFab::negate), - py::arg("region"), py::arg("nghost")=0, - "Negates the value of each cell in the valid region of\n" - "the MultiFab that also intersects the Box region. The value\n" - "of nghost specifies the number of cells in the boundary region\n" - "that should be modified." - ) - .def("negate", - py::overload_cast< Box const &, int, int, int >(&MultiFab::negate), - py::arg("region"), py::arg("comp"), py::arg("num_comp"), py::arg("nghost")=0, - "Identical to the previous version of negate(), with the\n" - "restriction that the subregion is further constrained to\n" - "the intersection with Box region." - ) - - /* static (standalone) simple math functions */ - .def("dot", - [](MultiFab const & self, int comp, MultiFab const & y, int y_comp, int numcomp, int nghost, bool local) { - return MultiFab::Dot(self, comp, y, y_comp, numcomp, nghost, local); - }, - py::arg("comp"), - py::arg("y"), py::arg("y_comp"), - py::arg("numcomp"), py::arg("nghost"), py::arg("local")=false, - "Returns the dot product of self with another MultiFab." - ) - .def("dot", - [](MultiFab const & self, int comp, int numcomp, int nghost, bool local) { - return MultiFab::Dot(self, comp, numcomp, nghost, local); - }, - py::arg("comp"), - py::arg("numcomp"), py::arg("nghost"), py::arg("local")=false, - "Returns the dot product with itself." - ) - /** TODO: Bind iMultiFab - .def("dot", - [](MultiFab const& self, const iMultiFab& mask, int comp, MultiFab const& y, int y_comp, int numcomp, int nghost, bool local) { - return MultiFab::Dot(mask, self, comp, y, y_comp, numcomp, nghost, local); - }, - py::arg("mask"), py::arg("comp"), py::arg("y"), py::arg("y_comp"), - py::arg("numcomp"), py::arg("nghost"), py::arg("local")=false, - "Returns the dot product of self with another MultiFab where the mask is valid." - ) - */ - - .def("add", - [](MultiFab & self, MultiFab const & src, int srccomp, int comp, int numcomp, int nghost) { - MultiFab::Add(self, src, srccomp, comp, numcomp, nghost); - }, - py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "Add src to self including nghost ghost cells.\n" - "The two MultiFabs MUST have the same underlying BoxArray." - ) - .def("add", - [](MultiFab & self, MultiFab const & src, int srccomp, int comp, int numcomp, IntVect const & nghost) { - MultiFab::Add(self, src, srccomp, comp, numcomp, nghost); - }, - py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "Add src to self including nghost ghost cells.\n" - "The two MultiFabs MUST have the same underlying BoxArray." - ) - - .def("subtract", - [](MultiFab & self, MultiFab const & src, int srccomp, int comp, int numcomp, int nghost) { - MultiFab::Subtract(self, src, srccomp, comp, numcomp, nghost); - }, - py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "Subtract src from self including nghost ghost cells.\n" - "The two MultiFabs MUST have the same underlying BoxArray." - ) - .def("subtract", - [](MultiFab & self, MultiFab const & src, int srccomp, int comp, int numcomp, IntVect const & nghost) { - MultiFab::Subtract(self, src, srccomp, comp, numcomp, nghost); - }, - py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "Subtract src from self including nghost ghost cells.\n" - "The two MultiFabs MUST have the same underlying BoxArray." - ) - - .def("multiply", - [](MultiFab & self, MultiFab const & src, int srccomp, int comp, int numcomp, int nghost) { - MultiFab::Multiply(self, src, srccomp, comp, numcomp, nghost); - }, - py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "Multiply self by src including nghost ghost cells.\n" - "The two MultiFabs MUST have the same underlying BoxArray." - ) - .def("multiply", - [](MultiFab & self, MultiFab const & src, int srccomp, int comp, int numcomp, IntVect const & nghost) { - MultiFab::Multiply(self, src, srccomp, comp, numcomp, nghost); - }, - py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "Multiply self by src including nghost ghost cells.\n" - "The two MultiFabs MUST have the same underlying BoxArray." - ) - - .def("divide", - [](MultiFab & self, MultiFab const & src, int srccomp, int comp, int numcomp, int nghost) { - MultiFab::Divide(self, src, srccomp, comp, numcomp, nghost); - }, - py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "Divide self by src including nghost ghost cells.\n" - "The two MultiFabs MUST have the same underlying BoxArray." - ) - .def("divide", - [](MultiFab & self, MultiFab const & src, int srccomp, int comp, int numcomp, IntVect const & nghost) { - MultiFab::Divide(self, src, srccomp, comp, numcomp, nghost); - }, - py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "Divide self by src including nghost ghost cells.\n" - "The two MultiFabs MUST have the same underlying BoxArray." - ) - - .def("swap", - [](MultiFab & self, MultiFab & src, int srccomp, int comp, int numcomp, int nghost) { - MultiFab::Swap(self, src, srccomp, comp, numcomp, nghost); - }, - py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "Swap from src to self including nghost ghost cells.\n" - "The two MultiFabs MUST have the same underlying BoxArray.\n" - "The swap is local." - ) - .def("swap", - [](MultiFab & self, MultiFab & src, int srccomp, int comp, int numcomp, IntVect const & nghost) { - MultiFab::Swap(self, src, srccomp, comp, numcomp, nghost); - }, - py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "Swap from src to self including nghost ghost cells.\n" - "The two MultiFabs MUST have the same underlying BoxArray.\n" - "The swap is local." - ) - - .def("saxpy", - [](MultiFab & self, Real a, MultiFab const & src, int srccomp, int comp, int numcomp, int nghost) { - MultiFab::Saxpy(self, a, src, srccomp, comp, numcomp, nghost); - }, - py::arg("a"), py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "self += a * src" - ) - - .def("xpay", - [](MultiFab & self, Real a, MultiFab const & src, int srccomp, int comp, int numcomp, int nghost) { - MultiFab::Xpay(self, a, src, srccomp, comp, numcomp, nghost); - }, - py::arg("a"), py::arg("src"), py::arg("srccomp"), py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "self = src + a * self" - ) - - .def("lin_comb", - [](MultiFab & self, Real a, MultiFab const & x, int x_comp, Real b, MultiFab const & y, int y_comp, int comp, int numcomp, int nghost) { - MultiFab::LinComb(self, a, x, x_comp, b, y, y_comp, comp, numcomp, nghost); - }, - py::arg("a"), py::arg("x"), py::arg("x_comp"), - py::arg("b"), py::arg("y"), py::arg("y_comp"), - py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "self = a * x + b * y" - ) - - .def("add_product", - [](MultiFab & self, MultiFab const & src1, int comp1, MultiFab const & src2, int comp2, int comp, int numcomp, int nghost) { - MultiFab::AddProduct(self, src1, comp1, src2, comp2, comp, numcomp, nghost); - }, - py::arg("src1"), py::arg("comp1"), - py::arg("src2"), py::arg("comp2"), - py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "self += src1 * src2" - ) - .def("add_product", - [](MultiFab & self, MultiFab const & src1, int comp1, MultiFab const & src2, int comp2, int comp, int numcomp, IntVect const & nghost) { - MultiFab::AddProduct(self, src1, comp1, src2, comp2, comp, numcomp, nghost); - }, - py::arg("src1"), py::arg("comp1"), - py::arg("src2"), py::arg("comp2"), - py::arg("comp"), py::arg("numcomp"), py::arg("nghost"), - "self += src1 * src2" - ) - - /* simple data validity checks */ - .def("contains_nan", - py::overload_cast< bool >(&MultiFab::contains_nan, py::const_), - py::arg("local")=false - ) - .def("contains_nan", - py::overload_cast< int, int, int, bool >(&MultiFab::contains_nan, py::const_), - py::arg("scomp"), py::arg("ncomp"), py::arg("ngrow")=0, py::arg("local")=false - ) - .def("contains_nan", - py::overload_cast< int, int, IntVect const &, bool >(&MultiFab::contains_nan, py::const_), - py::arg("scomp"), py::arg("ncomp"), py::arg("ngrow"), py::arg("local")=false - ) - - .def("contains_inf", - py::overload_cast< bool >(&MultiFab::contains_inf, py::const_), - py::arg("local")=false - ) - .def("contains_inf", - py::overload_cast< int, int, int, bool >(&MultiFab::contains_inf, py::const_), - py::arg("scomp"), py::arg("ncomp"), py::arg("ngrow")=0, py::arg("local")=false - ) - .def("contains_inf", - py::overload_cast< int, int, IntVect const &, bool >(&MultiFab::contains_inf, py::const_), - py::arg("scomp"), py::arg("ncomp"), py::arg("ngrow"), py::arg("local")=false - ) - - .def("box_array", &MultiFab::boxArray) - .def("dm", &MultiFab::DistributionMap) - .def_property_readonly("n_comp", &MultiFab::nComp) - .def_property_readonly("n_grow_vect", &MultiFab::nGrowVect) - - /* masks & ownership */ - // TODO: - // - OverlapMask -> std::unique_ptr - // - OwnerMask -> std::unique_ptr - - /* Syncs */ - .def("average_sync", &MultiFab::AverageSync) - .def("weighted_sync", &MultiFab::WeightedSync) - //.def("override_sync", py::overload_cast< iMultiFab const &, Periodicity const & >(&MultiFab::OverrideSync)) - - /* Init & Finalize */ - .def_static("initialize", &MultiFab::Initialize) - .def_static("finalize", &MultiFab::Finalize) - ; - + make_MultiFab(py_MultiFab, "MultiFab"); m.def("copy_mfab", py::overload_cast< MultiFab &, MultiFab const &, int, int, int, int >(&MultiFab::Copy), py::arg("dst"), py::arg("src"), py::arg("srccomp"), py::arg("dstcomp"), py::arg("numcomp"), py::arg("nghost")) .def("copy_mfab", py::overload_cast< MultiFab &, MultiFab const &, int, int, int, IntVect const & >(&MultiFab::Copy), py::arg("dst"), py::arg("src"), py::arg("srccomp"), py::arg("dstcomp"), py::arg("numcomp"), py::arg("nghost")); diff --git a/src/Base/iMultiFab.cpp b/src/Base/iMultiFab.cpp new file mode 100644 index 00000000..ba144866 --- /dev/null +++ b/src/Base/iMultiFab.cpp @@ -0,0 +1,23 @@ +/* Copyright 2021-2022 The AMReX Community + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#include "pyAMReX.H" +#include "MultiFab.H" + +#include +#include +#include + + +void init_iMultiFab(py::module &m) +{ + using namespace amrex; + + py::class_< iMultiFab, FabArray > py_iMultiFab(m, "iMultiFab"); + make_MultiFab(py_iMultiFab, "iMultiFab"); + + m.def("copy_mfab", py::overload_cast< iMultiFab &, iMultiFab const &, int, int, int, int >(&iMultiFab::Copy), py::arg("dst"), py::arg("src"), py::arg("srccomp"), py::arg("dstcomp"), py::arg("numcomp"), py::arg("nghost")) + .def("copy_mfab", py::overload_cast< iMultiFab &, iMultiFab const &, int, int, int, IntVect const & >(&iMultiFab::Copy), py::arg("dst"), py::arg("src"), py::arg("srccomp"), py::arg("dstcomp"), py::arg("numcomp"), py::arg("nghost")); +} diff --git a/src/amrex/extensions/MultiFab.py b/src/amrex/extensions/MultiFab.py index 76e0dbe5..8cae31d6 100644 --- a/src/amrex/extensions/MultiFab.py +++ b/src/amrex/extensions/MultiFab.py @@ -45,7 +45,8 @@ def mf_to_numpy(self, copy=False, order="F"): mf = self if copy: - mf = amr.MultiFab( + mf_type = type(self) # MultiFab or iMultiFab + mf = mf_type( self.box_array(), self.dm(), self.n_comp, @@ -156,7 +157,8 @@ def copy_multifab(amr, self): amrex.MultiFab A copy of this MultiFab. """ - mf = amr.MultiFab( + mf_type = type(self) # MultiFab or iMultiFab + mf = mf_type( self.box_array(), self.dm(), self.n_comp, @@ -657,3 +659,19 @@ def register_MultiFab_extension(amr): amr.MultiFab.shape_with_ghosts = property(shape_with_ghosts) amr.MultiFab.__getitem__ = __getitem__ amr.MultiFab.__setitem__ = __setitem__ + + # iMultiFab + amr.iMultiFab.__iter__ = lambda imfab: amr.MFIter(imfab) + + amr.iMultiFab.to_numpy = mf_to_numpy + amr.iMultiFab.to_cupy = mf_to_cupy + amr.iMultiFab.to_xp = mf_to_xp + + amr.iMultiFab.copy = lambda self: copy_multifab(amr, self) + amr.iMultiFab.copy.__doc__ = copy_multifab.__doc__ + + amr.iMultiFab.imesh = imesh + amr.iMultiFab.shape = property(shape) + amr.iMultiFab.shape_with_ghosts = property(shape_with_ghosts) + amr.iMultiFab.__getitem__ = __getitem__ + amr.iMultiFab.__setitem__ = __setitem__ diff --git a/src/pyAMReX.cpp b/src/pyAMReX.cpp index fd241527..36ce03d0 100644 --- a/src/pyAMReX.cpp +++ b/src/pyAMReX.cpp @@ -6,6 +6,7 @@ #include "pyAMReX.H" #include +#include #define STRINGIFY(x) #x #define MACRO_STRINGIFY(x) STRINGIFY(x) @@ -24,14 +25,17 @@ void init_BoxArray(py::module &); void init_CoordSys(py::module&); void init_Dim3(py::module&); void init_DistributionMapping(py::module&); +void init_FabArray(py::module &); void init_FArrayBox(py::module&); void init_Geometry(py::module&); +void init_iMultiFab(py::module&); void init_IndexType(py::module &); void init_IntVect(py::module &); +void init_MFInfo(py::module &); #ifdef AMREX_USE_MPI void init_MPMD(py::module &); #endif -void init_MultiFab(py::module &); +void init_MultiFab(py::module &, py::class_< amrex::MFIter >&); void init_ParallelDescriptor(py::module &); void init_ParmParse(py::module &); void init_ParticleContainer(py::module &); @@ -70,9 +74,12 @@ PYBIND11_MODULE(amrex_3d_pybind, m) { BoxArray Dim3 FArrayBox + iMultiFab IntVect IndexType RealVect + MFInfo + MFItInfo MultiFab ParallelDescriptor Particle @@ -111,7 +118,11 @@ PYBIND11_MODULE(amrex_3d_pybind, m) { init_DistributionMapping(m); init_BaseFab(m); init_FArrayBox(m); - init_MultiFab(m); + py::class_< amrex::MFIter > py_MFIter(m, "MFIter", py::dynamic_attr()); + init_FabArray(m); + init_MFInfo(m); + init_iMultiFab(m); + init_MultiFab(m, py_MFIter); init_ParallelDescriptor(m); init_PODVector(m); diff --git a/tests/conftest.py b/tests/conftest.py index d963af22..1c6a6776 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -138,3 +138,45 @@ def __exit__(self, exc_type, exc_value, traceback): with MfabDeviceContextManager() as mfab: yield mfab + + +@pytest.fixture(scope="function", params=list(itertools.product([1, 3], [0, 1]))) +def imfab(boxarr, distmap, request): + class iMfabContextManager: + def __enter__(self): + num_components = request.param[0] + num_ghost = request.param[1] + self.imfab = amr.iMultiFab(boxarr, distmap, num_components, num_ghost) + self.imfab.set_val(0, 0, num_components) + return self.imfab + + def __exit__(self, exc_type, exc_value, traceback): + self.imfab.clear() + del self.imfab + + with iMfabContextManager() as imfab: + yield imfab + + +@pytest.fixture(scope="function", params=list(itertools.product([1, 3], [0, 1]))) +def imfab_device(boxarr, distmap, request): + class iMfabDeviceContextManager: + def __enter__(self): + num_components = request.param[0] + num_ghost = request.param[1] + self.imfab = amr.iMultiFab( + boxarr, + distmap, + num_components, + num_ghost, + amr.MFInfo().set_arena(amr.The_Device_Arena()), + ) + self.imfab.set_val(0, 0, num_components) + return self.imfab + + def __exit__(self, exc_type, exc_value, traceback): + self.imfab.clear() + del self.imfab + + with iMfabDeviceContextManager() as imfab: + yield imfab diff --git a/tests/test_imultifab.py b/tests/test_imultifab.py new file mode 100644 index 00000000..da0b94af --- /dev/null +++ b/tests/test_imultifab.py @@ -0,0 +1,486 @@ +# -*- coding: utf-8 -*- + +import math + +import numpy as np +import pytest + +import amrex.space3d as amr + + +def test_imfab_numpy(imfab): + # finest active MR level, get from a + # simulation's AmrMesh object, e.g.: + # finest_level = sim.finest_level + finest_level = 0 # no MR + + # iterate over mesh-refinement levels + for lev in range(finest_level + 1): + # get an existing MultiFab, e.g., + # from a simulation: + # imfab = sim.get_field(lev=lev) + # Config = sim.extension.Config + + # grow (aka guard/ghost/halo) regions + ngv = imfab.n_grow_vect + + # get every local block of the field + for mfi in imfab: + # global index box w/ guards + bx = mfi.tilebox().grow(ngv) + print(bx) + + # numpy/cupy representation: non- + # copying view, w/ guard/ghost + field = imfab.array(mfi).to_xp() + + # notes on indexing in field: + # - numpy uses locally zero-based indexing + # - layout is F_CONTIGUOUS by default, just like AMReX + + field[()] = 42 + + # finest active MR level, get from a + # simulation's AmrMesh object, e.g.: + # finest_level = sim.finest_level + finest_level = 0 # no MR + + # iterate over mesh-refinement levels + for lev in range(finest_level + 1): + # get an existing MultiFab, e.g., + # from a simulation: + # imfab = sim.get_field(lev=lev) + # Config = sim.extension.Config + + field_list = imfab.to_xp() + + for field in field_list: + field[()] = 42 + + # finest active MR level, get from a + # simulation's AmrMesh object, e.g.: + # finest_level = sim.finest_level + finest_level = 0 # no MR + + # iterate over mesh-refinement levels + for lev in range(finest_level + 1): + # get an existing MultiFab, e.g., + # from a simulation: + # imfab = sim.get_field(lev=lev) + # Config = sim.extension.Config + + # Using global indexing + # Set all valid cells (and internal ghost cells) + imfab[...] = 42 + + # Set a range of cells. Indices are in Fortran order. + # First dimension, sets from first lower guard cell to first upper guard cell. + # - Imaginary indices refer to the guard cells, negative lower, positive upper. + # Second dimension, sets all valid cells. + # Third dimension, sets all valid and ghost cells + # - The empty tuple is used to specify the range to include all valid and ghost cells. + # Components dimension, sets first component. + imfab[-1j:2j, :, (), 0] = 42 + + # Get a range of cells + # Get the data along the valid cells in the first dimension (gathering data across blocks + # and processors), at the first upper guard cell in the second dimensionn, and cell 2 of + # the third (with 2 being relative to 0 which is the lower end of the valid cells of the full domain). + # Note that in an MPI context, this is a global operation, so caution is required when + # scaling to large numbers of processors. + if imfab.n_grow_vect.max > 0: + mfslice = imfab[:, 1j, 2] + # The assignment is to the last valid cell of the second dimension. + imfab[:, -1, 2] = 2 * mfslice + + +@pytest.mark.skipif(amr.Config.have_gpu, reason="This test only runs on CPU") +def test_imfab_loop_slow(imfab): + ngv = imfab.n_grow_vect + print(f"\n imfab={imfab}, imfab.n_grow_vect={ngv}") + + for mfi in imfab: + bx = mfi.tilebox().grow(ngv) + marr = imfab.array(mfi) + + # print(imfab.num_comp) + # print(imfab.size) + # print(marr.size) + # print(marr.nComp) + + # index by index assignment + # notes: + # - this is AMReX Array4, F-order indices + # - even though we iterate by fastest varying index, + # such loops are naturally very slow in Python + three_comps = imfab.num_comp == 3 + if three_comps: + for i, j, k in bx: + # print(i,j,k) + marr[i, j, k, 0] = 10 * i + marr[i, j, k, 1] = 10 * j + marr[i, j, k, 2] = 10 * k + else: + for i, j, k in bx: + # print(i,j,k) + marr[i, j, k] = 10 * i + + # numpy representation: non-copying view, zero-indexed, + # includes the guard/ghost region + marr_np = marr.to_numpy() + + # check the values at start/end are the same: first component + assert marr_np[0, 0, 0, 0] == marr[bx.small_end] + assert marr_np[-1, -1, -1, 0] == marr[bx.big_end] + # same check, but for all components + for n in range(imfab.num_comp): + small_end_comp = list(bx.small_end) + [n] + big_end_comp = list(bx.big_end) + [n] + assert marr_np[0, 0, 0, n] == marr[small_end_comp] + assert marr_np[-1, -1, -1, n] == marr[big_end_comp] + + # all components and all indices set at once to 42 + marr_np[()] = 42 + + # values in start & end still match? + assert marr_np[0, 0, 0, 0] == marr[bx.small_end] + assert marr_np[-1, -1, -1, -1] == marr[bx.big_end] + + # all values for all indices match between multifab & numpy view? + for n in range(imfab.num_comp): + for i, j, k in bx: + assert marr[i, j, k, n] == 42 + + +def test_imfab_loop(imfab): + ngv = imfab.n_grow_vect + print(f"\n imfab={imfab}, imfab.n_grow_vect={ngv}") + + for mfi in imfab: + bx = mfi.tilebox().grow(ngv) + marr = imfab.array(mfi) + + # note: offset from index space in numpy + # in numpy, we start indices from zero, not small_end + + # numpy/cupy representation: non-copying view, including the + # guard/ghost region + marr_xp = marr.to_xp() + + marr_xp[()] = 10 # TODO: fill with index value or so as in test_imfab_loop_slow + + def iv2s(iv, comp): + return tuple(iv) + (comp,) + + # check the values at start/end are the same: first component + for n in range(imfab.num_comp): + assert marr_xp[0, 0, 0, n] == 10 + assert marr_xp[-1, -1, -1, n] == marr_xp[iv2s(bx.big_end - bx.small_end, n)] + + # now we do some faster assignments, using range based access + # This should fail as out-of-bounds, but does not. + # Does NumPy/CuPy not check array access for non-owned views? + # marr_xp[24:200, :, :, :] = 42. + + # all components and all indices set at once to 42 + marr_xp[()] = 42 + + # values in start & end still match? + for n in range(imfab.num_comp): + assert marr_xp[0, 0, 0, n] == 42 + assert marr_xp[-1, -1, -1, n] == marr_xp[iv2s(bx.big_end - bx.small_end, n)] + + +def test_imfab_simple(imfab): + assert imfab.is_all_cell_centered + # assert(all(not imfab.is_nodal(i) for i in [-1, 0, 1, 2])) # -1?? + assert all(not imfab.is_nodal(i) for i in [0, 1, 2]) + + for i in range(imfab.num_comp): + imfab.set_val(-10 * (i + 1), i, 1) + imfab.abs(0, imfab.num_comp) + for i in range(imfab.num_comp): + assert imfab.max(i) == (10 * (i + 1)) # Assert: None == 10 for i=0 + assert imfab.min(i) == (10 * (i + 1)) + + imfab.plus(20, 0, imfab.num_comp) + for i in range(imfab.num_comp): + np.testing.assert_allclose(imfab.max(i), 20 + (10 * (i + 1))) + np.testing.assert_allclose(imfab.min(i), 20 + (10 * (i + 1))) + + imfab.mult(10, 0, imfab.num_comp) + for i in range(imfab.num_comp): + np.testing.assert_allclose(imfab.max(i), 10 * (20 + (10 * (i + 1)))) + np.testing.assert_allclose(imfab.min(i), 10 * (20 + (10 * (i + 1)))) + + imfab.negate(0, imfab.num_comp) + for i in range(imfab.num_comp): + np.testing.assert_allclose(imfab.max(i), -10 * (20 + (10 * (i + 1)))) + np.testing.assert_allclose(imfab.min(i), -10 * (20 + (10 * (i + 1)))) + + +@pytest.mark.parametrize("nghost", [0, 1]) +def test_imfab_ops(boxarr, distmap, nghost): + src = amr.MultiFab(boxarr, distmap, 3, nghost) + dst = amr.MultiFab(boxarr, distmap, 1, nghost) + + src.set_val(10, 0, 1) + src.set_val(20, 1, 1) + src.set_val(30, 2, 1) + dst.set_val(0, 0, 1) + + dst.add(src, 2, 0, 1, nghost) + dst.subtract(src, 1, 0, 1, nghost) + dst.multiply(src, 0, 0, 1, nghost) + dst.divide(src, 1, 0, 1, nghost) + + print(dst.min(0)) + np.testing.assert_allclose(dst.min(0), 5) + np.testing.assert_allclose(dst.max(0), 5) + + +def test_imfab_mfiter(imfab): + assert iter(imfab).is_valid + assert iter(imfab).length == 8 + + cnt = 0 + for _mfi in imfab: + cnt += 1 + + assert iter(imfab).length == cnt + + +@pytest.mark.skipif( + amr.Config.gpu_backend != "CUDA", reason="Requires AMReX_GPU_BACKEND=CUDA" +) +def test_imfab_ops_cuda_numba(imfab_device): + # https://numba.pydata.org/numba-doc/dev/cuda/cuda_array_interface.html + from numba import cuda + + ngv = imfab_device.n_grow_vect + + # assign 3: define kernel + @cuda.jit + def set_to_three(array): + i, j, k = cuda.grid(3) + if i < array.shape[0] and j < array.shape[1] and k < array.shape[2]: + array[i, j, k] = 3 + + # assign 3: loop through boxes and launch kernels + for mfi in imfab_device: + bx = mfi.tilebox().grow(ngv) # noqa + marr = imfab_device.array(mfi) + marr_numba = cuda.as_cuda_array(marr) + + # kernel launch + threadsperblock = (4, 4, 4) + blockspergrid = tuple( + [math.ceil(s / b) for s, b in zip(marr_numba.shape, threadsperblock)] + ) + set_to_three[blockspergrid, threadsperblock](marr_numba) + + # Check results + shape = 32**3 * 8 + sum_threes = imfab_device.sum_unique(comp=0, local=False) + assert sum_threes == shape * 3 + + +@pytest.mark.skipif( + amr.Config.gpu_backend != "CUDA", reason="Requires AMReX_GPU_BACKEND=CUDA" +) +def test_imfab_ops_cuda_cupy(imfab_device): + # https://docs.cupy.dev/en/stable/user_guide/interoperability.html + import cupy as cp + import cupyx.profiler + + # AMReX -> cupy + ngv = imfab_device.n_grow_vect + print(f"\n imfab_device={imfab_device}, imfab_device.n_grow_vect={ngv}") + + # assign 3 + with cupyx.profiler.time_range("assign 3 [()]", color_id=0): + for mfi in imfab_device: + bx = mfi.tilebox().grow(ngv) # noqa + marr_cupy = imfab_device.array(mfi).to_cupy(order="C") + # print(marr_cupy.shape) # 1, 32, 32, 32 + # print(marr_cupy.dtype) # float64 + # performance: + # https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074 + + # write and read into the marr_cupy + marr_cupy[()] = 3 + + # verify result with a .sum_unique + with cupyx.profiler.time_range("verify 3", color_id=0): + shape = 32**3 * 8 + # print(imfab_device.shape) + sum_threes = imfab_device.sum_unique(comp=0, local=False) + assert sum_threes == shape * 3 + + # assign 2 + with cupyx.profiler.time_range("assign 2 (set_val)", color_id=1): + imfab_device.set_val(2) + with cupyx.profiler.time_range("verify 2", color_id=1): + sum_twos = imfab_device.sum_unique(comp=0, local=False) + assert sum_twos == shape * 2 + + # assign 5 + with cupyx.profiler.time_range("assign 5 (ones-like)", color_id=2): + + def set_to_five(mm): + xp = cp.get_array_module(mm) + assert xp.__name__ == "cupy" + mm = xp.ones_like(mm) * 10 + mm /= 2 + return mm + + for mfi in imfab_device: + bx = mfi.tilebox().grow(ngv) # noqa + marr_cupy = imfab_device.array(mfi).to_cupy(order="F") + # print(marr_cupy.shape) # 32, 32, 32, 1 + # print(marr_cupy.dtype) # float64 + # performance: + # https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074 + + # write and read into the marr_cupy + fives_cp = set_to_five(marr_cupy) + marr_cupy[()] = 0 + marr_cupy += fives_cp + + # verify + with cupyx.profiler.time_range("verify 5", color_id=2): + sum = imfab_device.sum_unique(comp=0, local=False) + assert sum == shape * 5 + + # assign 7 + with cupyx.profiler.time_range("assign 7 (fuse)", color_id=3): + + @cp.fuse(kernel_name="set_to_seven") + def set_to_seven(x): + x[...] = 7 + + for mfi in imfab_device: + bx = mfi.tilebox().grow(ngv) # noqa + marr_cupy = imfab_device.array(mfi).to_cupy(order="C") + + # write and read into the marr_cupy + set_to_seven(marr_cupy) + + # verify + with cupyx.profiler.time_range("verify 7", color_id=3): + sum = imfab_device.sum_unique(comp=0, local=False) + assert sum == shape * 7 + + # TODO: @jit.rawkernel() + + +@pytest.mark.skipif( + amr.Config.gpu_backend != "CUDA", reason="Requires AMReX_GPU_BACKEND=CUDA" +) +def test_imfab_ops_cuda_pytorch(imfab_device): + # https://docs.cupy.dev/en/stable/user_guide/interoperability.html#pytorch + import torch + + # assign 3: loop through boxes and launch kernel + for mfi in imfab_device: + marr = imfab_device.array(mfi) + marr_torch = torch.as_tensor(marr, device="cuda") + marr_torch[:, :, :] = 3 + + # Check results + shape = 32**3 * 8 + sum_threes = imfab_device.sum_unique(comp=0, local=False) + assert sum_threes == shape * 3 + + +@pytest.mark.skipif( + amr.Config.gpu_backend != "CUDA", reason="Requires AMReX_GPU_BACKEND=CUDA" +) +def test_imfab_ops_cuda_cuml(imfab_device): + pass + # https://github.com/rapidsai/cuml + # https://github.com/rapidsai/cudf + # maybe better for particles as a dataframe test + # import cudf + # import cuml + + # AMReX -> RAPIDSAI cuML + # arr_cuml = ... + # assert(arr_cuml.__cuda_array_interface__['data'][0] == arr.__cuda_array_interface__['data'][0]) + # TODO + + +@pytest.mark.skipif( + amr.Config.gpu_backend != "CUDA", reason="Requires AMReX_GPU_BACKEND=CUDA" +) +def test_imfab_dtoh_copy(imfab_device): + class MfabPinnedContextManager: + def __enter__(self): + self.imfab = amr.MultiFab( + imfab_device.box_array(), + imfab_device.dm(), + imfab_device.n_comp, + imfab_device.n_grow_vect, + amr.MFInfo().set_arena(amr.The_Pinned_Arena()), + ) + return self.imfab + + def __exit__(self, exc_type, exc_value, traceback): + self.imfab.clear() + del self.imfab + + with MfabPinnedContextManager() as imfab_host: + imfab_host.set_val(42) + + amr.dtoh_memcpy(imfab_host, imfab_device) + + # assert all are 0 on host + host_min = imfab_host.min(0) + host_max = imfab_host.max(0) + assert host_min == host_max + assert host_max == 0 + + dev_val = 11 + imfab_host.set_val(dev_val) + amr.htod_memcpy(imfab_device, imfab_host) + + # assert all are 11 on device + for n in range(imfab_device.n_comp): + assert imfab_device.min(comp=n) == dev_val + assert imfab_device.max(comp=n) == dev_val + + # numpy bindings (w/ copy) + local_boxes_host = imfab_device.to_numpy(copy=True) + assert max([np.max(box) for box in local_boxes_host]) == dev_val + del local_boxes_host + + # numpy bindings (w/ copy) + for mfi in imfab_device: + marr = imfab_device.array(mfi).to_numpy(copy=True) + assert np.min(marr) >= dev_val + assert np.max(marr) <= dev_val + + # cupy bindings (w/o copy) + import cupy as cp + + local_boxes_device = imfab_device.to_cupy() + assert max([cp.max(box) for box in local_boxes_device]) == dev_val + + +def test_imfab_copy(imfab): + # write to imfab + imfab.set_val(42) + for i in range(imfab.num_comp): + np.testing.assert_allclose(imfab.max(i), 42) + + # copy + new_imfab = imfab.copy() + + # write to old imfab + imfab.set_val(1) + for i in range(imfab.num_comp): + np.testing.assert_allclose(imfab.max(i), 1) + + # check new imfab is the original data + for i in range(new_imfab.num_comp): + np.testing.assert_allclose(new_imfab.max(i), 42)