From 45b9f2bdb18b3fd9913066bcbe2e50ccfae15ccc Mon Sep 17 00:00:00 2001 From: gitRasheed Date: Sun, 14 Jun 2026 20:33:35 +0100 Subject: [PATCH 1/2] Return PM.matrix native results as NumPy arrays --- src/nns/_nnscore_bindings.cpp | 24 ++++++++++++----- src/nns/pm_matrix.py | 27 +++++++++---------- .../test_native_original_src_coverage.py | 4 +++ 3 files changed, 33 insertions(+), 22 deletions(-) diff --git a/src/nns/_nnscore_bindings.cpp b/src/nns/_nnscore_bindings.cpp index a046d196..59ba869b 100644 --- a/src/nns/_nnscore_bindings.cpp +++ b/src/nns/_nnscore_bindings.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -16,6 +17,7 @@ namespace { using Vector = nb::ndarray, nb::c_contig>; using IntVector = nb::ndarray, nb::c_contig>; +using Matrix = nb::ndarray, nb::f_contig>; std::size_t checked_size(const Vector& x, const char* name) { const std::size_t n = x.shape(0); @@ -25,6 +27,14 @@ std::size_t checked_size(const Vector& x, const char* name) { return n; } +Matrix matrix_from_column_major_vector(std::vector&& values, std::size_t dim) { + auto* storage = new std::vector(std::move(values)); + nb::capsule owner(storage, [](void* p) noexcept { + delete static_cast*>(p); + }); + return Matrix(storage->data(), {dim, dim}, owner, {1, static_cast(dim)}); +} + void check_same_size(const Vector& x, const Vector& y, const char* x_name, const char* y_name) { if (checked_size(x, x_name) != checked_size(y, y_name)) { throw std::invalid_argument(std::string(x_name) + " and " + y_name + " must have the same length."); @@ -136,14 +146,14 @@ nb::dict pm_matrix_dict(double degree_lpm, throw std::invalid_argument("target length must equal d."); } checked_flat_matrix_size(variable, n, d, "variable"); - const nns::PMMatrixResult result = nns::pm_matrix(degree_lpm, degree_upm, target.data(), - variable.data(), n, d, pop_adj, norm); + nns::PMMatrixResult result = nns::pm_matrix(degree_lpm, degree_upm, target.data(), + variable.data(), n, d, pop_adj, norm); nb::dict out; - out["cupm"] = result.cupm; - out["dupm"] = result.dupm; - out["dlpm"] = result.dlpm; - out["clpm"] = result.clpm; - out["cov.matrix"] = result.cov; + out["cupm"] = matrix_from_column_major_vector(std::move(result.cupm), result.dim); + out["dupm"] = matrix_from_column_major_vector(std::move(result.dupm), result.dim); + out["dlpm"] = matrix_from_column_major_vector(std::move(result.dlpm), result.dim); + out["clpm"] = matrix_from_column_major_vector(std::move(result.clpm), result.dim); + out["cov.matrix"] = matrix_from_column_major_vector(std::move(result.cov), result.dim); out["dim"] = result.dim; return out; } diff --git a/src/nns/pm_matrix.py b/src/nns/pm_matrix.py index a3655800..9a888bed 100644 --- a/src/nns/pm_matrix.py +++ b/src/nns/pm_matrix.py @@ -56,21 +56,11 @@ def pm_matrix( ) dim = int(native_result["dim"]) result: PMMatrixResult = { - "cupm": np.asarray(native_result["cupm"], dtype=np.float64).reshape( - (dim, dim), order="F" - ), - "dupm": np.asarray(native_result["dupm"], dtype=np.float64).reshape( - (dim, dim), order="F" - ), - "dlpm": np.asarray(native_result["dlpm"], dtype=np.float64).reshape( - (dim, dim), order="F" - ), - "clpm": np.asarray(native_result["clpm"], dtype=np.float64).reshape( - (dim, dim), order="F" - ), - "cov.matrix": np.asarray(native_result["cov.matrix"], dtype=np.float64).reshape( - (dim, dim), order="F" - ), + "cupm": _native_matrix(native_result["cupm"], dim), + "dupm": _native_matrix(native_result["dupm"], dim), + "dlpm": _native_matrix(native_result["dlpm"], dim), + "clpm": _native_matrix(native_result["clpm"], dim), + "cov.matrix": _native_matrix(native_result["cov.matrix"], dim), } if resolved_names is not None: result["names"] = resolved_names @@ -116,6 +106,13 @@ def pm_matrix( return result +def _native_matrix(value: Any, dim: int) -> NDArray[np.float64]: + matrix = np.asarray(value, dtype=np.float64) + if matrix.shape == (dim, dim): + return cast(NDArray[np.float64], matrix) + return cast(NDArray[np.float64], matrix.reshape((dim, dim), order="F")) + + def _resolve_names(names: Sequence[str] | None, n_cols: int) -> list[str] | None: if names is None: return None diff --git a/tests/invariants/test_native_original_src_coverage.py b/tests/invariants/test_native_original_src_coverage.py index b9a50e73..1d66edad 100644 --- a/tests/invariants/test_native_original_src_coverage.py +++ b/tests/invariants/test_native_original_src_coverage.py @@ -89,6 +89,10 @@ def test_direct_native_partial_moment_smoke(native: ModuleType) -> None: ) assert native_pm["dim"] == 2 assert set(native_pm) >= {"cupm", "dupm", "dlpm", "clpm", "cov.matrix", "dim"} + for key in ("cupm", "dupm", "dlpm", "clpm", "cov.matrix"): + assert isinstance(native_pm[key], np.ndarray) + assert native_pm[key].shape == (2, 2) + assert native_pm[key].flags.f_contiguous def test_direct_native_fast_lm_smoke(native: ModuleType) -> None: From 8a974e7659ee4a1f70a17bab4cdfd3f1c8e79ebf Mon Sep 17 00:00:00 2001 From: gitRasheed Date: Sun, 14 Jun 2026 20:39:31 +0100 Subject: [PATCH 2/2] Fix PM.matrix native helper typing --- src/nns/pm_matrix.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/nns/pm_matrix.py b/src/nns/pm_matrix.py index 9a888bed..946e8110 100644 --- a/src/nns/pm_matrix.py +++ b/src/nns/pm_matrix.py @@ -109,7 +109,7 @@ def pm_matrix( def _native_matrix(value: Any, dim: int) -> NDArray[np.float64]: matrix = np.asarray(value, dtype=np.float64) if matrix.shape == (dim, dim): - return cast(NDArray[np.float64], matrix) + return matrix return cast(NDArray[np.float64], matrix.reshape((dim, dim), order="F"))