bench_gridder.py 1.97 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11
from time import time

import matplotlib.pyplot as plt
import numpy as np

import nifty5 as ift

np.random.seed(40)

N0s, a0s, b0s, c0s = [], [], [], []

12
for ii in range(10, 26):
13 14
    fovx = 0.0001
    fovy = 0.0002
Martin Reinecke's avatar
Martin Reinecke committed
15 16
    nu = 1024
    nv = 1024
Martin Reinecke's avatar
Martin Reinecke committed
17
    N = int(2**ii)
18
    print('N = {}'.format(N))
Martin Reinecke's avatar
Martin Reinecke committed
19 20
    nchan=16
    nrow=N//nchan
21
    freq = 1e9+1e6*np.arange(nchan)
Martin Reinecke's avatar
Martin Reinecke committed
22 23
    uvw = np.random.rand(nrow, 3) - 0.5
    vis = (np.random.randn(N) + 1j*np.random.randn(N)).reshape((nrow,nchan))
24 25 26

    uvspace = ift.RGSpace((nu, nv))

Martin Reinecke's avatar
Martin Reinecke committed
27
    visspace = ift.UnstructuredDomain((N//nchan,nchan))
28

29
    img = np.random.randn(nu*nv)
30 31 32 33
    img = img.reshape((nu, nv))
    img = ift.from_global_data(uvspace, img)

    t0 = time()
34
    GM = ift.GridderMaker(uvspace, eps=1e-7, uvw=uvw,
35
                          freq=freq, fovx=fovx, fovy=fovy,
Martin Reinecke's avatar
Martin Reinecke committed
36
                          flags=np.zeros((N//nchan, nchan), dtype=np.bool))
37
    vis = ift.from_global_data(visspace, vis)
Martin Reinecke's avatar
Martin Reinecke committed
38
    op = GM.getFull().adjoint
39 40 41 42 43
    t1 = time()
    op(img).to_global_data()
    t2 = time()
    op.adjoint(vis).to_global_data()
    t3 = time()
44
    print(t2-t1, t3-t2)
45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61
    N0s.append(N)
    a0s.append(t1 - t0)
    b0s.append(t2 - t1)
    c0s.append(t3 - t2)

print('Measure rest operator')
sc = ift.StatCalculator()
op = GM.getRest().adjoint
for _ in range(10):
    t0 = time()
    res = op(img)
    sc.add(time() - t0)
t_fft = sc.mean
print('FFT shape', res.shape)

plt.scatter(N0s, a0s, label='Gridder mr')
plt.legend()
Martin Reinecke's avatar
Martin Reinecke committed
62
# no idea why this is necessary, but if it is omitted, the range is wrong
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
63
plt.ylim(min(a0s), max(a0s))
64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80
plt.ylabel('time [s]')
plt.title('Initialization')
plt.loglog()
plt.savefig('bench0.png')
plt.close()

plt.scatter(N0s, b0s, color='k', marker='^', label='Gridder mr times')
plt.scatter(N0s, c0s, color='k', label='Gridder mr adjoint times')
plt.axhline(sc.mean, label='FFT')
plt.axhline(sc.mean + np.sqrt(sc.var))
plt.axhline(sc.mean - np.sqrt(sc.var))
plt.legend()
plt.ylabel('time [s]')
plt.title('Apply')
plt.loglog()
plt.savefig('bench1.png')
plt.close()