steady-rans-surrogates / code /ezflow_v3 /baselines /Transolver-main /Car-Design-ShapeNetCar /dataset /dataset.py
| import torch | |
| import vtk | |
| import os | |
| import itertools | |
| import random | |
| import numpy as np | |
| from torch_geometric import nn as nng | |
| from sklearn.neighbors import NearestNeighbors | |
| from torch_geometric.data import Data, Dataset | |
| from torch_geometric.utils import k_hop_subgraph, subgraph | |
| from vtk.util.numpy_support import vtk_to_numpy | |
| def load_unstructured_grid_data(file_name): | |
| reader = vtk.vtkUnstructuredGridReader() | |
| reader.SetFileName(file_name) | |
| reader.Update() | |
| output = reader.GetOutput() | |
| return output | |
| def unstructured_grid_data_to_poly_data(unstructured_grid_data): | |
| filter = vtk.vtkDataSetSurfaceFilter() | |
| filter.SetInputData(unstructured_grid_data) | |
| filter.Update() | |
| poly_data = filter.GetOutput() | |
| return poly_data, filter | |
| def get_sdf(target, boundary): | |
| nbrs = NearestNeighbors(n_neighbors=1).fit(boundary) | |
| dists, indices = nbrs.kneighbors(target) | |
| neis = np.array([boundary[i[0]] for i in indices]) | |
| dirs = (target - neis) / (dists + 1e-8) | |
| return dists.reshape(-1), dirs | |
| def get_normal(unstructured_grid_data): | |
| poly_data, surface_filter = unstructured_grid_data_to_poly_data(unstructured_grid_data) | |
| # visualize_poly_data(poly_data, surface_filter) | |
| # poly_data.GetPointData().SetScalars(None) | |
| normal_filter = vtk.vtkPolyDataNormals() | |
| normal_filter.SetInputData(poly_data) | |
| normal_filter.SetAutoOrientNormals(1) | |
| normal_filter.SetConsistency(1) | |
| # normal_filter.SetSplitting(0) | |
| normal_filter.SetComputeCellNormals(1) | |
| normal_filter.SetComputePointNormals(0) | |
| normal_filter.Update() | |
| ''' | |
| normal_filter.SetComputeCellNormals(0) | |
| normal_filter.SetComputePointNormals(1) | |
| normal_filter.Update() | |
| #visualize_poly_data(poly_data, surface_filter, normal_filter) | |
| poly_data.GetPointData().SetNormals(normal_filter.GetOutput().GetPointData().GetNormals()) | |
| p2c = vtk.vtkPointDataToCellData() | |
| p2c.ProcessAllArraysOn() | |
| p2c.SetInputData(poly_data) | |
| p2c.Update() | |
| unstructured_grid_data.GetCellData().SetNormals(p2c.GetOutput().GetCellData().GetNormals()) | |
| #visualize_poly_data(poly_data, surface_filter, p2c) | |
| ''' | |
| unstructured_grid_data.GetCellData().SetNormals(normal_filter.GetOutput().GetCellData().GetNormals()) | |
| c2p = vtk.vtkCellDataToPointData() | |
| # c2p.ProcessAllArraysOn() | |
| c2p.SetInputData(unstructured_grid_data) | |
| c2p.Update() | |
| unstructured_grid_data = c2p.GetOutput() | |
| # return unstructured_grid_data | |
| normal = vtk_to_numpy(c2p.GetOutput().GetPointData().GetNormals()).astype(np.double) | |
| # print(np.max(np.max(np.abs(normal), axis=1)), np.min(np.max(np.abs(normal), axis=1))) | |
| normal /= (np.max(np.abs(normal), axis=1, keepdims=True) + 1e-8) | |
| normal /= (np.linalg.norm(normal, axis=1, keepdims=True) + 1e-8) | |
| if np.isnan(normal).sum() > 0: | |
| print(np.isnan(normal).sum()) | |
| print("recalculate") | |
| return get_normal(unstructured_grid_data) # re-calculate | |
| # print(normal) | |
| return normal | |
| def visualize_poly_data(poly_data, surface_filter, normal_filter=None): | |
| if normal_filter is not None: | |
| mask = vtk.vtkMaskPoints() | |
| mask.SetInputData(normal_filter.GetOutput()) | |
| # mask.RandomModeOn() | |
| mask.Update() | |
| arrow = vtk.vtkArrowSource() | |
| arrow.Update() | |
| glyph = vtk.vtkGlyph3D() | |
| glyph.SetInputData(mask.GetOutput()) | |
| glyph.SetSourceData(arrow.GetOutput()) | |
| glyph.SetVectorModeToUseNormal() | |
| glyph.SetScaleFactor(0.1) | |
| glyph.Update() | |
| norm_mapper = vtk.vtkPolyDataMapper() | |
| norm_mapper.SetInputData(normal_filter.GetOutput()) | |
| glyph_mapper = vtk.vtkPolyDataMapper() | |
| glyph_mapper.SetInputData(glyph.GetOutput()) | |
| norm_actor = vtk.vtkActor() | |
| norm_actor.SetMapper(norm_mapper) | |
| glyph_actor = vtk.vtkActor() | |
| glyph_actor.SetMapper(glyph_mapper) | |
| glyph_actor.GetProperty().SetColor(1, 0, 0) | |
| norm_render = vtk.vtkRenderer() | |
| norm_render.AddActor(norm_actor) | |
| norm_render.SetBackground(0, 1, 0) | |
| glyph_render = vtk.vtkRenderer() | |
| glyph_render.AddActor(glyph_actor) | |
| glyph_render.AddActor(norm_actor) | |
| glyph_render.SetBackground(0, 0, 1) | |
| scalar_range = poly_data.GetScalarRange() | |
| mapper = vtk.vtkDataSetMapper() | |
| mapper.SetInputConnection(surface_filter.GetOutputPort()) | |
| mapper.SetScalarRange(scalar_range) | |
| actor = vtk.vtkActor() | |
| actor.SetMapper(mapper) | |
| renderer = vtk.vtkRenderer() | |
| renderer.AddActor(actor) | |
| renderer.SetBackground(1, 1, 1) # Set background to white | |
| renderer_window = vtk.vtkRenderWindow() | |
| renderer_window.AddRenderer(renderer) | |
| if normal_filter is not None: | |
| renderer_window.AddRenderer(norm_render) | |
| renderer_window.AddRenderer(glyph_render) | |
| renderer_window.Render() | |
| interactor = vtk.vtkRenderWindowInteractor() | |
| interactor.SetRenderWindow(renderer_window) | |
| interactor.Initialize() | |
| interactor.Start() | |
| def get_datalist(root, samples, norm=False, coef_norm=None, savedir=None, preprocessed=False): | |
| dataset = [] | |
| mean_in, mean_out = 0, 0 | |
| std_in, std_out = 0, 0 | |
| for k, s in enumerate(samples): | |
| if preprocessed and savedir is not None: | |
| save_path = os.path.join(savedir, s) | |
| if not os.path.exists(save_path): | |
| continue | |
| init = np.load(os.path.join(save_path, 'x.npy')) | |
| target = np.load(os.path.join(save_path, 'y.npy')) | |
| pos = np.load(os.path.join(save_path, 'pos.npy')) | |
| surf = np.load(os.path.join(save_path, 'surf.npy')) | |
| edge_index = np.load(os.path.join(save_path, 'edge_index.npy')) | |
| else: | |
| file_name_press = os.path.join(root, os.path.join(s, 'quadpress_smpl.vtk')) | |
| file_name_velo = os.path.join(root, os.path.join(s, 'hexvelo_smpl.vtk')) | |
| if not os.path.exists(file_name_press) or not os.path.exists(file_name_velo): | |
| continue | |
| unstructured_grid_data_press = load_unstructured_grid_data(file_name_press) | |
| unstructured_grid_data_velo = load_unstructured_grid_data(file_name_velo) | |
| velo = vtk_to_numpy(unstructured_grid_data_velo.GetPointData().GetVectors()) | |
| press = vtk_to_numpy(unstructured_grid_data_press.GetPointData().GetScalars()) | |
| points_velo = vtk_to_numpy(unstructured_grid_data_velo.GetPoints().GetData()) | |
| points_press = vtk_to_numpy(unstructured_grid_data_press.GetPoints().GetData()) | |
| edges_press = get_edges(unstructured_grid_data_press, points_press, cell_size=4) | |
| edges_velo = get_edges(unstructured_grid_data_velo, points_velo, cell_size=8) | |
| sdf_velo, normal_velo = get_sdf(points_velo, points_press) | |
| sdf_press = np.zeros(points_press.shape[0]) | |
| normal_press = get_normal(unstructured_grid_data_press) | |
| surface = {tuple(p) for p in points_press} | |
| exterior_indices = [i for i, p in enumerate(points_velo) if tuple(p) not in surface] | |
| velo_dict = {tuple(p): velo[i] for i, p in enumerate(points_velo)} | |
| pos_ext = points_velo[exterior_indices] | |
| pos_surf = points_press | |
| sdf_ext = sdf_velo[exterior_indices] | |
| sdf_surf = sdf_press | |
| normal_ext = normal_velo[exterior_indices] | |
| normal_surf = normal_press | |
| velo_ext = velo[exterior_indices] | |
| velo_surf = np.array([velo_dict[tuple(p)] if tuple(p) in velo_dict else np.zeros(3) for p in pos_surf]) | |
| press_ext = np.zeros([len(exterior_indices), 1]) | |
| press_surf = press | |
| init_ext = np.c_[pos_ext, sdf_ext, normal_ext] | |
| init_surf = np.c_[pos_surf, sdf_surf, normal_surf] | |
| target_ext = np.c_[velo_ext, press_ext] | |
| target_surf = np.c_[velo_surf, press_surf] | |
| surf = np.concatenate([np.zeros(len(pos_ext)), np.ones(len(pos_surf))]) | |
| pos = np.concatenate([pos_ext, pos_surf]) | |
| init = np.concatenate([init_ext, init_surf]) | |
| target = np.concatenate([target_ext, target_surf]) | |
| edge_index = get_edge_index(pos, edges_press, edges_velo) | |
| if savedir is not None: | |
| save_path = os.path.join(savedir, s) | |
| if not os.path.exists(save_path): | |
| os.makedirs(save_path) | |
| np.save(os.path.join(save_path, 'x.npy'), init) | |
| np.save(os.path.join(save_path, 'y.npy'), target) | |
| np.save(os.path.join(save_path, 'pos.npy'), pos) | |
| np.save(os.path.join(save_path, 'surf.npy'), surf) | |
| np.save(os.path.join(save_path, 'edge_index.npy'), edge_index) | |
| surf = torch.tensor(surf) | |
| pos = torch.tensor(pos) | |
| x = torch.tensor(init) | |
| y = torch.tensor(target) | |
| edge_index = torch.tensor(edge_index) | |
| if norm and coef_norm is None: | |
| if k == 0: | |
| old_length = init.shape[0] | |
| mean_in = init.mean(axis=0) | |
| mean_out = target.mean(axis=0) | |
| else: | |
| new_length = old_length + init.shape[0] | |
| mean_in += (init.sum(axis=0) - init.shape[0] * mean_in) / new_length | |
| mean_out += (target.sum(axis=0) - init.shape[0] * mean_out) / new_length | |
| old_length = new_length | |
| data = Data(pos=pos, x=x, y=y, surf=surf.bool(), edge_index=edge_index) | |
| # data = Data(pos=pos, x=x, y=y, surf=surf.bool()) | |
| dataset.append(data) | |
| if norm and coef_norm is None: | |
| for k, data in enumerate(dataset): | |
| if k == 0: | |
| old_length = data.x.numpy().shape[0] | |
| std_in = ((data.x.numpy() - mean_in) ** 2).sum(axis=0) / old_length | |
| std_out = ((data.y.numpy() - mean_out) ** 2).sum(axis=0) / old_length | |
| else: | |
| new_length = old_length + data.x.numpy().shape[0] | |
| std_in += (((data.x.numpy() - mean_in) ** 2).sum(axis=0) - data.x.numpy().shape[ | |
| 0] * std_in) / new_length | |
| std_out += (((data.y.numpy() - mean_out) ** 2).sum(axis=0) - data.x.numpy().shape[ | |
| 0] * std_out) / new_length | |
| old_length = new_length | |
| std_in = np.sqrt(std_in) | |
| std_out = np.sqrt(std_out) | |
| for data in dataset: | |
| data.x = ((data.x - mean_in) / (std_in + 1e-8)).float() | |
| data.y = ((data.y - mean_out) / (std_out + 1e-8)).float() | |
| coef_norm = (mean_in, std_in, mean_out, std_out) | |
| dataset = (dataset, coef_norm) | |
| elif coef_norm is not None: | |
| for data in dataset: | |
| data.x = ((data.x - coef_norm[0]) / (coef_norm[1] + 1e-8)).float() | |
| data.y = ((data.y - coef_norm[2]) / (coef_norm[3] + 1e-8)).float() | |
| return dataset | |
| def get_edges(unstructured_grid_data, points, cell_size=4): | |
| edge_indeces = set() | |
| cells = vtk_to_numpy(unstructured_grid_data.GetCells().GetData()).reshape(-1, cell_size + 1) | |
| for i in range(len(cells)): | |
| for j, k in itertools.product(range(1, cell_size + 1), repeat=2): | |
| edge_indeces.add((cells[i][j], cells[i][k])) | |
| edge_indeces.add((cells[i][k], cells[i][j])) | |
| edges = [[], []] | |
| for u, v in edge_indeces: | |
| edges[0].append(tuple(points[u])) | |
| edges[1].append(tuple(points[v])) | |
| return edges | |
| def get_edge_index(pos, edges_press, edges_velo): | |
| indices = {tuple(pos[i]): i for i in range(len(pos))} | |
| edges = set() | |
| for i in range(len(edges_press[0])): | |
| edges.add((indices[edges_press[0][i]], indices[edges_press[1][i]])) | |
| for i in range(len(edges_velo[0])): | |
| edges.add((indices[edges_velo[0][i]], indices[edges_velo[1][i]])) | |
| edge_index = np.array(list(edges)).T | |
| return edge_index | |
| def get_induced_graph(data, idx, num_hops): | |
| subset, sub_edge_index, _, _ = k_hop_subgraph(node_idx=idx, num_hops=num_hops, edge_index=data.edge_index, | |
| relabel_nodes=True) | |
| return Data(x=data.x[subset], y=data.y[idx], edge_index=sub_edge_index) | |
| def pc_normalize(pc): | |
| centroid = torch.mean(pc, axis=0) | |
| pc = pc - centroid | |
| m = torch.max(torch.sqrt(torch.sum(pc ** 2, axis=1))) | |
| pc = pc / m | |
| return pc | |
| def get_shape(data, max_n_point=8192, normalize=True, use_height=False): | |
| surf_indices = torch.where(data.surf)[0].tolist() | |
| if len(surf_indices) > max_n_point: | |
| surf_indices = np.array(random.sample(range(len(surf_indices)), max_n_point)) | |
| shape_pc = data.pos[surf_indices].clone() | |
| if normalize: | |
| shape_pc = pc_normalize(shape_pc) | |
| if use_height: | |
| gravity_dim = 1 | |
| height_array = shape_pc[:, gravity_dim:gravity_dim + 1] - shape_pc[:, gravity_dim:gravity_dim + 1].min() | |
| shape_pc = torch.cat((shape_pc, height_array), axis=1) | |
| return shape_pc | |
| def create_edge_index_radius(data, r, max_neighbors=32): | |
| data.edge_index = nng.radius_graph(x=data.pos, r=r, loop=True, max_num_neighbors=max_neighbors) | |
| # print(f'r = {r}, #edges = {data.edge_index.size(1)}') | |
| return data | |
| class GraphDataset(Dataset): | |
| def __init__(self, datalist, use_height=False, use_cfd_mesh=True, r=None): | |
| super().__init__() | |
| self.datalist = datalist | |
| self.use_height = use_height | |
| if not use_cfd_mesh: | |
| assert r is not None | |
| for i in range(len(self.datalist)): | |
| self.datalist[i] = create_edge_index_radius(self.datalist[i], r) | |
| def len(self): | |
| return len(self.datalist) | |
| def get(self, idx): | |
| data = self.datalist[idx] | |
| shape = get_shape(data, use_height=self.use_height) | |
| return self.datalist[idx], shape | |
| if __name__ == '__main__': | |
| import numpy as np | |
| file_name = '1a0bc9ab92c915167ae33d942430658c' | |
| root = '/data/PDE_data/mlcfd_data/training_data' | |
| save_path = '/data/PDE_data/mlcfd_data/preprocessed_data/param0/' + file_name | |
| file_name_press = 'param0/' + file_name + '/quadpress_smpl.vtk' | |
| file_name_velo = 'param0/' + file_name + '/hexvelo_smpl.vtk' | |
| file_name_press = os.path.join(root, file_name_press) | |
| file_name_velo = os.path.join(root, file_name_velo) | |
| unstructured_grid_data_press = load_unstructured_grid_data(file_name_press) | |
| unstructured_grid_data_velo = load_unstructured_grid_data(file_name_velo) | |
| velo = vtk_to_numpy(unstructured_grid_data_velo.GetPointData().GetVectors()) | |
| press = vtk_to_numpy(unstructured_grid_data_press.GetPointData().GetScalars()) | |
| points_velo = vtk_to_numpy(unstructured_grid_data_velo.GetPoints().GetData()) | |
| points_press = vtk_to_numpy(unstructured_grid_data_press.GetPoints().GetData()) | |
| edges_press = get_edges(unstructured_grid_data_press, points_press, cell_size=4) | |
| edges_velo = get_edges(unstructured_grid_data_velo, points_velo, cell_size=8) | |
| sdf_velo, normal_velo = get_sdf(points_velo, points_press) | |
| sdf_press = np.zeros(points_press.shape[0]) | |
| normal_press = get_normal(unstructured_grid_data_press) | |
| surface = {tuple(p) for p in points_press} | |
| exterior_indices = [i for i, p in enumerate(points_velo) if tuple(p) not in surface] | |
| velo_dict = {tuple(p): velo[i] for i, p in enumerate(points_velo)} | |
| pos_ext = points_velo[exterior_indices] | |
| pos_surf = points_press | |
| sdf_ext = sdf_velo[exterior_indices] | |
| sdf_surf = sdf_press | |
| normal_ext = normal_velo[exterior_indices] | |
| normal_surf = normal_press | |
| velo_ext = velo[exterior_indices] | |
| velo_surf = np.array([velo_dict[tuple(p)] if tuple(p) in velo_dict else np.zeros(3) for p in pos_surf]) | |
| press_ext = np.zeros([len(exterior_indices), 1]) | |
| press_surf = press | |
| init_ext = np.c_[pos_ext, sdf_ext, normal_ext] | |
| init_surf = np.c_[pos_surf, sdf_surf, normal_surf] | |
| target_ext = np.c_[velo_ext, press_ext] | |
| target_surf = np.c_[velo_surf, press_surf] | |
| surf = np.concatenate([np.zeros(len(pos_ext)), np.ones(len(pos_surf))]) | |
| pos = np.concatenate([pos_ext, pos_surf]) | |
| init = np.concatenate([init_ext, init_surf]) | |
| target = np.concatenate([target_ext, target_surf]) | |
| edge_index = get_edge_index(pos, edges_press, edges_velo) | |
| data = Data(pos=torch.tensor(pos), edge_index=torch.tensor(edge_index)) | |
| data = create_edge_index_radius(data, r=0.2) | |
| x, y = data.edge_index | |
| import torch_geometric | |
| print(max(torch_geometric.utils.degree(x)), max(torch_geometric.utils.degree(y))) | |
| print(points_velo.shape, points_press.shape) | |
| print(surf.shape, pos.shape, init.shape, target.shape, edge_index.shape) | |