Coverage for src/lsqfitgp/_special/_bernoulli.py: 64%

1# lsqfitgp/_special/_bernoulli.py 

2# 

3# Copyright (c) 2022, 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 functools 3 ctx1bca

21 

22import jax 3 ctx1bca

23import numpy 3 ctx1bca

24from jax import numpy as jnp 3 ctx1bca

25from jax.scipy import special as jspecial 3 ctx1bca

26from scipy import special 3 ctx1bca

27 

28 

29def periodic_bernoulli(n, x): 3 ctx1bca

30 n = int(n) 1 ctx1a

31 bernoulli = special.bernoulli(n) 1 ctx1a

32 k = numpy.arange(n + 1) 1 ctx1a

33 binom = special.binom(n, k) 1 ctx1a

34 coeffs = binom[::-1] * bernoulli 1 ctx1a

35 x = x % 1 1 ctx1a

36 cond = x < 0.5 1 ctx1a

37 x = jnp.where(cond, x, 1 - x) 1 ctx1a

38 out = jnp.polyval(coeffs, x) 1 ctx1a

39 if n % 2 == 1: 1 ctx1a

40 out = out * jnp.where(cond, 1, -1) 1 ctx1a

41 return out 1 ctx1a

42 

43 

44@functools.partial(jax.custom_jvp, nondiff_argnums=(0,)) 3 ctx1bca

45def scaled_periodic_bernoulli(n, x): 3 ctx1bca

46 """Periodic Bernoulli polynomial scaled such that B_n(0) = ζ(n).""" 

47 tau = 2 * jnp.pi 

48 lognorm = n * jnp.log(tau) - jspecial.gammaln(n + 1) 

49 norm = jnp.exp(lognorm) / 2 

50 cond = n < 60 

51 smalls = norm * periodic_bernoulli(n if cond else 1, x) 

52 # n -> ∞: -Re e^2πix / i^n 

53 # don't use 1j ** n because it is very inaccurate 

54 arg = tau * x 

55 sign = jnp.where((n // 2) % 2, 1, -1) 

56 larges = sign * jnp.where(n % 2, jnp.sin(arg), jnp.cos(arg)) 

57 return jnp.where(cond, smalls, larges) 

58 

59 

60@scaled_periodic_bernoulli.defjvp 3 ctx1bca

61def _scaled_periodic_bernoulli_jvp(n, primals, tangents): 3 ctx1bca

62 (x,) = primals 

63 (xt,) = tangents 

64 primal = scaled_periodic_bernoulli(n, x) 

65 tangent = 2 * jnp.pi * scaled_periodic_bernoulli(n - 1, x) * xt 

66 return primal, tangent