Feature/solver inspect geometry

tests/test_007_mpi_lossy_cavity.py

@@ -86,19 +86,19 @@ class TestMPILossyCavity:
                     -6.04105997e+01 ,-3.06532160e+01 ,-1.17749936e+01 ,-3.12574866e+00,
                     -7.35339521e-01 ,-1.13085658e-01 , 7.18247535e-01 , 8.73829036e-02])
 
-    gridLogs = {'use_mesh_refinement': False, 'Nx': 60, 'Ny': 60, 'Nz': 140, 'dx': 0.00866666634877522, 
-                'dy': 0.00866666634877522, 'dz': 0.005714285799435207, 
-                'xmin': -0.25999999046325684, 'xmax': 0.25999999046325684, 
-                'ymin': -0.25999999046325684, 'ymax': 0.25999999046325684, 
-                'zmin': -0.25, 'zmax': 0.550000011920929, 
+    gridLogs = {'use_mesh_refinement': False, 'Nx': 60, 'Ny': 60, 'Nz': 140, 'dx': 0.00866666634877522,
+                'dy': 0.00866666634877522, 'dz': 0.005714285799435207,
+                'xmin': -0.25999999046325684, 'xmax': 0.25999999046325684,
+                'ymin': -0.25999999046325684, 'ymax': 0.25999999046325684,
+                'zmin': -0.25, 'zmax': 0.550000011920929,
                 'stl_solids': {'cavity': 'tests/stl/007_vacuum_cavity.stl', 'shell': 'tests/stl/007_lossymetal_shell.stl'},
-                'stl_materials': {'cavity': 'vacuum', 'shell': [30, 1.0, 30]}, 
+                'stl_materials': {'cavity': 'vacuum', 'shell': [30, 1.0, 30]},
                 'gridInitializationTime': 0}
-    
+
     solverLogs = {'use_gpu': False, 'use_mpi': False, 'background': 'pec',
-                'bc_low': ['pec', 'pec', 'pec'], 'bc_high': ['pec', 'pec', 'pec'], 
+                'bc_low': ['pec', 'pec', 'pec'], 'bc_high': ['pec', 'pec', 'pec'],
                 'dt': 6.970326728398966e-12, 'solverInitializationTime': 0}
-    
+
     wakeSolverLogs = {'ti': 2.8516132094735135e-09, 'q': 1e-09, 'sigmaz': 0.1, 'beta': 1.0,
                        'xsource': 0.0, 'ysource': 0.0, 'xtest': 0.0, 'ytest': 0.0, 'chargedist': None,
                        'skip_cells': 10, 'results_folder': 'tests/007_results/', 'wakelength': 10.0, 'simulationTime': 0}
@@ -252,29 +252,6 @@ def test_mpi_plot1D(self):
                 xscale='linear', yscale='linear',
                 off_screen=True, title=self.img_folder+'Ez1d', n=3000)
 
-    @pytest.mark.skipif(not flag_plot_3D, reason="Requires interactive plotting")
-    def test_mpi_plot3D(self):
-        # Plot Abs Electric field on domain
-        # disabled when mpi = True
-        global solver
-        solver.plot3D('E', component='Abs',
-                cmap='rainbow', clim=[0, 500],
-                add_stl=['cavity', 'shell'], stl_opacity=0.1,
-                clip_interactive=True, clip_normal='-y')
-
-    @pytest.mark.skipif(not flag_plot_3D, reason="Requires interactive plotting")
-    def test_mpi_plot3DonSTL(self):
-        # Plot Abs Electric field on STL solid `cavity`
-        # disabled when mpi = True
-        global solver
-        solver.plot3DonSTL('E', component='Abs',
-                        cmap='rainbow', clim=[0, 500],
-                        stl_with_field='cavity', field_opacity=1.0,
-                        stl_transparent='shell', stl_opacity=0.1, stl_colors='white',
-                        clip_plane=True, clip_normal='-y', clip_origin=[0,0,0],
-                        off_screen=False, zoom=1.2, title=self.img_folder+'Ez3d')
-
-
     def test_mpi_wakefield(self):
         # Reset fields
         global solver
@@ -343,7 +320,7 @@ def test_long_impedance(self):
             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"
             assert np.cumsum(np.abs(wake.Z))[-1] == pytest.approx(250910.51090497518, 0.1), "Abs Impedance cumsum failed"
-    
+
     def test_log_file(self):
         # Helper function to compare nested dicts with float tolerance
         def assert_dict_allclose(d1, d2, rtol=1e-6, atol=1e-12, path=""):
@@ -379,7 +356,7 @@ def assert_dict_allclose(d1, d2, rtol=1e-6, atol=1e-12, path=""):
 
         global solver
         # Exclude timing info from comparison as they can vary between runs
