dataset_pemsbay.py

import pickle
from torch.utils.data import DataLoader, Dataset
import pandas as pd
import numpy as np
import torch
import torchcde
from utils import get_randmask, get_block_mask
from scipy.interpolate import interp1d


import os
os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE"


def sample_mask(shape, p=0.0015, p_noise=0.05, max_seq=1, min_seq=1, rng=None):
    if rng is None:
        rand = np.random.random
        randint = np.random.randint
    else:
        rand = rng.random
        randint = rng.integers
    mask = rand(shape) < p
    for col in range(mask.shape[1]):
        idxs = np.flatnonzero(mask[:, col])
        if not len(idxs):
            continue
        fault_len = min_seq
        if max_seq > min_seq:
            fault_len = fault_len + int(randint(max_seq - min_seq))
        idxs_ext = np.concatenate([np.arange(i, i + fault_len) for i in idxs])
        idxs = np.unique(idxs_ext)
        idxs = np.clip(idxs, 0, shape[0] - 1)
        mask[idxs, col] = True
    mask = mask | (rand(mask.shape) < p_noise)
    return mask.astype('uint8')


class PemsBAY_Dataset_original(Dataset):
    def __init__(self, eval_length=24, mode="train", val_len=0.1, test_len=0.2, missing_pattern='block',
                 is_interpolate=False, target_strategy='random'):
        self.eval_length = eval_length
        self.is_interpolate = is_interpolate
        self.target_strategy = target_strategy
        self.mode = mode
        path = "./data/pems_bay/pems_meanstd.pk"
        with open(path, "rb") as f:
            self.train_mean, self.train_std = pickle.load(f)

        # create data for batch
        self.use_index = []
        self.cut_length = []

        df = pd.read_hdf("./data/pems_bay/pems_bay.h5")




        ob_mask = (df.values != 0.).astype('uint8')
        SEED = 9101112
        self.rng = np.random.default_rng(SEED)
        if missing_pattern == 'block':
            eval_mask = sample_mask(shape=(52116, 325), p=0.0015, p_noise=0.05, min_seq=12, max_seq=12 * 4, rng=self.rng)
        elif missing_pattern == 'point':
            eval_mask = sample_mask(shape=(52116, 325), p=0., p_noise=0.25, max_seq=12, min_seq=12 * 4, rng=self.rng)
        gt_mask = (1-(eval_mask | (1-ob_mask))).astype('uint8')

        val_start = int((1 - val_len - test_len) * len(df))
        test_start = int((1 - test_len) * len(df))
        c_data = (
             (df.fillna(0).values - self.train_mean) / self.train_std
        ) * ob_mask


########################################################################
        # impute the  original missing data
        totoal_mask = 1 - gt_mask
        content = c_data * gt_mask

        y_hat = []
        total_mask = totoal_mask
        total_mask[0] = 0
        total_mask[total_mask.shape[0] - 1] = 0
        mask_seq = [i for i in range(total_mask.shape[0])]
        mask_seq = np.array(mask_seq)
        for kk in range(total_mask.shape[1]):
            x = []
            y = []
            for ii in range(total_mask.shape[0]):
                if total_mask[ii, kk] == 0:
                    x.append(mask_seq[ii])
                    y.append(content[ii, kk])
            f = interp1d(x, y)
            y_hatt = f(mask_seq)
            y_hat.append(y_hatt)

        data_seq = np.transpose(np.array(y_hat))

