Read about our upcoming Code of Conduct on this issue

bench.py 752 Bytes
Newer Older
1
import sys
Ashwin Vishnu's avatar
Ashwin Vishnu committed
2
import numpy as np
3
import pyperf
Ashwin Vishnu's avatar
Ashwin Vishnu committed
4
from mpi4py import MPI
5
6
from mpi4py_fft.mpifft import PFFT
from mpi4py_fft.distarray import newDistArray
Ashwin Vishnu's avatar
Ashwin Vishnu committed
7
8


9
10
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
11
12


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


19
def create_arrayX(o):
20
    u = newDistArray(o, False)
21
22
23
24
25
    u[:] = np.random.random(u.shape).astype(u.dtype)
    return u


def create_arrayK(o):
26
    u_hat = newDistArray(o, True)
27
28
29
30
31
32
33
34
35
36
    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)