Add support for thickness-based entity ID mapping in preprocessing

src/vasp/automatedPreprocessing/automated_preprocessing.py

@@ -27,7 +27,7 @@
 
 from vasp.automatedPreprocessing.preprocessing_common import generate_mesh, distance_to_spheres_solid_thickness, \
     dist_sphere_spheres, convert_xml_mesh_to_hdf5, convert_vtu_mesh_to_xdmf, edge_length_evaluator, \
-    check_flatten_boundary
+    check_flatten_boundary, map_thickness_to_mesh, update_entity_ids_by_thickness
 from vasp.simulations.simulation_common import load_mesh_and_data
 
 
@@ -54,7 +54,7 @@ def run_pre_processing(input_model, verbose_print, smoothing_method, smoothing_f
                        remove_all, solid_thickness, solid_thickness_parameters, mesh_format, flow_rate_factor,
                        solid_side_wall_id, interface_fsi_id, solid_outer_wall_id, fluid_volume_id, solid_volume_id,
                        mesh_generation_retries, no_solid, extract_branch, branch_group_ids, branch_ids_offset,
-                       distance_method):
+                       distance_method, thickness_to_entity_id_mapping):
     """
     Automatically generate mesh of surface model in .vtu and .xml format, including prescribed
     flow rates at inlet and outlet based on flow network model.
@@ -100,6 +100,7 @@ def run_pre_processing(input_model, verbose_print, smoothing_method, smoothing_f
         branch_group_ids (list): Specify group IDs to extract for the branch.
         branch_ids_offset (int): Set offset for marking solid mesh IDs when extracting a branch.
         distance_method (str): Change between 'euclidean' and 'geodesic' distance measure
+        thickness_to_entity_id_mapping (dict): Mapping of thickness ranges to entity IDs
     """
     # Get paths
     input_model_path = Path(input_model)
@@ -539,6 +540,15 @@ def try_generate_mesh(distance_to_sphere, number_of_sublayers_fluid, number_of_s
         assert mesh.GetNumberOfPoints() > 0, "No points in mesh, try to remesh."
         assert remeshed_surface.GetNumberOfPoints() > 0, "No points in surface mesh, try to remesh."
 
+        if thickness_to_entity_id_mapping and solid_thickness in ("variable", "painted"):
+            # Map thickness values to input mesh
+            print("--- Mapping thickness values to mesh...")
+            mesh = map_thickness_to_mesh(mesh, distance_to_sphere)
+
+            # Update entity IDs based on thickness values
+            print("--- Updating entity IDs based on thickness values...")
+            mesh = update_entity_ids_by_thickness(mesh, thickness_to_entity_id_mapping, solid_volume_id)
+
         if mesh_format in ("xml", "hdf5"):
             write_mesh(compress_mesh, str(file_name_surface_name), str(file_name_vtu_mesh), str(file_name_xml_mesh),
                        mesh, remeshed_surface)
@@ -892,6 +902,13 @@ def read_command_line(input_path=None):
     parser.add_argument('-bo', '--branch-ids-offset', default=1000, type=int,
                         help="Set the offset for marking solid mesh IDs when extracting a branch.")
 
+    parser.add_argument("-tm", "--thickness-to-entity-id-mapping", type=float, nargs='+', default=[],
+                        help="Mapping of thickness ranges to entity IDs in the format: min1 max1 id1 min2 max2 id2 ..."
+                             "Default is no mapping."
+                             "For example, --thickness-to-entity-id-mapping 0.2 0.3 100 0.3 0.4 200 will map the "
+                             "cells with thickness between 0.2 and 0.3 to entity ID 100, and the cells with thickness "
+                             "between 0.3 and 0.4 to entity ID 200.")
+
     # Parse path to get default values
     if required:
         args = parser.parse_args()
@@ -911,6 +928,13 @@ def read_command_line(input_path=None):
     elif args.solid_thickness == "painted" and len(args.solid_thickness_parameters) != 3:
         raise ValueError("ERROR: solid thickness parameters for 'painted' thickness should be three floats.")
 
+    # Parse the thickness to entity ID mapping
+    args_mapping = args.thickness_to_entity_id_mapping
+    if len(args_mapping) % 3 != 0:
+        raise ValueError("Thickness to entity ID mapping must be a multiple of 3: min, max, and ID.")
+    thickness_to_entity_id_mapping = \
+        {(args_mapping[i], args_mapping[i + 1]): int(args_mapping[i + 2]) for i in range(0, len(args_mapping), 3)}
+
     if args.verbosity:
         print()
         print("--- VERBOSE MODE ACTIVATED ---")
@@ -944,7 +968,7 @@ def verbose_print(*args):
                 solid_volume_id=args.solid_volume_id, mesh_generation_retries=args.mesh_generation_retries,
                 no_solid=args.no_solid, extract_branch=args.extract_branch,
                 branch_group_ids=args.branch_group_ids, branch_ids_offset=args.branch_ids_offset,
-                distance_method=args.distance_method)
+                distance_method=args.distance_method, thickness_to_entity_id_mapping=thickness_to_entity_id_mapping)
 
 
 def main_meshing():

src/vasp/automatedPreprocessing/preprocessing_common.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 from typing import Union
 
