make test work with mpi

lehner · Oct 6, 2023 · fa40d60 · fa40d60
1 parent 2a3bc02
commit fa40d60
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 4 deletions.
diff --git a/lib/gpt/core/stencil/matrix_vector.py b/lib/gpt/core/stencil/matrix_vector.py
@@ -40,7 +40,6 @@ def __init__(self, lat_matrix, lat_vector, points, code, code_parallel_block_siz
         self.code_parallel_block_size = code_parallel_block_size
         if code_parallel_block_size is None:
             code_parallel_block_size = len(code)
-        assert lat_matrix.grid == lat_vector.grid
         self.obj = cgpt.stencil_matrix_vector_create(
             lat_matrix.v_obj[0],
             lat_vector.v_obj[0],

diff --git a/tests/core/stencil.py b/tests/core/stencil.py
@@ -165,9 +165,15 @@ def stencil_cshift(src, direction):
 rng.cnormal(src)
 cov = g.covariant.shift(U, boundary_phases=[1.0, 1.0, 1.0, 1.0])
 for mu in range(4):
+    pad_U = g.padded_local_fields(U, list(evec[mu]))
+    pad_src = g.padded_local_fields(src, list(evec[mu]))
+
+    p_U = pad_U(U)
+    p_src = pad_src(src)
+
     st = g.stencil.matrix_vector(
-        m,
-        v,
+        p_U[0],
+        p_src,
         [(0, 0, 0, 0), evec[mu], nevec[mu]],
         [
             {
@@ -200,10 +206,16 @@ def stencil_cshift(src, direction):
     def lap(dst, src):
         dst @= -2.0 * src + cov.forward[mu] * src + cov.backward[mu] * src
 
+    p_dst = g.lattice(p_src)
+    dst = g.lattice(src)
+
     ref = g.lattice(src)
     stv = g.lattice(src)
+
     lap(ref, src)
-    st(U, [stv, src])
+    st(p_U, [p_dst, p_src])
+    pad_src.extract(stv, p_dst)
+
     eps2 = g.norm2(stv - ref)
     g.message(f"Stencil covariant laplace versus cshift version: {eps2}")
     assert eps2 < 1e-25