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

45 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-10-24 11:55 -0700

1from typing import Optional 

2 

3import numpy as np 

4from itertools import compress 

5 

6from numpy.fft import ifftshift, fft, ifft 

7 

8 

9def 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 

19def 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 

29def 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 

46def 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