From a25c84d8ba4a6151a615c5edf3c29d2443a6b33f Mon Sep 17 00:00:00 2001 From: Andrew Myers Date: Wed, 16 Apr 2025 12:17:00 -0700 Subject: [PATCH 1/4] Add wrappers to invalidate particles and set id/cpu from Python --- cmake/dependencies/AMReX.cmake | 2 +- src/Particle/ParticleContainer.cpp | 72 +++++++++++++++++++++++++----- tests/test_particleTile.py | 23 +++++++++- 3 files changed, 83 insertions(+), 14 deletions(-) diff --git a/cmake/dependencies/AMReX.cmake b/cmake/dependencies/AMReX.cmake index 375ebddc..671351da 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 "development" 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..8ca4f52b 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,66 @@ 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 = (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 +106,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..74dab886 100644 --- a/tests/test_particleTile.py +++ b/tests/test_particleTile.py @@ -156,8 +156,27 @@ 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([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])) + + 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])) + 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([1, 1, 1, 1, 1])) assert np.array_equal(cpus, np.array([100, 100, 100, 100, 100])) From 6ca03bbfc5d2245ff6c8227681e96ac56694bf22 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 16 Apr 2025 19:18:15 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_particleTile.py | 42 +++++++++++++++++++++++++++++++------- 1 file changed, 35 insertions(+), 7 deletions(-) diff --git a/tests/test_particleTile.py b/tests/test_particleTile.py index 74dab886..ae44d867 100644 --- a/tests/test_particleTile.py +++ b/tests/test_particleTile.py @@ -156,26 +156,54 @@ def test_ptile_aos_3d(): def test_ptile_aos(): - idcpu = np.array([9223372036871553124, 9223372036871553124, 9223372036871553124, 9223372036871553124, 9223372036871553124], 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])) + assert amr.is_valid(idcpu[0]) idcpu[0] = amr.make_invalid(idcpu[0]) - assert(not amr.is_valid(idcpu[0])) + assert not amr.is_valid(idcpu[0]) idcpu[0] = amr.make_valid(idcpu[0]) - assert(amr.is_valid(idcpu[0])) - assert(idcpu[0] == 9223372036871553124) + assert amr.is_valid(idcpu[0]) + 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])) + assert np.array_equal( + idcpu, + np.array( + [ + 9223372036871553124, + 9223372036871553124, + 9223372036871553124, + 9223372036871553124, + 9223372036871553124, + ] + ), + ) - 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([1, 1, 1, 1, 1])) From 5a79696cd52bc45b0f2a6c80d4f1b269830757d7 Mon Sep 17 00:00:00 2001 From: Andrew Myers Date: Sat, 3 May 2025 09:06:53 -0700 Subject: [PATCH 3/4] fix cast for windows --- src/Particle/ParticleContainer.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Particle/ParticleContainer.cpp b/src/Particle/ParticleContainer.cpp index 8ca4f52b..095f43cd 100644 --- a/src/Particle/ParticleContainer.cpp +++ b/src/Particle/ParticleContainer.cpp @@ -15,7 +15,7 @@ namespace using namespace amrex; py::object pack_ids (py::array_t idcpus, - py::array_t ids) + py::array_t ids) { if (idcpus.ndim() != 1) { throw std::runtime_error("Input should be 1-D NumPy array"); @@ -28,7 +28,7 @@ namespace int N = idcpus.shape()[0]; for (int i = 0; i < N; i++) { uint64_t* idcpus_ptr = (uint64_t*) buf.ptr; - amrex::Long* ids_ptr = (long*) buf2.ptr; + amrex::Long* ids_ptr = (amrex::Long*) buf2.ptr; particle_impl::pack_id(idcpus_ptr[i], ids_ptr[i]); } return py::cast(Py_None); From ed5ad3a16790803ad2867b2344d58313009ddbb2 Mon Sep 17 00:00:00 2001 From: Andrew Myers Date: Mon, 5 May 2025 18:31:55 -0700 Subject: [PATCH 4/4] apply suggestions from code review --- cmake/dependencies/AMReX.cmake | 2 +- src/Particle/ParticleContainer.cpp | 8 ++++++-- tests/test_particleTile.py | 6 ++++++ 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/cmake/dependencies/AMReX.cmake b/cmake/dependencies/AMReX.cmake index 671351da..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 "development" +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 095f43cd..421fcdb6 100644 --- a/src/Particle/ParticleContainer.cpp +++ b/src/Particle/ParticleContainer.cpp @@ -23,7 +23,9 @@ namespace auto buf = idcpus.request(); auto buf2 = ids.request(); - if (buf.size != buf2.size) throw std::runtime_error("sizes do not match!"); + 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++) { @@ -43,7 +45,9 @@ namespace auto buf = idcpus.request(); auto buf2 = cpus.request(); - if (buf.size != buf2.size) throw std::runtime_error("sizes do not match!"); + 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++) { diff --git a/tests/test_particleTile.py b/tests/test_particleTile.py index ae44d867..6c81732f 100644 --- a/tests/test_particleTile.py +++ b/tests/test_particleTile.py @@ -176,6 +176,12 @@ def test_ptile_aos(): 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)