bench.py 693 Bytes
Newer Older
1
import sys
Ashwin Vishnu's avatar
Ashwin Vishnu committed
2
import numpy as np
3
import perf
Ashwin Vishnu's avatar
Ashwin Vishnu committed
4
5
6
7
from mpi4py import MPI
from mpi4py_fft.mpifft import PFFT, Function


8
9
def init2d(N=1024, slab=True):
    return PFFT(MPI.COMM_WORLD, (N, N), slab=slab, axes=(0, 1), dtype=np.float)
Ashwin Vishnu's avatar
Ashwin Vishnu committed
10
11


12
13
def init3d(N=128, slab=True):
    return PFFT(MPI.COMM_WORLD, (N, N, N), slab=slab, axes=(0, 1, 2), dtype=np.float)
Ashwin Vishnu's avatar
Ashwin Vishnu committed
14
15


16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def create_arrayX(o):
    u = Function(o, False)
    u[:] = np.random.random(u.shape).astype(u.dtype)
    return u


def create_arrayK(o):
    u_hat = Function(o, True)
    u_hat[:] = np.random.random(u_hat.shape).astype(u_hat.dtype)
    return u_hat


def fft(o, u, u_hat):
    o.forward(u, u_hat)


def ifft(o, u, u_hat):
    o.backward(u_hat, u)