@@ -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