new demo

demo_wstack.py

+import nifty_gridder as ng
+import numpy as np
+
+
+def _l2error(a, b):
+    return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))
+
+
+def explicit_gridder(uvw,freq,ms,nxdirty,nydirty,xpixsize,ypixsize):
+    speedoflight = 299792458.
+
+    x, y = np.meshgrid(*[-ss/2 + np.arange(ss) for ss in [nxdirty, nydirty]],
+                       indexing='ij')
+    x *= xpixsize
+    y *= ypixsize
+    res = np.zeros((nxdirty, nydirty))+0j
+
+    eps = x**2+y**2
+    nm1 = -eps/(np.sqrt(1.-eps)+1.);
+    n = nm1+1
+    for row in range(ms.shape[0]):
+        for chan in range(ms.shape[1]):
+            phase = freq[chan]/speedoflight*(x*uvw[row, 0] + y*uvw[row, 1] + uvw[row, 2]*nm1)
+            res += (ms[row,chan]*np.exp(2j*np.pi*phase))
+    return res/n
+
+
+def test_against_wdft(nrow, nchan, nxdirty, nydirty, fov, epsilon, nthreads, test_against_explicit):
+    print("\n\nTesting gridding/degridding with {} rows and {} frequency channels".format(nrow, nchan))
+    print("Dirty image has {}x{} pixels, FOV={} degrees".format(nxdirty,nydirty,fov))
+    print("Requested accuracy: {}".format(epsilon))
+    print("Number of threads: {}".format(nthreads))
+
+    single = epsilon>5e-6
+
+    speedoflight = 299792458.
+    np.random.seed(40)
+    xpixsize = fov*np.pi/180/nxdirty
+    ypixsize = fov*np.pi/180/nydirty
+    f0 = 1e9
+    freq = f0 + np.arange(nchan)*(f0/nchan)
+    uvw = (np.random.rand(nrow, 3)-0.5)/(xpixsize*f0/speedoflight)
+    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
+    tdirty = np.random.rand(nxdirty, nydirty)-0.5 + 1j*(np.random.rand(nxdirty, nydirty)-0.5)
+
+    single = epsilon>5e-6
+    if single:
+        print("\nCalling single-precision functions")
+        uvw = uvw.astype("f4")
+        freq = freq.astype("f4")
+        ms = ms.astype("c8")
+        tdirty = tdirty.astype("c8")
+        gridder, degridder = ng.full_gridding_f, ng.full_degridding_f
+    else:
+        print("\nCalling double-precision functions")
+        gridder, degridder = ng.full_gridding, ng.full_degridding
+
+    if test_against_explicit:
+        print("\nTesting against explicit transform (potentially VERY slow!)...")
+        truth = explicit_gridder(uvw,freq,ms,nxdirty,nydirty,xpixsize,ypixsize)
+        res = gridder(uvw,freq,ms,None,nxdirty,nydirty,xpixsize,ypixsize,epsilon,nthreads)
+        print("L2 error between explicit transform and gridder:",_l2error(truth,res))
+
+    # test adjointness
+    print("\nTesting adjointness of the gridding/degridding operation")
+    adj1 = np.vdot(gridder(uvw,freq,ms,None,nxdirty,nydirty,xpixsize,ypixsize,epsilon,nthreads),tdirty)
+    adj2 = np.vdot(ms,degridder(uvw,freq,tdirty,None,xpixsize,ypixsize,epsilon,nthreads))
+    print("adjointness test:", np.abs(adj1-adj2)/np.maximum(np.abs(adj1),np.abs(adj2)))
+
+
+test_against_wdft(100, 20, 64, 130, 0.5, 1e-12, 3, True)
+from time import time
+t0=time()
+test_against_wdft(1000,300, 1024, 1024, 2., 1e-12, 4, False)
+print(time()-t0)