dataset.py

from pathlib import Path
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader, random_split, TensorDataset
import h5py
import numpy as np
import torch
import tqdm



class DatasetOneStep(Dataset):

    def __init__(self, label, phase):
        self.x_org, self.y_org = DatasetOneStep._feature_gen(label)  # shape (N, 49, ?, 50, 50)
        self.whole_dataset = TensorDataset(self.x_org, self.y_org)

        dataset_size = len(self.x_org)
        train_size = 1240
        test_size = dataset_size - train_size
        train_dataset, test_dataset = random_split(
            self.whole_dataset,
            [train_size, test_size],
            generator=torch.Generator().manual_seed(777))
        print(f'>>>>>>>train cases: {len(train_dataset)}, test cases: {len(test_dataset)}<<<<<<<')

        if phase == 'train':
            self.dataset = train_dataset
        elif phase == 'test':
            self.dataset = test_dataset

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, item):
        return self.dataset[item]

    def _merge_dimensions(self, dataset):
        x, y = zip(*dataset)  # Unzip the dataset
        x = torch.cat(x).view(-1, x[0].size(1), *x[0].size()[2:])
        y = torch.cat(y).view(-1, y[0].size(1), *y[0].size()[2:])
        print(f'||||||x shape: {x.shape}, y shape: {y.shape}||||||')
        return TensorDataset(x, y)

    @staticmethod
    def _feature_gen(label_type):
        data = DatasetOneStep._pre_process()

        whole_s = torch.tensor(data['sG_3d']).float().squeeze().unsqueeze(2)  # (N, T, 1, 50, 50)
        whole_p = torch.tensor(data['pres_3d']).float().squeeze().unsqueeze(2)  # (N, T, 1, 50, 50)
        label_s = whole_s[:, 1:, :, :, :]
        label_p = whole_p[:, 1:, :, :, :]
        print(f'----> State variables shape: Sg {label_s.shape}; P {label_p.shape}')

        # print(label_s.shape)
        num_samples, num_steps, _, height, width = label_s.shape

        # previous state
        # previous_s = whole_s[:, :-1, :, :, :]
        # previous_p = whole_p[:, :-1, :, :, :]
        previous_s = whole_s[:, :-1, :, :, :].clone()
        previous_p = whole_p[:, :-1, :, :, :].clone()
        # noise_s = torch.normal(mean=0., std=0.02, size=previous_s.shape)
        # noise_p = torch.normal(mean=0., std=0.02, size=previous_p.shape)
        # previous_s[:, 1:, :, :, :] += noise_s[:, 1:, :, :, :]
        # previous_p[:, 1:, :, :, :] += noise_p[:, 1:, :, :, :]

        # spatial variable
        perm = torch.tensor(data['perm_3d']).float().squeeze().unsqueeze(1).expand(-1, num_steps, -1, -1).unsqueeze(2)  # (N, 50, 50)
        poro = torch.tensor(data['poro_3d']).float().squeeze().unsqueeze(1).expand(-1, num_steps, -1, -1).unsqueeze(2)  # (N, 50, 50)
        pore = torch.tensor(data['pore_vol']).float().squeeze().unsqueeze(1).expand(-1, num_steps, -1, -1).unsqueeze(2)  # (N, 50, 50)

        grid_xcoord = torch.tensor(data['grid_xcoord']).float().squeeze().unsqueeze(0).unsqueeze(0).expand(num_samples, num_steps, -1, -1).unsqueeze(2)  # (N, 50, 50)
        grid_zcoord = torch.tensor(data['grid_zcoord']).float().squeeze().unsqueeze(0).unsqueeze(0).expand(num_samples, num_steps, -1, -1).unsqueeze(2)  # (N, 50, 50)

        well = torch.tensor(data['Cummulative_InjPerf']).float().squeeze().unsqueeze(2)[:, 1:, :, :, :]  # (N, T, 50, 50)

        # scalar
        salinity = torch.tensor(data['Salinity']).float().unsqueeze(1).unsqueeze(1).unsqueeze(1).expand(-1, num_steps, height, width).unsqueeze(2)  # (N, )
        temperature = torch.tensor(data['Temperature']).float().unsqueeze(1).unsqueeze(1).unsqueeze(1).expand(-1, num_steps, height, width).unsqueeze(2)  # (N, )
        thickness = torch.tensor(data['Res_thicknes']).float().unsqueeze(1).unsqueeze(1).unsqueeze(1).expand(-1, num_steps, height, width).unsqueeze(2)  # (N, )
        rate = torch.tensor(data['Injection_rate']).float().unsqueeze(1).unsqueeze(1).unsqueeze(1).expand(-1, num_steps, height, width).unsqueeze(2)  # (N, )

        # time series
        delta_time = torch.diff(torch.tensor(data['timeStep_ts'])).float()
        time_step = delta_time.unsqueeze(0).unsqueeze(-1).unsqueeze(-1).expand(num_samples, -1, height, width).unsqueeze(2)  # (T)
        injection_reminder = torch.tensor(data['qInj_ts']).float().unsqueeze(-1).unsqueeze(-1).expand(-1, -1, height, width).unsqueeze(2)[:, 1:, :, :, :]  # (N, T)

        # all the tensors are in shape: (N, T, 1, 50, 50)
        # inputs = torch.cat((perm, poro, pore, well, grid_xcoord, grid_zcoord, salinity, temperature, thickness, time_step), dim=2)  # should be (N, T, x, 50, 50)
        # print(f'----> Input features shape: {inputs.shape}')

        if label_type == 'saturation':
            inputs = torch.cat((previous_s, perm, poro, pore, well, grid_xcoord, grid_zcoord, injection_reminder, salinity, temperature, thickness, time_step), dim=2)
            print(f'----> Input features shape: {inputs.shape}. State variables shape: Sg {label_s.shape}')
            return inputs[:, :30, :, :, :], label_s[:, :30, :, :, :]  # [N, T, C, 50, 50]

        elif label_type == 'pressure':
            inputs = torch.cat((previous_p, perm, poro, pore, well, salinity, temperature, thickness, time_step), dim=2)
            inputs = inputs[:, :59, :, :, :]
            label_p = label_p[:, :59, :, :, :]
            print(f'----> Input features shape: {inputs.shape}. State variables shape: P {label_p.shape}')
            return inputs, label_p

    @staticmethod
    def _pre_process():
        """ min max norm all data """
        data_dict = DatasetOneStep._load_from_h5()
        normed_data = {}

        field_feature_list = ['perm_3d', 'poro_3d', 'pore_vol', 'Cummulative_InjPerf', 'grid_xcoord', 'grid_zcoord', 'qInj_ts']  # (499, 49, 50, 1, 50)
        scalar_feature_list = ['Salinity', 'Temperature', 'Res_thicknes', 'Injection_rate', 'timeStep_ts']  # (499,)
        label_list = ['sG_3d', 'pres_3d']  # (499, 49, 50, 1, 50)

        whole_list = field_feature_list + scalar_feature_list + label_list
        for key_name in whole_list:
            arr = data_dict.get(key_name).astype(np.float32)
            if key_name == 'sG_3d':
                normed_data.update({key_name: arr})
            else:
                norm_arr = DatasetOneStep._min_max_norm(arr)
                normed_data.update({key_name: norm_arr})
        return normed_data

    @staticmethod
    def _load_from_h5():
        """Load all data from h5."""
        root_path = r'G:\Transformers\Dataset\new'
        h5_path = Path(root_path, 'Rad2d_1300cases_10kmLength.h5')

        # print(f'Load simulation data from h5 format...')
        # Init h5_keys
        h5_keys = ['Injection_rate', 'Salinity', 'Temperature', 'perm_3d', 'Res_thicknes', 'pore_vol', 'poro_3d',
                   'pres_3d', 'sG_3d', 'timeStep_ts', 'Cummulative_InjPerf', 'grid_xcoord', 'grid_zcoord', 'qInj_ts']

        # If h5 file already save all the numpy array
        if Path(h5_path).exists():

            data = {}

            # Get all the numpy array first
            with h5py.File(h5_path, 'r') as hf:
                for key in h5_keys:
                    if key in hf.keys():
                        val = hf.get(name=key)[:]
                        data.update({key: val})
        return data

    @staticmethod
    def _min_max_norm(input_arr):
        """ normalize input numpy arrays """
        v_min, v_max = np.min(input_arr), np.max(input_arr)
        normed_arr = (input_arr - v_min) / (v_max - v_min)
        return normed_arr

    @staticmethod
    def _stand_norm(input_arr, save_name):
        mu = np.mean(input_arr)
        sigma = np.std(input_arr)
        norm_arr = (input_arr - mu) / sigma
        np.savetxt(f'G:\\Transformers\\checkpoint\\norm_records\\[radial]{save_name}.txt', np.array([mu, sigma]),
                   delimiter='\n')
        return norm_arr















