Split_study_perod_data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import datetime as dt\n",
    "import os\n",
    "\n",
    "os.chdir(\"./data/stocks_data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>A</th>\n",
       "      <th>AAL</th>\n",
       "      <th>AAP</th>\n",
       "      <th>AAPL</th>\n",
       "      <th>ABBV</th>\n",
       "      <th>ABMD</th>\n",
       "      <th>ABT</th>\n",
       "      <th>ACGL</th>\n",
       "      <th>ACN</th>\n",
       "      <th>ACS</th>\n",
       "      <th>...</th>\n",
       "      <th>XEL</th>\n",
       "      <th>XOM</th>\n",
       "      <th>XRAY</th>\n",
       "      <th>XRX</th>\n",
       "      <th>XYL</th>\n",
       "      <th>YUM</th>\n",
       "      <th>ZBH</th>\n",
       "      <th>ZBRA</th>\n",
       "      <th>ZION</th>\n",
       "      <th>ZTS</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Date</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1996-01-02</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.242501</td>\n",
       "      <td>NaN</td>\n",
       "      <td>7.250</td>\n",
       "      <td>4.985257</td>\n",
       "      <td>2.569444</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>7.238386</td>\n",
       "      <td>8.375014</td>\n",
       "      <td>5.400020</td>\n",
       "      <td>31.252378</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>15.388889</td>\n",
       "      <td>11.960916</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1996-01-03</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.242501</td>\n",
       "      <td>NaN</td>\n",
       "      <td>7.375</td>\n",
       "      <td>5.000181</td>\n",
       "      <td>2.541667</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>7.329781</td>\n",
       "      <td>8.400988</td>\n",
       "      <td>5.468374</td>\n",
       "      <td>31.223751</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>15.666667</td>\n",
       "      <td>11.696796</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1996-01-04</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.238255</td>\n",
       "      <td>NaN</td>\n",
       "      <td>7.500</td>\n",
       "      <td>5.074811</td>\n",
       "      <td>2.555556</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>7.238386</td>\n",
       "      <td>8.452925</td>\n",
       "      <td>5.451284</td>\n",
       "      <td>30.565504</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>14.000000</td>\n",
       "      <td>11.206285</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1996-01-05</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.258543</td>\n",
       "      <td>NaN</td>\n",
       "      <td>7.500</td>\n",
       "      <td>4.970332</td>\n",
       "      <td>2.500000</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>7.293225</td>\n",
       "      <td>8.660674</td>\n",
       "      <td>5.365843</td>\n",
       "      <td>30.222082</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>13.444444</td>\n",
       "      <td>11.244011</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1996-01-08</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.261373</td>\n",
       "      <td>NaN</td>\n",
       "      <td>7.125</td>\n",
       "      <td>5.059889</td>\n",
       "      <td>2.500000</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>7.256669</td>\n",
       "      <td>8.764545</td>\n",
       "      <td>5.331663</td>\n",
       "      <td>30.164845</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>13.444444</td>\n",
       "      <td>11.130821</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 791 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "             A  AAL  AAP      AAPL  ABBV   ABMD       ABT      ACGL  ACN  ACS  \\\n",
       "Date                                                                            \n",
       "1996-01-02 NaN  NaN  NaN  0.242501   NaN  7.250  4.985257  2.569444  NaN  NaN   \n",
       "1996-01-03 NaN  NaN  NaN  0.242501   NaN  7.375  5.000181  2.541667  NaN  NaN   \n",
       "1996-01-04 NaN  NaN  NaN  0.238255   NaN  7.500  5.074811  2.555556  NaN  NaN   \n",
       "1996-01-05 NaN  NaN  NaN  0.258543   NaN  7.500  4.970332  2.500000  NaN  NaN   \n",
       "1996-01-08 NaN  NaN  NaN  0.261373   NaN  7.125  5.059889  2.500000  NaN  NaN   \n",
       "\n",
       "            ...       XEL       XOM      XRAY        XRX  XYL  YUM  ZBH  \\\n",
       "Date        ...                                                           \n",
       "1996-01-02  ...  7.238386  8.375014  5.400020  31.252378  NaN  NaN  NaN   \n",
       "1996-01-03  ...  7.329781  8.400988  5.468374  31.223751  NaN  NaN  NaN   \n",
       "1996-01-04  ...  7.238386  8.452925  5.451284  30.565504  NaN  NaN  NaN   \n",
       "1996-01-05  ...  7.293225  8.660674  5.365843  30.222082  NaN  NaN  NaN   \n",
       "1996-01-08  ...  7.256669  8.764545  5.331663  30.164845  NaN  NaN  NaN   \n",
       "\n",
       "                 ZBRA       ZION  ZTS  \n",
       "Date                                   \n",
       "1996-01-02  15.388889  11.960916  NaN  \n",
       "1996-01-03  15.666667  11.696796  NaN  \n",
       "1996-01-04  14.000000  11.206285  NaN  \n",
       "1996-01-05  13.444444  11.244011  NaN  \n",
       "1996-01-08  13.444444  11.130821  NaN  \n",
       "\n",
       "[5 rows x 791 columns]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adj_close = pd.read_parquet(\"adj_close.parquet\")\n",
    "return_tommorrow = pd.read_parquet(\"return_tomorrow.parquet\")\n",
    "\n",
    "trade_dates = adj_close.index.to_list()[:-1]\n",
    "\n",
    "adj_close.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7129\n",
      "7000\n"
     ]
    }
   ],
   "source": [
    "# generate all date_ranges\n",
    "print(len(trade_dates))\n",
    "\n",
    "# drop some dates at the begining to make the number of dates dividable by 250\n",
    "n_drop = len(trade_dates) % 250\n",
    "if n_drop != 0:\n",
    "    trade_dates = trade_dates[n_drop:]\n",
    "\n",
    "print(len(trade_dates))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "25\n"
     ]
    }
   ],
   "source": [
    "# generate date_ranges, each with length 1000 and step 250\n",
    "ranges = []\n",
    "for start in range(0, len(trade_dates) - 1000 + 1, 250):\n",
    "    end = start + 1000\n",
    "    ranges.append(trade_dates[start:end])\n",
    "print(len(ranges))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_single_period(dates_range):\n",
    "    period_close = adj_close.loc[dates_range]\n",
    "    period_close = period_close.drop(columns=period_close.columns[period_close.count() < 600])\n",
    "    period_close = period_close.stack().reset_index()\n",
    "    period_close.columns = [\"Date\", \"Ticker\", \"Adj_Close\"]\n",
    "\n",
    "    period_rtn = return_tommorrow.loc[dates_range].stack().reset_index()\n",
    "    period_rtn.columns = [\"Date\", \"Ticker\", \"Return_tomorrow\"]\n",
    "    \n",
    "    period_data = pd.merge(period_close, period_rtn, on=[\"Date\", \"Ticker\"], how=\"left\")\n",
    "    period_data = period_data.dropna(how=\"any\", axis=0)\n",
    "    if period_data.Date.nunique() != len(dates_range):\n",
    "        print(\"Error: date range not match\")\n",
    "    return period_data.sort_values(by=[\"Ticker\", \"Date\"]).reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/25 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 25/25 [00:16<00:00,  1.49it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "save_path = \"./single_period_data\"\n",
    "os.makedirs(save_path, exist_ok=True)\n",
    "\n",
    "for i, period_dates in tqdm(enumerate(ranges), total=len(ranges)):\n",
    "    data = generate_single_period(period_dates)\n",
    "    data.to_parquet(os.path.join(save_path, f\"period_{i}.parquet\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.11.5"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}