xref: /petsc/src/dm/impls/composite/packm.c (revision 647798804180922dfb980ca29d2b49fa46af1416) 
1 #define PETSCDM_DLL
2 
3 #include "packimpl.h" /*I   "petscdm.h"   I*/
4 
5 #undef __FUNCT__
6 #define __FUNCT__ "DMGetMatrix_Composite_Nest"
7 static PetscErrorCode DMGetMatrix_Composite_Nest(DM dm,const MatType mtype,Mat *J)
8 {
9   const DM_Composite           *com = (DM_Composite*)dm->data;
10   const struct DMCompositeLink *rlink,*clink;
11   PetscErrorCode               ierr;
12   IS                           *isg;
13   Mat                          *submats;
14   PetscInt                     i,j,n;
15 
16   PetscFunctionBegin;
17   n = com->nDM + com->nredundant; /* Total number of entries */
18 
19   /* Explicit index sets are not required for MatCreateNest, but getting them here allows MatNest to do consistency
20    * checking and allows ISEqual to compare by identity instead of by contents. */
21   ierr = DMCompositeGetGlobalISs(dm,&isg);CHKERRQ(ierr);
22 
23   /* Get submatrices */
24   ierr = PetscMalloc(n*n*sizeof(Mat),&submats);CHKERRQ(ierr);
25   for (i=0,rlink=com->next; rlink; i++,rlink=rlink->next) {
26     for (j=0,clink=com->next; clink; j++,clink=clink->next) {
27       Mat sub = PETSC_NULL;
28       if (i == j) {
29         switch (rlink->type) {
30         case DMCOMPOSITE_ARRAY:
31           ierr = MatCreateMPIDense(((PetscObject)dm)->comm,rlink->n,clink->n,PETSC_DETERMINE,PETSC_DETERMINE,PETSC_NULL,&sub);CHKERRQ(ierr);
32           break;
33         case DMCOMPOSITE_DM:
34           ierr = DMGetMatrix(rlink->dm,PETSC_NULL,&sub);CHKERRQ(ierr);
35           break;
36         default: SETERRQ(((PetscObject)dm)->comm,PETSC_ERR_SUP,"unknown type");
37         }
38       } else if (com->FormCoupleLocations) {
39         SETERRQ(((PetscObject)dm)->comm,PETSC_ERR_SUP,"Cannot manage off-diagonal parts yet");
40       }
41       submats[i*n+j] = sub;
42     }
43   }
44 
45   ierr = MatCreateNest(((PetscObject)dm)->comm,n,isg,n,isg,submats,J);CHKERRQ(ierr);
46 
47   /* Disown references */
48   for (i=0; i<n; i++) {ierr = ISDestroy(isg[i]);CHKERRQ(ierr);}
49   ierr = PetscFree(isg);CHKERRQ(ierr);
50 
51   for (i=0; i<n*n; i++) {
52     if (submats[i]) {ierr = MatDestroy(submats[i]);CHKERRQ(ierr);}
53   }
54   ierr = PetscFree(submats);CHKERRQ(ierr);
55   PetscFunctionReturn(0);
56 }
57 
58 #undef __FUNCT__
59 #define __FUNCT__ "DMGetMatrix_Composite_AIJ"
60 static PetscErrorCode DMGetMatrix_Composite_AIJ(DM dm,const MatType mtype,Mat *J)
61 {
62   PetscErrorCode         ierr;
63   DM_Composite           *com = (DM_Composite*)dm->data;
64   struct DMCompositeLink *next = com->next;
65   PetscInt               m,*dnz,*onz,i,j,mA;
66   Mat                    Atmp;
67   PetscMPIInt            rank;
68   PetscBool              dense = PETSC_FALSE,prealloc_only = PETSC_FALSE;
69 
70   PetscFunctionBegin;
71   /* use global vector to determine layout needed for matrix */
72   m = com->n;
73 
74   ierr = MatCreate(((PetscObject)dm)->comm,J);CHKERRQ(ierr);
75   ierr = MatSetSizes(*J,m,m,PETSC_DETERMINE,PETSC_DETERMINE);CHKERRQ(ierr);
76   ierr = MatSetType(*J,mtype);CHKERRQ(ierr);
77 
78   /*
79      Extremely inefficient but will compute entire Jacobian for testing
80   */
81   ierr = PetscOptionsGetBool(((PetscObject)dm)->prefix,"-dmcomposite_dense_jacobian",&dense,PETSC_NULL);CHKERRQ(ierr);
82   if (dense) {
83     PetscInt    rstart,rend,*indices;
84     PetscScalar *values;
85 
86     mA    = com->N;
87     ierr = MatMPIAIJSetPreallocation(*J,mA,PETSC_NULL,mA-m,PETSC_NULL);CHKERRQ(ierr);
88     ierr = MatSeqAIJSetPreallocation(*J,mA,PETSC_NULL);CHKERRQ(ierr);
89 
90     ierr = MatGetOwnershipRange(*J,&rstart,&rend);CHKERRQ(ierr);
91     ierr = PetscMalloc2(mA,PetscScalar,&values,mA,PetscInt,&indices);CHKERRQ(ierr);
92     ierr = PetscMemzero(values,mA*sizeof(PetscScalar));CHKERRQ(ierr);
93     for (i=0; i<mA; i++) indices[i] = i;
94     for (i=rstart; i<rend; i++) {
95       ierr = MatSetValues(*J,1,&i,mA,indices,values,INSERT_VALUES);CHKERRQ(ierr);
96     }
97     ierr = PetscFree2(values,indices);CHKERRQ(ierr);
98     ierr = MatAssemblyBegin(*J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
99     ierr = MatAssemblyEnd(*J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
100     PetscFunctionReturn(0);
101   }
102 
103   ierr = MPI_Comm_rank(((PetscObject)dm)->comm,&rank);CHKERRQ(ierr);
104   ierr = MatPreallocateInitialize(((PetscObject)dm)->comm,m,m,dnz,onz);CHKERRQ(ierr);
105   /* loop over packed objects, handling one at at time */
106   next = com->next;
107   while (next) {
108     if (next->type == DMCOMPOSITE_ARRAY) {
109       if (rank == next->rank) {  /* zero the "little" block */
110         for (j=com->rstart+next->rstart; j<com->rstart+next->rstart+next->n; j++) {
111           for (i=com->rstart+next->rstart; i<com->rstart+next->rstart+next->n; i++) {
112             ierr = MatPreallocateSet(j,1,&i,dnz,onz);CHKERRQ(ierr);
113           }
114         }
115       }
116     } else if (next->type == DMCOMPOSITE_DM) {
117       PetscInt       nc,rstart,*ccols,maxnc;
118       const PetscInt *cols,*rstarts;
119       PetscMPIInt    proc;
120 
121       ierr = DMGetMatrix(next->dm,mtype,&Atmp);CHKERRQ(ierr);
122       ierr = MatGetOwnershipRange(Atmp,&rstart,PETSC_NULL);CHKERRQ(ierr);
123       ierr = MatGetOwnershipRanges(Atmp,&rstarts);CHKERRQ(ierr);
124       ierr = MatGetLocalSize(Atmp,&mA,PETSC_NULL);CHKERRQ(ierr);
125 
126       maxnc = 0;
127       for (i=0; i<mA; i++) {
128         ierr  = MatGetRow(Atmp,rstart+i,&nc,PETSC_NULL,PETSC_NULL);CHKERRQ(ierr);
129         ierr  = MatRestoreRow(Atmp,rstart+i,&nc,PETSC_NULL,PETSC_NULL);CHKERRQ(ierr);
130         maxnc = PetscMax(nc,maxnc);
131       }
132       ierr = PetscMalloc(maxnc*sizeof(PetscInt),&ccols);CHKERRQ(ierr);
133       for (i=0; i<mA; i++) {
134         ierr = MatGetRow(Atmp,rstart+i,&nc,&cols,PETSC_NULL);CHKERRQ(ierr);
135         /* remap the columns taking into how much they are shifted on each process */
136         for (j=0; j<nc; j++) {
137           proc = 0;
138           while (cols[j] >= rstarts[proc+1]) proc++;
139           ccols[j] = cols[j] + next->grstarts[proc] - rstarts[proc];
140         }
141         ierr = MatPreallocateSet(com->rstart+next->rstart+i,nc,ccols,dnz,onz);CHKERRQ(ierr);
142         ierr = MatRestoreRow(Atmp,rstart+i,&nc,&cols,PETSC_NULL);CHKERRQ(ierr);
143       }
144       ierr = PetscFree(ccols);CHKERRQ(ierr);
145       ierr = MatDestroy(Atmp);CHKERRQ(ierr);
146     } else {
147       SETERRQ(((PetscObject)dm)->comm,PETSC_ERR_SUP,"Cannot handle that object type yet");
148     }
149     next = next->next;
150   }
151   if (com->FormCoupleLocations) {
152     ierr = (*com->FormCoupleLocations)(dm,PETSC_NULL,dnz,onz,__rstart,__nrows,__start,__end);CHKERRQ(ierr);
153   }
154   ierr = MatMPIAIJSetPreallocation(*J,0,dnz,0,onz);CHKERRQ(ierr);
155   ierr = MatSeqAIJSetPreallocation(*J,0,dnz);CHKERRQ(ierr);
156   ierr = MatPreallocateFinalize(dnz,onz);CHKERRQ(ierr);
157 
158   ierr = PetscOptionsGetBool(((PetscObject)dm)->prefix,"-dmcomposite_preallocate_only",&prealloc_only,PETSC_NULL);CHKERRQ(ierr);
159   if (prealloc_only) PetscFunctionReturn(0);
160 
161   next = com->next;
162   while (next) {
163     if (next->type == DMCOMPOSITE_ARRAY) {
164       if (rank == next->rank) {
165         for (j=com->rstart+next->rstart; j<com->rstart+next->rstart+next->n; j++) {
166           for (i=com->rstart+next->rstart; i<com->rstart+next->rstart+next->n; i++) {
167             ierr = MatSetValue(*J,j,i,0.0,INSERT_VALUES);CHKERRQ(ierr);
168           }
169         }
170       }
171     } else if (next->type == DMCOMPOSITE_DM) {
172       PetscInt          nc,rstart,row,maxnc,*ccols;
173       const PetscInt    *cols,*rstarts;
174       const PetscScalar *values;
175       PetscMPIInt       proc;
176 
177       ierr = DMGetMatrix(next->dm,mtype,&Atmp);CHKERRQ(ierr);
178       ierr = MatGetOwnershipRange(Atmp,&rstart,PETSC_NULL);CHKERRQ(ierr);
179       ierr = MatGetOwnershipRanges(Atmp,&rstarts);CHKERRQ(ierr);
180       ierr = MatGetLocalSize(Atmp,&mA,PETSC_NULL);CHKERRQ(ierr);
181       maxnc = 0;
182       for (i=0; i<mA; i++) {
183         ierr  = MatGetRow(Atmp,rstart+i,&nc,PETSC_NULL,PETSC_NULL);CHKERRQ(ierr);
184         ierr  = MatRestoreRow(Atmp,rstart+i,&nc,PETSC_NULL,PETSC_NULL);CHKERRQ(ierr);
185         maxnc = PetscMax(nc,maxnc);
186       }
187       ierr = PetscMalloc(maxnc*sizeof(PetscInt),&ccols);CHKERRQ(ierr);
188       for (i=0; i<mA; i++) {
189         ierr = MatGetRow(Atmp,rstart+i,&nc,(const PetscInt **)&cols,&values);CHKERRQ(ierr);
190         for (j=0; j<nc; j++) {
191           proc = 0;
192           while (cols[j] >= rstarts[proc+1]) proc++;
193           ccols[j] = cols[j] + next->grstarts[proc] - rstarts[proc];
194         }
195         row  = com->rstart+next->rstart+i;
196         ierr = MatSetValues(*J,1,&row,nc,ccols,values,INSERT_VALUES);CHKERRQ(ierr);
197         ierr = MatRestoreRow(Atmp,rstart+i,&nc,(const PetscInt **)&cols,&values);CHKERRQ(ierr);
198       }
199       ierr = PetscFree(ccols);CHKERRQ(ierr);
200       ierr = MatDestroy(Atmp);CHKERRQ(ierr);
201     } else {
202       SETERRQ(((PetscObject)dm)->comm,PETSC_ERR_SUP,"Cannot handle that object type yet");
203     }
204     next = next->next;
205   }
206   if (com->FormCoupleLocations) {
207     PetscInt __rstart;
208     ierr = MatGetOwnershipRange(*J,&__rstart,PETSC_NULL);CHKERRQ(ierr);
209     ierr = (*com->FormCoupleLocations)(dm,*J,PETSC_NULL,PETSC_NULL,__rstart,0,0,0);CHKERRQ(ierr);
210   }
211   ierr = MatAssemblyBegin(*J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
212   ierr = MatAssemblyEnd(*J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
213   PetscFunctionReturn(0);
214 }
215 
216 #undef __FUNCT__
217 #define __FUNCT__ "DMGetMatrix_Composite"
218 PetscErrorCode  DMGetMatrix_Composite(DM dm,const MatType mtype,Mat *J)
219 {
220   PetscErrorCode         ierr;
221   PetscBool              usenest;
222   ISLocalToGlobalMapping ltogmap,ltogmapb;
223 
224   PetscFunctionBegin;
225   ierr = PetscStrcmp(mtype,MATNEST,&usenest);CHKERRQ(ierr);
226   if (usenest) {
227     ierr = DMGetMatrix_Composite_Nest(dm,mtype,J);CHKERRQ(ierr);
228   } else {
229     ierr = DMGetMatrix_Composite_AIJ(dm,mtype?mtype:MATAIJ,J);CHKERRQ(ierr);
230   }
231 
232   ierr = DMGetLocalToGlobalMapping(dm,&ltogmap);CHKERRQ(ierr);
233   ierr = DMGetLocalToGlobalMappingBlock(dm,&ltogmapb);CHKERRQ(ierr);
234   ierr = MatSetLocalToGlobalMapping(*J,ltogmap,ltogmap);CHKERRQ(ierr);
235   ierr = MatSetLocalToGlobalMappingBlock(*J,ltogmapb,ltogmapb);CHKERRQ(ierr);
236   PetscFunctionReturn(0);
237 }
238