-        solver.logger.grid["gridInitializationTime"] = 0 
+        solver.logger.grid["gridInitializationTime"] = 0
         solver.logger.solver["solverInitializationTime"] = 0
         solver.logger.wakeSolver["simulationTime"] = 0
         self.solverLogs['use_mpi'] = use_mpi

tests/test_008_3D_plotting.py

@@ -0,0 +1,165 @@
+import os
+import sys
+import numpy as np
+import pyvista as pv
+import matplotlib.pyplot as plt
+
+sys.path.append('../wakis')
+
+from tqdm import tqdm
+from scipy.constants import c
+
+from wakis import SolverFIT3D
+from wakis import GridFIT3D
+from wakis.sources import Beam
+
+import pytest
+
+# Turn true when running local
+flag_plot_3D = False
+flag_offscreen = False
+
+@pytest.mark.slow
+class Test3Dplotting:
+
+    img_folder = 'tests/008_img/'
+
+    def test_simulation(self):
+
+        # ---------- Domain setup ---------
+
+        # Geometry & Materials
+        solid_1 = 'tests/stl/007_vacuum_cavity.stl'
+        solid_2 = 'tests/stl/007_lossymetal_shell.stl'
+
+        stl_solids = {'cavity': solid_1,
+                    'shell': solid_2
+                    }
+
+        stl_materials = {'cavity': 'vacuum',
+                        'shell': [30, 1.0, 30] #[eps_r, mu_r, sigma[S/m]]
+                        }
+
+        stl_colors = {'cavity': 'vacuum',
+                    'shell': [1, 1, 1]}
+
+        # Extract domain bounds from geometry
+        solids = pv.read(solid_1) + pv.read(solid_2)
+        xmin, xmax, ymin, ymax, ZMIN, ZMAX = solids.bounds
+
+        # Number of mesh cells
+        Nx = 60
+        Ny = 60
+        NZ = 140
+
+        grid = GridFIT3D(xmin, xmax, ymin, ymax, ZMIN, ZMAX,
+                        Nx, Ny, NZ,
+                        use_mpi=False, # Enables MPI subdivision of the domain
+                        stl_solids=stl_solids,
+                        stl_materials=stl_materials,
+                        stl_colors=stl_colors,
+                        stl_scale=1.0,
+                        stl_rotate=[0,0,0],
+                        stl_translate=[0,0,0],
+                        verbose=1)
+
+        # ------------ Beam source & Wake ----------------
+        # Beam parameters
+        sigmaz = 10e-2      #[m] -> 2 GHz
+        q = 1e-9            #[C]
+        beta = 1.0          # beam beta
+        xs = 0.             # x source position [m]
+        ys = 0.             # y source position [m]
+        ti = 3*sigmaz/c     # injection time [s]
+
+        beam = Beam(q=q, sigmaz=sigmaz, beta=beta,
+                    xsource=xs, ysource=ys, ti=ti)
+
+        # ----------- Solver & Simulation ----------
+        # boundary conditions
+        bc_low=['pec', 'pec', 'pec']
+        bc_high=['pec', 'pec', 'pec']
+
+        # Solver setup
+        global solver
+        solver = SolverFIT3D(grid,
+                            bc_low=bc_low,
+                            bc_high=bc_high,
+                            use_stl=True,
+                            use_mpi=False, # Activate MPI
+                            bg='pec' # Background material
+                            )
+
+        # -------------- Output folder ---------------------
+        if not os.path.exists(self.img_folder):
+            os.mkdir(self.img_folder)
+
+        # -------------- Custom time loop  -----------------
+        Nt = 3000
+        for n in tqdm(range(Nt)):
+            beam.update(solver, n*solver.dt)
+            solver.one_step()
+
+    @pytest.mark.skipif(not flag_plot_3D, reason="Requires interactive plotting")
+    def test_grid_inspect(self):
+        # Plot grid and imported solids
+        global solver
+        solver.grid.inspect(add_stl=['cavity', 'shell'],
+                            stl_opacity=0.1, off_screen=flag_offscreen,
+                            anti_aliasing='ssaa')
+
+    @pytest.mark.skipif(not flag_plot_3D, reason="Requires interactive plotting")
+    def test_grid_plot_solids(self):
+        # Plot only imported solids
+        global solver
+        solver.grid.plot_solids(bounding_box=True,
+                                show_grid=True,
+                                opacity=1,
+                                specular=0.5,
+                                smooth_shading=False,
+                                off_screen=flag_offscreen,)
+
+    @pytest.mark.skipif(not flag_plot_3D, reason="Requires interactive plotting")
+    def test_grid_stl_mask(self):
+        # Plot STL solid masks in the grid
+        global solver
+        solver.grid.plot_stl_mask(stl_solid='cavity',
+                                  cmap='viridis',
+                                  add_stl='all',
+                                  stl_opacity=0.5,
+                                  smooth_shading=False,
+                                  anti_aliasing='ssaa',
+                                  ymax=0.0,
+                                  off_screen=flag_offscreen,)
+
+    @pytest.mark.skipif(not flag_plot_3D, reason="Requires interactive plotting")
+    def test_solver_inspect(self):
+        # Plot imported solids and beam source and integraiton path
+        global solver
+        pl = solver.inspect(window_size=(1200,800), off_screen=flag_offscreen,
+                            specular=0.,opacity=1, inactive_opacity=0.1,
+                            add_silhouette=True,)
+        if flag_offscreen:
+            pl.export_html(self.img_folder+'solver_inspect.html')
+            pl.screenshot(self.img_folder+'solver_inspect.png')
+
+    @pytest.mark.skipif(not flag_plot_3D, reason="Requires interactive plotting")
+    def test_plot3D(self):
+        # Plot Abs Electric field on domain
+        global solver
+        solver.plot3D('E', component='Abs',
+                cmap='rainbow', clim=[0, 500],
+                add_stl=['cavity', 'shell'], stl_opacity=0.1,
+                clip_interactive=True, clip_normal='-y',
+                off_screen=flag_offscreen)
+
+    @pytest.mark.skipif(not flag_plot_3D, reason="Requires interactive plotting")
+    def test_plot3DonSTL(self):
+        # Plot Abs Electric field on STL solid `cavity`
+        global solver
+        solver.plot3DonSTL('E', component='Abs',
+                        cmap='rainbow', clim=[0, 500],
+                        stl_with_field='cavity', field_opacity=1.0,
+                        stl_transparent='shell', stl_opacity=0.1, stl_colors='white',
+                        clip_plane=True, clip_normal='-y', clip_origin=[0,0,0],
+                        off_screen=flag_offscreen, zoom=1.2, title=self.img_folder+'Ez3d')

