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