cudaCAC/tests/cuda_unit_tests/test_forces.cu

#include <cmath>
#include <cuda_runtime.h>
#include <gtest/gtest.h>
#include <vector>

// Include your header files
#include "forces.cuh"
#include "kernel_config.cuh"
#include "potentials/pair_potentials.cuh"
#include "precision.hpp"

class CudaForceKernelTest : public ::testing::Test {
protected:
  void SetUp() override {
    // Set up CUDA device
    cudaError_t err = cudaSetDevice(0);
    ASSERT_EQ(err, cudaSuccess) << "Failed to set CUDA device";
  }

  void TearDown() override {
    // Clean up any remaining GPU memory
    cudaDeviceReset();
  }

  // Helper function to check CUDA errors
  void checkCudaError(cudaError_t err, const std::string &operation) {
    ASSERT_EQ(err, cudaSuccess)
        << "CUDA error in " << operation << ": " << cudaGetErrorString(err);
  }

  // Helper function to allocate and copy data to GPU
  template <typename T>
  T *allocateAndCopyToGPU(const std::vector<T> &host_data) {
    T *device_ptr;
    size_t size = host_data.size() * sizeof(T);
    checkCudaError(cudaMalloc(&device_ptr, size), "cudaMalloc");
    checkCudaError(
        cudaMemcpy(device_ptr, host_data.data(), size, cudaMemcpyHostToDevice),
        "cudaMemcpy H2D");
    return device_ptr;
  }

  // Helper function to copy data from GPU and free GPU memory
  template <typename T>
  std::vector<T> copyFromGPUAndFree(T *device_ptr, size_t count) {
    std::vector<T> host_data(count);
    size_t size = count * sizeof(T);
    checkCudaError(
        cudaMemcpy(host_data.data(), device_ptr, size, cudaMemcpyDeviceToHost),
        "cudaMemcpy D2H");
    checkCudaError(cudaFree(device_ptr), "cudaFree");
    return host_data;
  }

  // Helper function to run the force calculation kernel
  std::vector<float4>
  run_force_calculation(int n_particles, const std::vector<float4> &positions,
                        const std::vector<real> &box_dimensions) {
    std::vector<float4> force_energies(n_particles,
                                       make_float4(0.0, 0.0, 0.0, 0.0));

    KernelConfig kernel_config = get_launch_config(n_particles);
    float4 *d_positions = allocateAndCopyToGPU(positions);
    float4 *d_force_energies = allocateAndCopyToGPU(force_energies);
    real *d_box_len = allocateAndCopyToGPU(box_dimensions);

    std::vector<PairPotentials> potentials = {LennardJones(1.0, 1.0, 3.0)};
    CAC::launch_force_kernels(d_positions, d_force_energies, n_particles,
                              d_box_len, potentials, kernel_config.blocks,
                              kernel_config.threads);

    checkCudaError(cudaGetLastError(), "kernel launch");
    checkCudaError(cudaDeviceSynchronize(), "kernel execution");

    std::vector<float4> result_force_energies =
        copyFromGPUAndFree(d_force_energies, n_particles);

    checkCudaError(cudaFree(d_positions), "cudaFree positions");
    checkCudaError(cudaFree(d_box_len), "cudaFree box_len");

    return result_force_energies;
  }
};

TEST_F(CudaForceKernelTest, BasicFunctionalityTest) {
  const int n_particles = 2;
  const real tolerance = 1e-5;

  // Set up test data - simple 2x2 grid of particles
  std::vector<float4> positions = {
      make_float4(0.0, 0.0, 0.0, 0.0), // particle 0
      make_float4(0.5, 0.0, 0.0, 0.0), // particle 1
  };

  std::vector<real> box_dimensions = {10.0, 10.0, 10.0};

  auto result_force_energies =
      run_force_calculation(n_particles, positions, box_dimensions);

  // Verify results - forces should be non-zero and energies should be
  // calculated
  bool has_nonzero_force = false;
  bool has_nonzero_energy = false;

  for (int i = 0; i < n_particles; i++) {
    if (std::abs(result_force_energies[i].x) > tolerance ||
        std::abs(result_force_energies[i].y) > tolerance ||
        std::abs(result_force_energies[i].z) > tolerance) {
      has_nonzero_force = true;
    }
    if (std::abs(result_force_energies[i].w) > tolerance) {
      has_nonzero_energy = true;
    }
  }

  EXPECT_TRUE(has_nonzero_force)
      << "Expected non-zero forces between particles";
  EXPECT_TRUE(has_nonzero_energy)
      << "Expected non-zero energies for particles ";
}

TEST_F(CudaForceKernelTest, PeriodicBoundaryConditionsTest) {
  const int n_particles = 2;
  const real tolerance = 1e-5;

  // Place particles near opposite edges of a small box
  std::vector<float4> positions = {
      make_float4(0.1, 0.0, 0.0, 0.0), // particle 0 near left edge
      make_float4(4.9, 0.0, 0.0, 0.0)  // particle 1 near right edge
  };
  std::vector<real> box_dimensions = {5.0, 5.0, 5.0}; // Small box to test PBC

  auto result_force_energies =
      run_force_calculation(n_particles, positions, box_dimensions);

  // With PBC, particles should interact as if they're close (distance ~0.2)
  // rather than far apart (distance ~4.8)
  EXPECT_GT(std::abs(result_force_energies[0].x), tolerance)
      << "Expected significant force due to PBC";
}

TEST_F(CudaForceKernelTest, SingleParticleTest) {
  const int n_particles = 1;

  std::vector<float4> positions = {make_float4(0.0, 0.0, 0.0, 0.0)};
  std::vector<real> box_dimensions = {10.0, 10.0, 10.0};

  auto result_force_energies =
      run_force_calculation(n_particles, positions, box_dimensions);
  // Single particle should have zero force and energy
  EXPECT_NEAR(result_force_energies[0].x, 0.0, 1e-10);
  EXPECT_NEAR(result_force_energies[0].y, 0.0, 1e-10);
  EXPECT_NEAR(result_force_energies[0].z, 0.0, 1e-10);
  EXPECT_NEAR(result_force_energies[0].w, 0.0, 1e-10);
}

TEST_F(CudaForceKernelTest, ForceSymmetryTest) {
  const int n_particles = 2;
  const real tolerance = 1e-5;

  std::vector<float4> positions = {
      make_float4(0.0, 0.0, 0.0, 0.0), // particle 0
      make_float4(1.5, 0.0, 0.0, 0.0)  // particle 1
  };
  std::vector<real> box_dimensions = {10.0, 10.0, 10.0};

  auto result_force_energies =
      run_force_calculation(n_particles, positions, box_dimensions);

  // Newton's third law: forces should be equal and opposite
  EXPECT_NEAR(result_force_energies[0].x, -result_force_energies[1].x,
              tolerance)
      << "Force x-components should be opposite";
  EXPECT_NEAR(result_force_energies[0].y, -result_force_energies[1].y,
              tolerance)
      << "Force y-components should be opposite";
  EXPECT_NEAR(result_force_energies[0].z, -result_force_energies[1].z,
              tolerance)
      << "Force z-components should be opposite";

  // Energies should be equal for symmetric particles
  EXPECT_NEAR(result_force_energies[0].w, result_force_energies[1].w, tolerance)
      << "Energies should be equal";
}