Update stencil2coo kernels for shifts > 1 (#442)

Fix the indexing in all kernels in `stencil2coo_kernels.py` so that the
method `toarray` works for `StencilMatrix` objects in 1D/2D/3D with
`shift` larger than 1.

The unit tests for `toarray` in `test_stencil_matrix.py` have been
parametrized to run with `shift` values of 1 and 2.

---------

Co-authored-by: Yaman Güçlü <yaman.guclu@gmail.com>

psydac/linalg/kernels/stencil2coo_kernels.py

@@ -19,7 +19,7 @@ def stencil2coo_1d_C(A:'T[:,:]', data:'T[:]', rows:'int64[:]', cols:'int64[:]',
         for j1 in range(ncl1):
             value = A[i1+pp1,j1]
             if abs(value) == 0.0:continue
-            J = ((I//cm1)*dm1+j1-dp1)%nc1
+            J = ((I*dm1//cm1)+j1-dp1)%nc1
             rows[nnz] = I
             cols[nnz] = J
             data[nnz] = value
@@ -38,7 +38,7 @@ def stencil2coo_1d_F(A:'T[:,:]', data:'T[:]', rows:'int64[:]', cols:'int64[:]',
         for j1 in range(ncl1):
             value = A[i1+pp1,j1]
             if abs(value) == 0.0:continue
-            J = ((I//cm1)*dm1+j1-dp1)%nc1
+            J = ((I*dm1//cm1)+j1-dp1)%nc1
             rows[nnz] = I
             cols[nnz] = J
             data[nnz] = value
@@ -66,8 +66,8 @@ def stencil2coo_2d_C(A:'T[:,:,:,:]', data:'T[:]', rows:'int64[:]', cols:'int64[:
                 for j2 in range(ncl2):
                     value = A[i1+pp1,i2+pp2,j1,j2]
                     if abs(value) == 0.0:continue
-                    jj1 = ((ii1//cm1)*dm1+j1-dp1)%nc1
-                    jj2 = ((ii2//cm2)*dm2+j2-dp2)%nc2
+                    jj1 = ((ii1*dm1//cm1)+j1-dp1)%nc1
+                    jj2 = ((ii2*dm2//cm2)+j2-dp2)%nc2
 
                     J   = jj1*nc2 + jj2
 
@@ -97,8 +97,8 @@ def stencil2coo_2d_F(A:'T[:,:,:,:]', data:'T[:]', rows:'int64[:]', cols:'int64[:
                 for j2 in range(ncl2):
                     value = A[i1+pp1,i2+pp2,j1,j2]
                     if abs(value) == 0.0:continue
-                    jj1 = ((ii1//cm1)*dm1+j1-dp1)%nc1
-                    jj2 = ((ii2//cm2)*dm2+j2-dp2)%nc2
+                    jj1 = ((ii1*dm1//cm1)+j1-dp1)%nc1
+                    jj2 = ((ii2*dm2//cm2)+j2-dp2)%nc2
 
                     J   = jj2*nc1 + jj1
 
@@ -132,9 +132,9 @@ def stencil2coo_3d_C(A:'T[:,:,:,:,:,:]', data:'T[:]', rows:'int64[:]', cols:'int
                         for j3 in range(ncl3):
                             value = A[i1+pp1,i2+pp2,i3+pp3,j1,j2,j3]
                             if abs(value) == 0.0:continue
-                            jj1 = ((ii1//cm1)*dm1+j1-dp1)%nc1
-                            jj2 = ((ii2//cm2)*dm2+j2-dp2)%nc2
-                            jj3 = ((ii3//cm3)*dm3+j3-dp3)%nc3
+                            jj1 = ((ii1*dm1//cm1)+j1-dp1)%nc1
+                            jj2 = ((ii2*dm2//cm2)+j2-dp2)%nc2
+                            jj3 = ((ii3*dm3//cm3)+j3-dp3)%nc3
 
                             J   = jj1*nc2*nc3 + jj2*nc3 + jj3
 
@@ -170,9 +170,9 @@ def stencil2coo_3d_F(A:'T[:,:,:,:,:,:]', data:'T[:]', rows:'int64[:]', cols:'int
                         for j3 in range(ncl3):
                             value = A[i1+pp1,i2+pp2,i3+pp3,j1,j2,j3]
                             if abs(value) == 0.0:continue
-                            jj1 = ((ii1//cm1)*dm1+j1-dp1)%nc1
-                            jj2 = ((ii2//cm2)*dm2+j2-dp2)%nc2
-                            jj3 = ((ii3//cm3)*dm3+j3-dp3)%nc3
+                            jj1 = ((ii1*dm1//cm1)+j1-dp1)%nc1
+                            jj2 = ((ii2*dm2//cm2)+j2-dp2)%nc2
+                            jj3 = ((ii3*dm3//cm3)+j3-dp3)%nc3
 
                             J   = jj3*nc1*nc2 + jj2*nc1 + jj1
 

psydac/linalg/tests/test_stencil_matrix.py

@@ -433,7 +433,7 @@ def test_stencil_matrix_2d_serial_spurious_entries( dtype, p1, p2, s1, s2, P1, P
 @pytest.mark.parametrize('dtype', [float])
 @pytest.mark.parametrize('n1', [7])
 @pytest.mark.parametrize('p1', [1,2])
-@pytest.mark.parametrize('s1', [1])
+@pytest.mark.parametrize('s1', [1, 2])
 @pytest.mark.parametrize('P1', [True, False])
 def test_stencil_matrix_1d_serial_toarray( dtype, n1, p1, s1, P1):
     # Select non-zero values based on diagonal index
@@ -489,8 +489,8 @@ def test_stencil_matrix_1d_serial_toarray( dtype, n1, p1, s1, P1):
 @pytest.mark.parametrize('n2', [8, 7])
 @pytest.mark.parametrize('p1', [1, 2])
 @pytest.mark.parametrize('p2', [1, 3])
-@pytest.mark.parametrize('s1', [1])
-@pytest.mark.parametrize('s2', [1])
+@pytest.mark.parametrize('s1', [1, 2])
+@pytest.mark.parametrize('s2', [1, 2])
 @pytest.mark.parametrize('P1', [True, False])
 @pytest.mark.parametrize('P2', [True, False])
 def test_stencil_matrix_2d_serial_toarray( dtype, n1, n2, p1, p2, s1, s2, P1, P2):
@@ -2191,7 +2191,7 @@ def test_stencil_matrix_2d_serial_backend_switch(dtype, n1, n2, p1, p2, s1, s2,
 @pytest.mark.parametrize('dtype', [float, complex])
 @pytest.mark.parametrize('n1', [20, 67])
 @pytest.mark.parametrize('p1', [1, 2, 3])
-@pytest.mark.parametrize('sh1', [1])
+@pytest.mark.parametrize('sh1', [1, 2])
 @pytest.mark.parametrize('P1', [True, False])
 @pytest.mark.parallel
 def test_stencil_matrix_1d_parallel_toarray(dtype, n1, p1, sh1, P1):