Finalize cell list code and add tests

tests/cuda_unit_tests/test_neighbor_list.cu

@@ -2,6 +2,8 @@
 #include "box.hpp"
 #include "neighbor_list.cuh"
 #include "gtest/gtest.h"
+#include <thrust/device_vector.h>
+#include <thrust/host_vector.h>
 
 TEST(CellListTest, Constructor) {
   // Test case 1: Simple case
@@ -27,6 +29,68 @@ TEST(CellListTest, GetCellIndex) {
   int expected_index = z * cell_list.grid_size.x * cell_list.grid_size.y +
                        y * cell_list.grid_size.x + x;
 
-  EXPECT_EQ(cell_list.get_cell_index_from_cell_coords({x, y, z}),
-            expected_index);
+  EXPECT_EQ(
+      CellList::get_cell_index_from_cell_coords({x, y, z}, cell_list.grid_size),
+      expected_index);
+}
+
+TEST(CellListTest, AssignParticlesToCells) {
+  Box box(0, 2, 0, 1, 0, 1); // Create a 2x1x1 box
+  float cutoff = 1.0;
+  CellList cell_list(4, box, cutoff); // 4 particles
+
+  // Particle positions
+  thrust::host_vector<float3> h_positions(4);
+  h_positions[0] = make_float3(0.1, 0.1, 0.1); // Cell 0
+  h_positions[1] = make_float3(1.1, 0.1, 0.1); // Cell 1
+  h_positions[2] = make_float3(0.2, 0.1, 0.1); // Cell 0
+  h_positions[3] = make_float3(1.2, 0.1, 0.1); // Cell 1
+
+  thrust::device_vector<float3> d_positions = h_positions;
+
+  cell_list.assign_particles_to_cells(
+      thrust::raw_pointer_cast(d_positions.data()));
+
+  // Copy results back to host
+  thrust::host_vector<int> h_particle_cells(cell_list.n_particles);
+  thrust::host_vector<int> h_sorted_particles(cell_list.n_particles);
+  thrust::host_vector<int> h_cell_starts(cell_list.total_cells);
+  thrust::host_vector<int> h_cell_count(cell_list.total_cells);
+
+  cudaMemcpy(h_particle_cells.data(), cell_list.particle_cells,
+             cell_list.n_particles * sizeof(int), cudaMemcpyDeviceToHost);
+  cudaMemcpy(h_sorted_particles.data(), cell_list.sorted_particles,
+             cell_list.n_particles * sizeof(int), cudaMemcpyDeviceToHost);
+  cudaMemcpy(h_cell_starts.data(), cell_list.cell_starts,
+             cell_list.total_cells * sizeof(int), cudaMemcpyDeviceToHost);
+  cudaMemcpy(h_cell_count.data(), cell_list.cell_count,
+             cell_list.total_cells * sizeof(int), cudaMemcpyDeviceToHost);
+
+  // Expected values
+  // Particles 0 and 2 are in cell 0, particles 1 and 3 are in cell 1
+  // Sorted by cell: (0, 2, 1, 3)
+  // Cell 0: particles 0, 2
+  // Cell 1: particles 1, 3
+
+  // Verify particle_cells (after sorting by key, this will be sorted)
+  // Original particle_cells: [0, 1, 0, 1]
+  // After sort_by_key, particle_cells will be sorted: [0, 0, 1, 1]
+  thrust::host_vector<int> expected_particle_cells = {0, 0, 1, 1};
+  EXPECT_EQ(h_particle_cells, expected_particle_cells);
+
+  // Verify sorted_particles (original indices sorted by cell)
+  thrust::host_vector<int> expected_sorted_particles = {0, 2, 1, 3};
+  EXPECT_EQ(h_sorted_particles, expected_sorted_particles);
+
+  // Verify cell_starts
+  // Cell 0 starts at index 0 in sorted_particles
+  // Cell 1 starts at index 2 in sorted_particles
+  thrust::host_vector<int> expected_cell_starts = {0, 2};
+  EXPECT_EQ(h_cell_starts, expected_cell_starts);
+
+  // Verify cell_count
+  // Cell 0 has 2 particles
+  // Cell 1 has 2 particles
+  thrust::host_vector<int> expected_cell_count = {2, 2};
+  EXPECT_EQ(h_cell_count, expected_cell_count);
 }