## MVP 1 FILE! from .step_feature_detection import ( load_step_model, load_step_unit, detect_faces, detect_circular_holes, detect_non_circular_holes, detect_fillets, mode_thickness, detect_all_edges, detect_all_faces, ) import os, json import pandas as pd import gmsh from OCP.Bnd import Bnd_Box from OCP.BRepBndLib import BRepBndLib import joblib from OCP.TopExp import TopExp_Explorer from OCP.TopAbs import TopAbs_FACE from OCP.TopoDS import TopoDS import torch from torch import nn import numpy as np import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D # noqa: F401 import plotly.graph_objects as go # --------------------------------------------------------------------- # Matplotlib static PNG plotter (unchanged) # --------------------------------------------------------------------- def plot_mesh( nodes, elements, title, highlight_faces=None, face_to_nodes=None, save_path="mesh.png", ): """ Static 3D line-plot of the surface mesh for PNG export. nodes: (N,3) elements: list[(etype, conn)], conn is array of 0-based node indices """ fig = plt.figure(figsize=(7, 7)) ax = fig.add_subplot(111, projection="3d") ax.view_init(elev=20, azim=45) # normal mesh (black) for (etype, conn) in elements: for e in conn: pts = nodes[e][:, :3] pts = np.vstack([pts, pts[0]]) ax.plot(pts[:, 0], pts[:, 1], pts[:, 2], "k-", linewidth=0.4) # optional red overlay for selected faces if highlight_faces and face_to_nodes: for sf in highlight_faces: if sf not in face_to_nodes: continue face_nodes = face_to_nodes[sf] for fn in face_nodes: pts = nodes[fn][:, :3] pts = np.vstack([pts, pts[0]]) ax.plot(pts[:, 0], pts[:, 1], pts[:, 2], "r-", linewidth=1.2) ax.set_title(title) ax.set_axis_off() plt.tight_layout() plt.savefig(save_path, dpi=300) plt.close() # --------------------------------------------------------------------- # Face list helper for GUI (unchanged) # --------------------------------------------------------------------- def extract_faces_for_gui(step_path): """ Load the STEP file, detect faces, and return a list of label/metadata dicts that the GUI can use. """ shape = load_step_model(step_path) faces = detect_all_faces(shape) face_list = [] for f in faces: center = f["center"] normal = f["normal"] area = f.get("area", 0.0) label = f"Face {f['face_id']} | normal={np.round(normal,3)} | area={area:.2f}" face_list.append( { "label": label, "face_id": f["face_id"], "center": list(center), "normal": list(normal), "area": float(area), } ) return face_list # --------------------------------------------------------------------- # Helpers for Plotly 3D mesh # --------------------------------------------------------------------- def mesh_edges_from_triangles(triangles): """ Given (M,3) triangle connectivity, return unique edges as (E,2). """ edges = set() for tri in triangles: a, b, c = tri edges.add(tuple(sorted((a, b)))) edges.add(tuple(sorted((b, c)))) edges.add(tuple(sorted((c, a)))) return np.array(list(edges), dtype=int) def make_plotly_mesh( nodes, triangles, title="Refined Mesh", view_mode="wireframe", refined_tri_mask=None, vertex_refinement_values=None, ): """ Build an interactive Plotly figure. Parameters ---------- nodes : (N,3) ndarray triangles : (M,3) ndarray of 0-based node indices view_mode : {"wireframe", "highlight", "heatmap"} - "wireframe": light-blue surface + black edges - "highlight": base surface + red overlay on refined triangles - "heatmap": surface colored by vertex_refinement_values refined_tri_mask : (M,) bool array True for triangles in locally-refined region (used in "highlight"). vertex_refinement_values : (N,) array Per-node refinement intensity used for "heatmap". Higher values → more refined → hotter colors. """ nodes = np.asarray(nodes) triangles = np.asarray(triangles, dtype=int) edges = mesh_edges_from_triangles(triangles) n_nodes = nodes.shape[0] if refined_tri_mask is None: refined_tri_mask = np.zeros(triangles.shape[0], dtype=bool) if vertex_refinement_values is None: vertex_refinement_values = np.zeros(n_nodes, dtype=float) fig = go.Figure() # ---------------------------------------- # Base surface for all modes # ---------------------------------------- if view_mode == "heatmap": # Heatmap using per-vertex refinement values fig.add_trace( go.Mesh3d( x=nodes[:, 0], y=nodes[:, 1], z=nodes[:, 2], i=triangles[:, 0], j=triangles[:, 1], k=triangles[:, 2], intensity=vertex_refinement_values, colorscale="Turbo", showscale=True, flatshading=True, opacity=0.9, showlegend=False, name="Local mesh size heatmap", ) ) else: # Simple uniform surface color fig.add_trace( go.Mesh3d( x=nodes[:, 0], y=nodes[:, 1], z=nodes[:, 2], i=triangles[:, 0], j=triangles[:, 1], k=triangles[:, 2], color="lightblue", flatshading=True, lighting=dict( ambient=0.6, diffuse=0.8, specular=0.2, roughness=0.9 ), lightposition=dict(x=0, y=0, z=1), opacity=0.45 if view_mode == "wireframe" else 0.55, showlegend=False, showscale=False, name="Surface", ) ) # ---------------------------------------- # Highlight overlay (local refinement in red) # ---------------------------------------- if view_mode == "highlight" and refined_tri_mask.any(): ref_tris = triangles[refined_tri_mask] fig.add_trace( go.Mesh3d( x=nodes[:, 0], y=nodes[:, 1], z=nodes[:, 2], i=ref_tris[:, 0], j=ref_tris[:, 1], k=ref_tris[:, 2], color="red", opacity=0.9, flatshading=True, showscale=False, name="Locally refined", ) ) # ---------------------------------------- # Wireframe edges for all modes # ---------------------------------------- edge_width = 2 if view_mode != "heatmap" else 1 edge_color = "black" if view_mode != "heatmap" else "rgba(0,0,0,0.4)" for e0, e1 in edges: fig.add_trace( go.Scatter3d( x=[nodes[e0, 0], nodes[e1, 0]], y=[nodes[e0, 1], nodes[e1, 1]], z=[nodes[e0, 2], nodes[e1, 2]], mode="lines", line=dict(color=edge_color, width=edge_width), showlegend=False, ) ) fig.update_layout( title=title, scene=dict(aspectmode="data"), margin=dict(l=0, r=0, t=30, b=0), ) return fig # --------------------------------------------------------------------- # Main pipeline # --------------------------------------------------------------------- def get_result( step_path, thickness, load_face, load_direction, load_scale, selected_face_label=None, faces_json=None, view_mode="wireframe", # NEW optional flag, default keeps current behavior ): # ----------------------------------------------------------------- # STEP → JSON features # ----------------------------------------------------------------- def step_to_json(step_path, thickness, part_id=None): shape = load_step_model(step_path) part_id = part_id or os.path.splitext(os.path.basename(step_path))[0] circular_holes = detect_circular_holes(shape) fillets = detect_fillets(shape) faces = detect_all_faces(shape) print(f"Detected {len(faces)} faces") features = { "part_id": part_id, "thickness": thickness, "circular_holes": circular_holes, "fillets": fillets, "faces": faces, } out_path = "output/example.json" with open(out_path, "w") as f: json.dump(features, f, indent=2) step_to_json(step_path, thickness) # ----------------------------------------------------------------- # Load geometry into gmsh # ----------------------------------------------------------------- gmsh.open(step_path) gmsh.model.occ.synchronize() surfaces = gmsh.model.getEntities(dim=2) surface_info = [] for dim, tag in surfaces: mass_props = gmsh.model.occ.getMass(dim, tag) com = gmsh.model.occ.getCenterOfMass(dim, tag) surface_info.append({"tag": tag, "center": com, "area": mass_props}) with open("output/example.json") as f: features = json.load(f) faces = features["faces"] num_fillet = len(features["fillets"]) num_hole = len(features["circular_holes"]) num_faces = len(features["faces"]) num_plane = num_faces - num_fillet - num_hole smallest_fillet_r = min( (f["radius"] for f in features["fillets"]), default=0.0 ) smallest_hole_r = min( (h["radius"] for h in features["circular_holes"]), default=0.0 ) print( f"num_fillet: {num_fillet}, num_hole: {num_hole}, num_plane: {num_plane}" ) print( f"smallest_fillet_r: {smallest_fillet_r}, smallest_hole_r: {smallest_hole_r}" ) # ----------------------------------------------------------------- # Bounding box → max dimension # ----------------------------------------------------------------- def detect_maxsize(shape): bbox = Bnd_Box() BRepBndLib.Add_s(shape, bbox) xmin, ymin, zmin, xmax, ymax, zmax = bbox.Get() size_x = xmax - xmin size_y = ymax - ymin size_z = zmax - zmin maxsize = max(size_x, size_y, size_z) return maxsize, (size_x, size_y, size_z) shape = load_step_model(step_path) maxsize, (sx, sy, sz) = detect_maxsize(shape) print("maxsize:", maxsize) # ----------------------------------------------------------------- # Heuristic face roles (bottom/left/front/back) # ----------------------------------------------------------------- def find_the_faces(step_path): shape = load_step_model(step_path) faces_local = detect_all_faces(shape) z_vals = [f["center"][2] for f in faces_local] x_vals = [f["center"][0] for f in faces_local] y_vals = [f["center"][1] for f in faces_local] idx_bottom = z_vals.index(min(z_vals)) idx_left = x_vals.index(min(x_vals)) idx_front = y_vals.index(min(y_vals)) idx_back = y_vals.index(max(y_vals)) return { "bottom": faces_local[idx_bottom], "left": faces_local[idx_left], "front": faces_local[idx_front], "back": faces_local[idx_back], "bottom_idx": idx_bottom, "left_idx": idx_left, "front_idx": idx_front, "back_idx": idx_back, } face_dict = find_the_faces(step_path) part_row = { "num_fillet": num_fillet, "num_hole": num_hole, "num_plane": num_plane, "smallest_fillet_r": smallest_fillet_r, "smallest_hole_r": smallest_hole_r, "maxsize": maxsize, "thickness": thickness, "load_face": load_face, "load_direction": load_direction, "load_scale": load_scale, } df_part = pd.DataFrame([part_row]) # classify face types face_type_dict = {} for hole in features["circular_holes"]: for f in hole["faces"]: face_type_dict[f["face_id"]] = "CIRC_HOLE" for fillet in features["fillets"]: for f in fillet["faces"]: face_type_dict[f["face_id"]] = "FILLET" for f in features["faces"]: if f["face_id"] not in face_type_dict: face_type_dict[f["face_id"]] = "FACE" # ----------------------------------------------------------------- # Choose load/fixed faces (GUI-aware) # ----------------------------------------------------------------- load_face_data = None fix_face_data = None if selected_face_label is not None and faces_json is not None and faces_json != "": faces_from_gui = json.loads(faces_json) label_to_face = {f["label"]: f for f in faces_from_gui} if selected_face_label not in label_to_face: raise ValueError( f"Selected load face '{selected_face_label}' not found in faces_json." ) load_face_data = label_to_face[selected_face_label] load_center = np.array(load_face_data["center"], dtype=float) all_faces_gui = list(label_to_face.values()) distances = [ np.linalg.norm(np.array(f["center"], dtype=float) - load_center) for f in all_faces_gui ] fix_face_data = all_faces_gui[int(np.argmax(distances))] else: if load_face == "bend_bottom": load_face_data = face_dict["bottom"] fix_face_data = face_dict["left"] elif load_face == "tension": load_face_data = face_dict["back"] fix_face_data = face_dict["front"] else: raise ValueError( f"No GUI selection and unsupported load_face value '{load_face}'." ) load_face_center = np.array(load_face_data["center"]) load_face_normal = np.array(load_face_data["normal"]) fix_face_center = np.array(fix_face_data["center"]) fix_face_normal = np.array(fix_face_data["normal"]) # ----------------------------------------------------------------- # BBox per CAD face # ----------------------------------------------------------------- def get_face_bbox_dict(shape): bbox_dict = {} exp = TopExp_Explorer(shape, TopAbs_FACE) face_id = 0 while exp.More(): face_shape = exp.Current() face = TopoDS.Face_s(face_shape) bbox = Bnd_Box() BRepBndLib.Add_s(face, bbox) xmin, ymin, zmin, xmax, ymax, zmax = bbox.Get() bbox_dx = xmax - xmin bbox_dy = ymax - ymin bbox_dz = zmax - zmin bbox_dict[face_id] = (bbox_dx, bbox_dy, bbox_dz) face_id += 1 exp.Next() return bbox_dict shape = load_step_model(step_path) bbox_dict = get_face_bbox_dict(shape) # ----------------------------------------------------------------- # Build per-face feature table # ----------------------------------------------------------------- rows1 = [] for f in faces: row = {} row["face_id"] = f["face_id"] row["ftype"] = face_type_dict[f["face_id"]] row["cx"], row["cy"], row["cz"] = f["center"] row["nx"], row["ny"], row["nz"] = f["normal"] row["area"] = f.get("area") if f["face_id"] == load_face_data["face_id"]: row["role"] = "load" elif f["face_id"] == fix_face_data["face_id"]: row["role"] = "fixed" else: row["role"] = "free" dx, dy, dz = bbox_dict[f["face_id"]] row["bbox_dx"] = dx row["bbox_dy"] = dy row["bbox_dz"] = dz fc = np.array(f["center"]) row["dist_load"] = np.linalg.norm(fc - np.array(load_face_center)) row["nlx"], row["nly"], row["nlz"] = load_face_normal row["dist_fixed"] = np.linalg.norm(fc - np.array(fix_face_center)) row["nfx"], row["nfy"], row["nfz"] = fix_face_normal row["load_face"] = load_face row["load_direction"] = load_direction row["load_scale"] = load_scale row["maxsize"] = maxsize row["thickness"] = thickness rows1.append(row) df_feature = pd.DataFrame(rows1) df_feature = df_feature.fillna(0.0) df_part = df_part.fillna(0.0) # ----------------------------------------------------------------- # Model 1: refine_needed classifier # ----------------------------------------------------------------- class RefineClassifier(nn.Module): def __init__(self, in_dim): super().__init__() self.encoder = nn.Sequential( nn.Linear(in_dim, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2), nn.Linear(256, 128), nn.BatchNorm1d(128), nn.ReLU(), nn.Dropout(0.1), nn.Linear(128, 64), nn.ReLU(), ) self.head = nn.Linear(64, 1) def forward(self, x): h = self.encoder(x) return self.head(h) preprocessor1 = joblib.load("saved_model/preprocessor1.pkl") feature_cols1 = [ c for c in df_feature.columns if c not in [ "part_id", "refine_needed", "converged", "local_size", "global_size_initial", "global_size_final", "load_id", ] ] X1 = preprocessor1.transform(df_feature[feature_cols1]) X1_t = torch.tensor( X1.toarray() if hasattr(X1, "toarray") else X1, dtype=torch.float32 ) model1 = RefineClassifier(X1_t.shape[1]) state_dict = torch.load("saved_model/model1.pt", map_location="cpu") model1.load_state_dict(state_dict) model1.eval() with torch.no_grad(): logits1 = model1(X1_t) probs1 = torch.sigmoid(logits1).cpu().numpy().flatten() df_feature["refine_prob"] = probs1 top2 = df_feature.sort_values("refine_prob", ascending=False).head(2) df_feature.drop(columns=["refine_prob"], inplace=True) # ----------------------------------------------------------------- # Model 2: part-level convergence # ----------------------------------------------------------------- preprocessor2 = joblib.load("saved_model/preprocessor2.pkl") feature_cols = [ col for col in df_feature.columns if col not in [ "part_id", "load_id", "refine_needed", "converged", "local_size", "global_size_initial", "global_size_final", ] ] part_level_cols = [ col for col in df_part.columns if col not in [ "part_id", "converged", "load_id", "top2_features", "global_size_initial", "global_size_final", "local_refine_needed", "local_size", ] ] def extract_supervised_features_model2(): part_row_local = df_part.iloc[0] feat_rows = df_feature # fixed fixed_row = feat_rows[feat_rows.role == "fixed"] if len(fixed_row): X_fixed = preprocessor1.transform(fixed_row[feature_cols]) with torch.no_grad(): fixed_prob = torch.sigmoid( model1(torch.tensor(X_fixed, dtype=torch.float32)) ).item() fixed_feats = list(fixed_row.iloc[0][feature_cols]) + [fixed_prob] else: fixed_feats = [0] * (len(feature_cols) + 1) # load load_row = feat_rows[feat_rows.role == "load"] if len(load_row): X_load = preprocessor1.transform(load_row[feature_cols]) with torch.no_grad(): load_prob = torch.sigmoid( model1(torch.tensor(X_load, dtype=torch.float32)) ).item() load_feats = list(load_row.iloc[0][feature_cols]) + [load_prob] else: load_feats = [0] * (len(feature_cols) + 1) # free free_rows = feat_rows[feat_rows.role == "free"] if len(free_rows): X_free = preprocessor1.transform(free_rows[feature_cols]) with torch.no_grad(): free_probs = torch.sigmoid( model1(torch.tensor(X_free, dtype=torch.float32)) ).cpu().numpy().flatten() best_idx = np.argmax(free_probs) free_feats = ( list(free_rows.iloc[best_idx][feature_cols]) + [free_probs[best_idx]] ) else: free_feats = [0] * (len(feature_cols) + 1) global_feats = list(part_row_local[part_level_cols]) return fixed_feats + load_feats + free_feats + global_feats col_names = ( [f"fixed_{c}" for c in feature_cols] + ["fixed_refine_prob"] + [f"load_{c}" for c in feature_cols] + ["load_refine_prob"] + [f"free_{c}" for c in feature_cols] + ["free_refine_prob"] + part_level_cols ) row = extract_supervised_features_model2() df_sup = pd.DataFrame([row], columns=col_names) missing_cols = set(preprocessor2.feature_names_in_) - set(df_sup.columns) for c in missing_cols: df_sup[c] = 0.0 df_sup = df_sup[preprocessor2.feature_names_in_] X2 = preprocessor2.transform(df_sup) print("X2 contains NaN?", np.isnan(X2).any()) X2_t = torch.tensor( X2.toarray() if hasattr(X2, "toarray") else X2, dtype=torch.float32 ) class ConvergeClassifier(nn.Module): def __init__(self, in_dim): super().__init__() self.encoder = nn.Sequential( nn.Linear(in_dim, 64), nn.BatchNorm1d(64), nn.ReLU(), nn.Dropout(0.2), nn.Linear(64, 32), nn.ReLU(), ) self.head = nn.Linear(32, 2) def forward(self, x): h = self.encoder(x) return self.head(h) model2 = ConvergeClassifier(X2_t.shape[1]) state_dict = torch.load("saved_model/model2.pt", map_location="cpu") model2.load_state_dict(state_dict) model2.eval() with torch.no_grad(): logits2 = model2(X2_t) probs2 = torch.softmax(logits2, dim=1).cpu().numpy() preds2 = probs2.argmax(1) print("Model 2 (part-level converge) predictions:") print("Probability not converged:", probs2[:, 0]) print("Probability converged:", probs2[:, 1]) print("Predicted class (0: unconverged, 1: converged):", preds2) # ----------------------------------------------------------------- # Model 3: local size regressor # ----------------------------------------------------------------- class LocalRefineRegressor(nn.Module): def __init__(self, in_dim): super().__init__() self.encoder = nn.Sequential( nn.Linear(in_dim, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2), nn.Linear(256, 128), nn.BatchNorm1d(128), nn.ReLU(), nn.Dropout(0.1), nn.Linear(128, 64), nn.ReLU(), ) self.head = nn.Linear(64, 1) def forward(self, x): h = self.encoder(x) return self.head(h) preprocessor3 = joblib.load("saved_model/preprocessor3.pkl") feature_cols3 = [ c for c in df_feature.columns if c not in [ "part_id", "refine_needed", "converged", "local_size", "global_size_initial", "global_size_final", "load_id", ] ] X3 = preprocessor3.transform(df_feature[feature_cols3]) X3_t = torch.tensor( X3.toarray() if hasattr(X3, "toarray") else X3, dtype=torch.float32 ) model3 = LocalRefineRegressor(X3_t.shape[1]) state_dict = torch.load("saved_model/model3.pt", map_location="cpu") model3.load_state_dict(state_dict) model3.eval() with torch.no_grad(): logits1 = model1(X1_t) probs1 = torch.sigmoid(logits1).cpu().numpy().flatten() df_feature["refine_prob"] = probs1 top2 = df_feature.sort_values("refine_prob", ascending=False).head(2) with torch.no_grad(): local_size_pred = model3(X3_t).cpu().numpy().flatten() df_feature["predicted_local_size"] = local_size_pred top2_with_local_size = top2.copy() top2_with_local_size["predicted_local_size"] = df_feature.loc[ top2.index, "predicted_local_size" ] print("Top-2 faces needing refinement and their predicted local sizes:") print( top2_with_local_size[ ["face_id", "ftype", "predicted_local_size"] ].to_string(index=False) ) df_feature.drop(columns=["refine_prob", "predicted_local_size"], inplace=True) # ----------------------------------------------------------------- # Model 4: global size regressor # ----------------------------------------------------------------- preprocessor4 = joblib.load("saved_model/preprocessor4.pkl") feature_cols = [ col for col in df_feature.columns if col not in [ "part_id", "load_id", "refine_needed", "converged", "local_size", "global_size_initial", "global_size_final", ] ] part_level_cols = [ col for col in df_part.columns if col not in [ "part_id", "converged", "load_id", "top2_features", "global_size_initial", "global_size_final", "local_refine_needed", "local_size", ] ] def extract_supervised_features_model4(): part_row_local = df_part.iloc[0] feat_rows = df_feature fixed_row = feat_rows[feat_rows.role == "fixed"] if len(fixed_row): X_fixed = preprocessor1.transform(fixed_row[feature_cols]) with torch.no_grad(): fixed_prob = torch.sigmoid( model1(torch.tensor(X_fixed, dtype=torch.float32)) ).item() fixed_feats = list(fixed_row.iloc[0][feature_cols]) + [fixed_prob] else: fixed_feats = [0] * (len(feature_cols) + 1) load_row = feat_rows[feat_rows.role == "load"] if len(load_row): X_load = preprocessor1.transform(load_row[feature_cols]) with torch.no_grad(): load_prob = torch.sigmoid( model1(torch.tensor(X_load, dtype=torch.float32)) ).item() load_feats = list(load_row.iloc[0][feature_cols]) + [load_prob] else: load_feats = [0] * (len(feature_cols) + 1) free_rows = feat_rows[feat_rows.role == "free"] if len(free_rows): X_free = preprocessor1.transform(free_rows[feature_cols]) with torch.no_grad(): free_probs = torch.sigmoid( model1(torch.tensor(X_free, dtype=torch.float32)) ).cpu().numpy().flatten() best_idx = np.argmax(free_probs) free_feats = ( list(free_rows.iloc[best_idx][feature_cols]) + [free_probs[best_idx]] ) else: free_feats = [0] * (len(feature_cols) + 1) global_feats = list(part_row_local[part_level_cols]) return fixed_feats + load_feats + free_feats + global_feats col_names = ( [f"fixed_{c}" for c in feature_cols] + ["fixed_refine_prob"] + [f"load_{c}" for c in feature_cols] + ["load_refine_prob"] + [f"free_{c}" for c in feature_cols] + ["free_refine_prob"] + part_level_cols ) row = extract_supervised_features_model4() df_sup = pd.DataFrame([row], columns=col_names) missing_cols = set(preprocessor4.feature_names_in_) - set(df_sup.columns) for c in missing_cols: df_sup[c] = 0.0 df_sup = df_sup[preprocessor4.feature_names_in_] X4 = preprocessor4.transform(df_sup) X4_t = torch.tensor( X4.toarray() if hasattr(X4, "toarray") else X4, dtype=torch.float32 ) class GlobalSizeRegressor(nn.Module): def __init__(self, in_dim): super().__init__() self.encoder = nn.Sequential( nn.Linear(in_dim, 64), nn.BatchNorm1d(64), nn.ReLU(), nn.Dropout(0.2), nn.Linear(64, 32), nn.ReLU(), ) self.head = nn.Linear(32, 1) def forward(self, x): h = self.encoder(x) return self.head(h) model4 = GlobalSizeRegressor(X4_t.shape[1]) state_dict = torch.load("saved_model/model4.pt", map_location="cpu") model4.load_state_dict(state_dict) model4.eval() with torch.no_grad(): pred4 = model4(X4_t).item() print("Model 4 (part-level global_size) prediction:") print(f"Predicted global_size_final: {pred4:.4f}") # ----------------------------------------------------------------- # Human-readable sentence # ----------------------------------------------------------------- prob_converged = probs2[0, 1] pred_class = preds2[0] global_size = pred4 face1, face2 = top2_with_local_size.itertuples(index=False) if prob_converged > 0.7: converge_text = "very likely to converge" elif prob_converged > 0.45: converge_text = "likely to converge" else: converge_text = "unlikely to converge" sentence = ( f"The model predicts that the simulation is {converge_text} " f"(confidence: {prob_converged:.2%}). " f"To improve accuracy, local refinement will be applied on a {face1.ftype.lower()} (ID {face1.face_id}) " f"and a {face2.ftype.lower()} (ID {face2.face_id}), " f"with predicted mesh sizes {face1.predicted_local_size:.3f} mm and {face2.predicted_local_size:.3f} mm respectively. " f"The overall global mesh size is set to {global_size:.3f} mm." ) print(sentence) # ----------------------------------------------------------------- # Baseline global-only mesh (for PNG) # ----------------------------------------------------------------- gmsh.model.mesh.clear() gmsh.option.setNumber("Mesh.Optimize", 1) gmsh.option.setNumber("Mesh.OptimizeNetgen", 1) gmsh.option.setNumber("Mesh.Smoothing", 10) gmsh.option.setNumber("Mesh.Algorithm", 5) gmsh.option.setNumber("Mesh.MeshSizeFromCurvature", 0) gmsh.option.setNumber("Mesh.MeshSizeFromPoints", 0) gmsh.option.setNumber("Mesh.CharacteristicLengthMin", 0) gmsh.option.setNumber("Mesh.CharacteristicLengthMax", float(global_size)) gmsh.model.mesh.generate(2) node_tags_b, node_coords_b, _ = gmsh.model.mesh.getNodes() nodes_b = node_coords_b.reshape(-1, 3) types_b, elem_tags_list_b, node_tags_list_b = gmsh.model.mesh.getElements( dim=2 ) elements_b = [] for etype, ntags in zip(types_b, node_tags_list_b): name, dim, order, nnode, *_ = gmsh.model.mesh.getElementProperties(etype) if nnode in (3, 4): conn = ntags.reshape(-1, nnode).astype(int) - 1 elements_b.append((etype, conn)) plot_mesh( nodes_b, elements_b, title="Baseline Mesh (Global Only)", highlight_faces=None, face_to_nodes=None, save_path="output/mesh_global_only.png", ) # ----------------------------------------------------------------- # Refined mesh with background field # ----------------------------------------------------------------- gmsh.model.mesh.clear() gmsh.option.setNumber("Mesh.Optimize", 1) gmsh.option.setNumber("Mesh.OptimizeNetgen", 1) gmsh.option.setNumber("Mesh.Smoothing", 10) gmsh.option.setNumber("Mesh.Algorithm", 5) gmsh.option.setNumber("Mesh.MeshSizeFromCurvature", 0) gmsh.option.setNumber("Mesh.CharacteristicLengthMin", global_size / 4) gmsh.option.setNumber("Mesh.CharacteristicLengthMax", float(global_size)) face_ids = top2_with_local_size["face_id"].tolist() local_sizes = top2_with_local_size["predicted_local_size"].round(3).tolist() # Collect curves for distance curve_ids = [] for fid in face_ids: surf = fid - 1 boundary = gmsh.model.getBoundary([(2, surf)], oriented=False) for (dim, ctag) in boundary: if dim == 1: curve_ids.append(ctag) # Distance field f_dist = gmsh.model.mesh.field.add("Distance") gmsh.model.mesh.field.setNumbers(f_dist, "CurvesList", curve_ids) gmsh.model.mesh.field.setNumber(f_dist, "Sampling", 200) # Threshold field min_local = min(local_sizes) * 2 f_th = gmsh.model.mesh.field.add("Threshold") gmsh.model.mesh.field.setNumber(f_th, "InField", f_dist) gmsh.model.mesh.field.setNumber(f_th, "SizeMin", float(min_local)) gmsh.model.mesh.field.setNumber(f_th, "SizeMax", float(global_size)) gmsh.model.mesh.field.setNumber(f_th, "DistMin", 0.0) gmsh.model.mesh.field.setNumber(f_th, "DistMax", global_size * 2.0) # Min field (clamped) f_min = gmsh.model.mesh.field.add("Min") gmsh.model.mesh.field.setNumbers(f_min, "FieldsList", [f_th]) gmsh.model.mesh.field.setAsBackgroundMesh(f_min) gmsh.model.occ.synchronize() gmsh.model.mesh.generate(2) gmsh.model.mesh.optimize("Netgen") gmsh.model.occ.synchronize() # ----------------------------------------------------------------- # Extract FE mesh nodes & triangles (compact indexing) # ----------------------------------------------------------------- all_node_tags, all_node_coords, _ = gmsh.model.mesh.getNodes() all_nodes = all_node_coords.reshape(-1, 3) types_all, elem_tags_all, elem_nodes_all = gmsh.model.mesh.getElements() used_node_tags = np.unique( np.concatenate([arr for arr in elem_nodes_all]) ) used_idx = used_node_tags - 1 nodes_r = all_nodes[used_idx] tag_to_new_index = {tag: i for i, tag in enumerate(used_node_tags)} triangle_list = [] for etype, conn in zip(types_all, elem_nodes_all): name, dim, order, nnode, *_ = gmsh.model.mesh.getElementProperties(etype) if dim != 2: continue if nnode == 3: raw = conn.reshape(-1, 3) remapped = np.vectorize(tag_to_new_index.get)(raw) triangle_list.extend(remapped) elif nnode == 4: raw = conn.reshape(-1, 4) remapped = np.vectorize(tag_to_new_index.get)(raw) for quad in remapped: triangle_list.append([quad[0], quad[1], quad[2]]) triangle_list.append([quad[0], quad[2], quad[3]]) triangle_array = np.array(triangle_list, dtype=int) print("Triangles extracted:", triangle_array.shape) print("FE MESH NODES:", len(nodes_r)) print("FE MESH TRIANGLES:", len(triangle_array)) # For PNG plotter elements_r = [(0, triangle_array)] # ----------------------------------------------------------------- # Mapping gmsh surfaces → local node indices (for PNG highlight) # ----------------------------------------------------------------- face_to_nodes = {} for (dim, sid) in gmsh.model.getEntities(dim=2): etypes_s, elem_tags_s, ntags_s = gmsh.model.mesh.getElements( dim=2, tag=sid ) surface_conn_local = [] for etype, ntags in zip(etypes_s, ntags_s): _, _, _, nnode, *_ = gmsh.model.mesh.getElementProperties(etype) if nnode in (3, 4): raw = ntags.reshape(-1, nnode) # map gmsh node tags → compact indices mapped = np.vectorize(tag_to_new_index.get)(raw) surface_conn_local.extend(mapped) face_to_nodes[sid] = surface_conn_local highlight_gmsh_ids = [fid - 1 for fid in face_ids] # ----------------------------------------------------------------- # Local refinement mask (which triangles are "near" refined faces) # ----------------------------------------------------------------- refined_nodes_set = set() for sid in highlight_gmsh_ids: if sid in face_to_nodes: for conn in face_to_nodes[sid]: for n in conn: refined_nodes_set.add(int(n)) refined_tri_mask = np.zeros(triangle_array.shape[0], dtype=bool) for t_idx, tri in enumerate(triangle_array): if any(int(v) in refined_nodes_set for v in tri): refined_tri_mask[t_idx] = True # ----------------------------------------------------------------- # Heatmap: per-vertex "local mesh size" (smaller → more refined) # ----------------------------------------------------------------- n_nodes = nodes_r.shape[0] n_tris = triangle_array.shape[0] tri_sizes = np.zeros(n_tris, dtype=float) for t_idx, (a, b, c) in enumerate(triangle_array): pa, pb, pc = nodes_r[[a, b, c]] area = 0.5 * np.linalg.norm(np.cross(pb - pa, pc - pa)) # characteristic length ~ sqrt(area) tri_sizes[t_idx] = np.sqrt(area + 1e-16) vertex_sum = np.zeros(n_nodes, dtype=float) vertex_count = np.zeros(n_nodes, dtype=float) for t_idx, tri in enumerate(triangle_array): s = tri_sizes[t_idx] for v in tri: vertex_sum[v] += s vertex_count[v] += 1.0 vertex_size = vertex_sum / np.maximum(vertex_count, 1.0) smin, smax = float(vertex_size.min()), float(vertex_size.max()) if smax > smin: # invert so small elements → high intensity vertex_refinement_values = 1.0 - (vertex_size - smin) / (smax - smin) else: vertex_refinement_values = np.zeros_like(vertex_size) # ----------------------------------------------------------------- # Save refined mesh PNG with red outlines on refined faces # ----------------------------------------------------------------- plot_mesh( nodes_r, elements_r, title="Refined Mesh (Local + Global)", highlight_faces=highlight_gmsh_ids, face_to_nodes=face_to_nodes, save_path="output/mesh_refined.png", ) gmsh.write("output/mesh_refined.msh") # ----------------------------------------------------------------- # Interactive Plotly figure (wireframe / highlight / heatmap) # ----------------------------------------------------------------- interactive_fig = make_plotly_mesh( nodes_r, triangle_array, title="Refined Mesh (Local + Global)", view_mode=view_mode, refined_tri_mask=refined_tri_mask, vertex_refinement_values=vertex_refinement_values, ) msh_path = "output/mesh_refined.msh" return ( sentence, "output/mesh_global_only.png", "output/mesh_refined.png", interactive_fig, msh_path, )