Coverage for src/lsqfitgp/_kernels/_wendland.py: 100%

1# lsqfitgp/_kernels/_wendland.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 

20from jax import numpy as jnp 6 ctx1efabcd

21 

22from .. import _jaxext 6 ctx1efabcd

23from .._Kernel import stationarykernel, isotropickernel 6 ctx1efabcd

24 

25def _wendland_derivable(k=0, **_): 6 ctx1efabcd

26 return k 4 ctx1abcd

27 

28def _wendland_maxdim(k=0, alpha=1): 6 ctx1efabcd

29 with _jaxext.skipifabstract(): 4 ctx1abcd

30 return int(jnp.floor(2 * alpha - 1)) 4 ctx1abcd

31 

32@isotropickernel(input='posabs', derivable=_wendland_derivable, maxdim=_wendland_maxdim) 6 ctx1efabcd

33def Wendland(r, k=0, alpha=1): 6 ctx1efabcd

34 """ 

35 Wendland kernel. 

36  

37 .. math:: 

38 k(r) &= \\frac1{B(2k+1,\\nu)} 

39 \\int_r^\\infty \\mathrm du\\, (u^2 - r^2)^k (1 - u)_+^{\\nu-1}, \\\\ 

40 \\quad k &\\in \\mathbb N,\\ \\nu = k + \\alpha,\\ \\alpha \\ge 1. 

41  

42 An isotropic kernel with finite support. The covariance is nonzero only 

43 when the distance between the points is less than 1. Parameter :math:`k \\in \\{0, 

44 1, 2, 3\\}` sets the differentiability, while the maximum dimensionality the 

45 kernel can be used in is :math:`\\lfloor 2\\alpha-1 \\rfloor`. Default is 

46 :math:`k = 0` (non derivable), :math:`\\alpha = 1` (can be used only in 

47 1D). 

48  

49 Reference: Gneiting (2002), Wendland (2004, p. 128), Rasmussen and Williams 

50 (2006, p. 87), Porcu, Furrer and Nychka (2020, p. 4). 

51  

52 """ 

53 

54 # TODO compute the kernel only on the nonzero points. 

55 

56 # TODO find the nonzero points in O(nlogn) instead of O(n^2) by sorting 

57 # the inputs, and output a sparse matrix => on second thought this should 

58 # be a general mechanism implemented in GP that gives sparse x and y to 

59 # the kernel 

60 

61 with _jaxext.skipifabstract(): 4 ctx1abcd

62 D = _wendland_maxdim(k, alpha) 4 ctx1abcd

63 assert D >= 1, D 4 ctx1abcd

64 

65 if k == 0: 4 ctx1abcd

66 poly = [ 4 ctx1abcd

67 [1], 

68 ] 

69 elif k == 1: 4 ctx1abcd

70 poly = [ 4 ctx1abcd

71 [1, 1], 

72 [1], 

73 ] 

74 elif k == 2: 4 ctx1abcd

75 poly = [ 4 ctx1abcd

76 [1/3, 4/3, 1], 

77 [1, 2], 

78 [1], 

79 ] 

80 elif k == 3: 4 ctx1abcd

81 poly = [ 4 ctx1abcd

82 [1/15, 3/5, 23/15, 1], 

83 [2/5, 12/5, 3], 

84 [1, 3], 

85 [1], 

86 ] 

87 else: 

88 raise NotImplementedError 4 ctx1abcd

89 

90 nu = k + alpha 4 ctx1abcd

91 coeffs = jnp.array([jnp.polyval(jnp.array(pj), nu) for pj in poly]) 4 ctx1abcd

92 poly = jnp.polyval(coeffs, r) 4 ctx1abcd

93 return jnp.where(r < 1, (1 - r) ** (nu + k) * poly, 0) 4 ctx1abcd

94 

95 

96# adapted from the "GP-circular" example in the PyMC documentation 

97 

98# TODO maxdim actually makes sense only for isotropic. I need a way to say 

99# structured/non structured. Maybe all this should just live in the test suite. 

100 

101# TODO Any stationary kernel supported on [0, pi] would be fine combined with 

102# the geodesic distance. Use the generic wendland kernel. Options: 

103# 1) add here the parameters of Wendland 

104# 2) add a distance option in stationary to use the angular distance, then 

105# let the user apply it to wendland => the problem is that the user would 

106# need to be careful with positiveness, while wendland checks it for him 

107 

108@stationarykernel(derivable=1, maxdim=1, input='posabs') 6 ctx1efabcd

109def Circular(delta, tau=4, c=1/2): 6 ctx1efabcd

110 """ 

111 Circular kernel. 

112  

113 .. math:: k(x, y) &= W_c(d_{\\text{geo}}(x, y)), \\\\ 

114 W_c(t) &= \\left(1 + \\tau\\frac tc\\right) 

115 \\left(1 - \\frac tc\\right)^\\tau_+, 

116 \\quad c \\in (0, 1/2], \\tau \\ge 4, \\\\ 

117 d_{\\text{geo}}(x, y) &= \\arccos\\cos(2\\pi(x-y)). 

118  

119 It is a stationary periodic kernel with period 1. 

120  

121 Reference: Padonou and Roustant (2016). 

122 """ 

123 with _jaxext.skipifabstract(): 4 ctx1abcd

124 assert tau >= 4, tau 4 ctx1abcd

125 assert 0 < c <= 1/2, c 4 ctx1abcd

126 x = delta % 1 4 ctx1abcd

127 t = jnp.minimum(x, 1 - x) 4 ctx1abcd

128 return (1 + tau * t / c) * jnp.maximum(1 - t / c, 0) ** tau 4 ctx1abcd