Feature/solver inspect geometry

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
 

wakis/gridFIT3D.py

@@ -12,6 +12,7 @@
 
 from .field import Field
 from .logger import Logger
+from .materials import material_colors
 
 try:
     from mpi4py import MPI
@@ -165,8 +166,8 @@ def __init__(self, xmin=None, xmax=None,
             if verbose:
                 print('Applying mesh refinement...')
             if self.snap_points is None and stl_solids is not None:
-                self.compute_snap_points(snap_solids=snap_solids, snap_tol=snap_tol)
-            self.refine_xyz_axis(method=refinement_method, tol=refinement_tol)
+                self._compute_snap_points(snap_solids=snap_solids, snap_tol=snap_tol)
+            self._refine_xyz_axis(method=refinement_method, tol=refinement_tol)
 
         if verbose:
             print(f'Generating grid with {self.Nx*self.Ny*self.Nz} mesh cells...')
@@ -183,31 +184,29 @@ def __init__(self, xmin=None, xmax=None,
             self.NZ = None
             self.Z = None
             if imported_mpi:
-                self.mpi_initialize()
+                self._mpi_initialize()
                 if self.verbose:
                     print(f"MPI initialized for {self.rank} of {self.size}")
             else:
                 raise ImportError("[!] mpi4py is required when use_mpi=True but was not found")
 
         # grid G and tilde grid ~G, lengths and inverse areas
-        self.compute_grid()
+        self._compute_grid()
 
         # tolerance for stl import tol*min(dx,dy,dz)
         if verbose:
             print('Importing STL solids...')
         self.stl_tol = stl_tol
         if stl_solids is not None:
-            self.mark_cells_in_stl()
-            if stl_colors is None:
-                self.assign_colors()
+            self._mark_cells_in_stl()
 
         if verbose:
             print(f'Total grid initialization time: {time.time() - t0} s')
 
         self.gridInitializationTime = time.time()-t0
         self.update_logger(['gridInitializationTime'])
 
-    def compute_grid(self):
+    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())
 
@@ -246,7 +245,7 @@ def compute_grid(self):
         self.itA.field_z = np.divide(1.0, aux, out=np.zeros_like(aux), where=aux!=0)
         del aux
 
-    def mpi_initialize(self):
+    def _mpi_initialize(self):
         comm = MPI.COMM_WORLD  # Get MPI communicator
         self.comm = comm
         self.rank = self.comm.Get_rank()
@@ -338,7 +337,25 @@ def _prepare_stl_dicts(self):
                 stl_translate[key] = self.stl_translate
             self.stl_translate = stl_translate
 
-    def mark_cells_in_stl(self):
+        if type(self.stl_colors) is not dict:
+            if self.stl_colors is None:
+                self._assign_colors()
+            elif self.stl_colors is str: # single color for all solids
+                stl_colors = {}
+                for key in self.stl_solids.keys():
+                    stl_colors[key] = self.stl_colors
+                self.stl_colors = stl_colors
+            elif type(self.stl_colors) is list:
+                stl_colors = {}
+                try:
+                    for i, key in enumerate(self.stl_solids.keys()):
+                        stl_colors[key] = self.stl_colors[i]
+                    self.stl_colors = stl_colors
+                except IndexError:
+                    raise Exception('If `stl_colors` is a list, it must have the same length as `stl_solids`.')
+                    self._assign_colors()
+
+    def _mark_cells_in_stl(self):
         # Obtain masks with grid cells inside each stl solid
         stl_tolerance = np.min([self.dx, self.dy, self.dz])*self.stl_tol
         for key in self.stl_solids.keys():
@@ -378,7 +395,7 @@ def read_stl(self, key):
 
         return surf
 
-    def compute_snap_points(self, snap_solids=None, snap_tol=1e-8):
+    def _compute_snap_points(self, snap_solids=None, snap_tol=1e-8):
         if self.verbose > 1:
             print(' * Calculating snappy points...')
         # Support for user-defined stl_keys as list
@@ -521,7 +538,7 @@ def loss_function(x, x0, is_snap):
         # transform back to [xmin, xmax]
         return result.x*(xmax-xmin)+xmin
 
-    def refine_xyz_axis(self, method='insert', tol=1e-6):
+    def _refine_xyz_axis(self, method='insert', tol=1e-6):
         '''Refine the grid in the x, y, z axis
         using the snap points extracted from the stl solids.
         The snap points are used to refine the grid '''
@@ -551,7 +568,7 @@ def refine_xyz_axis(self, method='insert', tol=1e-6):
         if self.verbose:
             print(f"Refined grid: Nx = {len(self.x)}, Ny ={len(self.y)}, Nz = {len(self.z)}")
 
-    def assign_colors(self):
+    def _assign_colors(self):
         '''Classify colors assigned to each solid
         based on the categories in `material_colors` dict
         inside `materials.py`
@@ -564,15 +581,15 @@ def assign_colors(self):
                 self.stl_colors[key] = mat
             elif len(mat) == 2:
                 if mat[0] is np.inf:  #eps_r
-                    self.stl_colors[key] = 'pec'
+                    self.stl_colors[key] = material_colors['pec']
                 elif mat[0] > 1.0:    #eps_r
-                    self.stl_colors[key] = 'dielectric'
+                    self.stl_colors[key] = material_colors['dielectric']
                 else:
-                    self.stl_colors[key] = 'vacuum'
+                    self.stl_colors[key] = material_colors['vacuum']
             elif len(mat) == 3:
-                self.stl_colors[key] = 'lossy metal'
+                self.stl_colors[key] = material_colors['lossy metal']
             else:
-                self.stl_colors[key] = 'other'
+                self.stl_colors[key] = material_colors['other']
 
     def _add_logo_widget(self, pl):
         """Add packaged logo via importlib.resources (Python 3.9+)."""
@@ -592,7 +609,8 @@ def _add_logo_widget(self, pl):
                     print(f'[!] Could not add logo widget: {e}')
 
     def plot_solids(self, bounding_box=False, show_grid=False, anti_aliasing=None,
-                    opacity=1.0, specular=0.5, offscreen=False, **kwargs):
+                    opacity=1.0, specular=0.5, smooth_shading=False, 
+                    offscreen=False, **kwargs):
         """
         Generates a 3D visualization of the imported STL geometries using PyVista.
 
@@ -627,22 +645,16 @@ def plot_solids(self, bounding_box=False, show_grid=False, anti_aliasing=None,
 
         """
 
-        from .materials import material_colors
-
         pl = pv.Plotter()
         pl.add_mesh(self.grid, opacity=0., name='grid', show_scalar_bar=False)
         for key in self.stl_solids:
-            try:
-                color = material_colors[self.stl_colors[key]] # match library e.g. 'vacuum'
-            except KeyError:
-                color = self.stl_colors[key] # specifies color e.g. 'tab:red'
-
-            if self.stl_colors[key] == 'vacuum' or self.stl_materials[key] == 'vacuum':
+            color = self.stl_colors[key]
+            if self.stl_materials[key] == 'vacuum':
                 _opacity = 0.3
             else:
                 _opacity = opacity
             pl.add_mesh(self.read_stl(key), color=color,
-                        opacity=_opacity, specular=specular, smooth_shading=True,
+                        opacity=_opacity, specular=specular, smooth_shading=smooth_shading,
                         **kwargs)
 
         pl.set_background('mistyrose', top='white')
@@ -668,8 +680,8 @@ def plot_solids(self, bounding_box=False, show_grid=False, anti_aliasing=None,
 
     def plot_stl_mask(self, stl_solid, cmap='viridis', bounding_box=True, show_grid=True,
                       add_stl='all', stl_opacity=0., stl_colors=None,
-                      xmax=None, ymax=None, zmax=None,
-                      anti_aliasing='ssaa', offscreen=False):
+                      xmax=None, ymax=None, zmax=None, 
+                      anti_aliasing='ssaa', smooth_shading=False, offscreen=False):
 
         """
         Interactive 3D visualization of the structured grid mask and imported STL geometries.
@@ -704,7 +716,7 @@ def plot_stl_mask(self, stl_solid, cmap='viridis', bounding_box=True, show_grid=
             Color(s) of the STL surfaces:
             * str   → single color for all STL surfaces
             * list  → per-solid colors, in order
-            * dict  → mapping from STL key to color (using `material_colors`)
+            * dict  → mapping from STL key to color 
             * None  → use default colors from `self.stl_colors`
         xmax, ymax, zmax : float, optional
             Initial clipping positions along each axis. If None, use the
@@ -726,7 +738,6 @@ def plot_stl_mask(self, stl_solid, cmap='viridis', bounding_box=True, show_grid=
         * A static domain bounding box can be added for reference.
         * Camera, background, and lighting are pre-configured for clarity
         """
-        from .materials import material_colors
         if stl_colors is None:
             stl_colors = self.stl_colors
 
@@ -773,30 +784,21 @@ def update_clip(val, axis="x"):
                     for i, key in enumerate(self.stl_solids):
                         surf = self.read_stl(key)
                         if type(stl_colors) is dict:
-                            if type(stl_colors[key]) is str:
-                                pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=dict(color=material_colors[stl_colors[key]]), name=key)
-                            else:
-                                pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=dict(color=stl_colors[key]), name=key)
+                            pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=dict(color=stl_colors[key]), name=key)
                         elif type(stl_colors) is list:
                             pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=dict(color=stl_colors[i]), name=key)
                         else:
                             pl.add_mesh(surf, color='white', opacity=stl_opacity, silhouette=True, name=key)
                 else: #add 1 selected stl solid
                     key = add_stl
                     surf = self.read_stl(key)
-                    if type(stl_colors[key]) is str:
-                        pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=dict(color=material_colors[stl_colors[key]]), name=key)
-                    else:
-                        pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=dict(color=stl_colors[key]), name=key)
+                    pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=dict(color=stl_colors[key]), name=key)
 
             elif type(add_stl) is list: #add selected list of stl solids
                 for i, key in enumerate(add_stl):
                     surf = self.read_stl(key)
                     if type(stl_colors[key]) is dict:
-                        if type(stl_colors) is str:
-                            pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=dict(color=material_colors[stl_colors[key]]), name=key)
-                        else:
-                            pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=dict(color=stl_colors[key]), name=key)
+                        pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=dict(color=stl_colors[key]), name=key)
                     elif type(stl_colors) is list:
                         pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=dict(color=stl_colors[i]), name=key)
                     else:
@@ -847,7 +849,7 @@ def update_clip(val, axis="x"):
             pl.show()
 
     def inspect(self, add_stl=None, stl_opacity=0.5, stl_colors=None,
-                anti_aliasing='ssaa', offscreen=False):
+                anti_aliasing='ssaa', smooth_shading=True, offscreen=False):
 
         '''3D plot using pyvista to visualize
         the structured grid and
@@ -862,7 +864,6 @@ def inspect(self, add_stl=None, stl_opacity=0.5, stl_colors=None,
         stl_colors: str or list of str, default 'white'
             Color of the stl surfaces
         '''
-        from .materials import material_colors
         if stl_colors is None:
             stl_colors = self.stl_colors
 
@@ -885,38 +886,29 @@ def clip(widget):
                     key = add_stl
                     surf = self.read_stl(key)
                     surf = surf.clip_box(widget.bounds, invert=False)
-                    if type(stl_colors[key]) is str:
-                        pl.add_mesh(surf, color=material_colors[self.stl_colors[key]], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key)
-                    else:
-                        pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key)
+                    pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=True, smooth_shading=smooth_shading, name=key)
 
                 elif type(add_stl) is list: #add selected list of stl solids
                     for i, key in enumerate(add_stl):
                         surf = self.read_stl(key)
                         surf = surf.clip_box(widget.bounds, invert=False)
                         if type(stl_colors) is dict:
-                            if type(stl_colors[key]) is str:
-                                pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key)
-                            else:
-                                pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key)
+                            pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=True, smooth_shading=smooth_shading, name=key)
                         elif type(stl_colors) is list:
-                            pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key)
+                            pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=True, smooth_shading=smooth_shading, name=key)
                         else:
-                            pl.add_mesh(surf, color='white', opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key)
+                            pl.add_mesh(surf, color='white', opacity=stl_opacity, silhouette=True, smooth_shading=smooth_shading, name=key)
 
             else: #add all stl solids
                 for i, key in enumerate(self.stl_solids):
                     surf = self.read_stl(key)
                     surf = surf.clip_box(widget.bounds, invert=False)
                     if type(stl_colors) is dict:
-                        if type(stl_colors[key]) is str:
-                            pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key)
-                        else:
-                            pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key)
+                        pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=True, smooth_shading=smooth_shading, name=key)
                     elif type(stl_colors) is list:
-                        pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key)
+                        pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=True, smooth_shading=smooth_shading, name=key)
                     else:
-                        pl.add_mesh(surf, color='white', opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key)
+                        pl.add_mesh(surf, color='white', opacity=stl_opacity, silhouette=True, smooth_shading=smooth_shading, name=key)
 
         _ = pl.add_box_widget(callback=clip, rotation_enabled=False)
 
@@ -1023,7 +1015,7 @@ def load_from_h5(self, filename):
         self.zmin, self.zmax = self.z[0], self.z[-1]
 
         # recommpute grid and L, iA, tL, itA
-        self.compute_grid()
+        self._compute_grid()
 
         # asign masks to grid.cell_data
         with h5py.File(filename, 'r') as hf:

wakis/logger.py

@@ -24,10 +24,13 @@ def __init__(self):
         self.solver = {}
         self.wakeSolver = {}
 
-    def save_logs(self):
+    def save_logs(self, results_folder=None):
         """
         Save all logs (grid, solver, wakeSolver) into log-file inside the results folder.
         """
+        if results_folder is not None:
+            self.wakeSolver["results_folder"] = results_folder
+
         logfile = os.path.join(self.wakeSolver["results_folder"], "wakis.log")
 
         # Write sections