Repository: scipy/scipy
File: scipy/signal/tests/test_signaltools.py (Python)

Code before:
                N = (N, N)
            elif len(N) != 2 or xp.any(xp.asarray(N) <= 0):
                raise ValueError("When given as a tuple, N must hold exactly "
                                "two positive integers")

            Xf = sp_fft.fft2(x, N, axes=(0, 1))
            h1 = xp.zeros(N[0], dtype=Xf.dtype)
            h2 = xp.zeros(N[1], dtype=Xf.dtype)
            for h in (h1, h2):
                N1 = h.shape[0]
                if N1 % 2 == 0:
                    h[0] = h[N1 // 2] = 1
                    h[1:N1 // 2] = 2
                else:
                    h[0] = 1
                    h[1:(N1 + 1) // 2] = 2

            h = h1[:, xp.newaxis] * h2[xp.newaxis, :]
            k = x.ndim
            while k > 2:
                h = h[:, xp.newaxis]
                k -= 1
            x = sp_fft.ifft2(Xf * h, axes=(0, 1))
            return x

        x = xp.arange(xp.prod(shape)).reshape(shape)
        xh_old = _hilbert2(x)
        xh = hilbert2(x)
        xp_assert_close(xh_old, xh)


class TestEnvelope:
    """Unit tests for function `._signaltools.envelope()`. """

    @staticmethod
    def assert_close(actual, desired, msg, xp):
        a_r_tol = ({'atol': 1e-12, 'rtol': 1e-12}
                   if xp_default_dtype(xp) == xp.float64
                   else {'atol': 1e-5, 'rtol': 1e-5}
        )

        """Little helper to compare to arrays with proper tolerances"""
        xp_assert_close(actual, desired, **a_r_tol, err_msg=msg)

    def test_envelope_invalid_parameters(self, xp):
        """For `envelope()` Raise all exceptions that are used to verify function
        parameters. """
        with pytest.raises(ValueError,
                           match=r"Invalid parameter axis=2 for z.shape=.*"):
            envelope(np.ones(3), axis=2)
        with pytest.raises(ValueError,

Review comment (suggestion):
```suggestion
        x = xp.arange(shape[0] * shape[1]).reshape(shape)
```

Code after:
        Z = xp.asarray([4.0, 2, 2, 3, 0], dtype=dt_r)
        Zr_a = xp.asarray([4.0, 0, 0, 6, 0, 0, 0, 0], dtype=dt_r)
        z = sp_fft.irfft(Z)
        n = z.shape[0]

        # the reference envelope:
        ze2_0, zr_0 = xp.unstack(envelope(z, (1, 3), residual='all', squared=True))
        self.assert_close(sp_fft.rfft(ze2_0),
                          xp.asarray([4, 2, 0, 0, 0], dtype=dt_c),
                          msg="Envelope calculation error", xp=xp)
        self.assert_close(sp_fft.rfft(zr_0),
                          xp.asarray([4, 0, 0, 3, 0], dtype=dt_c),
                          msg="Residual calculation error", xp=xp)

        ze_1, zr_1 = xp.unstack(envelope(z, (1, 3), residual='all', squared=False))
        self.assert_close(ze_1**2, ze2_0,
                          msg="Unsquared versus Squared envelope calculation error",
                          xp=xp)
        self.assert_close(zr_1, zr_0,
                          msg="Unsquared versus Squared residual calculation error",
                          xp=xp)

        ze2_2, zr_2 = xp.unstack(
            envelope(z, (1, 3), residual='all', squared=True, n_out=3*n)
        )
        self.assert_close(ze2_2[::3], ze2_0,
                          msg="3x up-sampled envelope calculation error", xp=xp)
        self.assert_close(zr_2[::3], zr_0,
                          msg="3x up-sampled residual calculation error", xp=xp)

        ze2_3, zr_3 = xp.unstack(envelope(z, (1, 3), residual='lowpass', squared=True))
        self.assert_close(ze2_3, ze2_0,
                          msg="`residual='lowpass'` envelope calculation error", xp=xp)
        self.assert_close(sp_fft.rfft(zr_3),
                          xp.asarray([4, 0, 0, 0, 0], dtype=dt_c),
                          msg="`residual='lowpass'` residual calculation error", xp=xp)

        ze2_4 = envelope(z, (1, 3), residual=None, squared=True)
        self.assert_close(ze2_4, ze2_0,
                          msg="`residual=None` envelope calculation error", xp=xp)

        # compare complex analytic signal to real version
        Z_a = xp.asarray(Z, copy=True)
        Z_a[1:] *= 2
        z_a = sp_fft.ifft(Z_a, n=n)  # analytic signal of Z
        self.assert_close(xp.real(z_a), z,
                          msg="Reference analytic signal error", xp=xp)
        ze2_a, zr_a = xp.unstack(envelope(z_a, (1, 3), residual='all', squared=True))
        self.assert_close(ze2_a, xp.astype(ze2_0, dt_c),  # dtypes must match
                          msg="Complex envelope calculation error", xp=xp)
        self.assert_close(sp_fft.fft(zr_a), xp.asarray(Zr_a, dtype=dt_c),