Coverage for src/uqtils/example.py: 100%

1"""Examples for using the package.""" 

2# ruff: noqa: F841 

3# ruff: noqa: I001 

4 

5def normal_example(): 

6 """Sample and plot a normal pdf.""" 

7 # --8<-- [start:normal] 

8 import numpy as np 

9 import uqtils as uq 

10 

11 ndim = 3 

12 shape = (5, ndim) 

13 

14 means = np.random.randint(0, 10, size=shape).astype(np.float64) 

15 cov = np.eye(ndim) * 0.1 

16 

17 samples = uq.normal_sample(means, cov, size=1000) # (1000, 5, 3) 

18 pdfs = uq.normal_pdf(samples, means, cov) # (1000, 5) 

19 

20 fig, ax = uq.ndscatter(samples[:, 0, :]) 

21 # --8<-- [end:normal] 

22 

23 

24def gradient_example(): 

25 """Evaluate 1d and multivariate gradients.""" 

26 # --8<-- [start:gradient] 

27 import numpy as np 

28 import uqtils as uq 

29 

30 # 1d example 

31 def f(x): 

32 return np.sin(x) 

33 

34 x0 = 1.0 

35 df_dx = uq.approx_jac(f, x0) 

36 d2f_dx2 = uq.approx_hess(f, x0) 

37 

38 # Multivariate example 

39 n_in, n_out = 3, 2 

40 def f(x): 

41 x0, x1, x2 = [x[..., i] for i in range(n_in)] 

42 f0 = x0 * x1 + x2 

43 f1 = np.sin(x0)**2 + x2**3 

44 return np.concatenate((f0[..., np.newaxis], f1[..., np.newaxis]), axis=-1) 

45 

46 shape = (100, 5, n_in) 

47 x0 = np.random.rand(*shape) 

48 jac = uq.approx_jac(f, x0) # (100, 5, n_out, n_in) 

49 hess = uq.approx_hess(f, x0) # (100, 5, n_out, n_in, n_in) 

50 # --8<-- [end:gradient] 

51 

52 

53def mcmc_example(): 

54 """Sample from a logpdf distribution using MCMC.""" 

55 # --8<-- [start:mcmc] 

56 import numpy as np 

57 import uqtils as uq 

58 

59 def fun(x): 

60 mu = [1, 1] 

61 cov = [[0.5, -0.1], [-0.1, 0.5]] 

62 return uq.normal_pdf(x, mu, cov, logpdf=True) 

63 

64 nsamples, nwalkers, ndim = 1000, 4, 2 

65 x0 = np.random.randn(nwalkers, ndim) 

66 cov0 = np.eye(ndim) 

67 

68 samples, log_pdf, accepted = uq.dram(fun, x0, nsamples, cov0=cov0) 

69 

70 burn_in = int(0.1 * nsamples) 

71 samples = samples[burn_in:, ...].reshape((-1, ndim)) 

72 fig, ax = uq.ndscatter(samples, plot2d='hist') 

73 # --8<-- [end:mcmc] 

74 

75def sobol_example(): 

76 """Do Sobol' analysis on the Ishigami test function.""" 

77 # --8<-- [start:sobol] 

78 import numpy as np 

79 import uqtils as uq 

80 

81 model = lambda x: uq.ishigami(x)['y'] 

82 sampler = lambda shape: np.random.rand(*shape, 3) * (2 * np.pi) - np.pi 

83 n_samples = 1000 

84 

85 S1, ST = uq.sobol_sa(model, sampler, n_samples) 

86 # --8<-- [end:sobol]