xref: /petsc/src/dm/dt/interface/dtprob.c (revision bfe80ac4a46d58cb7760074b25f5e81b2f541d8a)
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 options database name for the probability distribution type
521 
522   Output Parameters:
523 + pdf     - The PDF of this type, or `NULL`
524 . cdf     - The CDF of this type, or `NULL`
525 - sampler - The PDF sampler of this type, or `NULL`
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 typedef enum {
624   NONE,
625   ASCII,
626   DRAW
627 } OutputType;
628 
629 static PetscErrorCode KSViewerCreate(PetscObject obj, OutputType *outputType, PetscViewer *viewer)
630 {
631   PetscViewerFormat format;
632   PetscOptions      options;
633   const char       *prefix;
634   PetscBool         flg;
635   MPI_Comm          comm;
636 
637   PetscFunctionBegin;
638   *outputType = NONE;
639   PetscCall(PetscObjectGetComm(obj, &comm));
640   PetscCall(PetscObjectGetOptionsPrefix(obj, &prefix));
641   PetscCall(PetscObjectGetOptions(obj, &options));
642   PetscCall(PetscOptionsCreateViewer(comm, options, prefix, "-ks_monitor", viewer, &format, &flg));
643   if (flg) {
644     PetscCall(PetscObjectTypeCompare((PetscObject)*viewer, PETSCVIEWERASCII, &flg));
645     if (flg) *outputType = ASCII;
646     PetscCall(PetscObjectTypeCompare((PetscObject)*viewer, PETSCVIEWERDRAW, &flg));
647     if (flg) *outputType = DRAW;
648     PetscCall(PetscViewerPushFormat(*viewer, format));
649   }
650   PetscFunctionReturn(PETSC_SUCCESS);
651 }
652 
653 static PetscErrorCode KSViewerDestroy(PetscViewer *viewer)
654 {
655   PetscFunctionBegin;
656   if (*viewer) {
657     PetscCall(PetscViewerFlush(*viewer));
658     PetscCall(PetscViewerPopFormat(*viewer));
659     PetscCall(PetscViewerDestroy(viewer));
660   }
661   PetscFunctionReturn(PETSC_SUCCESS);
662 }
663 
664 static PetscErrorCode PetscProbComputeKSStatistic_Internal(MPI_Comm comm, PetscInt n, PetscReal val[], PetscReal wgt[], PetscProbFunc cdf, PetscReal *alpha, OutputType outputType, PetscViewer viewer)
665 {
666 #if !defined(PETSC_HAVE_KS)
667   SETERRQ(comm, PETSC_ERR_SUP, "No support for Kolmogorov-Smirnov test.\nReconfigure using --download-ks");
668 #else
669   PetscDraw     draw;
670   PetscDrawLG   lg;
671   PetscDrawAxis axis;
672   const char   *names[2] = {"Analytic", "Empirical"};
673   char          title[PETSC_MAX_PATH_LEN];
674   PetscReal     Fn = 0., Dn = PETSC_MIN_REAL;
675   PetscMPIInt   size;
676 
677   PetscFunctionBegin;
678   PetscCallMPI(MPI_Comm_size(comm, &size));
679   PetscCheck(size == 1, comm, PETSC_ERR_SUP, "Parallel K-S test not yet supported");
680   if (outputType == DRAW) {
681     PetscCall(PetscViewerDrawGetDraw(viewer, 0, &draw));
682     PetscCall(PetscDrawLGCreate(draw, 2, &lg));
683     PetscCall(PetscDrawLGSetLegend(lg, names));
684   }
685   if (wgt) {
686     PetscReal *tmpv, *tmpw;
687     PetscInt  *perm;
688 
689     PetscCall(PetscMalloc3(n, &perm, n, &tmpv, n, &tmpw));
690     for (PetscInt i = 0; i < n; ++i) perm[i] = i;
691     PetscCall(PetscSortRealWithPermutation(n, val, perm));
692     for (PetscInt i = 0; i < n; ++i) {
693       tmpv[i] = val[perm[i]];
694       tmpw[i] = wgt[perm[i]];
695     }
696     for (PetscInt i = 0; i < n; ++i) {
697       val[i] = tmpv[i];
698       wgt[i] = tmpw[i];
699     }
700     PetscCall(PetscFree3(perm, tmpv, tmpw));
701   } else PetscCall(PetscSortReal(n, val));
702   // Compute empirical cumulative distribution F_n and discrepancy D_n
703   for (PetscInt p = 0; p < n; ++p) {
704     const PetscReal x = val[p];
705     const PetscReal w = wgt ? wgt[p] : 1. / n;
706     PetscReal       F, vals[2];
707 
708     Fn += w;
709     PetscCall(cdf(&x, NULL, &F));
710     Dn = PetscMax(PetscAbsReal(Fn - F), Dn);
711     switch (outputType) {
712     case ASCII:
713       PetscCall(PetscViewerASCIIPrintf(viewer, "x: %g F: %g Fn: %g Dn: %.2g\n", x, F, Fn, Dn));
714       break;
715     case DRAW:
716       vals[0] = F;
717       vals[1] = Fn;
718       PetscCall(PetscDrawLGAddCommonPoint(lg, x, vals));
719       break;
720     case NONE:
721       break;
722     }
723   }
724   if (outputType == DRAW) {
725     PetscCall(PetscDrawLGGetAxis(lg, &axis));
726     PetscCall(PetscSNPrintf(title, PETSC_MAX_PATH_LEN, "Kolmogorov-Smirnov Test (Dn %.2g)", Dn));
727     PetscCall(PetscDrawAxisSetLabels(axis, title, "x", "CDF(x)"));
728     PetscCall(PetscDrawLGDraw(lg));
729     PetscCall(PetscDrawLGDestroy(&lg));
730   }
731   *alpha = KSfbar((int)n, (double)Dn);
732   if (outputType == ASCII) PetscCall(PetscViewerASCIIPrintf(viewer, "KSfbar(%" PetscInt_FMT ", %.2g): %g\n", n, Dn, *alpha));
733   PetscFunctionReturn(PETSC_SUCCESS);
734 #endif
735 }
736 
737 /*@C
738   PetscProbComputeKSStatistic - Compute the Kolmogorov-Smirnov statistic for the empirical distribution for an input vector, compared to an analytic CDF.
739 
740   Collective
741 
742   Input Parameters:
743 + v   - The data vector, blocksize is the sample dimension
744 - cdf - The analytic CDF
745 
746   Output Parameter:
747 . alpha - The KS statistic
748 
749   Level: advanced
750 
751   Notes:
752   The Kolmogorov-Smirnov statistic for a given cumulative distribution function $F(x)$ is
753 
754   $$
755   D_n = \sup_x \left| F_n(x) - F(x) \right|
756   $$
757 
758   where $\sup_x$ is the supremum of the set of distances, and the empirical distribution function $F_n(x)$ is discrete, and given by
759 
760   $$
761   F_n =  # of samples <= x / n
762   $$
763 
764   The empirical distribution function $F_n(x)$ is discrete, and thus had a ``stairstep''
765   cumulative distribution, making $n$ the number of stairs. Intuitively, the statistic takes
766   the largest absolute difference between the two distribution functions across all $x$ values.
767 
768   The goodness-of-fit test, or Kolmogorov-Smirnov test, is constructed using the Kolmogorov
769   distribution. It rejects the null hypothesis at level $\alpha$ if
770 
771   $$
772   \sqrt{n} D_{n} > K_{\alpha},
773   $$
774 
775   where $K_\alpha$ is found from
776 
777   $$
778   \operatorname{Pr}(K \leq K_{\alpha}) = 1 - \alpha.
779   $$
780 
781   This means that getting a small alpha says that we have high confidence that the data did not come
782   from the input distribution, so we say that it rejects the null hypothesis.
783 
784 .seealso: `PetscProbComputeKSStatisticWeighted()`, `PetscProbComputeKSStatisticMagnitude()`, `PetscProbFunc`
785 @*/
786 PetscErrorCode PetscProbComputeKSStatistic(Vec v, PetscProbFunc cdf, PetscReal *alpha)
787 {
788   PetscViewer        viewer     = NULL;
789   OutputType         outputType = NONE;
790   const PetscScalar *val;
791   PetscInt           n;
792 
793   PetscFunctionBegin;
794   PetscCall(KSViewerCreate((PetscObject)v, &outputType, &viewer));
795   PetscCall(VecGetLocalSize(v, &n));
796   PetscCall(VecGetArrayRead(v, &val));
797   PetscCheck(!PetscDefined(USE_COMPLEX), PetscObjectComm((PetscObject)v), PETSC_ERR_SUP, "K-S test does not support complex");
798   PetscCall(PetscProbComputeKSStatistic_Internal(PetscObjectComm((PetscObject)v), n, (PetscReal *)val, NULL, cdf, alpha, outputType, viewer));
799   PetscCall(VecRestoreArrayRead(v, &val));
800   PetscCall(KSViewerDestroy(&viewer));
801   PetscFunctionReturn(PETSC_SUCCESS);
802 }
803 
804 /*@C
805   PetscProbComputeKSStatisticWeighted - Compute the Kolmogorov-Smirnov statistic for the weighted empirical distribution for an input vector, compared to an analytic CDF.
806 
807   Collective
808 
809   Input Parameters:
810 + v   - The data vector, blocksize is the sample dimension
811 . w   - The vector of weights for each sample, instead of the default 1/n
812 - cdf - The analytic CDF
813 
814   Output Parameter:
815 . alpha - The KS statistic
816 
817   Level: advanced
818 
819   Notes:
820   The Kolmogorov-Smirnov statistic for a given cumulative distribution function $F(x)$ is
821 
822   $$
823   D_n = \sup_x \left| F_n(x) - F(x) \right|
824   $$
825 
826   where $\sup_x$ is the supremum of the set of distances, and the empirical distribution function $F_n(x)$ is discrete, and given by
827 
828   $$
829   F_n =  # of samples <= x / n
830   $$
831 
832   The empirical distribution function $F_n(x)$ is discrete, and thus had a ``stairstep''
833   cumulative distribution, making $n$ the number of stairs. Intuitively, the statistic takes
834   the largest absolute difference between the two distribution functions across all $x$ values.
835 
836   The goodness-of-fit test, or Kolmogorov-Smirnov test, is constructed using the Kolmogorov
837   distribution. It rejects the null hypothesis at level $\alpha$ if
838 
839   $$
840   \sqrt{n} D_{n} > K_{\alpha},
841   $$
842 
843   where $K_\alpha$ is found from
844 
845   $$
846   \operatorname{Pr}(K \leq K_{\alpha}) = 1 - \alpha.
847   $$
848 
849   This means that getting a small alpha says that we have high confidence that the data did not come
850   from the input distribution, so we say that it rejects the null hypothesis.
851 
852 .seealso: `PetscProbComputeKSStatistic()`, `PetscProbComputeKSStatisticMagnitude()`, `PetscProbFunc`
853 @*/
854 PetscErrorCode PetscProbComputeKSStatisticWeighted(Vec v, Vec w, PetscProbFunc cdf, PetscReal *alpha)
855 {
856   PetscViewer        viewer     = NULL;
857   OutputType         outputType = NONE;
858   const PetscScalar *val, *wgt;
859   PetscInt           n;
860 
861   PetscFunctionBegin;
862   PetscCall(KSViewerCreate((PetscObject)v, &outputType, &viewer));
863   PetscCall(VecGetLocalSize(v, &n));
864   PetscCall(VecGetArrayRead(v, &val));
865   PetscCall(VecGetArrayRead(w, &wgt));
866   PetscCheck(!PetscDefined(USE_COMPLEX), PetscObjectComm((PetscObject)v), PETSC_ERR_SUP, "K-S test does not support complex");
867   PetscCall(PetscProbComputeKSStatistic_Internal(PetscObjectComm((PetscObject)v), n, (PetscReal *)val, (PetscReal *)wgt, cdf, alpha, outputType, viewer));
868   PetscCall(VecRestoreArrayRead(v, &val));
869   PetscCall(VecRestoreArrayRead(w, &wgt));
870   PetscCall(KSViewerDestroy(&viewer));
871   PetscFunctionReturn(PETSC_SUCCESS);
872 }
873 
874 /*@C
875   PetscProbComputeKSStatisticMagnitude - Compute the Kolmogorov-Smirnov statistic for the empirical distribution for the magnitude over each block of an input vector, compared to an analytic CDF.
876 
877   Collective
878 
879   Input Parameters:
880 + v   - The data vector, blocksize is the sample dimension
881 - cdf - The analytic CDF
882 
883   Output Parameter:
884 . alpha - The KS statistic
885 
886   Level: advanced
887 
888   Notes:
889   The Kolmogorov-Smirnov statistic for a given cumulative distribution function $F(x)$ is
890 
891   $$
892   D_n = \sup_x \left| F_n(x) - F(x) \right|
893   $$
894 
895   where $\sup_x$ is the supremum of the set of distances, and the empirical distribution function $F_n(x)$ is discrete, and given by
896 
897   $$
898   F_n =  # of samples <= x / n
899   $$
900 
901   The empirical distribution function $F_n(x)$ is discrete, and thus had a ``stairstep''
902   cumulative distribution, making $n$ the number of stairs. Intuitively, the statistic takes
903   the largest absolute difference between the two distribution functions across all $x$ values.
904 
905   The goodness-of-fit test, or Kolmogorov-Smirnov test, is constructed using the Kolmogorov
906   distribution. It rejects the null hypothesis at level $\alpha$ if
907 
908   $$
909   \sqrt{n} D_{n} > K_{\alpha},
910   $$
911 
912   where $K_\alpha$ is found from
913 
914   $$
915   \operatorname{Pr}(K \leq K_{\alpha}) = 1 - \alpha.
916   $$
917 
918   This means that getting a small alpha says that we have high confidence that the data did not come
919   from the input distribution, so we say that it rejects the null hypothesis.
920 
921 .seealso: `PetscProbComputeKSStatistic()`, `PetscProbComputeKSStatisticWeighted()`, `PetscProbFunc`
922 @*/
923 PetscErrorCode PetscProbComputeKSStatisticMagnitude(Vec v, PetscProbFunc cdf, PetscReal *alpha)
924 {
925   PetscViewer        viewer     = NULL;
926   OutputType         outputType = NONE;
927   const PetscScalar *a;
928   PetscReal         *speed;
929   PetscInt           n, dim;
930 
931   PetscFunctionBegin;
932   PetscCall(KSViewerCreate((PetscObject)v, &outputType, &viewer));
933   // Convert to a scalar value
934   PetscCall(VecGetLocalSize(v, &n));
935   PetscCall(VecGetBlockSize(v, &dim));
936   n /= dim;
937   PetscCall(PetscMalloc1(n, &speed));
938   PetscCall(VecGetArrayRead(v, &a));
939   for (PetscInt p = 0; p < n; ++p) {
940     PetscReal mag = 0.;
941 
942     for (PetscInt d = 0; d < dim; ++d) mag += PetscSqr(PetscRealPart(a[p * dim + d]));
943     speed[p] = PetscSqrtReal(mag);
944   }
945   PetscCall(VecRestoreArrayRead(v, &a));
946   // Compute statistic
947   PetscCall(PetscProbComputeKSStatistic_Internal(PetscObjectComm((PetscObject)v), n, speed, NULL, cdf, alpha, outputType, viewer));
948   PetscCall(PetscFree(speed));
949   PetscCall(KSViewerDestroy(&viewer));
950   PetscFunctionReturn(PETSC_SUCCESS);
951 }
952