ArborX two particle sets neighbours

lid_driven_cavity_neighbours_arborx.cpp

/*

             *****************************************************
             *****************************************************
             *****************************************************
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *** O O O O O O O O O O O O O O O O O O O O O O O ***
             *****************************************************
             *****************************************************
             *****************************************************

  O = Fluid particles
  * = Boundary particles
 */
#include <Kokkos_Core.hpp>
#include <ArborX.hpp>
#include <iostream>
#include <fstream>
#include <random>    // For std::mt19937, std::uniform_int_distribution
#include <set>       // Optional: To avoid duplicates


using Point = ArborX::Point<3>;

// ------------------ MyValues + AccessTraits ------------------
template<typename MemorySpace>
struct MyValues {
    Kokkos::View<Point*, MemorySpace> fluid;
    Kokkos::View<Point*, MemorySpace> boundary;
};

struct Value {
    Point point;
    int index;
    char mesh_type; // 0 = fluid, 1 = boundary
};

template<typename MemorySpace>
struct ArborX::AccessTraits<MyValues<MemorySpace>>
{
    using Self = MyValues<MemorySpace>;
    using memory_space = MemorySpace;  // required

    // static function for size
    static KOKKOS_FUNCTION auto size(Self const &self) -> std::size_t
    {
        return self.fluid.size() + self.boundary.size();
    }

    // static function for get
    static KOKKOS_FUNCTION auto get(Self const &self, int index) -> Value
    {
        if(index < self.fluid.size())
            return Value{self.fluid(index), index, 0};
        index -= self.fluid.size();
        return Value{self.boundary(index), index, 1};
    }
};


struct MyIndexableGetter {
    KOKKOS_FUNCTION auto operator()(Value const& value) const {
        return value.point;
    }
};


// ------------------ Callback ------------------
struct PairIndexMesh {
    int index;
    char mesh_type;
};

struct MyCallback {
    template<typename Predicate, typename Value, typename OutputFunctor>
    KOKKOS_FUNCTION void operator()(Predicate const &predicate,
                                    Value const &value,
                                    OutputFunctor const &out) const
    {
        out(PairIndexMesh{value.index, value.mesh_type});
    }
};


