Fixed style

Author: ycastorium

lib/nx_signal/windows.ex

@@ -5,7 +5,6 @@ defmodule NxSignal.Windows do
   import Nx.Defn
 
   @pi :math.pi()
-  @eps 1.0e-7
 
   @doc """
   Rectangular window.
@@ -302,44 +301,46 @@ defmodule NxSignal.Windows do
        otherwise produces a symmetric window. Defaults to `true`
     * `:type` - the output type for the window. Defaults to `{:f, 32}`
     * `:beta` - Shape parameter for the window. As beta increases, the window becomes more focused in frequency domain. Defaults to 12.0.
-    * `:name` - the axis name. Defaults to `nil`
+    * `:axis_name` - the axis name. Defaults to `nil`
 
   ## Examples
       iex> NxSignal.Windows.kaiser(n: 4, beta: 12.0, is_periodic: true)
       #Nx.Tensor<
         f32[4]
-        [5.2776191296288744e-5, 0.21566684544086456, 1.0, 0.21566684544086456]
+        [5.2776191296288744e-5, 0.21566666662693024, 1.0, 0.21566666662693024]
       >
 
       iex> NxSignal.Windows.kaiser(n: 5, beta: 12.0, is_periodic: true)
       #Nx.Tensor<
         f32[5]
-        [5.2776191296288744e-5, 0.10171464085578918, 0.7929376363754272, 0.7929376363754272, 0.10171464085578918]
+        [5.2776191296288744e-5, 0.10171464085578918, 0.7929369807243347, 0.7929369807243347, 0.10171464085578918]
       >
 
       iex> NxSignal.Windows.kaiser(n: 4, beta: 12.0, is_periodic: false)
       #Nx.Tensor<
         f32[4]
-        [5.2776191296288744e-5, 0.5188400149345398, 0.5188400149345398, 5.2776191296288744e-5]
+        [5.2776191296288744e-5, 0.5188394784927368, 0.5188390612602234, 5.2776191296288744e-5]
       >
   """
   @doc type: :windowing
   defn kaiser(opts \\ []) do
-    opts = keyword!(opts, [:n, :name, beta: 12.0, is_periodic: true, type: {:f, 32}])
+    opts =
+      keyword!(opts, [:n, :axis_name, eps: 1.0e-7, beta: 12.0, is_periodic: true, type: {:f, 32}])
+
     {l, opts} = pop_window_size(opts)
-    name = opts[:name]
+    name = opts[:axis_name]
     type = opts[:type]
     beta = opts[:beta]
+    eps = opts[:eps]
     is_periodic = opts[:is_periodic]
 
     window_length = if is_periodic, do: l + 1, else: l
-    alpha = (window_length - 1) / 2.0
 
-    n = Nx.iota({window_length}, names: [name], type: type)
-    ratio = (n - alpha) / alpha
-    r = beta * Nx.sqrt(1 - Nx.pow(ratio, 2) + @eps)
+    ratio = Nx.linspace(-1, 1, n: window_length, endpoint: true, type: type) |> Nx.rename([name])
+    sqrt_arg = Nx.max(1 - ratio ** 2, eps)
+    r = beta * Nx.sqrt(sqrt_arg)
 
-    window = i0(r) / i0(beta)
+    window = kaiser_bessel_i0(r) / kaiser_bessel_i0(beta)
 
     if is_periodic do
       Nx.slice(window, [0], [l])
@@ -348,19 +349,19 @@ defmodule NxSignal.Windows do
     end
   end
 
-  defnp i0(x) do
+  defnp kaiser_bessel_i0(x) do
     abs_x = Nx.abs(x)
 
     small_x_result =
-      1.0 +
-        Nx.pow(abs_x, 2) / 4.0 +
-        Nx.pow(abs_x, 4) / 64.0 +
-        Nx.pow(abs_x, 6) / 2304.0 +
-        Nx.pow(abs_x, 8) / 147_456.0
+      1 +
+        abs_x ** 2 / 4 +
+        abs_x ** 4 / 64 +
+        abs_x ** 6 / 2304 +
+        abs_x ** 8 / 147_456
 
     large_x_result =
       Nx.exp(abs_x) / Nx.sqrt(2 * Nx.Constants.pi() * abs_x) *
-        (1.0 + 1.0 / (8.0 * abs_x) + 9.0 / (128.0 * Nx.pow(abs_x, 2)))
+        (1 + 1 / (8 * abs_x) + 9 / (128 * Nx.pow(abs_x, 2)))
 
     Nx.select(abs_x < 3.75, small_x_result, large_x_result)
   end