t3d-surface-hole/bld_quasistatic_simulation/fem_tetrahedral.py

import bpy
import bmesh
import os
import sys
import math
import numpy as np
import shutil
from mathutils import Vector
from ..project_config import PROJECT_CONFIG
import datetime
import json


def boundary_condition_mesh_callback(scene, context):
  items = [("","empty","")]
  object_list = bpy.context.scene.objects
  for i in object_list:
      if i.type == 'MESH':
        items.append((i.name, i.name, i.name))
  return items

def custom_normalize_array(arr, min_value=0, max_value=1, threshold=0.0007):
    small_values = [x for x in arr if x <= threshold]
    
    if small_values:
        min_val_small = min(small_values)
        max_val_small = max(small_values)
        normalized_small = [
            min_value + (x - min_val_small) / (max_val_small - min_val_small) * (max_value - min_value) for x in small_values
        ]
    else:
        normalized_small = []

    large_values = [x for x in arr if x > threshold]
    if large_values:
        min_val_large = min(large_values)
        max_val_large = max(large_values)
        normalized_large = [
            min_value + ((math.log(1 + x - threshold) / math.log(1 + max_val_large - threshold)) * (max_value - min_value)) for x in large_values
        ]
    else:
        normalized_large = []

    normalized_arr = normalized_small + normalized_large
    return normalized_arr

def get_neighbors(mesh, num_vertices):
    neighbors = [set() for _ in range(num_vertices)]
    for polygon in mesh.polygons:
        for i in polygon.vertices:
            if i < num_vertices:
                for j in polygon.vertices:
                    if i != j and j < num_vertices:
                        neighbors[i].add(j)
    return neighbors


# 进行颜色的平均

def average_colors(normalized_stress, neighbors, num_vertices, iterations=1, weight=0.1):
    colors = np.array(normalized_stress[:num_vertices])
    for _ in range(iterations):
        for i, color in enumerate(colors):
            neighbor_colors = np.array([colors[j] for j in neighbors[i]])
            average_color = np.mean(neighbor_colors)
            colors[i] = (1 - weight) * color + weight * average_color
    return colors

class FEMTetrahedralOperator(bpy.types.Operator):
    bl_idname: str = "designauto.quasistatic_simulation_fem_tetrahedral"
    bl_label: str = "FEM仿真"
    bl_options = {"REGISTER", "UNDO"}
    
    # YM: bpy.props.FloatProperty(name="Young's modulus", default=1e5)
    # PR: bpy.props.FloatProperty(name="Poisson's ratio", default=0.3)
    # density: bpy.props.FloatProperty(name="density", default=1e3)
    
    hole_height: bpy.props.FloatProperty(name="孔洞深度", default=0.1)
    max_hole_radius: bpy.props.FloatProperty(name="最大半径", default=0.4)
    min_hole_radius: bpy.props.FloatProperty(name="最小半径", default=0.2)
    
    # nbc1: bpy.props.EnumProperty(name="NBC 1", items=boundary_condition_mesh_callback)
    # nbcv1: bpy.props.FloatProperty(name="NBC va1 1", default=0)
    
    # nbc2: bpy.props.EnumProperty(name="NBC 2", items=boundary_condition_mesh_callback)
    # nbcv2: bpy.props.FloatProperty(name="NBC val 2", default=0)
    
    # nbc3: bpy.props.EnumProperty(name="NBC 3", items=boundary_condition_mesh_callback)
    # nbcv3: bpy.props.FloatProperty(name="NBC val 3", default=0)
    
    # dbc1: bpy.props.EnumProperty(name="DBC 1", items=boundary_condition_mesh_callback)
    # dbc2: bpy.props.EnumProperty(name="DBC 2", items=boundary_condition_mesh_callback)
    # dbc3: bpy.props.EnumProperty(name="DBC 3", items=boundary_condition_mesh_callback)
    
    
    def execute(self, context):
        working_object = context.active_object
        bpyscene = bpy.context.scene
        working_name = 'fem'
        workplace = PROJECT_CONFIG.workplace_dir_path
        executable_path = PROJECT_CONFIG.executable_dir_path
        fem_mesh_path = os.path.join(workplace, "parameters", "femmodel.obj")
        import dapy_qss_fem_tetrahedral
        bpy.ops.export_scene.obj(filepath=fem_mesh_path, use_selection=True, axis_forward='Y', axis_up='Z', use_edges=False, use_animation=False, use_materials=False, use_uvs=False, use_normals=False,
                                 use_mesh_modifiers=False, use_nurbs=False, use_smooth_groups=False, use_vertex_groups=False, use_blen_objects=False, use_smooth_groups_bitflags=False)
        os.chdir(executable_path)
        surface_mesh, mat_deformed_coordinates, \
            vtx_displacement, vtx_stress = \
                dapy_qss_fem_tetrahedral.QuasistaticSimulationByFEMTetrahedral(self.hole_height,self.max_hole_radius,self.min_hole_radius)
       
        

        
        working_object.hide_set(True)
        new_mesh = bpy.data.meshes.new("mesh-" + working_object.name + "-hole")
        new_mesh.from_pydata(surface_mesh.mat_coordinates.tolist(), [], surface_mesh.mat_faces.tolist())
        new_mesh.update()
        new_object = bpy.data.objects.new(working_object.name + "-hole", new_mesh)
        #new_stress_mesh = bpy.data.meshes.new("mesh-" + working_name + "-stress")
        new_stress_mesh = working_object.data
        new_stress_mesh.update()
        new_stress_object = bpy.data.objects.new(working_name + "-stress", new_stress_mesh)
        
        normalized_stress = custom_normalize_array(vtx_stress, min_value=0.2, max_value=1, threshold=0.0007)
        if not new_stress_mesh.vertex_colors:
            new_stress_mesh.vertex_colors.new()
        num_vertices = len(new_stress_mesh.vertices)
        neighbors = get_neighbors(new_stress_mesh, num_vertices)
        normalized_stress = average_colors(normalized_stress, neighbors, num_vertices, iterations=10, weight=0.2)
        for polygon in new_stress_mesh.polygons:
            for i, index in enumerate(polygon.vertices):
                loop_index = polygon.loop_indices[i]
                new_stress_mesh.vertex_colors.active.data[loop_index].color = (normalized_stress[index], 0, 0, 1)
        found = False
        for c in bpy.data.collections:
            if c.name == 'designauto':
                found = True
                new_collection = c
        if not found:
            new_collection = bpy.data.collections.new('designauto')
            bpy.context.scene.collection.children.link(new_collection)
            
        new_collection.objects.link(new_object)
        new_collection.objects.link(new_stress_object)
        working_object.select_set(True)
        return {'FINISHED'}
        
    
    def invoke(self, context, event):
        return context.window_manager.invoke_props_dialog(self)
        # if context.active_object.type == 'MESH':
            
        #     return context.window_manager.invoke_props_dialog(self)
        # else:
        #     self.report({'ERROR'}, "Selected object is not a mesh!")
        #     return {'CANCELLED'}