1# -------------------------------------------------------------------- 2 3class DMStagStencilType(object): 4 """Stencil types.""" 5 STAR = DMSTAG_STENCIL_STAR 6 BOX = DMSTAG_STENCIL_BOX 7 NONE = DMSTAG_STENCIL_NONE 8 9 10class DMStagStencilLocation(object): 11 """Stencil location types.""" 12 NULLLOC = DMSTAG_NULL_LOCATION 13 BACK_DOWN_LEFT = DMSTAG_BACK_DOWN_LEFT 14 BACK_DOWN = DMSTAG_BACK_DOWN 15 BACK_DOWN_RIGHT = DMSTAG_BACK_DOWN_RIGHT 16 BACK_LEFT = DMSTAG_BACK_LEFT 17 BACK = DMSTAG_BACK 18 BACK_RIGHT = DMSTAG_BACK_RIGHT 19 BACK_UP_LEFT = DMSTAG_BACK_UP_LEFT 20 BACK_UP = DMSTAG_BACK_UP 21 BACK_UP_RIGHT = DMSTAG_BACK_UP_RIGHT 22 DOWN_LEFT = DMSTAG_DOWN_LEFT 23 DOWN = DMSTAG_DOWN 24 DOWN_RIGHT = DMSTAG_DOWN_RIGHT 25 LEFT = DMSTAG_LEFT 26 ELEMENT = DMSTAG_ELEMENT 27 RIGHT = DMSTAG_RIGHT 28 UP_LEFT = DMSTAG_UP_LEFT 29 UP = DMSTAG_UP 30 UP_RIGHT = DMSTAG_UP_RIGHT 31 FRONT_DOWN_LEFT = DMSTAG_FRONT_DOWN_LEFT 32 FRONT_DOWN = DMSTAG_FRONT_DOWN 33 FRONT_DOWN_RIGHT = DMSTAG_FRONT_DOWN_RIGHT 34 FRONT_LEFT = DMSTAG_FRONT_LEFT 35 FRONT = DMSTAG_FRONT 36 FRONT_RIGHT = DMSTAG_FRONT_RIGHT 37 FRONT_UP_LEFT = DMSTAG_FRONT_UP_LEFT 38 FRONT_UP = DMSTAG_FRONT_UP 39 FRONT_UP_RIGHT = DMSTAG_FRONT_UP_RIGHT 40 41# -------------------------------------------------------------------- 42 43 44cdef class DMStag(DM): 45 """A DM object representing a "staggered grid" or a structured cell complex.""" 46 47 StencilType = DMStagStencilType 48 StencilLocation = DMStagStencilLocation 49 50 def create( 51 self, 52 dim: int, 53 dofs: tuple[int, ...] | None = None, 54 sizes: tuple[int, ...] | None = None, 55 boundary_types: tuple[DM.BoundaryType | int | str | bool, ...] | None = None, 56 stencil_type: StencilType | None = None, 57 stencil_width: int | None = None, 58 proc_sizes: tuple[int, ...] | None = None, 59 ownership_ranges: tuple[Sequence[int], ...] | None = None, 60 comm: Comm | None = None, 61 setUp: bool | None = False) -> Self: 62 """Create a DMDA object. 63 64 Collective. 65 66 Creates an object to manage data living on the elements and vertices / 67 the elements, faces, and vertices / the elements, faces, edges, and 68 vertices of a parallelized regular 1D / 2D / 3D grid. 69 70 Parameters 71 ---------- 72 dim 73 The number of dimensions. 74 dofs 75 The number of degrees of freedom per vertex, element (1D); vertex, 76 face, element (2D); or vertex, edge, face, element (3D). 77 sizes 78 The number of elements in each dimension. 79 boundary_types 80 The boundary types. 81 stencil_type 82 The ghost/halo stencil type. 83 stencil_width 84 The width of the ghost/halo region. 85 proc_sizes 86 The number of processes in x, y, z dimensions. 87 ownership_ranges 88 Local x, y, z element counts, of length equal to ``proc_sizes``, 89 summing to ``sizes``. 90 comm 91 MPI communicator, defaults to `Sys.getDefaultComm`. 92 setUp 93 Whether to call the setup routine after creating the object. 94 95 See Also 96 -------- 97 petsc.DMStagCreate1d, petsc.DMStagCreate2d, petsc.DMStagCreate3d 98 petsc.DMSetUp 99 100 """ 101 # ndim 102 cdef PetscInt ndim = asInt(dim) 103 104 # sizes 105 cdef object gsizes = sizes 106 cdef PetscInt nsizes=PETSC_DECIDE, M=1, N=1, P=1 107 if sizes is not None: 108 nsizes = asStagDims(gsizes, &M, &N, &P) 109 assert(nsizes==ndim) 110 111 # dofs 112 cdef object cdofs = dofs 113 cdef PetscInt ndofs=PETSC_DECIDE, dof0=1, dof1=0, dof2=0, dof3=0 114 if dofs is not None: 115 ndofs = asDofs(cdofs, &dof0, &dof1, &dof2, &dof3) 116 assert(ndofs==ndim+1) 117 118 # boundary types 119 cdef PetscDMBoundaryType btx = DM_BOUNDARY_NONE 120 cdef PetscDMBoundaryType bty = DM_BOUNDARY_NONE 121 cdef PetscDMBoundaryType btz = DM_BOUNDARY_NONE 122 asBoundary(boundary_types, &btx, &bty, &btz) 123 124 # stencil 125 cdef PetscInt swidth = 0 126 if stencil_width is not None: 127 swidth = asInt(stencil_width) 128 cdef PetscDMStagStencilType stype = DMSTAG_STENCIL_NONE 129 if stencil_type is not None: 130 stype = asStagStencil(stencil_type) 131 132 # comm 133 cdef MPI_Comm ccomm = def_Comm(comm, PETSC_COMM_DEFAULT) 134 135 # proc sizes 136 cdef object psizes = proc_sizes 137 cdef PetscInt nprocs=PETSC_DECIDE, m=PETSC_DECIDE, n=PETSC_DECIDE, p=PETSC_DECIDE 138 if proc_sizes is not None: 139 nprocs = asStagDims(psizes, &m, &n, &p) 140 assert(nprocs==ndim) 141 142 # ownership ranges 143 cdef PetscInt *lx = NULL, *ly = NULL, *lz = NULL 144 if ownership_ranges is not None: 145 nranges = asStagOwnershipRanges(ownership_ranges, ndim, &m, &n, &p, &lx, &ly, &lz) 146 assert(nranges==ndim) 147 148 # create 149 cdef PetscDM newda = NULL 150 if dim == 1: 151 CHKERR(DMStagCreate1d(ccomm, btx, M, dof0, dof1, stype, swidth, lx, &newda)) 152 if dim == 2: 153 CHKERR(DMStagCreate2d(ccomm, btx, bty, M, N, m, n, dof0, dof1, dof2, stype, swidth, lx, ly, &newda)) 154 if dim == 3: 155 CHKERR(DMStagCreate3d(ccomm, btx, bty, btz, M, N, P, m, n, p, dof0, dof1, dof2, dof3, stype, swidth, lx, ly, lz, &newda)) 156 CHKERR(PetscCLEAR(self.obj)); self.dm = newda 157 if setUp: 158 CHKERR(DMSetUp(self.dm)) 159 return self 160 161 # Setters 162 163 def setStencilWidth(self, swidth: int) -> None: 164 """Set elementwise stencil width. 165 166 Logically collective. 167 168 The width value is not used when `StencilType.NONE` is specified. 169 170 Parameters 171 ---------- 172 swidth 173 Stencil/halo/ghost width in elements. 174 175 See Also 176 -------- 177 petsc.DMStagSetStencilWidth 178 179 """ 180 cdef PetscInt sw = asInt(swidth) 181 CHKERR(DMStagSetStencilWidth(self.dm, sw)) 182 183 def setStencilType(self, stenciltype: StencilType | str) -> None: 184 """Set elementwise ghost/halo stencil type. 185 186 Logically collective. 187 188 Parameters 189 ---------- 190 stenciltype 191 The elementwise ghost stencil type. 192 193 See Also 194 -------- 195 getStencilType, petsc.DMStagSetStencilType 196 197 """ 198 cdef PetscDMStagStencilType stype = asStagStencil(stenciltype) 199 CHKERR(DMStagSetStencilType(self.dm, stype)) 200 201 def setBoundaryTypes( 202 self, 203 boundary_types: tuple[DM.BoundaryType | int | str | bool, ...]) -> None: 204 """Set the boundary types. 205 206 Logically collective. 207 208 Parameters 209 ---------- 210 boundary_types 211 Boundary types for one/two/three dimensions. 212 213 See Also 214 -------- 215 getBoundaryTypes, petsc.DMStagSetBoundaryTypes 216 217 """ 218 cdef PetscDMBoundaryType btx = DM_BOUNDARY_NONE 219 cdef PetscDMBoundaryType bty = DM_BOUNDARY_NONE 220 cdef PetscDMBoundaryType btz = DM_BOUNDARY_NONE 221 asBoundary(boundary_types, &btx, &bty, &btz) 222 CHKERR(DMStagSetBoundaryTypes(self.dm, btx, bty, btz)) 223 224 def setDof(self, dofs: tuple[int, ...]) -> None: 225 """Set DOFs/stratum. 226 227 Logically collective. 228 229 Parameters 230 ---------- 231 dofs 232 The number of points per 0-cell (vertex/node), 1-cell (element in 233 1D, edge in 2D and 3D), 2-cell (element in 2D, face in 3D), or 234 3-cell (element in 3D). 235 236 See Also 237 -------- 238 petsc.DMStagSetDOF 239 240 """ 241 cdef tuple gdofs = tuple(dofs) 242 cdef PetscInt dof0=1, dof1=0, dof2=0, dof3=0 243 asDofs(gdofs, &dof0, &dof1, &dof2, &dof3) 244 CHKERR(DMStagSetDOF(self.dm, dof0, dof1, dof2, dof3)) 245 246 def setGlobalSizes(self, sizes: tuple[int, ...]) -> None: 247 """Set global element counts in each dimension. 248 249 Logically collective. 250 251 Parameters 252 ---------- 253 sizes 254 Global elementwise size in the one/two/three dimensions. 255 256 See Also 257 -------- 258 petsc.DMStagSetGlobalSizes 259 260 """ 261 cdef tuple gsizes = tuple(sizes) 262 cdef PetscInt M=1, N=1, P=1 263 asStagDims(gsizes, &M, &N, &P) 264 CHKERR(DMStagSetGlobalSizes(self.dm, M, N, P)) 265 266 def setProcSizes(self, sizes: tuple[int, ...]) -> None: 267 """Set the number of processes in each dimension in the global process grid. 268 269 Logically collective. 270 271 Parameters 272 ---------- 273 sizes 274 Number of processes in one/two/three dimensions. 275 276 See Also 277 -------- 278 petsc.DMStagSetNumRanks 279 280 """ 281 cdef tuple psizes = tuple(sizes) 282 cdef PetscInt m=PETSC_DECIDE, n=PETSC_DECIDE, p=PETSC_DECIDE 283 asStagDims(psizes, &m, &n, &p) 284 CHKERR(DMStagSetNumRanks(self.dm, m, n, p)) 285 286 def setOwnershipRanges(self, ranges: tuple[Sequence[int], ...]) -> None: 287 """Set elements per process in each dimension. 288 289 Logically collective. 290 291 Parameters 292 ---------- 293 ranges 294 Element counts for each process in one/two/three dimensions. 295 296 See Also 297 -------- 298 getOwnershipRanges, petsc.DMStagSetOwnershipRanges 299 300 """ 301 cdef PetscInt dim=0, m=PETSC_DECIDE, n=PETSC_DECIDE, p=PETSC_DECIDE 302 cdef PetscInt *lx = NULL, *ly = NULL, *lz = NULL 303 CHKERR(DMGetDimension(self.dm, &dim)) 304 CHKERR(DMStagGetNumRanks(self.dm, &m, &n, &p)) 305 asStagOwnershipRanges(ranges, dim, &m, &n, &p, &lx, &ly, &lz) 306 CHKERR(DMStagSetOwnershipRanges(self.dm, lx, ly, lz)) 307 308 # Getters 309 310 def getDim(self) -> int: 311 """Return the number of dimensions. 312 313 Not collective. 314 315 """ 316 return self.getDimension() 317 318 def getEntriesPerElement(self) -> int: 319 """Return the number of entries per element in the local representation. 320 321 Not collective. 322 323 This is the natural block size for most local operations. 324 325 See Also 326 -------- 327 petsc.DMStagGetEntriesPerElement 328 329 """ 330 cdef PetscInt epe=0 331 CHKERR(DMStagGetEntriesPerElement(self.dm, &epe)) 332 return toInt(epe) 333 334 def getStencilWidth(self) -> int: 335 """Return elementwise stencil width. 336 337 Not collective. 338 339 See Also 340 -------- 341 petsc.DMStagGetStencilWidth 342 343 """ 344 cdef PetscInt swidth=0 345 CHKERR(DMStagGetStencilWidth(self.dm, &swidth)) 346 return toInt(swidth) 347 348 def getDof(self) -> tuple[int, ...]: 349 """Get number of DOFs associated with each stratum of the grid. 350 351 Not collective. 352 353 See Also 354 -------- 355 petsc.DMStagGetDOF 356 357 """ 358 cdef PetscInt dim=0, dof0=0, dof1=0, dof2=0, dof3=0 359 CHKERR(DMStagGetDOF(self.dm, &dof0, &dof1, &dof2, &dof3)) 360 CHKERR(DMGetDimension(self.dm, &dim)) 361 return toDofs(dim+1, dof0, dof1, dof2, dof3) 362 363 def getCorners(self) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: 364 """Return starting element index, width and number of partial elements. 365 366 Not collective. 367 368 The returned value is calculated excluding ghost points. 369 370 The number of extra partial elements is either ``1`` or ``0``. The 371 value is ``1`` on right, top, and front non-periodic domain 372 ("physical") boundaries, in the x, y, and z dimensions respectively, 373 and otherwise ``0``. 374 375 See Also 376 -------- 377 getGhostCorners, petsc.DMStagGetCorners, petsc.DMGetDimension 378 379 """ 380 cdef PetscInt dim=0, x=0, y=0, z=0, m=0, n=0, p=0, nExtrax=0, nExtray=0, nExtraz=0 381 CHKERR(DMGetDimension(self.dm, &dim)) 382 CHKERR(DMStagGetCorners(self.dm, &x, &y, &z, &m, &n, &p, &nExtrax, &nExtray, &nExtraz)) 383 return (asInt(x), asInt(y), asInt(z))[:<Py_ssize_t>dim], (asInt(m), asInt(n), asInt(p))[:<Py_ssize_t>dim], (asInt(nExtrax), asInt(nExtray), asInt(nExtraz))[:<Py_ssize_t>dim] 384 385 def getGhostCorners(self) -> tuple[tuple[int, ...], tuple[int, ...]]: 386 """Return starting element index and width of local region. 387 388 Not collective. 389 390 See Also 391 -------- 392 getCorners, petsc.DMStagGetGhostCorners 393 394 """ 395 cdef PetscInt dim=0, x=0, y=0, z=0, m=0, n=0, p=0 396 CHKERR(DMGetDimension(self.dm, &dim)) 397 CHKERR(DMStagGetGhostCorners(self.dm, &x, &y, &z, &m, &n, &p)) 398 return (asInt(x), asInt(y), asInt(z))[:<Py_ssize_t>dim], (asInt(m), asInt(n), asInt(p))[:<Py_ssize_t>dim] 399 400 def getLocalSizes(self) -> tuple[int, ...]: 401 """Return local elementwise sizes in each dimension. 402 403 Not collective. 404 405 The returned value is calculated excluding ghost points. 406 407 See Also 408 -------- 409 getGlobalSizes, petsc.DMStagGetLocalSizes 410 411 """ 412 cdef PetscInt dim=0, m=PETSC_DECIDE, n=PETSC_DECIDE, p=PETSC_DECIDE 413 CHKERR(DMGetDimension(self.dm, &dim)) 414 CHKERR(DMStagGetLocalSizes(self.dm, &m, &n, &p)) 415 return toStagDims(dim, m, n, p) 416 417 def getGlobalSizes(self) -> tuple[int, ...]: 418 """Return global element counts in each dimension. 419 420 Not collective. 421 422 See Also 423 -------- 424 getLocalSizes, petsc.DMStagGetGlobalSizes 425 426 """ 427 cdef PetscInt dim=0, m=PETSC_DECIDE, n=PETSC_DECIDE, p=PETSC_DECIDE 428 CHKERR(DMGetDimension(self.dm, &dim)) 429 CHKERR(DMStagGetGlobalSizes(self.dm, &m, &n, &p)) 430 return toStagDims(dim, m, n, p) 431 432 def getProcSizes(self) -> tuple[int, ...]: 433 """Return number of processes in each dimension. 434 435 Not collective. 436 437 See Also 438 -------- 439 petsc.DMStagGetNumRanks 440 441 """ 442 cdef PetscInt dim=0, m=PETSC_DECIDE, n=PETSC_DECIDE, p=PETSC_DECIDE 443 CHKERR(DMGetDimension(self.dm, &dim)) 444 CHKERR(DMStagGetNumRanks(self.dm, &m, &n, &p)) 445 return toStagDims(dim, m, n, p) 446 447 def getStencilType(self) -> str: 448 """Return elementwise ghost/halo stencil type. 449 450 Not collective. 451 452 See Also 453 -------- 454 setStencilType, petsc.DMStagGetStencilType 455 456 """ 457 cdef PetscDMStagStencilType stype = DMSTAG_STENCIL_BOX 458 CHKERR(DMStagGetStencilType(self.dm, &stype)) 459 return toStagStencil(stype) 460 461 def getOwnershipRanges(self) -> tuple[Sequence[int], ...]: 462 """Return elements per process in each dimension. 463 464 Not collective. 465 466 See Also 467 -------- 468 setOwnershipRanges, petsc.DMStagGetOwnershipRanges 469 470 """ 471 cdef PetscInt dim=0, m=0, n=0, p=0 472 cdef const PetscInt *lx = NULL, *ly = NULL, *lz = NULL 473 CHKERR(DMGetDimension(self.dm, &dim)) 474 CHKERR(DMStagGetNumRanks(self.dm, &m, &n, &p)) 475 CHKERR(DMStagGetOwnershipRanges(self.dm, &lx, &ly, &lz)) 476 return toStagOwnershipRanges(dim, m, n, p, lx, ly, lz) 477 478 def getBoundaryTypes(self) -> tuple[str, ...]: 479 """Return boundary types in each dimension. 480 481 Not collective. 482 483 See Also 484 -------- 485 setBoundaryTypes, petsc.DMStagGetBoundaryTypes 486 487 """ 488 cdef PetscInt dim=0 489 cdef PetscDMBoundaryType btx = DM_BOUNDARY_NONE 490 cdef PetscDMBoundaryType bty = DM_BOUNDARY_NONE 491 cdef PetscDMBoundaryType btz = DM_BOUNDARY_NONE 492 CHKERR(DMGetDimension(self.dm, &dim)) 493 CHKERR(DMStagGetBoundaryTypes(self.dm, &btx, &bty, &btz)) 494 return toStagBoundaryTypes(dim, btx, bty, btz) 495 496 def getIsFirstRank(self) -> tuple[int, ...]: 497 """Return whether this process is first in each dimension in the process grid. 498 499 Not collective. 500 501 See Also 502 -------- 503 petsc.DMStagGetIsFirstRank 504 505 """ 506 cdef PetscBool rank0=PETSC_FALSE, rank1=PETSC_FALSE, rank2=PETSC_FALSE 507 cdef PetscInt dim=0 508 CHKERR(DMGetDimension(self.dm, &dim)) 509 CHKERR(DMStagGetIsFirstRank(self.dm, &rank0, &rank1, &rank2)) 510 return toStagDims(dim, rank0, rank1, rank2) 511 512 def getIsLastRank(self) -> tuple[int, ...]: 513 """Return whether this process is last in each dimension in the process grid. 514 515 Not collective. 516 517 See Also 518 -------- 519 petsc.DMStagGetIsLastRank 520 521 """ 522 cdef PetscBool rank0=PETSC_FALSE, rank1=PETSC_FALSE, rank2=PETSC_FALSE 523 cdef PetscInt dim=0 524 CHKERR(DMGetDimension(self.dm, &dim)) 525 CHKERR(DMStagGetIsLastRank(self.dm, &rank0, &rank1, &rank2)) 526 return toStagDims(dim, rank0, rank1, rank2) 527 528 # Coordinate-related functions 529 530 def setUniformCoordinatesExplicit( 531 self, 532 xmin: float = 0, 533 xmax: float = 1, 534 ymin: float = 0, 535 ymax: float = 1, 536 zmin: float = 0, 537 zmax: float = 1) -> None: 538 """Set coordinates to be a uniform grid, storing all values. 539 540 Collective. 541 542 Parameters 543 ---------- 544 xmin 545 The minimum global coordinate value in the x dimension. 546 xmax 547 The maximum global coordinate value in the x dimension. 548 ymin 549 The minimum global coordinate value in the y dimension. 550 ymax 551 The maximum global coordinate value in the y dimension. 552 zmin 553 The minimum global coordinate value in the z dimension. 554 zmax 555 The maximum global coordinate value in the z dimension. 556 557 See Also 558 -------- 559 setUniformCoordinatesProduct, setUniformCoordinates 560 petsc.DMStagSetUniformCoordinatesExplicit 561 562 """ 563 cdef PetscReal _xmin = asReal(xmin), _xmax = asReal(xmax) 564 cdef PetscReal _ymin = asReal(ymin), _ymax = asReal(ymax) 565 cdef PetscReal _zmin = asReal(zmin), _zmax = asReal(zmax) 566 CHKERR(DMStagSetUniformCoordinatesExplicit(self.dm, _xmin, _xmax, _ymin, _ymax, _zmin, _zmax)) 567 568 def setUniformCoordinatesProduct( 569 self, 570 xmin: float = 0, 571 xmax: float = 1, 572 ymin: float = 0, 573 ymax: float = 1, 574 zmin: float = 0, 575 zmax: float = 1) -> None: 576 """Create uniform coordinates, as a product of 1D arrays. 577 578 Collective. 579 580 The per-dimension 1-dimensional `DMStag` objects that comprise the 581 product always have active 0-cells (vertices, element boundaries) and 582 1-cells (element centers). 583 584 Parameters 585 ---------- 586 xmin 587 The minimum global coordinate value in the x dimension. 588 xmax 589 The maximum global coordinate value in the x dimension. 590 ymin 591 The minimum global coordinate value in the y dimension. 592 ymax 593 The maximum global coordinate value in the y dimension. 594 zmin 595 The minimum global coordinate value in the z dimension. 596 zmax 597 The maximum global coordinate value in the z dimension. 598 599 See Also 600 -------- 601 setUniformCoordinatesExplicit, setUniformCoordinates 602 petsc.DMStagSetUniformCoordinatesProduct 603 604 """ 605 cdef PetscReal _xmin = asReal(xmin), _xmax = asReal(xmax) 606 cdef PetscReal _ymin = asReal(ymin), _ymax = asReal(ymax) 607 cdef PetscReal _zmin = asReal(zmin), _zmax = asReal(zmax) 608 CHKERR(DMStagSetUniformCoordinatesProduct(self.dm, _xmin, _xmax, _ymin, _ymax, _zmin, _zmax)) 609 610 def setUniformCoordinates( 611 self, 612 xmin: float = 0, 613 xmax: float = 1, 614 ymin: float = 0, 615 ymax: float = 1, 616 zmin: float = 0, 617 zmax: float = 1) -> None: 618 """Set the coordinates to be a uniform grid.. 619 620 Collective. 621 622 Local coordinates are populated, linearly extrapolated to ghost cells, 623 including those outside the physical domain. This is also done in case 624 of periodic boundaries, meaning that the same global point may have 625 different coordinates in different local representations, which are 626 equivalent assuming a periodicity implied by the arguments to this 627 function, i.e., two points are equivalent if their difference is a 628 multiple of ``xmax-xmin`` in the x dimension, ``ymax-ymin`` in the y 629 dimension, and ``zmax-zmin`` in the z dimension. 630 631 Parameters 632 ---------- 633 xmin 634 The minimum global coordinate value in the x dimension. 635 xmax 636 The maximum global coordinate value in the x dimension. 637 ymin 638 The minimum global coordinate value in the y dimension. 639 ymax 640 The maximum global coordinate value in the y dimension. 641 zmin 642 The minimum global coordinate value in the z dimension. 643 zmax 644 The maximum global coordinate value in the z dimension. 645 646 See Also 647 -------- 648 setUniformCoordinatesExplicit, setUniformCoordinatesProduct 649 petsc.DMStagSetUniformCoordinates 650 651 """ 652 cdef PetscReal _xmin = asReal(xmin), _xmax = asReal(xmax) 653 cdef PetscReal _ymin = asReal(ymin), _ymax = asReal(ymax) 654 cdef PetscReal _zmin = asReal(zmin), _zmax = asReal(zmax) 655 CHKERR(DMStagSetUniformCoordinates(self.dm, _xmin, _xmax, _ymin, _ymax, _zmin, _zmax)) 656 657 def setCoordinateDMType(self, dmtype: DM.Type) -> None: 658 """Set the type to store coordinates. 659 660 Logically collective. 661 662 Parameters 663 ---------- 664 dmtype 665 The type to store coordinates. 666 667 See Also 668 -------- 669 petsc.DMStagSetCoordinateDMType 670 671 """ 672 cdef PetscDMType cval = NULL 673 dmtype = str2bytes(dmtype, &cval) 674 CHKERR(DMStagSetCoordinateDMType(self.dm, cval)) 675 676 # Location slot related functions 677 678 def getLocationSlot(self, loc: StencilLocation, c: int) -> int: 679 """Return index to use in accessing raw local arrays. 680 681 Not collective. 682 683 Parameters 684 ---------- 685 loc 686 Location relative to an element. 687 c 688 Component. 689 690 See Also 691 -------- 692 petsc.DMStagGetLocationSlot 693 694 """ 695 cdef PetscInt slot=0 696 cdef PetscInt comp=asInt(c) 697 cdef PetscDMStagStencilLocation sloc = asStagStencilLocation(loc) 698 CHKERR(DMStagGetLocationSlot(self.dm, sloc, comp, &slot)) 699 return toInt(slot) 700 701 def getProductCoordinateLocationSlot(self, loc: StencilLocation) -> None: 702 """Return slot for use with local product coordinate arrays. 703 704 Not collective. 705 706 Parameters 707 ---------- 708 loc 709 The grid location. 710 711 See Also 712 -------- 713 petsc.DMStagGetProductCoordinateLocationSlot 714 715 """ 716 cdef PetscInt slot=0 717 cdef PetscDMStagStencilLocation sloc = asStagStencilLocation(loc) 718 CHKERR(DMStagGetProductCoordinateLocationSlot(self.dm, sloc, &slot)) 719 return toInt(slot) 720 721 def getLocationDof(self, loc: StencilLocation) -> int: 722 """Return number of DOFs associated with a given point on the grid. 723 724 Not collective. 725 726 Parameters 727 ---------- 728 loc 729 The grid point. 730 731 See Also 732 -------- 733 petsc.DMStagGetLocationDOF 734 735 """ 736 cdef PetscInt dof=0 737 cdef PetscDMStagStencilLocation sloc = asStagStencilLocation(loc) 738 CHKERR(DMStagGetLocationDOF(self.dm, sloc, &dof)) 739 return toInt(dof) 740 741 # Random other functions 742 743 def migrateVec(self, Vec vec, DM dmTo, Vec vecTo) -> None: 744 """Transfer a vector between two ``DMStag`` objects. 745 746 Collective. 747 748 Currently only implemented to migrate global vectors to global vectors. 749 750 Parameters 751 ---------- 752 vec 753 The source vector. 754 dmTo 755 The compatible destination object. 756 vecTo 757 The destination vector. 758 759 See Also 760 -------- 761 petsc.DMStagMigrateVec 762 763 """ 764 CHKERR(DMStagMigrateVec(self.dm, vec.vec, dmTo.dm, vecTo.vec)) 765 766 def createCompatibleDMStag(self, dofs: tuple[int, ...]) -> DM: 767 """Create a compatible ``DMStag`` with different DOFs/stratum. 768 769 Collective. 770 771 Parameters 772 ---------- 773 dofs 774 The number of DOFs on the strata in the new `DMStag`. 775 776 See Also 777 -------- 778 petsc.DMStagCreateCompatibleDMStag 779 780 """ 781 cdef tuple gdofs = tuple(dofs) 782 cdef PetscInt dof0=1, dof1=0, dof2=0, dof3=0 783 asDofs(gdofs, &dof0, &dof1, &dof2, &dof3) 784 cdef PetscDM newda = NULL 785 CHKERR(DMStagCreateCompatibleDMStag(self.dm, dof0, dof1, dof2, dof3, &newda)) 786 cdef DM newdm = type(self)() 787 CHKERR(PetscCLEAR(newdm.obj)); newdm.dm = newda 788 return newdm 789 790 def VecSplitToDMDA( 791 self, 792 Vec vec, 793 loc: StencilLocation, 794 c: int) -> tuple[DMDA, Vec]: 795 """Return ``DMDA``, ``Vec`` from a subgrid of a ``DMStag``, its ``Vec``. 796 797 Collective. 798 799 If a ``c`` value of ``-k`` is provided, the first ``k`` DOFs for that 800 position are extracted, padding with zero values if needed. If a 801 non-negative value is provided, a single DOF is extracted. 802 803 Parameters 804 ---------- 805 vec 806 The ``Vec`` object. 807 loc 808 Which subgrid to extract. 809 c 810 Which component to extract. 811 812 See Also 813 -------- 814 petsc.DMStagVecSplitToDMDA 815 816 """ 817 cdef PetscInt pc = asInt(c) 818 cdef PetscDMStagStencilLocation sloc = asStagStencilLocation(loc) 819 cdef PetscDM pda = NULL 820 cdef PetscVec pdavec = NULL 821 CHKERR(DMStagVecSplitToDMDA(self.dm, vec.vec, sloc, pc, &pda, &pdavec)) 822 cdef DM da = DMDA() 823 CHKERR(PetscCLEAR(da.obj)); da.dm = pda 824 cdef Vec davec = Vec() 825 CHKERR(PetscCLEAR(davec.obj)); davec.vec = pdavec 826 return (da, davec) 827 828 def getVecArray(self, Vec vec) -> None: 829 """Not implemented.""" 830 raise NotImplementedError('getVecArray for DMStag not yet implemented in petsc4py') 831 832 def get1dCoordinatecArrays(self) -> None: 833 """Not implemented.""" 834 raise NotImplementedError('get1dCoordinatecArrays for DMStag not yet implemented in petsc4py') 835 836 property dim: 837 """The dimension.""" 838 def __get__(self) -> int: 839 return self.getDim() 840 841 property dofs: 842 """The number of DOFs associated with each stratum of the grid.""" 843 def __get__(self) -> tuple[int, ...]: 844 return self.getDof() 845 846 property entries_per_element: 847 """The number of entries per element in the local representation.""" 848 def __get__(self) -> int: 849 return self.getEntriesPerElement() 850 851 property global_sizes: 852 """Global element counts in each dimension.""" 853 def __get__(self) -> tuple[int, ...]: 854 return self.getGlobalSizes() 855 856 property local_sizes: 857 """Local elementwise sizes in each dimension.""" 858 def __get__(self) -> tuple[int, ...]: 859 return self.getLocalSizes() 860 861 property proc_sizes: 862 """The number of processes in each dimension in the global decomposition.""" 863 def __get__(self) -> tuple[int, ...]: 864 return self.getProcSizes() 865 866 property boundary_types: 867 """Boundary types in each dimension.""" 868 def __get__(self) -> tuple[str, ...]: 869 return self.getBoundaryTypes() 870 871 property stencil_type: 872 """Stencil type.""" 873 def __get__(self) -> str: 874 return self.getStencilType() 875 876 property stencil_width: 877 """Elementwise stencil width.""" 878 def __get__(self) -> int: 879 return self.getStencilWidth() 880 881 property corners: 882 """The lower left corner and size of local region in each dimension.""" 883 def __get__(self) -> tuple[tuple[int, ...], tuple[int, ...]]: 884 return self.getCorners() 885 886 property ghost_corners: 887 """The lower left corner and size of local region in each dimension.""" 888 def __get__(self) -> tuple[tuple[int, ...], tuple[int, ...]]: 889 return self.getGhostCorners() 890 891 892# -------------------------------------------------------------------- 893 894del DMStagStencilType 895del DMStagStencilLocation 896 897# -------------------------------------------------------------------- 898