From 215d2211dd4688017890bac5dfd31a1746947c41 Mon Sep 17 00:00:00 2001 From: Xiaoyan Wang Date: Mon, 24 Nov 2025 01:44:15 -0500 Subject: [PATCH 1/2] Array2D bindings to support passing data to/from numpy --- genmetaballs/src/cuda/bindings.cu | 20 +++++++++++++ genmetaballs/src/cuda/core/utils.cuh | 3 ++ .../src/genmetaballs/core/__init__.py | 3 +- tests/python_tests/test_utils.py | 30 ++++++++++++++++++- 4 files changed, 54 insertions(+), 2 deletions(-) diff --git a/genmetaballs/src/cuda/bindings.cu b/genmetaballs/src/cuda/bindings.cu index b07e8ad..f9f24f2 100644 --- a/genmetaballs/src/cuda/bindings.cu +++ b/genmetaballs/src/cuda/bindings.cu @@ -1,5 +1,6 @@ #include #include +#include #include #include @@ -79,4 +80,23 @@ 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))"); + nb::class_>(utils, "FloatArray2D") + // TODO: switch to the array_api in future nanobind release + // https://nanobind.readthedocs.io/en/latest/api_extra.html#_CPPv4N8nanobind9array_apiE + .def( + "numpy", + [](const Array2D& self) { + return nb::ndarray( + self.data(), {self.num_rows(), self.num_cols()}); + }, + nb::rv_policy::reference_internal) + .def_static("from_array", + [](const nb::ndarray, nb::c_contig>& array) { + return Array2D(array.data(), array.shape(0), array.shape(1)); + }) + .def_prop_ro("num_rows", &Array2D::num_rows) + .def_prop_ro("num_cols", &Array2D::num_cols) + .def_prop_ro("ndim", &Array2D::ndim) + .def_prop_ro("size", &Array2D::size); + } // NB_MODULE(_genmetaballs_bindings) diff --git a/genmetaballs/src/cuda/core/utils.cuh b/genmetaballs/src/cuda/core/utils.cuh index 43e5707..3f6aeef 100644 --- a/genmetaballs/src/cuda/core/utils.cuh +++ b/genmetaballs/src/cuda/core/utils.cuh @@ -56,6 +56,9 @@ public: CUDA_CALLABLE constexpr auto size() const noexcept { return data_view_.size(); } + CUDA_CALLABLE constexpr T* data() const noexcept { + return data_view_.data_handle(); + } }; // class Array2D // Type deduction guide diff --git a/genmetaballs/src/genmetaballs/core/__init__.py b/genmetaballs/src/genmetaballs/core/__init__.py index 23c7077..a09878d 100644 --- a/genmetaballs/src/genmetaballs/core/__init__.py +++ b/genmetaballs/src/genmetaballs/core/__init__.py @@ -3,9 +3,10 @@ TwoParameterConfidence, ZeroParameterConfidence, ) -from genmetaballs._genmetaballs_bindings.utils import sigmoid +from genmetaballs._genmetaballs_bindings.utils import FloatArray2D, sigmoid __all__ = [ + "FloatArray2D", "ZeroParameterConfidence", "TwoParameterConfidence", "geometry", diff --git a/tests/python_tests/test_utils.py b/tests/python_tests/test_utils.py index 41d2714..55cc0d3 100644 --- a/tests/python_tests/test_utils.py +++ b/tests/python_tests/test_utils.py @@ -2,7 +2,7 @@ import pytest from scipy.special import expit -from genmetaballs.core import sigmoid +from genmetaballs.core import FloatArray2D, sigmoid NUM_RNG_SEEDS_PER_TEST = 5 NUM_N_VALUES_PER_TEST = 5 @@ -58,3 +58,31 @@ def test_sigmoid_edge_cases(x: float) -> None: assert np.isclose(actual, expected, rtol=1e-5, atol=1e-6) assert actual >= 0.0 assert actual <= 1.0 + + +def test_float_array2d_creation_and_view(): + """Test creation of Array2D from a numpy array.""" + rows, cols = 4, 5 + data = np.arange(rows * cols, dtype=np.float32).reshape((rows, cols)) + array_2d = FloatArray2D.from_array(data) + + assert array_2d.num_rows == rows + assert array_2d.num_cols == cols + assert array_2d.ndim == 2 + + # then try converting back to numpy array via view + data_view = array_2d.numpy() + print(type(data_view)) + assert np.allclose(data, data_view) + + # check that the view is writable and changes reflect back to original data + data_view[0, 0] = 999.0 + assert np.isclose(data[0, 0], 999.0) + + +def test_create_invalid_array2d(): + """Test that creating Array2D with invalid dimensions raises errors.""" + data = np.arange(12, dtype=np.float32).reshape((3, 4)) + + with pytest.raises(TypeError): + FloatArray2D.from_array(data.reshape((3, 4, 1))) # not 2D From 16e874ea56432e48f8285b8d000b4d43568d3a6c Mon Sep 17 00:00:00 2001 From: Xiaoyan Wang Date: Mon, 24 Nov 2025 01:46:21 -0500 Subject: [PATCH 2/2] swap definition order --- genmetaballs/src/cuda/bindings.cu | 8 ++++---- tests/python_tests/test_utils.py | 1 - 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/genmetaballs/src/cuda/bindings.cu b/genmetaballs/src/cuda/bindings.cu index f9f24f2..2a94a2a 100644 --- a/genmetaballs/src/cuda/bindings.cu +++ b/genmetaballs/src/cuda/bindings.cu @@ -81,6 +81,10 @@ NB_MODULE(_genmetaballs_bindings, m) { utils.def("sigmoid", sigmoid, nb::arg("x"), "Compute the sigmoid function: 1 / (1 + exp(-x))"); nb::class_>(utils, "FloatArray2D") + .def_static("from_array", + [](const nb::ndarray, nb::c_contig>& array) { + return Array2D(array.data(), array.shape(0), array.shape(1)); + }) // TODO: switch to the array_api in future nanobind release // https://nanobind.readthedocs.io/en/latest/api_extra.html#_CPPv4N8nanobind9array_apiE .def( @@ -90,10 +94,6 @@ NB_MODULE(_genmetaballs_bindings, m) { self.data(), {self.num_rows(), self.num_cols()}); }, nb::rv_policy::reference_internal) - .def_static("from_array", - [](const nb::ndarray, nb::c_contig>& array) { - return Array2D(array.data(), array.shape(0), array.shape(1)); - }) .def_prop_ro("num_rows", &Array2D::num_rows) .def_prop_ro("num_cols", &Array2D::num_cols) .def_prop_ro("ndim", &Array2D::ndim) diff --git a/tests/python_tests/test_utils.py b/tests/python_tests/test_utils.py index 55cc0d3..bdfbca4 100644 --- a/tests/python_tests/test_utils.py +++ b/tests/python_tests/test_utils.py @@ -72,7 +72,6 @@ def test_float_array2d_creation_and_view(): # then try converting back to numpy array via view data_view = array_2d.numpy() - print(type(data_view)) assert np.allclose(data, data_view) # check that the view is writable and changes reflect back to original data