dynamics_model_dataset.py

import torch
from torch.utils.data import Dataset
import h5py
import numpy as np
import pathlib

# Per-worker HDF5 chunk cache cap. Total RAM = num_workers × this value.
_MAX_RDCC_NBYTES = 4 * 1024**3  # 4 GB


def getDatasetName(system, name):
    filename = name
    filename += "_" + str(system.system_type)
    filename += "_dt_" + str(system.dt)
    filename += "_T_" + str(system.T)
    filename += "_tau_" + str(system.tau)
    filename += ".hdf5"
    return filename


def append_to_dataset(path, data, log_name):
    """
    Stores the full input / state values in one long array
        and the specific indices that should appear in each
    :param path: full path to the hdf5 file
    :param data: dictionary storing numpy arrays in expected format
    :param log_name: name of the log file, manages what indexes are for what log
    It becomes the name of a dataset in the hdf5 file itself
    :return:
    """
    if not path.exists():
        # create new file from scratch
        with h5py.File(path, "w") as f:
            num_points, input_dim = data["data"]["input"].shape
            state_dim = data["data"]["state"].shape[1]
            check_valid_dim = data["data"]["check_valid"].shape[1]

            data_chunk_size = 100
            data_grp = f.create_group("data")
            input = data_grp.create_dataset(
                "input",
                data["data"]["input"].shape,
                maxshape=(None, input_dim),
                chunks=(data_chunk_size, input_dim),
                compression="lzf",
            )
            state = data_grp.create_dataset(
                "state",
                data["data"]["state"].shape,
                maxshape=(None, state_dim),
                chunks=(data_chunk_size, state_dim),
                compression="lzf",
            )
            check_valid = data_grp.create_dataset(
                "check_valid",
                data["data"]["check_valid"].shape,
                maxshape=(None, check_valid_dim),
                chunks=(data_chunk_size, check_valid_dim),
                compression="lzf",
            )
            meta_data = data_grp.create_dataset(
                "meta_data",
                data["data"]["meta_data"].shape,
                maxshape=(None, 300),
                chunks=(data_chunk_size, 300),
                compression="lzf",
                dtype=np.int32,
            )
            is_valid = data_grp.create_dataset(
                "is_valid",
                data["data"]["is_valid"].shape,
                maxshape=(None, 1),
                chunks=(data_chunk_size, 1),
                compression="lzf",
                dtype=np.bool,
            )
            split = data_grp.create_dataset(
                "split",
                data["data"]["split"].shape,
                maxshape=(None, 1),
                chunks=(data_chunk_size, 1),
                compression="lzf",
                dtype=np.int32,
            )

            index_chunk_size = 100
            train_indexes = f.create_dataset(
                "train",
                data["train"].shape,
                maxshape=(None, 1),
                chunks=(index_chunk_size, 1),
                compression="lzf",
                dtype=np.int32,
            )
            val_indexes = f.create_dataset(
                "validation",
                data["validation"].shape,
                maxshape=(None, 1),
                chunks=(index_chunk_size, 1),
                compression="lzf",
                dtype=np.int32,
            )
            f.create_dataset(log_name, data=(0, num_points))

            input[:] = data["data"]["input"].astype(np.float32)
            state[:] = data["data"]["state"].astype(np.float32)
            check_valid[:] = data["data"]["check_valid"].astype(np.float32)
            meta_data[:] = data["data"]["meta_data"].astype(np.int32)
            is_valid[:] = data["data"]["is_valid"].astype(np.bool)
            split[:] = data["data"]["split"].astype(np.int32)
            val_indexes[:] = data["validation"].astype(np.int32)
            train_indexes[:] = data["train"].astype(np.int32)
    else:
        # append to the path
        new_points = data["data"]["input"].shape[0]
        train_new_points = data["train"].shape[0]
        val_new_points = data["validation"].shape[0]
        print(
            f"appending to compressed pickle {path.name} got val {val_new_points} and train {train_new_points}"
        )
        with h5py.File(path, "a") as f:
            old_max_index = f["data/input"].shape[0]

            # create space for all data
            f["data/input"].resize(new_points + old_max_index, axis=0)
            f["data/state"].resize(new_points + old_max_index, axis=0)
            f["data/check_valid"].resize(new_points + old_max_index, axis=0)
            f["data/meta_data"].resize(new_points + old_max_index, axis=0)
            f["data/is_valid"].resize(new_points + old_max_index, axis=0)
            f["data/split"].resize(new_points + old_max_index, axis=0)
            # create space for new indexes in train
            if train_new_points > 0:
                f["train"].resize(train_new_points + f["train"].shape[0], axis=0)
            # create space for new indexes in val
            if val_new_points > 0:
                f["validation"].resize(
                    val_new_points + f["validation"].shape[0], axis=0
                )

            # add in the new data
            f["data/input"][-new_points:] = data["data"]["input"].astype(np.float32)
            f["data/state"][-new_points:] = data["data"]["state"].astype(np.float32)
            f["data/check_valid"][-new_points:] = data["data"]["check_valid"].astype(
                np.float32
            )
            f["data/meta_data"][-new_points:] = data["data"]["meta_data"].astype(
                np.int32
            )
            f["data/is_valid"][-new_points:] = data["data"]["is_valid"].astype(np.bool)
            f["data/split"][-new_points:] = data["data"]["split"].astype(np.int32)

            # append the new indexes that we got
            if train_new_points > 0:
                f["train"][-train_new_points:] = (data["train"] + old_max_index).astype(
                    np.int32
                )
            if val_new_points > 0:
                f["validation"][-val_new_points:] = (
                    data["validation"] + old_max_index
                ).astype(np.int32)

            f.create_dataset(log_name, data=(old_max_index, f["data/input"].shape[0]))


