xref: /petsc/src/snes/utils/dmlocalsnes.c (revision b41ce5d507ea9a58bfa83cf403107a702e77a67d)
1 #include <petsc/private/dmimpl.h>
2 #include <petsc/private/snesimpl.h>   /*I "petscsnes.h" I*/
3 
4 typedef struct {
5   PetscErrorCode (*residuallocal)(DM,Vec,Vec,void*);
6   PetscErrorCode (*jacobianlocal)(DM,Vec,Mat,Mat,void*);
7   PetscErrorCode (*boundarylocal)(DM,Vec,void*);
8   void *residuallocalctx;
9   void *jacobianlocalctx;
10   void *boundarylocalctx;
11 } DMSNES_Local;
12 
13 static PetscErrorCode DMSNESDestroy_DMLocal(DMSNES sdm)
14 {
15   PetscErrorCode ierr;
16 
17   PetscFunctionBegin;
18   ierr = PetscFree(sdm->data);CHKERRQ(ierr);
19   PetscFunctionReturn(0);
20 }
21 
22 static PetscErrorCode DMSNESDuplicate_DMLocal(DMSNES oldsdm,DMSNES sdm)
23 {
24   PetscErrorCode ierr;
25 
26   PetscFunctionBegin;
27   ierr = PetscNewLog(sdm,(DMSNES_Local**)&sdm->data);CHKERRQ(ierr);
28   if (oldsdm->data) {
29     ierr = PetscMemcpy(sdm->data,oldsdm->data,sizeof(DMSNES_Local));CHKERRQ(ierr);
30   }
31   PetscFunctionReturn(0);
32 }
33 
34 static PetscErrorCode DMLocalSNESGetContext(DM dm,DMSNES sdm,DMSNES_Local **dmlocalsnes)
35 {
36   PetscErrorCode ierr;
37 
38   PetscFunctionBegin;
39   *dmlocalsnes = NULL;
40   if (!sdm->data) {
41     ierr = PetscNewLog(dm,(DMSNES_Local**)&sdm->data);CHKERRQ(ierr);
42 
43     sdm->ops->destroy   = DMSNESDestroy_DMLocal;
44     sdm->ops->duplicate = DMSNESDuplicate_DMLocal;
45   }
46   *dmlocalsnes = (DMSNES_Local*)sdm->data;
47   PetscFunctionReturn(0);
48 }
49 
50 static PetscErrorCode SNESComputeFunction_DMLocal(SNES snes,Vec X,Vec F,void *ctx)
51 {
52   PetscErrorCode ierr;
53   DM             dm;
54   DMSNES_Local   *dmlocalsnes = (DMSNES_Local*)ctx;
55   Vec            Xloc,Floc;
56 
57   PetscFunctionBegin;
58   PetscValidHeaderSpecific(snes,SNES_CLASSID,1);
59   PetscValidHeaderSpecific(X,VEC_CLASSID,2);
60   PetscValidHeaderSpecific(F,VEC_CLASSID,3);
61   ierr = SNESGetDM(snes,&dm);CHKERRQ(ierr);
62   ierr = DMGetLocalVector(dm,&Xloc);CHKERRQ(ierr);
63   ierr = DMGetLocalVector(dm,&Floc);CHKERRQ(ierr);
64   ierr = VecZeroEntries(Xloc);CHKERRQ(ierr);
65   if (dmlocalsnes->boundarylocal) {ierr = (*dmlocalsnes->boundarylocal)(dm,Xloc,dmlocalsnes->boundarylocalctx);CHKERRQ(ierr);}
66   ierr = VecZeroEntries(Floc);CHKERRQ(ierr);
67   ierr = DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
68   ierr = DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
69   CHKMEMQ;
70   ierr = (*dmlocalsnes->residuallocal)(dm,Xloc,Floc,dmlocalsnes->residuallocalctx);CHKERRQ(ierr);
71   CHKMEMQ;
72   ierr = VecZeroEntries(F);CHKERRQ(ierr);
73   ierr = DMLocalToGlobalBegin(dm,Floc,ADD_VALUES,F);CHKERRQ(ierr);
74   ierr = DMLocalToGlobalEnd(dm,Floc,ADD_VALUES,F);CHKERRQ(ierr);
75   ierr = DMRestoreLocalVector(dm,&Floc);CHKERRQ(ierr);
76   ierr = DMRestoreLocalVector(dm,&Xloc);CHKERRQ(ierr);
77   PetscFunctionReturn(0);
78 }
79 
80 static PetscErrorCode SNESComputeJacobian_DMLocal(SNES snes,Vec X,Mat A,Mat B,void *ctx)
81 {
82   PetscErrorCode ierr;
83   DM             dm;
84   DMSNES_Local   *dmlocalsnes = (DMSNES_Local*)ctx;
85   Vec            Xloc;
86 
87   PetscFunctionBegin;
88   ierr = SNESGetDM(snes,&dm);CHKERRQ(ierr);
89   if (dmlocalsnes->jacobianlocal) {
90     ierr = DMGetLocalVector(dm,&Xloc);CHKERRQ(ierr);
91     ierr = VecZeroEntries(Xloc);CHKERRQ(ierr);
92     if (dmlocalsnes->boundarylocal) {ierr = (*dmlocalsnes->boundarylocal)(dm,Xloc,dmlocalsnes->boundarylocalctx);CHKERRQ(ierr);}
93     ierr = DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
94     ierr = DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
95     CHKMEMQ;
96     ierr = (*dmlocalsnes->jacobianlocal)(dm,Xloc,A,B,dmlocalsnes->jacobianlocalctx);CHKERRQ(ierr);
97     CHKMEMQ;
98     ierr = DMRestoreLocalVector(dm,&Xloc);CHKERRQ(ierr);
99   } else {
100     MatFDColoring fdcoloring;
101     ierr = PetscObjectQuery((PetscObject)dm,"DMDASNES_FDCOLORING",(PetscObject*)&fdcoloring);CHKERRQ(ierr);
102     if (!fdcoloring) {
103       ISColoring coloring;
104 
105       ierr = DMCreateColoring(dm,dm->coloringtype,&coloring);CHKERRQ(ierr);
106       ierr = MatFDColoringCreate(B,coloring,&fdcoloring);CHKERRQ(ierr);
107       ierr = ISColoringDestroy(&coloring);CHKERRQ(ierr);
108       switch (dm->coloringtype) {
109       case IS_COLORING_GLOBAL:
110         ierr = MatFDColoringSetFunction(fdcoloring,(PetscErrorCode (*)(void))SNESComputeFunction_DMLocal,dmlocalsnes);CHKERRQ(ierr);
111         break;
112       default: SETERRQ1(PetscObjectComm((PetscObject)snes),PETSC_ERR_SUP,"No support for coloring type '%s'",ISColoringTypes[dm->coloringtype]);
113       }
114       ierr = PetscObjectSetOptionsPrefix((PetscObject)fdcoloring,((PetscObject)dm)->prefix);CHKERRQ(ierr);
115       ierr = MatFDColoringSetFromOptions(fdcoloring);CHKERRQ(ierr);
116       ierr = MatFDColoringSetUp(B,coloring,fdcoloring);CHKERRQ(ierr);
117       ierr = PetscObjectCompose((PetscObject)dm,"DMDASNES_FDCOLORING",(PetscObject)fdcoloring);CHKERRQ(ierr);
118       ierr = PetscObjectDereference((PetscObject)fdcoloring);CHKERRQ(ierr);
119 
120       /* The following breaks an ugly reference counting loop that deserves a paragraph. MatFDColoringApply() will call
121        * VecDuplicate() with the state Vec and store inside the MatFDColoring. This Vec will duplicate the Vec, but the
122        * MatFDColoring is composed with the DM. We dereference the DM here so that the reference count will eventually
123        * drop to 0. Note the code in DMDestroy() that exits early for a negative reference count. That code path will be
124        * taken when the PetscObjectList for the Vec inside MatFDColoring is destroyed.
125        */
126       ierr = PetscObjectDereference((PetscObject)dm);CHKERRQ(ierr);
127     }
128     ierr  = MatFDColoringApply(B,fdcoloring,X,snes);CHKERRQ(ierr);
129   }
130   /* This will be redundant if the user called both, but it's too common to forget. */
131   if (A != B) {
132     ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
133     ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
134   }
135   PetscFunctionReturn(0);
136 }
137 
138 /*@C
139    DMSNESSetFunctionLocal - set a local residual evaluation function. This function is called with local vector
140       containing the local vector information PLUS ghost point information. It should compute a result for all local
141       elements and DMSNES will automatically accumulate the overlapping values.
142 
143    Logically Collective
144 
145    Input Arguments:
146 +  dm - DM to associate callback with
147 .  func - local residual evaluation
148 -  ctx - optional context for local residual evaluation
149 
150    Level: beginner
151 
152 .seealso: DMSNESSetFunction(), DMDASNESSetJacobianLocal(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
153 @*/
154 PetscErrorCode DMSNESSetFunctionLocal(DM dm,PetscErrorCode (*func)(DM,Vec,Vec,void*),void *ctx)
155 {
156   PetscErrorCode ierr;
157   DMSNES         sdm;
158   DMSNES_Local   *dmlocalsnes;
159 
160   PetscFunctionBegin;
161   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
162   ierr = DMGetDMSNESWrite(dm,&sdm);CHKERRQ(ierr);
163   ierr = DMLocalSNESGetContext(dm,sdm,&dmlocalsnes);CHKERRQ(ierr);
164 
165   dmlocalsnes->residuallocal    = func;
166   dmlocalsnes->residuallocalctx = ctx;
167 
168   ierr = DMSNESSetFunction(dm,SNESComputeFunction_DMLocal,dmlocalsnes);CHKERRQ(ierr);
169   if (!sdm->ops->computejacobian) {  /* Call us for the Jacobian too, can be overridden by the user. */
170     ierr = DMSNESSetJacobian(dm,SNESComputeJacobian_DMLocal,dmlocalsnes);CHKERRQ(ierr);
171   }
172   PetscFunctionReturn(0);
173 }
174 
175 /*@C
176    DMSNESSetBoundaryLocal - set a local boundary value function. This function is called with local vector
177       containing the local vector information PLUS ghost point information. It should insert values into the local
178       vector that do not come from the global vector, such as essential boundary condition data.
179 
180    Logically Collective
181 
182    Input Arguments:
183 +  dm - DM to associate callback with
184 .  func - local boundary value evaluation
185 -  ctx - optional context for local boundary value evaluation
186 
187    Level: intermediate
188 
189 .seealso: DMSNESSetFunctionLocal(), DMDASNESSetJacobianLocal()
190 @*/
191 PetscErrorCode DMSNESSetBoundaryLocal(DM dm,PetscErrorCode (*func)(DM,Vec,void*),void *ctx)
192 {
193   PetscErrorCode ierr;
194   DMSNES         sdm;
195   DMSNES_Local   *dmlocalsnes;
196 
197   PetscFunctionBegin;
198   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
199   ierr = DMGetDMSNESWrite(dm,&sdm);CHKERRQ(ierr);
200   ierr = DMLocalSNESGetContext(dm,sdm,&dmlocalsnes);CHKERRQ(ierr);
201 
202   dmlocalsnes->boundarylocal    = func;
203   dmlocalsnes->boundarylocalctx = ctx;
204 
205   PetscFunctionReturn(0);
206 }
207 
208 /*@C
209    DMSNESSetJacobianLocal - set a local Jacobian evaluation function
210 
211    Logically Collective
212 
213    Input Arguments:
214 +  dm - DM to associate callback with
215 .  func - local Jacobian evaluation
216 -  ctx - optional context for local Jacobian evaluation
217 
218    Level: beginner
219 
220 .seealso: DMSNESSetJacobian(), DMDASNESSetJacobian(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
221 @*/
222 PetscErrorCode DMSNESSetJacobianLocal(DM dm,PetscErrorCode (*func)(DM,Vec,Mat,Mat,void*),void *ctx)
223 {
224   PetscErrorCode ierr;
225   DMSNES         sdm;
226   DMSNES_Local   *dmlocalsnes;
227 
228   PetscFunctionBegin;
229   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
230   ierr = DMGetDMSNESWrite(dm,&sdm);CHKERRQ(ierr);
231   ierr = DMLocalSNESGetContext(dm,sdm,&dmlocalsnes);CHKERRQ(ierr);
232 
233   dmlocalsnes->jacobianlocal    = func;
234   dmlocalsnes->jacobianlocalctx = ctx;
235 
236   ierr = DMSNESSetJacobian(dm,SNESComputeJacobian_DMLocal,dmlocalsnes);CHKERRQ(ierr);
237   PetscFunctionReturn(0);
238 }
239 
240