fixed test_numba.py in parallel

mfem · Mar 27, 2023 · 4767c62 · 4767c62
1 parent 0901e5e
commit 4767c62
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/test/test_numba.py b/test/test_numba.py
@@ -13,14 +13,15 @@
 if len(sys.argv) > 1 and sys.argv[1] == '-p':
     import mfem.par as mfem
     use_parallel = True
-    from mfem.common.mpi_debug import nicePrint as print
+    from mfem.common.mpi_debug import nicePrint
     from mpi4py import MPI
     myid = MPI.COMM_WORLD.rank
 
 else:
     import mfem.ser as mfem
     use_parallel = False
     myid = 0
+    nicePrint = print
 
 
 class s_coeff(mfem.PyCoefficient):
@@ -321,23 +322,23 @@ def m_func5(p):
     end = time.time()
     a3.Finalize()
     M3 = a3.SpMat()
-    print("Numba (simpler interface) (matrix)", end - start)
+    nicePrint("Numba (simpler interface) (matrix)", end - start)
 
     start = time.time()
     m_func4_complex.SetTime(3.0)
     a4.Assemble()
     end = time.time()
     a4.Finalize()
     M4 = a4.SpMat()
-    print("Numba (complex dependency as complex) (matrix)", end - start)
+    nicePrint("Numba (complex dependency as complex) (matrix)", end - start)
 
     start = time.time()
     m_func4_split.SetTime(3.0)
     a5.Assemble()
     end = time.time()
     a5.Finalize()
     M5 = a5.SpMat()
-    print("Numba (complex dependency as decomposed) (matrix)", end - start)
+    nicePrint("Numba (complex dependency as decomposed) (matrix)", end - start)
 
     #from mfem.commmon.sparse_utils import sparsemat_to_scipycsr
     #csr1 = sparsemat_to_scipycsr(M1, float)
@@ -360,7 +361,7 @@ def compare_mat(M1o, M2o):
     compare_mat(M1, M5)
 
     # print(m_func3.SpaceDimension())
-    print("PASSED")
+    nicePrint("PASSED")
 
 
 if __name__ == '__main__':