xref: /petsc/src/dm/dt/interface/dtprob.c (revision 98d129c30f3ee9fdddc40fdbc5a989b7be64f888)
1 #include <petscdt.h> /*I "petscdt.h" I*/
2 
3 #include <petscvec.h>
4 #include <petscdraw.h>
5 #include <petsc/private/petscimpl.h>
6 
7 const char *const DTProbDensityTypes[] = {"constant", "gaussian", "maxwell_boltzmann", "DTProb Density", "DTPROB_DENSITY_", NULL};
8 
9 /*@
10   PetscPDFMaxwellBoltzmann1D - PDF for the Maxwell-Boltzmann distribution in 1D
11 
12   Not Collective
13 
14   Input Parameters:
15 + x     - Speed in $[0, \infty]$
16 - dummy - Unused
17 
18   Output Parameter:
19 . p - The probability density at `x`
20 
21   Level: beginner
22 
23 .seealso: `PetscCDFMaxwellBoltzmann1D()`
24 @*/
25 PetscErrorCode PetscPDFMaxwellBoltzmann1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
26 {
27   p[0] = PetscSqrtReal(2. / PETSC_PI) * PetscExpReal(-0.5 * PetscSqr(x[0]));
28   return PETSC_SUCCESS;
29 }
30 
31 /*@
32   PetscCDFMaxwellBoltzmann1D - CDF for the Maxwell-Boltzmann distribution in 1D
33 
34   Not Collective
35 
36   Input Parameters:
37 + x     - Speed in $[0, \infty]$
38 - dummy - Unused
39 
40   Output Parameter:
41 . p - The probability density at `x`
42 
43   Level: beginner
44 
45 .seealso: `PetscPDFMaxwellBoltzmann1D()`
46 @*/
47 PetscErrorCode PetscCDFMaxwellBoltzmann1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
48 {
49   p[0] = PetscErfReal(x[0] / PETSC_SQRT2);
50   return PETSC_SUCCESS;
51 }
52 
53 /*@
54   PetscPDFMaxwellBoltzmann2D - PDF for the Maxwell-Boltzmann distribution in 2D
55 
56   Not Collective
57 
58   Input Parameters:
59 + x     - Speed in $[0, \infty]$
60 - dummy - Unused
61 
62   Output Parameter:
63 . p - The probability density at `x`
64 
65   Level: beginner
66 
67 .seealso: `PetscCDFMaxwellBoltzmann2D()`
68 @*/
69 PetscErrorCode PetscPDFMaxwellBoltzmann2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
70 {
71   p[0] = x[0] * PetscExpReal(-0.5 * PetscSqr(x[0]));
72   return PETSC_SUCCESS;
73 }
74 
75 /*@
76   PetscCDFMaxwellBoltzmann2D - CDF for the Maxwell-Boltzmann distribution in 2D
77 
78   Not Collective
79 
80   Input Parameters:
81 + x     - Speed in $[0, \infty]$
82 - dummy - Unused
83 
84   Output Parameter:
85 . p - The probability density at `x`
86 
87   Level: beginner
88 
89 .seealso: `PetscPDFMaxwellBoltzmann2D()`
90 @*/
91 PetscErrorCode PetscCDFMaxwellBoltzmann2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
92 {
93   p[0] = 1. - PetscExpReal(-0.5 * PetscSqr(x[0]));
94   return PETSC_SUCCESS;
95 }
96 
97 /*@
98   PetscPDFMaxwellBoltzmann3D - PDF for the Maxwell-Boltzmann distribution in 3D
99 
100   Not Collective
101 
102   Input Parameters:
103 + x     - Speed in $[0, \infty]$
104 - dummy - Unused
105 
106   Output Parameter:
107 . p - The probability density at `x`
108 
109   Level: beginner
110 
111 .seealso: `PetscCDFMaxwellBoltzmann3D()`
112 @*/
113 PetscErrorCode PetscPDFMaxwellBoltzmann3D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
114 {
115   p[0] = PetscSqrtReal(2. / PETSC_PI) * PetscSqr(x[0]) * PetscExpReal(-0.5 * PetscSqr(x[0]));
116   return PETSC_SUCCESS;
117 }
118 
119 /*@
120   PetscCDFMaxwellBoltzmann3D - CDF for the Maxwell-Boltzmann distribution in 3D
121 
122   Not Collective
123 
124   Input Parameters:
125 + x     - Speed in $[0, \infty]$
126 - dummy - Unused
127 
128   Output Parameter:
129 . p - The probability density at `x`
130 
131   Level: beginner
132 
133 .seealso: `PetscPDFMaxwellBoltzmann3D()`
134 @*/
135 PetscErrorCode PetscCDFMaxwellBoltzmann3D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
136 {
137   p[0] = PetscErfReal(x[0] / PETSC_SQRT2) - PetscSqrtReal(2. / PETSC_PI) * x[0] * PetscExpReal(-0.5 * PetscSqr(x[0]));
138   return PETSC_SUCCESS;
139 }
140 
141 /*@
142   PetscPDFGaussian1D - PDF for the Gaussian distribution in 1D
143 
144   Not Collective
145 
146   Input Parameters:
147 + x     - Coordinate in $[-\infty, \infty]$
148 - scale - Scaling value
149 
150   Output Parameter:
151 . p - The probability density at `x`
152 
153   Level: beginner
154 
155 .seealso: `PetscPDFMaxwellBoltzmann3D()`
156 @*/
157 PetscErrorCode PetscPDFGaussian1D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
158 {
159   const PetscReal sigma = scale ? scale[0] : 1.;
160   p[0]                  = PetscSqrtReal(1. / (2. * PETSC_PI)) * PetscExpReal(-0.5 * PetscSqr(x[0] / sigma)) / sigma;
161   return PETSC_SUCCESS;
162 }
163 
164 PetscErrorCode PetscCDFGaussian1D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
165 {
166   const PetscReal sigma = scale ? scale[0] : 1.;
167   p[0]                  = 0.5 * (1. + PetscErfReal(x[0] / PETSC_SQRT2 / sigma));
168   return PETSC_SUCCESS;
169 }
170 
171 /*@
172   PetscPDFSampleGaussian1D - Sample uniformly from a Gaussian distribution in 1D
173 
174   Not Collective
175 
176   Input Parameters:
177 + p     - A uniform variable on $[0, 1]$
178 - dummy - ignored
179 
180   Output Parameter:
181 . x - Coordinate in $[-\infty, \infty]$
182 
183   Level: beginner
184 
185   Note:
186   See <http://www.mimirgames.com/articles/programming/approximations-of-the-inverse-error-function> and
187   <https://stackoverflow.com/questions/27229371/inverse-error-function-in-c>
188 
189 .seealso: `PetscPDFGaussian2D()`
190 @*/
191 PetscErrorCode PetscPDFSampleGaussian1D(const PetscReal p[], const PetscReal dummy[], PetscReal x[])
192 {
193   const PetscReal q       = 2 * p[0] - 1.;
194   const PetscInt  maxIter = 100;
195   PetscReal       ck[100], r = 0.;
196   PetscInt        k, m;
197 
198   PetscFunctionBeginHot;
199   /* Transform input to [-1, 1] since the code below computes the inverse error function */
200   for (k = 0; k < maxIter; ++k) ck[k] = 0.;
201   ck[0] = 1;
202   r     = ck[0] * (PetscSqrtReal(PETSC_PI) / 2.) * q;
203   for (k = 1; k < maxIter; ++k) {
204     const PetscReal temp = 2. * k + 1.;
205 
206     for (m = 0; m <= k - 1; ++m) {
207       PetscReal denom = (m + 1.) * (2. * m + 1.);
208 
209       ck[k] += (ck[m] * ck[k - 1 - m]) / denom;
210     }
211     r += (ck[k] / temp) * PetscPowReal((PetscSqrtReal(PETSC_PI) / 2.) * q, 2. * k + 1.);
212   }
213   /* Scale erfinv() by \sqrt{\pi/2} */
214   x[0] = PetscSqrtReal(PETSC_PI * 0.5) * r;
215   PetscFunctionReturn(PETSC_SUCCESS);
216 }
217 
218 /*@
219   PetscPDFGaussian2D - PDF for the Gaussian distribution in 2D
220 
221   Not Collective
222 
223   Input Parameters:
224 + x     - Coordinate in $[-\infty, \infty]^2$
225 - dummy - ignored
226 
227   Output Parameter:
228 . p - The probability density at `x`
229 
230   Level: beginner
231 
232 .seealso: `PetscPDFSampleGaussian2D()`, `PetscPDFMaxwellBoltzmann3D()`
233 @*/
234 PetscErrorCode PetscPDFGaussian2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
235 {
236   p[0] = (1. / PETSC_PI) * PetscExpReal(-0.5 * (PetscSqr(x[0]) + PetscSqr(x[1])));
237   return PETSC_SUCCESS;
238 }
239 
240 /*@
241   PetscPDFSampleGaussian2D - Sample uniformly from a Gaussian distribution in 2D
242   <https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform>
243 
244   Not Collective
245 
246   Input Parameters:
247 + p     - A uniform variable on $[0, 1]^2$
248 - dummy - ignored
249 
250   Output Parameter:
251 . x - Coordinate in $[-\infty, \infty]^2 $
252 
253   Level: beginner
254 
255 .seealso: `PetscPDFGaussian2D()`, `PetscPDFMaxwellBoltzmann3D()`
256 @*/
257 PetscErrorCode PetscPDFSampleGaussian2D(const PetscReal p[], const PetscReal dummy[], PetscReal x[])
258 {
259   const PetscReal mag = PetscSqrtReal(-2.0 * PetscLogReal(p[0]));
260   x[0]                = mag * PetscCosReal(2.0 * PETSC_PI * p[1]);
261   x[1]                = mag * PetscSinReal(2.0 * PETSC_PI * p[1]);
262   return PETSC_SUCCESS;
263 }
264 
265 /*@
266   PetscPDFGaussian3D - PDF for the Gaussian distribution in 3D
267 
268   Not Collective
269 
270   Input Parameters:
271 + x     - Coordinate in $[-\infty, \infty]^3$
272 - dummy - ignored
273 
274   Output Parameter:
275 . p - The probability density at `x`
276 
277   Level: beginner
278 
279 .seealso: `PetscPDFSampleGaussian3D()`, `PetscPDFMaxwellBoltzmann3D()`
280 @*/
281 PetscErrorCode PetscPDFGaussian3D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
282 {
283   p[0] = (1. / PETSC_PI * PetscSqrtReal(PETSC_PI)) * PetscExpReal(-0.5 * (PetscSqr(x[0]) + PetscSqr(x[1]) + PetscSqr(x[2])));
284   return PETSC_SUCCESS;
285 }
286 
287 /*@
288   PetscPDFSampleGaussian3D - Sample uniformly from a Gaussian distribution in 3D
289   <https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform>
290 
291   Not Collective
292 
293   Input Parameters:
294 + p     - A uniform variable on $[0, 1]^3$
295 - dummy - ignored
296 
297   Output Parameter:
298 . x - Coordinate in $[-\infty, \infty]^3$
299 
300   Level: beginner
301 
302 .seealso: `PetscPDFGaussian3D()`, `PetscPDFMaxwellBoltzmann3D()`
303 @*/
304 PetscErrorCode PetscPDFSampleGaussian3D(const PetscReal p[], const PetscReal dummy[], PetscReal x[])
305 {
306   PetscCall(PetscPDFSampleGaussian1D(p, dummy, x));
307   PetscCall(PetscPDFSampleGaussian2D(&p[1], dummy, &x[1]));
308   return PETSC_SUCCESS;
309 }
310 
311 /*@
312   PetscPDFConstant1D - PDF for the uniform distribution in 1D
313 
314   Not Collective
315 
316   Input Parameters:
317 + x     - Coordinate in $[-1, 1]$
318 - dummy - Unused
319 
320   Output Parameter:
321 . p - The probability density at `x`
322 
323   Level: beginner
324 
325 .seealso: `PetscCDFConstant1D()`, `PetscPDFSampleConstant1D()`, `PetscPDFConstant2D()`, `PetscPDFConstant3D()`
326 @*/
327 PetscErrorCode PetscPDFConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
328 {
329   p[0] = x[0] > -1. && x[0] <= 1. ? 0.5 : 0.;
330   return PETSC_SUCCESS;
331 }
332 
333 /*@
334   PetscCDFConstant1D - CDF for the uniform distribution in 1D
335 
336   Not Collective
337 
338   Input Parameters:
339 + x     - Coordinate in $[-1, 1]$
340 - dummy - Unused
341 
342   Output Parameter:
343 . p - The cumulative probability at `x`
344 
345   Level: beginner
346 
347 .seealso: `PetscPDFConstant1D()`, `PetscPDFSampleConstant1D()`, `PetscCDFConstant2D()`, `PetscCDFConstant3D()`
348 @*/
349 PetscErrorCode PetscCDFConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
350 {
351   p[0] = x[0] <= -1. ? 0. : (x[0] > 1. ? 1. : 0.5 * (x[0] + 1.));
352   return PETSC_SUCCESS;
353 }
354 
355 /*@
356   PetscPDFSampleConstant1D - Sample uniformly from a uniform distribution on [-1, 1] in 1D
357 
358   Not Collective
359 
360   Input Parameters:
361 + p     - A uniform variable on $[0, 1]$
362 - dummy - Unused
363 
364   Output Parameter:
365 . x - Coordinate in $[-1, 1]$
366 
367   Level: beginner
368 
369 .seealso: `PetscPDFConstant1D()`, `PetscCDFConstant1D()`, `PetscPDFSampleConstant2D()`, `PetscPDFSampleConstant3D()`
370 @*/
371 PetscErrorCode PetscPDFSampleConstant1D(const PetscReal p[], const PetscReal dummy[], PetscReal x[])
372 {
373   x[0] = 2. * p[0] - 1.;
374   return PETSC_SUCCESS;
375 }
376 
377 /*@
378   PetscPDFConstant2D - PDF for the uniform distribution in 2D
379 
380   Not Collective
381 
382   Input Parameters:
383 + x     - Coordinate in $[-1, 1]^2$
384 - dummy - Unused
385 
386   Output Parameter:
387 . p - The probability density at `x`
388 
389   Level: beginner
390 
391 .seealso: `PetscCDFConstant2D()`, `PetscPDFSampleConstant2D()`, `PetscPDFConstant1D()`, `PetscPDFConstant3D()`
392 @*/
393 PetscErrorCode PetscPDFConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
394 {
395   p[0] = x[0] > -1. && x[0] <= 1. && x[1] > -1. && x[1] <= 1. ? 0.25 : 0.;
396   return PETSC_SUCCESS;
397 }
398 
399 /*@
400   PetscCDFConstant2D - CDF for the uniform distribution in 2D
401 
402   Not Collective
403 
404   Input Parameters:
405 + x     - Coordinate in $[-1, 1]^2$
406 - dummy - Unused
407 
408   Output Parameter:
409 . p - The cumulative probability at `x`
410 
411   Level: beginner
412 
413 .seealso: `PetscPDFConstant2D()`, `PetscPDFSampleConstant2D()`, `PetscCDFConstant1D()`, `PetscCDFConstant3D()`
414 @*/
415 PetscErrorCode PetscCDFConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
416 {
417   p[0] = x[0] <= -1. || x[1] <= -1. ? 0. : (x[0] > 1. ? 1. : 0.5 * (x[0] + 1.)) * (x[1] > 1. ? 1. : 0.5 * (x[1] + 1.));
418   return PETSC_SUCCESS;
419 }
420 
421 /*@
422   PetscPDFSampleConstant2D - Sample uniformly from a uniform distribution on $[-1, 1]^2$ in 2D
423 
424   Not Collective
425 
426   Input Parameters:
427 + p     - Two uniform variables on $[0, 1]$
428 - dummy - Unused
429 
430   Output Parameter:
431 . x - Coordinate in $[-1, 1]^2$
432 
433   Level: beginner
434 
435 .seealso: `PetscPDFConstant2D()`, `PetscCDFConstant2D()`, `PetscPDFSampleConstant1D()`, `PetscPDFSampleConstant3D()`
436 @*/
437 PetscErrorCode PetscPDFSampleConstant2D(const PetscReal p[], const PetscReal dummy[], PetscReal x[])
438 {
439   x[0] = 2. * p[0] - 1.;
440   x[1] = 2. * p[1] - 1.;
441   return PETSC_SUCCESS;
442 }
443 
444 /*@
445   PetscPDFConstant3D - PDF for the uniform distribution in 3D
446 
447   Not Collective
448 
449   Input Parameters:
450 + x     - Coordinate in $[-1, 1]^3$
451 - dummy - Unused
452 
453   Output Parameter:
454 . p - The probability density at `x`
455 
456   Level: beginner
457 
458 .seealso: `PetscCDFConstant3D()`, `PetscPDFSampleConstant3D()`, `PetscPDFSampleConstant1D()`, `PetscPDFSampleConstant2D()`
459 @*/
460 PetscErrorCode PetscPDFConstant3D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
461 {
462   p[0] = x[0] > -1. && x[0] <= 1. && x[1] > -1. && x[1] <= 1. && x[2] > -1. && x[2] <= 1. ? 0.125 : 0.;
463   return PETSC_SUCCESS;
464 }
465 
466 /*@
467   PetscCDFConstant3D - CDF for the uniform distribution in 3D
468 
469   Not Collective
470 
471   Input Parameters:
472 + x     - Coordinate in $[-1, 1]^3$
473 - dummy - Unused
474 
475   Output Parameter:
476 . p - The cumulative probability at `x`
477 
478   Level: beginner
479 
480 .seealso: `PetscPDFConstant3D()`, `PetscPDFSampleConstant3D()`, `PetscCDFConstant1D()`, `PetscCDFConstant2D()`
481 @*/
482 PetscErrorCode PetscCDFConstant3D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
483 {
484   p[0] = x[0] <= -1. || x[1] <= -1. || x[2] <= -1. ? 0. : (x[0] > 1. ? 1. : 0.5 * (x[0] + 1.)) * (x[1] > 1. ? 1. : 0.5 * (x[1] + 1.)) * (x[2] > 1. ? 1. : 0.5 * (x[2] + 1.));
485   return PETSC_SUCCESS;
486 }
487 
488 /*@
489   PetscPDFSampleConstant3D - Sample uniformly from a uniform distribution on $[-1, 1]^3$ in 3D
490 
491   Not Collective
492 
493   Input Parameters:
494 + p     - Three uniform variables on $[0, 1]$
495 - dummy - Unused
496 
497   Output Parameter:
498 . x - Coordinate in $[-1, 1]^3$
499 
500   Level: beginner
501 
502 .seealso: `PetscPDFConstant3D()`, `PetscCDFConstant3D()`, `PetscPDFSampleConstant1D()`, `PetscPDFSampleConstant2D()`
503 @*/
504 PetscErrorCode PetscPDFSampleConstant3D(const PetscReal p[], const PetscReal dummy[], PetscReal x[])
505 {
506   x[0] = 2. * p[0] - 1.;
507   x[1] = 2. * p[1] - 1.;
508   x[2] = 2. * p[2] - 1.;
509   return PETSC_SUCCESS;
510 }
511 
512 /*@C
513   PetscProbCreateFromOptions - Return the probability distribution specified by the arguments and options
514 
515   Not Collective
516 
517   Input Parameters:
518 + dim    - The dimension of sample points
519 . prefix - The options prefix, or `NULL`
520 - name   - The option name for the probability distribution type
521 
522   Output Parameters:
523 + pdf     - The PDF of this type
524 . cdf     - The CDF of this type
525 - sampler - The PDF sampler of this type
526 
527   Level: intermediate
528 
529 .seealso: `PetscProbFunc`, `PetscPDFMaxwellBoltzmann1D()`, `PetscPDFGaussian1D()`, `PetscPDFConstant1D()`
530 @*/
531 PetscErrorCode PetscProbCreateFromOptions(PetscInt dim, const char prefix[], const char name[], PetscProbFunc *pdf, PetscProbFunc *cdf, PetscProbFunc *sampler)
532 {
533   DTProbDensityType den = DTPROB_DENSITY_GAUSSIAN;
534 
535   PetscFunctionBegin;
536   PetscOptionsBegin(PETSC_COMM_SELF, prefix, "PetscProb Options", "DT");
537   PetscCall(PetscOptionsEnum(name, "Method to compute PDF <constant, gaussian>", "", DTProbDensityTypes, (PetscEnum)den, (PetscEnum *)&den, NULL));
538   PetscOptionsEnd();
539 
540   if (pdf) {
541     PetscAssertPointer(pdf, 4);
542     *pdf = NULL;
543   }
544   if (cdf) {
545     PetscAssertPointer(cdf, 5);
546     *cdf = NULL;
547   }
548   if (sampler) {
549     PetscAssertPointer(sampler, 6);
550     *sampler = NULL;
551   }
552   switch (den) {
553   case DTPROB_DENSITY_CONSTANT:
554     switch (dim) {
555     case 1:
556       if (pdf) *pdf = PetscPDFConstant1D;
557       if (cdf) *cdf = PetscCDFConstant1D;
558       if (sampler) *sampler = PetscPDFSampleConstant1D;
559       break;
560     case 2:
561       if (pdf) *pdf = PetscPDFConstant2D;
562       if (cdf) *cdf = PetscCDFConstant2D;
563       if (sampler) *sampler = PetscPDFSampleConstant2D;
564       break;
565     case 3:
566       if (pdf) *pdf = PetscPDFConstant3D;
567       if (cdf) *cdf = PetscCDFConstant3D;
568       if (sampler) *sampler = PetscPDFSampleConstant3D;
569       break;
570     default:
571       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Dimension %" PetscInt_FMT " not supported for density type %s", dim, DTProbDensityTypes[den]);
572     }
573     break;
574   case DTPROB_DENSITY_GAUSSIAN:
575     switch (dim) {
576     case 1:
577       if (pdf) *pdf = PetscPDFGaussian1D;
578       if (cdf) *cdf = PetscCDFGaussian1D;
579       if (sampler) *sampler = PetscPDFSampleGaussian1D;
580       break;
581     case 2:
582       if (pdf) *pdf = PetscPDFGaussian2D;
583       if (sampler) *sampler = PetscPDFSampleGaussian2D;
584       break;
585     case 3:
586       if (pdf) *pdf = PetscPDFGaussian3D;
587       if (sampler) *sampler = PetscPDFSampleGaussian3D;
588       break;
589     default:
590       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Dimension %" PetscInt_FMT " not supported for density type %s", dim, DTProbDensityTypes[den]);
591     }
592     break;
593   case DTPROB_DENSITY_MAXWELL_BOLTZMANN:
594     switch (dim) {
595     case 1:
596       if (pdf) *pdf = PetscPDFMaxwellBoltzmann1D;
597       if (cdf) *cdf = PetscCDFMaxwellBoltzmann1D;
598       break;
599     case 2:
600       if (pdf) *pdf = PetscPDFMaxwellBoltzmann2D;
601       if (cdf) *cdf = PetscCDFMaxwellBoltzmann2D;
602       break;
603     case 3:
604       if (pdf) *pdf = PetscPDFMaxwellBoltzmann3D;
605       if (cdf) *cdf = PetscCDFMaxwellBoltzmann3D;
606       break;
607     default:
608       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Dimension %" PetscInt_FMT " not supported for density type %s", dim, DTProbDensityTypes[den]);
609     }
610     break;
611   default:
612     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "Density type %s is not supported", DTProbDensityTypes[PetscMax(0, PetscMin(den, DTPROB_NUM_DENSITY))]);
613   }
614   PetscFunctionReturn(PETSC_SUCCESS);
615 }
616 
617 #ifdef PETSC_HAVE_KS
618 EXTERN_C_BEGIN
619   #include <KolmogorovSmirnovDist.h>
620 EXTERN_C_END
621 #endif
622 
623 /*@C
624   PetscProbComputeKSStatistic - Compute the Kolmogorov-Smirnov statistic for the empirical distribution for an input vector, compared to an analytic CDF.
625 
626   Collective
627 
628   Input Parameters:
629 + v   - The data vector, blocksize is the sample dimension
630 - cdf - The analytic CDF
631 
632   Output Parameter:
633 . alpha - The KS statisic
634 
635   Level: advanced
636 
637   Notes:
638   The Kolmogorov-Smirnov statistic for a given cumulative distribution function $F(x)$ is
639 
640   $$
641   D_n = \sup_x \left| F_n(x) - F(x) \right|
642   $$
643 
644   where $\sup_x$ is the supremum of the set of distances, and the empirical distribution function $F_n(x)$ is discrete, and given by
645 
646   $$
647   F_n =  # of samples <= x / n
648   $$
649 
650   The empirical distribution function $F_n(x)$ is discrete, and thus had a ``stairstep''
651   cumulative distribution, making $n$ the number of stairs. Intuitively, the statistic takes
652   the largest absolute difference between the two distribution functions across all $x$ values.
653 
654 .seealso: `PetscProbFunc`
655 @*/
656 PetscErrorCode PetscProbComputeKSStatistic(Vec v, PetscProbFunc cdf, PetscReal *alpha)
657 {
658 #if !defined(PETSC_HAVE_KS)
659   SETERRQ(PetscObjectComm((PetscObject)v), PETSC_ERR_SUP, "No support for Kolmogorov-Smirnov test.\nReconfigure using --download-ks");
660 #else
661   PetscViewer        viewer = NULL;
662   PetscViewerFormat  format;
663   PetscDraw          draw;
664   PetscDrawLG        lg;
665   PetscDrawAxis      axis;
666   const PetscScalar *a;
667   PetscReal         *speed, Dn = PETSC_MIN_REAL;
668   PetscBool          isascii = PETSC_FALSE, isdraw = PETSC_FALSE, flg;
669   PetscInt           n, p, dim, d;
670   PetscMPIInt        size;
671   const char        *names[2] = {"Analytic", "Empirical"};
672   char               title[PETSC_MAX_PATH_LEN];
673   PetscOptions       options;
674   const char        *prefix;
675   MPI_Comm           comm;
676 
677   PetscFunctionBegin;
678   PetscCall(PetscObjectGetComm((PetscObject)v, &comm));
679   PetscCall(PetscObjectGetOptionsPrefix((PetscObject)v, &prefix));
680   PetscCall(PetscObjectGetOptions((PetscObject)v, &options));
681   PetscCall(PetscOptionsGetViewer(comm, options, prefix, "-ks_monitor", &viewer, &format, &flg));
682   if (flg) {
683     PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
684     PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
685   }
686   if (isascii) {
687     PetscCall(PetscViewerPushFormat(viewer, format));
688   } else if (isdraw) {
689     PetscCall(PetscViewerPushFormat(viewer, format));
690     PetscCall(PetscViewerDrawGetDraw(viewer, 0, &draw));
691     PetscCall(PetscDrawLGCreate(draw, 2, &lg));
692     PetscCall(PetscDrawLGSetLegend(lg, names));
693   }
694 
695   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)v), &size));
696   PetscCall(VecGetLocalSize(v, &n));
697   PetscCall(VecGetBlockSize(v, &dim));
698   n /= dim;
699   /* TODO Parallel algorithm is harder */
700   if (size == 1) {
701     PetscCall(PetscMalloc1(n, &speed));
702     PetscCall(VecGetArrayRead(v, &a));
703     for (p = 0; p < n; ++p) {
704       PetscReal mag = 0.;
705 
706       for (d = 0; d < dim; ++d) mag += PetscSqr(PetscRealPart(a[p * dim + d]));
707       speed[p] = PetscSqrtReal(mag);
708     }
709     PetscCall(PetscSortReal(n, speed));
710     PetscCall(VecRestoreArrayRead(v, &a));
711     for (p = 0; p < n; ++p) {
712       const PetscReal x = speed[p], Fn = ((PetscReal)p) / n;
713       PetscReal       F, vals[2];
714 
715       PetscCall(cdf(&x, NULL, &F));
716       Dn = PetscMax(PetscAbsReal(Fn - F), Dn);
717       if (isascii) PetscCall(PetscViewerASCIIPrintf(viewer, "x: %g F: %g Fn: %g Dn: %.2g\n", x, F, Fn, Dn));
718       if (isdraw) {
719         vals[0] = F;
720         vals[1] = Fn;
721         PetscCall(PetscDrawLGAddCommonPoint(lg, x, vals));
722       }
723     }
724     if (speed[n - 1] < 6.) {
725       const PetscReal k = (PetscInt)(6. - speed[n - 1]) + 1, dx = (6. - speed[n - 1]) / k;
726       for (p = 0; p <= k; ++p) {
727         const PetscReal x = speed[n - 1] + p * dx, Fn = 1.0;
728         PetscReal       F, vals[2];
729 
730         PetscCall(cdf(&x, NULL, &F));
731         Dn = PetscMax(PetscAbsReal(Fn - F), Dn);
732         if (isascii) PetscCall(PetscViewerASCIIPrintf(viewer, "x: %g F: %g Fn: %g Dn: %.2g\n", x, F, Fn, Dn));
733         if (isdraw) {
734           vals[0] = F;
735           vals[1] = Fn;
736           PetscCall(PetscDrawLGAddCommonPoint(lg, x, vals));
737         }
738       }
739     }
740     PetscCall(PetscFree(speed));
741   }
742   if (isdraw) {
743     PetscCall(PetscDrawLGGetAxis(lg, &axis));
744     PetscCall(PetscSNPrintf(title, PETSC_MAX_PATH_LEN, "Kolmogorov-Smirnov Test (Dn %.2g)", Dn));
745     PetscCall(PetscDrawAxisSetLabels(axis, title, "x", "CDF(x)"));
746     PetscCall(PetscDrawLGDraw(lg));
747     PetscCall(PetscDrawLGDestroy(&lg));
748   }
749   if (viewer) {
750     PetscCall(PetscViewerFlush(viewer));
751     PetscCall(PetscViewerPopFormat(viewer));
752     PetscCall(PetscOptionsRestoreViewer(&viewer));
753   }
754   *alpha = KSfbar((int)n, (double)Dn);
755   PetscFunctionReturn(PETSC_SUCCESS);
756 #endif
757 }
758