fix: mpi tests_007 passing, fixed typo Nz->NZ

Author: elenafuengar

tests/test_007_mpi_lossy_cavity.py

@@ -1,4 +1,5 @@
-import os, sys
+import os
+import sys
 import numpy as np
 import pyvista as pv
 import matplotlib.pyplot as plt
@@ -15,6 +16,9 @@
 
 import pytest
 
+# Run with:
+# mpiexec -n 2 python -m pytest --color=yes -v -s tests/test_007_mpi_lossy_cavity.py
+
 # Turn true when running local
 flag_plot_3D = True
 
@@ -92,7 +96,6 @@ def test_mpi_import(self):
             from mpi4py import MPI
 
             comm = MPI.COMM_WORLD  # Get MPI communicator
-            rank = comm.Get_rank()  # Process ID
             size = comm.Get_size()  # Total number of MPI processes
             if size > 1:
                 use_mpi = True
@@ -184,10 +187,12 @@ def test_mpi_simulation(self):
                 beam.update(solver, n*solver.dt)
                 solver.mpi_one_step()
 
-            Ez = solver.mpi_gather('Ez', x=int(Nx/2), y=int(Ny/2), z=np.s_[::5])
+            Ez = solver.mpi_gather('Ez', x=int(Nx/2), y=int(Ny/2))
             if solver.rank == 0:
                 #print(Ez)
-                assert np.allclose(Ez, self.Ez, rtol=0.1), "Electric field Ez samples MPI failed"
+                print(len(Ez))
+                assert len(Ez) == NZ, "Electric field Ez samples length mismatch"
+                assert np.allclose(Ez[np.s_[::5]], self.Ez, rtol=0.1), "Electric field Ez samples MPI failed"
         else:
             Nt = 3000
             for n in tqdm(range(Nt)):
@@ -197,9 +202,9 @@ def test_mpi_simulation(self):
 
             Ez = solver.E[int(Nx/2), int(Ny/2), np.s_[::5], 'z']
             #print(Ez)
+            assert len(solver.E[int(Nx/2), int(Ny/2), :, 'z']) == NZ, "Electric field Ez samples length mismatch"
             assert np.allclose(Ez, self.Ez, rtol=0.1), "Electric field Ez samples failed"
 
-
     def test_mpi_gather_asField(self):
         # Plot inspect after mpi gather
         global solver
@@ -292,10 +297,13 @@ def test_long_wake_potential(self):
         if use_mpi:
             if solver.rank == 0:
                 #print(wake.WP[::50])
+                print(len(wake.WP))
+                assert len(wake.WP) == 5195, "Wake potential samples length mismatch"
                 assert np.allclose(wake.WP[::50], self.WP, rtol=0.1), "Wake potential samples failed"
                 assert np.cumsum(np.abs(wake.WP))[-1] == pytest.approx(184.43818552913254, 0.1), "Wake potential cumsum MPI failed"
         else:
             #print(wake.WP[::50])
+            assert len(wake.WP) == 5195, "Wake potential samples length mismatch"
             assert np.allclose(wake.WP[::50], self.WP, rtol=0.1), "Wake potential samples failed"
             assert np.cumsum(np.abs(wake.WP))[-1] == pytest.approx(184.43818552913254, 0.1), "Wake potential cumsum MPI failed"
 
@@ -305,12 +313,15 @@ def test_long_impedance(self):
         if use_mpi:
             if solver.rank == 0:
                 #print(wake.Z[::20])
+                print(len(wake.Z))
+                assert len(wake.Z) == 998, "Impedance samples length mismatch"
                 assert np.allclose(np.abs(wake.Z)[::20], np.abs(self.Z), rtol=0.1), "Abs Impedance samples MPI failed"
                 assert np.allclose(np.real(wake.Z)[::20], np.real(self.Z), rtol=0.1), "Real Impedance samples MPI failed"
                 assert np.allclose(np.imag(wake.Z)[::20], np.imag(self.Z), rtol=0.1), "Imag Impedance samples MPI failed"
                 assert np.cumsum(np.abs(wake.Z))[-1] == pytest.approx(250910.51090497518, 0.1), "Abs Impedance cumsum MPI failed"
         else:
             #print(wake.Z[::20])