#################################################################
        if mode == 'train':
            self.observed_mask = ob_mask[:val_start]
            self.gt_mask = gt_mask[:val_start]
            self.observed_data = c_data[:val_start]
            self.observed_data_interpolation = data_seq[:val_start]
        elif mode == 'valid':
            self.observed_mask = ob_mask[val_start: test_start]
            self.gt_mask = gt_mask[val_start: test_start]
            self.observed_data = c_data[val_start: test_start]
            self.observed_data_interpolation= data_seq[val_start: test_start]
        elif mode == 'test':
            self.observed_mask = ob_mask[test_start:]
            self.gt_mask = gt_mask[test_start:]
            self.observed_data = c_data[test_start:]
            self.observed_data_interpolation = data_seq[test_start:]
        else:
            self.observed_mask = ob_mask
            self.gt_mask = gt_mask
            self.observed_data = c_data
            self.observed_data_interpolation = data_seq



        current_length = len(self.observed_mask) - eval_length + 1

        if mode == "test":
            n_sample = len(self.observed_data) // eval_length
            c_index = np.arange(
                0, 0 + eval_length * n_sample, eval_length
            )
            self.use_index += c_index.tolist()
            self.cut_length += [0] * len(c_index)
            if len(self.observed_data) % eval_length != 0:
                self.use_index += [current_length - 1]
                self.cut_length += [eval_length - len(self.observed_data) % eval_length]
        elif mode != "test":
            self.use_index = np.arange(current_length)
            self.cut_length = [0] * len(self.use_index)

    def __getitem__(self, org_index):
        index = self.use_index[org_index]
        ob_data = self.observed_data[index: index + self.eval_length]
        ob_data_interpolation = self.observed_data_interpolation[index: index + self.eval_length]
        ob_mask = self.observed_mask[index: index + self.eval_length]
        ob_mask_t = torch.tensor(ob_mask).float()
        gt_mask = self.gt_mask[index: index + self.eval_length]
        if self.mode != 'train':
            cond_mask = torch.tensor(gt_mask).to(torch.float32)
        else:
            if self.target_strategy != 'random':
                cond_mask = get_block_mask(ob_mask_t, target_strategy=self.target_strategy)
            else:
                cond_mask = get_randmask(ob_mask_t)
            cond_mask = torch.tensor(gt_mask).to(torch.float32)


        s = {
            "observed_data": ob_data,
            "observed_data_interpolation": ob_data_interpolation,
            "observed_mask": ob_mask,
            "gt_mask": gt_mask,
            "hist_mask": self.observed_mask,
            "timepoints": np.arange(self.eval_length),
            "cut_length": self.cut_length[org_index],
            "cond_mask": cond_mask
        }
        if self.is_interpolate:
            tmp_data = torch.tensor(ob_data).to(torch.float64)
            itp_data = torch.where(cond_mask == 0, float('nan'), tmp_data).to(torch.float32)
            itp_data = torchcde.linear_interpolation_coeffs(
                itp_data.permute(1, 0).unsqueeze(-1)).squeeze(-1).permute(1, 0)
            s["coeffs"] = itp_data.numpy()
        return s

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


class PemsBAY_Dataset(Dataset):
    def __init__(self, eval_length=24, mode="train", val_len=0.1, test_len=0.2, missing_pattern='block',
                 is_interpolate=False, target_strategy='random'):
        self.eval_length = eval_length
        self.is_interpolate = is_interpolate
        self.target_strategy = target_strategy
        self.mode = mode
        path = "./data/pems_bay/pems_meanstd.pk"
        with open(path, "rb") as f:
            self.train_mean, self.train_std = pickle.load(f)

        # create data for batch
        self.use_index = []
        self.cut_length = []

        df = pd.read_hdf("./data/pems_bay/pems_bay.h5")

        ob_mask = (df.values != 0.).astype('uint8')
        SEED = 9101112
        self.rng = np.random.default_rng(SEED)
        if missing_pattern == 'block':
            eval_mask = sample_mask(shape=(52116, 325), p=0.0015, p_noise=0.05, min_seq=12, max_seq=12 * 4,
                                    rng=self.rng)
        elif missing_pattern == 'point':
            eval_mask = sample_mask(shape=(52116, 325), p=0., p_noise=0.25, max_seq=12, min_seq=12 * 4, rng=self.rng)
        gt_mask = (1 - (eval_mask | (1 - ob_mask))).astype('uint8')

        val_start = int((1 - val_len - test_len) * len(df))
        test_start = int((1 - test_len) * len(df))
        c_data = (
                         (df.fillna(0).values - self.train_mean) / self.train_std
                 ) * ob_mask

        ########################################################################
        # impute by the vae missing data

        data_seq1 = np.load('restore-sequences.npz')['all_predicted_train']
        data_seq1 = (data_seq1 - self.train_mean) / self.train_std

        data_seq = np.where(gt_mask==0, data_seq1, c_data)


        #################################################################
        if mode == 'train':
            self.observed_mask = ob_mask[:val_start]
            self.gt_mask = gt_mask[:val_start]
            self.observed_data = c_data[:val_start]
            self.observed_data_interpolation = data_seq[:val_start]
        elif mode == 'valid':
            self.observed_mask = ob_mask[val_start: test_start]
            self.gt_mask = gt_mask[val_start: test_start]
            self.observed_data = c_data[val_start: test_start]
            self.observed_data_interpolation = data_seq[val_start: test_start]
        elif mode == 'test':
            self.observed_mask = ob_mask[test_start:]
            self.gt_mask = gt_mask[test_start:]
            self.observed_data = c_data[test_start:]
            self.observed_data_interpolation = data_seq[test_start:]
        current_length = len(self.observed_mask) - eval_length + 1

        if mode == "test":
            n_sample = len(self.observed_data) // eval_length
            c_index = np.arange(
                0, 0 + eval_length * n_sample, eval_length
            )
            self.use_index += c_index.tolist()
            self.cut_length += [0] * len(c_index)
            if len(self.observed_data) % eval_length != 0:
                self.use_index += [current_length - 1]
                self.cut_length += [eval_length - len(self.observed_data) % eval_length]
        elif mode != "test":
            self.use_index = np.arange(current_length)
            self.cut_length = [0] * len(self.use_index)

    def __getitem__(self, org_index):
        index = self.use_index[org_index]
        ob_data = self.observed_data[index: index + self.eval_length]
        ob_data_interpolation = self.observed_data_interpolation[index: index + self.eval_length]
        ob_mask = self.observed_mask[index: index + self.eval_length]
        ob_mask_t = torch.tensor(ob_mask).float()
        gt_mask = self.gt_mask[index: index + self.eval_length]
        if self.mode != 'train':
            cond_mask = torch.tensor(gt_mask).to(torch.float32)
        else:
            if self.target_strategy != 'random':
                cond_mask = get_block_mask(ob_mask_t, target_strategy=self.target_strategy)
            else:
                cond_mask = get_randmask(ob_mask_t)
            cond_mask = torch.tensor(gt_mask).to(torch.float32)

        s = {
            "observed_data": ob_data,
            "observed_data_interpolation": ob_data_interpolation,
            "observed_mask": ob_mask,
            "gt_mask": gt_mask,
            "hist_mask": self.observed_mask,
            "timepoints": np.arange(self.eval_length),
            "cut_length": self.cut_length[org_index],
            "cond_mask": cond_mask
        }
        if self.is_interpolate:
            tmp_data = torch.tensor(ob_data).to(torch.float64)
            itp_data = torch.where(cond_mask == 0, float('nan'), tmp_data).to(torch.float32)
            itp_data = torchcde.linear_interpolation_coeffs(
                itp_data.permute(1, 0).unsqueeze(-1)).squeeze(-1).permute(1, 0)
            s["coeffs"] = itp_data.numpy()
        return s

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

