1# -------------------------------------------------------------------- 2 3cdef extern from * nogil: 4 5 PetscErrorCode DMCompositeCreate(MPI_Comm, PetscDM*) 6 PetscErrorCode DMCompositeAddDM(PetscDM, PetscDM) 7 PetscErrorCode DMCompositeGetNumberDM(PetscDM, PetscInt*) 8 PetscErrorCode DMCompositeScatterArray(PetscDM, PetscVec, PetscVec*) 9 PetscErrorCode DMCompositeGatherArray(PetscDM, PetscInsertMode, PetscVec, PetscVec*) 10 PetscErrorCode DMCompositeGetEntriesArray(PetscDM, PetscDM*) 11 PetscErrorCode DMCompositeGetAccessArray(PetscDM, PetscVec, PetscInt, const PetscInt*, PetscVec*) 12 PetscErrorCode DMCompositeRestoreAccessArray(PetscDM, PetscVec, PetscInt, const PetscInt*, PetscVec*) 13 PetscErrorCode DMCompositeGetGlobalISs(PetscDM, PetscIS**) 14 PetscErrorCode DMCompositeGetLocalISs(PetscDM, PetscIS**) 15 PetscErrorCode DMCompositeGetISLocalToGlobalMappings(PetscDM, PetscLGMap**) 16 17cdef class _DMComposite_access: 18 cdef PetscDM dm 19 cdef PetscVec gvec 20 cdef PetscInt nlocs 21 cdef PetscInt *locs 22 cdef PetscVec *vecs 23 cdef object locs_mem 24 cdef object vecs_mem 25 cdef object access 26 27 def __cinit__(self, DM dm, Vec gvec, locs=None): 28 self.dm = dm.dm 29 CHKERR(PetscINCREF(<PetscObject*>&self.dm)) 30 self.gvec = gvec.vec 31 CHKERR(PetscINCREF(<PetscObject*>&self.gvec)) 32 if locs is None: 33 CHKERR(DMCompositeGetNumberDM(self.dm, &self.nlocs)) 34 locs = arange(0, <long>self.nlocs, 1) 35 self.locs_mem = iarray_i(locs, &self.nlocs, &self.locs) 36 self.vecs_mem = oarray_p(empty_p(self.nlocs), NULL, <void**>&self.vecs) 37 self.access = None 38 39 def __dealloc__(self): 40 CHKERR(DMDestroy(&self.dm)) 41 CHKERR(VecDestroy(&self.gvec)) 42 43 def __enter__(self): 44 cdef Py_ssize_t n = self.nlocs 45 CHKERR(DMCompositeGetAccessArray(self.dm, self.gvec, self.nlocs, self.locs, self.vecs)) 46 self.access = [ref_Vec(self.vecs[i]) for i from 0 <= i < n] 47 return tuple(self.access) 48 49 def __exit__(self, *exc): 50 cdef Py_ssize_t i, n = self.nlocs 51 for i from 0 <= i < n: (<Vec>self.access[i]).vec = NULL 52 CHKERR(DMCompositeRestoreAccessArray(self.dm, self.gvec, self.nlocs, self.locs, self.vecs)) 53 self.access = None 54