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

1# lsqfitgp/_kernels/_zeta.py 

2# 

3# Copyright (c) 2023, 2024, 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 functools 6 ctx1fabcde

21 

22from jax import numpy as jnp 6 ctx1fabcde

23 

24from .. import _special 6 ctx1fabcde

25from .. import _jaxext 6 ctx1fabcde

26from .. import _Kernel 6 ctx1fabcde

27 

28def check_nu(nu): 6 ctx1fabcde

29 with _jaxext.skipifabstract(): 5 ctx1abcde

30 assert 0 <= nu < jnp.inf, nu 5 ctx1abcde

31 

32def zeta_derivable(*, nu): 6 ctx1fabcde

33 check_nu(nu) 5 ctx1abcde

34 with _jaxext.skipifabstract(): 5 ctx1abcde

35 return int(max(0, jnp.ceil(nu) - 1)) 5 ctx1abcde

36 

37@_Kernel.crosskernel(bases=(_Kernel.AffineSpan, _Kernel.StationaryKernel,), maxdim=1, derivable=zeta_derivable) 6 ctx1fabcde

38def Zeta(delta, *, nu, **_): 6 ctx1fabcde

39 r""" 

40  

41 Zeta kernel. 

42  

43 .. math:: 

44 k(\Delta) 

45 &= \frac{\Re F(\Delta, s)}{\zeta(s)} = 

46 \qquad (s = 1 + 2 \nu, \quad \nu \ge 0) \\ 

47 &= \frac1{\zeta(s)} \sum_{k=1}^\infty 

48 \frac {\cos(2\pi k\Delta)} {k^s} = \\ 

49 &= -(-1)^{s/2} 

50 \frac {(2\pi)^s} {2s!} 

51 \frac {\tilde B_s(\Delta)} {\zeta(s)} 

52 \quad \text{for even integer $s$.} 

53  

54 It is equivalent to fitting with a Fourier series of period 1 with 

55 independent priors on the coefficients with mean zero and variance 

56 :math:`1/(\zeta(s)k^s)` for the :math:`k`-th term. Analogously to 

57 :class:`Matern`, the process is :math:`\lceil\nu\rceil - 1` times 

58 derivable, and the highest derivative is continuous iff :math:`\nu\bmod 1 

59 \ge 1/2`. 

60  

61 The :math:`k = 0` term is not included in the summation, so the mean of the 

62 process over one period is forced to be zero. 

63 

64 Reference: Petrillo (2022). 

65  

66 """ 

67 check_nu(nu) 5 ctx1abcde

68 s = 1 + 2 * nu 5 ctx1abcde

69 nupos = _special.periodic_zeta(delta, s) / _special.zeta(s) 5 ctx1abcde

70 nuzero = jnp.where(delta % 1, 0, 1) 5 ctx1abcde

71 return jnp.where(s > 1, nupos, nuzero) 5 ctx1abcde

72 

73 # return -(-1) ** (s // 2) * _special.scaled_periodic_bernoulli(s, delta) / jspecial.zeta(s, 1) 

74 

75 

76 

77 

78@_Kernel.kernel(maxdim=1, derivable=False) 6 ctx1fabcde

79def ZetaFourier(k, q, *, nu, lloc, rloc, lscale, rscale, offset, ampl): 6 ctx1fabcde

80 check_nu(nu) 5 ctx1abcde

81 s = 1 + 2 * nu 5 ctx1abcde

82 lorder = jnp.ceil(k / 2) 5 ctx1abcde

83 rorder = jnp.ceil(q / 2) 5 ctx1abcde

84 lodd = k % 2 5 ctx1abcde

85 rodd = q % 2 5 ctx1abcde

86 var = ampl / (lorder ** s * _special.zeta(s)) 5 ctx1abcde

87 arg = 2 * jnp.pi * lorder * (lloc / lscale - rloc / rscale) 5 ctx1abcde

88 return jnp.where(lorder == rorder, 5 ctx1abcde

89 jnp.where(lodd == rodd, 

90 jnp.where(lorder, var * jnp.cos(arg), offset), 

91 var * jnp.sin(arg) * jnp.where(lodd, 1, -1), 

92 ), 

93 0, 

94 ) 

95 

96def crosszeta_derivable(*, nu, **_): 6 ctx1fabcde

97 return 0, zeta_derivable(nu=nu) 5 ctx1abcde

98 

99@_Kernel.crosskernel(bases=(_Kernel.PreservedBySwap, _Kernel.CrossKernel), maxdim=1, derivable=crosszeta_derivable) 6 ctx1fabcde

100def CrossZetaFourier(k, y, *, nu, lloc, rloc, lscale, rscale, offset, ampl): 6 ctx1fabcde

101 check_nu(nu) 5 ctx1abcde

102 s = 1 + 2 * nu 5 ctx1abcde

103 order = jnp.ceil(k / 2) 5 ctx1abcde

104 odd = k % 2 5 ctx1abcde

105 var = ampl / (order ** s * _special.zeta(s)) 5 ctx1abcde

106 arg = 2 * jnp.pi * order * (lloc / lscale + (y - rloc) / rscale) 5 ctx1abcde

107 return jnp.where(odd, 5 ctx1abcde

108 var * jnp.sin(arg), 

109 jnp.where(order, var * jnp.cos(arg), offset), 

110 ) 

111 

112fourier_doc = r""" 6 ctx1fabcde

113 

114Compute the Fourier series transform of the function. 

115 

116.. math:: 

117 

118 T(f)(k) = \begin{cases} 

119 \frac2T \int_0^T \mathrm dx\, f(x) 

120 \cos\left(\frac{2\pi}T \frac k2 x\right) 

121 & \text{if $k$ is even} \\ 

122 \frac2T \int_0^T \mathrm dx\, f(x) 

123 \sin\left(\frac{2\pi}T \frac{k+1}2 x\right) 

124 & \text{if $k$ is odd} 

125 \end{cases} 

126  

127The period :math:`T` is 1. 

128 

129""" 

130 

131def fourier_argparser(do): 6 ctx1fabcde

132 return do if do else None 5 ctx1abcde

133 

134def translkw(*, dynkw, **initkw): 6 ctx1fabcde

135 return dict(**dynkw, **initkw) 5 ctx1abcde

136 

137Zeta.make_linop_family('fourier', ZetaFourier, CrossZetaFourier, translkw=translkw, doc=fourier_doc, argparser=fourier_argparser) 6 ctx1fabcde

138