def get_dataloader_original(batch_size, device, val_len=0.1, test_len=0.2, missing_pattern='block',
                   is_interpolate=False, num_workers=4, target_strategy='random'):
    dataset = PemsBAY_Dataset_original(mode="train", val_len=val_len, test_len=test_len, missing_pattern=missing_pattern,
                             is_interpolate=is_interpolate, target_strategy=target_strategy)
    train_loader = DataLoader(
        dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True
    )

    dataset_test_train = PemsBAY_Dataset_original(mode="ttest-train", val_len=val_len, test_len=test_len, missing_pattern=missing_pattern,
                             is_interpolate=is_interpolate, target_strategy=target_strategy)
    train_test_loader = DataLoader(
        dataset_test_train, batch_size=batch_size, num_workers=num_workers, shuffle=False
    )


    dataset_test = PemsBAY_Dataset_original(mode="test", val_len=val_len, test_len=test_len, missing_pattern=missing_pattern,
                                  is_interpolate=is_interpolate, target_strategy=target_strategy)
    test_loader = DataLoader(
        dataset_test, batch_size=batch_size, num_workers=num_workers, shuffle=False
    )



    dataset_valid = PemsBAY_Dataset_original(mode="valid", val_len=val_len, test_len=test_len, missing_pattern=missing_pattern,
                                   is_interpolate=is_interpolate, target_strategy=target_strategy)
    valid_loader = DataLoader(
        dataset_valid, batch_size=batch_size, num_workers=num_workers, shuffle=False
    )

    scaler = torch.from_numpy(dataset.train_std).to(device).float()
    mean_scaler = torch.from_numpy(dataset.train_mean).to(device).float()

    return train_loader, valid_loader, test_loader, train_test_loader, valid_loader, scaler, mean_scaler



def get_dataloader(batch_size, device, val_len=0.1, test_len=0.2, missing_pattern='block',
                   is_interpolate=False, num_workers=4, target_strategy='random'):
    dataset = PemsBAY_Dataset(mode="train", val_len=val_len, test_len=test_len, missing_pattern=missing_pattern,
                             is_interpolate=is_interpolate, target_strategy=target_strategy)
    train_loader = DataLoader(
        dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True
    )
    dataset_test = PemsBAY_Dataset(mode="test", val_len=val_len, test_len=test_len, missing_pattern=missing_pattern,
                                  is_interpolate=is_interpolate, target_strategy=target_strategy)
    test_loader = DataLoader(
        dataset_test, batch_size=batch_size, num_workers=num_workers, shuffle=False
    )
    dataset_valid = PemsBAY_Dataset(mode="valid", val_len=val_len, test_len=test_len, missing_pattern=missing_pattern,
                                   is_interpolate=is_interpolate, target_strategy=target_strategy)
    valid_loader = DataLoader(
        dataset_valid, batch_size=batch_size, num_workers=num_workers, shuffle=False
    )

    scaler = torch.from_numpy(dataset.train_std).to(device).float()
    mean_scaler = torch.from_numpy(dataset.train_mean).to(device).float()

    return train_loader, valid_loader, test_loader, scaler, mean_scaler