sagemath · vbraun · Dec 8, 2024 · Oct 21, 2024 · Nov 8, 2024 · Nov 8, 2024
diff --git a/src/sage/modules/meson.build b/src/sage/modules/meson.build
@@ -76,7 +76,10 @@ foreach name, pyx : extension_data
   )
 endforeach
 
-extension_data_cpp = {'vector_mod2_dense': files('vector_mod2_dense.pyx')}
+extension_data_cpp = {
+  'numpy_util' : files('numpy_util.pyx'),
+  'vector_mod2_dense': files('vector_mod2_dense.pyx'),
+}
 
 foreach name, pyx : extension_data_cpp
   py.extension_module(

diff --git a/src/sage/modules/numpy_util.pyx b/src/sage/modules/numpy_util.pyx
@@ -0,0 +1,65 @@
+# sage.doctest: optional - numpy
+r"""
+Utility functions for numpy.
+"""
+
+cimport numpy as np
+import numpy as np
+from sage.libs.m4ri cimport *
+from libc.stdint cimport uintptr_t
+
+
+ctypedef fused numpy_integral:
+    np.int8_t
+    np.int32_t
+    np.int64_t
+
+
+cdef set_mzd_from_numpy_unsafe(mzd_t* entries, np.ndarray[numpy_integral, ndim=1] x):
+    """
+    Internal function.
+    Caller are responsible for checking the two arrays have the same length.
+    """
+    for i in range(len(x)):
+        mzd_write_bit(entries, 0, i, x[i] & 1)
+
+
+def set_mzd_from_numpy(uintptr_t entries_addr, Py_ssize_t degree, x):
+    """
+    Set the entries in ``entries`` from numpy array ``x``.
+
+    INPUT:
+
+    - ``entries_addr`` -- must be a ``mzd_t*`` casted to ``uintptr_t``; the casting
+      is necessary to pass it through Python boundary because of lazy import
+
+    - ``degree`` -- the length of the array
+
+    - ``x`` -- a numpy array of integers or booleans, or any other object (in which
+      case this function will return ``False``)
+
+    OUTPUT: ``True`` if successful, ``False`` otherwise. May throw ``ValueError``.
+    """
+    cdef Py_ssize_t i
+    cdef np.ndarray[np.npy_bool, ndim=1] x_bool
+    cdef mzd_t* entries = <mzd_t*>entries_addr
+    if isinstance(x, np.ndarray):
+        if x.ndim != 1:
+            raise ValueError("numpy array must have dimension 1")
+        if x.shape[0] != degree:
+            raise ValueError("numpy array must have the right length")
+        if x.dtype == np.int8:
+            set_mzd_from_numpy_unsafe(entries, <np.ndarray[np.int8_t, ndim=1]>x)
+            return True
+        if x.dtype == np.int32:
+            set_mzd_from_numpy_unsafe(entries, <np.ndarray[np.int32_t, ndim=1]>x)
+            return True
+        if x.dtype == np.int64:
+            set_mzd_from_numpy_unsafe(entries, <np.ndarray[np.int64_t, ndim=1]>x)
+            return True
+        if x.dtype == np.bool_:
+            x_bool = x
+            for i in range(degree):
+                mzd_write_bit(entries, 0, i, x_bool[i])
+            return True
+    return False
diff --git a/src/sage/modules/vector_mod2_dense.pyx b/src/sage/modules/vector_mod2_dense.pyx
@@ -46,6 +46,7 @@ from sage.rings.rational cimport Rational
 from sage.structure.element cimport Element, Vector
 from sage.structure.richcmp cimport rich_to_bool
 cimport sage.modules.free_module_element as free_module_element
+from libc.stdint cimport uintptr_t
 
 from sage.libs.m4ri cimport *
 
@@ -192,8 +193,44 @@ cdef class Vector_mod2_dense(free_module_element.FreeModuleElement):
             TypeError: can...t initialize vector from nonzero non-list
             sage: (GF(2)**0).zero_vector()
             ()
+
+        Check construction from numpy arrays::
+
+            sage: # needs numpy
+            sage: import numpy
+            sage: VS = VectorSpace(GF(2),3)
+            sage: VS(numpy.array([0,-3,7], dtype=numpy.int8))
+            (0, 1, 1)
+            sage: VS(numpy.array([0,-3,7], dtype=numpy.int32))
+            (0, 1, 1)
+            sage: VS(numpy.array([0,-3,7], dtype=numpy.int64))
+            (0, 1, 1)
+            sage: VS(numpy.array([False,True,False], dtype=bool))
+            (0, 1, 0)
+            sage: VS(numpy.array([[1]]))
+            Traceback (most recent call last):
+            ...
+            ValueError: numpy array must have dimension 1
+            sage: VS(numpy.array([1,2,3,4]))
+            Traceback (most recent call last):
+            ...
+            ValueError: numpy array must have the right length
+
+        Make sure it's reasonably fast::
+
+            sage: # needs numpy
+            sage: import numpy
+            sage: VS = VectorSpace(GF(2),2*10^7)
+            sage: v = VS(numpy.random.randint(0, 1, size=VS.dimension()))  # around 300ms
         """
-        cdef Py_ssize_t i
+        try:
+            import numpy
+        except ImportError:
+            pass
+        else:
+            from .numpy_util import set_mzd_from_numpy
+            if set_mzd_from_numpy(<uintptr_t>self._entries, self._degree, x):
+                return
         if isinstance(x, (list, tuple)):
             if len(x) != self._degree:
                 raise TypeError("x must be a list of the right length")