ENH: add filter windows

Author: ev-br

torch_np/__init__.py

@@ -14,7 +14,7 @@
 
 inf = float("inf")
 nan = float("nan")
-from math import pi  # isort: skip
+from math import pi, e  # isort: skip
 
 False_ = asarray(False, bool_)
 True_ = asarray(True, bool_)

torch_np/_funcs.py

@@ -1898,3 +1898,41 @@ def __getitem__(self, item):
 
 index_exp = IndexExpression(maketuple=True)
 s_ = IndexExpression(maketuple=False)
+
+
+# ### Filter windows ###
+
+
+@normalizer
+def hamming(M: ArrayLike):
+    M = int(M.item())
+    dtype = _dtypes_impl.default_float_dtype
+    return torch.hamming_window(M, periodic=False, dtype=dtype)
+
+
+@normalizer
+def hanning(M: ArrayLike):
+    M = int(M.item())
+    dtype = _dtypes_impl.default_float_dtype
+    return torch.hann_window(M, periodic=False, dtype=dtype)
+
+
+@normalizer
+def kaiser(M: ArrayLike, beta):
+    M = int(M.item())
+    dtype = _dtypes_impl.default_float_dtype
+    return torch.kaiser_window(M, beta=beta, periodic=False, dtype=dtype)
+
+
+@normalizer
+def blackman(M: ArrayLike):
+    M = int(M.item())
+    dtype = _dtypes_impl.default_float_dtype
+    return torch.blackman_window(M, periodic=False, dtype=dtype)
+
+
+@normalizer
+def bartlett(M: ArrayLike):
+    M = int(M.item())
+    dtype = _dtypes_impl.default_float_dtype
+    return torch.bartlett_window(M, periodic=False, dtype=dtype)

torch_np/tests/numpy_tests/lib/test_function_base.py

@@ -26,14 +26,15 @@
 
 # FIXME: make from torch_np
 from numpy.lib import (
-    bartlett, blackman,
-    delete, digitize, extract, gradient, hamming, hanning,
-    insert, interp, kaiser, msort, piecewise, place,
+    delete, digitize, extract, gradient,
+    insert, interp, msort, piecewise, place,
     select, setxor1d, trapz, trim_zeros, unwrap, vectorize
     )
 from torch_np._detail._util import normalize_axis_tuple
 
 from torch_np import corrcoef, cov, i0, angle, sinc, diff, meshgrid, unique
+from torch_np import flipud, hamming, hanning, kaiser, blackman, bartlett
+
 
 def get_mat(n):
     data = np.arange(n)
@@ -1701,7 +1702,6 @@ def test_period(self):
         assert sm_discont.dtype == wrap_uneven.dtype
 
 
-@pytest.mark.xfail(reason='TODO: implement')
 @pytest.mark.parametrize(
     "dtype", np.typecodes["AllInteger"] + np.typecodes["Float"]
 )
@@ -1716,7 +1716,7 @@ def test_hanning(self, dtype: str, M: int) -> None:
         assert w.dtype == ref_dtype
 
         # check symmetry
-        assert_equal(w, flipud(w))
+        assert_allclose(w, flipud(w), atol=1e-15)
 
         # check known value
         if scalar < 1:
@@ -1734,7 +1734,7 @@ def test_hamming(self, dtype: str, M: int) -> None:
         assert w.dtype == ref_dtype
 
         # check symmetry
-        assert_equal(w, flipud(w))
+        assert_allclose(w, flipud(w), atol=1e-15)
 
         # check known value
         if scalar < 1:
@@ -1752,7 +1752,7 @@ def test_bartlett(self, dtype: str, M: int) -> None:
         assert w.dtype == ref_dtype
 
         # check symmetry
-        assert_equal(w, flipud(w))
+        assert_allclose(w, flipud(w), atol=1e-15)
 
         # check known value
         if scalar < 1:
@@ -1770,7 +1770,7 @@ def test_blackman(self, dtype: str, M: int) -> None:
         assert w.dtype == ref_dtype
 
         # check symmetry
-        assert_equal(w, flipud(w))
+        assert_allclose(w, flipud(w), atol=1e-15)
 
         # check known value
         if scalar < 1: