Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor get win #232

Draft
wants to merge 8 commits into
base: master
Choose a base branch
from
Draft

Refactor get win #232

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
d5089de
begin to refactor get_win
waltsims Dec 20, 2023
9a70d5f
remove unused imports
waltsims Dec 20, 2023
3ee7ecc
remove for loop
waltsims Dec 20, 2023
f8c20cb
further simplify
waltsims Dec 20, 2023
2789f59
clean up
waltsims Dec 20, 2023
90289d4
rename for clarity
waltsims Dec 20, 2023
0bfd4c3
minor refacotring of Tukey
waltsims Dec 20, 2023
27171f8
ruff fix
waltsims Dec 20, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
203 changes: 95 additions & 108 deletions kwave/utils/signals.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -51,8 +52,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,
Expand Down Expand Up @@ -96,55 +95,24 @@ 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
Expand All @@ -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)
Expand All @@ -184,16 +156,19 @@ def cosine_series(n: int, N: int, coeffs: List[float]) -> 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)
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
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_}')

Expand All @@ -205,37 +180,78 @@ 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:
# Compute the radius and create a linear window
L = np.max(N)
radius = (L - 1) / 2
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)

# 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)

return window


# 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)

# create the 2D window using the outer product
win = (win_y * win_x.T).T
windows1D = [_win1D(dim, type_, param=param)[0] for dim in N]

# # create the 2D window using the outer product
win = np.outer(*windows1D)

# trim the window if required
N = N - 1 * (1 - np.array(symmetric).astype(int))
Expand All @@ -246,42 +262,12 @@ 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)
win_z, _ = get_win(N[2], type_, param=param)

# create the 2D window using the outer product
win_2D = (win_x * win_z.T)
[win_x, win_y, win_z] = [_win1D(dim, type_, param=param)[0] for dim in N]

# 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))
Expand All @@ -294,6 +280,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)
Expand Down
Loading