Coverage for src/lsqfitgp/_fastraniter.py: 95%

1# lsqfitgp/_fastraniter.py 

2# 

3# Copyright (c) 2020, 2022, 2023, 2024, 2026, Giacomo Petrillo 

4# 

5# This file is part of lsqfitgp. 

6# 

7# lsqfitgp is free software: you can redistribute it and/or modify 

8# it under the terms of the GNU General Public License as published by 

9# the Free Software Foundation, either version 3 of the License, or 

10# (at your option) any later version. 

11# 

12# lsqfitgp is distributed in the hope that it will be useful, 

13# but WITHOUT ANY WARRANTY; without even the implied warranty of 

14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

15# GNU General Public License for more details. 

16# 

17# You should have received a copy of the GNU General Public License 

18# along with lsqfitgp. If not, see <http://www.gnu.org/licenses/>. 

19 

20import itertools 3 ctx1cba

21 

22import gvar 3 ctx1cba

23import numpy 3 ctx1cba

24 

25from lsqfitgp import _linalg 3 ctx1cba

26 

27 

28def _toslice(s): 3 ctx1cba

29 if isinstance(s, slice): 2 ctx1ba

30 return s 2 ctx1ba

31 if isinstance(s, int): 31 ↛ 33line 31 didn't jump to line 33 because the condition on line 31 was always true2 ctx1ba

32 return slice(s, s + 1) 2 ctx1ba

33 msg = f'cannot convert {s!r} to slice' 

34 raise TypeError(msg) 

35 

36 

37def raniter(mean, cov, n=None, eps=None, rng=None): 3 ctx1cba

38 """ 

39 Take random samples from a multivariate Gaussian. 

40 

41 This generator mimics the interface of `gvar.raniter`, but takes as input 

42 the mean and covariance separately instead of a collection of gvars. 

43 

44 Parameters 

45 ---------- 

46 mean : scalar, array, or dictionary of scalars/arrays 

47 The mean of the Gaussian distribution. 

48 cov : scalar, array, or dictionary of scalars/arrays 

49 The covariance matrix. If ``mean`` is a dictionary, ``cov`` must be a 

50 dictionary with pair of keys from ``mean`` as keys. 

51 n : int, optional 

52 The maximum number of iterations. Default unlimited. 

53 eps : float, optional 

54 Used to correct the eigenvalues of the covariance matrix to handle 

55 non-positivity due to roundoff, relative to the largest eigenvalue. 

56 Default is number of variables times floating point epsilon. 

57 rng : seed or random generator, optional 

58 ``rng`` is passed through `numpy.random.default_rng` to produce a random 

59 number generator. 

60 

61 Yields 

62 ------ 

63 samp : scalar, array, or dictionary of scalars/arrays 

64 The random sample in the same format of ``mean``. 

65 

66 Examples 

67 -------- 

68 >>> mean = {'a': np.arange(3)} 

69 >>> cov = {('a', 'a'): np.eye(3)} 

70 >>> for sample in lgp.raniter(mean, cov, 3): 

71 >>> print(sample) 

72 """ # noqa: DOC403 

73 # convert mean and cov to 1d and 2d arrays 

74 if hasattr(mean, 'keys'): # a dict or gvar.BufferDict 2 ctx1ba

75 if not hasattr(mean, 'buf'): 2 ctx1ba

76 mean = gvar.BufferDict(mean) 2 ctx1ba

77 flatmean = mean.buf 2 ctx1ba

78 squarecov = numpy.empty((len(flatmean), len(flatmean))) 2 ctx1ba

79 for k1 in mean: 2 ctx1ba

80 slic1 = _toslice(mean.slice(k1)) 2 ctx1ba

81 for k2 in mean: 2 ctx1ba

82 slic2 = _toslice(mean.slice(k2)) 2 ctx1ba

83 sqshape = (slic1.stop - slic1.start, slic2.stop - slic2.start) 2 ctx1ba

84 squarecov[slic1, slic2] = cov[k1, k2].reshape(sqshape) 2 ctx1ba

85 else: # an array or scalar 

86 mean = numpy.array(mean, copy=False) 2 ctx1ba

87 cov = numpy.array(cov, copy=False) 2 ctx1ba

88 flatmean = mean.reshape(-1) 2 ctx1ba

89 squarecov = cov.reshape(len(flatmean), len(flatmean)) 2 ctx1ba

90 

91 # decompose the covariance matrix 

92 try: 2 ctx1ba

93 covdec = _linalg.Chol(squarecov, epsrel='auto' if eps is None else eps) 2 ctx1ba

94 except numpy.linalg.LinAlgError as e: 1 ctx1a

95 msg = f'covariance matrix not positive definite with eps={eps}' 1 ctx1a

96 raise numpy.linalg.LinAlgError(msg) from e 1 ctx1a

97 

98 # get random number generator 

99 rng = numpy.random.default_rng(rng) 2 ctx1ba

100 

101 # take samples 

102 iterable = itertools.count() if n is None else range(n) 2 ctx1ba

103 for _ in iterable: 2 ctx1ba

104 iidsamp = rng.standard_normal(covdec.m) 2 ctx1ba

105 samp = flatmean + covdec.correlate(iidsamp) 2 ctx1ba

106 

107 # pack the samples with the original shape 

108 if hasattr(mean, 'keys'): 2 ctx1ba

109 samp = gvar.BufferDict(mean, buf=samp) 2 ctx1ba

110 else: 

111 samp = samp.reshape(mean.shape) if mean.shape else samp.item 2 ctx1ba

112 

113 yield samp 2 ctx1ba

114 

115 

116def sample(*args, **kw): 3 ctx1cba

117 """Shortcut for ``next(raniter(..., n=1))``.""" 

118 return next(raniter(*args, n=1, **kw)) 2 ctx1ba