refactor(assign_tags_to_cell): move into utilities module

kieran-mackle · Feb 10, 2024 · 193483a · 193483a
1 parent 77fb7ce
commit 193483a
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 36 deletions.
diff --git a/hypervehicle/components/component.py b/hypervehicle/components/component.py
@@ -10,14 +10,15 @@
 from typing import Callable, Union
 from hypervehicle.geometry import Vector3
 from gdtk.geom.sgrid import StructuredGrid
+from hypervehicle.utilities import assign_tags_to_cell
+from hypervehicle.utilities import parametricSurfce2stl, parametricSurfce2vtk
 from typing import Callable, Union, Optional
 from hypervehicle.geometry import (
     CurvedPatch,
     RotatedPatch,
     MirroredPatch,
     OffsetPatchFunction,
 )
-from hypervehicle.utilities import parametricSurfce2stl, parametricSurfce2vtk
 
 
 class AbstractComponent:
@@ -88,15 +89,6 @@ def analyse(self):
         pass
 
 
-def AssignTags2Celle(patch, length):
-    # Creates a tag vector for a given patch
-
-    tags_definition = {"FreeStream": 1, "Inlet": 2, "Outlet": 3, "Nozzle": 4}
-
-    tags = np.ones(length) * tags_definition[patch.tag]
-    return tags
-
-
 class Component(AbstractComponent):
     def __init__(
         self,
@@ -197,8 +189,8 @@ def mesh(self):
                 # Generate surfaces
                 self.surface()
 
-        # Combine all surface data
-        surface_data = np.concatenate([s[1].data for s in self.surfaces.items()])
+            # Combine all surface data
+            surface_data = np.concatenate([s[1].data for s in self.surfaces.items()])
 
             # Create nominal STL mesh
             self._mesh = mesh.Mesh(surface_data)
@@ -330,7 +322,7 @@ def wrapper(key: str, patch):
             )
 
             # Assign tags to the cells
-            tags = AssignTags2Celle(patch, len(cell_ids))
+            tags = assign_tags_to_cell(patch, len(cell_ids))
 
             return (key, vertices, cell_ids, tags)
 

diff --git a/hypervehicle/utilities.py b/hypervehicle/utilities.py
@@ -21,7 +21,8 @@ def create_cells(
     j_clustering_func: callable = None,
 ):
     """
-    Generates a list of vertices and a corrosponding list of index triplets, each pinting the vertices of a single cell
+    Generates a list of vertices and a corrosponding list of index triplets,
+    each pinting the vertices of a single cell
 
     Parameters
     ----------
@@ -51,7 +52,8 @@ def create_cells(
     Returns
     ----------
     vertices: numpy 2D array, each row contains (x,y,z) coordinates of one vertex
-    cell_ids: numpy 2D array, each row contains 3 indecies (row number in vertices) of a single cell
+    cell_ids: numpy 2D array, each row contains 3 indecies (row number in vertices)
+    of a single cell
 
     """
     # TODO - allow different ni and nj discretisation
