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