Coverage for kwave/utils/mathutils.py: 18%

1 from typing import Optional

2

3 import numpy as np

4 from itertools import compress

5

6 from numpy.fft import ifftshift, fft, ifft

7

8

9 def largest_prime_factor(n):

10 i = 2

11 while i * i <= n:

12 if n % i:

13 i += 1

14 else:

15 n //= i

16 return n

17

18

19 def rwh_primes(n):

20 # https://stackoverflow.com/questions/2068372/fastest-way-to-list-all-primes-below-n-in-python/3035188#3035188

21 """ Returns a list of primes < n for n > 2 """

22 sieve = bytearray([True]) * (n//2+1)

23 for i in range(1, int(n**0.5)//2+1):

24 if sieve[i]:

25 sieve[2*i*(i+1)::2*i+1] = bytearray((n//2-2*i*(i+1))//(2*i+1)+1)

26 return [2, *compress(range(3, n, 2), sieve[1:])]

27

28

29 def check_divisible(number: float, divider: float) -> bool:

30 """

31 Checks whether number is divisible by divider without any remainder

32 Why do we need such a function? -> Because due to floating point precision we

33 experience rounding errors while using standard modulo operator with floating point numbers

34 Args:

35 number: Number that's supposed to be divided

36 divider: Divider that should devide the number

37

38 Returns:

39

40 """

41 result = number / divider

42 after_decimal = result % 1

43 return after_decimal == 0

44

45

46 def fourier_shift(

47 data: np.ndarray,

48 shift: float,

49 shift_dim: Optional[int] = None

50) -> np.ndarray:

51 if shift_dim is None:

52 shift_dim = data.ndim - 1

53 if (shift_dim == 1) and (data.shape[1] == 1):

54 # row vector

55 shift_dim = 0

56 else:

57 # subtract 1 in order to keep function interface compatible with matlab

58 shift_dim -= 1

59

60 try:

61 N = data.shape[shift_dim]

62 except IndexError:

63 print('aaa')

64 if N % 2 == 0:

65 # grid dimension has an even number of points

66 k_vec = (2 * np.pi) * ( np.arange(-N // 2, N // 2) / N)

67 else:

68 # grid dimension has an odd number of points

69 k_vec = (2 * np.pi) * ( np.arange(-(N -1) // 2, N // 2 + 1) / N)

70

71 # force middle value to be zero in case 1/N is a recurring number and the

72 # series doesn't give exactly zero

73 k_vec[N // 2] = 0

74

75 # put the wavenumber vector in the correct orientation for use with bsxfun

76 reshape_dims_to = [1] * data.ndim

77 if 0 <= shift_dim <= 3:

78 reshape_dims_to[shift_dim] = -1

79 k_vec = np.reshape(k_vec, reshape_dims_to)

80 else:

81 raise ValueError('Input dim must be 0, 1, 2 or 3.')

82

83 # shift the input using a Fourier interpolant

84 part_1 = ifftshift(np.exp(1j * k_vec * shift))

85 part_2 = fft(data, axis=shift_dim)

86 part_1_times_2 = part_1 * part_2

87 result = ifft(part_1_times_2, axis=shift_dim).real

88 return result