xref: /libCEED/python/ceed_vector.py (revision b19271b6eb0791edc605fd8a4a305de5b22bda53)
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 contextlib
21from .ceed_constants import MEM_HOST, USE_POINTER, COPY_VALUES, NORM_2, scalar_types
22
23# ------------------------------------------------------------------------------
24
25
26class Vector():
27    """Ceed Vector: storing and manipulating vectors."""
28
29    # Constructor
30    def __init__(self, ceed, size):
31        # CeedVector object
32        self._pointer = ffi.new("CeedVector *")
33
34        # Reference to Ceed
35        self._ceed = ceed
36
37        # libCEED call
38        err_code = lib.CeedVectorCreate(
39            self._ceed._pointer[0], size, self._pointer)
40        self._ceed._check_error(err_code)
41
42    # Destructor
43    def __del__(self):
44        # libCEED call
45        err_code = lib.CeedVectorDestroy(self._pointer)
46        self._ceed._check_error(err_code)
47
48    # Representation
49    def __repr__(self):
50        return "<CeedVector instance at " + hex(id(self)) + ">"
51
52    # String conversion for print() to stdout
53    def __str__(self):
54        """View a Vector via print()."""
55
56        # libCEED call
57        fmt = ffi.new("char[]", "%f".encode('ascii'))
58        with tempfile.NamedTemporaryFile() as key_file:
59            with open(key_file.name, 'r+') as stream_file:
60                stream = ffi.cast("FILE *", stream_file)
61
62                err_code = lib.CeedVectorView(self._pointer[0], fmt, stream)
63                self._ceed._check_error(err_code)
64
65                stream_file.seek(0)
66                out_string = stream_file.read()
67
68        return out_string
69
70    # Set Vector's data array
71    def set_array(self, array, memtype=MEM_HOST, cmode=COPY_VALUES):
72        """Set the array used by a Vector, freeing any previously allocated
73           array if applicable.
74
75           Args:
76             *array: Numpy or Numba array to be used
77             **memtype: memory type of the array being passed, default CEED_MEM_HOST
78             **cmode: copy mode for the array, default CEED_COPY_VALUES"""
79
80        # Store array reference if needed
81        if cmode == USE_POINTER:
82            self._array_reference = array
83        else:
84            self._array_reference = None
85
86        # Setup the numpy array for the libCEED call
87        array_pointer = ffi.new("CeedScalar *")
88        if memtype == MEM_HOST:
89            array_pointer = ffi.cast(
90                "CeedScalar *",
91                array.__array_interface__['data'][0])
92        else:
93            array_pointer = ffi.cast(
94                "CeedScalar *",
95                array.__cuda_array_interface__['data'][0])
96
97        # libCEED call
98        err_code = lib.CeedVectorSetArray(
99            self._pointer[0], memtype, cmode, array_pointer)
100        self._ceed._check_error(err_code)
101
102    # Get Vector's data array
103    def get_array(self, memtype=MEM_HOST):
104        """Get read/write access to a Vector via the specified memory type.
105
106           Args:
107             **memtype: memory type of the array being passed, default CEED_MEM_HOST
108
109           Returns:
110             *array: Numpy or Numba array"""
111
112        # Retrieve the length of the array
113        length_pointer = ffi.new("CeedInt *")
114        err_code = lib.CeedVectorGetLength(self._pointer[0], length_pointer)
115        self._ceed._check_error(err_code)
116
117        # Setup the pointer's pointer
118        array_pointer = ffi.new("CeedScalar **")
119
120        # libCEED call
121        err_code = lib.CeedVectorGetArray(
122            self._pointer[0], memtype, array_pointer)
123        self._ceed._check_error(err_code)
124
125        # Return array created from buffer
126        if memtype == MEM_HOST:
127            # Create buffer object from returned pointer
128            buff = ffi.buffer(
129                array_pointer[0],
130                ffi.sizeof("CeedScalar") *
131                length_pointer[0])
132            # return Numpy array
133            return np.frombuffer(buff, dtype=scalar_types[lib.CEED_SCALAR_TYPE])
134        else:
135            # CUDA array interface
136            # https://numba.pydata.org/numba-doc/latest/cuda/cuda_array_interface.html
137            import numba.cuda as nbcuda
138            desc = {
139                'shape': (length_pointer[0]),
140                'typestr': '>f8',
141                'data': (int(ffi.cast("intptr_t", array_pointer[0])), False),
142                'version': 2
143            }
144            # return Numba array
145            return nbcuda.from_cuda_array_interface(desc)
146
147    # Get Vector's data array in read-only mode
148    def get_array_read(self, memtype=MEM_HOST):
149        """Get read-only access to a Vector via the specified memory type.
150
151           Args:
152             **memtype: memory type of the array being passed, default CEED_MEM_HOST
153
154           Returns:
155             *array: Numpy or Numba array"""
156
157        # Retrieve the length of the array
158        length_pointer = ffi.new("CeedInt *")
159        err_code = lib.CeedVectorGetLength(self._pointer[0], length_pointer)
160        self._ceed._check_error(err_code)
161
162        # Setup the pointer's pointer
163        array_pointer = ffi.new("CeedScalar **")
164
165        # libCEED call
166        err_code = lib.CeedVectorGetArrayRead(
167            self._pointer[0], memtype, array_pointer)
168        self._ceed._check_error(err_code)
169
170        # Return array created from buffer
171        if memtype == MEM_HOST:
172            # Create buffer object from returned pointer
173            buff = ffi.buffer(
174                array_pointer[0],
175                ffi.sizeof("CeedScalar") *
176                length_pointer[0])
177            # return read only Numpy array
178            ret = np.frombuffer(buff, dtype=scalar_types[lib.CEED_SCALAR_TYPE])
179            ret.flags['WRITEABLE'] = False
180            return ret
181        else:
182            # CUDA array interface
183            # https://numba.pydata.org/numba-doc/latest/cuda/cuda_array_interface.html
184            import numba.cuda as nbcuda
185            desc = {
186                'shape': (length_pointer[0]),
187                'typestr': '>f8',
188                'data': (int(ffi.cast("intptr_t", array_pointer[0])), False),
189                'version': 2
190            }
191            # return read only Numba array
192            return nbcuda.from_cuda_array_interface(desc)
193
194    # Get Vector's data array in write-only mode
195    def get_array_write(self, memtype=MEM_HOST):
196        """Get write-only access to a Vector via the specified memory type.
197           All old values should be considered invalid.
198
199           Args:
200             **memtype: memory type of the array being passed, default CEED_MEM_HOST
201
202           Returns:
203             *array: Numpy or Numba array"""
204
205        # Retrieve the length of the array
206        length_pointer = ffi.new("CeedInt *")
207        err_code = lib.CeedVectorGetLength(self._pointer[0], length_pointer)
208        self._ceed._check_error(err_code)
209
210        # Setup the pointer's pointer
211        array_pointer = ffi.new("CeedScalar **")
212
213        # libCEED call
214        err_code = lib.CeedVectorGetArrayWrite(
215            self._pointer[0], memtype, array_pointer)
216        self._ceed._check_error(err_code)
217
218        # Return array created from buffer
219        if memtype == MEM_HOST:
220            # Create buffer object from returned pointer
221            buff = ffi.buffer(
222                array_pointer[0],
223                ffi.sizeof("CeedScalar") *
224                length_pointer[0])
225            # return Numpy array
226            return np.frombuffer(buff, dtype=scalar_types[lib.CEED_SCALAR_TYPE])
227        else:
228            # CUDA array interface
229            # https://numba.pydata.org/numba-doc/latest/cuda/cuda_array_interface.html
230            import numba.cuda as nbcuda
231            desc = {
232                'shape': (length_pointer[0]),
233                'typestr': '>f8',
234                'data': (int(ffi.cast("intptr_t", array_pointer[0])), False),
235                'version': 2
236            }
237            # return Numba array
238            return nbcuda.from_cuda_array_interface(desc)
239
240    # Restore the Vector's data array
241    def restore_array(self):
242        """Restore an array obtained using get_array()."""
243
244        # Setup the pointer's pointer
245        array_pointer = ffi.new("CeedScalar **")
246
247        # libCEED call
248        err_code = lib.CeedVectorRestoreArray(self._pointer[0], array_pointer)
249        self._ceed._check_error(err_code)
250
251    # Restore an array obtained using getArrayRead
252    def restore_array_read(self):
253        """Restore an array obtained using get_array_read()."""
254
255        # Setup the pointer's pointer
256        array_pointer = ffi.new("CeedScalar **")
257
258        # libCEED call
259        err_code = lib.CeedVectorRestoreArrayRead(
260            self._pointer[0], array_pointer)
261        self._ceed._check_error(err_code)
262
263    @contextlib.contextmanager
264    def array(self, *shape, memtype=MEM_HOST):
265        """Context manager for array access.
266
267        Args:
268          shape (tuple): shape of returned numpy.array
269          **memtype: memory type of the array being passed, default CEED_MEM_HOST
270
271
272        Returns:
273          np.array: writable view of vector
274
275        Examples:
276          Constructing the identity inside a libceed.Vector:
277
278          >>> vec = ceed.Vector(16)
279          >>> with vec.array(4, 4) as x:
280          >>>     x[...] = np.eye(4)
281        """
282        x = self.get_array(memtype=memtype)
283        if shape:
284            x = x.reshape(shape)
285        yield x
286        self.restore_array()
287
288    @contextlib.contextmanager
289    def array_read(self, *shape, memtype=MEM_HOST):
290        """Context manager for read-only array access.
291
292        Args:
293          shape (tuple): shape of returned numpy.array
294          **memtype: memory type of the array being passed, default CEED_MEM_HOST
295
296        Returns:
297          np.array: read-only view of vector
298
299        Examples:
300          Viewing contents of a reshaped libceed.Vector view:
301
302          >>> vec = ceed.Vector(6)
303          >>> vec.set_value(1.3)
304          >>> with vec.array_read(2, 3) as x:
305          >>>     print(x)
306        """
307        x = self.get_array_read(memtype=memtype)
308        if shape:
309            x = x.reshape(shape)
310        yield x
311        self.restore_array_read()
312
313    @contextlib.contextmanager
314    def array_write(self, *shape, memtype=MEM_HOST):
315        """Context manager for write-only array access.
316           All old values should be considered invalid.
317
318        Args:
319          shape (tuple): shape of returned numpy.array
320          **memtype: memory type of the array being passed, default CEED_MEM_HOST
321
322        Returns:
323          np.array: write-only view of vector
324
325        Examples:
326          Viewing contents of a reshaped libceed.Vector view:
327
328          >>> vec = ceed.Vector(6)
329          >>> vec.set_value(1.3)
330          >>> with vec.array_read(2, 3) as x:
331          >>>     print(x)
332        """
333        x = self.get_array_write(memtype=memtype)
334        if shape:
335            x = x.reshape(shape)
336        yield x
337        self.restore_array()
338
339    # Get the length of a Vector
340    def get_length(self):
341        """Get the length of a Vector.
342
343           Returns:
344             length: length of the Vector"""
345
346        length_pointer = ffi.new("CeedInt *")
347
348        # libCEED call
349        err_code = lib.CeedVectorGetLength(self._pointer[0], length_pointer)
350        self._ceed._check_error(err_code)
351
352        return length_pointer[0]
353
354    # Get the length of a Vector
355    def __len__(self):
356        """Get the length of a Vector.
357
358           Returns:
359             length: length of the Vector"""
360
361        length_pointer = ffi.new("CeedInt *")
362
363        # libCEED call
364        err_code = lib.CeedVectorGetLength(self._pointer[0], length_pointer)
365        self._ceed._check_error(err_code)
366
367        return length_pointer[0]
368
369    # Set the Vector to a given constant value
370    def set_value(self, value):
371        """Set the Vector to a constant value.
372
373           Args:
374             value: value to be used"""
375
376        # libCEED call
377        err_code = lib.CeedVectorSetValue(self._pointer[0], value)
378        self._ceed._check_error(err_code)
379
380    # Sync the Vector to a specified memtype
381    def sync_array(self, memtype=MEM_HOST):
382        """Sync the Vector to a specified memtype.
383
384           Args:
385             **memtype: memtype to be synced"""
386
387        # libCEED call
388        err_code = lib.CeedVectorSyncArray(self._pointer[0], memtype)
389        self._ceed._check_error(err_code)
390
391    # Compute the norm of a vector
392    def norm(self, normtype=NORM_2):
393        """Get the norm of a Vector.
394
395           Args:
396             **normtype: type of norm to be computed"""
397
398        norm_pointer = ffi.new("CeedScalar *")
399
400        # libCEED call
401        err_code = lib.CeedVectorNorm(self._pointer[0], normtype, norm_pointer)
402        self._ceed._check_error(err_code)
403
404        return norm_pointer[0]
405
406    # Take the reciprocal of a vector
407    def reciprocal(self):
408        """Take the reciprocal of a Vector."""
409
410        # libCEED call
411        err_code = lib.CeedVectorReciprocal(self._pointer[0])
412        self._ceed._check_error(err_code)
413
414        return self
415
416    # Compute self = alpha self
417    def scale(self, alpha):
418        """Compute self = alpha self."""
419
420        # libCEED call
421        err_code = lib.CeedVectorScale(self._pointer[0], alpha)
422        self._ceed._check_error(err_code)
423
424        return self
425
426    # Compute self = alpha x + self
427    def axpy(self, alpha, x):
428        """Compute self = alpha x + self."""
429
430        # libCEED call
431        err_code = lib.CeedVectorAXPY(self._pointer[0], alpha, x._pointer[0])
432        self._ceed._check_error(err_code)
433
434        return self
435
436    # Compute the pointwise multiplication self = x .* y
437    def pointwise_mult(self, x, y):
438        """Compute the pointwise multiplication self = x .* y."""
439
440        # libCEED call
441        err_code = lib.CeedVectorPointwiseMult(
442            self._pointer[0], x._pointer[0], y._pointer[0]
443        )
444        self._ceed._check_error(err_code)
445
446        return self
447
448    def _state(self):
449        """Return the modification state of the Vector.
450
451        State is incremented each time the Vector is mutated, and is odd whenever a
452        mutable reference has not been returned.
453        """
454
455        state_pointer = ffi.new("uint64_t *")
456        err_code = lib.CeedVectorGetState(self._pointer[0], state_pointer)
457        self._ceed._check_error(err_code)
458        return state_pointer[0]
459
460# ------------------------------------------------------------------------------
461
462
463class _VectorWrap(Vector):
464    """Wrap a CeedVector pointer in a Vector object."""
465
466    # Constructor
467    def __init__(self, ceed, pointer):
468        # CeedVector object
469        self._pointer = pointer
470
471        # Reference to Ceed
472        self._ceed = ceed
473
474# ------------------------------------------------------------------------------
475