diff --git a/cmake/dependencies/AMReX.cmake b/cmake/dependencies/AMReX.cmake index 375ebddc..db95e0b3 100644 --- a/cmake/dependencies/AMReX.cmake +++ b/cmake/dependencies/AMReX.cmake @@ -86,7 +86,7 @@ option(pyAMReX_amrex_internal "Download & build AMReX" ON) set(pyAMReX_amrex_repo "https://github.com/AMReX-Codes/amrex.git" CACHE STRING "Repository URI to pull and build AMReX from if(pyAMReX_amrex_internal)") -set(pyAMReX_amrex_branch "25.04" +set(pyAMReX_amrex_branch "793ea9f717590d66d178a86298b82aac244b77a7" CACHE STRING "Repository branch for pyAMReX_amrex_repo if(pyAMReX_amrex_internal)") diff --git a/src/Particle/ParticleContainer.cpp b/src/Particle/ParticleContainer.cpp index 8fc10726..421fcdb6 100644 --- a/src/Particle/ParticleContainer.cpp +++ b/src/Particle/ParticleContainer.cpp @@ -1,6 +1,6 @@ /* Copyright 2022 The AMReX Community * - * Authors: Ryan Sandberg, Axel Huebl + * Authors: Ryan Sandberg, Axel Huebl, Andrew Myers * License: BSD-3-Clause-LBNL */ #include "ParticleContainer.H" @@ -14,21 +14,70 @@ namespace { using namespace amrex; - // Note - this function MUST be consistent with AMReX_Particle.H - Long unpack_id (uint64_t idcpu) { - Long r = 0; + py::object pack_ids (py::array_t idcpus, + py::array_t ids) + { + if (idcpus.ndim() != 1) { + throw std::runtime_error("Input should be 1-D NumPy array"); + } + + auto buf = idcpus.request(); + auto buf2 = ids.request(); + if (buf.size != buf2.size) { + throw std::runtime_error("sizes do not match!"); + } + + int N = idcpus.shape()[0]; + for (int i = 0; i < N; i++) { + uint64_t* idcpus_ptr = (uint64_t*) buf.ptr; + amrex::Long* ids_ptr = (amrex::Long*) buf2.ptr; + particle_impl::pack_id(idcpus_ptr[i], ids_ptr[i]); + } + return py::cast(Py_None); + } - uint64_t sign = idcpu >> 63; // extract leftmost sign bit - uint64_t val = ((idcpu >> 24) & 0x7FFFFFFFFF); // extract next 39 id bits + py::object pack_cpus (py::array_t idcpus, + py::array_t cpus) + { + if (idcpus.ndim() != 1) { + throw std::runtime_error("Input should be 1-D NumPy array"); + } + + auto buf = idcpus.request(); + auto buf2 = cpus.request(); + if (buf.size != buf2.size) { + throw std::runtime_error("sizes do not match!"); + } + + int N = idcpus.shape()[0]; + for (int i = 0; i < N; i++) { + uint64_t* idcpus_ptr = (uint64_t*) buf.ptr; + int* cpus_ptr = (int*) buf2.ptr; + particle_impl::pack_cpu(idcpus_ptr[i], cpus_ptr[i]); + } + return py::cast(Py_None); + } - Long lval = static_cast(val); // bc we take - - r = (sign) ? lval : -lval; - return r; + Long unpack_id (uint64_t idcpu) { + return particle_impl::unpack_id(idcpu); } - // Note - this function MUST be consistent with AMReX_Particle.H int unpack_cpu (uint64_t idcpu) { - return static_cast(idcpu & 0x00FFFFFF); + return particle_impl::unpack_cpu(idcpu); + } + + uint64_t make_invalid (uint64_t idcpu) { + particle_impl::make_invalid(idcpu); + return idcpu; + } + + uint64_t make_valid (uint64_t idcpu) { + particle_impl::make_valid(idcpu); + return idcpu; + } + + bool is_valid (const uint64_t idcpu) { + return particle_impl::is_valid(idcpu); } } @@ -61,6 +110,11 @@ void init_ParticleContainer(py::module& m) { init_ParticleContainer_WarpX(m); // for particle idcpu arrays + m.def("pack_ids", &pack_ids); + m.def("pack_cpus", &pack_cpus); m.def("unpack_ids", py::vectorize(unpack_id)); m.def("unpack_cpus", py::vectorize(unpack_cpu)); + m.def("make_invalid", make_invalid); + m.def("make_valid", make_valid); + m.def("is_valid", is_valid); } diff --git a/tests/test_particleTile.py b/tests/test_particleTile.py index 6854275b..6c81732f 100644 --- a/tests/test_particleTile.py +++ b/tests/test_particleTile.py @@ -156,8 +156,61 @@ def test_ptile_aos_3d(): def test_ptile_aos(): - idcpu = np.array([100, 100, 100, 100, 100], dtype=np.uint64) + idcpu = np.array( + [ + 9223372036871553124, + 9223372036871553124, + 9223372036871553124, + 9223372036871553124, + 9223372036871553124, + ], + dtype=np.uint64, + ) + ids = amr.unpack_ids(idcpu) + cpus = amr.unpack_cpus(idcpu) + assert np.array_equal(ids, np.array([1, 1, 1, 1, 1])) + assert np.array_equal(cpus, np.array([100, 100, 100, 100, 100])) + + assert amr.is_valid(idcpu[0]) + idcpu[0] = amr.make_invalid(idcpu[0]) + assert not amr.is_valid(idcpu[0]) + idcpu[0] = amr.make_valid(idcpu[0]) + assert amr.is_valid(idcpu[0]) + + # the leftmost bit stores the sign of the id + # the next 39 store its absolute value + # the rightmost 24 then store the cpu number + # using this scheme, id = 1 cpu = 100 + # corresponds to 9223372036871553124 + assert idcpu[0] == 9223372036871553124 + + idcpu = np.array([0, 0, 0, 0, 0], dtype=np.uint64) + amr.pack_ids(idcpu, np.array([1, 1, 1, 1, 1], dtype=np.int64)) + amr.pack_cpus(idcpu, np.array([100, 100, 100, 100, 100], dtype=np.int32)) + print(idcpu) + assert np.array_equal( + idcpu, + np.array( + [ + 9223372036871553124, + 9223372036871553124, + 9223372036871553124, + 9223372036871553124, + 9223372036871553124, + ] + ), + ) + + idcpu = np.array( + [ + 9223372036871553124, + 9223372036871553124, + 9223372036871553124, + 9223372036871553124, + 9223372036871553124, + ] + ) ids = amr.unpack_ids(idcpu) cpus = amr.unpack_cpus(idcpu) - assert np.array_equal(ids, np.array([0, 0, 0, 0, 0])) + assert np.array_equal(ids, np.array([1, 1, 1, 1, 1])) assert np.array_equal(cpus, np.array([100, 100, 100, 100, 100]))