fft_stress.py 6.13 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
import numpy as np
Martin Reinecke's avatar
Martin Reinecke committed
2
import ducc_0_1.fft as fft
Martin Reinecke's avatar
Martin Reinecke committed
3
4


5
6
7
rng = np.random.default_rng(42)


Martin Reinecke's avatar
Martin Reinecke committed
8
9
10
11
12
def _l2error(a, b, axes):
    return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))/np.log2(np.max([2,np.prod(np.take(a.shape,axes))]))


def fftn(a, axes=None, inorm=0, out=None, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
13
14
    return fft.c2c(a, axes=axes, forward=True, inorm=inorm,
                   out=out, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
15
16
17


def ifftn(a, axes=None, inorm=0, out=None, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
18
19
    return fft.c2c(a, axes=axes, forward=False, inorm=inorm,
                   out=out, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
20
21
22


def rfftn(a, axes=None, inorm=0, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
23
24
    return fft.r2c(a, axes=axes, forward=True, inorm=inorm,
                   nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
25
26
27


def irfftn(a, axes=None, lastsize=0, inorm=0, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
28
29
    return fft.c2r(a, axes=axes, lastsize=lastsize, forward=False,
                   inorm=inorm, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
30
31
32
33
34


nthreads = 0


Martin Reinecke's avatar
Martin Reinecke committed
35
def update_err(err, name, value, shape):
Martin Reinecke's avatar
Martin Reinecke committed
36
37
38
    if name in err and err[name] >= value:
        return err
    err[name] = value
Martin Reinecke's avatar
Martin Reinecke committed
39
    print(shape)
Martin Reinecke's avatar
Martin Reinecke committed
40
41
42
43
44
45
46
    for (nm, v) in err.items():
        print("{}: {}".format(nm, v))
    print()
    return err


def test(err):
47
    ndim = rng.integers(1, 5)
Martin Reinecke's avatar
Martin Reinecke committed
48
    axlen = int((2**20)**(1./ndim))
49
    shape = rng.integers(1, axlen, ndim)
Martin Reinecke's avatar
Martin Reinecke committed
50
    axes = np.arange(ndim)
51
52
    rng.shuffle(axes)
    nax = rng.integers(1, ndim+1)
Martin Reinecke's avatar
Martin Reinecke committed
53
54
    axes = axes[:nax]
    lastsize = shape[axes[-1]]
55
    a = rng.random(shape)-0.5 + 1j*rng.random(shape)-0.5j
Martin Reinecke's avatar
Martin Reinecke committed
56
57
58
    a_32 = a.astype(np.complex64)
    b = ifftn(fftn(a, axes=axes, nthreads=nthreads), axes=axes, inorm=2,
              nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
59
    err = update_err(err, "cmax", _l2error(a, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
60
61
    b = ifftn(fftn(a.real, axes=axes, nthreads=nthreads), axes=axes, inorm=2,
              nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
62
    err = update_err(err, "cmax", _l2error(a.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
63
64
    b = fftn(ifftn(a.real, axes=axes, nthreads=nthreads), axes=axes, inorm=2,
             nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
65
    err = update_err(err, "cmax", _l2error(a.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
66
67
    b = ifftn(fftn(a.astype(np.complex64), axes=axes, nthreads=nthreads),
              axes=axes, inorm=2, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
68
    err = update_err(err, "cmaxf", _l2error(a.astype(np.complex64), b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
69
70
    b = irfftn(rfftn(a.real, axes=axes, nthreads=nthreads), axes=axes, inorm=2,
               lastsize=lastsize, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
71
    err = update_err(err, "rmax", _l2error(a.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
72
73
    b = irfftn(rfftn(a.real.astype(np.float32), axes=axes, nthreads=nthreads),
               axes=axes, inorm=2, lastsize=lastsize, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
74
    err = update_err(err, "rmaxf", _l2error(a.real.astype(np.float32), b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
75
76
    b = fft.separable_hartley(
        fft.separable_hartley(a.real, axes=axes, nthreads=nthreads),
Martin Reinecke's avatar
Martin Reinecke committed
77
        axes=axes, inorm=2, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
78
    err = update_err(err, "hmax", _l2error(a.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
79
80
    b = fft.genuine_hartley(
        fft.genuine_hartley(a.real, axes=axes, nthreads=nthreads),
Martin Reinecke's avatar
Martin Reinecke committed
81
        axes=axes, inorm=2, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
82
    err = update_err(err, "hmax", _l2error(a.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
83
84
    b = fft.separable_hartley(
            fft.separable_hartley(
Martin Reinecke's avatar
Martin Reinecke committed
85
86
                a.real.astype(np.float32), axes=axes, nthreads=nthreads),
            axes=axes, inorm=2, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
87
    err = update_err(err, "hmaxf", _l2error(a.real.astype(np.float32), b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
88
89
    b = fft.genuine_hartley(
            fft.genuine_hartley(a.real.astype(np.float32), axes=axes,
Martin Reinecke's avatar
Martin Reinecke committed
90
91
                                        nthreads=nthreads),
            axes=axes, inorm=2, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
92
    err = update_err(err, "hmaxf", _l2error(a.real.astype(np.float32), b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
93
    if all(a.shape[i] > 1 for i in axes):
Martin Reinecke's avatar
Martin Reinecke committed
94
95
        b = fft.dct(
            fft.dct(a.real, axes=axes, nthreads=nthreads, type=1),
Martin Reinecke's avatar
Martin Reinecke committed
96
            axes=axes, type=1, nthreads=nthreads, inorm=2)
Martin Reinecke's avatar
Martin Reinecke committed
97
        err = update_err(err, "c1max", _l2error(a.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
98
99
        b = fft.dct(
            fft.dct(a_32.real, axes=axes, nthreads=nthreads, type=1),
Martin Reinecke's avatar
Martin Reinecke committed
100
            axes=axes, type=1, nthreads=nthreads, inorm=2)
Martin Reinecke's avatar
Martin Reinecke committed
101
        err = update_err(err, "c1maxf", _l2error(a_32.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
102
103
    b = fft.dct(
        fft.dct(a.real, axes=axes, nthreads=nthreads, type=2),
Martin Reinecke's avatar
Martin Reinecke committed
104
        axes=axes, type=3, nthreads=nthreads, inorm=2)
Martin Reinecke's avatar
Martin Reinecke committed
105
    err = update_err(err, "c23max", _l2error(a.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
106
107
    b = fft.dct(
        fft.dct(a_32.real, axes=axes, nthreads=nthreads, type=2),
Martin Reinecke's avatar
Martin Reinecke committed
108
        axes=axes, type=3, nthreads=nthreads, inorm=2)
Martin Reinecke's avatar
Martin Reinecke committed
109
    err = update_err(err, "c23maxf", _l2error(a_32.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
110
111
    b = fft.dct(
        fft.dct(a.real, axes=axes, nthreads=nthreads, type=4),
Martin Reinecke's avatar
Martin Reinecke committed
112
        axes=axes, type=4, nthreads=nthreads, inorm=2)
Martin Reinecke's avatar
Martin Reinecke committed
113
    err = update_err(err, "c4max", _l2error(a.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
114
115
    b = fft.dct(
        fft.dct(a_32.real, axes=axes, nthreads=nthreads, type=4),
Martin Reinecke's avatar
Martin Reinecke committed
116
        axes=axes, type=4, nthreads=nthreads, inorm=2)
Martin Reinecke's avatar
Martin Reinecke committed
117
    err = update_err(err, "c4maxf", _l2error(a_32.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
118
119
    b = fft.dst(
        fft.dst(a.real, axes=axes, nthreads=nthreads, type=1),
Martin Reinecke's avatar
Martin Reinecke committed
120
        axes=axes, type=1, nthreads=nthreads, inorm=2)
Martin Reinecke's avatar
Martin Reinecke committed
121
    err = update_err(err, "s1max", _l2error(a.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
122
123
    b = fft.dst(
        fft.dst(a_32.real, axes=axes, nthreads=nthreads, type=1),
Martin Reinecke's avatar
Martin Reinecke committed
124
        axes=axes, type=1, nthreads=nthreads, inorm=2)
Martin Reinecke's avatar
Martin Reinecke committed
125
    err = update_err(err, "s1maxf", _l2error(a_32.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
126
127
    b = fft.dst(
        fft.dst(a.real, axes=axes, nthreads=nthreads, type=2),
Martin Reinecke's avatar
Martin Reinecke committed
128
        axes=axes, type=3, nthreads=nthreads, inorm=2)
Martin Reinecke's avatar
Martin Reinecke committed
129
    err = update_err(err, "s23max", _l2error(a.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
130
131
    b = fft.dst(
        fft.dst(a_32.real, axes=axes, nthreads=nthreads, type=2),
Martin Reinecke's avatar
Martin Reinecke committed
132
        axes=axes, type=3, nthreads=nthreads, inorm=2)
Martin Reinecke's avatar
Martin Reinecke committed
133
    err = update_err(err, "s23maxf", _l2error(a_32.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
134
135
    b = fft.dst(
        fft.dst(a.real, axes=axes, nthreads=nthreads, type=4),
Martin Reinecke's avatar
Martin Reinecke committed
136
        axes=axes, type=4, nthreads=nthreads, inorm=2)
Martin Reinecke's avatar
Martin Reinecke committed
137
    err = update_err(err, "s4max", _l2error(a.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
138
139
    b = fft.dst(
        fft.dst(a_32.real, axes=axes, nthreads=nthreads, type=4),
Martin Reinecke's avatar
Martin Reinecke committed
140
        axes=axes, type=4, nthreads=nthreads, inorm=2)
Martin Reinecke's avatar
Martin Reinecke committed
141
    err = update_err(err, "s4maxf", _l2error(a_32.real, b, axes), shape)
Martin Reinecke's avatar
Martin Reinecke committed
142
143
144
145
146


err = dict()
while True:
    test(err)