xref: /libCEED/python/ceed_vector.py (revision 7f1dc7b91d4f5ca67db10df8f4f0b60afe166f53)
1# Copyright (c) 2017, Lawrence Livermore National Security, LLC. Produced at
2# the Lawrence Livermore National Laboratory. LLNL-CODE-734707. All Rights
3# reserved. See files LICENSE and NOTICE for details.
4#
5# This file is part of CEED, a collection of benchmarks, miniapps, software
6# libraries and APIs for efficient high-order finite element and spectral
7# element discretizations for exascale applications. For more information and
8# source code availability see http://github.com/ceed.
9#
10# The CEED research is supported by the Exascale Computing Project 17-SC-20-SC,
11# a collaborative effort of two U.S. Department of Energy organizations (Office
12# of Science and the National Nuclear Security Administration) responsible for
13# the planning and preparation of a capable exascale ecosystem, including
14# software, applications, hardware, advanced system engineering and early
15# testbed platforms, in support of the nation's exascale computing imperative.
16
17from _ceed_cffi import ffi, lib
18import tempfile
19import numpy as np
20import numba.cuda as nbcuda
21import contextlib
22from .ceed_constants import MEM_HOST, COPY_VALUES, NORM_2
23
24# ------------------------------------------------------------------------------
25
26
27class Vector():
28    """Ceed Vector: storing and manipulating vectors."""
29
30    # Attributes
31    _ceed = ffi.NULL
32    _pointer = ffi.NULL
33
34    # Constructor
35    def __init__(self, ceed, size):
36        # CeedVector object
37        self._pointer = ffi.new("CeedVector *")
38
39        # Reference to Ceed
40        self._ceed = ceed
41
42        # libCEED call
43        lib.CeedVectorCreate(self._ceed._pointer[0], size, self._pointer)
44
45    # Destructor
46    def __del__(self):
47        # libCEED call
48        lib.CeedVectorDestroy(self._pointer)
49
50    # Representation
51    def __repr__(self):
52        return "<CeedVector instance at " + hex(id(self)) + ">"
53
54    # String conversion for print() to stdout
55    def __str__(self):
56        """View a Vector via print()."""
57
58        # libCEED call
59        fmt = ffi.new("char[]", "%f".encode('ascii'))
60        with tempfile.NamedTemporaryFile() as key_file:
61            with open(key_file.name, 'r+') as stream_file:
62                stream = ffi.cast("FILE *", stream_file)
63
64                lib.CeedVectorView(self._pointer[0], fmt, stream)
65
66                stream_file.seek(0)
67                out_string = stream_file.read()
68
69        return out_string
70
71    # Set Vector's data array
72    def set_array(self, array, memtype=MEM_HOST, cmode=COPY_VALUES):
73        """Set the array used by a Vector, freeing any previously allocated
74           array if applicable.
75
76           Args:
77             *array: Numpy or Numba array to be used
78             **memtype: memory type of the array being passed, default CEED_MEM_HOST
79             **cmode: copy mode for the array, default CEED_COPY_VALUES"""
80
81        # Setup the numpy array for the libCEED call
82        array_pointer = ffi.new("CeedScalar *")
83        if memtype == MEM_HOST:
84            array_pointer = ffi.cast(
85                "CeedScalar *",
86                array.__array_interface__['data'][0])
87        else:
88            array_pointer = ffi.cast(
89                "CeedScalar *",
90                array.__cuda_array_interface__['data'][0])
91
92        # libCEED call
93        lib.CeedVectorSetArray(self._pointer[0], memtype, cmode, array_pointer)
94
95    # Get Vector's data array
96    def get_array(self, memtype=MEM_HOST):
97        """Get read/write access to a Vector via the specified memory type.
98
99           Args:
100             **memtype: memory type of the array being passed, default CEED_MEM_HOST
101
102           Returns:
103             *array: Numpy or Numba array"""
104
105        # Retrieve the length of the array
106        length_pointer = ffi.new("CeedInt *")
107        lib.CeedVectorGetLength(self._pointer[0], length_pointer)
108
109        # Setup the pointer's pointer
110        array_pointer = ffi.new("CeedScalar **")
111
112        # libCEED call
113        lib.CeedVectorGetArray(self._pointer[0], memtype, array_pointer)
114
115        # Return array created from buffer
116        if memtype == MEM_HOST:
117            # Create buffer object from returned pointer
118            buff = ffi.buffer(
119                array_pointer[0],
120                ffi.sizeof("CeedScalar") *
121                length_pointer[0])
122            # return Numpy array
123            return np.frombuffer(buff, dtype="float64")
124        else:
125            # CUDA array interface
126            # https://numba.pydata.org/numba-doc/latest/cuda/cuda_array_interface.html
127            desc = {
128                'shape': (length_pointer[0]),
129                'typestr': '>f8',
130                'data': (int(ffi.cast("intptr_t", array_pointer[0])), False),
131                'version': 2
132            }
133            # return Numba array
134            return nbcuda.from_cuda_array_interface(desc)
135
136    # Get Vector's data array in read-only mode
137    def get_array_read(self, memtype=MEM_HOST):
138        """Get read-only access to a Vector via the specified memory type.
139
140           Args:
141             **memtype: memory type of the array being passed, default CEED_MEM_HOST
142
143           Returns:
144             *array: Numpy or Numba array"""
145
146        # Retrieve the length of the array
147        length_pointer = ffi.new("CeedInt *")
148        lib.CeedVectorGetLength(self._pointer[0], length_pointer)
149
150        # Setup the pointer's pointer
151        array_pointer = ffi.new("CeedScalar **")
152
153        # libCEED call
154        lib.CeedVectorGetArrayRead(self._pointer[0], memtype, array_pointer)
155
156        # Return array created from buffer
157        if memtype == MEM_HOST:
158            # Create buffer object from returned pointer
159            buff = ffi.buffer(
160                array_pointer[0],
161                ffi.sizeof("CeedScalar") *
162                length_pointer[0])
163            # return read only Numpy array
164            ret = np.frombuffer(buff, dtype="float64")
165            ret.flags['WRITEABLE'] = False
166            return ret
167        else:
168            # CUDA array interface
169            # https://numba.pydata.org/numba-doc/latest/cuda/cuda_array_interface.html
170            desc = {
171                'shape': (length_pointer[0]),
172                'typestr': '>f8',
173                'data': (int(ffi.cast("intptr_t", array_pointer[0])), False),
174                'version': 2
175            }
176            # return read only Numba array
177            return nbcuda.from_cuda_array_interface(desc)
178
179    # Restore the Vector's data array
180    def restore_array(self):
181        """Restore an array obtained using get_array()."""
182
183        # Setup the pointer's pointer
184        array_pointer = ffi.new("CeedScalar **")
185
186        # libCEED call
187        lib.CeedVectorRestoreArray(self._pointer[0], array_pointer)
188
189    # Restore an array obtained using getArrayRead
190    def restore_array_read(self):
191        """Restore an array obtained using get_array_read()."""
192
193        # Setup the pointer's pointer
194        array_pointer = ffi.new("CeedScalar **")
195
196        # libCEED call
197        lib.CeedVectorRestoreArrayRead(self._pointer[0], array_pointer)
198
199    @contextlib.contextmanager
200    def array(self, *shape, memtype=MEM_HOST):
201        """Context manager for array access.
202
203        Args:
204          shape (tuple): shape of returned numpy.array
205          **memtype: memory type of the array being passed, default CEED_MEM_HOST
206
207
208        Returns:
209          np.array: writable view of vector
210
211        Examples:
212          Constructing the identity inside a libceed.Vector:
213
214          >>> vec = ceed.Vector(16)
215          >>> with vec.array(4, 4) as x:
216          >>>     x[...] = np.eye(4)
217        """
218        x = self.get_array(memtype=memtype)
219        if shape:
220            x = x.reshape(shape)
221        yield x
222        self.restore_array()
223
224    @contextlib.contextmanager
225    def array_read(self, *shape, memtype=MEM_HOST):
226        """Context manager for read-only array access.
227
228        Args:
229          shape (tuple): shape of returned numpy.array
230          **memtype: memory type of the array being passed, default CEED_MEM_HOST
231
232        Returns:
233          np.array: read-only view of vector
234
235        Examples:
236          Constructing the identity inside a libceed.Vector:
237
238          >>> vec = ceed.Vector(6)
239          >>> vec.set_value(1.3)
240          >>> with vec.array_read(2, 3) as x:
241          >>>     print(x)
242        """
243        x = self.get_array_read(memtype=memtype)
244        if shape:
245            x = x.reshape(shape)
246        yield x
247        self.restore_array_read()
248
249    # Get the length of a Vector
250    def get_length(self):
251        """Get the length of a Vector.
252
253           Returns:
254             length: length of the Vector"""
255
256        length_pointer = ffi.new("CeedInt *")
257
258        # libCEED call
259        lib.CeedVectorGetLength(self._pointer[0], length_pointer)
260
261        return length_pointer[0]
262
263    # Get the length of a Vector
264    def __len__(self):
265        """Get the length of a Vector.
266
267           Returns:
268             length: length of the Vector"""
269
270        length_pointer = ffi.new("CeedInt *")
271
272        # libCEED call
273        lib.CeedVectorGetLength(self._pointer[0], length_pointer)
274
275        return length_pointer[0]
276
277    # Set the Vector to a given constant value
278    def set_value(self, value):
279        """Set the Vector to a constant value.
280
281           Args:
282             value: value to be used"""
283
284        # libCEED call
285        lib.CeedVectorSetValue(self._pointer[0], value)
286
287    # Sync the Vector to a specified memtype
288    def sync_array(self, memtype=MEM_HOST):
289        """Sync the Vector to a specified memtype.
290
291           Args:
292             **memtype: memtype to be synced"""
293
294        # libCEED call
295        lib.CeedVectorSyncArray(self._pointer[0], memtype)
296
297    # Compute the norm of a vector
298    def norm(self, normtype=NORM_2):
299        """Get the norm of a Vector.
300
301           Args:
302             **normtype: type of norm to be computed"""
303
304        norm_pointer = ffi.new("CeedScalar *")
305
306        # libCEED call
307        lib.CeedVectorNorm(self._pointer[0], normtype, norm_pointer)
308
309        return norm_pointer[0]
310
311# ------------------------------------------------------------------------------
312
313
314class _VectorWrap(Vector):
315    """Wrap a CeedVector pointer in a Vector object."""
316
317    # Constructor
318    def __init__(self, ceed, pointer):
319        # CeedVector object
320        self._pointer = pointer
321
322        # Reference to Ceed
323        self._ceed = ceed
324
325# ------------------------------------------------------------------------------
326