Rearranging all attributes for particles in a bunch is not possible ("sorting")

[aliemen]

Problem

Assume I want to sort particles in the container and computed an index array to do so. Currently this is only possible if the attribute name (which should be rearranged) is hardcoded, like this:

hash_type hash(n_part);
fill(hash); // saves argsort of some kind

view_type R_view = bunch_m->R.getView();
view_type tmp("tmp", n_part);
for (i = 0, ..., n_part) {
	tmp(i) = R_view(hash(i));
}
for (i = 0, ..., n_part) {
	R_view(i) = tmp(i);
}

The problem is that this does not rearrange all the other attributes and therefore does not really sort the bunch --> in this case, e.g. particle velocities are mixed with position.

Solution

ParticleAttribBase does not contain any information about the type of the attribute, which is why simply calling forAllAttributes(...) does not give access to the data of the attribute itself. However, one can use pack/unpack to rearrange them:

• Use pack to put particles in their own internal buffer (field buf_m in class ParticleAttrib).
• Usually unpack is used to append particles coming from other ranks at the end of the bunch. However, instead of appending, one can simply overwrite the values by modify unpack slightly: virtual void unpack(size_type, const bool overwrite = false) = 0;.

This is how the implementation would look like:

template <typename T, class... Properties>
void ParticleAttrib<T, Properties...>::unpack(size_type nrecvs, bool overwrite) {
    auto size          = dview_m.extent(0);
+   size_type required = overwrite ? nrecvs : (*(this->localNum_mp) + nrecvs);
    if (size < required) {
        int overalloc = Comm->getDefaultOverallocation();
        this->resize(required * overalloc);
    }

+   size_type count   = overwrite ? 0 : *(this->localNum_mp); // Changed this!
    using policy_type = Kokkos::RangePolicy<execution_space>;
    Kokkos::parallel_for(
        "ParticleAttrib::unpack()", policy_type(0, nrecvs),
        KOKKOS_CLASS_LAMBDA(const size_t i) { dview_m(count + i) = buf_m(i); });
    Kokkos::fence();
}

This should neither trigger the resizing, nor should it have any impact on performance.

Usage

Now, assuming the index array (of type ippl::hash_type) mentioned before was already calculated, one can rearrange all attributed like this:

bunch_m->template forAllAttributes([&]<typename Attribute>(Attribute*& attribute) {
    // Ensure indices are in the correct memory space --> copies data ONLY when different memory spaces, so should be efficient
    using memory_space    = typename Attribute::memory_space;
    auto indices_device = Kokkos::create_mirror_view_and_copy(memory_space{}, indices);

    attribute->pack(hash_array);
    attribute->unpack(localNumParticles, true);
});

_{Note: If you know the "index" of the attribute, you could also hardcode into this "loop" to skip the current one to save computation time (for example, since charge might be constant for every particle). This works either by comparing the pointers of your attribute and the functional Parameter or by skipping the i-th attribute (the iteration order of forAllAttributes should be preserved).}

Possible Improvements

• Currently, it is only possible to rearrange all particles at once or only an interval starting at index 0. However, one could pass a custom "beginIndex" to the unpack function. That would allow rearranging custom connected intervals of particle indices. E.g. only sorting particles from i = 100 up to i = 200.
• It would also be possible to pass a hash_type again, but the only use case I would see for that is implementing a "swapping" function for two particles.

[aliemen]

Update: I created a test file (test/particle/TestSorting.cpp) that initializes random charge values $0.5 \leq q \leq 10.5$, sorts an index array (of type ippl::hash_type) with bunch.q being the key and rearranges all attributes according to that index array (using the forAllAttribute... as explained in my problem description).

I tested LandauDamping both on CPU/GPU on multiple ranks and didn't see any problems. I haven't tested the rest of alpine yet, but I am confident that it does not affect other calls of unpack(...) (since the change to the function is quite small).