class DatasetNonAR(Dataset):

    def __init__(self, label, test_phase):
        if not test_phase:
            x_org, y_org = DatasetNonAR._feature_gen(label)  # shape (N, T, c, 50, 50)
            x = x_org[:-381]
            y = y_org[:-381]

            self.x = x
            self.y = y

            print(f'After resampling, x-y pairs shape: {self.x.shape}, {self.y.shape}')

            self.x = self.x.reshape(-1, x.size(2), x.size(3), x.size(4))
            self.y = self.y.reshape(-1, y.size(2), y.size(3), y.size(4))
        elif test_phase:
            self.x_org, self.y_org = DatasetNonAR._feature_gen(label)  # shape (N, T, c, 50, 50)
            self.x = self.x_org[-381:]
            self.y = self.y_org[-381:]
            # don't view it because we need to evaluate case by case

    def __len__(self):
        return len(self.y)

    def __getitem__(self, item):
        return self.x[item], self.y[item]

    @staticmethod
    def _feature_gen(label_type):
        data = DatasetNonAR._pre_process()

        label_s = torch.tensor(data['sG_3d']).float().squeeze().unsqueeze(2)  # (N, T, 1, 50, 50)
        label_p = torch.tensor(data['pres_3d']).float().squeeze().unsqueeze(2)  # (N, T, 1, 50, 50)
        print(f'----> Original state variables shape: Sg {label_s.shape}; P {label_p.shape}')
        num_samples, num_steps, _, height, width = label_s.shape

        # spatial variable
        perm = torch.tensor(data['perm_3d']).float().squeeze().unsqueeze(1).expand(-1, num_steps, -1, -1).unsqueeze(2)  # (N, 50, 50)
        poro = torch.tensor(data['poro_3d']).float().squeeze().unsqueeze(1).expand(-1, num_steps, -1, -1).unsqueeze(2)  # (N, 50, 50)
        pore = torch.tensor(data['pore_vol']).float().squeeze().unsqueeze(1).expand(-1, num_steps, -1, -1).unsqueeze(2)  # (N, 50, 50)
        rate_perf = torch.tensor(data['InjRate_Perf']).float().squeeze().unsqueeze(1).expand(-1, num_steps, -1, -1).unsqueeze(2)  # (N, 50, 50)
        well = torch.tensor(data['Cummulative_InjPerf']).float().squeeze().unsqueeze(2)  # (N, T, 50, 50)

        # scalar
        salinity = torch.tensor(data['Salinity']).float().unsqueeze(1).unsqueeze(1).unsqueeze(1).expand(-1, num_steps, height, width).unsqueeze(2)  # (N, )
        temperature = torch.tensor(data['Temperature']).float().unsqueeze(1).unsqueeze(1).unsqueeze(1).expand(-1, num_steps, height, width).unsqueeze(2)  # (N, )
        thickness = torch.tensor(data['Res_thicknes']).float().unsqueeze(1).unsqueeze(1).unsqueeze(1).expand(-1, num_steps, height, width).unsqueeze(2)  # (N, )
        perforation = torch.tensor(data['Perforations']).float().unsqueeze(1).unsqueeze(1).unsqueeze(1).expand(-1, num_steps, height, width).unsqueeze(2)  # (N, )
        rate = torch.tensor(data['Injection_rate']).float().unsqueeze(1).unsqueeze(1).unsqueeze(1).expand(-1, num_steps, height, width).unsqueeze(2)  # (N, )

        # time series
        time_step = torch.tensor(data['timeStep_ts']).float().unsqueeze(0).unsqueeze(-1).unsqueeze(-1).expand(num_samples, -1, height, width).unsqueeze(2)  # (T)
        # injection_reminder = torch.tensor(data['qInj_ts']).float().unsqueeze(-1).unsqueeze(-1).expand(-1, -1, height, width).unsqueeze(2)  # (N, T)

        # all the tensors are in shape: (N, T, 1, 50, 50)

        if label_type == 'saturation':
            inputs = torch.cat((perm, poro, pore, well, rate_perf, rate, perforation, salinity, temperature, thickness, time_step), dim=2)
            inputs = inputs[:, 1:, :, :, :]  # 75    # ::3    # 31
            label_s = label_s[:, 1:, :, :, :]
            print(f'----> Input features shape: {inputs.shape}. State variables shape: Sg {label_s.shape}')
            return inputs, label_s  # [N, T, C, 50, 50]

        elif label_type == 'pressure':
            inputs = torch.cat((perm, poro, pore, well, rate_perf, rate, perforation, salinity, temperature, thickness, time_step), dim=2)
            inputs = inputs[:, 1:, :, :, :]
            label_p = label_p[:, 1:, :, :, :]
            print(f'----> Input features shape: {inputs.shape}. State variables shape: P {label_p.shape}')
            return inputs, label_p

    @staticmethod
    def _pre_process():
        """ min max norm all data """
        data_dict = DatasetNonAR._load_from_h5()
        normed_data = {}

        field_feature_list = ['perm_3d', 'poro_3d', 'pore_vol', 'Cummulative_InjPerf']
        scalar_feature_list = ['Salinity', 'Temperature', 'Res_thicknes', 'Injection_rate', 'timeStep_ts', 'Perforations', 'InjRate_Perf']
        label_list = ['sG_3d', 'pres_3d']

        whole_list = field_feature_list + scalar_feature_list + label_list
        for key_name in whole_list:
            arr = data_dict.get(key_name).astype(np.float32)
            if key_name == 'sG_3d':
                normed_data.update({key_name: arr})
            elif key_name == 'timeStep_ts':
                # norm_arr = np.log1p(arr)
                norm_arr = DatasetNonAR._min_max_norm(arr)
                normed_data.update({key_name: norm_arr})
            else:
                norm_arr = DatasetNonAR._min_max_norm(arr)
                normed_data.update({key_name: norm_arr})
        return normed_data

    @staticmethod
    def _load_from_h5():
        """Load all data from h5."""
        root_path = r'G:\Transformers\Dataset\new'
        h5_path = Path(root_path, 'Rad2d_1300cases_10kmLength.h5')

        # print(f'Load simulation data from h5 format...')
        # Init h5_keys
        h5_keys = ['Injection_rate', 'Salinity', 'Temperature', 'perm_3d', 'Res_thicknes', 'pore_vol', 'poro_3d',
                   'pres_3d', 'sG_3d', 'timeStep_ts', 'Cummulative_InjPerf', 'Perforations', 'InjRate_Perf']

        # If h5 file already save all the numpy array
        if Path(h5_path).exists():

            data = {}

            # Get all the numpy array first
            with h5py.File(h5_path, 'r') as hf:
                for key in h5_keys:
                    if key in hf.keys():
                        val = hf.get(name=key)[:]
                        data.update({key: val})
        return data

    @staticmethod
    def _min_max_norm(input_arr):
        """ normalize input numpy arrays """
        v_min, v_max = np.min(input_arr), np.max(input_arr)
        normed_arr = (input_arr - v_min) / (v_max - v_min)
        return normed_arr



