Spaces:
Runtime error
Runtime error
| import os | |
| import numpy as np | |
| import torch | |
| from torch.utils.data import Dataset, DataLoader | |
| import pyvista as pv | |
| class Data_loader(Dataset): | |
| def __init__(self, cfg, split, epoch_seed=None, mode='train'): | |
| """ | |
| data_dir: parent directory | |
| split: list of int, e.g. [0,1,2,3,4] for train, [5] for val, [6] for test | |
| num_points: number of points to sample per geometry | |
| epoch_seed: seed for random sampling (for training) | |
| mode: 'train', 'val', or 'test' | |
| """ | |
| self.data_dir = cfg.data_dir | |
| self.split = split | |
| self.num_points = cfg.num_points | |
| self.epoch_seed = epoch_seed | |
| self.mode = mode | |
| self.cfg = cfg | |
| self.meshes = [] | |
| self.mesh_names = [] | |
| for idx in split: | |
| folder = f"{cfg.data_folder}_{idx}" | |
| vtp_file = os.path.join(self.data_dir,folder, f"{folder}.vtp") | |
| if not os.path.exists(vtp_file): | |
| raise FileNotFoundError(f"{vtp_file} not found.") | |
| mesh = pv.read(vtp_file) | |
| self.meshes.append(mesh) | |
| self.mesh_names.append(folder) | |
| # For validation chunking | |
| self.val_indices = None | |
| self.val_chunk_ptr = 0 | |
| def set_epoch(self, epoch): | |
| self.epoch_seed = epoch | |
| self.val_indices = None | |
| self.val_chunk_ptr = 0 | |
| def __len__(self): | |
| if self.mode == 'train': | |
| return len(self.meshes) | |
| elif self.mode == 'val': | |
| return len(self.meshes) | |
| elif self.mode == 'test': | |
| # Number of chunks = total points in all val meshes // num_points + remainder chunk | |
| total = 0 | |
| for mesh in self.meshes: | |
| return len(self.meshes) | |
| else: | |
| raise ValueError(f"Unknown mode: {self.mode}") | |
| def __getitem__(self, idx): | |
| if self.mode == 'train' or self.mode == 'val': | |
| # Each item is a geometry, sample num_points randomly | |
| mesh = self.meshes[idx] | |
| n_pts = mesh.points.shape[0] | |
| rng = np.random.default_rng(self.epoch_seed+idx) | |
| indices = rng.choice(n_pts, self.num_points, replace=False) | |
| pos = mesh.points[indices] | |
| target = mesh["pressure"][indices] | |
| pos = torch.tensor(pos, dtype=torch.float32) | |
| target = torch.tensor(target, dtype=torch.float32).unsqueeze(-1) | |
| if self.cfg.normalization == "std_norm": | |
| target = (target - self.cfg.press_mean) / self.cfg.press_std | |
| if self.cfg.pos_embed_sincos: | |
| input_pos_mins = torch.tensor(self.cfg.input_pos_mins) | |
| input_pos_maxs = torch.tensor(self.cfg.input_pos_maxs) | |
| pos = 1000*(pos - input_pos_mins) / (input_pos_maxs - input_pos_mins) | |
| return {"input_pos": pos, "output_feat": target ,"data_id": self.mesh_names[idx]} | |
| elif self.mode == 'test': | |
| # For each mesh in test, scramble all points and return the full mesh | |
| mesh = self.meshes[idx] | |
| n_pts = mesh.points.shape[0] | |
| rng = np.random.default_rng(self.epoch_seed+idx) | |
| indices = rng.permutation(n_pts) | |
| pos = mesh.points[indices] | |
| target = mesh["pressure"][indices] | |
| pos = torch.tensor(pos, dtype=torch.float32) | |
| target = torch.tensor(target, dtype=torch.float32).unsqueeze(-1) | |
| if self.cfg.normalization == "std_norm": | |
| target = (target - self.cfg.press_mean) / self.cfg.press_std | |
| if self.cfg.pos_embed_sincos: | |
| input_pos_mins = torch.tensor(self.cfg.input_pos_mins) | |
| input_pos_maxs = torch.tensor(self.cfg.input_pos_maxs) | |
| pos = 1000*(pos - input_pos_mins) / (input_pos_maxs - input_pos_mins) | |
| return {"input_pos": pos, "output_feat": target ,"data_id": self.mesh_names[idx],"physical_coordinates":mesh.points[indices]} | |
| else: | |
| raise ValueError(f"Unknown mode: {self.mode}") | |
| def get_dataloaders(cfg): | |
| #with open(os.path.join(cfg.splits_file, "train_design_ids.txt")) as f: | |
| # train_split = [int(line.strip().split('_')[-1]) for line in f if line.strip()] | |
| # with open(os.path.join(cfg.splits_file, "val_design_ids.txt")) as f: | |
| # val_split = [int(line.strip().split('_')[-1]) for line in f if line.strip()] | |
| # with open(os.path.join(cfg.splits_file, "test_design_ids.txt")) as f: | |
| # test_split = [int(line.strip().split('_')[-1]) for line in f if line.strip()] | |
| # Create a list from 0 to 399 | |
| all_indices = set(range(400)) | |
| # Remove numbers present in test_split | |
| test_split = [396, 397, 398, 399, 4, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 5, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 6, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 7, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 8, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 9, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] | |
| train_split = sorted(list(all_indices - set(test_split))) | |
| print("Indices from 0 to 399 not in train_split:", test_split) | |
| train_dataset = Data_loader(cfg, train_split, mode='train') | |
| val_dataset = Data_loader(cfg, test_split, mode='val') ##?? | |
| test_dataset = Data_loader(cfg, test_split, mode='test') | |
| train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True) | |
| val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False) | |
| test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False) | |
| return train_loader, val_loader, test_loader | |