From 8ee6305f1f1bb75b517ec484f1d9a02ef680c0b0 Mon Sep 17 00:00:00 2001
From: Casper da Costa-Luis <casper.dcl@physics.org>
Date: Sun, 17 Jan 2021 21:16:41 +0000
Subject: [PATCH] add `zeros`, `from_numpy` helpers

- deprecate `vector`
- add tests
---
 cuvec/__init__.py     |  9 +++++----
 cuvec/helpers.py      | 24 ++++++++++++++++++++++++
 cuvec/pycuvec.py      |  5 -----
 tests/test_cuvec.py   | 23 +++++++----------------
 tests/test_pycuvec.py | 22 ++++++++++++++++++++++
 5 files changed, 58 insertions(+), 25 deletions(-)
 create mode 100644 cuvec/helpers.py
 create mode 100644 tests/test_pycuvec.py

diff --git a/cuvec/__init__.py b/cuvec/__init__.py
index 5c19e26..697e39a 100644
--- a/cuvec/__init__.py
+++ b/cuvec/__init__.py
@@ -19,9 +19,9 @@
     # config
     'cmake_prefix', 'include_path',
     # functions
-    'dev_sync', 'vector',
+    'dev_sync', 'from_numpy', 'zeros',
     # data
-    'typecodes'] # yapf: disable
+    'typecodes', 'vec_types'] # yapf: disable
 
 from pathlib import Path
 
@@ -29,11 +29,12 @@
 
 try:
     from .cuvec import dev_sync
-    from .pycuvec import typecodes, vector
 except ImportError as err:
     from warnings import warn
     warn(str(err), UserWarning)
-    dev_sync = vector = None
+else:
+    from .helpers import from_numpy, zeros
+    from .pycuvec import typecodes, vec_types
 
 # for use in `cmake -DCMAKE_PREFIX_PATH=...`
 cmake_prefix = Path(resource_filename(__name__, "cmake")).resolve()
diff --git a/cuvec/helpers.py b/cuvec/helpers.py
new file mode 100644
index 0000000..daf3dee
--- /dev/null
+++ b/cuvec/helpers.py
@@ -0,0 +1,24 @@
+"""Useful helper functions."""
+from collections.abc import Sequence
+
+import numpy as np
+
+from .pycuvec import vec_types
+
+
+def zeros(shape, dtype="float32"):
+    """
+    Returns a new `Vector_*` of the specified shape and data type
+    (`cuvec` equivalent of `numpy.zeros`).
+    """
+    return vec_types[np.dtype(dtype)](shape if isinstance(shape, Sequence) else (shape,))
+
+
+def from_numpy(arr):
+    """
+    Returns a new `Vector_*` of the specified shape and data type
+    (`cuvec` equivalent of `numpy.copy`).
+    """
+    res = zeros(arr.shape, arr.dtype)
+    np.asarray(res)[:] = arr[:]
+    return res
diff --git a/cuvec/pycuvec.py b/cuvec/pycuvec.py
index 3e2fa3c..964afa4 100644
--- a/cuvec/pycuvec.py
+++ b/cuvec/pycuvec.py
@@ -1,6 +1,5 @@
 """Thin wrappers around `cuvec` C++/CUDA module"""
 import array
-from collections.abc import Sequence
 
 import numpy as np
 
@@ -31,7 +30,3 @@
     np.dtype('uint64'): Vector_Q,
     np.dtype('float32'): Vector_f,
     np.dtype('float64'): Vector_d}
-
-
-def vector(shape, dtype=np.float32):
-    return vec_types[np.dtype(dtype)](shape if isinstance(shape, Sequence) else (shape,))
diff --git a/tests/test_cuvec.py b/tests/test_cuvec.py
index c4f3e04..415f7c7 100644
--- a/tests/test_cuvec.py
+++ b/tests/test_cuvec.py
@@ -4,19 +4,19 @@
 import cuvec
 
 
-@mark.parametrize("vtype", list(cuvec.typecodes))
-def test_Vector_asarray(vtype):
-    """vtype(char): any of bBhHiIqQfd"""
-    v = getattr(cuvec.cuvec, f"Vector_{vtype}")((1, 2, 3))
-    assert str(v) == f"Vector_{vtype}((1, 2, 3))"
+@mark.parametrize("tp", list(cuvec.typecodes))
+def test_Vector_asarray(tp):
+    """tp(char): any of bBhHiIqQfd"""
+    v = getattr(cuvec.cuvec, f"Vector_{tp}")((1, 2, 3))
+    assert str(v) == f"Vector_{tp}((1, 2, 3))"
     a = np.asarray(v)
     assert not a.any()
     a[0, 0] = 42
     b = np.asarray(v)
     assert (b[0, 0] == 42).all()
     assert not b[1:, 1:].any()
-    assert a.dtype.char == vtype
-    del a, b
+    assert a.dtype.char == tp
+    del a, b, v
 
 
 def test_Vector_strides():
@@ -25,12 +25,3 @@ def test_Vector_strides():
     a = np.asarray(v)
     assert a.shape == shape
     assert a.strides == (473344, 1376, 4)
-
-
-def test_vector():
-    shape = 127, 344, 344
-    a = np.asarray(cuvec.vector(shape, "i"))
-    assert a.dtype == np.int32
-
-    a = np.asarray(cuvec.vector(shape, "d"))
-    assert a.dtype == np.float64
diff --git a/tests/test_pycuvec.py b/tests/test_pycuvec.py
new file mode 100644
index 0000000..3b88ad2
--- /dev/null
+++ b/tests/test_pycuvec.py
@@ -0,0 +1,22 @@
+import numpy as np
+from pytest import mark
+
+import cuvec
+
+
+@mark.parametrize("spec,result", [("i", np.int32), ("d", np.float64)])
+def test_zeros(spec, result):
+    shape = 127, 344, 344
+    a = np.asarray(cuvec.zeros(shape, spec))
+    assert a.dtype == result
+    assert a.shape == shape
+    assert not a.any()
+
+
+def test_from_numpy():
+    shape = 127, 344, 344
+    a = np.random.random(shape)
+    b = np.asarray(cuvec.from_numpy(a))
+    assert a.shape == b.shape
+    assert a.dtype == b.dtype
+    assert (a == b).all()