Skip to content

Commit

Permalink
Vectorized functions for unpacking the 64-bit amrex particle ids (#165)
Browse files Browse the repository at this point in the history
* vectorized functions for unpacking the 64-bit amrex particle ids
* remove C++ from Python file
* remove more C++
  • Loading branch information
atmyers committed Aug 2, 2023
1 parent f561142 commit e5bc881
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 0 deletions.
20 changes: 20 additions & 0 deletions src/Particle/ArrayOfStructs.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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<Long>(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<int>(cpuid & 0x00FFFFFF);
}

/** CPU: __array_interface__ v3
*
* https://numpy.org/doc/stable/reference/arrays.interface.html
Expand Down Expand Up @@ -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));
}
8 changes: 8 additions & 0 deletions tests/test_particleTile.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]))

0 comments on commit e5bc881

Please sign in to comment.