update

alvgaona · alvgaona · commit 1603573a6509 · 2022-10-11T14:13:53.000-03:00
diff --git a/docs/source/signal.windows.rst b/docs/source/signal.windows.rst
@@ -0,0 +1,17 @@
+.. role:: hidden
+    :class: hidden-section
+
+torch.signal.windows
+=============
+
+The torch.signal.windows module, modeled after SciPy's `special <https://docs.scipy.org/doc/scipy/reference/signal.windows.html>`_ module.
+
+.. automodule:: torch.signal.windows
+.. currentmodule:: torch.signal.windows
+
+Functions
+-----------------------
+
+.. autofunction:: cosine_window
+.. autofunction:: exponential_window
+.. autofunction:: gaussian_window
diff --git a/torch/__init__.py b/torch/__init__.py
@@ -846,6 +846,7 @@ def _assert(condition, message):
 from torch import futures as futures
 from torch import nested as nested
 from torch import nn as nn
+from torch.signal import windows as windows
 from torch import optim as optim
 import torch.optim._multi_tensor
 from torch import multiprocessing as multiprocessing
diff --git a/torch/signal/__init__.py b/torch/signal/__init__.py
@@ -1,7 +1,7 @@
-from .windows import cosine_window, exponential_window, chebyshev_window
+from .windows import cosine_window, exponential_window, gaussian_window
 
 __all__ = [
     'cosine_window',
     'exponential_window',
-    'chebyshev_window'
+    'gaussian_window',
 ]
diff --git a/torch/signal/windows/__init__.py b/torch/signal/windows/__init__.py
@@ -1,7 +1,7 @@
-from .windows import cosine_window, exponential_window, chebyshev_window
+from .windows import cosine_window, exponential_window, gaussian_window
 
 __all__ = [
     'cosine_window',
     'exponential_window',
-    'chebyshev_window'
+    'gaussian_window',
 ]
diff --git a/torch/signal/windows/windows.py b/torch/signal/windows/windows.py
@@ -6,91 +6,73 @@
 from torch.types import _dtype, _device, _layout
 from torch.fft import fft
 
-__all__ = ['cosine_window']
+__all__ = [
+    'cosine_window',
+    'exponential_window',
+    'gaussian_window',
+]
 
 
-def _window_function_checks(function_name: str, window_length: int, dtype: _dtype, layout: _layout, device: _device):
+def _window_function_checks(function_name: str, window_length: int, dtype: _dtype, layout: _layout):
     def is_floating_type(t: _dtype) -> bool:
         return t == torch.float32 or t == torch.bfloat16 or t == torch.float64 or t == torch.float16
 
     def is_complex_type(t: _dtype) -> bool:
         return t == torch.complex64 or t == torch.complex128 or t == torch.complex32
 
     if window_length < 0:
-        raise RuntimeError(f'{function_name} requires non-negative window_length, got window_length: {window_length}')
-    if layout is torch.sparse:
-        raise RuntimeError(f'{function_name} is not implemented for sparse types, got layout: {layout}')
+        raise RuntimeError(f'{function_name} requires non-negative window_length, got window_length={window_length}')
+    if layout is torch.sparse_coo:
+        raise RuntimeError(f'{function_name} is not implemented for sparse types, got: {layout}')
     if not is_floating_type(dtype) and not is_complex_type(dtype):
         raise RuntimeError(f'{function_name} expects floating point dtypes, got: {dtype}')
 
 
-def chebyshev_window(window_length: int,
-                     attenuation: float,
-                     periodic: bool = True,
-                     dtype: _dtype = None,
-                     layout: _layout = torch.strided,
-                     device: _device = None) -> Tensor:
-    _window_function_checks('chebyshev_window', window_length, dtype, layout, device)
-
-    if window_length == 0:
-        return torch.empty((0,), dtype=dtype, layout=layout, device=device)
-
-    if window_length == 1:
-        return torch.ones((1,), dtype=dtype, layout=layout, device=device)
-
-    if not periodic:
-        window_length += 1
-
-    k = torch.arange(window_length, dtype=dtype, layout=layout, device=device)
-
-    order = window_length - 1
-    beta = np.cosh(1.0 / order * np.arccosh(np.power(10, attenuation / 20.0)))
-
-    x = beta * torch.cos(torch.pi * k / window_length)
-    window = torch.special.chebyshev_polynomial_t(x, order) / np.power(10, attenuation / 20.0)
-
-    if window_length % 2 != 0:
-        window = torch.real(fft(window))
-        n = (window_length + 1) // 2
-        window = torch.concat((torch.flip(window[1:n], (0,)), window[:n]))
-    else:
-        window = window * torch.exp(1.j * torch.pi / window_length * torch.arange(window_length))
-        window = torch.real(fft(window))
-        n = window_length // 2 + 1
-        window = torch.concat((torch.flip(window[1:n], (0,)), window[1:n]))
-
-    window /= torch.max(window)
-
-    return window if periodic else window[:window_length - 1]
-
-
 def exponential_window(window_length: int,
                        periodic: bool = True,
                        center: float = None,
                        tau: float = 1.0,
                        dtype: _dtype = None,
                        layout: _layout = torch.strided,
-                       device: _device = None) -> Tensor:
+                       device: _device = None,
+                       requires_grad: bool = False) -> Tensor:
     """r
-    Computes a window with a simple cosine waveform.
+    Computes a window with an exponential form. The window
+    is also known as Poisson window.
+
+    The exponential window is defined as follows:
+
+    .. math::
+        w(n) = \exp{-\frac{|n - center|}{\tau}}
 
     Args:
-        window_length:
+        window_length: the length of the output window. In other words, the number of points of the cosine window.
+        periodic: If `True`, returns a periodic window suitable for use in spectral analysis. If `False`,
+        returns a symmetric window suitable for use in filter design.
+        center: this value defines where the center of the window will be located. In other words, at which
+        sample the peak of the window can be found.
+        tau: the decay value. For `center = 0`, it's suggested to use `tau = -(M - 1) / ln(x)`, if `` is
+        the fraction of the window remaining at the end.
 
     Keyword args:
         {dtype}
+        layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only
+              `torch.strided` (dense layout) is supported.
         {device}
-
+        {requires_grad}
     """