@@ -165,6 +167,20 @@ def parametricSurfce2stl(
 
     Parameters
     ----------
+    parametric_surface : Any
+        The parametric surface object.
+
+    si : float, optional
+        The clustering in the i-direction. The default is 1.0.
+
+    sj : float, optional
+        The clustering in the j-direction. The default is 1.0.
+
+    triangles_per_edge : int
+        The resolution for the stl object.
+
+    mirror_y : bool, optional
+        Create mirror image about x-z plane. The default is False.
         parametric_surface : Any
             The parametric surface object.
         si : float, optional
@@ -217,20 +233,24 @@ def parametricSurfce2vtk(
 
     Parameters
     ----------
-        parametric_surface : Any
-            The parametric surface object.
-        si : float, optional
-            The clustering in the i-direction. The default is 1.0.
-        sj : float, optional
-            The clustering in the j-direction. The default is 1.0.
-        triangles_per_edge : int
-            The resolution for the stl object.
-        mirror_y : bool, optional
-            Create mirror image about x-z plane. The default is False.
+    parametric_surface : Any
+        The parametric surface object.
+
+    si : float, optional
+        The clustering in the i-direction. The default is 1.0.
+
+    sj : float, optional
+        The clustering in the j-direction. The default is 1.0.
+
+    triangles_per_edge : int
+        The resolution for the stl object.
+
+    mirror_y : bool, optional
+        Create mirror image about x-z plane. The default is False.
 
     Returns
     ----------
-    cells
+    tuple[vertices, cell_ids]
     """
     # Generate the mesh vertices and cell index list
     vertices, cell_ids = create_cells(
@@ -255,7 +275,7 @@ def parametricSurfce2vtk(
 
 
 def assess_inertial_properties(vehicle, component_densities: Dict[str, float]):
-    """
+    """Return the inertial properties of a vehicle.
 
     Parameters
     ----------
@@ -407,7 +427,6 @@ def dvdp(
         write_nominal_stl : bool, optional
             A boolean flag to write the nominal geometry STL(s) to file. The
             default is True.
-
         nominal_stl_prefix : str, optional
             The prefix to append when writing STL files for the nominal geometry.
             If None, no prefix will be used. The default is None.
@@ -763,7 +782,7 @@ def append_sensitivities_to_tri(
 
     Examples
     ---------
-    >>> dp_files = ['wing_0_body_width_sensitivity.csv',
+     >>> dp_files = ['wing_0_body_width_sensitivity.csv',
                     'wing_1_body_width_sensitivity.csv']
     """
     # Check outdir
@@ -1032,3 +1051,11 @@ def print_banner():
     tprint("Hypervehicle", "tarty4")
     p = art("airplane2")
     print(f" {p}               {p}               {p}               {p}")
+
+
+def assign_tags_to_cell(patch, length):
+    """Assign tags to cells."""
+    # Creates a tag vector for a given patch
+    tags_definition = {"FreeStream": 1, "Inlet": 2, "Outlet": 3, "Nozzle": 4}
+    tags = np.ones(length) * tags_definition[patch.tag]
+    return tags
diff --git a/hypervehicle/vehicle.py b/hypervehicle/vehicle.py
@@ -300,7 +300,12 @@ def transform(self, transformations: List[Tuple[str, Any]]) -> None:
             component.surfaces = None
             component.mesh = None
 
-    def to_file(self, prefix: str = None, file_type: str = "stl") -> None:
+    def to_file(
+        self,
+        prefix: str = None,
+        merge: Union[bool, List[str]] = False,
+        file_type: str = "stl",
+    ) -> None:
         """Writes the vehicle components to output file. If analysis results are
         present, they will also be written to file, either as CSV, or using
         the Numpy tofile method.
@@ -315,16 +320,18 @@ def to_file(self, prefix: str = None, file_type: str = "stl") -> None:
             component name tag is available, the Vehicle name will be used.
             The default is None.
 
-        file_type: str
-            Defines the output file format to be written to. Can be stl or
-            vtk. Default file_type is 'stl'
-
         merge : [bool, list[str]], optional
             Merge components of the vehicle into a single STL file. The merge
             argument can either be a boolean (with True indicating to merge all
             components of the vehicle), or a list of the component names to
             merge. This functionality depends on PyMesh. The default is False.
 
+        file_type: str
+            Defines the output file format to be written to. Can be stl or
+            vtk. Default file_type is 'stl'
+
+
+
         See Also
         --------
         utilities.merge_stls
@@ -432,7 +439,7 @@ def to_file(self, prefix: str = None, file_type: str = "stl") -> None:
                 end="\n",
             )
 
-    def to_stl(self, prefix: str = None) -> None:
+    def to_stl(self, prefix: str = None, merge: Union[bool, List[str]] = False) -> None:
         """Writes the vehicle components to STL. If analysis results are
         present, they will also be written to file, either as CSV, or using
         the Numpy tofile method.
@@ -445,9 +452,15 @@ def to_stl(self, prefix: str = None) -> None:
             provided will take precedence. If no prefix is specified, and no
             component name tag is available, the Vehicle name will be used.
             The default is None.
+
+        merge : [bool, list[str]], optional
+            Merge components of the vehicle into a single STL file. The merge
+            argument can either be a boolean (with True indicating to merge all
+            components of the vehicle), or a list of the component names to
+            merge. This functionality depends on PyMesh. The default is False.
         """
 
-        self.to_file(prefix, file_type="stl")
+        self.to_file(prefix, merge=merge, file_type="stl")
 
     def to_vtk(self, prefix: str = None) -> None:
         """Writes the vehicle components to VTK. If analysis results are