test_fft.py 8.18 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
import ducc0.fft as fft
Martin Reinecke's avatar
Martin Reinecke committed
2
# import pyfftw
Martin Reinecke's avatar
Martin Reinecke committed
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import numpy as np
import pytest
from numpy.testing import assert_

pmp = pytest.mark.parametrize

shapes1D = ((10,), (127,))
shapes2D = ((128, 128), (128, 129), (1, 129), (129, 1))
shapes3D = ((32, 17, 39),)
shapes = shapes1D+shapes2D+shapes3D
len1D = range(1, 2048)


def _l2error(a, b):
    return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))


def _assert_close(a, b, epsilon):
    err = _l2error(a, b)
    if (err >= epsilon):
        print("Error: {} > {}".format(err, epsilon))
    assert_(err<epsilon)


def fftn(a, axes=None, inorm=0, out=None, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
28
    return fft.c2c(a, axes=axes, forward=True, inorm=inorm,
Martin Reinecke's avatar
Martin Reinecke committed
29
30
31
32
                           out=out, nthreads=nthreads)


def ifftn(a, axes=None, inorm=0, out=None, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
33
    return fft.c2c(a, axes=axes, forward=False, inorm=inorm,
Martin Reinecke's avatar
Martin Reinecke committed
34
35
36
37
                           out=out, nthreads=nthreads)


def rfftn(a, axes=None, inorm=0, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
38
    return fft.r2c(a, axes=axes, forward=True, inorm=inorm,
Martin Reinecke's avatar
Martin Reinecke committed
39
40
41
42
                           nthreads=nthreads)


def irfftn(a, axes=None, lastsize=0, inorm=0, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
43
    return fft.c2r(a, axes=axes, lastsize=lastsize, forward=False,
Martin Reinecke's avatar
Martin Reinecke committed
44
45
46
47
                           inorm=inorm, nthreads=nthreads)


def rfft_scipy(a, axis, inorm=0, out=None, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
48
    return fft.r2r_fftpack(a, axes=(axis,), real2hermitian=True,
Martin Reinecke's avatar
Martin Reinecke committed
49
50
51
52
53
                                   forward=True, inorm=inorm, out=out,
                                   nthreads=nthreads)


def irfft_scipy(a, axis, inorm=0, out=None, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
54
    return fft.r2r_fftpack(a, axes=(axis,), real2hermitian=False,
Martin Reinecke's avatar
Martin Reinecke committed
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
                                   forward=False, inorm=inorm, out=out,
                                   nthreads=nthreads)

tol = {np.float32: 6e-7, np.float64: 1.5e-15, np.longfloat: 1e-18}
ctype = {np.float32: np.complex64, np.float64: np.complex128, np.longfloat: np.longcomplex}

import platform
on_wsl = "microsoft" in platform.uname()[3].lower()
true_long_double = (np.longfloat != np.float64 and not on_wsl)
dtypes = [np.float32, np.float64,
          pytest.param(np.longfloat, marks=pytest.mark.xfail(
              not true_long_double,
              reason="Long double doesn't offer more precision"))]


@pmp("len", len1D)
@pmp("inorm", [0, 1, 2])
@pmp("dtype", dtypes)
def test1D(len, inorm, dtype):
74
    rng = np.random.default_rng(42)
75
    a = rng.random(len)-0.5 + 1j*rng.random(len)-0.5j
Martin Reinecke's avatar
Martin Reinecke committed
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
    a = a.astype(ctype[dtype])
    eps = tol[dtype]
    assert_(_l2error(a, ifftn(fftn(a, inorm=inorm), inorm=2-inorm)) < eps)
    assert_(_l2error(a.real, ifftn(fftn(a.real, inorm=inorm), inorm=2-inorm))
            < eps)
    assert_(_l2error(a.real, fftn(ifftn(a.real, inorm=inorm), inorm=2-inorm))
            < eps)
    assert_(_l2error(a.real, irfftn(rfftn(a.real, inorm=inorm),
                                    inorm=2-inorm, lastsize=len)) < eps)
    tmp = a.copy()
    assert_(ifftn(fftn(tmp, out=tmp, inorm=inorm), out=tmp, inorm=2-inorm)
            is tmp)
    assert_(_l2error(tmp, a) < eps)


@pmp("shp", shapes)
@pmp("nthreads", (0, 1, 2))
Martin Reinecke's avatar
Martin Reinecke committed
93
94
@pmp("inorm", [0, 1, 2])
def test_fftn(shp, nthreads, inorm):
95
    rng = np.random.default_rng(42)
96
    a = rng.random(shp)-0.5 + 1j*rng.random(shp)-0.5j
Martin Reinecke's avatar
Martin Reinecke committed
97
98
    assert_(_l2error(a, ifftn(fftn(a, nthreads=nthreads, inorm=inorm),
                              nthreads=nthreads, inorm=2-inorm)) < 1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
99
    a = a.astype(np.complex64)
Martin Reinecke's avatar
Martin Reinecke committed
100
101
    assert_(_l2error(a, ifftn(fftn(a, nthreads=nthreads, inorm=inorm),
                              nthreads=nthreads, inorm=2-inorm)) < 5e-7)
Martin Reinecke's avatar
Martin Reinecke committed
102
103
104
105


@pmp("shp", shapes2D)
@pmp("axes", ((0,), (1,), (0, 1), (1, 0)))
Martin Reinecke's avatar
Martin Reinecke committed
106
107
@pmp("inorm", [0, 1, 2])
def test_fftn2D(shp, axes, inorm):
108
    rng = np.random.default_rng(42)
109
    a = rng.random(shp)-0.5 + 1j*rng.random(shp)-0.5j
Martin Reinecke's avatar
Martin Reinecke committed
110
111
    assert_(_l2error(a, ifftn(fftn(a, axes=axes, inorm=inorm),
                              axes=axes, inorm=2-inorm)) < 1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
112
    a = a.astype(np.complex64)
Martin Reinecke's avatar
Martin Reinecke committed
113
114
    assert_(_l2error(a, ifftn(fftn(a, axes=axes, inorm=inorm),
                              axes=axes, inorm=2-inorm)) < 5e-7)
Martin Reinecke's avatar
Martin Reinecke committed
115
116
117
118


@pmp("shp", shapes)
def test_rfftn(shp):
119
    rng = np.random.default_rng(42)
120
    a = rng.random(shp)-0.5
Martin Reinecke's avatar
Martin Reinecke committed
121
122
123
124
    tmp1 = rfftn(a)
    tmp2 = fftn(a)
    part = tuple(slice(0,tmp1.shape[i]) for i in range(tmp1.ndim))
    assert_(_l2error(tmp1, tmp2[part]) < 1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
125
    a = a.astype(np.float32)
Martin Reinecke's avatar
Martin Reinecke committed
126
127
128
129
    tmp1 = rfftn(a)
    tmp2 = fftn(a)
    part = tuple(slice(0,tmp1.shape[i]) for i in range(tmp1.ndim))
    assert_(_l2error(tmp1, tmp2[part]) < 5e-7)
Martin Reinecke's avatar
Martin Reinecke committed
130
131


Martin Reinecke's avatar
Martin Reinecke committed
132
133
134
# @pmp("shp", shapes)
# def test_rfft_scipy(shp):
#     for i in range(len(shp)):
135
#         a = rng.random(shp)-0.5
Martin Reinecke's avatar
Martin Reinecke committed
136
137
138
139
#         assert_(_l2error(pyfftw.interfaces.scipy_fftpack.rfft(a, axis=i),
#                          rfft_scipy(a, axis=i)) < 1e-15)
#         assert_(_l2error(pyfftw.interfaces.scipy_fftpack.irfft(a, axis=i),
#                          irfft_scipy(a, axis=i, inorm=2)) < 1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
140
141
142
143
144


@pmp("shp", shapes2D)
@pmp("axes", ((0,), (1,), (0, 1), (1, 0)))
def test_rfftn2D(shp, axes):
145
    rng = np.random.default_rng(42)
146
    a = rng.random(shp)-0.5
Martin Reinecke's avatar
Martin Reinecke committed
147
148
149
150
    tmp1 = rfftn(a,axes=axes)
    tmp2 = fftn(a,axes=axes)
    part = tuple(slice(0,tmp1.shape[i]) for i in range(tmp1.ndim))
    assert_(_l2error(tmp1, tmp2[part]) < 1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
151
    a = a.astype(np.float32)
Martin Reinecke's avatar
Martin Reinecke committed
152
153
154
155
    tmp1 = rfftn(a,axes=axes)
    tmp2 = fftn(a,axes=axes)
    part = tuple(slice(0,tmp1.shape[i]) for i in range(tmp1.ndim))
    assert_(_l2error(tmp1, tmp2[part]) < 5e-7)
Martin Reinecke's avatar
Martin Reinecke committed
156
157
158
159


@pmp("shp", shapes)
def test_identity(shp):
160
    rng = np.random.default_rng(42)
161
    a = rng.random(shp)-0.5 + 1j*rng.random(shp)-0.5j
Martin Reinecke's avatar
Martin Reinecke committed
162
163
164
165
166
167
168
169
170
171
172
173
    assert_(_l2error(ifftn(fftn(a), inorm=2), a) < 1.5e-15)
    assert_(_l2error(ifftn(fftn(a.real), inorm=2), a.real) < 1.5e-15)
    assert_(_l2error(fftn(ifftn(a.real), inorm=2), a.real) < 1.5e-15)
    tmp = a.copy()
    assert_(ifftn(fftn(tmp, out=tmp), inorm=2, out=tmp) is tmp)
    assert_(_l2error(tmp, a) < 1.5e-15)
    a = a.astype(np.complex64)
    assert_(_l2error(ifftn(fftn(a), inorm=2), a) < 6e-7)


@pmp("shp", shapes)
def test_identity_r(shp):
174
    rng = np.random.default_rng(42)
175
    a = rng.random(shp)-0.5
Martin Reinecke's avatar
Martin Reinecke committed
176
177
178
179
180
181
182
183
184
185
186
    b = a.astype(np.float32)
    for ax in range(a.ndim):
        n = a.shape[ax]
        assert_(_l2error(irfftn(rfftn(a, (ax,)), (ax,), lastsize=n, inorm=2),
                         a) < 1e-15)
        assert_(_l2error(irfftn(rfftn(b, (ax,)), (ax,), lastsize=n, inorm=2),
                         b) < 5e-7)


@pmp("shp", shapes)
def test_identity_r2(shp):
187
    rng = np.random.default_rng(42)
188
    a = rng.random(shp)-0.5 + 1j*rng.random(shp)-0.5j
Martin Reinecke's avatar
Martin Reinecke committed
189
190
191
192
193
194
    a = rfftn(irfftn(a))
    assert_(_l2error(rfftn(irfftn(a), inorm=2), a) < 1e-15)


@pmp("shp", shapes2D+shapes3D)
def test_genuine_hartley(shp):
195
    rng = np.random.default_rng(42)
196
    a = rng.random(shp)-0.5
Martin Reinecke's avatar
Martin Reinecke committed
197
    v1 = fft.genuine_hartley(a)
Martin Reinecke's avatar
Martin Reinecke committed
198
199
200
201
202
203
204
    v2 = fftn(a.astype(np.complex128))
    v2 = v2.real+v2.imag
    assert_(_l2error(v1, v2) < 1e-15)


@pmp("shp", shapes)
def test_hartley_identity(shp):
205
    rng = np.random.default_rng(42)
206
    a = rng.random(shp)-0.5
Martin Reinecke's avatar
Martin Reinecke committed
207
    v1 = fft.separable_hartley(fft.separable_hartley(a))/a.size
Martin Reinecke's avatar
Martin Reinecke committed
208
209
210
211
212
    assert_(_l2error(a, v1) < 1e-15)


@pmp("shp", shapes)
def test_genuine_hartley_identity(shp):
213
    rng = np.random.default_rng(42)
214
    a = rng.random(shp)-0.5
Martin Reinecke's avatar
Martin Reinecke committed
215
    v1 = fft.genuine_hartley(fft.genuine_hartley(a))/a.size
Martin Reinecke's avatar
Martin Reinecke committed
216
217
    assert_(_l2error(a, v1) < 1e-15)
    v1 = a.copy()
Martin Reinecke's avatar
Martin Reinecke committed
218
219
    assert_(fft.genuine_hartley(
        fft.genuine_hartley(v1, out=v1), inorm=2, out=v1) is v1)
Martin Reinecke's avatar
Martin Reinecke committed
220
221
222
223
224
225
    assert_(_l2error(a, v1) < 1e-15)


@pmp("shp", shapes2D+shapes3D)
@pmp("axes", ((0,), (1,), (0, 1), (1, 0)))
def test_genuine_hartley_2D(shp, axes):
226
    rng = np.random.default_rng(42)
227
    a = rng.random(shp)-0.5
Martin Reinecke's avatar
Martin Reinecke committed
228
    assert_(_l2error(fft.genuine_hartley(fft.genuine_hartley(
Martin Reinecke's avatar
Martin Reinecke committed
229
230
231
232
233
234
235
236
        a, axes=axes), axes=axes, inorm=2), a) < 1e-15)


@pmp("len", len1D)
@pmp("inorm", [0, 1])  # inorm==2 not needed, tested via inverse
@pmp("type", [1, 2, 3, 4])
@pmp("dtype", dtypes)
def testdcst1D(len, inorm, type, dtype):
237
    rng = np.random.default_rng(42)
238
    a = (rng.random(len)-0.5).astype(dtype)
Martin Reinecke's avatar
Martin Reinecke committed
239
240
241
242
    eps = tol[dtype]
    itp = (0, 1, 3, 2, 4)
    itype = itp[type]
    if type != 1 or len > 1:  # there are no length-1 type 1 DCTs
Martin Reinecke's avatar
Martin Reinecke committed
243
244
        _assert_close(a, fft.dct(fft.dct(a, inorm=inorm, type=type), inorm=2-inorm, type=itype), eps)
    _assert_close(a, fft.dst(fft.dst(a, inorm=inorm, type=type), inorm=2-inorm, type=itype), eps)