class DynamicsModelDataset(Dataset):
    def __init__(
        self,
        dataset_file_path,
        type,
        init_length,
        traj_length,
        rdcc_nbytes=None,
        rdcc_nslots=None,
    ):
        self.dataset_file_path = pathlib.Path(dataset_file_path)
        if not self.dataset_file_path.exists():
            raise RuntimeError(f"Dataset not found at path {dataset_file_path}")
        with h5py.File(dataset_file_path, "r") as f:
            self.type = type
            self.len = f[type].len()
            total_rows = f["data/input"].shape[0]
            input_dim = f["data/input"].shape[1]
            state_dim = f["data/state"].shape[1]
            chunks = f["data/input"].chunks
            chunk_size = chunks[0] if chunks is not None else 1
        self.init_length = int(init_length)
        self.traj_length = int(traj_length)
        self._total_rows = total_rows
        num_chunks = -(-total_rows // chunk_size)
        self._rdcc_nbytes = (
            rdcc_nbytes
            if rdcc_nbytes is not None
            else min(total_rows * (input_dim + state_dim) * 4, _MAX_RDCC_NBYTES)
        )
        # h5py uses open addressing, so the slot table needs headroom above the
        # number of chunks to keep lookup O(1). ~10× avoids collisions; a prime
        # multiplier (11) reduces clustering from the modular hash function.
        self._rdcc_nslots = (
            rdcc_nslots if rdcc_nslots is not None else num_chunks * 11
        )

    def __len__(self):
        # return self.inputs.shape[0]
        return self.len

    def _open_hdf5(self):
        self._hf = h5py.File(
            self.dataset_file_path,
            "r",
            rdcc_nbytes=self._rdcc_nbytes,
            rdcc_nslots=self._rdcc_nslots,
        )

    def __getitem__(self, idx):
        if not hasattr(self, "_hf"):
            self._open_hdf5()
        dataset_index = int(self._hf[self.type][idx].item())
        inputs = torch.from_numpy(
            self._hf["data/input"][
                dataset_index - self.init_length : dataset_index + self.traj_length,
                :,
            ]
        )
        states = torch.from_numpy(
            self._hf["data/state"][
                dataset_index - self.init_length : dataset_index + self.traj_length + 1,
                :,
            ]
        )

        return inputs, states


class DynamicsModelFullDataset(DynamicsModelDataset):
    def __init__(
        self,
        dataset_file_path,
        type,
        init_length,
        traj_length,
        rdcc_nbytes=None,
        rdcc_nslots=None,
    ):
        super().__init__(
            dataset_file_path, type, init_length, traj_length, rdcc_nbytes, rdcc_nslots
        )
        if rdcc_nbytes is None:
            with h5py.File(dataset_file_path, "r") as f:
                check_valid_dim = f["data/check_valid"].shape[1]
            self._rdcc_nbytes = min(
                self._rdcc_nbytes + self._total_rows * check_valid_dim * 4,
                _MAX_RDCC_NBYTES,
            )

    def __getitem__(self, idx):
        if not hasattr(self, "_hf"):
            self._open_hdf5()

        dataset_index = int(self._hf[self.type][idx].item())
        inputs = torch.from_numpy(
            self._hf["data/input"][
                dataset_index - self.init_length : dataset_index + self.traj_length,
                :,
            ]
        )
        states = torch.from_numpy(
            self._hf["data/state"][
                dataset_index - self.init_length : dataset_index + self.traj_length + 1,
                :,
            ]
        )
        check_valid_states = self._hf["data/check_valid"][
            dataset_index - self.init_length : dataset_index + self.traj_length + 1
        ]
        meta_data = self._hf["data/meta_data"][dataset_index]

        return inputs, states, check_valid_states, meta_data