xref: /petsc/src/dm/impls/swarm/tutorials/ex1.c (revision bef158480efac06de457f7a665168877ab3c2fd7)
1 static char help[] = "Example program demonstrating projection between particle and finite element spaces\n\n";
2 
3 #include "petscdmplex.h"
4 #include "petscds.h"
5 #include "petscdmswarm.h"
6 #include "petscksp.h"
7 
8 static PetscErrorCode crd_func(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf_dummy, PetscScalar *u, void *actx)
9 {
10   PetscInt i;
11   PetscFunctionBeginUser;
12   for (i = 0; i < dim; ++i) u[i] = x[i];
13   PetscFunctionReturn(0);
14 }
15 
16 int main(int argc, char **argv)
17 {
18   DM               dm, crddm, sw;
19   PetscFE          fe;
20   KSP              ksp;
21   Mat              M_p, M;
22   Vec              f, rho, rhs, crd_vec;
23   PetscInt         dim = 2, Nc = 1, timestep = 0, N, i, idx[3];
24   PetscInt         Np = 10, p, field = 0, zero = 0, bs, cells[] = {40, 20, 20};
25   PetscReal        time = 0.0,  norm;
26   PetscBool        removePoints = PETSC_TRUE, flg;
27   const PetscReal  *xx, *vv;
28   PetscReal        *wq, *coords, lo[] = {-1,-1,-1}, hi[]={1,1,1}, h[3];
29   PetscDataType    dtype;
30   PetscBool        interpolate = PETSC_TRUE;
31   PetscErrorCode   ierr;
32   PetscErrorCode   (*initu[1])(PetscInt, PetscReal, const PetscReal [], PetscInt, PetscScalar [], void *);
33 
34   ierr = PetscInitialize(&argc, &argv, NULL,help);if (ierr) return ierr;
35   /* get options */
36   ierr = PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Options for Fokker-Plank collision operator", "none");CHKERRQ(ierr);
37   ierr = PetscOptionsRangeInt("-dim", "dim (2 or 3)", "ex1.c", dim, &dim, NULL,2,3);CHKERRQ(ierr);
38   i = dim;  ierr = PetscOptionsIntArray("-dm_plex_box_faces", "Number of cells of grid", "ex1.c", cells, &i, &flg);CHKERRQ(ierr);
39   i = dim;  ierr = PetscOptionsRealArray("-dm_plex_box_lower", "Low corner of grid", "ex1.c", lo, &i, &flg);CHKERRQ(ierr);
40   i = dim;  ierr = PetscOptionsRealArray("-dm_plex_box_upper", "High corner of grid", "ex1.c", hi, &i, &flg);CHKERRQ(ierr);
41   for (i=0;i<dim;i++) {
42     h[i] = (hi[i] - lo[i])/cells[i];
43     ierr = PetscPrintf(PETSC_COMM_SELF," lo = %g hi = %g n = %D h = %g\n",lo[i],hi[i],cells[i],h[i]);CHKERRQ(ierr);
44   }
45   ierr = PetscOptionsEnd();CHKERRQ(ierr);
46   /* Create a mesh */
47   ierr = DMPlexCreateBoxMesh(PETSC_COMM_SELF, dim, PETSC_FALSE, NULL, NULL, NULL, NULL, interpolate, &dm);CHKERRQ(ierr);
48   ierr = PetscObjectSetName((PetscObject)dm, "Potential Grid");CHKERRQ(ierr);
49   ierr = DMViewFromOptions(dm, NULL, "-dm_view");CHKERRQ(ierr);
50   ierr = PetscFECreateDefault(PETSC_COMM_SELF, dim, Nc, PETSC_FALSE, "", PETSC_DECIDE, &fe);CHKERRQ(ierr);
51   ierr = PetscFESetFromOptions(fe);CHKERRQ(ierr);
52   ierr = PetscObjectSetName((PetscObject)fe, "fe");CHKERRQ(ierr);
53   ierr = DMSetField(dm, field, NULL, (PetscObject)fe);CHKERRQ(ierr);
54   ierr = DMCreateDS(dm);CHKERRQ(ierr);
55   ierr = PetscFEDestroy(&fe);CHKERRQ(ierr);
56   /* Create particle swarm */
57   ierr = DMCreate(PETSC_COMM_SELF, &sw);CHKERRQ(ierr);
58   ierr = DMSetType(sw, DMSWARM);CHKERRQ(ierr);
59   ierr = DMSetDimension(sw, dim);CHKERRQ(ierr);
60   ierr = DMSwarmSetType(sw, DMSWARM_PIC);CHKERRQ(ierr);
61   ierr = DMSwarmSetCellDM(sw, dm);CHKERRQ(ierr);
62   ierr = DMSwarmRegisterPetscDatatypeField(sw, "w_q", Nc, PETSC_SCALAR);CHKERRQ(ierr);
63   ierr = DMSwarmFinalizeFieldRegister(sw);CHKERRQ(ierr);
64   ierr = DMSwarmSetLocalSizes(sw, Np, zero);CHKERRQ(ierr);
65   ierr = DMSetFromOptions(sw);CHKERRQ(ierr);
66   ierr = DMSwarmGetField(sw, "w_q", &bs, &dtype, (void**)&wq);CHKERRQ(ierr);
67   ierr = DMSwarmGetField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void**)&coords);CHKERRQ(ierr);
68   for (p=0;p<Np;p++) {
69     coords[p*2+0]  = -PetscCosReal((PetscReal)(p+1)/(PetscReal)(Np+1) * PETSC_PI);
70     coords[p*2+1] =   PetscSinReal((PetscReal)(p+1)/(PetscReal)(Np+1) * PETSC_PI);
71     wq[p]          = 1.0;
72   }
73   ierr = DMSwarmRestoreField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void**)&coords);CHKERRQ(ierr);
74   ierr = DMSwarmRestoreField(sw, "w_q", &bs, &dtype, (void**)&wq);CHKERRQ(ierr);
75   ierr = DMSwarmMigrate(sw, removePoints);CHKERRQ(ierr);
76   ierr = PetscObjectSetName((PetscObject)sw, "Particle Grid");CHKERRQ(ierr);
77   ierr = DMViewFromOptions(sw, NULL, "-swarm_view");CHKERRQ(ierr);
78   /* Project particles to field */
79   /* This gives M f = \int_\Omega \phi f, which looks like a rhs for a PDE */
80   ierr = DMCreateMassMatrix(sw, dm, &M_p);CHKERRQ(ierr);
81   ierr = DMCreateGlobalVector(dm, &rho);CHKERRQ(ierr);
82   ierr = PetscObjectSetName((PetscObject)rho, "rho");CHKERRQ(ierr);
83   ierr = DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f);CHKERRQ(ierr);
84   ierr = PetscObjectSetName((PetscObject)f, "weights");CHKERRQ(ierr);
85   ierr = MatMultTranspose(M_p, f, rho);CHKERRQ(ierr);
86 
87   /* Visualize mesh field */
88   ierr = DMSetOutputSequenceNumber(dm, timestep, time);CHKERRQ(ierr);
89   ierr = VecViewFromOptions(rho, NULL, "-rho_view");CHKERRQ(ierr);
90 
91   /* create coordinate DM */
92   ierr = DMClone(dm, &crddm);CHKERRQ(ierr);
93   ierr = PetscFECreateDefault(PETSC_COMM_SELF, dim, dim, PETSC_FALSE, "", PETSC_DECIDE, &fe);CHKERRQ(ierr);
94   ierr = PetscFESetFromOptions(fe);CHKERRQ(ierr);
95   ierr = DMSetField(crddm, field, NULL, (PetscObject)fe);CHKERRQ(ierr);
96   ierr = DMCreateDS(crddm);CHKERRQ(ierr);
97   ierr = PetscFEDestroy(&fe);CHKERRQ(ierr);
98   /* project coordiantes to vertices */
99   ierr = DMCreateGlobalVector(crddm, &crd_vec);CHKERRQ(ierr);
100   initu[0] = crd_func;
101   ierr = DMProjectFunction(crddm, 0.0, initu, NULL, INSERT_ALL_VALUES, crd_vec);CHKERRQ(ierr);
102   ierr = VecViewFromOptions(crd_vec, NULL, "-coord_view");CHKERRQ(ierr);
103   /* iterate over mesh data and get indices */
104   ierr = VecGetArrayRead(crd_vec,&xx);CHKERRQ(ierr);
105   ierr = VecGetArrayRead(rho,&vv);CHKERRQ(ierr);
106   ierr = VecGetLocalSize(rho,&N);CHKERRQ(ierr);
107   for (p=0;p<N;p++) {
108     for (i=0;i<dim;i++) idx[i] = (PetscInt)((xx[p*dim+i] - lo[i])/h[i] + 1.e-8);
109     ierr = PetscPrintf(PETSC_COMM_SELF,"(%D,%D) = %g\n",idx[0],idx[1],vv[p]);CHKERRQ(ierr);
110     /* access grid data here */
111   }
112   ierr = VecRestoreArrayRead(crd_vec,&xx);CHKERRQ(ierr);
113   ierr = VecRestoreArrayRead(rho,&vv);CHKERRQ(ierr);
114   ierr = VecDestroy(&crd_vec);CHKERRQ(ierr);
115   /* Project field to particles */
116   /*   This gives f_p = M_p^+ M f */
117   ierr = DMCreateMassMatrix(dm, dm, &M);CHKERRQ(ierr);
118   ierr = DMCreateGlobalVector(dm, &rhs);CHKERRQ(ierr);
119   if (0) {
120     ierr = MatMult(M, rho, rhs);CHKERRQ(ierr);  /* this is what you would do for an FE solve */
121   } else {
122     ierr = VecCopy(rho, rhs);CHKERRQ(ierr); /* Indentity: M^1 M rho */
123   }
124   ierr = KSPCreate(PETSC_COMM_WORLD, &ksp);CHKERRQ(ierr);
125   ierr = KSPSetOptionsPrefix(ksp, "ftop_");CHKERRQ(ierr);
126   ierr = KSPSetFromOptions(ksp);CHKERRQ(ierr);
127   ierr = KSPSetOperators(ksp, M_p, M_p);CHKERRQ(ierr);
128   ierr = KSPSolveTranspose(ksp, rhs, f);CHKERRQ(ierr);
129   ierr = KSPDestroy(&ksp);CHKERRQ(ierr);
130   ierr = VecDestroy(&rhs);CHKERRQ(ierr);
131 
132   /* Visualize particle field */
133   ierr = DMSetOutputSequenceNumber(sw, timestep, time);CHKERRQ(ierr);
134   ierr = VecViewFromOptions(f, NULL, "-weights_view");CHKERRQ(ierr);
135   ierr = VecNorm(f,NORM_1,&norm);CHKERRQ(ierr);
136   ierr = PetscPrintf(PETSC_COMM_SELF,"Total number density = %g\n", norm);CHKERRQ(ierr);
137   /* Cleanup */
138   ierr = DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f);CHKERRQ(ierr);
139   ierr = MatDestroy(&M);CHKERRQ(ierr);
140   ierr = MatDestroy(&M_p);CHKERRQ(ierr);
141   ierr = VecDestroy(&rho);CHKERRQ(ierr);
142   ierr = DMDestroy(&sw);CHKERRQ(ierr);
143   ierr = DMDestroy(&dm);CHKERRQ(ierr);
144   ierr = DMDestroy(&crddm);CHKERRQ(ierr);
145   ierr = DMDestroy(&dm);CHKERRQ(ierr);
146   ierr = PetscFinalize();
147   /* PetscFunctionReturn(0); */
148   return ierr;
149 }
150 
151 /*TEST
152 
153   build:
154     requires: !complex
155 
156   test:
157     suffix: 0
158     requires: double
159     args: -dm_plex_box_faces 4,2 -dm_plex_box_lower -2.0,0.0 -dm_plex_box_upper 2.0,2.0 -petscspace_degree 2 -ftop_ksp_type lsqr -ftop_pc_type none -dm_view -swarm_view
160     filter: grep -v DM_ | grep -v atomic
161     filter_output: grep -v atomic
162 
163 TEST*/
164