Repository: equinor/GraphSPME (language: python)

Issue:
Validate input to `prec_sparse` in Python
In particular
- number of columns of `x` should equal both dimensions of `Z` (also implying Z is a square matrix)
- `Z` should be of a specific sparsity-type


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index b23e60d..843936e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -237,3 +237,4 @@ dmypy.json
 # Cython debug symbols
 cython_debug/
 
+eigen
diff --git a/python-package/pyproject.toml b/python-package/pyproject.toml
index 6191684..5b370cd 100644
--- a/python-package/pyproject.toml
+++ b/python-package/pyproject.toml
@@ -9,7 +9,7 @@ build-backend = "setuptools.build_meta"
 
 [tool.cibuildwheel]
 test-command = "cd {project}/python-package && pytest tests"
-test-skip = "*universal2:arm64"
+test-skip = "pp* *universal2:arm64 cp310-manylinux_i686 *musllinux* cp310-win32"
 before-all = ["git clone https://gitlab.com/libeigen/eigen.git"]
 
 [tool.cibuildwheel.windows]
diff --git a/python-package/setup.py b/python-package/setup.py
index 4a6fa9d..994d939 100644
--- a/python-package/setup.py
+++ b/python-package/setup.py
@@ -1,7 +1,6 @@
 import os
-import sysconfig
 from glob import glob
-from setuptools import setup
+from setuptools import find_packages, setup
 
 from pybind11.setup_helpers import Pybind11Extension
 
@@ -9,7 +8,7 @@
 
 ext_modules = [
     Pybind11Extension(
-        "graph_spme",
+        "_graphspme",
         sorted(glob("src/*.cpp")),
         cxx_std=14,
         include_dirs=[
@@ -31,6 +30,9 @@
     description="High dimensional precision matrix estimation with a known graphical structure",
     long_description="",
     ext_modules=ext_modules,
+    install_requires=["numpy", "scipy"],
+    package_dir={"": "src"},
+    packages=find_packages(where="src"),
     extras_require={"test": "pytest"},
     zip_safe=False,
     python_requires=">=3.6",
diff --git a/python-package/src/graphspme/__init__.py b/python-package/src/graphspme/__init__.py
new file mode 100644
index 0000000..fff0cd4
--- /dev/null
+++ b/python-package/src/graphspme/__init__.py
@@ -0,0 +1,42 @@
+import sys
+
+import numpy
+import scipy.sparse
+
+from _graphspme import _prec_sparse, _cov_shrink_spd
+
+
+if sys.version_info >= (3, 7):
+    import numpy.typing
+
+    NDArray = numpy.typing.NDArray
+else:
+    NDArray = numpy.ndarray
+
+
+def prec_sparse(
+    x: NDArray,
+    Z: scipy.sparse.csr_matrix,
+    cov_shrinkage: bool = True,
+) -> scipy.sparse.csr_matrix:
+    if not scipy.sparse.isspmatrix_csr(Z):
+        raise ValueError(
+            "Z matrix is not on csr (Compressed Sparse Row) format. "
+            "This is usually solved by adding format='csr' to the"
+            "matrix initialization, e.g. identity(n=10, format='csr')"
+        )
+    if len(x.shape) != 2:
+        raise ValueError(
+            f"x should be a two dimensional matrix, but has shape: {x.shape}"
+        )
+    _, n = x.shape
+    if Z.shape != (n, n):
+        raise ValueError(
+            "If x has shape (p,n) then Z should have shape (n,n), "
+            f"but got x.shape = {x.shape} and Z.shape = {Z.shape}"
+        )
+    return _prec_sparse(x, Z, cov_shrinkage)
+
+
+def cov_shrink_spd(x: NDArray) -> NDArray:
+    return _cov_shrink_spd(x)
diff --git a/python-package/src/pybind.cpp b/python-package/src/pybind.cpp
index 513d8d3..d5e0934 100644
--- a/python-package/src/pybind.cpp
+++ b/python-package/src/pybind.cpp
@@ -8,12 +8,13 @@
 
 namespace py = pybind11;
 
-PYBIND11_MODULE(graph_spme, m) {
+PYBIND11_MODULE(_graphspme, m)
+{
     m.doc() = R"pbdoc(
         GraphSPME
         ---------
 
-        .. currentmodule:: graph_spme
+        .. currentmodule:: _graphspme
 
         .. autosummary::
            :toctree: _generate
@@ -21,10 +22,22 @@ PYBIND11_MODULE(graph_spme, m) {
            prec_sparse
     )pbdoc";
 
-    m.def("prec_sparse", &prec_sparse, R"pbdoc(
+    m.def("_prec_sparse", &prec_sparse, R"pbdoc(
         prec_sparse
     )pbdoc");
 
+    m.def("_cov_shrink_spd", &cov_shrink_spd, R"pbdoc(
+        cov_shrink_spd
+    )pbdoc");
+
+    m.def("_sparse_matrix_inverse", &sparse_matrix_inverse, R"pbdoc(
+        sparse_matrix_inverse
+    )pbdoc");
+
+    m.def("_cov_ml", &cov_ml, R"pbdoc(
+        cov_ml
+    )pbdoc");
+
 #ifdef VERSION_INFO
     m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO);
 #else


Test patch:
diff --git a/python-package/tests/test_dummy.py b/python-package/tests/test_dummy.py
deleted file mode 100644
index f3b5d69..0000000
--- a/python-package/tests/test_dummy.py
+++ /dev/null
@@ -1,3 +0,0 @@
-
-def test_dummy():
-    assert True
diff --git a/python-package/tests/test_prec_sparse.py b/python-package/tests/test_prec_sparse.py
new file mode 100644
index 0000000..a546d08
--- /dev/null
+++ b/python-package/tests/test_prec_sparse.py
@@ -0,0 +1,32 @@
+import pytest
+import numpy as np
+from scipy import sparse
+
+import graphspme
+
+
+def test_prec_sparse_non_csr():
+    x = np.tile(np.ones([4]), (4, 1))
+    Z = sparse.diags([1, 1, 1, 1], 0, format="csc")
+    with pytest.raises(
+        ValueError, match=r".*Z matrix is not on csr \(Compressed Sparse Row\) format.*"
+    ):
+        graphspme.prec_sparse(x, Z)
+
+
+@pytest.mark.parametrize(
+    "x", [np.ones(shape=(5,)), np.random.normal(0, 1, size=(5, 3, 5))]
+)
+def test_prec_sparse_x_wrong_dimensions(x):
+    Z = sparse.diags([1, 1, 1], 0, format="csr")
+    with pytest.raises(ValueError, match=r".*x should be a two dimensional matrix.*"):
+        graphspme.prec_sparse(x, Z)
+
+
+def test_prec_sparse_dimension_mismatch():
+    x = np.tile(np.ones([4]), (4, 1))
+    Z = sparse.diags([1, 1, 1], 0, format="csr")
+    with pytest.raises(
+        ValueError, match=r".*but got x.shape = \(4, 4\) and Z.shape = \(3, 3\)"
+    ):
+        graphspme.prec_sparse(x, Z)