1# -------------------------------------------------------------------- 2 3cdef class DMComposite(DM): 4 """A DM object that is used to manage data for a collection of DMs.""" 5 6 def create(self, comm: Comm | None = None) -> Self: 7 """Create a composite object. 8 9 Collective. 10 11 Parameters 12 ---------- 13 comm 14 MPI communicator, defaults to `Sys.getDefaultComm`. 15 16 See Also 17 -------- 18 petsc.DMCompositeCreate 19 20 """ 21 cdef MPI_Comm ccomm = def_Comm(comm, PETSC_COMM_DEFAULT) 22 cdef PetscDM newdm = NULL 23 CHKERR(DMCompositeCreate(ccomm, &newdm)) 24 CHKERR(PetscCLEAR(self.obj)) 25 self.dm = newdm 26 return self 27 28 def addDM(self, DM dm, *args: DM) -> None: 29 """Add a DM vector to the composite. 30 31 Collective. 32 33 Parameters 34 ---------- 35 dm 36 The DM object. 37 *args 38 Additional DM objects. 39 40 See Also 41 -------- 42 petsc.DMCompositeAddDM 43 44 """ 45 CHKERR(DMCompositeAddDM(self.dm, dm.dm)) 46 cdef object item 47 for item in args: 48 dm = <DM?> item 49 CHKERR(DMCompositeAddDM(self.dm, dm.dm)) 50 51 def getNumber(self) -> int: 52 """Get number of sub-DMs contained in the composite. 53 54 Not collective. 55 56 See Also 57 -------- 58 petsc.DMCompositeGetNumberDM 59 60 """ 61 cdef PetscInt n = 0 62 CHKERR(DMCompositeGetNumberDM(self.dm, &n)) 63 return toInt(n) 64 getNumberDM = getNumber 65 66 def getEntries(self) -> list[DM]: 67 """Return sub-DMs contained in the composite. 68 69 Not collective. 70 71 See Also 72 -------- 73 petsc.DMCompositeGetEntriesArray 74 75 """ 76 cdef PetscInt i, n = 0 77 cdef PetscDM *cdms = NULL 78 CHKERR(DMCompositeGetNumberDM(self.dm, &n)) 79 cdef object unused = oarray_p(empty_p(<PetscInt>n), NULL, <void**>&cdms) 80 CHKERR(DMCompositeGetEntriesArray(self.dm, cdms)) 81 cdef DM entry = None 82 cdef list entries = [] 83 for i from 0 <= i < n: 84 entry = subtype_DM(cdms[i])() 85 entry.dm = cdms[i] 86 CHKERR(PetscINCREF(entry.obj)) 87 entries.append(entry) 88 return entries 89 90 def scatter(self, Vec gvec, lvecs: Sequence[Vec]) -> None: 91 """Scatter coupled global vector into split local vectors. 92 93 Collective. 94 95 Parameters 96 ---------- 97 gvec 98 The global vector. 99 lvecs 100 Array of local vectors. 101 102 See Also 103 -------- 104 gather, petsc.DMCompositeScatterArray 105 106 """ 107 cdef PetscInt i, n = 0 108 CHKERR(DMCompositeGetNumberDM(self.dm, &n)) 109 cdef PetscVec *clvecs = NULL 110 cdef object unused = oarray_p(empty_p(<PetscInt>n), NULL, <void**>&clvecs) 111 for i from 0 <= i < n: 112 clvecs[i] = (<Vec?>lvecs[<Py_ssize_t>i]).vec 113 CHKERR(DMCompositeScatterArray(self.dm, gvec.vec, clvecs)) 114 115 def gather(self, Vec gvec, imode: InsertModeSpec, lvecs: Sequence[Vec]) -> None: 116 """Gather split local vectors into a coupled global vector. 117 118 Collective. 119 120 Parameters 121 ---------- 122 gvec 123 The global vector. 124 imode 125 The insertion mode. 126 lvecs 127 The individual sequential vectors. 128 129 See Also 130 -------- 131 scatter, petsc.DMCompositeGatherArray 132 133 """ 134 cdef PetscInsertMode cimode = insertmode(imode) 135 cdef PetscInt i, n = 0 136 CHKERR(DMCompositeGetNumberDM(self.dm, &n)) 137 cdef PetscVec *clvecs = NULL 138 cdef object unused = oarray_p(empty_p(<PetscInt>n), NULL, <void**>&clvecs) 139 for i from 0 <= i < n: 140 clvecs[i] = (<Vec?>lvecs[<Py_ssize_t>i]).vec 141 CHKERR(DMCompositeGatherArray(self.dm, cimode, gvec.vec, clvecs)) 142 143 def getGlobalISs(self) -> list[IS]: 144 """Return the index sets for each composed object in the composite. 145 146 Collective. 147 148 These could be used to extract a subset of vector entries for a 149 "multi-physics" preconditioner. 150 151 Use `getLocalISs` for index sets in the packed local numbering, and 152 `getLGMaps` for to map local sub-DM (including ghost) indices to packed 153 global indices. 154 155 See Also 156 -------- 157 petsc.DMCompositeGetGlobalISs 158 159 """ 160 cdef PetscInt i, n = 0 161 cdef PetscIS *cis = NULL 162 CHKERR(DMCompositeGetNumberDM(self.dm, &n)) 163 CHKERR(DMCompositeGetGlobalISs(self.dm, &cis)) 164 cdef object isets = [ref_IS(cis[i]) for i from 0 <= i < n] 165 for i from 0 <= i < n: 166 CHKERR(ISDestroy(&cis[i])) 167 CHKERR(PetscFree(cis)) 168 return isets 169 170 def getLocalISs(self) -> list[IS]: 171 """Return index sets for each component of a composite local vector. 172 173 Not collective. 174 175 To get the composite global indices at all local points (including 176 ghosts), use `getLGMaps`. 177 178 To get index sets for pieces of the composite global vector, use 179 `getGlobalISs`. 180 181 See Also 182 -------- 183 petsc.DMCompositeGetLocalISs 184 185 """ 186 cdef PetscInt i, n = 0 187 cdef PetscIS *cis = NULL 188 CHKERR(DMCompositeGetNumberDM(self.dm, &n)) 189 CHKERR(DMCompositeGetLocalISs(self.dm, &cis)) 190 cdef object isets = [ref_IS(cis[i]) for i from 0 <= i < n] 191 for i from 0 <= i < n: 192 CHKERR(ISDestroy(&cis[i])) 193 CHKERR(PetscFree(cis)) 194 return isets 195 196 def getLGMaps(self) -> list[LGMap]: 197 """Return a local-to-global mapping for each DM in the composite. 198 199 Collective. 200 201 Note that this includes all the ghost points that individual ghosted 202 DMDA may have. 203 204 See Also 205 -------- 206 petsc.DMCompositeGetISLocalToGlobalMappings 207 208 """ 209 cdef PetscInt i, n = 0 210 cdef PetscLGMap *clgm = NULL 211 CHKERR(DMCompositeGetNumberDM(self.dm, &n)) 212 CHKERR(DMCompositeGetISLocalToGlobalMappings(self.dm, &clgm)) 213 cdef object lgms = [ref_LGMap(clgm[i]) for i from 0 <= i < n] 214 for i from 0 <= i < n: 215 CHKERR(ISLocalToGlobalMappingDestroy(&clgm[i])) 216 CHKERR(PetscFree(clgm)) 217 return lgms 218 219 def getAccess(self, Vec gvec, locs: Sequence[int] | None = None) -> Any: 220 """Get access to the individual vectors from the global vector. 221 222 Not collective. 223 224 Use via `with` context manager (PEP 343). 225 226 Parameters 227 ---------- 228 gvec 229 The global vector. 230 locs 231 Indices of vectors wanted, or `None` to get all vectors. 232 233 """ 234 return _DMComposite_access(self, gvec, locs) 235