Add wrappers to invalidate particles and set id/cpu from Python (#435)

To see how this is used, see
[tests/test_particleTile.py](https://github.com/AMReX-Codes/pyamrex/pull/435/files#diff-ecc0338faa7592dbb084e945148b2aaa8f982e06dbc8cdcf72dcdf751a02dacb)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: atmyers

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)")
 

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<uint64_t> idcpus,
+                         py::array_t<amrex::Long> 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>(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<uint64_t> idcpus,
+                          py::array_t<int> 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>(Py_None);
+    }
 
-        Long lval = static_cast<Long>(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<int>(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);
 }

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]))