diff --git a/genmetaballs/src/cuda/bindings.cu b/genmetaballs/src/cuda/bindings.cu index 7494660..dff316f 100644 --- a/genmetaballs/src/cuda/bindings.cu +++ b/genmetaballs/src/cuda/bindings.cu @@ -4,6 +4,7 @@ #include #include +#include "core/blender.cuh" #include "core/confidence.cuh" #include "core/geometry.cuh" #include "core/utils.cuh" @@ -64,17 +65,26 @@ NB_MODULE(_genmetaballs_bindings, m) { .def_rw("direction", &Ray::direction); /* - * Confidence module bindings + * Confidence submodule bindings */ nb::module_ confidence = m.def_submodule("confidence"); nb::class_(confidence, "ZeroParameterConfidence") .def(nb::init<>()) - .def("get_confidence", &ZeroParameterConfidence::get_confidence); + .def("get_confidence", &ZeroParameterConfidence::get_confidence, nb::arg("sumexpd"), + "Get the confidence value for a given sumexpd") + .def("__repr__", + [](const ZeroParameterConfidence& c) { return nb::str("ZeroParameterConfidence()"); }); nb::class_(confidence, "TwoParameterConfidence") .def(nb::init()) - .def("get_confidence", &TwoParameterConfidence::get_confidence); + .def_ro("beta4", &TwoParameterConfidence::beta4) + .def_ro("beta5", &TwoParameterConfidence::beta5) + .def("get_confidence", &TwoParameterConfidence::get_confidence, nb::arg("sumexpd"), + "Get the confidence value for a given sumexpd") + .def("__repr__", [](const TwoParameterConfidence& c) { + return nb::str("TwoParameterConfidence(beta4={}, beta5={})").format(c.beta4, c.beta5); + }); /* * Utils module bindings @@ -83,6 +93,32 @@ NB_MODULE(_genmetaballs_bindings, m) { nb::module_ utils = m.def_submodule("utils"); utils.def("sigmoid", sigmoid, nb::arg("x"), "Compute the sigmoid function: 1 / (1 + exp(-x))"); + // blender submodule + nb::module_ blender = m.def_submodule("blender"); + nb::class_(blender, "FourParameterBlender") + .def(nb::init()) + .def_ro("beta1", &FourParameterBlender::beta1) + .def_ro("beta2", &FourParameterBlender::beta2) + .def_ro("beta3", &FourParameterBlender::beta3) + .def_ro("eta", &FourParameterBlender::eta) + .def("blend", &FourParameterBlender::blend, nb::arg("t"), nb::arg("d"), + "Blend two values with (t,d)") + .def("__repr__", [](const FourParameterBlender& b) { + return nb::str("FourParameterBlender(beta1={}, beta2={}, beta3={}, eta={})") + .format(b.beta1, b.beta2, b.beta3, b.eta); + }); + + nb::class_(blender, "ThreeParameterBlender") + .def(nb::init()) + .def_ro("beta1", &ThreeParameterBlender::beta1) + .def_ro("beta2", &ThreeParameterBlender::beta2) + .def_ro("eta", &ThreeParameterBlender::eta) + .def("blend", &ThreeParameterBlender::blend, nb::arg("t"), nb::arg("d"), + "Blend two values with (t,d)") + .def("__repr__", [](const ThreeParameterBlender& b) { + return nb::str("ThreeParameterBlender(beta1={}, beta2={}, eta={})") + .format(b.beta1, b.beta2, b.eta); + }); bind_array2d(utils, "CPUFloatArray2D"); bind_array2d(utils, "GPUFloatArray2D"); diff --git a/genmetaballs/src/cuda/core/blender.cuh b/genmetaballs/src/cuda/core/blender.cuh index 5ce8c7a..774fab0 100644 --- a/genmetaballs/src/cuda/core/blender.cuh +++ b/genmetaballs/src/cuda/core/blender.cuh @@ -1,14 +1,27 @@ #pragma once -#include "fmb.h" -#include "geometry.h" +#include +#include + +#include "utils.cuh" + +struct FourParameterBlender { + float beta1; + float beta2; + float beta3; + float eta; + + CUDA_CALLABLE __forceinline__ float blend(float t, float d) const { + return expf((beta1 * d * sigmoid((beta3 / eta) * t)) - ((beta2 / eta) * t)); + } +}; struct ThreeParameterBlender { float beta1; float beta2; float eta; - CUDA_CALLABLE __forceinline__ // TODO inline? - float - blend(float t, float d, const FMB& fmb, const Ray& ray) const; + CUDA_CALLABLE __forceinline__ float blend(float t, float d) const { + return expf((beta1 * d) - ((beta2 / eta) * t)); + } }; diff --git a/genmetaballs/src/cuda/core/forward.cu b/genmetaballs/src/cuda/core/forward.cu index a71caf6..1f28793 100644 --- a/genmetaballs/src/cuda/core/forward.cu +++ b/genmetaballs/src/cuda/core/forward.cu @@ -34,7 +34,7 @@ __global__ render_kernel(const Getter fmb_getter, const Blender blender, float w0 = 0.0f, tf = 0.0f, sumexpd = 0.0f; for (const auto& fmb : fmb_getter->get_metaballs(ray)) { const auto& [t, d] = Intersector::intersect(fmb, ray); - w = blender->blend(t, d, fmb, ray); + w = blender->blend(t, d); sumexpd += exp(d); tf += t; w0 += w; diff --git a/genmetaballs/src/genmetaballs/core/__init__.py b/genmetaballs/src/genmetaballs/core/__init__.py index e2c7b76..89c845b 100644 --- a/genmetaballs/src/genmetaballs/core/__init__.py +++ b/genmetaballs/src/genmetaballs/core/__init__.py @@ -1,4 +1,8 @@ from genmetaballs._genmetaballs_bindings import geometry +from genmetaballs._genmetaballs_bindings.blender import ( + FourParameterBlender, + ThreeParameterBlender, +) from genmetaballs._genmetaballs_bindings.confidence import ( TwoParameterConfidence, ZeroParameterConfidence, @@ -27,4 +31,6 @@ def array2d_float(data, device) -> CPUFloatArray2D | GPUFloatArray2D: "TwoParameterConfidence", "geometry", "sigmoid", + "FourParameterBlender", + "ThreeParameterBlender", ] diff --git a/tests/cpp_tests/test_blender.cu b/tests/cpp_tests/test_blender.cu new file mode 100644 index 0000000..8b692a4 --- /dev/null +++ b/tests/cpp_tests/test_blender.cu @@ -0,0 +1,164 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/blender.cuh" +#include "core/utils.cuh" + +template +__global__ void blender_kernel(const float* t, const float* d, float* blended, uint32_t n, + Blender blender) { + uint32_t i = threadIdx.x + (blockIdx.x * blockDim.x); + if (i < n) { + blended[i] = blender.blend(t[i], d[i]); + } +} + +constexpr uint32_t GRID_DIM = 256; +constexpr uint32_t BLOCK_DIM = 1024; + +template +std::vector gpu_blend(const std::vector& t_vec, const std::vector& d_vec, + Blender blender) { + auto n = static_cast(t_vec.size()); + auto nbytes = n * sizeof(float); + float *d_t = nullptr, *d_d = nullptr, *d_blended = nullptr; + std::vector result(n); + + CUDA_CHECK(cudaMalloc(&d_t, nbytes)); + CUDA_CHECK(cudaMalloc(&d_d, nbytes)); + CUDA_CHECK(cudaMalloc(&d_blended, nbytes)); + CUDA_CHECK(cudaMemcpy(d_t, t_vec.data(), nbytes, cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(d_d, d_vec.data(), nbytes, cudaMemcpyHostToDevice)); + + auto block_dim = BLOCK_DIM; + auto grid_dim = (n + block_dim - 1) / block_dim; + if (grid_dim > GRID_DIM) + grid_dim = GRID_DIM; + + blender_kernel<<>>(d_t, d_d, d_blended, n, blender); + + CUDA_CHECK(cudaMemcpy(result.data(), d_blended, nbytes, cudaMemcpyDeviceToHost)); + CUDA_CHECK(cudaFree(d_t)); + CUDA_CHECK(cudaFree(d_d)); + CUDA_CHECK(cudaFree(d_blended)); + return result; +} + +constexpr int NUM_RNG_SEEDS_PER_TEST = 5; +constexpr int NUM_N_VALUES_PER_TEST = 5; +constexpr uint32_t MASTER_SEED = 42; + +static std::vector blender_test_sizes() { + std::vector sizes; + for (int k = 0; k < NUM_N_VALUES_PER_TEST; ++k) + sizes.push_back(1 << (4 + k)); // 2^(4+k): [16, 32, 64, 128, 256] + return sizes; +} + +struct BlenderCase { + float beta1, beta2, beta3, eta; + const char* name; +}; + +static std::vector blender_cases() { + return { + {1.0F, 0.5F, 0.2F, 2.0F, "case1"}, + {-2.0F, 1.0F, -1.0F, 1.5F, "case2"}, + {0.0F, 0.0F, 1.0F, 1.0F, "case3"}, + {0.5F, -0.5F, 0.8F, 0.5F, "case4"}, + }; +} + +// Smoke test for FourParameterBlender +TEST(GpuBlenderTest, Blender_GPU_Smoke_FourParameter) { + auto sizes = blender_test_sizes(); + std::mt19937 master_gen(MASTER_SEED); + std::uniform_int_distribution seed_dist(0, std::numeric_limits::max()); + std::vector seeds(NUM_RNG_SEEDS_PER_TEST); + for (auto& s : seeds) + s = seed_dist(master_gen); + + for (int size_idx = 0; size_idx < static_cast(sizes.size()); ++size_idx) { + int N = sizes[size_idx]; + + for (const auto& blend_case : blender_cases()) { + for (uint32_t test_seed : seeds) { + SCOPED_TRACE(testing::Message() << "N=" << N << ", seed=" << test_seed + << ", blend_type=" << blend_case.name); + + std::mt19937 rng(test_seed); + std::uniform_real_distribution tdist(0.0F, 10.0F); + std::uniform_real_distribution ddist(0.0F, 10.0F); + + std::vector t_vec(N), d_vec(N); + for (int i = 0; i < N; ++i) { + t_vec[i] = tdist(rng); + d_vec[i] = ddist(rng); + } + + FourParameterBlender blender{blend_case.beta1, blend_case.beta2, blend_case.beta3, + blend_case.eta}; + + std::vector actual = gpu_blend(t_vec, d_vec, blender); + + ASSERT_EQ(actual.size(), static_cast(N)); + } + } + } +} + +// Smoke test for ThreeParameterBlender +struct ThreeParamBlenderCase { + float beta1, beta2, eta; + const char* name; +}; + +static std::vector threeparam_blender_cases() { + return { + {1.0F, 0.5F, 2.0F, "three_case1"}, + {-2.0F, 1.0F, 1.5F, "three_case2"}, + {0.0F, 0.0F, 1.0F, "three_case3"}, + {0.5F, -0.5F, 0.5F, "three_case4"}, + }; +} + +TEST(GpuBlenderTest, Blender_GPU_Smoke_ThreeParameter) { + auto sizes = blender_test_sizes(); + std::mt19937 master_gen(MASTER_SEED); + std::uniform_int_distribution seed_dist(0, std::numeric_limits::max()); + std::vector seeds(NUM_RNG_SEEDS_PER_TEST); + for (auto& s : seeds) + s = seed_dist(master_gen); + + for (int size_idx = 0; size_idx < static_cast(sizes.size()); ++size_idx) { + int N = sizes[size_idx]; + for (const auto& blend_case : threeparam_blender_cases()) { + for (uint32_t test_seed : seeds) { + SCOPED_TRACE(testing::Message() << "N=" << N << ", seed=" << test_seed + << ", blend_type=" << blend_case.name); + + std::mt19937 rng(test_seed); + std::uniform_real_distribution tdist(0.0F, 10.0F); + std::uniform_real_distribution ddist(0.0F, 10.0F); + + std::vector t_vec(N), d_vec(N); + for (int i = 0; i < N; ++i) { + t_vec[i] = tdist(rng); + d_vec[i] = ddist(rng); + } + + ThreeParameterBlender blender{blend_case.beta1, blend_case.beta2, blend_case.eta}; + + std::vector actual = gpu_blend(t_vec, d_vec, blender); + + ASSERT_EQ(actual.size(), static_cast(N)); + } + } + } +} diff --git a/tests/python_tests/test_blender.py b/tests/python_tests/test_blender.py new file mode 100644 index 0000000..4cb6793 --- /dev/null +++ b/tests/python_tests/test_blender.py @@ -0,0 +1,129 @@ +import numpy as np +import pytest +from scipy.special import expit + +from genmetaballs.core import ( + FourParameterBlender, + ThreeParameterBlender, +) + + +def ground_truth_four_parameter_blender( + beta1: float, + beta2: float, + beta3: float, + eta: float, + di: float | np.ndarray, + ti: float | np.ndarray, +) -> float | np.ndarray: + """ + Compute the four-parameter blender function for a single value or arrays. + + Implements: + wi = exp( (beta1 * di) * sigmoid((beta3/eta) * ti) - (beta2/eta) * ti ) + """ + sig = expit((beta3 / eta) * ti) + result = np.exp((beta1 * di * sig) - ((beta2 / eta) * ti)) + return result + + +def ground_truth_three_parameter_blender( + beta1: float, + beta2: float, + eta: float, + di: float | np.ndarray, + ti: float | np.ndarray, +) -> float | np.ndarray: + """ + Compute the three-parameter blender function for a single value or arrays. + + Implements: + wi = exp(beta1 * di - (beta2 / eta) * ti) + """ + result = np.exp((beta1 * di) - ((beta2 / eta) * ti)) + return result + + +NUM_RNG_SEEDS_PER_TEST = 5 +NUM_N_VALUES_PER_TEST = 5 +MASTER_SEED = 0 + +# Test data for different four-parameter blender configurations +BLENDER_TEST_CASES = [ + # Just use the parameter dict (no ground truth dict) + {"beta1": 1.0, "beta2": 0.5, "beta3": 0.2, "eta": 2.0}, + {"beta1": -2.0, "beta2": 1.0, "beta3": -1.0, "eta": 1.5}, + {"beta1": 0.0, "beta2": 0.0, "beta3": 1.0, "eta": 1.0}, + {"beta1": 0.5, "beta2": -0.5, "beta3": 0.8, "eta": 0.5}, +] + + +@pytest.mark.parametrize( + "rng_seed", np.random.default_rng(MASTER_SEED).integers(0, 2**32, size=NUM_RNG_SEEDS_PER_TEST) +) +@pytest.mark.parametrize("blender_kwargs", BLENDER_TEST_CASES) +def test_blender_single_value(rng_seed: int, blender_kwargs: dict) -> None: + """Test that FourParameterBlender computes correct blend values for a single value.""" + rng = np.random.default_rng(rng_seed) + blender = FourParameterBlender( + blender_kwargs["beta1"], + blender_kwargs["beta2"], + blender_kwargs["beta3"], + blender_kwargs["eta"], + ) + + di = rng.uniform(low=0, high=10.0, size=1).astype(np.float32).item() + ti = rng.uniform(low=0, high=10.0, size=1).astype(np.float32).item() + + expected = ground_truth_four_parameter_blender( + blender_kwargs["beta1"], + blender_kwargs["beta2"], + blender_kwargs["beta3"], + blender_kwargs["eta"], + di, + ti, + ) + + actual = blender.blend(t=ti, d=di) + + # check that the actual and expected values are close + assert np.isclose(actual, expected, rtol=1e-6) or (np.isnan(actual) and np.isnan(expected)) + # Check for finiteness (blend can in rare cases produce 0 or inf for extreme values) + assert np.isfinite(actual) + + # INSERT_YOUR_CODE + + +@pytest.mark.parametrize( + "rng_seed", np.random.default_rng(MASTER_SEED).integers(0, 2**32, size=NUM_RNG_SEEDS_PER_TEST) +) +@pytest.mark.parametrize( + "blender_kwargs", + [ + {"beta1": 1.0, "beta2": 0.5, "eta": 2.0}, + {"beta1": -2.0, "beta2": 1.0, "eta": 1.5}, + {"beta1": 0.0, "beta2": 0.0, "eta": 1.0}, + {"beta1": 0.5, "beta2": -0.5, "eta": 0.5}, + ], +) +def test_three_parameter_blender_single_value(rng_seed: int, blender_kwargs: dict) -> None: + """Test that ThreeParameterBlender computes correct blend values for a single value.""" + rng = np.random.default_rng(rng_seed) + blender = ThreeParameterBlender( + blender_kwargs["beta1"], blender_kwargs["beta2"], blender_kwargs["eta"] + ) + + di = rng.uniform(low=0, high=10.0, size=1).astype(np.float32).item() + ti = rng.uniform(low=0, high=10.0, size=1).astype(np.float32).item() + + # Expected from the definition: expf((beta1 * d) - ((beta2 / eta) * t)) + beta1 = blender_kwargs["beta1"] + beta2 = blender_kwargs["beta2"] + eta = blender_kwargs["eta"] + + expected = np.exp((beta1 * di) - ((beta2 / eta) * ti)) + actual = blender.blend(t=ti, d=di) + + # check close (rtol can be tight since it's directly a float formula) + assert np.isclose(actual, expected, rtol=1e-6) or (np.isnan(actual) and np.isnan(expected)) + assert np.isfinite(actual)