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