Skip to content

Commit

Permalink
ParticleContainer: More Pure SoA Init Functions
Browse files Browse the repository at this point in the history 
Depends on upstream.
  • Loading branch information
@ax3l
ax3l committed May 30, 2023
1 parent 2712b9d commit 9781ae7
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 29 deletions.
40 changes: 19 additions & 21 deletions src/Particle/ParticleContainer.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -131,11 +131,17 @@ void make_Iterators (py::module &m, std::string allocstr)
template <typename T_ParticleType, int T_NArrayReal=0, int T_NArrayInt=0>
void make_ParticleInitData (py::module &m) {
using ParticleType = T_ParticleType;
using ParticleInitData = ParticleInitType<ParticleType::NReal, ParticleType::NInt, T_NArrayReal, T_NArrayInt>;
// depends on https://github.com/AMReX-Codes/amrex/pull/3280
// using ParticleInitData = ParticleInitType<ParticleType, T_NArrayReal, T_NArrayInt>;

std::string const suffix = particle_type_suffix<T_ParticleType, T_NArrayReal, T_NArrayInt>();
//using ParticleInitData = typename ParticleContainerType::ParticleInitData;
using ParticleInitData = typename std::conditional<
ParticleType::is_soa_particle,
// SoA Particle: init w/o positions and id/cpu, but explicitly listed in define
ParticleInitTypeSoA<T_NArrayReal - AMREX_SPACEDIM, T_NArrayInt - 2>,
// legacy particle: init w/o positions and id/cpu, only implicitly listed in define
ParticleInitType<ParticleType::NReal, ParticleType::NInt, T_NArrayReal, T_NArrayInt>
>::type;

std::string const suffix = particle_type_suffix<ParticleType, T_NArrayReal, T_NArrayInt>();
auto const particle_init_data_type =
std::string("ParticleInitType_") + suffix;
auto py_particle_init_data = py::class_<ParticleInitData>(m, particle_init_data_type.c_str())
Expand Down Expand Up @@ -221,14 +227,16 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
// const IntVect* Nrep = nullptr);

// void InitFromBinaryFile (const std::string& file, int extradata);

// void InitFromBinaryMetaFile
// void InitRandom (Long icount, ULong iseed,
// const ParticleInitData& mass,
// bool serialize = false, RealBox bx = RealBox());

.def("Increment", &ParticleContainerType::Increment) // TODO pure SoA
//.def("IncrementWithTotal", &ParticleContainerType::IncrementWithTotal, py::arg("mf"), py::arg("level"), py::arg("local")=false) // TODO pure SoA
// TODO for pure SoA
// depends on https://github.com/AMReX-Codes/amrex/pull/3280
.def("InitRandom", py::overload_cast<Long, ULong, const ParticleInitData&, bool, RealBox>(&ParticleContainerType::InitRandom))
//.def("InitRandomPerBox", py::overload_cast<Long, ULong, const ParticleInitData&>(&ParticleContainerType::InitRandomPerBox)) // TODO pure SoA
.def("InitOnePerCell", &ParticleContainerType::InitOnePerCell) // TODO pure SoA

.def("Increment", &ParticleContainerType::Increment)
.def("IncrementWithTotal", &ParticleContainerType::IncrementWithTotal, py::arg("mf"), py::arg("level"), py::arg("local")=false)
.def("Redistribute", &ParticleContainerType::Redistribute, py::arg("lev_min")=0, py::arg("lev_max")=-1,
py::arg("nGrow")=0, py::arg("local")=0, py::arg("remove_negative")=true)
.def("SortParticlesByCell", &ParticleContainerType::SortParticlesByCell)
Expand Down Expand Up @@ -361,15 +369,6 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
// }
;

// TODO for pure SoA
// depends on https://github.com/AMReX-Codes/amrex/pull/3280
if constexpr (!T_ParticleType::is_soa_particle) {
py_pc
.def("InitRandom", py::overload_cast<Long, ULong, const ParticleInitData&, bool, RealBox>(&ParticleContainerType::InitRandom)) // TODO pure SoA
.def("InitRandomPerBox", py::overload_cast<Long, ULong, const ParticleInitData&>(&ParticleContainerType::InitRandomPerBox)) // TODO pure SoA
.def("InitOnePerCell", &ParticleContainerType::InitOnePerCell);
}

using iterator = amrex::ParIter_impl<ParticleType, T_NArrayReal, T_NArrayInt, Allocator>;
make_Iterators< false, iterator, Allocator >(m, allocstr);
using const_iterator = amrex::ParConstIter_impl<ParticleType, T_NArrayReal, T_NArrayInt, Allocator>;
Expand All @@ -383,8 +382,7 @@ void make_ParticleContainer_and_Iterators (py::module &m)
{
// TODO for pure SoA
// depends on https://github.com/AMReX-Codes/amrex/pull/3280
if constexpr (!T_ParticleType::is_soa_particle)
make_ParticleInitData<T_ParticleType, T_NArrayReal, T_NArrayInt>(m);
make_ParticleInitData<T_ParticleType, T_NArrayReal, T_NArrayInt>(m);

// see Src/Base/AMReX_GpuContainers.H
// !AMREX_USE_GPU: DefaultAllocator = std::allocator
Expand Down
16 changes: 8 additions & 8 deletions src/Particle/ParticleTile.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,11 @@ void make_ParticleTileData(py::module &m) {
std::to_string(NArrayInt);
py::class_<ParticleTileDataType>(m, particle_tile_data_type.c_str())
.def(py::init())
.def_readonly("m_size", &ParticleTileDataType::m_size)
//.def_readonly("m_size", &ParticleTileDataType::m_size)
.def_readonly("m_num_runtime_real", &ParticleTileDataType::m_num_runtime_real)
.def_readonly("m_num_runtime_int", &ParticleTileDataType::m_num_runtime_int)
.def("getSuperParticle", &ParticleTileDataType::template getSuperParticle<ParticleType>)
.def("setSuperParticle", &ParticleTileDataType::setSuperParticle)
//.def("getSuperParticle", &ParticleTileDataType::template getSuperParticle<ParticleType>)
//.def("setSuperParticle", &ParticleTileDataType::setSuperParticle)
// setter & getter
.def("__setitem__", [](ParticleTileDataType &pdt, int const v,
SuperParticleType const value) { pdt.setSuperParticle(value, v); })
Expand Down Expand Up @@ -73,11 +73,11 @@ void make_ParticleTile(py::module &m, std::string allocstr)
.def("GetStructOfArrays", py::overload_cast<>(&ParticleTileType::GetStructOfArrays),
py::return_value_policy::reference_internal)
.def("empty", &ParticleTileType::empty)
.def("size", &ParticleTileType::template size<ParticleType>)
.def("numParticles", &ParticleTileType::template numParticles<ParticleType>)
.def("numRealParticles", &ParticleTileType::template numRealParticles<ParticleType>)
.def("numNeighborParticles", &ParticleTileType::template numNeighborParticles<ParticleType>)
.def("numTotalParticles", &ParticleTileType::template numTotalParticles<ParticleType>)
//.def("size", &ParticleTileType::template size<ParticleType>)
//.def("numParticles", &ParticleTileType::template numParticles<ParticleType>)
//.def("numRealParticles", &ParticleTileType::template numRealParticles<ParticleType>)
//.def("numNeighborParticles", &ParticleTileType::template numNeighborParticles<ParticleType>)
//.def("numTotalParticles", &ParticleTileType::template numTotalParticles<ParticleType>)
.def("setNumNeighbors", &ParticleTileType::setNumNeighbors)
.def("getNumNeighbors", &ParticleTileType::getNumNeighbors)
.def("resize", &ParticleTileType::resize)
Expand Down

0 comments on commit 9781ae7

Please sign in to comment.