+    if dtype is None:
+        dtype = torch.get_default_dtype()
+
     _window_function_checks('exponential_window', window_length, dtype, layout, device)
 
     if window_length == 0:
-        return torch.empty((0,), dtype=dtype, layout=layout, device=device)
+        return torch.empty((0,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad)
 
     if window_length == 1:
-        return torch.ones((1,), dtype=dtype, layout=layout, device=device)
+        return torch.ones((1,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad)
 
-    if not periodic:
+    if periodic:
         window_length += 1
 
     if periodic and center is not None:
@@ -99,17 +81,18 @@ def exponential_window(window_length: int,
     if center is None:
         center = (window_length - 1) / 2
 
-    k = torch.arange(window_length, dtype=dtype, layout=layout, device=device)
+    k = torch.arange(window_length, dtype=dtype, layout=layout, device=device, requires_grad=requires_grad)
     window = torch.exp(-torch.abs(k - center) / tau)
 
-    return window if periodic else window[:window_length - 1]
+    return window[:-1] if periodic else window
 
 
 def cosine_window(window_length: int,
                   periodic: bool = True,
                   dtype: _dtype = None,
                   layout: _layout = torch.strided,
-                  device: _device = None) -> Tensor:
+                  device: _device = None,
+                  requires_grad: bool = False) -> Tensor:
     """r
     Computes a window with a simple cosine waveform.
 
@@ -119,7 +102,7 @@ def cosine_window(window_length: int,
     .. math::
         w(n) = \cos{(\frac{\pi n}{M}) - \frac{\pi}{2})} = \sin{(\frac{\pi n}{M})}
 
-    Where `M
+    Where `M` is the window length.
 
 
     Args:
@@ -134,20 +117,70 @@ def cosine_window(window_length: int,
         {device}
         {requires_grad}
     """
+    if dtype is None:
+        dtype = torch.get_default_dtype()
+
     _window_function_checks('cosine_window', window_length, dtype, layout, device)
 
     if window_length == 0:
-        return torch.empty((0,), dtype=dtype, layout=layout, device=device)
+        return torch.empty((0,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad)
 
     if window_length == 1:
-        return torch.ones((1,), dtype=dtype, layout=layout, device=device)
+        return torch.ones((1,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad)
 
-    if not periodic:
+    if periodic:
         window_length += 1
 
-    # k = torch.arange(window_length, dtype=dtype, layout=layout, device=device)
-    k = np.arange(0, window_length)
-    window = np.sin(np.pi / window_length * (k + .5))
-    window = torch.from_numpy(window)
+    k = torch.arange(window_length, dtype=dtype, layout=layout, device=device, requires_grad=requires_grad)
+    window = torch.sin(torch.pi / window_length * (k + .5))
+    return window[:-1] if periodic else window
+
+
+def gaussian_window(window_length: int,
+                    periodic: bool = True,
+                    std: float = 0.5,
+                    dtype: _dtype = None,
+                    layout: _layout = torch.strided,
+                    device: _device = None,
+                    requires_grad: bool = False) -> Tensor:
+    """r
+    Computes a window with a gaussian waveform.
+
+    The gaussian window is defined as follows:
+
+    .. math::
+        w(n) = \exp{-\frac{1}{2}\frac{n}{\sigma}^2}
+
+    Args:
+        window_length: the length of the output window. In other words, the number of points of the cosine window.
+        periodic: If `True`, returns a periodic window suitable for use in spectral analysis. If `False`,
+        returns a symmetric window suitable for use in filter design.
+        std: the standard deviation of the gaussian. It controls how narrow or wide the window is.
+
+
+    Keyword args:
+        {dtype}
+        layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only
+              `torch.strided` (dense layout) is supported.
+        {device}
+        {requires_grad}
+    """
+    if dtype is None:
+        dtype = torch.get_default_dtype()
+
+    _window_function_checks('cosine_window', window_length, dtype, layout, device)
+
+    if window_length == 0:
+        return torch.empty((0,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad)
+
+    if window_length == 1:
+        return torch.ones((1,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad)
+
+    if periodic:
+        window_length += 1
 
-    return window if periodic else window[:window_length - 1]
+    k = torch.arange(window_length, dtype=dtype, layout=layout, device=device, requires_grad=requires_grad)
+    k = k - (window_length - 1.0) / 2.0
+    sig2 = 2 * std * std
+    window = torch.exp(-k ** 2 / sig2)
+    return window[:-1] if periodic else window

Original file line number	Diff line number	Diff line change
`@@ -1,7 +1,7 @@`
`1`		`-from .windows import cosine_window, exponential_window, chebyshev_window`
	`1`	`+from .windows import cosine_window, exponential_window, gaussian_window`
`2`	`2`
`3`	`3`	`__all__ = [`
`4`	`4`	`'cosine_window',`
`5`	`5`	`'exponential_window',`
`6`		`- 'chebyshev_window'`
	`6`	`+ 'gaussian_window',`
`7`	`7`	`]`