// ------------------ Main Example ------------------
int main(int argc, char* argv[]) {
    Kokkos::ScopeGuard guard(argc, argv);

    using ExecutionSpace = Kokkos::DefaultExecutionSpace;
    using MemorySpace = ExecutionSpace::memory_space;
    ExecutionSpace exec{};

    // --- create fluid and boundary particles ---
    // Full grid with padding: 17 x 17
    float const spacing = 0.1f;
    int const padding = 3;
    int const nx_fluid = 11;
    int const ny_fluid = 11;
    int const n_fluid = nx_fluid * ny_fluid;

    int const nx_boundary = nx_fluid + 2 * padding;
    int const ny_boundary = ny_fluid + 2 * padding;
    int const n_boundary = 2 * nx_boundary * padding + 2 * ny_fluid * padding;


    Kokkos::View<Point*, MemorySpace> fluid("fluid", n_fluid);
    Kokkos::View<Point*, MemorySpace> boundary("boundary", n_boundary);

    // Initialize fluid block (center region)
    Kokkos::parallel_for(
                         "init_fluid",
                         Kokkos::MDRangePolicy<ExecutionSpace, Kokkos::Rank<2>>({0,0}, {nx_fluid, ny_fluid}),
                         KOKKOS_LAMBDA(int i, int j) {
                           int idx = i * ny_fluid + j;
                           float x = i * spacing;
                           float y = j * spacing;
                           fluid(idx) = Point{ float(x), float(y), float(0.) };
                           // fluid(idx) = Point{ x, y, 0.0f };
                         });

    int offset = 0;

    // --- Bottom boundary ---
    Kokkos::parallel_for("bottom_boundary",
                         Kokkos::MDRangePolicy<ExecutionSpace, Kokkos::Rank<2>>({0, 0}, {nx_boundary, padding}),
                         KOKKOS_LAMBDA(int i, int j) {
                           int idx = i * padding + j;
                           float x = (i - padding) * spacing;
                           float y = (j - padding) * spacing;
                           boundary(idx) = Point{ x, y, 0.0f };
                         });
    offset += nx_boundary * padding;

    // --- Top boundary ---
    Kokkos::parallel_for("top_boundary",
                         Kokkos::MDRangePolicy<ExecutionSpace, Kokkos::Rank<2>>({0, 0}, {nx_boundary, padding}),
                         KOKKOS_LAMBDA(int i, int j) {
                           int idx = offset + i * padding + j;
                           float x = (i - padding) * spacing;
                           float y = (ny_fluid + j) * spacing;
                           boundary(idx) = Point{ x, y, 0.0f };
                         });
    offset += nx_boundary * padding;

    // --- Left boundary ---
    Kokkos::parallel_for("left_boundary",
                         Kokkos::MDRangePolicy<ExecutionSpace, Kokkos::Rank<2>>({0, 0}, {padding, ny_fluid}),
                         KOKKOS_LAMBDA(int i, int j) {
                           int idx = offset + i * ny_fluid + j;
                           float x = (i - padding) * spacing;
                           float y = j * spacing;
                           boundary(idx) = Point{ x, y, 0.0f };
                         });
    offset += ny_fluid * padding;

    // --- Right boundary ---
    Kokkos::parallel_for("right_boundary",
                         Kokkos::MDRangePolicy<ExecutionSpace, Kokkos::Rank<2>>({0, 0}, {padding, ny_fluid}),
                         KOKKOS_LAMBDA(int i, int j) {
                           int idx = offset + i * ny_fluid + j;
                           float x = (nx_fluid + i) * spacing;
                           float y = j * spacing;
                           boundary(idx) = Point{ x, y, 0.0f };
                         });

    // --- build BVH ---
    MyValues<MemorySpace> all_particles{fluid, boundary};
    ArborX::BoundingVolumeHierarchy<MemorySpace, Value, MyIndexableGetter>
        bvh(exec, all_particles, MyIndexableGetter{});

    // --- query fluid neighbors for fluid particles ---
    Kokkos::View<PairIndexMesh*, MemorySpace> pairs("pairs", 0);
    Kokkos::View<int*, MemorySpace> offsets("offsets", 0);

    // using spheres of radius 1.0 for neighbors
    auto fluid_query = ArborX::Experimental::make_intersects(fluid, 2. * spacing + spacing * 0.5);
    bvh.query(exec, fluid_query, MyCallback{}, pairs, offsets);

    // --- copy results to host to print ---
    auto h_pairs = Kokkos::create_mirror_view(pairs);
    auto h_offsets = Kokkos::create_mirror_view(offsets);
    Kokkos::deep_copy(h_pairs, pairs);
    Kokkos::deep_copy(h_offsets, offsets);

    std::cout << "Fluid neighbors for fluid particles:\n";
    // int const n_fluid = h_offsets.extent(0) - 1;  // or however many fluid particles you have
    int const num_random_particles = 5;

    std::mt19937 rng(12345); // fixed seed for reproducibility
    std::uniform_int_distribution<int> dist(0, n_fluid - 1);

    std::set<int> selected_particles;
    while (selected_particles.size() < num_random_particles) {
        int i = dist(rng);
        selected_particles.insert(i); // avoid duplicates
    }

    std::cout << "Fluid neighbors for randomly selected fluid particles:\n";
    for (int i : selected_particles) {
        std::cout << "Fluid particle " << i << " neighbors: ";
        for (int j = h_offsets(i); j < h_offsets(i + 1); ++j) {
            std::cout << "(" << h_pairs(j).index << "," << int(h_pairs(j).mesh_type) << ") ";
        }
        std::cout << "\n";
    }

    // --- query fluid neighbors for boundary particles ---
    auto boundary_query = ArborX::Experimental::make_intersects(boundary, 2. * spacing + spacing * 0.5);
    bvh.query(exec, boundary_query, MyCallback{}, pairs, offsets);

    auto h1_pairs = Kokkos::create_mirror_view(pairs);
    auto h1_offsets = Kokkos::create_mirror_view(offsets);
    Kokkos::deep_copy(h1_pairs, pairs);
    Kokkos::deep_copy(h1_offsets, offsets);
    int const num_random_boundary = 5;

    std::mt19937 rng_boundary(54321); // different seed from fluid
    std::uniform_int_distribution<int> dist_boundary(0, n_boundary - 1);

    std::set<int> selected_boundary_particles;
    while (selected_boundary_particles.size() < num_random_boundary) {
      int i = dist_boundary(rng_boundary);
      selected_boundary_particles.insert(i); // avoid duplicates
    }

    std::cout << "\nFluid neighbors for randomly selected boundary particles:\n";
    for (int i : selected_boundary_particles) {
      std::cout << "Boundary particle " << i << " neighbors: ";
      for (int j = h1_offsets(i); j < h1_offsets(i + 1); ++j) {
        std::cout << "(" << h1_pairs(j).index << "," << int(h1_pairs(j).mesh_type) << ") ";
      }
      std::cout << "\n";
    }


    // Write the data to file for plotting
    std::ofstream fluid_file("fluid.csv");
    fluid_file << "id,x,y,z\n";
    for (int i = 0; i < n_fluid; ++i)
      {
        const auto &p = fluid(i);
        fluid_file << i << "," << p[0] << "," << p[1] << "," << p[2] << "\n";
      }

    std::ofstream boundary_file("boundary.csv");
    boundary_file << "id,x,y,z\n";
    for (int i = 0; i < n_boundary; ++i)
      {
        const auto &p = boundary(i);
        boundary_file << i << "," << p[0] << "," << p[1] << "," << p[2] << "\n";
      }


    std::ofstream fluid_nfile("fluid_neighbors.txt");
    for (int i : selected_particles) {
      fluid_nfile << i << ":\n";
      for (int j = h_offsets(i); j < h_offsets(i + 1); ++j) {
        fluid_nfile << "(" << h_pairs(j).index << " " << int(h_pairs(j).mesh_type) << ")\n";
      }
      std::cout << "\n";
    }

    std::ofstream boundary_nfile("boundary_neighbors.txt");
    for (int i : selected_boundary_particles) {
      {
        boundary_nfile << i << ":\n";
        for (int j = h1_offsets(i); j < h1_offsets(i + 1); ++j) {
          boundary_nfile << "(" << h1_pairs(j).index << " " << int(h1_pairs(j).mesh_type) << ")\n";
        }
        std::cout << "\n";
      }
    }

    return 0;
}