1# -------------------------------------------------------------------- 2 3cdef extern from * nogil: 4 5 ctypedef const char* PetscVecType "VecType" 6 PetscVecType VECSEQ 7 PetscVecType VECMPI 8 PetscVecType VECSTANDARD 9 PetscVecType VECSHARED 10 PetscVecType VECSEQVIENNACL 11 PetscVecType VECMPIVIENNACL 12 PetscVecType VECVIENNACL 13 PetscVecType VECSEQCUDA 14 PetscVecType VECMPICUDA 15 PetscVecType VECCUDA 16 PetscVecType VECSEQHIP 17 PetscVecType VECMPIHIP 18 PetscVecType VECHIP 19 PetscVecType VECNEST 20 PetscVecType VECSEQKOKKOS 21 PetscVecType VECMPIKOKKOS 22 PetscVecType VECKOKKOS 23 24 ctypedef enum PetscVecOption "VecOption": 25 VEC_IGNORE_OFF_PROC_ENTRIES 26 VEC_IGNORE_NEGATIVE_INDICES 27 28 PetscErrorCode VecView(PetscVec, PetscViewer) 29 PetscErrorCode VecDestroy(PetscVec*) 30 PetscErrorCode VecCreate(MPI_Comm, PetscVec*) 31 32 PetscErrorCode VecSetOptionsPrefix(PetscVec, char[]) 33 PetscErrorCode VecAppendOptionsPrefix(PetscVec, char[]) 34 PetscErrorCode VecGetOptionsPrefix(PetscVec, char*[]) 35 PetscErrorCode VecSetFromOptions(PetscVec) 36 PetscErrorCode VecSetUp(PetscVec) 37 38 PetscErrorCode VecCreateSeq(MPI_Comm, PetscInt, PetscVec*) 39 PetscErrorCode VecCreateSeqWithArray(MPI_Comm, PetscInt, PetscInt, PetscScalar[], PetscVec*) 40 PetscErrorCode VecCreateSeqCUDAWithArrays(MPI_Comm, PetscInt, PetscInt, PetscScalar[], PetscScalar[], PetscVec*) 41 PetscErrorCode VecCreateSeqHIPWithArrays(MPI_Comm, PetscInt, PetscInt, PetscScalar[], PetscScalar[], PetscVec*) 42 PetscErrorCode VecCreateSeqViennaCLWithArrays(MPI_Comm, PetscInt, PetscInt, PetscScalar[], PetscScalar[], PetscVec*) 43 PetscErrorCode VecCreateMPI(MPI_Comm, PetscInt, PetscInt, PetscVec*) 44 PetscErrorCode VecCreateMPIWithArray(MPI_Comm, PetscInt, PetscInt, PetscInt, PetscScalar[], PetscVec*) 45 PetscErrorCode VecCreateMPICUDAWithArrays(MPI_Comm, PetscInt, PetscInt, PetscInt, PetscScalar[], PetscScalar[], PetscVec*) 46 PetscErrorCode VecCreateMPIHIPWithArrays(MPI_Comm, PetscInt, PetscInt, PetscInt, PetscScalar[], PetscScalar[], PetscVec*) 47 PetscErrorCode VecCreateMPIViennaCLWithArrays(MPI_Comm, PetscInt, PetscInt, PetscInt, PetscScalar[], PetscScalar[], PetscVec*) 48 PetscErrorCode VecCreateGhost(MPI_Comm, PetscInt, PetscInt, PetscInt, PetscInt[], PetscVec*) 49 PetscErrorCode VecCreateGhostWithArray(MPI_Comm, PetscInt, PetscInt, PetscInt, PetscInt[], PetscScalar[], PetscVec*) 50 PetscErrorCode VecCreateGhostBlock(MPI_Comm, PetscInt, PetscInt, PetscInt, PetscInt, PetscInt[], PetscVec*) 51 PetscErrorCode VecCreateGhostBlockWithArray(MPI_Comm, PetscInt, PetscInt, PetscInt, PetscInt, PetscInt[], PetscScalar[], PetscVec*) 52 PetscErrorCode VecCreateShared(MPI_Comm, PetscInt, PetscInt, PetscVec*) 53 PetscErrorCode VecCreateNest(MPI_Comm, PetscInt, PetscIS[], PetscVec[], PetscVec*) 54 PetscErrorCode VecGetType(PetscVec, PetscVecType*) 55 PetscErrorCode VecSetType(PetscVec, PetscVecType) 56 PetscErrorCode VecSetOption(PetscVec, PetscVecOption, PetscBool) 57 PetscErrorCode VecSetSizes(PetscVec, PetscInt, PetscInt) 58 PetscErrorCode VecGetSize(PetscVec, PetscInt*) 59 PetscErrorCode VecGetLocalSize(PetscVec, PetscInt*) 60 PetscErrorCode VecSetBlockSize(PetscVec, PetscInt) 61 PetscErrorCode VecGetBlockSize(PetscVec, PetscInt*) 62 PetscErrorCode VecGetOwnershipRange(PetscVec, PetscInt*, PetscInt*) 63 PetscErrorCode VecGetOwnershipRanges(PetscVec, const PetscInt*[]) 64 65 PetscErrorCode VecCreateLocalVector(PetscVec, PetscVec*) 66 PetscErrorCode VecGetLocalVector(PetscVec, PetscVec) 67 PetscErrorCode VecRestoreLocalVector(PetscVec, PetscVec) 68 PetscErrorCode VecGetLocalVectorRead(PetscVec, PetscVec) 69 PetscErrorCode VecRestoreLocalVectorRead(PetscVec, PetscVec) 70 71 PetscErrorCode VecGetArrayWrite(PetscVec, PetscScalar*[]) 72 PetscErrorCode VecRestoreArrayWrite(PetscVec, PetscScalar*[]) 73 PetscErrorCode VecGetArrayRead(PetscVec, const PetscScalar*[]) 74 PetscErrorCode VecRestoreArrayRead(PetscVec, const PetscScalar*[]) 75 PetscErrorCode VecGetArray(PetscVec, PetscScalar*[]) 76 PetscErrorCode VecRestoreArray(PetscVec, PetscScalar*[]) 77 PetscErrorCode VecPlaceArray(PetscVec, PetscScalar[]) 78 PetscErrorCode VecResetArray(PetscVec) 79 PetscErrorCode VecGetArrayWriteAndMemType(PetscVec, PetscScalar*[], PetscMemType*) 80 PetscErrorCode VecRestoreArrayWriteAndMemType(PetscVec, PetscScalar*[]) 81 PetscErrorCode VecGetArrayReadAndMemType(PetscVec, const PetscScalar*[], PetscMemType*) 82 PetscErrorCode VecRestoreArrayReadAndMemType(PetscVec, const PetscScalar*[]) 83 PetscErrorCode VecGetArrayAndMemType(PetscVec, PetscScalar*[], PetscMemType*) 84 PetscErrorCode VecRestoreArrayAndMemType(PetscVec, PetscScalar*[]) 85 86 PetscErrorCode VecEqual(PetscVec, PetscVec, PetscBool*) 87 PetscErrorCode VecLoad(PetscVec, PetscViewer) 88 89 PetscErrorCode VecDuplicate(PetscVec, PetscVec*) 90 PetscErrorCode VecCopy(PetscVec, PetscVec) 91 PetscErrorCode VecFilter(PetscVec, PetscReal) 92 93 PetscErrorCode VecDuplicateVecs(PetscVec, PetscInt, PetscVec*[]) 94 PetscErrorCode VecDestroyVecs(PetscInt, PetscVec*[]) 95 96 PetscErrorCode VecGetValues(PetscVec, PetscInt, PetscInt[], PetscScalar[]) 97 98 PetscErrorCode VecSetValue(PetscVec, PetscInt, PetscScalar, PetscInsertMode) 99 PetscErrorCode VecSetValues(PetscVec, PetscInt, const PetscInt[], const PetscScalar[], PetscInsertMode) 100 PetscErrorCode VecSetValuesBlocked(PetscVec, PetscInt, const PetscInt[], const PetscScalar[], PetscInsertMode) 101 102 PetscErrorCode VecSetLocalToGlobalMapping(PetscVec, PetscLGMap) 103 PetscErrorCode VecGetLocalToGlobalMapping(PetscVec, PetscLGMap*) 104 PetscErrorCode VecSetValueLocal(PetscVec, PetscInt, PetscScalar, PetscInsertMode) 105 PetscErrorCode VecSetValuesLocal(PetscVec, PetscInt, const PetscInt[], const PetscScalar[], PetscInsertMode) 106 PetscErrorCode VecSetValuesBlockedLocal(PetscVec, PetscInt, const PetscInt[], const PetscScalar[], PetscInsertMode) 107 108 PetscErrorCode VecDot(PetscVec, PetscVec, PetscScalar*) 109 PetscErrorCode VecDotBegin(PetscVec, PetscVec, PetscScalar*) 110 PetscErrorCode VecDotEnd(PetscVec, PetscVec, PetscScalar*) 111 PetscErrorCode VecTDot(PetscVec, PetscVec, PetscScalar*) 112 PetscErrorCode VecTDotBegin(PetscVec, PetscVec, PetscScalar*) 113 PetscErrorCode VecTDotEnd(PetscVec, PetscVec, PetscScalar*) 114 PetscErrorCode VecMDot(PetscVec, PetscInt, PetscVec[], PetscScalar*) 115 PetscErrorCode VecMDotBegin(PetscVec, PetscInt, PetscVec[], PetscScalar*) 116 PetscErrorCode VecMDotEnd(PetscVec, PetscInt, PetscVec[], PetscScalar*) 117 PetscErrorCode VecMTDot(PetscVec, PetscInt, PetscVec[], PetscScalar*) 118 PetscErrorCode VecMTDotBegin(PetscVec, PetscInt, PetscVec[], PetscScalar*) 119 PetscErrorCode VecMTDotEnd(PetscVec, PetscInt, PetscVec[], PetscScalar*) 120 121 PetscErrorCode VecNorm(PetscVec, PetscNormType, PetscReal*) 122 PetscErrorCode VecNormBegin(PetscVec, PetscNormType, PetscReal*) 123 PetscErrorCode VecNormEnd(PetscVec, PetscNormType, PetscReal*) 124 PetscErrorCode VecDotNorm2(PetscVec, PetscVec, PetscScalar*, PetscReal*) 125 126 PetscErrorCode VecAssemblyBegin(PetscVec) 127 PetscErrorCode VecAssemblyEnd(PetscVec) 128 129 PetscErrorCode VecZeroEntries(PetscVec) 130 PetscErrorCode VecConjugate(PetscVec) 131 PetscErrorCode VecNormalize(PetscVec, PetscReal*) 132 PetscErrorCode VecSum(PetscVec, PetscScalar*) 133 PetscErrorCode VecMean(PetscVec, PetscScalar*) 134 PetscErrorCode VecMax(PetscVec, PetscInt*, PetscReal*) 135 PetscErrorCode VecMin(PetscVec, PetscInt*, PetscReal*) 136 PetscErrorCode VecScale(PetscVec, PetscScalar) 137 PetscErrorCode VecCopy(PetscVec, PetscVec) 138 PetscErrorCode VecSetRandom(PetscVec, PetscRandom) 139 PetscErrorCode VecSet(PetscVec, PetscScalar) 140 PetscErrorCode VecSwap(PetscVec, PetscVec) 141 PetscErrorCode VecAXPY(PetscVec, PetscScalar, PetscVec) 142 PetscErrorCode VecAXPBY(PetscVec, PetscScalar, PetscScalar, PetscVec) 143 PetscErrorCode VecAYPX(PetscVec, PetscScalar, PetscVec) 144 PetscErrorCode VecWAXPY(PetscVec, PetscScalar, PetscVec, PetscVec) 145 PetscErrorCode VecMAXPY(PetscVec, PetscInt, PetscScalar[], PetscVec[]) 146 PetscErrorCode VecPointwiseMax(PetscVec, PetscVec, PetscVec) 147 PetscErrorCode VecPointwiseMaxAbs(PetscVec, PetscVec, PetscVec) 148 PetscErrorCode VecPointwiseMin(PetscVec, PetscVec, PetscVec) 149 PetscErrorCode VecPointwiseMult(PetscVec, PetscVec, PetscVec) 150 PetscErrorCode VecPointwiseDivide(PetscVec, PetscVec, PetscVec) 151 PetscErrorCode VecMaxPointwiseDivide(PetscVec, PetscVec, PetscReal*) 152 PetscErrorCode VecShift(PetscVec, PetscScalar) 153 PetscErrorCode VecFilter(PetscVec, PetscReal) 154 PetscErrorCode VecReciprocal(PetscVec) 155 PetscErrorCode VecPermute(PetscVec, PetscIS, PetscBool) 156 PetscErrorCode VecExp(PetscVec) 157 PetscErrorCode VecLog(PetscVec) 158 PetscErrorCode VecSqrtAbs(PetscVec) 159 PetscErrorCode VecAbs(PetscVec) 160 161 PetscErrorCode VecStrideMin(PetscVec, PetscInt, PetscInt*, PetscReal*) 162 PetscErrorCode VecStrideMax(PetscVec, PetscInt, PetscInt*, PetscReal*) 163 PetscErrorCode VecStrideScale(PetscVec, PetscInt, PetscScalar) 164 PetscErrorCode VecStrideGather(PetscVec, PetscInt, PetscVec, PetscInsertMode) 165 PetscErrorCode VecStrideScatter(PetscVec, PetscInt, PetscVec, PetscInsertMode) 166 PetscErrorCode VecStrideNorm(PetscVec, PetscInt, PetscNormType, PetscReal*) 167 168 PetscErrorCode VecGhostGetLocalForm(PetscVec, PetscVec*) 169 PetscErrorCode VecGhostRestoreLocalForm(PetscVec, PetscVec*) 170 PetscErrorCode VecGhostUpdateBegin(PetscVec, PetscInsertMode, PetscScatterMode) 171 PetscErrorCode VecGhostUpdateEnd(PetscVec, PetscInsertMode, PetscScatterMode) 172 PetscErrorCode VecMPISetGhost(PetscVec, PetscInt, const PetscInt*) 173 PetscErrorCode VecGhostGetGhostIS(PetscVec, PetscIS*) 174 175 PetscErrorCode VecGetSubVector(PetscVec, PetscIS, PetscVec*) 176 PetscErrorCode VecRestoreSubVector(PetscVec, PetscIS, PetscVec*) 177 178 PetscErrorCode VecNestGetSubVecs(PetscVec, PetscInt*, PetscVec**) 179 PetscErrorCode VecNestSetSubVecs(PetscVec, PetscInt, PetscInt*, PetscVec*) 180 181 PetscErrorCode VecConcatenate(PetscInt nx, PetscVec[], PetscVec*, PetscIS*[]) 182 183 PetscErrorCode VecISAXPY(PetscVec, PetscIS, PetscScalar, PetscVec) 184 PetscErrorCode VecISSet(PetscVec, PetscIS, PetscScalar) 185 186 PetscErrorCode VecCUDAGetArrayRead(PetscVec, const PetscScalar*[]) 187 PetscErrorCode VecCUDAGetArrayWrite(PetscVec, PetscScalar*[]) 188 PetscErrorCode VecCUDAGetArray(PetscVec, PetscScalar*[]) 189 PetscErrorCode VecCUDARestoreArrayRead(PetscVec, const PetscScalar*[]) 190 PetscErrorCode VecCUDARestoreArrayWrite(PetscVec, PetscScalar*[]) 191 PetscErrorCode VecCUDARestoreArray(PetscVec, PetscScalar*[]) 192 193 PetscErrorCode VecHIPGetArrayRead(PetscVec, const PetscScalar*[]) 194 PetscErrorCode VecHIPGetArrayWrite(PetscVec, PetscScalar*[]) 195 PetscErrorCode VecHIPGetArray(PetscVec, PetscScalar*[]) 196 PetscErrorCode VecHIPRestoreArrayRead(PetscVec, const PetscScalar*[]) 197 PetscErrorCode VecHIPRestoreArrayWrite(PetscVec, PetscScalar*[]) 198 PetscErrorCode VecHIPRestoreArray(PetscVec, PetscScalar*[]) 199 200 PetscErrorCode VecBindToCPU(PetscVec, PetscBool) 201 PetscErrorCode VecBoundToCPU(PetscVec, PetscBool*) 202 PetscErrorCode VecGetOffloadMask(PetscVec, PetscOffloadMask*) 203 204 PetscErrorCode VecViennaCLGetCLContext(PetscVec, Py_uintptr_t*) 205 PetscErrorCode VecViennaCLGetCLQueue(PetscVec, Py_uintptr_t*) 206 PetscErrorCode VecViennaCLGetCLMemRead(PetscVec, Py_uintptr_t*) 207 PetscErrorCode VecViennaCLGetCLMemWrite(PetscVec, Py_uintptr_t*) 208 PetscErrorCode VecViennaCLRestoreCLMemWrite(PetscVec) 209 PetscErrorCode VecViennaCLGetCLMem(PetscVec, Py_uintptr_t*) 210 PetscErrorCode VecViennaCLRestoreCLMem(PetscVec) 211 212 PetscErrorCode VecCreateSeqCUDAWithArray(MPI_Comm, PetscInt, PetscInt, const PetscScalar*, PetscVec*) 213 PetscErrorCode VecCreateMPICUDAWithArray(MPI_Comm, PetscInt, PetscInt, PetscInt, const PetscScalar*, PetscVec*) 214 PetscErrorCode VecCreateSeqHIPWithArray(MPI_Comm, PetscInt, PetscInt, const PetscScalar*, PetscVec*) 215 PetscErrorCode VecCreateMPIHIPWithArray(MPI_Comm, PetscInt, PetscInt, PetscInt, const PetscScalar*, PetscVec*) 216 217cdef extern from * nogil: # custom.h 218 PetscErrorCode VecStrideSum(PetscVec, PetscInt, PetscScalar*) 219 PetscErrorCode VecGetCurrentMemType(PetscVec, PetscMemType*) 220 221# -------------------------------------------------------------------- 222 223cdef inline Vec ref_Vec(PetscVec vec): 224 cdef Vec ob = <Vec> Vec() 225 ob.vec = vec 226 CHKERR(PetscINCREF(ob.obj)) 227 return ob 228 229# -------------------------------------------------------------------- 230 231# unary operations 232 233cdef Vec vec_pos(Vec self): 234 cdef Vec vec = type(self)() 235 CHKERR(VecDuplicate(self.vec, &vec.vec)) 236 CHKERR(VecCopy(self.vec, vec.vec)) 237 return vec 238 239cdef Vec vec_neg(Vec self): 240 cdef Vec vec = <Vec> vec_pos(self) 241 CHKERR(VecScale(vec.vec, -1)) 242 return vec 243 244cdef Vec vec_abs(Vec self): 245 cdef Vec vec = <Vec> vec_pos(self) 246 CHKERR(VecAbs(vec.vec)) 247 return vec 248 249# inplace binary operations 250 251cdef Vec vec_iadd(Vec self, other): 252 cdef PetscScalar alpha = 1 253 cdef Vec vec 254 if isinstance(other, Vec): 255 alpha = 1; vec = other 256 CHKERR(VecAXPY(self.vec, alpha, vec.vec)) 257 elif isinstance(other, (tuple, list)): 258 other, vec = other 259 alpha = asScalar(other) 260 CHKERR(VecAXPY(self.vec, alpha, vec.vec)) 261 else: 262 alpha = asScalar(other) 263 CHKERR(VecShift(self.vec, alpha)) 264 return self 265 266cdef Vec vec_isub(Vec self, other): 267 cdef PetscScalar alpha = 1 268 cdef Vec vec 269 if isinstance(other, Vec): 270 alpha = 1; vec = other 271 CHKERR(VecAXPY(self.vec, -alpha, vec.vec)) 272 elif isinstance(other, (tuple, list)): 273 other, vec = other 274 alpha = asScalar(other) 275 CHKERR(VecAXPY(self.vec, -alpha, vec.vec)) 276 else: 277 alpha = asScalar(other) 278 CHKERR(VecShift(self.vec, -alpha)) 279 return self 280 281cdef Vec vec_imul(Vec self, other): 282 cdef PetscScalar alpha = 1 283 cdef Vec vec 284 if isinstance(other, Vec): 285 vec = other 286 CHKERR(VecPointwiseMult(self.vec, self.vec, vec.vec)) 287 else: 288 alpha = asScalar(other) 289 CHKERR(VecScale(self.vec, alpha)) 290 return self 291 292cdef Vec vec_idiv(Vec self, other): 293 cdef PetscScalar one = 1 294 cdef PetscScalar alpha = 1 295 cdef Vec vec 296 if isinstance(other, Vec): 297 vec = other 298 CHKERR(VecPointwiseDivide(self.vec, self.vec, vec.vec)) 299 else: 300 alpha = asScalar(other) 301 CHKERR(VecScale(self.vec, one/alpha)) 302 return self 303 304# binary operations 305 306cdef Vec vec_add(Vec self, other): 307 return vec_iadd(vec_pos(self), other) 308 309cdef Vec vec_sub(Vec self, other): 310 return vec_isub(vec_pos(self), other) 311 312cdef Vec vec_mul(Vec self, other): 313 return vec_imul(vec_pos(self), other) 314 315cdef Vec vec_div(Vec self, other): 316 return vec_idiv(vec_pos(self), other) 317 318cdef object vec_matmul(Vec self, other): 319 if isinstance(other, Vec): 320 return self.dot(other) 321 if isinstance(other, Mat): 322 result = other.createVecRight() 323 other.multTranspose(self, result) 324 return result 325 return NotImplemented 326 327# reflected binary operations 328 329cdef Vec vec_radd(Vec self, other): 330 return vec_add(self, other) 331 332cdef Vec vec_rsub(Vec self, other): 333 cdef Vec vec = <Vec> vec_sub(self, other) 334 CHKERR(VecScale(vec.vec, -1)) 335 return vec 336 337cdef Vec vec_rmul(Vec self, other): 338 return vec_mul(self, other) 339 340cdef Vec vec_rdiv(Vec self, other): 341 cdef Vec vec = <Vec> vec_div(self, other) 342 CHKERR(VecReciprocal(vec.vec)) 343 return vec 344 345# -------------------------------------------------------------------- 346 347cdef inline int Vec_Sizes(object size, object bsize, 348 PetscInt *b, PetscInt *n, PetscInt *N) except -1: 349 Sys_Sizes(size, bsize, b, n, N) 350 return 0 351 352# -------------------------------------------------------------------- 353 354ctypedef PetscErrorCode VecSetValuesFcn(PetscVec, 355 PetscInt, const PetscInt*, 356 const PetscScalar*, PetscInsertMode) 357 358cdef inline VecSetValuesFcn* vecsetvalues_fcn(int blocked, int local): 359 cdef VecSetValuesFcn *setvalues = NULL 360 if blocked and local: setvalues = VecSetValuesBlockedLocal 361 elif blocked: setvalues = VecSetValuesBlocked 362 elif local: setvalues = VecSetValuesLocal 363 else: setvalues = VecSetValues 364 return setvalues 365 366cdef inline int vecsetvalues(PetscVec V, 367 object oi, object ov, object oim, 368 int blocked, int local) except -1: 369 # block size 370 cdef PetscInt bs=1 371 if blocked: 372 CHKERR(VecGetBlockSize(V, &bs)) 373 if bs < 1: bs = 1 374 # indices and values 375 cdef PetscInt ni=0, nv=0 376 cdef PetscInt *i=NULL 377 cdef PetscScalar *v=NULL 378 cdef object unused1 = iarray_i(oi, &ni, &i) 379 cdef object unused2 = iarray_s(ov, &nv, &v) 380 if ni*bs != nv: raise ValueError( 381 "incompatible array sizes: ni=%d, nv=%d, bs=%d" % 382 (toInt(ni), toInt(nv), toInt(bs))) 383 # VecSetValuesXXX function and insert mode 384 cdef VecSetValuesFcn *setvalues = vecsetvalues_fcn(blocked, local) 385 cdef PetscInsertMode addv = insertmode(oim) 386 # actual call 387 CHKERR(setvalues(V, ni, i, v, addv)) 388 return 0 389 390cdef object vecgetvalues(PetscVec vec, object oindices, object values): 391 cdef PetscInt ni=0, nv=0 392 cdef PetscInt *i=NULL 393 cdef PetscScalar *v=NULL 394 cdef object indices = iarray_i(oindices, &ni, &i) 395 if values is None: 396 values = empty_s(ni) 397 values.shape = indices.shape 398 values = oarray_s(values, &nv, &v) 399 if (ni != nv): raise ValueError( 400 ("incompatible array sizes: " 401 "ni=%d, nv=%d") % (toInt(ni), toInt(nv))) 402 CHKERR(VecGetValues(vec, ni, i, v)) 403 return values 404 405# -------------------------------------------------------------------- 406 407cdef inline _Vec_buffer vec_getbuffer_r(Vec self): 408 cdef _Vec_buffer buf = _Vec_buffer(self) 409 buf.readonly = 1 410 return buf 411 412cdef inline _Vec_buffer vec_getbuffer_w(Vec self): 413 cdef _Vec_buffer buf = _Vec_buffer(self) 414 buf.readonly = 0 415 return buf 416 417cdef inline ndarray vec_getarray_r(Vec self): 418 return asarray(vec_getbuffer_r(self)) 419 420cdef inline ndarray vec_getarray_w(Vec self): 421 return asarray(vec_getbuffer_w(self)) 422 423cdef inline int vec_setarray(Vec self, object o) except -1: 424 cdef PetscInt na=0, nv=0, i=0 425 cdef PetscScalar *va=NULL, *vv=NULL 426 cdef ndarray ary = iarray_s(o, &na, &va) 427 CHKERR(VecGetLocalSize(self.vec, &nv)) 428 if (na != nv) and PyArray_NDIM(ary) > 0: raise ValueError( 429 "array size %d incompatible with vector local size %d" % 430 (toInt(na), toInt(nv))) 431 CHKERR(VecGetArray(self.vec, &vv)) 432 try: 433 if PyArray_NDIM(ary) == 0: 434 for i from 0 <= i < nv: 435 vv[i] = va[0] 436 else: 437 CHKERR(PetscMemcpy(vv, va, <size_t>nv*sizeof(PetscScalar))) 438 finally: 439 CHKERR(VecRestoreArray(self.vec, &vv)) 440 return 0 441 442cdef object vec_getitem(Vec self, object i): 443 cdef PetscInt N=0 444 if i is Ellipsis: 445 return asarray(self) 446 if isinstance(i, slice): 447 CHKERR(VecGetSize(self.vec, &N)) 448 start, stop, stride = i.indices(toInt(N)) 449 i = arange(start, stop, stride) 450 return vecgetvalues(self.vec, i, None) 451 452cdef int vec_setitem(Vec self, object i, object v) except -1: 453 cdef PetscInt N=0 454 if i is Ellipsis: 455 return vec_setarray(self, v) 456 if isinstance(i, slice): 457 CHKERR(VecGetSize(self.vec, &N)) 458 start, stop, stride = i.indices(toInt(N)) 459 i = arange(start, stop, stride) 460 vecsetvalues(self.vec, i, v, None, 0, 0) 461 return 0 462 463cdef vec_get_dlpack_ctx(Vec self): 464 cdef object ctx0 = self.get_attr('__dltensor_ctx__') 465 cdef PetscInt n = 0 466 cdef int64_t ndim = 1 467 cdef int64_t* shape_arr = NULL 468 cdef int64_t* strides_arr = NULL 469 cdef object s1 = None 470 cdef object s2 = None 471 cdef PetscInt devId = 0 472 cdef PetscMemType mtype = PETSC_MEMTYPE_HOST 473 if ctx0 is None: # First time in, create a linear memory view 474 s1 = oarray_p(empty_p(<PetscInt>ndim), NULL, <void**>&shape_arr) 475 s2 = oarray_p(empty_p(<PetscInt>ndim), NULL, <void**>&strides_arr) 476 CHKERR(VecGetLocalSize(self.vec, &n)) 477 shape_arr[0] = <int64_t>n 478 strides_arr[0] = 1 479 else: 480 (_, _, ndim, s1, s2) = ctx0 481 482 devType_ = {PETSC_MEMTYPE_HOST : kDLCPU, PETSC_MEMTYPE_CUDA : kDLCUDA, PETSC_MEMTYPE_HIP : kDLROCM} 483 CHKERR(VecGetCurrentMemType(self.vec, &mtype)) 484 dtype = devType_.get(mtype, kDLCPU) 485 if dtype != kDLCPU: 486 CHKERR(PetscObjectGetDeviceId(<PetscObject>self.vec, &devId)) 487 ctx0 = (dtype, devId, ndim, s1, s2) 488 self.set_attr('__dltensor_ctx__', ctx0) 489 return ctx0 490 491# -------------------------------------------------------------------- 492 493cdef int Vec_AcquireArray(PetscVec v, PetscScalar *a[], int ro) except -1 nogil: 494 if ro: CHKERR(VecGetArrayRead(v, <const PetscScalar**>a)) 495 else: CHKERR(VecGetArray(v, a)) 496 return 0 497 498cdef int Vec_ReleaseArray(PetscVec v, PetscScalar *a[], int ro) except -1 nogil: 499 if ro: CHKERR(VecRestoreArrayRead(v, <const PetscScalar**>a)) 500 else: CHKERR(VecRestoreArray(v, a)) 501 return 0 502 503cdef class _Vec_buffer: 504 505 cdef PetscVec vec 506 cdef PetscInt size 507 cdef PetscScalar *data 508 cdef bint readonly 509 cdef bint hasarray 510 511 def __cinit__(self, Vec vec, bint readonly=0): 512 cdef PetscVec v = vec.vec 513 CHKERR(PetscINCREF(<PetscObject*>&v)) 514 self.vec = v 515 self.size = 0 516 self.data = NULL 517 self.readonly = 1 if readonly else 0 518 self.hasarray = 0 519 520 def __dealloc__(self): 521 if self.hasarray and self.vec != NULL: 522 Vec_ReleaseArray(self.vec, &self.data, self.readonly) 523 CHKERR(VecDestroy(&self.vec)) 524 525 # 526 527 cdef int acquire(self) except -1 nogil: 528 if not self.hasarray and self.vec != NULL: 529 CHKERR(VecGetLocalSize(self.vec, &self.size)) 530 Vec_AcquireArray(self.vec, &self.data, self.readonly) 531 self.hasarray = 1 532 return 0 533 534 cdef int release(self) except -1 nogil: 535 if self.hasarray and self.vec != NULL: 536 self.size = 0 537 Vec_ReleaseArray(self.vec, &self.data, self.readonly) 538 self.hasarray = 0 539 return 0 540 541 # buffer interface (PEP 3118) 542 543 cdef int acquirebuffer(self, Py_buffer *view, int flags) except -1: 544 self.acquire() 545 PyPetscBuffer_FillInfo(view, <void*>self.data, self.size, 546 c's', self.readonly, flags) 547 view.obj = self 548 return 0 549 550 cdef int releasebuffer(self, Py_buffer *view) except -1: 551 PyPetscBuffer_Release(view) 552 self.release() 553 return 0 554 555 def __getbuffer__(self, Py_buffer *view, int flags): 556 self.acquirebuffer(view, flags) 557 558 def __releasebuffer__(self, Py_buffer *view): 559 self.releasebuffer(view) 560 561 # 'with' statement (PEP 343) 562 563 cdef object enter(self): 564 self.acquire() 565 return asarray(self) 566 567 cdef object exit(self): 568 self.release() 569 return None 570 571 def __enter__(self): 572 return self.enter() 573 574 def __exit__(self, *exc): 575 return self.exit() 576 577 # NumPy array interface (legacy) 578 579 property __array_interface__: 580 def __get__(self): 581 cdef PetscInt n = 0 582 if self.vec != NULL: 583 CHKERR(VecGetLocalSize(self.vec, &n)) 584 cdef object size = toInt(n) 585 cdef dtype descr = PyArray_DescrFromType(NPY_PETSC_SCALAR) 586 cdef str typestr = "=%c%d" % (descr.kind, descr.itemsize) 587 return dict(version=3, 588 data=self, 589 shape=(size,), 590 typestr=typestr) 591 592# -------------------------------------------------------------------- 593 594cdef class _Vec_LocalForm: 595 596 cdef Vec gvec 597 cdef Vec lvec 598 599 def __cinit__(self, Vec gvec): 600 self.gvec = gvec 601 self.lvec = Vec() 602 603 def __enter__(self): 604 cdef PetscVec gvec = self.gvec.vec 605 CHKERR(VecGhostGetLocalForm(gvec, &self.lvec.vec)) 606 return self.lvec 607 608 def __exit__(self, *exc): 609 cdef PetscVec gvec = self.gvec.vec 610 CHKERR(VecGhostRestoreLocalForm(gvec, &self.lvec.vec)) 611 self.lvec.vec = NULL 612