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