Skip to content

Commit e42d747

Browse files
authored
feat: Implements Kaiser Window (#27)
Adds Kaiser window
1 parent c167891 commit e42d747

File tree

1 file changed

+78
-0
lines changed

1 file changed

+78
-0
lines changed

lib/nx_signal/windows.ex

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -289,6 +289,84 @@ defmodule NxSignal.Windows do
289289
end
290290
end
291291

292+
@doc """
293+
Creates a Kaiser window of size `window_length`.
294+
295+
The Kaiser window is a taper formed by using a Bessel function.
296+
297+
## Options
298+
299+
* `:n` - The window length. Mandatory option.
300+
* `:is_periodic` - If `true`, produces a periodic window,
301+
otherwise produces a symmetric window. Defaults to `true`
302+
* `:type` - the output type for the window. Defaults to `{:f, 32}`
303+
* `:beta` - Shape parameter for the window. As beta increases, the window becomes more focused in frequency domain. Defaults to 12.0.
304+
* `:eps` - Epsilon value to avoid division by zero. Defaults to 1.0e-7.
305+
* `:axis_name` - the axis name. Defaults to `nil`
306+
307+
## Examples
308+
iex> NxSignal.Windows.kaiser(n: 4, beta: 12.0, is_periodic: true)
309+
#Nx.Tensor<
310+
f32[4]
311+
[5.2776191296288744e-5, 0.21566666662693024, 1.0, 0.21566666662693024]
312+
>
313+
314+
iex> NxSignal.Windows.kaiser(n: 5, beta: 12.0, is_periodic: true)
315+
#Nx.Tensor<
316+
f32[5]
317+
[5.2776191296288744e-5, 0.10171464085578918, 0.7929369807243347, 0.7929369807243347, 0.10171464085578918]
318+
>
319+
320+
iex> NxSignal.Windows.kaiser(n: 4, beta: 12.0, is_periodic: false)
321+
#Nx.Tensor<
322+
f32[4]
323+
[5.2776191296288744e-5, 0.5188394784927368, 0.5188390612602234, 5.2776191296288744e-5]
324+
>
325+
"""
326+
@doc type: :windowing
327+
defn kaiser(opts \\ []) do
328+
opts =
329+
keyword!(opts, [:n, :axis_name, eps: 1.0e-7, beta: 12.0, is_periodic: true, type: {:f, 32}])
330+
331+
{l, opts} = pop_window_size(opts)
332+
name = opts[:axis_name]
333+
type = opts[:type]
334+
beta = opts[:beta]
335+
eps = opts[:eps]
336+
is_periodic = opts[:is_periodic]
337+
338+
window_length = if is_periodic, do: l + 1, else: l
339+
340+
ratio = Nx.linspace(-1, 1, n: window_length, endpoint: true, type: type) |> Nx.rename([name])
341+
sqrt_arg = Nx.max(1 - ratio ** 2, eps)
342+
r = beta * Nx.sqrt(sqrt_arg)
343+
344+
window = kaiser_bessel_i0(r) / kaiser_bessel_i0(beta)
345+
346+
if is_periodic do
347+
Nx.slice(window, [0], [l])
348+
else
349+
window
350+
end
351+
end
352+
353+
defnp kaiser_bessel_i0(x) do
354+
abs_x = Nx.abs(x)
355+
356+
small_x_result =
357+
1 +
358+
abs_x ** 2 / 4 +
359+
abs_x ** 4 / 64 +
360+
abs_x ** 6 / 2304 +
361+
abs_x ** 8 / 147_456
362+
363+
large_x_result =
364+
Nx.exp(abs_x) / Nx.sqrt(2 * Nx.Constants.pi() * abs_x) *
365+
(1 + 1 / (8 * abs_x) + 9 / (128 * Nx.pow(abs_x, 2)))
366+
367+
Nx.select(abs_x < 3.75, small_x_result, large_x_result)
368+
end
369+
292370
deftransformp pop_window_size(opts) do
293371
{n, opts} = Keyword.pop(opts, :n)
294372

0 commit comments

Comments
 (0)