xref: /petsc/src/ts/tests/ex35.c (revision 377c486fdcd144dcf0fa6b5d4b628c68659cbe7a)
1 static char help[] = "Test of Colorized Scatter Plot.\n";
2 
3 #include "petscdraw.h"
4 #include "petscvec.h"
5 #include "petscis.h"
6 
7 typedef struct {
8   PetscInt  Np;          /* total number of particles */
9   PetscInt  dim;
10   PetscInt  dim_inp;
11 } AppCtx;
12 
13 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
14 {
15   PetscErrorCode ierr;
16 
17   PetscFunctionBeginUser;
18   options->dim  = 2;
19   options->dim_inp = 2;
20   options->Np   = 100;
21 
22   ierr = PetscOptionsBegin(comm, "", "Test of colorized scatter plot", "");CHKERRQ(ierr);
23   ierr = PetscOptionsInt("-Np", "Number of particles", "ex35.c", options->Np, &options->Np, PETSC_NULL);CHKERRQ(ierr);
24   ierr = PetscOptionsInt("-dim", "Number of dimensions", "ex35.c", options->dim_inp, &options->dim_inp, PETSC_NULL);CHKERRQ(ierr);
25   ierr = PetscOptionsEnd();CHKERRQ(ierr);
26   PetscFunctionReturn(0);
27 }
28 
29 /*
30   ref: http://www.mimirgames.com/articles/programming/approximations-of-the-inverse-error-function/
31 */
32 PetscReal erfinv(PetscReal x)
33 {
34   PetscReal      *ck, r = 0.;
35   PetscInt       k, m, maxIter=100;
36   PetscErrorCode ierr;
37 
38   ierr = PetscCalloc1(maxIter,&ck);CHKERRQ(ierr);
39   ck[0] = 1;
40   r = ck[0]*((PetscSqrtReal(PETSC_PI)/2.)*x);
41   for (k = 1; k < maxIter; ++k){
42     for (m = 0; m <= k-1; ++m){
43       PetscReal denom = (m+1.)*(2.*m+1.);
44       ck[k] += (ck[m]*ck[k-1-m])/denom;
45     }
46     PetscReal temp = 2.*k+1.;
47     r += (ck[k]/temp)*PetscPowReal((PetscSqrtReal(PETSC_PI)/2.)*x,2.*k+1.);
48   }
49   ierr = PetscFree(ck);
50   return r;
51 }
52 
53 int main(int argc, char **argv)
54 {
55   PetscInt          p, dim, Np;
56   PetscScalar       *randVecNums;
57   PetscReal         speed, value, *x, *v;
58   PetscRandom       rngx, rng1, rng2;
59   Vec               randVec, subvecvx, subvecvy;
60   IS                isvx, isvy;
61   AppCtx            user;
62   PetscDrawAxis     axis;
63   PetscDraw         positionDraw;
64   PetscDrawSP       positionDrawSP;
65   MPI_Comm          comm;
66   PetscErrorCode    ierr;
67 
68   ierr = PetscInitialize(&argc,&argv,NULL,help);if (ierr) return ierr;
69   comm = PETSC_COMM_WORLD;
70   ierr = ProcessOptions(comm, &user);CHKERRQ(ierr);
71 
72   Np = user.Np;
73   dim = user.dim;
74 
75   ierr = PetscMalloc2(Np*dim, &x, Np*dim, &v);CHKERRQ(ierr);
76 
77   ierr = PetscRandomCreate(comm, &rngx);CHKERRQ(ierr);
78   ierr = PetscRandomSetInterval(rngx, 0., 1.);CHKERRQ(ierr);
79   ierr = PetscRandomSetFromOptions(rngx);CHKERRQ(ierr);
80   ierr = PetscRandomSetSeed(rngx, 1034);CHKERRQ(ierr);
81   ierr = PetscRandomSeed(rngx);CHKERRQ(ierr);
82 
83   ierr = PetscRandomCreate(comm, &rng1);CHKERRQ(ierr);
84   ierr = PetscRandomSetInterval(rng1, 0., 1.);CHKERRQ(ierr);
85   ierr = PetscRandomSetFromOptions(rng1);CHKERRQ(ierr);
86   ierr = PetscRandomSetSeed(rng1, 3084);CHKERRQ(ierr);
87   ierr = PetscRandomSeed(rng1);CHKERRQ(ierr);
88 
89   ierr = PetscRandomCreate(comm, &rng2);CHKERRQ(ierr);
90   ierr = PetscRandomSetInterval(rng2, 0., 1.);CHKERRQ(ierr);
91   ierr = PetscRandomSetFromOptions(rng2);CHKERRQ(ierr);
92   ierr = PetscRandomSetSeed(rng2, 2397);CHKERRQ(ierr);
93   ierr = PetscRandomSeed(rng2);CHKERRQ(ierr);
94 
95   /* Set particle positions and velocities */
96   if (user.dim_inp == 1) {
97     for (p = 0; p < Np; ++p) {
98       PetscReal temp;
99       ierr = PetscRandomGetValueReal(rngx, &value);CHKERRQ(ierr);
100       x[p*dim] = value;
101       x[p*dim+1] = 0.;
102       temp = erfinv(2*value-1);
103       v[p*dim] = temp;
104       v[p*dim+1] = 0.;
105     }
106   } else if (user.dim_inp == 2) {
107     /*
108      Use Box-Muller to sample a distribution of velocities for the maxwellian.
109      https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform
110     */
111     ierr = VecCreate(comm,&randVec);
112     ierr = VecSetSizes(randVec,PETSC_DECIDE, Np*dim);
113     ierr = VecSetFromOptions(randVec);
114 
115     ierr = ISCreateStride(comm, Np*dim/2, 0, 2, &isvx);CHKERRQ(ierr);
116     ierr = ISCreateStride(comm, Np*dim/2, 1, 2, &isvy);CHKERRQ(ierr);
117     ierr = VecGetSubVector(randVec, isvx, &subvecvx);CHKERRQ(ierr);
118     ierr = VecGetSubVector(randVec, isvy, &subvecvy);CHKERRQ(ierr);
119     ierr = VecSetRandom(subvecvx, rng1);CHKERRQ(ierr);
120     ierr = VecSetRandom(subvecvy, rng2);CHKERRQ(ierr);
121     ierr = VecRestoreSubVector(randVec, isvx, &subvecvx);CHKERRQ(ierr);
122     ierr = VecRestoreSubVector(randVec, isvy, &subvecvy);CHKERRQ(ierr);
123     ierr = VecGetArray(randVec, &randVecNums);CHKERRQ(ierr);
124 
125     for (p = 0; p < Np; ++p) {
126       PetscReal u1, u2, mag, zx, zy;
127 
128       u1 = PetscRealPart(randVecNums[p*dim]);
129       u2 = PetscRealPart(randVecNums[p*dim+1]);
130 
131       x[p*dim] = u1;
132       x[p*dim+1] = u2;
133 
134       mag = PetscSqrtReal(-2.0 * PetscLogReal(u1));
135 
136       zx = mag * PetscCosReal(2*PETSC_PI*u2) + 0;
137       zy = mag * PetscSinReal(2*PETSC_PI*u2) + 0;
138 
139       v[p*dim] = zx;
140       v[p*dim+1] = zy;
141     }
142     ierr = ISDestroy(&isvx);CHKERRQ(ierr);
143     ierr = ISDestroy(&isvy);CHKERRQ(ierr);
144     ierr = VecDestroy(&subvecvx);CHKERRQ(ierr);
145     ierr = VecDestroy(&subvecvy);CHKERRQ(ierr);
146     ierr = VecDestroy(&randVec);CHKERRQ(ierr);
147   } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Do not support dimension %D", dim);
148 
149   ierr = PetscDrawCreate(comm, NULL, "monitor_particle_positions", 0,0,400,300, &positionDraw);CHKERRQ(ierr);
150   ierr = PetscDrawSetFromOptions(positionDraw);CHKERRQ(ierr);
151   ierr = PetscDrawSPCreate(positionDraw, 10, &positionDrawSP);CHKERRQ(ierr);
152   ierr = PetscDrawSPSetDimension(positionDrawSP,1);CHKERRQ(ierr);
153   ierr = PetscDrawSPGetAxis(positionDrawSP, &axis);CHKERRQ(ierr);
154   ierr = PetscDrawSPReset(positionDrawSP);CHKERRQ(ierr);
155   ierr = PetscDrawAxisSetLabels(axis, "Particles", "x", "y");CHKERRQ(ierr);
156   ierr = PetscDrawSetSave(positionDraw, "ex35_pos.ppm");CHKERRQ(ierr);
157   ierr = PetscDrawSPReset(positionDrawSP);CHKERRQ(ierr);
158   ierr = PetscDrawSPSetLimits(positionDrawSP, 0, 1, 0, 1);CHKERRQ(ierr);
159   for (p = 0; p < Np; ++p) {
160     speed = PetscSqrtReal(PetscSqr(v[p*dim]) + PetscSqr(v[p*dim+1]));
161     ierr = PetscDrawSPAddPointColorized(positionDrawSP, &x[p*dim], &x[p*dim+1], &speed);CHKERRQ(ierr);
162   }
163   ierr = PetscDrawSPDraw(positionDrawSP, PETSC_TRUE);CHKERRQ(ierr);
164   ierr = PetscDrawSave(positionDraw);CHKERRQ(ierr);
165 
166   ierr = PetscFree2(x, v);CHKERRQ(ierr);
167   ierr = PetscRandomDestroy(&rngx);CHKERRQ(ierr);
168   ierr = PetscRandomDestroy(&rng1);CHKERRQ(ierr);
169   ierr = PetscRandomDestroy(&rng2);CHKERRQ(ierr);
170 
171   ierr = PetscDrawSPDestroy(&positionDrawSP);CHKERRQ(ierr);
172   ierr = PetscDrawDestroy(&positionDraw);CHKERRQ(ierr);
173   ierr = PetscFinalize();
174   return ierr;
175 }
176 
177 /*TEST
178    test:
179      suffix: 1D
180      args: -Np 50\
181      -dim 1
182    test:
183      suffix: 2D
184      args: -Np 50\
185      -dim 2
186 TEST*/
187