xref: /libCEED/python/ceed_vector.py (revision e0dd3b27d19566a0742e776b14ca8ce6d5794b11)
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
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="float64")
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="float64")
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    # Restore the Vector's data array
195    def restore_array(self):
196        """Restore an array obtained using get_array()."""
197
198        # Setup the pointer's pointer
199        array_pointer = ffi.new("CeedScalar **")
200
201        # libCEED call
202        err_code = lib.CeedVectorRestoreArray(self._pointer[0], array_pointer)
203        self._ceed._check_error(err_code)
204
205    # Restore an array obtained using getArrayRead
206    def restore_array_read(self):
207        """Restore an array obtained using get_array_read()."""
208
209        # Setup the pointer's pointer
210        array_pointer = ffi.new("CeedScalar **")
211
212        # libCEED call
213        err_code = lib.CeedVectorRestoreArrayRead(
214            self._pointer[0], array_pointer)
215        self._ceed._check_error(err_code)
216
217    @contextlib.contextmanager
218    def array(self, *shape, memtype=MEM_HOST):
219        """Context manager for array access.
220
221        Args:
222          shape (tuple): shape of returned numpy.array
223          **memtype: memory type of the array being passed, default CEED_MEM_HOST
224
225
226        Returns:
227          np.array: writable view of vector
228
229        Examples:
230          Constructing the identity inside a libceed.Vector:
231
232          >>> vec = ceed.Vector(16)
233          >>> with vec.array(4, 4) as x:
234          >>>     x[...] = np.eye(4)
235        """
236        x = self.get_array(memtype=memtype)
237        if shape:
238            x = x.reshape(shape)
239        yield x
240        self.restore_array()
241
242    @contextlib.contextmanager
243    def array_read(self, *shape, memtype=MEM_HOST):
244        """Context manager for read-only array access.
245
246        Args:
247          shape (tuple): shape of returned numpy.array
248          **memtype: memory type of the array being passed, default CEED_MEM_HOST
249
250        Returns:
251          np.array: read-only view of vector
252
253        Examples:
254          Viewing contents of a reshaped libceed.Vector view:
255
256          >>> vec = ceed.Vector(6)
257          >>> vec.set_value(1.3)
258          >>> with vec.array_read(2, 3) as x:
259          >>>     print(x)
260        """
261        x = self.get_array_read(memtype=memtype)
262        if shape:
263            x = x.reshape(shape)
264        yield x
265        self.restore_array_read()
266
267    # Get the length of a Vector
268    def get_length(self):
269        """Get the length of a Vector.
270
271           Returns:
272             length: length of the Vector"""
273
274        length_pointer = ffi.new("CeedInt *")
275
276        # libCEED call
277        err_code = lib.CeedVectorGetLength(self._pointer[0], length_pointer)
278        self._ceed._check_error(err_code)
279
280        return length_pointer[0]
281
282    # Get the length of a Vector
283    def __len__(self):
284        """Get the length of a Vector.
285
286           Returns:
287             length: length of the Vector"""
288
289        length_pointer = ffi.new("CeedInt *")
290
291        # libCEED call
292        err_code = lib.CeedVectorGetLength(self._pointer[0], length_pointer)
293        self._ceed._check_error(err_code)
294
295        return length_pointer[0]
296
297    # Set the Vector to a given constant value
298    def set_value(self, value):
299        """Set the Vector to a constant value.
300
301           Args:
302             value: value to be used"""
303
304        # libCEED call
305        err_code = lib.CeedVectorSetValue(self._pointer[0], value)
306        self._ceed._check_error(err_code)
307
308    # Sync the Vector to a specified memtype
309    def sync_array(self, memtype=MEM_HOST):
310        """Sync the Vector to a specified memtype.
311
312           Args:
313             **memtype: memtype to be synced"""
314
315        # libCEED call
316        err_code = lib.CeedVectorSyncArray(self._pointer[0], memtype)
317        self._ceed._check_error(err_code)
318
319    # Compute the norm of a vector
320    def norm(self, normtype=NORM_2):
321        """Get the norm of a Vector.
322
323           Args:
324             **normtype: type of norm to be computed"""
325
326        norm_pointer = ffi.new("CeedScalar *")
327
328        # libCEED call
329        err_code = lib.CeedVectorNorm(self._pointer[0], normtype, norm_pointer)
330        self._ceed._check_error(err_code)
331
332        return norm_pointer[0]
333
334    # Take the reciprocal of a vector
335    def reciprocal(self):
336        """Take the reciprocal of a Vector."""
337
338        # libCEED call
339        err_code = lib.CeedVectorReciprocal(self._pointer[0])
340        self._ceed._check_error(err_code)
341
342        return self
343
344    # Compute self = alpha self
345    def scale(self, alpha):
346        """Compute self = alpha self."""
347
348        # libCEED call
349        err_code = lib.CeedVectorScale(self._pointer[0], alpha)
350        self._ceed._check_error(err_code)
351
352        return self
353
354    # Compute self = alpha x + self
355    def axpy(self, alpha, x):
356        """Compute self = alpha x + self."""
357
358        # libCEED call
359        err_code = lib.CeedVectorAXPY(self._pointer[0], alpha, x._pointer[0])
360        self._ceed._check_error(err_code)
361
362        return self
363
364    # Compute the pointwise multiplication self = x .* y
365    def pointwise_mult(self, x, y):
366        """Compute the pointwise multiplication self = x .* y."""
367
368        # libCEED call
369        err_code = lib.CeedVectorPointwiseMult(
370            self._pointer[0], x._pointer[0], y._pointer[0]
371        )
372        self._ceed._check_error(err_code)
373
374        return self
375
376    def _state(self):
377        """Return the modification state of the Vector.
378
379        State is incremented each time the Vector is mutated, and is odd whenever a
380        mutable reference has not been returned.
381        """
382
383        state_pointer = ffi.new("uint64_t *")
384        err_code = lib.CeedVectorGetState(self._pointer[0], state_pointer)
385        self._ceed._check_error(err_code)
386        return state_pointer[0]
387
388# ------------------------------------------------------------------------------
389
390
391class _VectorWrap(Vector):
392    """Wrap a CeedVector pointer in a Vector object."""
393
394    # Constructor
395    def __init__(self, ceed, pointer):
396        # CeedVector object
397        self._pointer = pointer
398
399        # Reference to Ceed
400        self._ceed = ceed
401
402# ------------------------------------------------------------------------------
403