stress.py 3.54 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
5

def _l2error(a, b):
Martin Reinecke's avatar
Martin Reinecke committed
6
7
    return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))

Martin Reinecke's avatar
Martin Reinecke committed
8

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

Martin Reinecke's avatar
Martin Reinecke committed
13

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

Martin Reinecke's avatar
Martin Reinecke committed
18

Martin Reinecke's avatar
Martin Reinecke committed
19
20
21
22
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
23

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

Martin Reinecke's avatar
Martin Reinecke committed
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

nthreads = 0


def update_err(err, name, value):
    if name not in err:
        err[name] = value
        print("{}: {}".format(name, value))
    else:
        if value > err[name]:
            err[name] = value
            print("{}: {}".format(name, value))
    return err


def test(err):
    ndim = np.random.randint(1, 5)
    axlen = int((2**20)**(1./ndim))
    shape = np.random.randint(1, axlen, ndim)
Martin Reinecke's avatar
Martin Reinecke committed
47
48
    axes = np.arange(ndim)
    np.random.shuffle(axes)
Martin Reinecke's avatar
Martin Reinecke committed
49
    nax = np.random.randint(1, ndim+1)
Martin Reinecke's avatar
Martin Reinecke committed
50
51
    axes = axes[:nax]
    lastsize = shape[axes[-1]]
Martin Reinecke's avatar
Martin Reinecke committed
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
    a = np.random.rand(*shape)-0.5 + 1j*np.random.rand(*shape)-0.5j
    b = ifftn(fftn(a, axes=axes, nthreads=nthreads), axes=axes, inorm=2,
              nthreads=nthreads)
    err = update_err(err, "cmax", _l2error(a, b))
    b = ifftn(fftn(a.real, axes=axes, nthreads=nthreads), axes=axes, inorm=2,
              nthreads=nthreads)
    err = update_err(err, "cmax", _l2error(a.real, b))
    b = fftn(ifftn(a.real, axes=axes, nthreads=nthreads), axes=axes, inorm=2,
             nthreads=nthreads)
    err = update_err(err, "cmax", _l2error(a.real, b))
    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))
    b = ifftn(fftn(a.astype(np.complex64), axes=axes, nthreads=nthreads),
              axes=axes, inorm=2, nthreads=nthreads)
    err = update_err(err, "cmaxf", _l2error(a.astype(np.complex64), b))
    b = irfftn(rfftn(a.real.astype(np.float32), axes=axes, nthreads=nthreads),
               axes=axes, inorm=2, lastsize=lastsize, nthreads=nthreads)
    err = update_err(err, "rmaxf", _l2error(a.real.astype(np.float32), b))
    b = pypocketfft.separable_hartley(
        pypocketfft.separable_hartley(a.real, axes=axes, nthreads=nthreads),
        axes=axes, inorm=2, nthreads=nthreads)
    err = update_err(err, "hmax", _l2error(a.real, b))
    b = pypocketfft.genuine_hartley(
        pypocketfft.genuine_hartley(a.real, axes=axes, nthreads=nthreads),
        axes=axes, inorm=2, nthreads=nthreads)
    err = update_err(err, "hmax", _l2error(a.real, b))
    b = pypocketfft.separable_hartley(
            pypocketfft.separable_hartley(
                a.real.astype(np.float32), axes=axes, nthreads=nthreads),
            axes=axes, inorm=2, nthreads=nthreads)
    err = update_err(err, "hmaxf", _l2error(a.real.astype(np.float32), b))
    b = pypocketfft.genuine_hartley(
            pypocketfft.genuine_hartley(a.real.astype(np.float32), axes=axes,
                                        nthreads=nthreads),
            axes=axes, inorm=2, nthreads=nthreads)
    err = update_err(err, "hmaxf", _l2error(a.real.astype(np.float32), b))

Martin Reinecke's avatar
Martin Reinecke committed
90

Martin Reinecke's avatar
Martin Reinecke committed
91
err = dict()
Martin Reinecke's avatar
Martin Reinecke committed
92
while True:
Martin Reinecke's avatar
Martin Reinecke committed
93
    test(err)