Simplify

test.py

@@ -4,7 +4,6 @@ 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
 
@@ -355,6 +354,8 @@ def test_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
 @pmp("nchan", (1, 5))
 @pmp('epsilon', (1e-3, 1e-6))
 def test_holo_from_correlations(nxdirty, nydirty, nchan, nrow, epsilon):
+    pytest.importorskip("scipy")
+    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
@@ -367,16 +368,15 @@ def test_holo_from_correlations(nxdirty, nydirty, nchan, nrow, epsilon):
     col = np.empty((n,))
     data = np.empty((n,))
     ind = np.arange(0, gridsize)
-    ycoord0 = ind % ny
-    xcoord0 = ind-ycoord0
-    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] = ind
-        foo = (ycoord0+dv) % ny + xcoord0
+    y0 = ind % ny
+    x0 = ind - y0
+    for ii, (du, dv) in enumerate(product(*(2*(range(-W+1, W),)))):
+        slc = slice(ii*gridsize, (ii+1)*gridsize)
+        foo = (y0+dv) % ny + x0
         yy = foo % ny
-        col[ilow:ihigh] = (foo-yy+du*ny) % gridsize + yy
+        col[slc] = (foo-yy+du*ny) % gridsize + yy
+        row[slc] = ind
+        data[slc] = ng.get_correlations(bl, conf, idx, du=du, dv=dv).ravel()
     mat = coo_matrix((data, (row, col)), shape=(gridsize, gridsize))
     res2 = mat.dot(grid.ravel()).reshape(nx, ny)
     assert_allclose(res1, res2)