implement dpnp.bartlett

Author: vtavana

dpnp/backend/extensions/window/bartlett.hpp

@@ -0,0 +1,67 @@
+//*****************************************************************************
+// Copyright (c) 2025, Intel Corporation
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+// - Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+// - Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+// THE POSSIBILITY OF SUCH DAMAGE.
+//*****************************************************************************
+
+#pragma once
+
+#include "common.hpp"
+#include <sycl/sycl.hpp>
+
+namespace dpnp::extensions::window::kernels
+{
+
+template <typename T>
+class BartlettFunctor
+{
+private:
+    T *data = nullptr;
+    const std::size_t N;
+
+public:
+    BartlettFunctor(T *data, const std::size_t N) : data(data), N(N) {}
+
+    void operator()(sycl::id<1> id) const
+    {
+        const auto i = id.get(0);
+
+        data[i] =
+            T(2) / (N - 1) * ((N - 1) / T(2) - sycl::fabs(i - (N - 1) / T(2)));
+    }
+};
+
+template <typename fnT, typename T>
+struct BartlettFactory
+{
+    fnT get()
+    {
+        if constexpr (std::is_floating_point_v<T>) {
+            return window_impl<T, BartlettFunctor>;
+        }
+        else {
+            return nullptr;
+        }
+    }
+};
+
+} // namespace dpnp::extensions::window::kernels

dpnp/backend/extensions/window/window_py.cpp

@@ -30,6 +30,7 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
+#include "bartlett.hpp"
 #include "blackman.hpp"
 #include "common.hpp"
 #include "hamming.hpp"
@@ -41,6 +42,7 @@ using window_ns::window_fn_ptr_t;
 
 namespace dpctl_td_ns = dpctl::tensor::type_dispatch;
 
+static window_fn_ptr_t bartlett_dispatch_vector[dpctl_td_ns::num_types];
 static window_fn_ptr_t blackman_dispatch_vector[dpctl_td_ns::num_types];
 static window_fn_ptr_t hamming_dispatch_vector[dpctl_td_ns::num_types];
 static window_fn_ptr_t hanning_dispatch_vector[dpctl_td_ns::num_types];
@@ -50,6 +52,21 @@ PYBIND11_MODULE(_window_impl, m)
     using arrayT = dpctl::tensor::usm_ndarray;
     using event_vecT = std::vector<sycl::event>;
 
