| import torch |
| import numpy as np |
| from torch.utils.data import Dataset |
| from scipy.interpolate import griddata |
| from scipy.linalg import qr |
|
|
| class CFDReconstructionDataset(Dataset): |
| def __init__(self, unsteady_path, mean_path, sensor_count=65, target_grid=(128, 256), dt=0.01): |
| super().__init__() |
| print("[Dataset] Initializing v3.0 (Sparse Mode: 5 returns)...") |
| self.sensor_count = sensor_count |
| self.H, self.W = target_grid |
| self.dt = dt |
| |
| try: |
| raw_unsteady = np.load(unsteady_path).astype(np.float32) |
| raw_mean = np.load(mean_path).astype(np.float32) |
| except Exception as e: |
| print(f"Error loading data: {e}") |
| raise |
|
|
| mean_cols = raw_mean.shape[1] if raw_mean.ndim > 1 else 1 |
| coords_mean = raw_mean[:, :2] |
| x_min, y_min = coords_mean.min(axis=0) |
| x_max, y_max = coords_mean.max(axis=0) |
| self.box_len = [x_max - x_min, y_max - y_min, 1.0] |
| |
| grid_x = np.linspace(x_min, x_max, self.W) |
| grid_y = np.linspace(y_min, y_max, self.H) |
| self.grid_X, self.grid_Y = np.meshgrid(grid_x, grid_y) |
| self.grid_coords = torch.stack([torch.tensor(self.grid_X), torch.tensor(self.grid_Y)], dim=-1).float() |
| self.grid_coords_norm = self.grid_coords.clone() |
| self.grid_coords_norm[..., 0] = 2 * (self.grid_coords[..., 0] - x_min) / (x_max - x_min) - 1 |
| self.grid_coords_norm[..., 1] = 2 * (self.grid_coords[..., 1] - y_min) / (y_max - y_min) - 1 |
|
|
| if mean_cols >= 6: mean_vals = raw_mean[:, 3:6] |
| elif mean_cols == 5: mean_vals = raw_mean[:, 2:5] |
| else: mean_vals = raw_mean[:, -3:] |
|
|
| self.mean_data = griddata(coords_mean, mean_vals, (self.grid_X, self.grid_Y), method='linear', fill_value=0) |
| self.mean_data = torch.from_numpy(self.mean_data).permute(2, 0, 1).float() |
|
|
| unsteady_cols = raw_unsteady.shape[-1] if raw_unsteady.ndim > 1 else raw_mean.shape[-1] |
| points_per_frame = raw_mean.shape[0] |
| total_elements = raw_unsteady.size |
| frame_size = points_per_frame * unsteady_cols |
| num_frames = total_elements // frame_size |
| raw_unsteady = raw_unsteady.flatten()[:num_frames * frame_size] |
| raw_unsteady = raw_unsteady.reshape(num_frames, points_per_frame, unsteady_cols) |
|
|
| process_frames = min(num_frames, 200) |
| data_list = [] |
| for i in range(process_frames): |
| frame_data = raw_unsteady[i] |
| if unsteady_cols >= 6: values = frame_data[:, 3:6] |
| elif unsteady_cols == 5: values = frame_data[:, 2:5] |
| else: values = frame_data[:, :3] |
| coords = frame_data[:, :2] if unsteady_cols >= 5 else coords_mean |
| grid_val = griddata(coords, values, (self.grid_X, self.grid_Y), method='linear', fill_value=0) |
| data_list.append(grid_val) |
| |
| self.data = np.stack(data_list, axis=0) |
| self.data = torch.from_numpy(self.data).permute(0, 3, 1, 2).float() |
|
|
| self.stats = {} |
| self.max_vals = torch.amax(self.data, dim=(0, 2, 3), keepdim=True) |
| self.min_vals = torch.amin(self.data, dim=(0, 2, 3), keepdim=True) |
| self.stats['max'] = self.max_vals |
| self.stats['min'] = self.min_vals |
| self.stats['dt'] = self.dt |
| self.stats['box_len'] = self.box_len |
| |
| denom = self.max_vals - self.min_vals |
| denom[denom < 1e-8] = 1.0 |
| self.data = 2 * (self.data - self.min_vals) / denom - 1 |
| self.mean_data = 2 * (self.mean_data - self.min_vals.squeeze(0)) / denom.squeeze(0) - 1 |
| |
| print(f"[Dataset] Selecting {sensor_count} sensors (QR Pivot)...") |
| self.sensor_indices = self.compute_qr_sensors(self.data, self.mean_data, sensor_count) |
|
|
| def compute_qr_sensors(self, data, mean, num_sensors): |
| T, C, H, W = data.shape |
| fluctuations = (data[:, 0, :, :] - mean[0, :, :]).reshape(T, -1).numpy().T |
| k = min(T, num_sensors + 10) |
| U, _, _ = np.linalg.svd(fluctuations, full_matrices=False) |
| Psi = U[:, :k] |
| _, _, P = qr(Psi.T, pivoting=True) |
| best_indices_flat = P[:num_sensors] |
| sensor_locs = [] |
| for idx in best_indices_flat: |
| y = idx // W |
| x = idx % W |
| sensor_locs.append((y, x)) |
| return sensor_locs |
|
|
| def denormalize(self, tensor): |
| return (tensor + 1) / 2 * (self.stats['max'].to(tensor.device) - self.stats['min'].to(tensor.device)) + self.stats['min'].to(tensor.device) |
|
|
| def __len__(self): |
| return max(0, self.data.shape[0] - 1) |
|
|
| def __getitem__(self, idx): |
| if idx >= self.data.shape[0] - 1: idx = self.data.shape[0] - 2 |
| target_t = self.data[idx] |
| target_next = self.data[idx + 1] |
| |
| vals_t, coords_t = [], [] |
| for y, x in self.sensor_indices: |
| vals_t.append(target_t[:, y, x]) |
| coords_t.append(self.grid_coords_norm[y, x]) |
| |
| vals_next = [] |
| for y, x in self.sensor_indices: |
| vals_next.append(target_next[:, y, x]) |
| |
| return (torch.stack(vals_t), torch.stack(coords_t), self.grid_coords_norm, |
| torch.stack(vals_next), self.mean_data) |
