Coverage for src/lsqfitgp/_gvarext/_jacobian.py: 100%

1# lsqfitgp/_gvarext/_jacobian.py 

2# 

3# Copyright (c) 2023, 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 gvar 6 ctx1efabcd

21import numpy 6 ctx1efabcd

22 

23def _getsvec(x): 6 ctx1efabcd

24 """ 

25 Get the sparse vector of derivatives of a GVar. 

26 """ 

27 if isinstance(x, gvar.GVar): 6 ctx1efabcd

28 return x.internaldata[1] 6 ctx1efabcd

29 else: 

30 return gvar.svec(0) 4 ctx1abcd

31 

32def _merge_svec(gvlist, start=None, stop=None): 6 ctx1efabcd

33 if start is None: 6 ctx1efabcd

34 return _merge_svec(gvlist, 0, len(gvlist)) 6 ctx1efabcd

35 n = stop - start 6 ctx1efabcd

36 if n <= 0: 6 ctx1efabcd

37 return gvar.svec(0) 4 ctx1abcd

38 if n == 1: 6 ctx1efabcd

39 return _getsvec(gvlist[start]) 6 ctx1efabcd

40 left = _merge_svec(gvlist, start, start + n // 2) 6 ctx1efabcd

41 right = _merge_svec(gvlist, start + n // 2, stop) 6 ctx1efabcd

42 return left.add(right, 1, 1) 6 ctx1efabcd

43 

44def jacobian(g): 6 ctx1efabcd

45 """ 

46 Extract the jacobian of gvars w.r.t. primary gvars. 

47  

48 Parameters 

49 ---------- 

50 g : array_like 

51 An array of numbers or gvars. 

52  

53 Returns 

54 ------- 

55 jac : array 

56 The shape is g.shape + (m,), where m is the total number of primary 

57 gvars that g depends on. 

58 indices : (m,) int array 

59 The indices that map the last axis of jac to primary gvars in the 

60 global covariance matrix. 

61 

62 See also 

63 -------- 

64 from_jacobian 

65 """ 

66 g = numpy.asarray(g) 6 ctx1efabcd

67 v = _merge_svec(g.flat) 6 ctx1efabcd

68 indices = v.indices() 6 ctx1efabcd

69 jac = numpy.zeros((g.size, len(indices)), float) 6 ctx1efabcd

70 for i, x in enumerate(g.flat): 6 ctx1efabcd

71 v = _getsvec(x) 6 ctx1efabcd

72 ind = numpy.searchsorted(indices, v.indices()) 6 ctx1efabcd

73 jac[i, ind] = v.values() 6 ctx1efabcd

74 jac = jac.reshape(g.shape + indices.shape) 6 ctx1efabcd

75 return jac, indices 6 ctx1efabcd

76 

77def from_jacobian(mean, jac, indices): 6 ctx1efabcd

78 """ 

79 Create new gvars from a jacobian w.r.t. primary gvars. 

80  

81 Parameters 

82 ---------- 

83 mean : array_like 

84 An array of numbers with the means of the new gvars. 

85 jac : mean.shape + (m,) array 

86 The derivatives of each new gvar w.r.t. m primary gvars. 

87 indices : (m,) int array 

88 The indices of the primary gvars. 

89  

90 Returns 

91 ------- 

92 g : mean.shape array 

93 The new gvars. 

94 

95 See also 

96 -------- 

97 jacobian 

98 """ 

99 cov = gvar.gvar.cov 6 ctx1efabcd

100 mean = numpy.asarray(mean) 6 ctx1efabcd

101 shape = mean.shape 6 ctx1efabcd

102 mean = mean.flat 6 ctx1efabcd

103 jac = numpy.asarray(jac) 6 ctx1efabcd

104 jac = jac.reshape(len(mean), len(indices)) 6 ctx1efabcd

105 g = numpy.zeros(len(mean), object) 6 ctx1efabcd

106 for i, jacrow in enumerate(jac): 6 ctx1efabcd

107 der = gvar.svec(len(indices)) 6 ctx1efabcd

108 der._assign(jacrow, indices) 6 ctx1efabcd

109 g[i] = gvar.GVar(mean[i], der, cov) 6 ctx1efabcd

110 return g.reshape(shape) 6 ctx1efabcd