+    {
+        window_ns::init_window_dispatch_vectors<
+            window_ns::kernels::BartlettFactory>(bartlett_dispatch_vector);
+
+        auto bartlett_pyapi = [&](sycl::queue &exec_q, const arrayT &result,
+                                  const event_vecT &depends = {}) {
+            return window_ns::py_window(exec_q, result, depends,
+                                        bartlett_dispatch_vector);
+        };
+
+        m.def("_bartlett", bartlett_pyapi, "Call Bartlett kernel",
+              py::arg("sycl_queue"), py::arg("result"),
+              py::arg("depends") = py::list());
+    }
+
     {
         window_ns::init_window_dispatch_vectors<
             window_ns::kernels::BlackmanFactory>(blackman_dispatch_vector);

dpnp/dpnp_iface_window.py

@@ -45,7 +45,7 @@
 import dpnp
 import dpnp.backend.extensions.window._window_impl as wi
 
-__all__ = ["blackman", "hamming", "hanning"]
+__all__ = ["bartlett", "blackman", "hamming", "hanning"]
 
 
 def _call_window_kernel(
@@ -174,6 +174,99 @@ def blackman(M, device=None, usm_type=None, sycl_queue=None):
     )
 
 
+def bartlett(M, device=None, usm_type=None, sycl_queue=None):
+    r"""
+    Return the Bartlett window.
+
+    The Bartlett window is very similar to a triangular window, except that the
+    end points are at zero. It is often used in signal processing for tapering
+    a signal, without generating too much ripple in the frequency domain.
+
+    For full documentation refer to :obj:`numpy.bartlett`.
+
+    Parameters
+    ----------
+    M : int
+        Number of points in the output window. If zero or less, an empty array
+        is returned.
+    device : {None, string, SyclDevice, SyclQueue, Device}, optional
+        An array API concept of device where the output array is created.
+        `device` can be ``None``, a oneAPI filter selector string, an instance
+        of :class:`dpctl.SyclDevice` corresponding to a non-partitioned SYCL
+        device, an instance of :class:`dpctl.SyclQueue`, or a
+        :class:`dpctl.tensor.Device` object returned by
+        :attr:`dpnp.ndarray.device`.
+
+        Default: ``None``.
+    usm_type : {None, "device", "shared", "host"}, optional
+        The type of SYCL USM allocation for the output array.
+
+        Default: ``None``.
+    sycl_queue : {None, SyclQueue}, optional
+        A SYCL queue to use for output array allocation and copying. The
+        `sycl_queue` can be passed as ``None`` (the default), which means
+        to get the SYCL queue from `device` keyword if present or to use
+        a default queue.
+
+        Default: ``None``.
+
+    Returns
+    -------
+    out : dpnp.ndarray of shape (M,)
+        The triangular window, with the maximum value normalized to one
+        (the value one appears only if the number of samples is odd), with the
+        first and last samples equal to zero.
+
+    See Also
+    --------
+    :obj:`dpnp.blackman` : Return the Blackman window.
+    :obj:`dpnp.hamming` : Return the Hamming window.
+    :obj:`dpnp.hanning` : Return the Hanning window.
+    :obj:`dpnp.kaiser` : Return the Kaiser window.
+
+    Notes
+    -----
+    The Bartlett window is defined as
+
+    .. math::  w(n) = \frac{2}{M-1} \left(\frac{M-1}{2} -
+               \left|n - \frac{M-1}{2}\right|\right)
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> np.bartlett(12)
+    array([0.        , 0.18181818, 0.36363636, 0.54545455, 0.72727273,
+        0.90909091, 0.90909091, 0.72727273, 0.54545455, 0.36363636,
+        0.18181818, 0.        ])
+
+    Creating the output array on a different device or with a
+    specified usm_type:
+
+    >>> x = np.bartlett(4) # default case
+    >>> x, x.device, x.usm_type
+    (array([0.        , 0.66666667, 0.66666667, 0.        ]),
+     Device(level_zero:gpu:0),
+     'device')
+
+    >>> y = np.bartlett(4, device="cpu")
+    >>> y, y.device, y.usm_type
+    (array([0.        , 0.66666667, 0.66666667, 0.        ]),
+     Device(opencl:cpu:0),
+     'device')
+
+    >>> z = np.bartlett(4, usm_type="host")
+    >>> z, z.device, z.usm_type
+    (array([0.        , 0.66666667, 0.66666667, 0.        ]),
+     Device(level_zero:gpu:0),
+     'host')
+
+    """
+
+    return _call_window_kernel(
+        M, wi._bartlett, device=device, usm_type=usm_type, sycl_queue=sycl_queue
+    )
+
+
 def hamming(M, device=None, usm_type=None, sycl_queue=None):
     r"""
     Return the Hamming window.

dpnp/tests/test_sycl_queue.py

@@ -53,6 +53,7 @@ def assert_sycl_queue_equal(result, expected):
     "func, arg, kwargs",
     [
         pytest.param("arange", [-25.7], {"stop": 10**8, "step": 15}),
+        pytest.param("bartlett", [10], {}),
         pytest.param("blackman", [10], {}),
         pytest.param("eye", [4, 2], {}),
         pytest.param("empty", [(2, 2)], {}),

dpnp/tests/test_usm_type.py

@@ -177,6 +177,7 @@ def test_array_creation_from_array(func, args, usm_type_x, usm_type_y):
     "func, arg, kwargs",
     [
         pytest.param("arange", [-25.7], {"stop": 10**8, "step": 15}),
+        pytest.param("bartlett", [10], {}),
         pytest.param("blackman", [10], {}),
         pytest.param("eye", [4, 2], {}),
         pytest.param("empty", [(3, 4)], {}),

dpnp/tests/test_window.py

@@ -7,7 +7,7 @@
 from .helper import assert_dtype_allclose
 
 
-@pytest.mark.parametrize("func", ["blackman", "hamming", "hanning"])
+@pytest.mark.parametrize("func", ["bartlett", "blackman", "hamming", "hanning"])
 @pytest.mark.parametrize(
     "M",
     [
@@ -32,7 +32,7 @@ def test_window(func, M):
     assert_dtype_allclose(result, expected)
 
 
-@pytest.mark.parametrize("func", ["blackman", "hamming", "hanning"])
+@pytest.mark.parametrize("func", ["bartlett", "blackman", "hamming", "hanning"])
 @pytest.mark.parametrize(
     "M",
     [

dpnp/tests/third_party/cupy/math_tests/test_window.py

@@ -10,8 +10,7 @@
     *testing.product(
         {
             "m": [0, 1, -1, 1024],
-            # TODO: add ["bartlett"] when supported
-            "name": ["blackman", "hamming", "hanning"],
+            "name": ["bartlett", "blackman", "hamming", "hanning"],
         }
     )
 )