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

tests/test_pre_processing.py

@@ -332,6 +332,86 @@ def test_mesh_model_with_variable_solid_thickness(tmpdir):
         f"VTU mesh has diameter {diameter_at_outlet} at outlet, expected {expected_diameter_at_outlet}"
 
 
+def test_mesh_model_with_thickness_to_entity_id_mapping(tmpdir):
+    """
+    Test meshing procedure with thickness-to-entity-ID mapping by verifying the count of entity IDs per range.
+    """
+    # Define test data paths
+    original_model_path = Path("tests/test_data/cylinder/cylinder.vtp")
+    sphere_file_path = original_model_path.with_name("stored_" + original_model_path.stem +
+                                                     "_variable_solid_thickness_distance_to_sphere_solid_thickness.vtp")
+
+    # Copy the original model to tmpdir
+    model_path = copy_original_model_to_tmpdir(original_model_path, tmpdir)
+
+    # Copy the sphere file to tmpdir
+    copied_sphere_file_path = model_path.with_name(model_path.stem + "_distance_to_sphere_solid_thickness.vtp")
+    copied_sphere_file_path.write_text(sphere_file_path.read_text())
+
+    # Define expected values
+    expected_num_points = 5687
+    expected_num_cells = 31335
+    expected_entity_id_counts = {
+        0: 21399,   # Entity ID 0 (fluid volume ID) should have 21399 cells
+        1: 7923,    # Entity ID 1 (solid volume ID) should have 7923 cells
+        2: 120,     # Entity ID 2 (inlet ID) should have 120 cells
+        3: 102,     # Entity ID 3 (outlet ID) should have 102 cells
+        11: 136,    # Entity ID 11 (solid sidewall ID) should have 136 cells
+        22: 1656,   # Entity ID 22 (interface FSI ID) should have 1656 cells
+        33: 1656,   # Entity ID 33 (solid outer wall ID) should have 1656 cells
+        100: 760,   # Entity ID 100 should have 760 cells
+        200: 932,   # Entity ID 200 should have 932 cells
+        300: 909,   # Entity ID 300 should have 909 cells
+        400: 1068,  # Entity ID 400 should have 1068 cells
+    }
+
+    # Get default input parameters
+    common_input = read_command_line(str(model_path))
+    common_input.update(
+        dict(
+            solid_thickness="variable",
+            solid_thickness_parameters=[0, 0.1, 0.2, 0.4],
+            meshing_method="diameter",
+            smoothing_method="no_smooth",
+            refine_region=False,
+            coarsening_factor=1.3,
+            visualize=False,
+            compress_mesh=False,
+            outlet_flow_extension_length=5,
+            inlet_flow_extension_length=5,
+            thickness_to_entity_id_mapping={
+                (0.21, 0.25): 100,  # Thickness range -> Entity ID
+                (0.25, 0.3): 200,
+                (0.3, 0.35): 300,
+                (0.35, 0.4): 400,
+            },
+        )
+    )
+
+    # Run pre processing and assert mesh sizes
+    model_path, mesh_vtu, mesh_hdf5 = run_pre_processing_with_common_input(model_path, common_input)
+    assert_mesh_sizes(mesh_vtu, mesh_hdf5, expected_num_points, expected_num_cells)
+
+    # Verify entity ID counts
+    cell_entity_ids = mesh_vtu.GetCellData().GetArray("CellEntityIds")
+
+    # Count occurrences of each entity ID
+    entity_id_counts = {}
+    for cell_id in range(mesh_vtu.GetNumberOfCells()):
+        entity_id = cell_entity_ids.GetValue(cell_id)
+        entity_id_counts[entity_id] = entity_id_counts.get(entity_id, 0) + 1
+
+    # Verify the counts match expectations
+    for entity_id, expected_count in expected_entity_id_counts.items():
+        actual_count = entity_id_counts.get(entity_id, 0)
+        assert actual_count == expected_count, \
+            f"Entity ID {entity_id} has {actual_count} cells, expected {expected_count}"
+
+    # Check for unexpected entity IDs
+    unexpected_ids = set(entity_id_counts) - set(expected_entity_id_counts)
+    assert not unexpected_ids, f"Unexpected entity IDs found: {unexpected_ids}"
+
+
 def test_xdmf_mesh_format(tmpdir):
     """
     Test meshing procedure with generated mesh in XDMF format.