diff --git a/screens/fields.py b/screens/fields.py index 2f0f80c..eb36d84 100644 --- a/screens/fields.py +++ b/screens/fields.py @@ -59,8 +59,9 @@ def phasor(indep, transform, linear_axis=None): ph0 = (indep[ph0_index] * transform * u.cycle).to_value(u.rad) dph = (np.diff(indep[ph01_index], axis=linear_axis) * transform * u.cycle).to_value(u.rad) - phasor = np.empty(np.broadcast(indep, transform).shape, complex) - phasor[ph0_index] = np.exp(1j * ph0) + phasor0 = np.exp(1j * ph0) + phasor = np.empty_like(phasor0, shape=np.broadcast(indep, transform).shape) + phasor[ph0_index] = phasor0 phasor[dph0_index] = np.exp(1j * dph) return np.cumprod(phasor, out=phasor, axis=linear_axis) diff --git a/screens/tests/test_phasor.py b/screens/tests/test_phasor.py index 8ebbaf2..4dca32a 100644 --- a/screens/tests/test_phasor.py +++ b/screens/tests/test_phasor.py @@ -7,50 +7,87 @@ from ..fields import phasor +NUMPY_LT_2_0 = np.__version__[0] < "2" + + class TestPhasor: + + f_type = np.float64 + c_type = np.complex128 + atol = 1e-10 + def setup_class(cls): - cls.fobs = 316 * u.MHz - f = np.arange(-8, 8) << u.MHz + cls.fobs = cls.f_type(316) * u.MHz + f = np.arange(-8, 8, dtype=cls.f_type) << u.MHz cls.f = cls.fobs + f - cls.t = (np.arange(-16, 16) << u.minute)[:, np.newaxis] - cls.tau = np.fft.fftfreq(cls.f.size, 1*u.MHz) - cls.fd = np.fft.fftfreq(cls.t.size, 1*u.minute) - cls.w_tau = 0.3 * u.us - cls.w_fd = 4 * u.mHz + cls.t = (np.arange(-16, 16, dtype=cls.f_type) << u.minute)[:, np.newaxis] + cls.tau = np.fft.fftfreq(cls.f.size, 1*u.MHz).astype(cls.f_type) + cls.fd = np.fft.fftfreq(cls.t.size, 1*u.minute).astype(cls.f_type) + cls.w_tau = cls.f_type(0.3) * u.us + cls.w_fd = cls.f_type(4) * u.mHz phase = (cls.f*cls.w_tau + cls.t*cls.w_fd)*u.cycle cls.dw = np.exp(1j*phase.to_value(u.radian)) - def test_linear(self): + def test_setup_dtype(self): + assert self.fobs.dtype == self.f_type + assert self.f.dtype == self.f_type + assert self.t.dtype == self.f_type + assert self.tau.dtype == self.f_type + assert self.fd.dtype == self.f_type + assert self.w_tau.dtype == self.f_type + assert self.w_fd.dtype == self.f_type + assert self.dw.dtype == self.c_type + + @pytest.mark.parametrize('linear_axis', (-1, "transform")) + def test_linear(self, linear_axis): t = self.t - self.t[0] fd = self.fd[:, np.newaxis, np.newaxis] full = phasor(t, fd) - linear = phasor(t, fd) - assert_allclose(full, linear, atol=1e-8, rtol=0) + assert full.dtype == self.c_type + linear = phasor(t, fd, linear_axis=linear_axis) + assert linear.dtype == self.c_type + assert_allclose(full, linear, atol=self.atol, rtol=0) @pytest.mark.parametrize('linear_axis', (None, -2, "transform")) def test_ft_t(self, linear_axis): expected = np.fft.ifft(self.dw, axis=0) + if NUMPY_LT_2_0: + expected = expected.astype(self.c_type) + else: + assert expected.dtype == self.c_type # Regular FT is always relative to start of array. t = self.t - self.t[0] phs = phasor(t, self.fd[:, np.newaxis, np.newaxis], linear_axis=linear_axis) + assert phs.dtype == self.c_type result = np.mean(self.dw*phs, axis=1) - assert_allclose(result, expected, atol=1e-8, rtol=0) + assert result.dtype == self.c_type + assert_allclose(result, expected, atol=self.atol, rtol=0) @pytest.mark.parametrize('linear_axis', (None, -1, "transform")) def test_ft_f(self, linear_axis): expected = np.fft.ifft(self.dw, axis=1) + if NUMPY_LT_2_0: + expected = expected.astype(self.c_type) + else: + assert expected.dtype == self.c_type # Regular FT is always relative to start of array. f = self.f - self.f[0] phs = phasor(f, self.tau[:, np.newaxis, np.newaxis], linear_axis=linear_axis) + assert phs.dtype == self.c_type result = np.mean(self.dw*phs, axis=-1).T - assert_allclose(result, expected, atol=1e-8, rtol=0) + assert result.dtype == self.c_type + assert_allclose(result, expected, atol=self.atol, rtol=0) @pytest.mark.parametrize('linear_axis_t,linear_axis_f', [(-2, -1), ("transform", "transform")]) def test_ft(self, linear_axis_t, linear_axis_f): expected = np.fft.ifft2(self.dw) + if NUMPY_LT_2_0: + expected = expected.astype(self.c_type) + else: + assert expected.dtype == self.c_type # Regular FT always has phase relative to start of array. t = self.t - self.t[0] f = self.f - self.f[0] @@ -58,5 +95,14 @@ def test_ft(self, linear_axis_t, linear_axis_f): linear_axis=linear_axis_t) phs_f = phasor(f, self.tau[:, np.newaxis, np.newaxis], linear_axis=linear_axis_f) + assert phs_t.dtype == self.c_type + assert phs_f.dtype == self.c_type result = np.mean(self.dw*phs_t*phs_f, axis=(-2, -1)) - assert_allclose(result, expected, atol=1e-8, rtol=0) + assert result.dtype == self.c_type + assert_allclose(result, expected, atol=self.atol, rtol=0) + + +class TestPhasorFloat32(TestPhasor): + f_type = np.float32 + c_type = np.complex64 + atol = 2e-5