Simplify loop

test.py

+from itertools import product
+
 import nifty_gridder as ng
 import numpy as np
 import pytest
 from numpy.testing import assert_, assert_allclose, assert_array_almost_equal
+from scipy.sparse import coo_matrix
 
 pmp = pytest.mark.parametrize
 
@@ -345,13 +348,13 @@ def test_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
     res1 /= n
     assert_(_l2error(res0, res1) < 1e-4)
 
+
 @pmp('nxdirty', (32, 16))
 @pmp('nydirty', (32, 64))
 @pmp("nrow", (100, 1000))
 @pmp("nchan", (1, 5))
 @pmp('epsilon', (1e-3, 1e-6))
 def test_holo_from_correlations(nxdirty, nydirty, nchan, nrow, epsilon):
-    from scipy.sparse import coo_matrix
     np.random.seed(42)
     bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
     nx, ny, gridsize = 2*nxdirty, 2*nydirty, 4*nxdirty*nydirty
@@ -363,21 +366,18 @@ def test_holo_from_correlations(nxdirty, nydirty, nchan, nrow, epsilon):
     row = np.zeros((n,))
     col = np.zeros((n,))
     data = np.zeros((n,))
-    ind = 0
     foo = np.arange(0, gridsize)
-    for ii in range(-W+1, W):
-        for jj in range(-W+1, W):
-            ilow = ind*gridsize
-            ihigh = (ind+1)*gridsize
-            data[ilow:ihigh] = ng.get_correlations(bl, conf, idx, du=ii, dv=jj).ravel()
-            row[ilow:ihigh] = foo
-            ycoord = foo % ny
-            xcoord = foo-ycoord
-            tmp = (ycoord+jj) % ny + xcoord
-            ycoord = tmp % ny
-            xcoord = tmp-ycoord
-            col[ilow:ihigh] = (xcoord + ii*ny) % gridsize + ycoord
-            ind += 1
+    for ii, (du, dv) in enumerate(product(range(-W+1, W), range(-W+1, W))):
+        ilow = ii*gridsize
+        ihigh = (ii+1)*gridsize
+        data[ilow:ihigh] = ng.get_correlations(bl, conf, idx, du=du, dv=dv).ravel()
+        row[ilow:ihigh] = foo
+        ycoord = foo % ny
+        xcoord = foo-ycoord
+        tmp = (ycoord+dv) % ny + xcoord
+        ycoord = tmp % ny
+        xcoord = tmp-ycoord
+        col[ilow:ihigh] = (xcoord + du*ny) % gridsize + ycoord
     mat = coo_matrix((data, (row, col)), shape=(gridsize, gridsize))
     res2 = mat.dot(grid.ravel()).reshape(nx, ny)
     assert_allclose(res1, res2)