+            assert len(wake.Z) == 998, "Impedance samples length mismatch"
             assert np.allclose(np.abs(wake.Z)[::20], np.abs(self.Z), rtol=0.1), "Abs Impedance samples failed"
             assert np.allclose(np.real(wake.Z)[::20], np.real(self.Z), rtol=0.1), "Real Impedance samples failed"
             assert np.allclose(np.imag(wake.Z)[::20], np.imag(self.Z), rtol=0.1), "Imag Impedance samples failed"

wakis/gridFIT3D.py

@@ -95,11 +95,7 @@ def __init__(self, xmin=None, xmax=None,
         if stl_solids is not None:
             self._prepare_stl_dicts()
 
-        # primal Grid G base axis x, y, z
-        self.x = x
-        self.y = y
-        self.z = z
-
+        # Grid data
         # generate from file
         if load_from_h5 is not None:
             t0 = time.time()
@@ -114,8 +110,11 @@ def __init__(self, xmin=None, xmax=None,
             return
 
         # generate from custom x,y,z arrays
-        elif self.x is not None and self.y is not None and self.z is not None:
+        elif x is not None and y is not None and z is not None:
             # allow user to set the grid axis manually
+            self.x = x
+            self.y = y
+            self.z = z
             self.Nx = len(self.x) - 1
             self.Ny = len(self.y) - 1
             self.Nz = len(self.z) - 1
@@ -146,7 +145,8 @@ def __init__(self, xmin=None, xmax=None,
 
         self.dx = np.min(np.diff(self.x))
         self.dy = np.min(np.diff(self.y))
-        self.dz = np.min(np.diff(self.z))
+        #self.dz = np.min(np.diff(self.z))
+        self.dz = (self.zmax - self.zmin)/self.Nz
 
         # refine self.x, self.y, self.z using snap points
         self.use_mesh_refinement = use_mesh_refinement
@@ -201,7 +201,6 @@ def __init__(self, xmin=None, xmax=None,
 
 
     def compute_grid(self):
-
         X, Y, Z = np.meshgrid(self.x, self.y, self.z, indexing='ij')
         self.grid = pv.StructuredGrid(X.transpose(), Y.transpose(), Z.transpose())
 
@@ -261,13 +260,13 @@ def mpi_initialize(self):
 
         # MPI subdomain quantities
         self.Nz = self.NZ // (self.size)
-        self.dz = (self.ZMAX - self.ZMIN) / self.Nz
+        self.dz = (self.ZMAX - self.ZMIN) / self.NZ
         self.zmin = self.rank * self.Nz * self.dz + self.ZMIN
         self.zmax = (self.rank+1) * self.Nz * self.dz + self.ZMIN
 
         if self.verbose:
             print(f"MPI rank {self.rank} of {self.size} initialized with \
-                                zmin={self.zmin}, zmax={self.zmax}, Nz={self.Nz}")
+                        zmin={self.zmin}, zmax={self.zmax}, Nz={self.Nz}")
         # Add ghost cells
         self.n_ghosts = 1
         if self.rank > 0:
@@ -341,9 +340,9 @@ def mark_cells_in_stl(self):
 
             # mark cells in stl [True == in stl, False == out stl]
             try:
-                select = self.grid.select_enclosed_points(surf, tolerance=tol)
+                select = self.grid.select_enclosed_points(surf, tolerance=stl_tolerance)
             except Exception:
-                select = self.grid.select_enclosed_points(surf, tolerance=tol, check_surface=False)
+                select = self.grid.select_enclosed_points(surf, tolerance=stl_tolerance, check_surface=False)
                 if self.verbose > 1:
                     print(f'[!] Warning: stl solid {key} may have issues with closed surfaces. Consider checking the STL file.')