dataset.py

import torch
from torch.utils.data import Dataset
import numpy as np
import os
from tqdm import tqdm, trange


class MCDataset(Dataset):
    def __init__(self, data_base="/media/aritra/F7B7-B745/final_data", num_images=10_000, seed=23, N=8, Js=[-1, 1]):
        self.data_base = data_base
        self.num_images = num_images
        self.seed = seed
        self.N = N
        self.Js = Js
        self.large_number = 1

        if self.num_images % self.large_number != 0:
            raise ValueError("num_images must be divisible by large_number")

        self.files = []
        tj = []
        self.data = []

        for J in tqdm(self.Js, desc="Loading MC Dataset"):
            folder = f"{self.data_base}/{self.num_images}_{self.seed}_{self.N}_{J}"
            files = os.listdir(folder)
            sorted_files = sorted(files, key=lambda x: float(x.split(".npy")[0]))
            
            self.files.extend([(folder, f) for f in sorted_files])
            tj.extend([(float(f.split(".npy")[0]), J) for f in sorted_files])
            self.data.extend([np.load(f"{folder}/{f}") for f in sorted_files])

        # self.data = np.array(self.data, dtype=np.float32) * 2 - 1  # Normalize
        self.data = np.array(self.data, dtype=np.bool)
        self.tj = np.array(tj, dtype=np.float32)

        self.data = self.data.reshape(-1, self.large_number, self.N, self.N)
        num_blocks = self.num_images // self.large_number
        self.tj = self.tj.repeat(num_blocks, axis=0)

        self.data = torch.from_numpy(self.data)
        self.tj = torch.from_numpy(self.tj)

        # print(f"{self.data.shape = }, {self.tj.shape = }")

    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.tj[idx], self.data[idx]