From d5089de1b49bd825a5c8c0ceb715d86a775dec27 Mon Sep 17 00:00:00 2001 From: Walter Simson Date: Tue, 19 Dec 2023 19:44:02 -0800 Subject: [PATCH 1/8] begin to refactor get_win --- kwave/utils/signals.py | 173 ++++++++++++++++++++--------------------- 1 file changed, 83 insertions(+), 90 deletions(-) diff --git a/kwave/utils/signals.py b/kwave/utils/signals.py index 21660c42..deb4683a 100644 --- a/kwave/utils/signals.py +++ b/kwave/utils/signals.py @@ -4,6 +4,8 @@ import numpy as np import scipy +from scipy.signal.windows import general_cosine +from scipy.signal import get_window from numpy.fft import ifftshift, fft, ifft from .conversion import freq2wavenumber @@ -51,8 +53,6 @@ def add_noise(signal: np.ndarray, snr: float, mode="rms"): def get_win(N: Union[int, List[int]], - # TODO: replace and refactor for scipy.signal.get_window - # https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.get_window.html#scipy.signal.get_window type_: str, # TODO change this to enum in the future plot_win: bool = False, param: Optional[float] = None, @@ -96,55 +96,23 @@ def cosine_series(n: int, N: int, coeffs: List[float]) -> np.ndarray: A numpy ndarray containing the calculated series. """ + # return general_cosine(N, coeffs) series = coeffs[0] for index in range(1, len(coeffs)): series = series + (-1) ** index * coeffs[index] * np.cos(index * 2 * np.pi * n / (N - 1)) return series.T - - # Check if N is either `int` or `list of ints` - # assert isinstance(N, int) or isinstance(N, list) or isinstance(N, np.ndarray) - N = np.array(N, dtype=int) - N = N if np.size(N) > 1 else int(N) - - # Check if symmetric is either `bool` or `list of bools` - # assert isinstance(symmetric, int) or isinstance(symmetric, list) - symmetric = np.array(symmetric, dtype=bool) - - # Set default value for `param` if type is one of the special ones - assert not plot_win, NotImplementedError('Plotting is not implemented.') - if type_ == 'Tukey': - if param is None: - param = 0.5 - param = np.clip(param, a_min=0, a_max=1) - elif type_ == 'Blackman': - if param is None: - param = 0.16 - param = np.clip(param, a_min=0, a_max=1) - elif type_ == 'Gaussian': - if param is None: - param = 0.5 - param = np.clip(param, a_min=0, a_max=0.5) - elif type_ == 'Kaiser': - if param is None: - param = 3 - param = np.clip(param, a_min=0, a_max=100) - - # if a non-symmetrical window is required, enlarge the window size (note, - # this expands each dimension individually if symmetric is a vector) - N = N + 1 * (1 - symmetric.astype(int)) - - # if a square window is required, replace grid sizes with smallest size and - # store a copy of the original size - if square and (N.size != 1): - N_orig = np.copy(N) - L = min(N) - N[:] = L - - # create the window - if N.size == 1: + + def _win1D(N: int, type_: str, param: Optional[float] = None) -> np.ndarray: + # TODO: replace and refactor for scipy.signal.get_window + # https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.get_window.html#scipy.signal.get_window + # TODO: what should this behaviour be if N is a list of ints? make windows of multiple lengths? n = np.arange(0, N) + # TODO: find failure cases in test suite when N is zero. + # assert np.all(N) > 1, 'Signal length N must be greater than 1' + if type_ == 'Bartlett': + # win = get_window(str.lower(type_), N, not symmetric) win = (2 / (N - 1) * ((N - 1) / 2 - abs(n - (N - 1) / 2))).T elif type_ == 'Bartlett-Hanning': win = (0.62 - 0.48 * abs(n / (N - 1) - 1 / 2) - 0.38 * np.cos(2 * np.pi * n / (N - 1))).T @@ -155,10 +123,14 @@ def cosine_series(n: int, N: int, coeffs: List[float]) -> np.ndarray: elif type_ == 'Blackman-Nuttall': win = cosine_series(n, N, [0.3635819, 0.4891775, 0.1365995, 0.0106411]) elif type_ == 'Cosine': - win = (np.cos(np.pi * n / (N - 1) - np.pi / 2)).T + # win = scipy.signal.windows.cosine(n) + # w = np.sin(np.pi / M * (np.arange(0, M) + .5)) + win = (np.cos(np.pi * (n / (N - 1) - 0.5))).T elif type_ == 'Flattop': + # win = get_window(str.lower(type_), N, not symmetric) win = cosine_series(n, N, [0.21557895, 0.41663158, 0.277263158, 0.083578947, 0.006947368]) elif type_ == 'Gaussian': + # win = get_window((str.lower(type_), param), N, not symmetric) win = (np.exp(-0.5 * ((n - (N - 1) / 2) / (param * (N - 1) / 2)) ** 2)).T elif type_ == 'HalfBand': win = np.ones(N) @@ -205,34 +177,76 @@ def cosine_series(n: int, N: int, coeffs: List[float]) -> np.ndarray: # calculate the coherent gain cg = win.sum() / N + return win, cg + + def rotate_win(type_: str, N: int) -> np.ndarray: + # create the window in one dimension using getWin recursively + L = np.max(N) + win_lin, _ = get_win(L, type_, param=param) + + # create the reference axis + radius = (L - 1) / 2 + ll = np.linspace(-radius, radius, L) + + # create the 3D window using rotation + axes = [np.linspace(-radius, radius, x) for x in N] # axes describe the axes of the grid + coords = ndgrid(*axes) # coords desribe the coordinates of all points in the grid + r = np.linalg.norm(coords, axis=0) + r[r > radius] = radius + + win_lin = np.squeeze(win_lin) + interp_func = scipy.interpolate.interp1d(ll, win_lin) + win = interp_func(r) + win[r <= radius] = interp_func(r[r <= radius]) + + return win + + # Check if symmetric is either `bool` or `list of bools` + symmetric = np.array(symmetric, dtype=bool) + + # Check if N is either `int` or `list of ints` + N = np.array(N, dtype=int) + N = N if np.size(N) > 1 else int(N) + + # if a non-symmetrical window is required, enlarge the window size (note, + # this expands each dimension individually if symmetric is a vector) + N = N + 1 * (1 - symmetric.astype(int)) + + # if a square window is required, replace grid sizes with smallest size and + # store a copy of the original size + if square and (N.size != 1): + N_orig = np.copy(N) + L = min(N) + N[:] = L + + # Set default value for `param` if type is one of the special ones + assert not plot_win, NotImplementedError('Plotting is not implemented.') + + # Define default parameters and clipping ranges for each window type + window_params = { + 'Tukey': {'default': 0.5, 'min': 0, 'max': 1}, + 'Blackman': {'default': 0.16, 'min': 0, 'max': 1}, + 'Gaussian': {'default': 0.5, 'min': 0, 'max': 0.5}, + 'Kaiser': {'default': 3, 'min': 0, 'max': 100} + } + + # Set and clip param based on the window type + if type_ in window_params: + param = window_params[type_]['default'] if param is None else param + param = np.clip(param, a_min=window_params[type_]['min'], a_max=window_params[type_]['max']) + + # create the window + if N.size == 1: + win, cg = _win1D(N, type_, param=param) elif N.size == 2: # create the 2D window if rotation: - - # create the window in one dimension using getWin recursively - L = max(N) - win_lin, _ = get_win(L, type_, param=param) - win_lin = np.squeeze(win_lin) - - # create the reference axis - radius = (L - 1) / 2 - ll = np.linspace(-radius, radius, L) - - # create the 2D window using rotation - xx = np.linspace(-radius, radius, N[0]) - yy = np.linspace(-radius, radius, N[1]) - [x, y] = ndgrid(xx, yy) - r = np.sqrt(x ** 2 + y ** 2) - r[r > radius] = radius - interp_func = scipy.interpolate.interp1d(ll, win_lin) - win = interp_func(r) - win[r <= radius] = interp_func(r[r <= radius]) - + win = rotate_win(type_, N) else: # create the window in each dimension using getWin recursively - win_x, _ = get_win(N[0], type_, param=param) - win_y, _ = get_win(N[1], type_, param=param) + win_x, _ = _win1D(N[0], type_, param=param) + win_y, _ = _win1D(N[1], type_, param=param) # create the 2D window using the outer product win = (win_y * win_x.T).T @@ -246,30 +260,8 @@ def cosine_series(n: int, N: int, coeffs: List[float]) -> np.ndarray: elif N.size == 3: # create the 3D window if rotation: - - # create the window in one dimension using getWin recursively - L = N.max() - win_lin, _ = get_win(L, type_, param=param) - - # create the reference axis - radius = (L - 1) / 2 - ll = np.linspace(-radius, radius, L) - - # create the 3D window using rotation - xx = np.linspace(-radius, radius, N[0]) - yy = np.linspace(-radius, radius, N[1]) - zz = np.linspace(-radius, radius, N[2]) - [x, y, z] = ndgrid(xx, yy, zz) - r = np.sqrt(x ** 2 + y ** 2 + z ** 2) - r[r > radius] = radius - - win_lin = np.squeeze(win_lin) - interp_func = scipy.interpolate.interp1d(ll, win_lin) - win = interp_func(r) - win[r <= radius] = interp_func(r[r <= radius]) - + win = rotate_win(type_, N) else: - # create the window in each dimension using getWin recursively win_x, _ = get_win(N[0], type_, param=param) win_y, _ = get_win(N[1], type_, param=param) @@ -294,6 +286,7 @@ def cosine_series(n: int, N: int, coeffs: List[float]) -> np.ndarray: # enlarge the window if required if square and (N.size != 1): + # TODO: use expand matrix or a padding opperation to abstract this logic away L = N[0] win_sq = win win = np.zeros(N_orig) From 9a70d5f271712dd8a3cefbd829cb5b8cffd17336 Mon Sep 17 00:00:00 2001 From: Walter Simson Date: Tue, 19 Dec 2023 19:45:12 -0800 Subject: [PATCH 2/8] remove unused imports --- kwave/utils/signals.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/kwave/utils/signals.py b/kwave/utils/signals.py index deb4683a..e979f049 100644 --- a/kwave/utils/signals.py +++ b/kwave/utils/signals.py @@ -4,8 +4,6 @@ import numpy as np import scipy -from scipy.signal.windows import general_cosine -from scipy.signal import get_window from numpy.fft import ifftshift, fft, ifft from .conversion import freq2wavenumber From 3ee7eccdf63466ae1d19c52b8baa9071c3ca8317 Mon Sep 17 00:00:00 2001 From: Walter Simson Date: Tue, 19 Dec 2023 20:02:49 -0800 Subject: [PATCH 3/8] remove for loop --- kwave/utils/signals.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/kwave/utils/signals.py b/kwave/utils/signals.py index e979f049..889bad29 100644 --- a/kwave/utils/signals.py +++ b/kwave/utils/signals.py @@ -265,13 +265,7 @@ def rotate_win(type_: str, N: int) -> np.ndarray: win_y, _ = get_win(N[1], type_, param=param) win_z, _ = get_win(N[2], type_, param=param) - # create the 2D window using the outer product - win_2D = (win_x * win_z.T) - - # create the 3D window - win = np.zeros((N[0], N[1], N[2])) - for index in range(0, N[1]): - win[:, index, :] = win_2D[:, :] * win_y[index] + win = win_x[:, np.newaxis] * win_z.T * win_y[np.newaxis, :] # trim the window if required N = N - 1 * (1 - np.array(symmetric).astype(int)) From f8c20cbb029cd1eedf2922f3f9fc8d2b1eb2ee97 Mon Sep 17 00:00:00 2001 From: Walter Simson Date: Tue, 19 Dec 2023 20:19:41 -0800 Subject: [PATCH 4/8] further simplify --- kwave/utils/signals.py | 31 ++++++++++++++++++++++++------- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/kwave/utils/signals.py b/kwave/utils/signals.py index 889bad29..86cfff82 100644 --- a/kwave/utils/signals.py +++ b/kwave/utils/signals.py @@ -100,6 +100,26 @@ def cosine_series(n: int, N: int, coeffs: List[float]) -> np.ndarray: series = series + (-1) ** index * coeffs[index] * np.cos(index * 2 * np.pi * n / (N - 1)) return series.T + def _multi_dim_window(N, type_, param, symmetric): + """ + Generates a multi-dimensional window. + + Args: + N: List of dimensions. + type_: Window type. + param: Parameter for specific window types. + symmetric: Symmetric flag for each dimension. + + Returns: + Multi-dimensional window. + """ + windows = [_win1D(dim, type_, param) for dim in N] + multi_dim_win = windows[0][0] + + for win in windows[1:]: + multi_dim_win = np.expand_dims(multi_dim_win, axis=-1) * win[0].T + + return multi_dim_win def _win1D(N: int, type_: str, param: Optional[float] = None) -> np.ndarray: # TODO: replace and refactor for scipy.signal.get_window # https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.get_window.html#scipy.signal.get_window @@ -243,11 +263,10 @@ def rotate_win(type_: str, N: int) -> np.ndarray: win = rotate_win(type_, N) else: # create the window in each dimension using getWin recursively - win_x, _ = _win1D(N[0], type_, param=param) - win_y, _ = _win1D(N[1], type_, param=param) - + windows1D = [_win1D(dim, type_, param=param)[0] for dim in N] + # create the 2D window using the outer product - win = (win_y * win_x.T).T + win = np.outer(*windows1D) # trim the window if required N = N - 1 * (1 - np.array(symmetric).astype(int)) @@ -261,9 +280,7 @@ def rotate_win(type_: str, N: int) -> np.ndarray: win = rotate_win(type_, N) else: # create the window in each dimension using getWin recursively - win_x, _ = get_win(N[0], type_, param=param) - win_y, _ = get_win(N[1], type_, param=param) - win_z, _ = get_win(N[2], type_, param=param) + [win_x, win_y, win_z] = [_win1D(dim, type_, param=param)[0] for dim in N] win = win_x[:, np.newaxis] * win_z.T * win_y[np.newaxis, :] From 2789f593efbf9ea08fa5e1a5b3918484d0c6ec9d Mon Sep 17 00:00:00 2001 From: Walter Simson Date: Tue, 19 Dec 2023 20:23:35 -0800 Subject: [PATCH 5/8] clean up --- kwave/utils/signals.py | 23 ++--------------------- 1 file changed, 2 insertions(+), 21 deletions(-) diff --git a/kwave/utils/signals.py b/kwave/utils/signals.py index 86cfff82..5f55245e 100644 --- a/kwave/utils/signals.py +++ b/kwave/utils/signals.py @@ -99,27 +99,8 @@ def cosine_series(n: int, N: int, coeffs: List[float]) -> np.ndarray: for index in range(1, len(coeffs)): series = series + (-1) ** index * coeffs[index] * np.cos(index * 2 * np.pi * n / (N - 1)) return series.T - - def _multi_dim_window(N, type_, param, symmetric): - """ - Generates a multi-dimensional window. - Args: - N: List of dimensions. - type_: Window type. - param: Parameter for specific window types. - symmetric: Symmetric flag for each dimension. - - Returns: - Multi-dimensional window. - """ - windows = [_win1D(dim, type_, param) for dim in N] - multi_dim_win = windows[0][0] - - for win in windows[1:]: - multi_dim_win = np.expand_dims(multi_dim_win, axis=-1) * win[0].T - - return multi_dim_win + def _win1D(N: int, type_: str, param: Optional[float] = None) -> np.ndarray: # TODO: replace and refactor for scipy.signal.get_window # https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.get_window.html#scipy.signal.get_window @@ -265,7 +246,7 @@ def rotate_win(type_: str, N: int) -> np.ndarray: # create the window in each dimension using getWin recursively windows1D = [_win1D(dim, type_, param=param)[0] for dim in N] - # create the 2D window using the outer product + # # create the 2D window using the outer product win = np.outer(*windows1D) # trim the window if required From 90289d49c5e50d0504fe2f6534b50084173d10ac Mon Sep 17 00:00:00 2001 From: Walter Simson Date: Tue, 19 Dec 2023 20:39:07 -0800 Subject: [PATCH 6/8] rename for clarity --- kwave/utils/signals.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/kwave/utils/signals.py b/kwave/utils/signals.py index 5f55245e..4fb5e043 100644 --- a/kwave/utils/signals.py +++ b/kwave/utils/signals.py @@ -179,26 +179,26 @@ def _win1D(N: int, type_: str, param: Optional[float] = None) -> np.ndarray: return win, cg def rotate_win(type_: str, N: int) -> np.ndarray: - # create the window in one dimension using getWin recursively + # Compute the radius and create a linear window L = np.max(N) - win_lin, _ = get_win(L, type_, param=param) - - # create the reference axis radius = (L - 1) / 2 - ll = np.linspace(-radius, radius, L) + linear_window = get_win(L, type_, param=param)[0].squeeze() + + # Create the reference axis for interpolation + reference_axis = np.linspace(-radius, radius, L) + interp_func = scipy.interpolate.interp1d(reference_axis, linear_window) + + # Create the grid + axes = [np.linspace(-radius, radius, dim_size) for dim_size in N] + grid = ndgrid(*axes) - # create the 3D window using rotation - axes = [np.linspace(-radius, radius, x) for x in N] # axes describe the axes of the grid - coords = ndgrid(*axes) # coords desribe the coordinates of all points in the grid - r = np.linalg.norm(coords, axis=0) - r[r > radius] = radius + # Compute radial distances and interpolate the window values + radial_distances = np.linalg.norm(grid, axis=0) + radial_distances_clipped = np.clip(radial_distances, a_min=None, a_max=radius) + window = interp_func(radial_distances_clipped) - win_lin = np.squeeze(win_lin) - interp_func = scipy.interpolate.interp1d(ll, win_lin) - win = interp_func(r) - win[r <= radius] = interp_func(r[r <= radius]) + return window - return win # Check if symmetric is either `bool` or `list of bools` symmetric = np.array(symmetric, dtype=bool) From 0bfd4c3b5e227abc08ad09e0a1445bd90724ad39 Mon Sep 17 00:00:00 2001 From: Walter Simson Date: Tue, 19 Dec 2023 21:08:27 -0800 Subject: [PATCH 7/8] minor refacotring of Tukey --- kwave/utils/signals.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/kwave/utils/signals.py b/kwave/utils/signals.py index 4fb5e043..7913e836 100644 --- a/kwave/utils/signals.py +++ b/kwave/utils/signals.py @@ -4,6 +4,7 @@ import numpy as np import scipy +from scipy.signal import get_window from numpy.fft import ifftshift, fft, ifft from .conversion import freq2wavenumber @@ -155,16 +156,18 @@ def _win1D(N: int, type_: str, param: Optional[float] = None) -> np.ndarray: elif type_ == 'Nuttall': win = cosine_series(n, N, [0.3635819, 0.4891775, 0.1365995, 0.0106411]) elif type_ == 'Rectangular': - win = np.ones(N) + win = get_window(str.lower(type_), N, not symmetric) elif type_ == 'Triangular': + # win = get_window(str.lower(type_[:-4]), N, not symmetric) win = (2 / N * (N / 2 - abs(n - (N - 1) / 2))).T elif type_ == 'Tukey': - win = np.ones((N, 1)) - index = np.arange(0, (N - 1) * param / 2 + 1e-8) + # win = get_window((str.lower(type_), param), N, not symmetric) + rise_func = lambda x: 0.5 * (1 + np.cos(2 * np.pi / param * (x - param / 2))) + win = np.ones((N,)) + index = np.arange((N - 1) * param / 2 + 1e-8) param = param * N - win[0: len(index)] = 0.5 * (1 + np.cos(2 * np.pi / param * (index - param / 2)))[:, None] - win[np.arange(-1, -len(index) - 1, -1)] = win[0:len(index)] - win = win.squeeze(axis=-1) + win[: len(index)] = rise_func(index) # left side + win[-len(index):] = np.flip(win[0:len(index)]) else: raise ValueError(f'Unknown window type: {type_}') From 27171f8b2a04f6fba458005e2a2615932fae5c4c Mon Sep 17 00:00:00 2001 From: Walter Simson Date: Tue, 19 Dec 2023 21:19:19 -0800 Subject: [PATCH 8/8] ruff fix --- kwave/utils/signals.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/kwave/utils/signals.py b/kwave/utils/signals.py index 7913e836..dd8a0937 100644 --- a/kwave/utils/signals.py +++ b/kwave/utils/signals.py @@ -162,7 +162,8 @@ def _win1D(N: int, type_: str, param: Optional[float] = None) -> np.ndarray: win = (2 / N * (N / 2 - abs(n - (N - 1) / 2))).T elif type_ == 'Tukey': # win = get_window((str.lower(type_), param), N, not symmetric) - rise_func = lambda x: 0.5 * (1 + np.cos(2 * np.pi / param * (x - param / 2))) + def rise_func(x): + return 0.5 * (1 + np.cos(2 * np.pi / param * (x - param / 2))) win = np.ones((N,)) index = np.arange((N - 1) * param / 2 + 1e-8) param = param * N