+import vtk
 import h5py
 import numpy as np
 import meshio
@@ -456,3 +457,112 @@ def check_flatten_boundary(num_inlets_outlets: int, mesh_path: Union[str, Path],
         vectorData.close()
         flat_in_out_mesh_path.unlink()
         print("No changes made to the mesh file")
+
+
+def map_thickness_to_mesh(mesh: vtkPolyData, surface: vtkPolyData) -> vtkPolyData:
+    """
+    Map the thickness values from a surface to the points in a mesh.
+
+    Args:
+        mesh (vtkPolyData): The input mesh.
+        surface (vtkPolyData): The surface with a "Thickness" array.
+
+    Returns:
+        vtkPolyData: The updated mesh with a mapped thickness array.
+    """
+    # Ensure the surface has the thickness array
+    thickness_array = surface.GetPointData().GetArray(distanceToSpheresArrayNameSolid)
+
+    # Create a new thickness array for the mesh
+    new_thickness_array = vtk.vtkFloatArray()
+    new_thickness_array.SetName(distanceToSpheresArrayNameSolid)
+    new_thickness_array.SetNumberOfComponents(1)
+    new_thickness_array.SetNumberOfTuples(mesh.GetNumberOfPoints())
+
+    # Setup a point locator for the surface
+    point_locator = vtk.vtkPointLocator()
+    point_locator.SetDataSet(surface)
+    point_locator.BuildLocator()
+
+    # Map the thickness values
+    for i in range(mesh.GetNumberOfPoints()):
+        point = mesh.GetPoint(i)
+        closest_point_id = point_locator.FindClosestPoint(point)
+        thickness_value = thickness_array.GetTuple1(closest_point_id)
+        new_thickness_array.SetTuple1(i, thickness_value)
+
+    # Add the new thickness array to the mesh
+    mesh.GetPointData().AddArray(new_thickness_array)
+    return mesh
+
+
+def update_entity_ids_by_thickness(mesh: vtkPolyData, entity_id_mapping: dict, volume_entity_id: int) -> vtkPolyData:
+    """
+    Update the entity IDs of cells in the mesh based on average thickness of their points.
+    Only update cells with the same entity ID as `volume_entity_id`.
+    Falls back to the original entity ID if no matching range is found.
+
+    Args:
+        mesh (vtkPolyData): The input mesh.
+        entity_id_mapping (dict): Mapping of thickness ranges to entity IDs.
+                                  Keys should be tuples defining ranges (min_thickness, max_thickness),
+                                  values should be the entity IDs to assign.
+        volume_entity_id (int): The cell entity ID that should be updated.
+
+    Returns:
+        vtkPolyData: Mesh with updated cell entity IDs.
+    """
+    # Fetch the thickness array
+    thickness_array = mesh.GetPointData().GetArray(distanceToSpheresArrayNameSolid)
+
+    # Fetch the original cell entity IDs
+    original_entity_ids_array = mesh.GetCellData().GetArray("CellEntityIds")
+
+    # Create a new array for updated cell entity IDs
+    updated_entity_ids_array = vtk.vtkIntArray()
+    updated_entity_ids_array.SetName("CellEntityIds")
+    updated_entity_ids_array.SetNumberOfTuples(mesh.GetNumberOfCells())
+
+    # Sort the entity_id_mapping by range to ensure proper assignment
+    sorted_mapping = sorted(entity_id_mapping.items(), key=lambda x: x[0])
+
+    # Iterate through each cell to compute average thickness
+    cell_point_ids = vtk.vtkIdList()
+    for cell_id in range(mesh.GetNumberOfCells()):
+        # Get the original entity ID for the current cell
+        original_entity_id = original_entity_ids_array.GetValue(cell_id)
+
+        # Skip cells that do not match the volume entity ID
+        if original_entity_id != volume_entity_id:
+            updated_entity_ids_array.SetValue(cell_id, original_entity_id)
+            continue
+
+        # Get the points of the current cell
+        mesh.GetCellPoints(cell_id, cell_point_ids)
+        point_ids = [cell_point_ids.GetId(i) for i in range(cell_point_ids.GetNumberOfIds())]
+
+        # Fetch thickness values for these points
+        thickness_values = [thickness_array.GetTuple1(pid) for pid in point_ids]
+
+        # Compute the average thickness for the cell
+        avg_thickness = sum(thickness_values) / len(thickness_values) if thickness_values else 0
+
+        # Determine the entity ID based on the average thickness and the mapping
+        new_entity_id = None
+        for (min_thickness, max_thickness), entity_id in sorted_mapping:
+            if min_thickness <= avg_thickness <= max_thickness:
+                new_entity_id = entity_id
+                break
+
+        # Fall back to the original entity ID if no match is found
+        if new_entity_id is None:
+            new_entity_id = original_entity_id
+
+        # Assign the determined or original entity ID to the updated array
+        updated_entity_ids_array.SetValue(cell_id, new_entity_id)
+
+    # Replace the original entity IDs array with the updated array
+    mesh.GetCellData().RemoveArray("CellEntityIds")
+    mesh.GetCellData().AddArray(updated_entity_ids_array)
+
+    return mesh