01_3_generate_data_hmm_gmm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate data from HMM model\n",
    "This is a copy of the notebook that generates data from GaussianHMM model, but model class is replaced with GMMHMM."
   ]
  },
  {
   "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 = \"./generated_data_set2\"\n",
    "GENERATING_MODEL_NAME = \"HMM_GMM\"\n",
    "N_TRAIN_SAMPLES_PER_MODEL = 30\n",
    "N_TEST_SAMPLES_PER_MODEL = 100\n",
    "MIN_SAMPLE_LEN = 50\n",
    "MAX_SAMPLE_LEN = 150\n",
    "\n",
    "FILENAME_SUFFIX = f\"{GENERATING_MODEL_NAME}_data_ntrain{N_TRAIN_SAMPLES_PER_MODEL}_min{MIN_SAMPLE_LEN}_max{MAX_SAMPLE_LEN}\"\n",
    "DESCRIPTION = \"\"\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_mix: int, weights_: np.ndarray, n_components: int, covariance_type: str, means_: np.ndarray, \n",
    "               covars_: np.ndarray, startprob_: np.ndarray, transmat_: np.ndarray) -> hmm.GMMHMM:\n",
    "    \"\"\"Create an instance of hmm.GMMHMM 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.GMMHMM(n_mix = n_mix, n_components= n_components, covariance_type= covariance_type) \n",
    "    model.means_ = means_\n",
    "    model.weights_ = weights_\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_mix = 2, n_components = 2, covariance_type = \"diag\", \n",
    "    weights_= np.array([[0.5, 0.5],[0.5, 0.5]]), # n_comp, n_mix\n",
    "    means_ = np.array([[[0.5],[0.5]], [[100],[50]]]),   # n_comp, n_mix, n_features\n",
    "    covars_ = np.array([[[0.001],[0.2]], [[10],[10]]]), # n_comp, n_mix, n_features\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_mix = 2, n_components = 2, covariance_type = \"diag\", \n",
    "    weights_= np.array([[0.5, 0.5],[0.5, 0.5]]), # n_comp, n_mix\n",
    "    means_ = np.array([[[0.7],[0.7]], [[90],[60]]]),   # n_comp, n_mix, n_features\n",
    "    covars_ = np.array([[[0.001],[0.2]], [[20],[10]]]), # n_comp, n_mix, n_features\n",
    "    startprob_ = np.array([0.7, 0.3]), \n",
    "    transmat_ = np.array([[0.9, 0.1],\n",
    "                          [0.5, 0.5]])\n",
    ")\n",
    "\n",
    "model3 = define_HMM(\n",
    "    n_mix = 2, n_components = 2, covariance_type = \"diag\", \n",
    "    weights_= np.array([[0.9, 0.1],[0.9, 0.1]]), # n_comp, n_mix\n",
    "    means_ = np.array([[[0.7],[0.7]], [[90],[60]]]),   # n_comp, n_mix, n_features\n",
    "    covars_ = np.array([[[0.001],[0.5]], [[0.001],[2]]]), # n_comp, n_mix, n_features\n",
    "    startprob_ = np.array([0.7, 0.3]), \n",
    "    transmat_ = np.array([[0.8, 0.2],\n",
    "                          [0.5, 0.5]])\n",
    ")\n",
    "\n",
    "model4 = define_HMM(\n",
    "    n_mix = 1, n_components= 1, covariance_type= \"diag\",\n",
    "    weights_ = np.array([[1.0]]),\n",
    "    means_ = np.array([[[100]]]),              # n_comp, n_mix, n_features\n",
    "    covars_ = np.array([[[20]]]),              # n_comp, n_mix, n_features\n",
    "    startprob_ = np.array([1]),\n",
    "    transmat_ = np.array([[1]]),\n",
    ")\n",
    "\n",
    "model5 = define_HMM(\n",
    "    n_mix = 2, n_components= 1, covariance_type= \"diag\",\n",
    "    weights_ = np.array([[0.3, 0.7]]),\n",
    "    means_ = np.array([[[100], [60]]]),              # n_comp, n_mix, n_features\n",
    "    covars_ = np.array([[[20], [5]]]),                      # n_comp, n_mix, n_features\n",
    "    startprob_ = np.array([1]),\n",
    "    transmat_ = np.array([[1]]),\n",
    ")\n",
    "\n",
    "model6 = define_HMM(\n",
    "    n_mix = 3, n_components= 1, covariance_type= \"diag\",\n",
    "    weights_ = np.array([[0.3, 0.3, 0.4]]),\n",
    "    means_ = np.array([[[100], [60], [5]]]),              # n_comp, n_mix, n_features\n",
    "    covars_ = np.array([[[20], [5], [0.001]]]),                      # n_comp, n_mix, n_features\n",
    "    startprob_ = np.array([1]),\n",
    "    transmat_ = np.array([[1]]),\n",
    ")\n",
    "\n",
    "model7 = define_HMM(\n",
    "    n_mix = 1, n_components= 5, covariance_type= \"diag\",\n",
    "    weights_ = np.array([[1.0],[1.0],[1.0],[1.0],[1.0]]),\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_mix = 2, n_components= 3, covariance_type= \"diag\",\n",
    "    weights_ = np.array([[0.2, 0.8], [0.2, 0.8], [0.5, 0.5]]),\n",
    "    means_ = np.array([[[0.1],[0.2]],[[50],[100]],[[100],[10]]]),\n",
    "    covars_ = np.array([[[0.001],[0.001]],[[5],[10]], [[5],[0.001]]]),\n",
    "    startprob_ = np.array([0.2, 0.6, 0.2]),\n",
    "    transmat_ = np.array([[0.6, 0.4, 0.0],\n",
    "                          [0.8, 0.1, 0.1],\n",
    "                          [0.7, 0.01, 0.29]\n",
    "                         ])\n",
    ")\n",
    "\n",
    "model9 = define_HMM(\n",
    "    n_mix = 1, n_components= 5, covariance_type= \"diag\",\n",
    "    weights_ = np.array([[1.0],[1.0],[1.0],[1.0],[1.0]]),\n",
    "    means_ = np.array([[[0.1]],[[0.1]],[[30]],[[300]],[[400]]]),\n",
    "    covars_ = np.array([[[0.001]],[[0.2]],[[5]],[[0.001]],[[20]]]),\n",
    "    startprob_ = np.array([0.6, 0.2, 0.0, 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.8, 0.03, 0.03, 0.04, 0.1]\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 = 50\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} ({GENERATING_MODEL_NAME})\")\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 =}, {model.n_mix = }\")\n",
    "    ax.grid()\n",
    "    \n",
    "plt.savefig(f\"./plots/models_{GENERATING_MODEL_NAME}_{time_index}.jpg\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model1.weights_.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate all data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_data(models_lst: list[hmm.BaseHMM], samples_per_model: int, max_sample_len: int, min_sample_len: int) -> tuple:\n",
    "    \"\"\"Generate data from list of hmm.BaseHMM instances.\n",
    "    Returns:\n",
    "        - labels_df: pd.DataFrame | sample_id | true_label | sample_len\n",
    "        - all_X_samples: a list with generated X samples\n",
    "        - all_Z_samples: a list with generated hidden states\n",
    "    \"\"\"\n",
    "    all_X_samples = []\n",
    "    all_Z_samples = []\n",
    "    true_labels = []\n",
    "    sample_lengths = []\n",
    "    sample_ids = list(range(samples_per_model*len(models_lst)))\n",
    "    len_step = (max_sample_len-min_sample_len)//10\n",
    "    possible_lenghts = range(min_sample_len, max_sample_len + 1, len_step)\n",
    "    \n",
    "    for i in range(len(models_lst)):\n",
    "        model = models_lst[i]\n",
    "        for j in range(samples_per_model):\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",
    "\n",
    "    labels_df = pd.DataFrame({\n",
    "        \"sample_id\": sample_ids,\n",
    "        \"true_label\": true_labels,\n",
    "        \"sample_len\": sample_lengths\n",
    "    })\n",
    "\n",
    "    return labels_df, all_X_samples, all_Z_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_per_model = N_TRAIN_SAMPLES_PER_MODEL + N_TEST_SAMPLES_PER_MODEL\n",
    "labels_df, all_X_samples, all_Z_samples = generate_data(models_lst, samples_per_model, MAX_SAMPLE_LEN, MIN_SAMPLE_LEN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_samples_ids = []\n",
    "test_samples_ids = []\n",
    "for label, sub_df in labels_df.groupby(\"true_label\"):\n",
    "    train_ids = sub_df.sample(N_TRAIN_SAMPLES_PER_MODEL).index.values\n",
    "    test_ids = sub_df.drop(train_ids).index.values\n",
    "    if len(test_ids) != N_TEST_SAMPLES_PER_MODEL:\n",
    "        raise Exception(f\"len(test_ids) {len(test_ids)} != N_TEST_SAMPLES_PER_MODEL {N_TEST_SAMPLES_PER_MODEL}\")\n",
    "    train_samples_ids.extend(train_ids)\n",
    "    test_samples_ids.extend(test_ids)\n",
    "    \n",
    "\n",
    "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}.\")"
   ]
  }
 ],
 "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
}