From 4ac671e0c96d948b0df8e3f28f2ef30a86e71cab Mon Sep 17 00:00:00 2001 From: maxlin-ipp Date: Fri, 27 Sep 2024 16:24:41 +0200 Subject: [PATCH] Update `stencil2coo` kernels for shifts > 1 (#442) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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ü --- psydac/linalg/kernels/stencil2coo_kernels.py | 24 ++++++++++---------- psydac/linalg/tests/test_stencil_matrix.py | 8 +++---- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/psydac/linalg/kernels/stencil2coo_kernels.py b/psydac/linalg/kernels/stencil2coo_kernels.py index 55f4fc874..81c278df6 100644 --- a/psydac/linalg/kernels/stencil2coo_kernels.py +++ b/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 diff --git a/psydac/linalg/tests/test_stencil_matrix.py b/psydac/linalg/tests/test_stencil_matrix.py index 404a0f970..cf2e602a6 100644 --- a/psydac/linalg/tests/test_stencil_matrix.py +++ b/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):