xref: /libCEED/python/ceed_qfunctioncontext.py (revision fe3941314958fee4f57f09c6529ef0f16a06d9c6)
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
22
23# ------------------------------------------------------------------------------
24
25
26class QFunctionContext():
27    """Ceed QFunction Context: stores Ceed QFunction user context data."""
28
29    # Attributes
30    _ceed = None
31    _pointer = ffi.NULL
32    _array_reference = None
33
34    # Constructor
35    def __init__(self, ceed):
36        # CeedQFunctionContext object
37        self._pointer = ffi.new("CeedQFunctionContext *")
38
39        # Reference to Ceed
40        self._ceed = ceed
41
42        # libCEED call
43        err_code = lib.CeedQFunctionContextCreate(
44            self._ceed._pointer[0], self._pointer)
45        self._ceed._check_error(err_code)
46
47    # Destructor
48    def __del__(self):
49        # libCEED call
50        err_code = lib.CeedQFunctionContextDestroy(self._pointer)
51        self._ceed._check_error(err_code)
52
53    # Representation
54    def __repr__(self):
55        return "<CeedQFunctionContext instance at " + hex(id(self)) + ">"
56
57    # String conversion for print() to stdout
58    def __str__(self):
59        """View a QFunction Context 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.CeedQFunctionContextView(
68                    self._pointer[0], stream)
69                self._ceed._check_error(err_code)
70
71                stream_file.seek(0)
72                out_string = stream_file.read()
73
74        return out_string
75
76    # Set QFunction Context's data
77    def set_data(self, data, memtype=MEM_HOST, cmode=COPY_VALUES):
78        """Set the data used by a QFunction Context, freeing any previously allocated
79           data if applicable.
80
81           Args:
82             *data: Numpy or Numba array to be used
83             **memtype: memory type of the array being passed, default CEED_MEM_HOST
84             **cmode: copy mode for the array, default CEED_COPY_VALUES"""
85
86        # Store array reference if needed
87        if cmode == USE_POINTER:
88            self._array_reference = data
89        else:
90            self._array_reference = None
91
92        # Setup the numpy array for the libCEED call
93        data_pointer = ffi.new("CeedScalar *")
94        if memtype == MEM_HOST:
95            data_pointer = ffi.cast(
96                "void *",
97                data.__array_interface__['data'][0])
98        else:
99            array_pointer = ffi.cast(
100                "void *",
101                data.__cuda_array_interface__['data'][0])
102
103        # libCEED call
104        err_code = lib.CeedQFunctionContextSetData(
105            self._pointer[0],
106            memtype,
107            cmode,
108            len(data) * ffi.sizeof("CeedScalar"),
109            data_pointer)
110        self._ceed._check_error(err_code)
111
112    # Get QFunction Context's data
113    def get_data(self, memtype=MEM_HOST):
114        """Get read/write access to a QFunction Context via the specified memory type.
115
116           Args:
117             **memtype: memory type of the array being passed, default CEED_MEM_HOST
118
119           Returns:
120             *data: Numpy or Numba array"""
121
122        # Retrieve the length of the array
123        size_pointer = ffi.new("CeedInt *")
124        err_code = lib.CeedQFunctionContextGetContextSize(
125            self._pointer[0], size_pointer)
126        self._ceed._check_error(err_code)
127
128        # Setup the pointer's pointer
129        data_pointer = ffi.new("CeedScalar **")
130
131        # libCEED call
132        err_code = lib.CeedQFunctionContextGetData(
133            self._pointer[0], memtype, data_pointer)
134        self._ceed._check_error(err_code)
135
136        # Return array created from buffer
137        if memtype == MEM_HOST:
138            # Create buffer object from returned pointer
139            buff = ffi.buffer(
140                data_pointer[0],
141                ffi.sizeof("CeedScalar") *
142                length_pointer[0])
143            # return Numpy array
144            return np.frombuffer(buff, dtype="float64")
145        else:
146            # CUDA array interface
147            # https://numba.pydata.org/numba-doc/latest/cuda/cuda_array_interface.html
148            import numba.cuda as nbcuda
149            desc = {
150                'shape': (length_pointer[0]),
151                'typestr': '>f8',
152                'data': (int(ffi.cast("intptr_t", data_pointer[0])), False),
153                'version': 2
154            }
155            # return Numba array
156            return nbcuda.from_cuda_array_interface(desc)
157
158    # Restore the QFunction Context's data
159    def restore_data(self):
160        """Restore an array obtained using get_data()."""
161
162        # Setup the pointer's pointer
163        data_pointer = ffi.new("CeedScalar **")
164
165        # libCEED call
166        err_code = lib.CeedQFunctionDataRestoreData(
167            self._pointer[0], data_pointer)
168        self._ceed._check_error(err_code)
169
170    @contextlib.contextmanager
171    def data(self, *shape, memtype=MEM_HOST):
172        """Context manager for array access.
173
174        Args:
175          shape (tuple): shape of returned numpy.array
176          **memtype: memory type of the data being passed, default CEED_MEM_HOST
177
178
179        Returns:
180          np.array: writable view of QFunction Context
181        """
182        x = self.get_data(memtype=memtype)
183        if shape:
184            x = x.reshape(shape)
185        yield x
186        self.restore_data()
187
188# ------------------------------------------------------------------------------
189