stress.py 6.29 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1 2 3
import numpy as np
import pypocketfft

Martin Reinecke's avatar
Martin Reinecke committed
4

Martin Reinecke's avatar
Martin Reinecke committed
5 6 7 8
np.random.seed(42)

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))]))
Martin Reinecke's avatar
Martin Reinecke committed
9

Martin Reinecke's avatar
Martin Reinecke committed
10

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

Martin Reinecke's avatar
Martin Reinecke committed
15

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

Martin Reinecke's avatar
Martin Reinecke committed
20

Martin Reinecke's avatar
Martin Reinecke committed
21 22 23 24
def rfftn(a, axes=None, inorm=0, nthreads=1):
    return pypocketfft.r2c(a, axes=axes, forward=True, inorm=inorm,
                           nthreads=nthreads)

Martin Reinecke's avatar
Martin Reinecke committed
25

Martin Reinecke's avatar
Martin Reinecke committed
26
def irfftn(a, axes=None, lastsize=0, inorm=0, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
27
    return pypocketfft.c2r(a, axes=axes, lastsize=lastsize, forward=False,
Martin Reinecke's avatar
Martin Reinecke committed
28 29
                           inorm=inorm, nthreads=nthreads)

Martin Reinecke's avatar
Martin Reinecke committed
30

Martin Reinecke's avatar
Martin Reinecke committed
31
nthreads = 1
Martin Reinecke's avatar
Martin Reinecke committed
32 33 34


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


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

Martin Reinecke's avatar
Martin Reinecke committed
141

Martin Reinecke's avatar
Martin Reinecke committed
142
err = dict()
Martin Reinecke's avatar
Martin Reinecke committed
143
while True:
Martin Reinecke's avatar
Martin Reinecke committed
144
    test(err)