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