python - add take_array

python/ceed_vector.py

@@ -101,6 +101,29 @@ def set_array(self, array, memtype=MEM_HOST, cmode=COPY_VALUES):
             self._pointer[0], memtype, cmode, array_pointer)
         self._ceed._check_error(err_code)
 
+    # Take Vector's data array
+    def take_array(self, memtype=MEM_HOST):
+        """Take ownership of the array set by set_array with USE_POINTER and remove the array from the Vector
+
+           Args:
+             **memtype: memory type of the array being taken, default CEED_MEM_HOST
+
+           Returns:
+             *array: array passed to set_array"""
+
+        # Setup the pointer's pointer
+        array_pointer = ffi.new("CeedScalar **")
+
+        # libCEED call
+        err_code = lib.CeedVectorTakeArray(
+            self._pointer[0], memtype, array_pointer)
+        self._ceed._check_error(err_code)
+
+        # Return array
+        array = self._array_reference
+        self._array_reference = None
+        return array
+
     # Get Vector's data array
     def get_array(self, memtype=MEM_HOST):
         """Get read/write access to a Vector via the specified memory type.

python/tests/test-1-vector.py

@@ -208,6 +208,33 @@ def test_108(ceed_resource, capsys):
 
     assert abs(norm - 9.) < TOL
 
+# -------------------------------------------------------------------------------
+# Test set_array and take_array
+# -------------------------------------------------------------------------------
+
+
+def test_109(ceed_resource, capsys):
+    ceed = libceed.Ceed(ceed_resource)
+
+    n = 10
+    x = ceed.Vector(n)
+
+    a = np.arange(0, n, dtype=ceed.scalar_type())
+    for i in range(n):
+        a[i] = -3.14 if i == 3 else 0
+    x.set_array(a, cmode=libceed.USE_POINTER)
+
+    # Verify correct array
+    a = x.take_array()
+    for i in range(n):
+        assert abs(a[i] + (3.14 if i == 3 else 0)) < TOL
+
+    # And ensure access removed in Vector
+    x.set_value(0)
+    with x.array() as b:
+        b[5] = 3.14
+    assert abs(a[5]) < TOL
+
 # -------------------------------------------------------------------------------
 # Test taking the reciprocal of a vector
 # -------------------------------------------------------------------------------

tests/t109-vector.c

@@ -1,6 +1,6 @@
 /// @file
-/// Test CeedVectorSetArray to remove array access
-/// \test Test CeedVectorSetArray to remove array access
+/// Test CeedVectorTakeArray to remove array access
+/// \test Test CeedVectorTakeArray to remove array access
 #include <ceed.h>
 #include <math.h>
 #include <stdio.h>