dataset_loader.py

import os
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import pyvista as pv
import json
import glob

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:
            # Find folder matching *_{idx}
            folder = os.path.join(self.data_dir, f"{idx}")
            if not os.path.exists(folder):
                raise FileNotFoundError(f"No folder matching '{idx}' found in {self.data_dir}")

            # Find file matching *_{idx}.vtp inside the folder
            vtp_files = glob.glob(os.path.join(folder, f"{idx}.vtp"))
            if not vtp_files:
                raise FileNotFoundError(f"No file matching '{idx}.vtp' found in {folder}")
            vtp_file = vtp_files[0]
            mesh = pv.read(vtp_file)
            self.meshes.append(mesh)
            self.mesh_names.append(os.path.splitext(os.path.basename(vtp_file))[0])
        # For validation chunking
        self.val_indices = None
        self.val_chunk_ptr = 0
        with open(cfg.json_file, "r") as f:
            self.json_data = json.load(f)
     
    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
            pos = torch.tensor(pos, dtype=torch.float32)
            pressure = torch.tensor( mesh["pressure"][indices], dtype=torch.float32).unsqueeze(-1)
            
            if self.cfg.normalization == "std_norm":
                target = (pressure - self.json_data["scalars"]["pressure"]["mean"]) / self.json_data["scalars"]["pressure"]["std"]

            if self.cfg.diff_input_velocity:
                inlet_x_vel = torch.tensor( mesh["inlet_x_velocity"], dtype=torch.float32).unsqueeze(-1)
                pos = torch.cat((pos,inlet_x_vel),dim = 1)
            
            if self.cfg.input_normalization == "shift_axis":
                coords = pos[:,:3].clone()

                # Shift x: set minimum x (front bumper) to 0
                coords[:, 0] = coords[:, 0] - coords[:, 0].min()

                # Shift z: set minimum z (ground) to 0
                coords[:, 2] = coords[:, 2] - coords[:, 2].min()

                # Shift y: center about 0 (left/right symmetry)
                y_center = (coords[:, 1].max() + coords[:, 1].min()) / 2.0
                coords[:, 1] = coords[:, 1] - y_center

                pos[:,:3] = coords
                
                
            if self.cfg.pos_embed_sincos:
                
                if self.cfg.diff_input_velocity:
                    raise Exception("pos_embed_sincos not supported with diff_input_velocity=True")
                
                input_pos_mins = torch.tensor(self.json_data["mesh_stats"]["min"])
                input_pos_maxs = torch.tensor(self.json_data["mesh_stats"]["max"])
                pos = 1000*(pos - input_pos_mins) / (input_pos_maxs - input_pos_mins)
                assert torch.all(pos >= 0) 
                assert torch.all(pos <= 1000)

            pos = pos[indices]
            
            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
            
            pos = torch.tensor(pos, dtype=torch.float32)
            pressure = torch.tensor( mesh["pressure"][indices], dtype=torch.float32).unsqueeze(-1)

            if self.cfg.normalization == "std_norm":
                target = (pressure - self.json_data["scalars"]["pressure"]["mean"]) / self.json_data["scalars"]["pressure"]["std"]
                
            if hasattr(self.cfg, "diff_input_velocity") and self.cfg.diff_input_velocity:
            
                inlet_x_vel = torch.tensor( mesh["inlet_x_velocity"], dtype=torch.float32).unsqueeze(-1)
                pos = torch.cat((pos,inlet_x_vel),dim = 1)
                
            if self.cfg.input_normalization == "shift_axis":
                
                coords = pos[:,:3].clone()

                # Shift x: set minimum x (front bumper) to 0
                coords[:, 0] = coords[:, 0] - coords[:, 0].min()

                # Shift z: set minimum z (ground) to 0
                coords[:, 2] = coords[:, 2] - coords[:, 2].min()

                # Shift y: center about 0 (left/right symmetry)
                y_center = (coords[:, 1].max() + coords[:, 1].min()) / 2.0
                coords[:, 1] = coords[:, 1] - y_center

                pos[:,:3] = coords
             
            if self.cfg.pos_embed_sincos:
                
                if hasattr(self.cfg, "diff_input_velocity") and self.cfg.diff_input_velocity:
                    raise Exception("pos_embed_sincos not supported with diff_input_velocity=True")
                
                input_pos_mins = torch.tensor(self.json_data["mesh_stats"]["min"])
                input_pos_maxs = torch.tensor(self.json_data["mesh_stats"]["max"])
                pos = 1000*(pos - input_pos_mins) / (input_pos_maxs - input_pos_mins)
                assert torch.all(pos >= 0) 
                assert torch.all(pos <= 1000)
            
            pos = pos[indices]
            
            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.txt")) as f:
        train_split = [line.strip() for line in f if line.strip()]
    with open(os.path.join(cfg.splits_file, "test.txt")) as f:
        val_split = [line.strip() for line in f if line.strip()]
    with open(os.path.join(cfg.splits_file, "test.txt")) as f:
        test_split = [line.strip() for line in f if line.strip()]
    print("Indices in test_split:", test_split[:5])  # Print first 5 indices for verification


    train_dataset = Data_loader(cfg, train_split, mode='train')
    val_dataset = Data_loader(cfg, val_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