diff --git a/python/ceed_vector.py b/python/ceed_vector.py index 06bd693ec6..2def121f6b 100644 --- a/python/ceed_vector.py +++ b/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. diff --git a/python/tests/test-1-vector.py b/python/tests/test-1-vector.py index 9838a35b30..8462689fab 100644 --- a/python/tests/test-1-vector.py +++ b/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 # ------------------------------------------------------------------------------- diff --git a/tests/t109-vector.c b/tests/t109-vector.c index c30273ebd8..61bd588324 100644 --- a/tests/t109-vector.c +++ b/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 #include #include