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