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