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