01_2_generate_data_hmm_gauss.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate data from HMM model\n",
    "Running this notebook will:\n",
    "\n",
    "- create 9 Hidden Markov Models\n",
    "- show and save plots of some data generated from these models\n",
    "- generate more X and Z samples and create a dataframe with true labels \n",
    "- save used models and generated data to pickle file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from hmmlearn import hmm\n",
    "import matplotlib.pyplot as plt\n",
    "import pickle\n",
    "from time import time\n",
    "import os\n",
    "\n",
    "DATA_DIR = \"./data_HMM\"\n",
    "GENERATING_MODEL_NAME = \"HMM\"\n",
    "N_TRAIN_SAMPLES_PER_MODEL = 25  # 9 models\n",
    "N_TEST_SAMPLES_PER_MODEL = 100\n",
    "MIN_SAMPLE_LEN = 128\n",
    "MAX_SAMPLE_LEN = 128\n",
    "\n",
    "DESCRIPTION = \"\"\n",
    "UNBALANCED = True\n",
    "unb = \"_unbal\" if UNBALANCED else \"\"\n",
    "FILENAME_SUFFIX = f\"{GENERATING_MODEL_NAME}_data_ntrain{N_TRAIN_SAMPLES_PER_MODEL}_min{MIN_SAMPLE_LEN}_max{MAX_SAMPLE_LEN}{unb}\"\n",
    "\n",
    "time_index = int(time())\n",
    "\n",
    "if not os.path.exists(DATA_DIR):\n",
    "    os.makedirs(DATA_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def define_HMM(n_components: int, covariance_type: str, means_: np.ndarray, \n",
    "               covars_: np.ndarray, startprob_: np.ndarray, transmat_: np.ndarray) -> hmm.GaussianHMM:\n",
    "    \"\"\"Create an instance of hmm.GaussianHMM and set all variables necessary to generate data.\"\"\"\n",
    "    # if sum(startprob_) != 1:\n",
    "    #     print(f\"{startprob_ = }\")\n",
    "    #     raise ValueError(f\"Sum of startprob_ (= {sum(startprob_)}) must be 1.\")\n",
    "    # for i, row in enumerate(transmat_):\n",
    "    #     if sum(row) != 1:\n",
    "    #         print(f\"{row = }\")\n",
    "    #         raise ValueError(f\"Sum of row {i} of transition matrix (= {sum(row)}) must be 1.\")\n",
    "\n",
    "    model = hmm.GaussianHMM(n_components= n_components, covariance_type= covariance_type) \n",
    "    model.means_ = means_\n",
    "    model.covars_ = covars_\n",
    "    model.startprob_ = startprob_\n",
    "    model.transmat_ = transmat_\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model1 = define_HMM(\n",
    "    n_components = 2, covariance_type = \"diag\", \n",
    "    means_ = np.array([[0.5],[100]]), \n",
    "    covars_ = np.array([[0.001],[10]]),\n",
    "    startprob_ = np.array([0.8, 0.2]), \n",
    "    transmat_ = np.array([[0.9, 0.1],\n",
    "                          [0.5, 0.5]])\n",
    ")\n",
    "\n",
    "model2 = define_HMM(\n",
    "    n_components = 2, covariance_type = \"diag\",\n",
    "    means_ = np.array([[0.5],[98]]),\n",
    "    covars_ = np.array([[0.001],[4]]),\n",
    "    startprob_ = np.array([0.7, 0.3]),\n",
    "    transmat_ = np.array([[0.8, 0.2],\n",
    "                          [0.6, 0.4]])\n",
    ")\n",
    "\n",
    "model3 = define_HMM(\n",
    "    n_components = 2, covariance_type = \"diag\",\n",
    "    means_ = np.array([[0.5],[50]]),\n",
    "    covars_ = np.array([[0.001],[5]]),\n",
    "    startprob_ = np.array([0.3, 0.7]),\n",
    "    transmat_ = np.array([[0.6, 0.4],\n",
    "                          [0.2, 0.8]])\n",
    ")\n",
    "\n",
    "model4 = define_HMM(\n",
    "    n_components= 1, covariance_type= \"diag\",\n",
    "    means_ = np.array([[100]]),\n",
    "    covars_ = np.array([[30]]),\n",
    "    startprob_ = np.array([1]),\n",
    "    transmat_ = np.array([[1]]),\n",
    ")\n",
    "\n",
    "model5 = define_HMM(\n",
    "    n_components= 1, covariance_type= \"diag\",\n",
    "    means_ = np.array([[50]]),\n",
    "    covars_ = np.array([[1]]),\n",
    "    startprob_ = np.array([1]),\n",
    "    transmat_ = np.array([[1]]),\n",
    ")\n",
    "\n",
    "model6 = define_HMM(\n",
    "    n_components= 1, covariance_type= \"diag\",\n",
    "    means_ = np.array([[100]]),\n",
    "    covars_ = np.array([[1]]),\n",
    "    startprob_ = np.array([1]),\n",
    "    transmat_ = np.array([[1]]),\n",
    ")\n",
    "\n",
    "model7 = define_HMM(\n",
    "    n_components= 5, covariance_type= \"diag\",\n",
    "    means_ = np.array([[0.1],[20],[30],[200],[300]]),\n",
    "    covars_ = np.array([[0.001],[0.2],[0.3],[0.5],[0.5]]),\n",
    "    startprob_ = np.array([0.6, 0.1, 0.1, 0.1, 0.1]),\n",
    "    transmat_ = np.array([[0.6, 0.0, 0.1, 0.1, 0.2],\n",
    "                          [0.1, 0.8, 0.05, 0.04, 0.01],\n",
    "                          [0.1, 0.05, 0.8, 0.04, 0.01],    \n",
    "                          [0.05, 0.2, 0.02, 0.7, 0.03], \n",
    "                          [0.3, 0.03, 0.03, 0.04, 0.6]\n",
    "                         ])\n",
    ")\n",
    "\n",
    "model8 = define_HMM(\n",
    "    n_components= 3, covariance_type= \"diag\",\n",
    "    means_ = np.array([[0.1],[50], [100]]),\n",
    "    covars_ = np.array([[0.001],[5], [5]]),\n",
    "    startprob_ = np.array([0.1, 0.7, 0.2]),\n",
    "    transmat_ = np.array([[0.6, 0.4, 0.0],\n",
    "                          [0.1, 0.8, 0.1],\n",
    "                          [0.7, 0.01, 0.29]\n",
    "                         ])\n",
    ")\n",
    "\n",
    "model9 = define_HMM(\n",
    "    n_components= 3, covariance_type= \"diag\",\n",
    "    means_ = np.array([[1],[30], [45]]),\n",
    "    covars_ = np.array([[0.001],[5], [5]]),\n",
    "    startprob_ = np.array([0.1, 0.7, 0.2]),\n",
    "    transmat_ = np.array([[0.5, 0.4, 0.1],\n",
    "                          [0.1, 0.6, 0.3],\n",
    "                          [0.2, 0.51, 0.29]\n",
    "                         ])\n",
    ")\n",
    "\n",
    "models_lst = [model1, model2, model3, model4, model5, model6, model7, model8, model9]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Show plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 100\n",
    "plt.rcParams[\"figure.figsize\"] = (40, 16)\n",
    "plt.rcParams['font.size'] = 20\n",
    "fig, axes = plt.subplots(3,3, sharey=\"all\")\n",
    "plt.suptitle(f\"Samples of length {n} (HMM)\")\n",
    "for i, ax in enumerate(axes.reshape(-1)):\n",
    "    model = models_lst[i]\n",
    "    try:\n",
    "        X,Z = model.sample(n)\n",
    "    except Exception as e:\n",
    "        print(f\"Error in model {i}: {e}\")\n",
    "        continue\n",
    "        \n",
    "    means = [model.means_[hidden_state] for hidden_state in Z]\n",
    "    ax.plot(X, color = \"black\")\n",
    "    ax.scatter(np.array(range(len(Z))), means, color = \"red\", s = 15)\n",
    "    ax.set_title(f\"models_lst[{i}]: {model.n_components =}\")\n",
    "    ax.grid()\n",
    "    \n",
    "plt.savefig(f\"./plots/models_{GENERATING_MODEL_NAME}_{time_index}.jpg\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate all data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "WEIGHTS =  [0.3, 0.05, 0.3, 0.05, 0.05, 0.1, 0.01, 0.03, 0.02]\n",
    "\n",
    "all_X_samples = []\n",
    "all_Z_samples = []\n",
    "true_labels = []\n",
    "sample_lengths = []\n",
    "sample_ids = []\n",
    "train_samples_ids = []\n",
    "test_samples_ids = []\n",
    "possible_lenghts = range(MIN_SAMPLE_LEN, MAX_SAMPLE_LEN + 1)\n",
    "\n",
    "total_train = len(models_lst)* N_TRAIN_SAMPLES_PER_MODEL\n",
    "total_test = len(models_lst)* N_TEST_SAMPLES_PER_MODEL\n",
    "id = 0\n",
    "for i in range(len(models_lst)):\n",
    "    model = models_lst[i]\n",
    "    curr_no_train = int(np.ceil(total_train * WEIGHTS[i])) if UNBALANCED else N_TRAIN_SAMPLES_PER_MODEL\n",
    "    curr_no_test = N_TEST_SAMPLES_PER_MODEL # test are balanced always\n",
    "    print(f\"Generating {curr_no_train} train and {curr_no_test} test samples from model {i}.\")\n",
    "    for _ in range(curr_no_train):\n",
    "        sample_len = np.random.choice(possible_lenghts)\n",
    "        X,Z = model.sample(sample_len)\n",
    "        all_X_samples.append(X)\n",
    "        all_Z_samples.append(Z)\n",
    "        sample_lengths.append(sample_len)\n",
    "        true_labels.append(i)\n",
    "        sample_ids.append(id)\n",
    "        train_samples_ids.append(id)\n",
    "        id += 1\n",
    "        \n",
    "    for _ in range(curr_no_test):\n",
    "        sample_len = np.random.choice(possible_lenghts)\n",
    "        X,Z = model.sample(sample_len)\n",
    "        all_X_samples.append(X)\n",
    "        all_Z_samples.append(Z)\n",
    "        sample_lengths.append(sample_len)\n",
    "        true_labels.append(i)\n",
    "        sample_ids.append(id)\n",
    "        test_samples_ids.append(id)\n",
    "        id += 1\n",
    "\n",
    "labels_df = pd.DataFrame({\n",
    "    \"sample_id\": sample_ids,\n",
    "    \"true_label\": true_labels,\n",
    "    \"sample_len\": sample_lengths\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "indices_split = {\n",
    "    \"train_samples_ids\": train_samples_ids, \n",
    "    \"test_samples_ids\": test_samples_ids}\n",
    "\n",
    "indices_splits_lst = [indices_split]\n",
    "print(len(indices_split[\"train_samples_ids\"]), len(indices_split[\"test_samples_ids\"]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_filename = f\"{FILENAME_SUFFIX}_{time_index}.pkl\"\n",
    "output_path = f\"{DATA_DIR}/{data_filename}\"\n",
    "\n",
    "metadata = {\n",
    "    \"N_TRAIN_SAMPLES_PER_MODEL\": N_TRAIN_SAMPLES_PER_MODEL,\n",
    "    \"N_TEST_SAMPLES_PER_MODEL\": N_TEST_SAMPLES_PER_MODEL,\n",
    "    \"MAX_SAMPLE_LEN\": MAX_SAMPLE_LEN,\n",
    "    \"MIN_SAMPLE_LEN\": MIN_SAMPLE_LEN\n",
    "}\n",
    "\n",
    "data = {\n",
    "    'generating_model': GENERATING_MODEL_NAME,\n",
    "    'data_filename': data_filename,\n",
    "    'time_index': time_index,\n",
    "    'models_lst': models_lst,\n",
    "    'labels_df': labels_df,\n",
    "    'all_X_samples': all_X_samples,\n",
    "    'all_Z_samples': all_Z_samples,\n",
    "    'indices_splits_lst': indices_splits_lst,\n",
    "    'metadata': metadata,\n",
    "    'description': DESCRIPTION\n",
    "}\n",
    "\n",
    "with open(output_path, 'wb') as f:\n",
    "    pickle.dump(data, f)\n",
    "    print(f\"Data saved to {output_path}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examine labels_dtf and samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_X_samples[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels_df"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "824dc94f4a21fb2b9b063d5374c6a5be71ea0704887dd3ffa4e3d703212d1775"
  },
  "kernelspec": {
   "display_name": "Python 3.9.12 ('hmm_dtw')",
   "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.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}