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
« prev ^ index » next coverage.py v6.5.0, created at 2022-10-24 11:55 -0700
1from typing import Optional
3import numpy as np
4from itertools import compress
6from numpy.fft import ifftshift, fft, ifft
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
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:])]
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
38 Returns:
40 """
41 result = number / divider
42 after_decimal = result % 1
43 return after_decimal == 0
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
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)
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
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.')
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