diff --git a/src/Particle/ArrayOfStructs.cpp b/src/Particle/ArrayOfStructs.cpp index 00aa7072..8ece653d 100644 --- a/src/Particle/ArrayOfStructs.cpp +++ b/src/Particle/ArrayOfStructs.cpp @@ -15,6 +15,23 @@ namespace { using namespace amrex; + // Note - this function MUST be consistent with AMReX_Particle.H + Long unpack_id (uint64_t cpuid) { + Long r = 0; + + uint64_t sign = cpuid >> 63; // extract leftmost sign bit + uint64_t val = ((cpuid >> 24) & 0x7FFFFFFFFF); // extract next 39 id bits + + Long lval = static_cast(val); // bc we take - + r = (sign) ? lval : -lval; + return r; + } + + // Note - this function MUST be consistent with AMReX_Particle.H + int unpack_cpu (uint64_t cpuid) { + return static_cast(cpuid & 0x00FFFFFF); + } + /** CPU: __array_interface__ v3 * * https://numpy.org/doc/stable/reference/arrays.interface.html @@ -156,4 +173,7 @@ void init_ArrayOfStructs(py::module& m) { make_ArrayOfStructs<0, 0> (m); // WarpX 22.07, ImpactX 22.07, HiPACE++ 22.07 make_ArrayOfStructs<1, 1> (m); // test in ParticleContainer make_ArrayOfStructs<2, 1> (m); // test + + m.def("unpack_ids", py::vectorize(unpack_id)); + m.def("unpack_cpus", py::vectorize(unpack_cpu)); } diff --git a/tests/test_particleTile.py b/tests/test_particleTile.py index 4171066f..2a1bdcf2 100644 --- a/tests/test_particleTile.py +++ b/tests/test_particleTile.py @@ -137,3 +137,11 @@ def test_ptile_aos(): print(aos) assert np.isclose(aos[0]["x"], 3.0) and np.isclose(aos[0]["y"], 0) assert np.isclose(aos[2][0], 0.0) and np.isclose(aos[2][2], 20) + + +def test_ptile_aos(): + cpuids = np.array([100, 100, 100, 100, 100], dtype=np.uint64) + ids = amr.unpack_ids(cpuids) + cpus = amr.unpack_cpus(cpuids) + assert np.array_equal(ids, np.array([0, 0, 0, 0, 0])) + assert np.array_equal(cpus, np.array([100, 100, 100, 100, 100]))