1# -------------------------------------------------------------------- 2 3cdef extern from * nogil: 4 5 ctypedef enum PetscDMDAStencilType"DMDAStencilType": 6 DMDA_STENCIL_STAR 7 DMDA_STENCIL_BOX 8 9 ctypedef enum PetscDMDAInterpolationType"DMDAInterpolationType": 10 DMDA_INTERPOLATION_Q0 "DMDA_Q0" 11 DMDA_INTERPOLATION_Q1 "DMDA_Q1" 12 13 ctypedef enum PetscDMDAElementType"DMDAElementType": 14 DMDA_ELEMENT_P1 15 DMDA_ELEMENT_Q1 16 17 PetscErrorCode DMDACreateND(MPI_Comm, 18 PetscInt, PetscInt, # dim, dof 19 PetscInt, PetscInt, PetscInt, # M, N, P 20 PetscInt, PetscInt, PetscInt, # m, n, p 21 PetscInt[], PetscInt[], PetscInt[], # lx, ly, lz 22 PetscDMBoundaryType, # bx 23 PetscDMBoundaryType, # by 24 PetscDMBoundaryType, # bz 25 PetscDMDAStencilType, # stencil type 26 PetscInt, # stencil width 27 PetscDM*) 28 29 PetscErrorCode DMDASetDof(PetscDM, PetscInt) 30 PetscErrorCode DMDASetSizes(PetscDM, PetscInt, PetscInt, PetscInt) 31 PetscErrorCode DMDASetNumProcs(PetscDM, PetscInt, PetscInt, PetscInt) 32 PetscErrorCode DMDASetBoundaryType(PetscDM, PetscDMBoundaryType, PetscDMBoundaryType, PetscDMBoundaryType) 33 PetscErrorCode DMDASetStencilType(PetscDM, PetscDMDAStencilType) 34 PetscErrorCode DMDASetStencilWidth(PetscDM, PetscInt) 35 36 PetscErrorCode DMDAGetInfo(PetscDM, 37 PetscInt*, 38 PetscInt*, PetscInt*, PetscInt*, 39 PetscInt*, PetscInt*, PetscInt*, 40 PetscInt*, PetscInt*, 41 PetscDMBoundaryType*, 42 PetscDMBoundaryType*, 43 PetscDMBoundaryType*, 44 PetscDMDAStencilType*) 45 PetscErrorCode DMDAGetCorners(PetscDM, 46 PetscInt*, PetscInt*, PetscInt*, 47 PetscInt*, PetscInt*, PetscInt*) 48 PetscErrorCode DMDAGetGhostCorners(PetscDM, 49 PetscInt*, PetscInt*, PetscInt*, 50 PetscInt*, PetscInt*, PetscInt*) 51 PetscErrorCode DMDAGetOwnershipRanges(PetscDM, 52 const PetscInt*[], 53 const PetscInt*[], 54 const PetscInt*[]) 55 56 PetscErrorCode DMDASetUniformCoordinates(PetscDM, 57 PetscReal, PetscReal, 58 PetscReal, PetscReal, 59 PetscReal, PetscReal) 60 PetscErrorCode DMGetBoundingBox(PetscDM, PetscReal[], PetscReal[]) 61 PetscErrorCode DMGetLocalBoundingBox(PetscDM, PetscReal[], PetscReal[]) 62 63 PetscErrorCode DMDACreateNaturalVector(PetscDM, PetscVec*) 64 PetscErrorCode DMDAGlobalToNaturalBegin(PetscDM, PetscVec, PetscInsertMode, PetscVec) 65 PetscErrorCode DMDAGlobalToNaturalEnd(PetscDM, PetscVec, PetscInsertMode, PetscVec) 66 PetscErrorCode DMDANaturalToGlobalBegin(PetscDM, PetscVec, PetscInsertMode, PetscVec) 67 PetscErrorCode DMDANaturalToGlobalEnd(PetscDM, PetscVec, PetscInsertMode, PetscVec) 68 69 PetscErrorCode DMDAGetAO(PetscDM, PetscAO*) 70 PetscErrorCode DMDAGetScatter(PetscDM, PetscScatter*, PetscScatter*) 71 72 PetscErrorCode DMDASetRefinementFactor(PetscDM, PetscInt, PetscInt, PetscInt) 73 PetscErrorCode DMDAGetRefinementFactor(PetscDM, PetscInt*, PetscInt*, PetscInt*) 74 PetscErrorCode DMDASetInterpolationType(PetscDM, PetscDMDAInterpolationType) 75 PetscErrorCode DMDAGetInterpolationType(PetscDM, PetscDMDAInterpolationType*) 76 PetscErrorCode DMDASetElementType(PetscDM, PetscDMDAElementType) 77 PetscErrorCode DMDAGetElementType(PetscDM, PetscDMDAElementType*) 78 PetscErrorCode DMDAGetElements(PetscDM, PetscInt*, PetscInt*, const PetscInt**) 79 PetscErrorCode DMDARestoreElements(PetscDM, PetscInt*, PetscInt*, const PetscInt**) 80 81 PetscErrorCode DMDASetFieldName(PetscDM, PetscInt, const char[]) 82 PetscErrorCode DMDAGetFieldName(PetscDM, PetscInt, const char*[]) 83 PetscErrorCode DMDASetCoordinateName(PetscDM, PetscInt, const char[]) 84 PetscErrorCode DMDAGetCoordinateName(PetscDM, PetscInt, const char*[]) 85 86# -------------------------------------------------------------------- 87 88cdef inline PetscDMDAStencilType asStencil(object stencil) \ 89 except <PetscDMDAStencilType>(-1): 90 if isinstance(stencil, str): 91 if stencil == "star": return DMDA_STENCIL_STAR 92 elif stencil == "box": return DMDA_STENCIL_BOX 93 else: raise ValueError("unknown stencil type: %s" % stencil) 94 return stencil 95 96cdef inline object toStencil(PetscDMDAStencilType stype): 97 if stype == DMDA_STENCIL_STAR: return "star" 98 elif stype == DMDA_STENCIL_BOX: return "box" 99 100cdef inline PetscDMDAInterpolationType dainterpolationtype(object itype) \ 101 except <PetscDMDAInterpolationType>(-1): 102 if (isinstance(itype, str)): 103 if itype in ("q0", "Q0"): return DMDA_INTERPOLATION_Q0 104 if itype in ("q1", "Q1"): return DMDA_INTERPOLATION_Q1 105 else: raise ValueError("unknown interpolation type: %s" % itype) 106 return itype 107 108cdef inline PetscDMDAElementType daelementtype(object etype) \ 109 except <PetscDMDAElementType>(-1): 110 if (isinstance(etype, str)): 111 if etype in ("p1", "P1"): return DMDA_ELEMENT_P1 112 if etype in ("q1", "Q1"): return DMDA_ELEMENT_Q1 113 else: raise ValueError("unknown element type: %s" % etype) 114 return etype 115 116cdef inline PetscErrorCode DMDAGetDim(PetscDM da, PetscInt *dim) noexcept nogil: 117 return DMDAGetInfo(da, dim, 118 NULL, NULL, NULL, 119 NULL, NULL, NULL, 120 NULL, NULL, 121 NULL, NULL, NULL, 122 NULL) 123 124cdef inline PetscInt asDims(dims, 125 PetscInt *_M, 126 PetscInt *_N, 127 PetscInt *_P) except? -1: 128 cdef PetscInt dim = PETSC_DECIDE 129 cdef object M=None, N=None, P=None 130 dims = tuple(dims) 131 dim = <PetscInt>len(dims) 132 if dim == 0: pass 133 elif dim == 1: M, = dims 134 elif dim == 2: M, N = dims 135 elif dim == 3: M, N, P = dims 136 if dim >= 1: _M[0] = asInt(M) 137 if dim >= 2: _N[0] = asInt(N) 138 if dim >= 3: _P[0] = asInt(P) 139 return dim 140 141cdef inline tuple toDims(PetscInt dim, 142 PetscInt M, 143 PetscInt N, 144 PetscInt P): 145 if dim == 0: return () 146 elif dim == 1: return (toInt(M),) 147 elif dim == 2: return (toInt(M), toInt(N)) 148 elif dim == 3: return (toInt(M), toInt(N), toInt(P)) 149 150cdef inline tuple asOwnershipRanges(object ownership_ranges, 151 PetscInt dim, 152 PetscInt *m, PetscInt *n, PetscInt *p, 153 PetscInt **_x, 154 PetscInt **_y, 155 PetscInt **_z): 156 cdef object ranges = list(ownership_ranges) 157 cdef PetscInt rdim = <PetscInt>len(ranges) 158 cdef PetscInt nlx=0, nly=0, nlz=0 159 if dim == PETSC_DECIDE: dim = rdim 160 elif dim != rdim: raise ValueError( 161 "number of dimensions %d and number ownership ranges %d" % 162 (toInt(dim), toInt(rdim))) 163 if dim >= 1: 164 ranges[0] = iarray_i(ranges[0], &nlx, _x) 165 if m[0] == PETSC_DECIDE: m[0] = nlx 166 elif m[0] != nlx: raise ValueError( 167 "ownership range size %d and number or processors %d" % 168 (toInt(nlx), toInt(m[0]))) 169 if dim >= 2: 170 ranges[1] = iarray_i(ranges[1], &nly, _y) 171 if n[0] == PETSC_DECIDE: n[0] = nly 172 elif n[0] != nly: raise ValueError( 173 "ownership range size %d and number or processors %d" % 174 (toInt(nly), toInt(n[0]))) 175 if dim >= 3: 176 ranges[2] = iarray_i(ranges[2], &nlz, _z) 177 if p[0] == PETSC_DECIDE: p[0] = nlz 178 elif p[0] != nlz: raise ValueError( 179 "ownership range size %d and number or processors %d" % 180 (toInt(nlz), toInt(p[0]))) 181 return tuple(ranges) 182 183cdef inline tuple toOwnershipRanges(PetscInt dim, 184 PetscInt m, PetscInt n, PetscInt p, 185 const PetscInt *lx, 186 const PetscInt *ly, 187 const PetscInt *lz): 188 # Returns tuple of arrays containing ownership ranges as Python arrays 189 ranges = [array_i(m, lx)] 190 if dim > 1: 191 ranges.append(array_i(n, ly)) 192 if dim > 2: 193 ranges.append(array_i(p, lz)) 194 return tuple(ranges) 195 196# -------------------------------------------------------------------- 197 198cdef class _DMDA_Vec_array(object): 199 200 cdef _Vec_buffer vecbuf 201 cdef readonly tuple starts, sizes 202 cdef readonly tuple shape, strides 203 cdef readonly ndarray array 204 205 def __cinit__(self, DMDA da, Vec vec, bint readonly, bint DOF=False): 206 # 207 cdef PetscInt dim=0, dof=0 208 CHKERR(DMDAGetInfo(da.dm, 209 &dim, NULL, NULL, NULL, NULL, NULL, NULL, 210 &dof, NULL, NULL, NULL, NULL, NULL)) 211 cdef PetscInt lxs=0, lys=0, lzs=0 212 cdef PetscInt lxm=0, lym=0, lzm=0 213 CHKERR(DMDAGetCorners(da.dm, 214 &lxs, &lys, &lzs, 215 &lxm, &lym, &lzm)) 216 cdef PetscInt gxs=0, gys=0, gzs=0 217 cdef PetscInt gxm=0, gym=0, gzm=0 218 CHKERR(DMDAGetGhostCorners(da.dm, 219 &gxs, &gys, &gzs, 220 &gxm, &gym, &gzm)) 221 # 222 cdef PetscInt n=0 223 CHKERR(VecGetLocalSize(vec.vec, &n)) 224 cdef PetscInt xs, ys, zs, xm, ym, zm 225 if (n == lxm*lym*lzm*dof): 226 xs, ys, zs = lxs, lys, lzs 227 xm, ym, zm = lxm, lym, lzm 228 elif (n == gxm*gym*gzm*dof): 229 xs, ys, zs = gxs, gys, gzs 230 xm, ym, zm = gxm, gym, gzm 231 else: 232 raise ValueError( 233 "Vector local size %d is not compatible " 234 "with DMDA local sizes %s" 235 % (<Py_ssize_t>n, toDims(dim, lxm, lym, lzm))) 236 # 237 cdef tuple starts = toDims(dim, xs, ys, zs) 238 cdef tuple sizes = toDims(dim, xm, ym, zm) 239 cdef Py_ssize_t k = <Py_ssize_t>sizeof(PetscScalar) 240 cdef Py_ssize_t f = <Py_ssize_t>dof 241 cdef Py_ssize_t d = <Py_ssize_t>dim 242 cdef tuple shape = toDims(dim, xm, ym, zm) 243 cdef tuple strides = (k*f, k*f*xm, k*f*xm*ym)[:d] 244 if DOF or f > 1: shape += (f,) 245 if DOF or f > 1: strides += (k,) 246 # 247 self.vecbuf = _Vec_buffer(vec, readonly) 248 self.starts = starts 249 self.sizes = sizes 250 self.shape = shape 251 self.strides = strides 252 253 cdef int acquire(self) except -1: 254 self.vecbuf.acquire() 255 if self.array is None: 256 self.array = asarray(self.vecbuf) 257 self.array.shape = self.shape 258 self.array.strides = self.strides 259 return 0 260 261 cdef int release(self) except -1: 262 self.vecbuf.release() 263 self.array = None 264 return 0 265 266 # 267 268 def __getitem__(self, index): 269 self.acquire() 270 index = adjust_index_exp(self.starts, index) 271 return self.array[index] 272 273 def __setitem__(self, index, value): 274 self.acquire() 275 index = adjust_index_exp(self.starts, index) 276 self.array[index] = value 277 278 # 'with' statement (PEP 343) 279 280 def __enter__(self): 281 self.acquire() 282 return self 283 284 def __exit__(self, *exc): 285 self.release() 286 return None 287 288 289cdef object adjust_index_exp(object starts, object index): 290 if not isinstance(index, tuple): 291 return adjust_index(starts[0], index) 292 index = list(index) 293 for i, start in enumerate(starts): 294 index[i] = adjust_index(start, index[i]) 295 index = tuple(index) 296 return index 297 298cdef object adjust_index(object lbound, object index): 299 if index is None: 300 return index 301 if index is Ellipsis: 302 return index 303 if isinstance(index, slice): 304 start = index.start 305 stop = index.stop 306 step = index.step 307 if start is not None: start -= lbound 308 if stop is not None: stop -= lbound 309 return slice(start, stop, step) 310 try: 311 return index - lbound 312 except TypeError: 313 return index 314 315# -------------------------------------------------------------------- 316