{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preprocess data for these Sine (simulation), Stocks, and Energy Dataset\n",
    "* each dataset should have dimension of (sample_size, number_of_features, sequence_length)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sine (with 5 features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def sine_data_generation (no, seq_len, dim):\n",
    "  \"\"\"\n",
    "  Sine data generation. Adapted from TimeGAN source code for sake of consistency.\n",
    "  \n",
    "  Args:\n",
    "    - no: the number of samples\n",
    "    - seq_len: sequence length of the time-series\n",
    "    - dim: feature dimensions\n",
    "    \n",
    "  Returns:\n",
    "    - data: generated data\n",
    "  \"\"\"  \n",
    "  # Initialize the output\n",
    "  data = list()\n",
    "\n",
    "  # Generate sine data\n",
    "  for i in range(no):      \n",
    "    # Initialize each time-series\n",
    "    temp = list()\n",
    "    # For each feature\n",
    "    for k in range(dim):\n",
    "      # Randomly drawn frequency and phase\n",
    "      freq = np.random.uniform(0, 0.1)            \n",
    "      phase = np.random.uniform(0, 0.1)\n",
    "          \n",
    "      # Generate sine signal based on the drawn frequency and phase\n",
    "      temp_data = [np.sin(freq * j + phase) for j in range(seq_len)] \n",
    "      temp.append(temp_data)\n",
    "        \n",
    "    # Align row/column\n",
    "    temp = np.transpose(np.asarray(temp))        \n",
    "    # Normalize to [0,1]\n",
    "    temp = (temp + 1)*0.5\n",
    "    # Stack the generated data\n",
    "    data.append(temp)\n",
    "                \n",
    "  return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "sine_data = sine_data_generation(20000, 50, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(20000, 50, 5)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sine_data = np.asarray(sine_data)\n",
    "sine_data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([20000, 5, 50])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "sine_data = np.transpose(sine_data, (0,2,1))\n",
    "sine_data = torch.tensor(sine_data)\n",
    "sine_data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([15000, 5, 50]) torch.Size([5000, 5, 50])\n"
     ]
    }
   ],
   "source": [
    "indices = np.arange(len(sine_data))\n",
    "np.random.shuffle(indices)\n",
    "sine_train = sine_data[indices[:int(len(sine_data)*0.75)]] # 75% train\n",
    "sine_test = sine_data[indices[int(len(sine_data)*0.75):]] # 25% test\n",
    "print(sine_train.shape, sine_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(sine_train, \"sine_train_5dim.pt\")\n",
    "torch.save(sine_test, \"sine_test_5dim.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([15000, 5, 50])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.load(\"sine_train_5dim.pt\").shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5000, 5, 50])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.load(\"sine_test_5dim.pt\").shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def real_data_loading (data_name, seq_len):\n",
    "  \"\"\"\n",
    "  Load and preprocess real-world datasets. \n",
    "  Adapted from TimeGAN source code for sake of consistency.\n",
    "  \n",
    "  Args:\n",
    "    - data_name: stock or energy\n",
    "    - seq_len: sequence length\n",
    "    \n",
    "  Returns:\n",
    "    - data: preprocessed data.\n",
    "  \"\"\"  \n",
    "  assert data_name in ['stock','energy']\n",
    "  \n",
    "  if data_name == 'stock':\n",
    "    ori_data = np.loadtxt('stock_data.csv', delimiter = \",\",skiprows = 1)\n",
    "  elif data_name == 'energy':\n",
    "    ori_data = np.loadtxt('energy_data.csv', delimiter = \",\",skiprows = 1)\n",
    "        \n",
    "  # Flip the data to make chronological data\n",
    "  ori_data = ori_data[::-1]\n",
    "    \n",
    "  # Preprocess the dataset\n",
    "  temp_data = []    \n",
    "  # Cut data by sequence length\n",
    "  for i in range(0, len(ori_data) - seq_len):\n",
    "    _x = ori_data[i:i + seq_len]\n",
    "    temp_data.append(_x)\n",
    "        \n",
    "  # Mix the datasets (to make it similar to i.i.d)\n",
    "  idx = np.random.permutation(len(temp_data))    \n",
    "  data = []\n",
    "  for i in range(len(temp_data)):\n",
    "    data.append(temp_data[idx[i]])\n",
    "    \n",
    "  return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "stocks = real_data_loading('stock', 24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3661, 6, 24])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stocks = np.asarray(stocks)\n",
    "stocks = torch.tensor(stocks)\n",
    "stocks = stocks.permute(0,2,1)\n",
    "stocks.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_test_split(data, train_percent=0.75):\n",
    "    indices = np.arange(len(data))\n",
    "    np.random.shuffle(indices)\n",
    "    train = data[indices[:int(len(data)*train_percent)]]\n",
    "    test = data[indices[int(len(data)*train_percent):]]\n",
    "    return train, test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "stocks_train, stocks_test = train_test_split(stocks, train_percent=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([2928, 6, 24]), torch.Size([733, 6, 24]))"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stocks_train.shape, stocks_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(stocks_train, \"stocks_train.pt\")\n",
    "torch.save(stocks_test, \"stocks_test.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "energy = real_data_loading('energy', 24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([19711, 28, 24])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "energy = torch.tensor(np.asarray(energy))\n",
    "energy = energy.permute(0,2,1)\n",
    "energy.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_train, energy_test = train_test_split(energy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([14783, 28, 24]), torch.Size([4928, 28, 24]))"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "energy_train.shape, energy_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(energy_train, \"energy_train.pt\")\n",
    "torch.save(energy_test, \"energy_test.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, save a .npy version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sine train (15000, 5, 50)\n",
      "sine test (5000, 5, 50)\n",
      "stocks train (2928, 6, 24)\n",
      "stocks test (733, 6, 24)\n",
      "energy train (14783, 28, 24)\n",
      "energy test (4928, 28, 24)\n"
     ]
    }
   ],
   "source": [
    "for name in [\"sine\", \"stocks\", \"energy\"]:\n",
    "    for split in [\"train\", \"test\"]:\n",
    "        data = torch.load(f\"{name}_{split}.pt\").numpy()\n",
    "        np.save(f\"{name}_{split}.npy\", data)\n",
    "        print(name, split, data.shape)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gen",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