wakis/geometry.py

@@ -68,8 +68,16 @@ def extract_materials_from_stp(stp_file):
 
     return stl_materials
 
-def extract_solids_from_stp(stp_file, path=''):
-    if path and not path.endswith('/'):
+def extract_solids_from_stp(stp_file, path=None):
+    '''
+    Extracts a mapping from solid names to STL file names from a STEP (.stp) file.  
+    Args:
+        stp_file (str): Path to the STEP file.
+        path (str) (optional): default: None, path to save the STL (.stl) files
+    Returns:
+        dict[str, str]: A dictionary mapping solid names to STL file names.
+    '''    
+    if path is not None and not path.endswith('/'):
         path += '/'
     solids, materials = extract_names_from_stp(stp_file)
     stl_solids = {}
@@ -85,7 +93,10 @@ def extract_solids_from_stp(stp_file, path=''):
         solid_re = re.sub(r'[^a-zA-Z0-9_-]', '-', solid)
         mat_re = re.sub(r'[^a-zA-Z0-9_-]', '-', mat)
         name = f'{str(i).zfill(3)}_{solid_re}_{mat_re}'
-        stl_solids[f'{str(i).zfill(3)}_{solid_re}'] = path + name + '.stl'
+        if path is not None:
+            stl_solids[f'{str(i).zfill(3)}_{solid_re}'] = path + name + '.stl'
+        else:
+            stl_solids[f'{str(i).zfill(3)}_{solid_re}'] = name + '.stl'
 
     return stl_solids
 
@@ -195,7 +206,7 @@ def get_stp_unit_scale(stp_file):
 
         return 1.0
 
-def generate_stl_solids_from_stp(stp_file, results_path=''):
+def generate_stl_solids_from_stp(stp_file, results_path=None):
     """
     Extracts solid objects from a STEP (.stp) file and exports them as STL files.
 
@@ -238,16 +249,13 @@ def generate_stl_solids_from_stp(stp_file, results_path=''):
         scaled_solids = [solid.scale(scale_factor) for solid in stp.objects[0]]
         stp.objects = [scaled_solids]
 
-    solid_dict = extract_solids_from_stp(stp_file)
-
-    if not results_path.endswith('/'):
-        results_path += '/'
+    solid_dict = extract_solids_from_stp(stp_file, results_path)
 
     print(f"Generating stl from file: {stp_file}... ")
     for i, obj in enumerate(stp.objects[0]):
         name = solid_dict[list(solid_dict.keys())[i]]
         print(name)
-        obj.exportStl(results_path+name)
+        obj.exportStl(name)
 
     return solid_dict