# ========================================================================================================================================================================


class OneStepCartesion(Dataset):

    def __init__(self, label, test_phase):
        if not test_phase:
            x_org, y_org = DatasetOneStepCartesion._feature_gen(label)  # shape (N, T, c, 50, 50)
            x = x_org[:-100]
            y = y_org[:-100]

            # x_rot = torch.rot90(x, k=1, dims=[3, 4])
            # y_rot = torch.rot90(y, k=1, dims=[3, 4])
            #
            # self.x = torch.cat((x, x_rot), dim=0)
            # self.y = torch.cat((y, y_rot), dim=0)

            self.x = x
            self.y = y

            print(f'After resampling, x-y pairs shape: {self.x.shape}, {self.y.shape}')

            self.x = self.x.reshape(-1, x.size(2), x.size(3), x.size(4))
            self.y = self.y.reshape(-1, y.size(2), y.size(3), y.size(4))
        elif test_phase:
            self.x_org, self.y_org = OneStepCartesion._feature_gen(label)  # shape (N, T, c, 50, 50)
            self.x = self.x_org[-100:]
            self.y = self.y_org[-100:]
            # don't view it because we need to evaluate case by case

    def __len__(self):
        return len(self.y)

    def __getitem__(self, item):
        return self.x[item], self.y[item]

    @staticmethod
    def _feature_gen(label_type):
        data = OneStepCartesion._pre_process()

        label_s = torch.tensor(data['saturation']).float().squeeze().unsqueeze(2)  # (N, T, 1, 50, 50)
        label_p = torch.tensor(data['pressure']).float().squeeze().unsqueeze(2)  # (N, T, 1, 50, 50)
        print(f'----> Original state variables shape: Sg {label_s.shape}; P {label_p.shape}')
        num_samples, num_steps, _, height, width = label_s.shape

        # spatial variable
        perm = torch.tensor(data['permeability']).float().squeeze().unsqueeze(2)  # (N, 50, 50)
        poro = torch.tensor(data['porosity']).float().squeeze().unsqueeze(2)  # (N, 50, 50)
        pos = torch.tensor(data['well_position']).float().squeeze().unsqueeze(2)  # (N, 50, 50)

        well = torch.tensor(data['bottomhole_pressure']).float().unsqueeze(-1).unsqueeze(-1).expand(-1, -1, height, width).unsqueeze(2)  # (N, T, )

        # scalar
        temperature = torch.tensor(data['temperature']).float().unsqueeze(-1).unsqueeze(-1).expand(-1, -1, height, width).unsqueeze(2)  # (N, 60)
        rate = torch.tensor(data['injection_rate']).float().unsqueeze(-1).unsqueeze(-1).expand(-1, -1, height, width).unsqueeze(2)

        # time series
        time_step = torch.tensor(data['time_step']).float().unsqueeze(-1).unsqueeze(-1).expand(-1, -1, height, width).unsqueeze(2)  # (T)
        # injection_reminder = torch.tensor(data['qInj_ts']).float().unsqueeze(-1).unsqueeze(-1).expand(-1, -1, height, width).unsqueeze(2)  # (N, T)

        # all the tensors are in shape: (N, T, 1, 50, 50)

        if label_type == 'saturation':
            inputs = torch.cat((perm, poro, pos, rate, well, temperature, time_step),dim=2)
            # inputs = inputs[:, 1:, :, :, :]  # 75    # ::3    # 31
            # label_s = label_s[:, 1:, :, :, :]
            print(f'----> Input features shape: {inputs.shape}. State variables shape: Sg {label_s.shape}')
            return inputs, label_s  # [N, T, C, 50, 50]

        elif label_type == 'pressure':
            inputs = torch.cat((perm, poro, pos, rate, well, temperature, time_step),dim=2)
            # inputs = inputs[:, 1:, :, :, :]
            # label_p = label_p[:, 1:, :, :, :]
            print(f'----> Input features shape: {inputs.shape}. State variables shape: P {label_p.shape}')
            return inputs, label_p

    @staticmethod
    def _pre_process():
        """ min max norm all data """
        data_dict = DatasetOneStepCartesion._load_from_h5()
        normed_data = {}

        field_feature_list = ['permeability', 'porosity', 'well_position']  # (b, t, 50, 1, 50)
        scalar_feature_list = ['injection_rate', 'time_step', 'bottomhole_pressure', 'temperature']  # (499,)
        label_list = ['saturation', 'pressure']  # (499, 49, 50, 1, 50)

        whole_list = field_feature_list + scalar_feature_list + label_list
        for key_name in whole_list:
            arr = data_dict.get(key_name).astype(np.float32)
            if key_name == 'saturation':
                normed_data.update({key_name: arr})
            else:
                norm_arr = DatasetOneStepCartesion._min_max_norm(arr)
                normed_data.update({key_name: norm_arr})
        return normed_data

    @staticmethod
    def _load_from_h5():
        """Load all data from h5."""
        root_path = r'G:\Transformers\Dataset\new'
        h5_path = Path(root_path, 'MyCartesion_dataset_new.hdf5')

        print(f'Load simulation data from h5 format...')
        # Init h5_keys
        h5_keys = ['permeability', 'porosity', 'well_position', 'injection_rate', 'time_step', 'bottomhole_pressure', 'temperature', 'pressure', 'saturation']

        # If h5 file already save all the numpy array
        if Path(h5_path).exists():

            data = {}

            # Get all the numpy array first
            with h5py.File(h5_path, 'r') as hf:
                for key in h5_keys:
                    if key in hf.keys():
                        val = hf.get(name=key)[:]
                        data.update({key: val})
                        # print(f'Complete loading {key}.')

        return data

    @staticmethod
    def _min_max_norm(input_arr):
        """ normalize input numpy arrays """
        v_min, v_max = np.min(input_arr), np.max(input_arr)
        print(f'max = {v_max}, min = {v_min}')
        normed_arr = (input_arr - v_min) / (v_max - v_min)
        return normed_arr