Code19_Longitudinal-biological-BMI-LMM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multiomics BMI Paper — Longitudinal Change of Biological BMI during the Arivale Program"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***by Kengo Watanabe***  \n",
    "\n",
    "This Jupyter Notebook (with Python 3 kernel) estimated longitudinal changes of the measured and omics-inferred BMIs during the Arivale program, using linear mixed model (LMM) with random effects for each participant.  \n",
    "\n",
    "Input files:  \n",
    "* Arivale logitudinal BMI predictions: 220805_Multiomics-BMI-NatMed1stRevision_BMI-longitudinal-LASSO_\\[MetBMI/ProtBMI/ChemBMI\\]-\\[BothSex/FemaleMale\\].tsv  \n",
    "* Arivale time-series BMI and blood omics: 210104_Biological-BMI-paper_data-cleaning-BMI-omics_time-series-\\[bmiDF/combiDF\\]-without-imputation_final-cohort.tsv  \n",
    "\n",
    "Output figures and tables:  \n",
    "* Figure 5  \n",
    "\n",
    "Original notebook (memo for my future tracing):  \n",
    "* dalek:\\[JupyterLab HOME\\]/220621_Multiomics-BMI-NatMedRevision/220806_Multiomics-BMI-NatMed1stRevision_BMI-longitudinal-LMM-ver2.ipynb  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# packages in environment at /opt/conda/envs/arivale-py3:\n",
      "#\n",
      "# Name                    Version                   Build  Channel\n",
      "_libgcc_mutex             0.1                 conda_forge    conda-forge\n",
      "_openmp_mutex             4.5                       1_gnu    conda-forge\n",
      "analytics                 0.1                      pypi_0    pypi\n",
      "argon2-cffi               21.1.0           py39h3811e60_0    conda-forge\n",
      "arivale-data-interface    0.1.0                    pypi_0    pypi\n",
      "async_generator           1.10                       py_0    conda-forge\n",
      "atk-1.0                   2.36.0               h3371d22_4    conda-forge\n",
      "attrs                     21.2.0             pyhd8ed1ab_0    conda-forge\n",
      "backcall                  0.2.0              pyh9f0ad1d_0    conda-forge\n",
      "backports                 1.0                        py_2    conda-forge\n",
      "backports.functools_lru_cache 1.6.4              pyhd8ed1ab_0    conda-forge\n",
      "biopython                 1.79             py39h3811e60_0    conda-forge\n",
      "bleach                    4.1.0              pyhd8ed1ab_0    conda-forge\n",
      "bokeh                     2.4.1            py39hf3d152e_1    conda-forge\n",
      "boto                      2.49.0                     py_0    conda-forge\n",
      "boto3                     1.19.2                   pypi_0    pypi\n",
      "botocore                  1.22.2                   pypi_0    pypi\n",
      "brotlipy                  0.7.0           py39h3811e60_1001    conda-forge\n",
      "bzip2                     1.0.8                h7f98852_4    conda-forge\n",
      "c-ares                    1.17.2               h7f98852_0    conda-forge\n",
      "ca-certificates           2022.6.15            ha878542_0    conda-forge\n",
      "cachecontrol              0.12.6                     py_0    conda-forge\n",
      "cairo                     1.16.0            h6cf1ce9_1008    conda-forge\n",
      "certifi                   2022.6.15        py39hf3d152e_0    conda-forge\n",
      "cffi                      1.14.6           py39h4bc2ebd_1    conda-forge\n",
      "chardet                   4.0.0            py39hf3d152e_1    conda-forge\n",
      "charset-normalizer        2.0.0              pyhd8ed1ab_0    conda-forge\n",
      "click                     8.0.3            py39hf3d152e_0    conda-forge\n",
      "colorama                  0.4.4              pyh9f0ad1d_0    conda-forge\n",
      "cryptography              35.0.0           py39h95dcef6_1    conda-forge\n",
      "cycler                    0.10.0                     py_2    conda-forge\n",
      "cython                    0.29.24          py39he80948d_0    conda-forge\n",
      "debugpy                   1.4.1            py39he80948d_0    conda-forge\n",
      "decorator                 4.4.2                      py_0    conda-forge\n",
      "defusedxml                0.7.1              pyhd8ed1ab_0    conda-forge\n",
      "english                   2020.7.0           pyhd8ed1ab_0    conda-forge\n",
      "entrypoints               0.3             py39hde42818_1002    conda-forge\n",
      "et_xmlfile                1.0.1                   py_1001    conda-forge\n",
      "expat                     2.4.1                h9c3ff4c_0    conda-forge\n",
      "font-ttf-dejavu-sans-mono 2.37                 hab24e00_0    conda-forge\n",
      "font-ttf-inconsolata      3.000                h77eed37_0    conda-forge\n",
      "font-ttf-source-code-pro  2.038                h77eed37_0    conda-forge\n",
      "font-ttf-ubuntu           0.83                 hab24e00_0    conda-forge\n",
      "fontconfig                2.13.1            hba837de_1005    conda-forge\n",
      "fonts-conda-ecosystem     1                             0    conda-forge\n",
      "fonts-conda-forge         1                             0    conda-forge\n",
      "freetype                  2.10.4               h0708190_1    conda-forge\n",
      "fribidi                   1.0.10               h36c2ea0_0    conda-forge\n",
      "future                    0.18.2           py39hf3d152e_3    conda-forge\n",
      "gdk-pixbuf                2.42.6               h04a7f16_0    conda-forge\n",
      "gettext                   0.19.8.1          h73d1719_1008    conda-forge\n",
      "giflib                    5.2.1                h36c2ea0_2    conda-forge\n",
      "graphite2                 1.3.13            h58526e2_1001    conda-forge\n",
      "graphviz                  2.49.1               h85b4f2f_0    conda-forge\n",
      "gtk2                      2.24.33              h539f30e_1    conda-forge\n",
      "gts                       0.7.6                h64030ff_2    conda-forge\n",
      "harfbuzz                  3.0.0                h83ec7ef_1    conda-forge\n",
      "hdmedians                 0.14.2           py39hce5d2b2_0    conda-forge\n",
      "icu                       68.1                 h58526e2_0    conda-forge\n",
      "idna                      3.1                pyhd3deb0d_0    conda-forge\n",
      "importlib-metadata        4.8.1            py39hf3d152e_0    conda-forge\n",
      "iniconfig                 1.1.1              pyh9f0ad1d_0    conda-forge\n",
      "interchange               2021.0.4           pyhd8ed1ab_0    conda-forge\n",
      "ipykernel                 6.4.2            py39hef51801_0    conda-forge\n",
      "ipython                   7.28.0           py39hef51801_0    conda-forge\n",
      "ipython_genutils          0.2.0                      py_1    conda-forge\n",
      "ipywidgets                7.6.5              pyhd8ed1ab_0    conda-forge\n",
      "jbig                      2.1               h7f98852_2003    conda-forge\n",
      "jedi                      0.18.0           py39hf3d152e_2    conda-forge\n",
      "jinja2                    3.0.2              pyhd8ed1ab_0    conda-forge\n",
      "jmespath                  0.10.0                   pypi_0    pypi\n",
      "joblib                    1.1.0              pyhd8ed1ab_0    conda-forge\n",
      "jpeg                      9d                   h36c2ea0_0    conda-forge\n",
      "jsonschema                4.1.2              pyhd8ed1ab_0    conda-forge\n",
      "jupyter_client            7.0.6              pyhd8ed1ab_0    conda-forge\n",
      "jupyter_core              4.8.1            py39hf3d152e_0    conda-forge\n",
      "jupyterlab_pygments       0.1.2              pyh9f0ad1d_0    conda-forge\n",
      "jupyterlab_widgets        1.0.2              pyhd8ed1ab_0    conda-forge\n",
      "kiwisolver                1.3.2            py39h1a9c180_0    conda-forge\n",
      "krb5                      1.19.2               hcc1bbae_2    conda-forge\n",
      "lcms2                     2.12                 hddcbb42_0    conda-forge\n",
      "ld_impl_linux-64          2.36.1               hea4e1c9_2    conda-forge\n",
      "lerc                      2.2.1                h9c3ff4c_0    conda-forge\n",
      "libblas                   3.9.0           12_linux64_openblas    conda-forge\n",
      "libcblas                  3.9.0           12_linux64_openblas    conda-forge\n",
      "libcurl                   7.79.1               h2574ce0_1    conda-forge\n",
      "libdeflate                1.7                  h7f98852_5    conda-forge\n",
      "libedit                   3.1.20191231         he28a2e2_2    conda-forge\n",
      "libev                     4.33                 h516909a_1    conda-forge\n",
      "libffi                    3.4.2                h9c3ff4c_4    conda-forge\n",
      "libgcc-ng                 11.2.0              h1d223b6_11    conda-forge\n",
      "libgd                     2.3.3                h6ad9fb6_0    conda-forge\n",
      "libgfortran-ng            11.2.0              h69a702a_11    conda-forge\n",
      "libgfortran5              11.2.0              h5c6108e_11    conda-forge\n",
      "libglib                   2.70.0               h174f98d_1    conda-forge\n",
      "libgomp                   11.2.0              h1d223b6_11    conda-forge\n",
      "libiconv                  1.16                 h516909a_0    conda-forge\n",
      "liblapack                 3.9.0           12_linux64_openblas    conda-forge\n",
      "libnghttp2                1.43.0               h812cca2_1    conda-forge\n",
      "libopenblas               0.3.18          pthreads_h8fe5266_0    conda-forge\n",
      "libpng                    1.6.37               h21135ba_2    conda-forge\n",
      "libpq                     13.3                 hd57d9b9_1    conda-forge\n",
      "librsvg                   2.52.2               hc3c00ef_0    conda-forge\n",
      "libsodium                 1.0.18               h36c2ea0_1    conda-forge\n",
      "libssh2                   1.10.0               ha56f1ee_2    conda-forge\n",
      "libstdcxx-ng              11.2.0              he4da1e4_11    conda-forge\n",
      "libtiff                   4.3.0                hf544144_1    conda-forge\n",
      "libtool                   2.4.6             h9c3ff4c_1008    conda-forge\n",
      "libuuid                   2.32.1            h7f98852_1000    conda-forge\n",
      "libwebp                   1.2.1                h3452ae3_0    conda-forge\n",
      "libwebp-base              1.2.1                h7f98852_0    conda-forge\n",
      "libxcb                    1.13              h7f98852_1003    conda-forge\n",
      "libxml2                   2.9.12               h72842e0_0    conda-forge\n",
      "libxslt                   1.1.33               h15afd5d_2    conda-forge\n",
      "libzlib                   1.2.11            h36c2ea0_1013    conda-forge\n",
      "lockfile                  0.12.2                     py_1    conda-forge\n",
      "lxml                      4.6.3            py39h107f48f_0    conda-forge\n",
      "lz4-c                     1.9.3                h9c3ff4c_1    conda-forge\n",
      "markupsafe                2.0.1            py39h3811e60_0    conda-forge\n",
      "matplotlib-base           3.4.3            py39h2fa2bec_1    conda-forge\n",
      "matplotlib-inline         0.1.3              pyhd8ed1ab_0    conda-forge\n",
      "matplotlib-venn           0.11.7             pyhd8ed1ab_0    conda-forge\n",
      "missingpy                 0.2.0                    pypi_0    pypi\n",
      "mistune                   0.8.4           py39h3811e60_1004    conda-forge\n",
      "more-itertools            8.10.0             pyhd8ed1ab_0    conda-forge\n",
      "mscorefonts               0.0.1                         3    conda-forge\n",
      "msgpack-python            1.0.2            py39h1a9c180_1    conda-forge\n",
      "natsort                   7.1.1              pyhd8ed1ab_0    conda-forge\n",
      "nbclient                  0.5.4              pyhd8ed1ab_0    conda-forge\n",
      "nbconvert                 6.2.0            py39hf3d152e_0    conda-forge\n",
      "nbformat                  5.1.3              pyhd8ed1ab_0    conda-forge\n",
      "ncurses                   6.2                  h58526e2_4    conda-forge\n",
      "neo4j-python-driver       4.3.7              pyhd8ed1ab_0    conda-forge\n",
      "neobolt                   1.7.17           py39h3811e60_2    conda-forge\n",
      "neotime                   1.7.4                      py_0    conda-forge\n",
      "nest-asyncio              1.5.1              pyhd8ed1ab_0    conda-forge\n",
      "networkx                  2.6.3              pyhd8ed1ab_0    conda-forge\n",
      "notebook                  6.4.5              pyha770c72_0    conda-forge\n",
      "numpy                     1.21.3           py39hdbf815f_0    conda-forge\n",
      "olefile                   0.46               pyh9f0ad1d_1    conda-forge\n",
      "openjpeg                  2.4.0                hb52868f_1    conda-forge\n",
      "openpyxl                  3.0.9              pyhd8ed1ab_0    conda-forge\n",
      "openssl                   1.1.1o               h166bdaf_0    conda-forge\n",
      "packaging                 21.0               pyhd8ed1ab_0    conda-forge\n",
      "pandas                    1.3.4            py39hde0f152_0    conda-forge\n",
      "pandoc                    2.15                 h7f98852_0    conda-forge\n",
      "pandocfilters             1.5.0              pyhd8ed1ab_0    conda-forge\n",
      "pango                     1.48.10              h54213e6_2    conda-forge\n",
      "parso                     0.8.2              pyhd8ed1ab_0    conda-forge\n",
      "patsy                     0.5.2              pyhd8ed1ab_0    conda-forge\n",
      "pcre                      8.45                 h9c3ff4c_0    conda-forge\n",
      "pexpect                   4.8.0              pyh9f0ad1d_2    conda-forge\n",
      "pickleshare               0.7.5           py39hde42818_1002    conda-forge\n",
      "pillow                    8.3.2            py39ha612740_0    conda-forge\n",
      "pip                       21.3.1             pyhd8ed1ab_0    conda-forge\n",
      "pixman                    0.40.0               h36c2ea0_0    conda-forge\n",
      "plotly                    5.3.1              pyhd8ed1ab_0    conda-forge\n",
      "pluggy                    1.0.0            py39hf3d152e_1    conda-forge\n",
      "prometheus_client         0.11.0             pyhd8ed1ab_0    conda-forge\n",
      "prompt-toolkit            3.0.21             pyha770c72_0    conda-forge\n",
      "prompt_toolkit            3.0.21               hd8ed1ab_0    conda-forge\n",
      "psycopg2                  2.9.1            py39h3811e60_0    conda-forge\n",
      "pthread-stubs             0.4               h36c2ea0_1001    conda-forge\n",
      "ptyprocess                0.7.0              pyhd3deb0d_0    conda-forge\n",
      "py                        1.10.0             pyhd3deb0d_0    conda-forge\n",
      "py2neo                    2021.2.3           pyhd8ed1ab_0    conda-forge\n",
      "pycparser                 2.20               pyh9f0ad1d_2    conda-forge\n",
      "pygments                  2.10.0             pyhd8ed1ab_0    conda-forge\n",
      "pygraphviz                1.7              py39h78163bd_0    conda-forge\n",
      "pyopenssl                 21.0.0             pyhd8ed1ab_0    conda-forge\n",
      "pyparsing                 2.4.7              pyh9f0ad1d_0    conda-forge\n",
      "pyrsistent                0.17.3           py39h3811e60_2    conda-forge\n",
      "pysam                     0.17.0           py39h051187c_0    bioconda\n",
      "pysocks                   1.7.1            py39hf3d152e_3    conda-forge\n",
      "pytabix                   0.1              py39h98c8e45_1    bioconda\n",
      "pytest                    6.2.5            py39hf3d152e_0    conda-forge\n",
      "python                    3.9.7           hb7a2778_3_cpython    conda-forge\n",
      "python-dateutil           2.8.2              pyhd8ed1ab_0    conda-forge\n",
      "python-graphviz           0.17               pyhaef67bd_0    conda-forge\n",
      "python_abi                3.9                      2_cp39    conda-forge\n",
      "pytz                      2021.3             pyhd8ed1ab_0    conda-forge\n",
      "pyvcf                     0.6.8           py39hde42818_1002    conda-forge\n",
      "pyyaml                    6.0              py39h3811e60_0    conda-forge\n",
      "pyzmq                     22.3.0           py39h37b5a0c_0    conda-forge\n",
      "readline                  8.1                  h46c0cb4_0    conda-forge\n",
      "requests                  2.26.0             pyhd8ed1ab_0    conda-forge\n",
      "s3transfer                0.5.0                    pypi_0    pypi\n",
      "scikit-bio                0.5.6            py39h16ac069_4    conda-forge\n",
      "scikit-learn              1.0.1            py39h7c5d8c9_0    conda-forge\n",
      "scipy                     1.7.1            py39hee8e79c_0    conda-forge\n",
      "seaborn                   0.11.2               hd8ed1ab_0    conda-forge\n",
      "seaborn-base              0.11.2             pyhd8ed1ab_0    conda-forge\n",
      "send2trash                1.8.0              pyhd8ed1ab_0    conda-forge\n",
      "setuptools                58.2.0           py39hf3d152e_0    conda-forge\n",
      "six                       1.16.0             pyh6c4a22f_0    conda-forge\n",
      "sqlite                    3.36.0               h9cd32fc_2    conda-forge\n",
      "statsmodels               0.13.0           py39hce5d2b2_0    conda-forge\n",
      "tenacity                  8.0.1              pyhd8ed1ab_0    conda-forge\n",
      "terminado                 0.12.1           py39hf3d152e_0    conda-forge\n",
      "testpath                  0.5.0              pyhd8ed1ab_0    conda-forge\n",
      "threadpoolctl             3.0.0              pyh8a188c0_0    conda-forge\n",
      "tk                        8.6.11               h27826a3_1    conda-forge\n",
      "toml                      0.10.2             pyhd8ed1ab_0    conda-forge\n",
      "tornado                   6.1              py39h3811e60_1    conda-forge\n",
      "traitlets                 5.1.0              pyhd8ed1ab_0    conda-forge\n",
      "typing_extensions         3.10.0.2           pyha770c72_0    conda-forge\n",
      "tzdata                    2021e                he74cb21_0    conda-forge\n",
      "urllib3                   1.26.7             pyhd8ed1ab_0    conda-forge\n",
      "venn                      0.1.3                    pypi_0    pypi\n",
      "wcwidth                   0.2.5              pyh9f0ad1d_2    conda-forge\n",
      "webencodings              0.5.1                      py_1    conda-forge\n",
      "wheel                     0.37.0             pyhd8ed1ab_1    conda-forge\n",
      "widgetsnbextension        3.5.1            py39hf3d152e_4    conda-forge\n",
      "wordcloud                 1.8.1                    pypi_0    pypi\n",
      "xlrd                      2.0.1              pyhd8ed1ab_3    conda-forge\n",
      "xorg-kbproto              1.0.7             h7f98852_1002    conda-forge\n",
      "xorg-libice               1.0.10               h7f98852_0    conda-forge\n",
      "xorg-libsm                1.2.3             hd9c2040_1000    conda-forge\n",
      "xorg-libx11               1.7.2                h7f98852_0    conda-forge\n",
      "xorg-libxau               1.0.9                h7f98852_0    conda-forge\n",
      "xorg-libxdmcp             1.1.3                h7f98852_0    conda-forge\n",
      "xorg-libxext              1.3.4                h7f98852_1    conda-forge\n",
      "xorg-libxrender           0.9.10            h7f98852_1003    conda-forge\n",
      "xorg-renderproto          0.11.1            h7f98852_1002    conda-forge\n",
      "xorg-xextproto            7.3.0             h7f98852_1002    conda-forge\n",
      "xorg-xproto               7.0.31            h7f98852_1007    conda-forge\n",
      "xz                        5.2.5                h516909a_1    conda-forge\n",
      "yaml                      0.2.5                h516909a_0    conda-forge\n",
      "zeromq                    4.3.4                h9c3ff4c_1    conda-forge\n",
      "zipp                      3.6.0              pyhd8ed1ab_0    conda-forge\n",
      "zlib                      1.2.11            h36c2ea0_1013    conda-forge\n",
      "zstd                      1.5.0                ha95c52a_0    conda-forge\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scipy.stats as stats\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import seaborn as sns\n",
    "#For Arial font\n",
    "#!conda install -c conda-forge -y mscorefonts\n",
    "##-> The below was also needed in matplotlib 3.4.2\n",
    "#import shutil\n",
    "#import matplotlib\n",
    "#shutil.rmtree(matplotlib.get_cachedir())\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "from IPython.display import display\n",
    "import time\n",
    "\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import statsmodels.api as sm\n",
    "import matplotlib.lines as mlines\n",
    "\n",
    "!conda list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Metabolomics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yvar = 'MetBMI'\n",
    "yvar_color = 'b'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1-1. Prepare the time-series DF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Import time-series biological BMI\n",
    "tempD = {'Sex-mixed '+yvar+' model':'BothSex', 'Sex-stratified '+yvar+' model':'FemaleMale'}\n",
    "tempD1 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    fileDir = './ExportData/'\n",
    "    ipynbName = '220805_Multiomics-BMI-NatMed1stRevision_BMI-longitudinal-LASSO_'\n",
    "    fileName = yvar+'-'+tempD[model_sex]+'.tsv'\n",
    "    tempDF = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    tempD1[model_sex] = tempDF\n",
    "    \n",
    "    print(model_sex)\n",
    "    display(tempDF)\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    print('')\n",
    "\n",
    "tsDF_BS = tempD1['Sex-mixed '+yvar+' model']\n",
    "tsDF_FM = tempD1['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1-2. Prepare DFs for LMM with the linear regresion spline for time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check measurement distribution\n",
    "sns.set(style='ticks', font='Arial', context='notebook')\n",
    "plt.figure(figsize=(4, 3))\n",
    "sns.distplot(tsDF_BS['days_in_program'], color=yvar_color)\n",
    "sns.despine()\n",
    "for knot in [0, 6, 12, 18]:\n",
    "    plt.axvline(x=365.25/12*knot, **{'linestyle':'--', 'color':'gray'})\n",
    "plt.axvspan(xmin=365.25/12*18, xmax=plt.gca().get_xlim()[1], facecolor='gray', alpha=0.2)\n",
    "plt.ylabel('Density')\n",
    "plt.xlabel('Days in program')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD1 = {}\n",
    "tempD2 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    print(model_sex)\n",
    "    tempDF = tempD[model_sex]\n",
    "    print('Time-series DF nrows before filtering:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select measurements for the longitudinal analysis\n",
    "    month_threshold = 18\n",
    "    tempDF = tempDF.loc[tempDF['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    print('Time-series DF nrows after filtering with the taget period:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select the cohort for the longitudinal analysis\n",
    "    ##Select participants who have 2 or more measuremnts in BMI\n",
    "    fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "    ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "    fileName = 'time-series-bmiDF-without-imputation_final-cohort.tsv'\n",
    "    tempDF1 = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF1 = tempDF1.set_index('KeyIndex')\n",
    "    tempDF1 = tempDF1.loc[tempDF1['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    tempS = tempDF1['public_client_id'].value_counts()\n",
    "    tempL1 = tempS.loc[tempS>=2].index.tolist()\n",
    "    ##Select participants who have 2 or more measuremnts in omics\n",
    "    fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "    ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "    fileName = 'time-series-combiDF-without-imputation_final-cohort.tsv'\n",
    "    tempDF1 = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF1 = tempDF1.set_index('KeyIndex')\n",
    "    tempDF1 = tempDF1.loc[tempDF1['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    tempS = tempDF1['public_client_id'].value_counts()\n",
    "    tempL2 = tempS.loc[tempS>=2].index.tolist()\n",
    "    ##Select participants who have 2 or more measuremnts in both BMI and omics\n",
    "    tempL = list(set(tempL1) & set(tempL2))\n",
    "    print('Participants who have 2 or more measurements in both BMI and omics:', len(tempL))\n",
    "    tempDF = tempDF.loc[tempDF['public_client_id'].isin(tempL)]\n",
    "    print('Time-series DF nrows after filtering with the number of measurements:',\n",
    "          len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select the participants who have all covariates (just in case)\n",
    "    tempDF = tempDF.dropna()\n",
    "    print('Time-series DF nrows after filtering with covariates:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Check measurement distribution\n",
    "    sns.set(style='ticks', font='Arial', context='talk')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    sns.distplot(tempDF['days_in_program'], color=yvar_color)\n",
    "    sns.despine()\n",
    "    for knot in [0, 6, 12, 18]:\n",
    "        plt.axvline(x=365.25/12*knot, **{'linestyle':'--', 'color':'gray'})\n",
    "    plt.axvspan(xmin=365.25/12*18, xmax=plt.gca().get_xlim()[1], facecolor='gray', alpha=0.2)\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.show()\n",
    "    \n",
    "    #Check dependent variable distribution\n",
    "    print('Skewness of the dependent variable:', stats.skew(tempDF['log_'+yvar]))\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    sns.distplot(tempDF['log_'+yvar], color=yvar_color)\n",
    "    sns.despine()\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel(yvar+' (log-scale)')\n",
    "    plt.show()\n",
    "    \n",
    "    #Add dummy variables for the linear regression spline for time\n",
    "    nknots = 3\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempL1 = []#2nd piece\n",
    "    tempL2 = []#3rd piece\n",
    "    for days in tempDF['days_in_program'].tolist():\n",
    "        if days < tempL[0]:#[0 month, 6 months)\n",
    "            tempL1.append(0.0)\n",
    "            tempL2.append(0.0)\n",
    "        elif days < tempL[1]:#[6 months, 12 months)\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(0.0)\n",
    "        elif days <= tempL[2]:#[12 months, 18 months]\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(days - tempL[1])\n",
    "        else:\n",
    "            tempL1.append('Error? Check time span setting.')\n",
    "            tempL2.append('Error? Check time span setting.')\n",
    "    tempDF['DaysP2'] = tempL1\n",
    "    tempDF['DaysP3'] = tempL2\n",
    "    \n",
    "    #Standardize numeric independent variables\n",
    "    scaler = StandardScaler(copy=True, with_mean=True, with_std=True)\n",
    "    tempL = ['days_in_program', 'DaysP2', 'DaysP3', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5']\n",
    "    scaler.fit(tempDF[tempL])\n",
    "    tempDF1 = pd.DataFrame(scaler.transform(tempDF[tempL]), index=tempDF.index,\n",
    "                           columns=['DaysZ', 'DaysP2Z', 'DaysP3Z', 'BaseAgeZ', 'PC1Z', 'PC2Z', 'PC3Z', 'PC4Z', 'PC5Z'])\n",
    "    tempDF = pd.merge(tempDF, tempDF1, left_index=True, right_index=True, how='inner')\n",
    "    \n",
    "    #Confirm examples\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    for variable in ['DaysZ', 'BaseAgeZ', 'PC1Z']:\n",
    "        sns.distplot(tempDF[variable], label=variable)\n",
    "    sns.despine()\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel(r'$Z$'+'-score')\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempS = (pd.Series(tempL) - scaler.mean_[0]) / scaler.scale_[0]#Z-score transformation\n",
    "    for knot in tempS.tolist():\n",
    "        plt.axvline(x=knot, **{'linestyle':'--', 'color':'gray'})\n",
    "    plt.legend(bbox_to_anchor=(1, 0.5), loc='center left', borderaxespad=1)\n",
    "    plt.show()\n",
    "    \n",
    "    display(tempDF)\n",
    "    print('')\n",
    "    \n",
    "    tempD1[model_sex] = tempDF\n",
    "    tempD2[model_sex] = scaler#Use later again during predictions\n",
    "\n",
    "#Update\n",
    "tsDF_BS = tempD1['Sex-mixed '+yvar+' model']\n",
    "scaler_BS = tempD2['Sex-mixed '+yvar+' model']\n",
    "tsDF_FM = tempD1['Sex-stratified '+yvar+' model']\n",
    "scaler_FM = tempD2['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1-3. LMM with random intercepts and random slopes for time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Days in the program is fitted as a numeric variable using linear regression splines (cf. Zubair, N. et al. Sci. Rep. 2019).  \n",
    "> ***–> Allows differences in the trajectory of BMI/bBMI change.***\n",
    ">\n",
    "> Random intercepts and random slopes for days in the program are used as random effects (cf. Zubair, N. et al. Sci. Rep. 2019).  \n",
    "> ***–> Allows individual differences in the baseline and longitudinal change over the program.***  \n",
    "> –> Random slope is applied not for the piecewise variable of time but only for the overall variable of time; otherwise, leads to \"LinAlgError: Singular matrix\" due to small sample size per most individuals. Also, it is not so weird to assume the random slope is consistent during the program.  \n",
    ">\n",
    "> ***Eliminate \"Underweight\" class in the baseline obesity class-stratified models.***  \n",
    "> <– Otherwise, leads to \"LinAlgError: Singular matrix\" due to small sample population.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD1 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    print(model_sex)\n",
    "    tempDF = tempD[model_sex]\n",
    "    \n",
    "    tempD2 = {'LMM1':'Overall LMM',\n",
    "              'LMM2':'Normal class LMM',\n",
    "              'LMM3':'Overweight class LMM',\n",
    "              'LMM4':'Obese class LMM'}\n",
    "    tempD3 = {'LMM1':tempDF,\n",
    "              'LMM2':tempDF.loc[tempDF['BaseBMI_class']=='Normal'],\n",
    "              'LMM3':tempDF.loc[tempDF['BaseBMI_class']=='Overweight'],\n",
    "              'LMM4':tempDF.loc[tempDF['BaseBMI_class']=='Obese']}\n",
    "    fe_formula = 'log_'+yvar+' ~ DaysZ + DaysP2Z + DaysP3Z'\\\n",
    "        ' + C(Sex) + BaseAgeZ + C(Season) + PC1Z + PC2Z + PC3Z + PC4Z + PC5Z'\n",
    "    for model_n in tempD2.keys():\n",
    "        print(' - '+tempD2[model_n])\n",
    "        #Fit LMM\n",
    "        model = sm.MixedLM.from_formula(formula=fe_formula, data=tempD3[model_n],\n",
    "                                        groups='public_client_id',\n",
    "                                        re_formula='DaysZ')#Random intercepts for each group are included as default\n",
    "        t_start = time.time()\n",
    "        #fit_res = model.fit(method=['bfgs', 'lbfgs', 'cg', 'powell', 'nm'])#Back up for failure of convergence\n",
    "        fit_res = model.fit(method=['powell'])#After checking convergence for all models, fix the method through the notebook\n",
    "        t_elapsed = time.time() - t_start\n",
    "        print('    - Elapsed time for fitting LMM:', round(t_elapsed//60), 'min', round(t_elapsed%60, 1), 'sec')\n",
    "        \n",
    "        tempD1[model_sex+' - '+model_n] = fit_res\n",
    "        \n",
    "        #Check the fitted models\n",
    "        display(fit_res.summary())\n",
    "        print('(Note that the estimates for fixed effects and the variances for random effects are shown in the single table.)')\n",
    "        ##Visualize b-coefs and g-coefs\n",
    "        tempDF1 = pd.DataFrame({'Coef':fit_res.params,\n",
    "                                'Coef_ci_l':fit_res.conf_int(alpha=0.05).iloc[:, 0],\n",
    "                                'Coef_ci_h':fit_res.conf_int(alpha=0.05).iloc[:, 1],\n",
    "                                'Pval':fit_res.pvalues})\n",
    "        tempDF1['Coef_ci'] = (tempDF1['Coef_ci_h'] - tempDF1['Coef_ci_l'])/2\n",
    "        tempDF1 = tempDF1.reset_index()\n",
    "        sns.set(style='ticks', font='Arial', context='notebook')\n",
    "        plt.figure(figsize=(4, 0.25*len(tempDF1)))\n",
    "        plt.errorbar(x=tempDF1['Coef'], y=tempDF1['index'], xerr=tempDF1['Coef_ci'],\n",
    "                     fmt='ok', ecolor='k', capsize=5)\n",
    "        plt.xlim(-0.1, 0.1)#Clip\n",
    "        sns.despine()\n",
    "        plt.xlabel(r'$\\beta$'+' or '+r'$\\gamma$'+' coefficient in LMM model\\n(Mean with 95% CI)')\n",
    "        plt.ylabel('')\n",
    "        plt.axvline(x=0, **{'linestyle':'--', 'color':'k'})\n",
    "        for row_i in range(len(tempDF1)):\n",
    "            if tempDF1['Pval'].iloc[row_i]<0.05:\n",
    "                plt.axhspan(ymin=row_i-0.5, ymax=row_i+0.5, facecolor=yvar_color, alpha=0.2, zorder=0)\n",
    "        plt.gca().invert_yaxis()\n",
    "        plt.margins(y=0.01, tight=True)\n",
    "        plt.show()\n",
    "        print('')\n",
    "    print('')\n",
    "\n",
    "LMM1_BS = tempD1['Sex-mixed '+yvar+' model - LMM1']\n",
    "LMM2_BS = tempD1['Sex-mixed '+yvar+' model - LMM2']\n",
    "LMM3_BS = tempD1['Sex-mixed '+yvar+' model - LMM3']\n",
    "LMM4_BS = tempD1['Sex-mixed '+yvar+' model - LMM4']\n",
    "LMM1_FM = tempD1['Sex-stratified '+yvar+' model - LMM1']\n",
    "LMM2_FM = tempD1['Sex-stratified '+yvar+' model - LMM2']\n",
    "LMM3_FM = tempD1['Sex-stratified '+yvar+' model - LMM3']\n",
    "LMM4_FM = tempD1['Sex-stratified '+yvar+' model - LMM4']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1-4. Estimated transition of biological BMI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> The days in program of the fitted values were different between individuals.  \n",
    "> –> In statsmodels, MixedLMResults.predict() method use only the fixed effects parameters for prediction.  \n",
    "> ***–> To evaluate the variability in the in-sample population, the rondom effects component should be added manually.***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD1 = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD2 = {'Sex-mixed '+yvar+' model':scaler_BS, 'Sex-stratified '+yvar+' model':scaler_FM}\n",
    "tempD3 = {'Sex-mixed '+yvar+' model':LMM1_BS, 'Sex-stratified '+yvar+' model':LMM1_FM}\n",
    "tempD4 = {'Sex-mixed '+yvar+' model':LMM2_BS, 'Sex-stratified '+yvar+' model':LMM2_FM}\n",
    "tempD5 = {'Sex-mixed '+yvar+' model':LMM3_BS, 'Sex-stratified '+yvar+' model':LMM3_FM}\n",
    "tempD6 = {'Sex-mixed '+yvar+' model':LMM4_BS, 'Sex-stratified '+yvar+' model':LMM4_FM}\n",
    "tempD4 = {'Normal':tempD4, 'Overweight':tempD5, 'Obese':tempD6}\n",
    "tempD5 = {}\n",
    "for model_sex in tempD1.keys():\n",
    "    #Prepare DF to impute days_in_program but maintain in-sample population\n",
    "    tempDF = tempD1[model_sex]\n",
    "    tempA1 = tempDF['public_client_id'].unique()\n",
    "    tempA2 = np.arange(0, 365.25/12*month_threshold + 1, 365.25/12*month_threshold/nknots)\n",
    "    tempL = ['Spring', 'Summer', 'Autumn', 'Winter']\n",
    "    tempDF = pd.DataFrame(data={'public_client_id':np.repeat(tempA1, len(tempA2)*len(tempL)),\n",
    "                                'days_in_program':np.tile(np.repeat(tempA2, len(tempL)), len(tempA1)),\n",
    "                                'Season':np.tile(tempL, len(tempA1)*len(tempA2))})\n",
    "    ##Add the covariates unique to the in-sample individuals\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempL = ['public_client_id', 'Sex', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5', 'BaseBMI_class']\n",
    "    tempDF1 = tempDF1[tempL]\n",
    "    tempDF1 = tempDF1.drop_duplicates('public_client_id', keep='first')\n",
    "    tempDF = pd.merge(tempDF, tempDF1, on='public_client_id', how='left')\n",
    "    ##Add dummy variables for the linear regression spline for time\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempL1 = []#2nd piecewise\n",
    "    tempL2 = []#3rd piecewise\n",
    "    for days in tempDF['days_in_program'].tolist():\n",
    "        if days < tempL[0]:#[0 month, 6 months)\n",
    "            tempL1.append(0.0)\n",
    "            tempL2.append(0.0)\n",
    "        elif days < tempL[1]:#[6 months, 12 months)\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(0.0)\n",
    "        elif days <= tempL[2]:#[12 months, 18 months]\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(days - tempL[1])\n",
    "        else:\n",
    "            tempL1.append('Error? Check time span setting.')\n",
    "            tempL2.append('Error? Check time span setting.')\n",
    "    tempDF['DaysP2'] = tempL1\n",
    "    tempDF['DaysP3'] = tempL2\n",
    "    ##Add a dummy index to use when merging later\n",
    "    tempDF['KeyIndex'] = tempDF['public_client_id'] + '_day' + tempDF['days_in_program'].astype(str)\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    ##Standardize numeric independent variables using the already fitted StandardScaler\n",
    "    scaler = tempD2[model_sex]\n",
    "    tempL = ['days_in_program', 'DaysP2', 'DaysP3', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5']\n",
    "    tempDF1 = pd.DataFrame(scaler.transform(tempDF[tempL]), index=tempDF.index,\n",
    "                           columns=['DaysZ', 'DaysP2Z', 'DaysP3Z', 'BaseAgeZ', 'PC1Z', 'PC2Z', 'PC3Z', 'PC4Z', 'PC5Z'])\n",
    "    tempDF = pd.merge(tempDF, tempDF1, left_index=True, right_index=True, how='inner')\n",
    "    \n",
    "    #Model 1 prediction\n",
    "    fit_res = tempD3[model_sex]\n",
    "    ##Predicted component from the fixed effects\n",
    "    tempS = fit_res.predict(tempDF)#Automatically recognize the column name in formula\n",
    "    ##Create design matrix for the random effects\n",
    "    tempDF1 = pd.DataFrame({'Intercept_re':np.repeat(1, len(tempDF)),\n",
    "                            'Slope_re':tempDF['DaysZ']})\n",
    "    tempD = fit_res.random_effects#Dictionary of key=group label and value=pd.Series\n",
    "    ##Predicted component from the random effects\n",
    "    tempL = [np.dot(tempDF1.iloc[row_i], tempD[group_n])\n",
    "             for (row_i, group_n) in enumerate(tempDF['public_client_id'].tolist())]\n",
    "    ##Total prediction for in-sample individual\n",
    "    tempS = tempS + tempL\n",
    "    ##Convert to original scale\n",
    "    tempDF[yvar+'_predicted1'] = np.e**tempS\n",
    "    \n",
    "    #Model 2-4 prediction\n",
    "    for model_i, bmi_class in enumerate(tempD4.keys(), start=2):\n",
    "        tempD = tempD4[bmi_class]\n",
    "        fit_res = tempD[model_sex]\n",
    "        ##Prepare the target class\n",
    "        tempDF1 = tempDF.loc[tempDF['BaseBMI_class']==bmi_class]\n",
    "        ##Predicted component from the fixed effects\n",
    "        tempS = fit_res.predict(tempDF1)#Automatically recognize the column name in formula\n",
    "        ##Create design matrix for the random effects\n",
    "        tempDF2 = pd.DataFrame({'Intercept_re':np.repeat(1, len(tempDF1)),\n",
    "                                'Slope_re':tempDF1['DaysZ']})\n",
    "        tempD = fit_res.random_effects#Dictionary of key=group label and value=pd.Series\n",
    "        ##Predicted component from the random effects\n",
    "        tempL = [np.dot(tempDF2.iloc[row_i], tempD[group_n])\n",
    "                 for (row_i, group_n) in enumerate(tempDF1['public_client_id'].tolist())]\n",
    "        ##Total prediction for in-sample individual\n",
    "        tempS = tempS + tempL\n",
    "        ##Convert to original scale\n",
    "        tempDF1[yvar+'_predicted'+str(model_i)] = np.e**tempS\n",
    "        ##Add the predictions to the working DF\n",
    "        tempDF = pd.merge(tempDF, tempDF1[yvar+'_predicted'+str(model_i)],\n",
    "                          left_index=True, right_index=True, how='outer')#NaN for out-of-class\n",
    "    \n",
    "    #Flatten effects of the dummy season variable\n",
    "    tempDF = tempDF.groupby(by=['public_client_id', 'days_in_program'])\n",
    "    tempDF = tempDF.agg({yvar+'_predicted1':'mean',\n",
    "                         yvar+'_predicted2':'mean',\n",
    "                         yvar+'_predicted3':'mean',\n",
    "                         yvar+'_predicted4':'mean'})\n",
    "    tempDF = tempDF.reset_index()\n",
    "    ##Recover the baseline BMI class to use later\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempL = ['public_client_id', 'BaseBMI_class']\n",
    "    tempDF1 = tempDF1[tempL].drop_duplicates('public_client_id', keep='first')\n",
    "    tempDF = pd.merge(tempDF, tempDF1, on='public_client_id', how='left')\n",
    "    tempDF['KeyIndex'] = tempDF['public_client_id'] + '_day' + tempDF['days_in_program'].astype(str)\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    \n",
    "    print(model_sex)\n",
    "    display(tempDF)\n",
    "    display(tempDF.describe(include='all'))\n",
    "    \n",
    "    tempD5[model_sex] = tempDF\n",
    "\n",
    "metLMM_BS = tempD5['Sex-mixed '+yvar+' model']\n",
    "metLMM_FM = tempD5['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check visually\n",
    "tempD1 = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD2 = {'Sex-mixed '+yvar+' model':metLMM_BS, 'Sex-stratified '+yvar+' model':metLMM_FM}\n",
    "for model_sex in tempD1.keys():\n",
    "    print(model_sex)\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempDF2 = tempD2[model_sex]\n",
    "    \n",
    "    #Plot1\n",
    "    print(' - Model 1: Overall LMM')\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(5, 3))\n",
    "    sns.lineplot(data=tempDF1, x='days_in_program', y=yvar,\n",
    "                 units='public_client_id', estimator=None, color='gray', lw=1, alpha=0.1, legend=None)\n",
    "    sns.scatterplot(data=tempDF1, x='days_in_program', y=yvar,\n",
    "                    color='gray', edgecolor='k', s=20, alpha=0.3)\n",
    "    p = sns.lineplot(data=tempDF2, x='days_in_program', y=yvar+'_predicted1',\n",
    "                     estimator='mean', ci=95, color=yvar_color)\n",
    "    p.set(xlim=(0, 365.25/12*month_threshold), xticks=np.arange(0, 365.25/12*month_threshold, 100),\n",
    "          ylim=(17.5, 42.5), yticks=np.arange(20, 40.1, 5))\n",
    "    sns.despine()\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.ylabel(yvar+' [kg/m'+r'$^2$'+']')\n",
    "    plt.show()\n",
    "    print('')\n",
    "    \n",
    "    #Plot2\n",
    "    print(' - Model 2-4: Baseline BMI class-stratified LMM')\n",
    "    tempD = {'Normal':'green', 'Overweight':'orange', 'Obese':'red'}\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(5, 3))\n",
    "    p = sns.scatterplot(data=tempDF1[tempDF1['BaseBMI_class']!='Underweight'],\n",
    "                        x='days_in_program', y=yvar,\n",
    "                        hue='BaseBMI_class', hue_order=tempD.keys(), palette=tempD,\n",
    "                        edgecolor='k', s=20, alpha=0.3)\n",
    "    for model_i, bmi_class in enumerate(tempD.keys(), start=2):\n",
    "        sns.lineplot(data=tempDF2[tempDF2['BaseBMI_class']==bmi_class],\n",
    "                     x='days_in_program', y=yvar+'_predicted'+str(model_i),\n",
    "                     color=tempD[bmi_class], estimator='mean', ci=95, ax=p)\n",
    "    p.set(xlim=(0, 365.25/12*month_threshold), xticks=np.arange(0, 365.25/12*month_threshold, 100),\n",
    "          ylim=(17.5, 42.5), yticks=np.arange(20, 40.1, 5))\n",
    "    sns.despine()\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.ylabel(yvar+' [kg/m'+r'$^2$'+']')\n",
    "    plt.legend(bbox_to_anchor=(1, 0.5), loc='center left', borderaxespad=1)\n",
    "    plt.show()\n",
    "    print('')\n",
    "    \n",
    "    print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Proteomics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yvar = 'ProtBMI'\n",
    "yvar_color = 'r'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2-1. Prepare the time-series DF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Import time-series biological BMI\n",
    "tempD = {'Sex-mixed '+yvar+' model':'BothSex', 'Sex-stratified '+yvar+' model':'FemaleMale'}\n",
    "tempD1 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    fileDir = './ExportData/'\n",
    "    ipynbName = '220805_Multiomics-BMI-NatMed1stRevision_BMI-longitudinal-LASSO_'\n",
    "    fileName = yvar+'-'+tempD[model_sex]+'.tsv'\n",
    "    tempDF = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    tempD1[model_sex] = tempDF\n",
    "    \n",
    "    print(model_sex)\n",
    "    display(tempDF)\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    print('')\n",
    "\n",
    "tsDF_BS = tempD1['Sex-mixed '+yvar+' model']\n",
    "tsDF_FM = tempD1['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2-2. Prepare DFs for LMM with the linear regresion spline for time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check measurement distribution\n",
    "sns.set(style='ticks', font='Arial', context='notebook')\n",
    "plt.figure(figsize=(4, 3))\n",
    "sns.distplot(tsDF_BS['days_in_program'], color=yvar_color)\n",
    "sns.despine()\n",
    "for knot in [0, 6, 12, 18]:\n",
    "    plt.axvline(x=365.25/12*knot, **{'linestyle':'--', 'color':'gray'})\n",
    "plt.axvspan(xmin=365.25/12*18, xmax=plt.gca().get_xlim()[1], facecolor='gray', alpha=0.2)\n",
    "plt.ylabel('Density')\n",
    "plt.xlabel('Days in program')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD1 = {}\n",
    "tempD2 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    print(model_sex)\n",
    "    tempDF = tempD[model_sex]\n",
    "    print('Time-series DF nrows before filtering:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select measurements for the longitudinal analysis\n",
    "    month_threshold = 18\n",
    "    tempDF = tempDF.loc[tempDF['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    print('Time-series DF nrows after filtering with the taget period:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select the cohort for the longitudinal analysis\n",
    "    ##Select participants who have 2 or more measuremnts in BMI\n",
    "    fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "    ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "    fileName = 'time-series-bmiDF-without-imputation_final-cohort.tsv'\n",
    "    tempDF1 = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF1 = tempDF1.set_index('KeyIndex')\n",
    "    tempDF1 = tempDF1.loc[tempDF1['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    tempS = tempDF1['public_client_id'].value_counts()\n",
    "    tempL1 = tempS.loc[tempS>=2].index.tolist()\n",
    "    ##Select participants who have 2 or more measuremnts in omics\n",
    "    fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "    ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "    fileName = 'time-series-combiDF-without-imputation_final-cohort.tsv'\n",
    "    tempDF1 = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF1 = tempDF1.set_index('KeyIndex')\n",
    "    tempDF1 = tempDF1.loc[tempDF1['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    tempS = tempDF1['public_client_id'].value_counts()\n",
    "    tempL2 = tempS.loc[tempS>=2].index.tolist()\n",
    "    ##Select participants who have 2 or more measuremnts in both BMI and omics\n",
    "    tempL = list(set(tempL1) & set(tempL2))\n",
    "    print('Participants who have 2 or more measurements in both BMI and omics:', len(tempL))\n",
    "    tempDF = tempDF.loc[tempDF['public_client_id'].isin(tempL)]\n",
    "    print('Time-series DF nrows after filtering with the number of measurements:',\n",
    "          len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select the participants who have all covariates (just in case)\n",
    "    tempDF = tempDF.dropna()\n",
    "    print('Time-series DF nrows after filtering with covariates:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Check measurement distribution\n",
    "    sns.set(style='ticks', font='Arial', context='talk')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    sns.distplot(tempDF['days_in_program'], color=yvar_color)\n",
    "    sns.despine()\n",
    "    for knot in [0, 6, 12, 18]:\n",
    "        plt.axvline(x=365.25/12*knot, **{'linestyle':'--', 'color':'gray'})\n",
    "    plt.axvspan(xmin=365.25/12*18, xmax=plt.gca().get_xlim()[1], facecolor='gray', alpha=0.2)\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.show()\n",
    "    \n",
    "    #Check dependent variable distribution\n",
    "    print('Skewness of the dependent variable:', stats.skew(tempDF['log_'+yvar]))\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    sns.distplot(tempDF['log_'+yvar], color=yvar_color)\n",
    "    sns.despine()\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel(yvar+' (log-scale)')\n",
    "    plt.show()\n",
    "    \n",
    "    #Add dummy variables for the linear regression spline for time\n",
    "    nknots = 3\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempL1 = []#2nd piece\n",
    "    tempL2 = []#3rd piece\n",
    "    for days in tempDF['days_in_program'].tolist():\n",
    "        if days < tempL[0]:#[0 month, 6 months)\n",
    "            tempL1.append(0.0)\n",
    "            tempL2.append(0.0)\n",
    "        elif days < tempL[1]:#[6 months, 12 months)\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(0.0)\n",
    "        elif days <= tempL[2]:#[12 months, 18 months]\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(days - tempL[1])\n",
    "        else:\n",
    "            tempL1.append('Error? Check time span setting.')\n",
    "            tempL2.append('Error? Check time span setting.')\n",
    "    tempDF['DaysP2'] = tempL1\n",
    "    tempDF['DaysP3'] = tempL2\n",
    "    \n",
    "    #Standardize numeric independent variables\n",
    "    scaler = StandardScaler(copy=True, with_mean=True, with_std=True)\n",
    "    tempL = ['days_in_program', 'DaysP2', 'DaysP3', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5']\n",
    "    scaler.fit(tempDF[tempL])\n",
    "    tempDF1 = pd.DataFrame(scaler.transform(tempDF[tempL]), index=tempDF.index,\n",
    "                           columns=['DaysZ', 'DaysP2Z', 'DaysP3Z', 'BaseAgeZ', 'PC1Z', 'PC2Z', 'PC3Z', 'PC4Z', 'PC5Z'])\n",
    "    tempDF = pd.merge(tempDF, tempDF1, left_index=True, right_index=True, how='inner')\n",
    "    \n",
    "    #Confirm examples\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    for variable in ['DaysZ', 'BaseAgeZ', 'PC1Z']:\n",
    "        sns.distplot(tempDF[variable], label=variable)\n",
    "    sns.despine()\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel(r'$Z$'+'-score')\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempS = (pd.Series(tempL) - scaler.mean_[0]) / scaler.scale_[0]#Z-score transformation\n",
    "    for knot in tempS.tolist():\n",
    "        plt.axvline(x=knot, **{'linestyle':'--', 'color':'gray'})\n",
    "    plt.legend(bbox_to_anchor=(1, 0.5), loc='center left', borderaxespad=1)\n",
    "    plt.show()\n",
    "    \n",
    "    display(tempDF)\n",
    "    print('')\n",
    "    \n",
    "    tempD1[model_sex] = tempDF\n",
    "    tempD2[model_sex] = scaler#Use later again during predictions\n",
    "\n",
    "#Update\n",
    "tsDF_BS = tempD1['Sex-mixed '+yvar+' model']\n",
    "scaler_BS = tempD2['Sex-mixed '+yvar+' model']\n",
    "tsDF_FM = tempD1['Sex-stratified '+yvar+' model']\n",
    "scaler_FM = tempD2['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2-3. LMM with random intercepts and random slopes for time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Days in the program is fitted as a numeric variable using linear regression splines (cf. Zubair, N. et al. Sci. Rep. 2019).  \n",
    "> ***–> Allows differences in the trajectory of BMI/bBMI change.***\n",
    ">\n",
    "> Random intercepts and random slopes for days in the program are used as random effects (cf. Zubair, N. et al. Sci. Rep. 2019).  \n",
    "> ***–> Allows individual differences in the baseline and longitudinal change over the program.***  \n",
    "> –> Random slope is applied not for the piecewise variable of time but only for the overall variable of time; otherwise, leads to \"LinAlgError: Singular matrix\" due to small sample size per most individuals. Also, it is not so weird to assume the random slope is consistent during the program.  \n",
    ">\n",
    "> ***Eliminate \"Underweight\" class in the baseline obesity class-stratified models.***  \n",
    "> <– Otherwise, leads to \"LinAlgError: Singular matrix\" due to small sample population.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD1 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    print(model_sex)\n",
    "    tempDF = tempD[model_sex]\n",
    "    \n",
    "    tempD2 = {'LMM1':'Overall LMM',\n",
    "              'LMM2':'Normal class LMM',\n",
    "              'LMM3':'Overweight class LMM',\n",
    "              'LMM4':'Obese class LMM'}\n",
    "    tempD3 = {'LMM1':tempDF,\n",
    "              'LMM2':tempDF.loc[tempDF['BaseBMI_class']=='Normal'],\n",
    "              'LMM3':tempDF.loc[tempDF['BaseBMI_class']=='Overweight'],\n",
    "              'LMM4':tempDF.loc[tempDF['BaseBMI_class']=='Obese']}\n",
    "    fe_formula = 'log_'+yvar+' ~ DaysZ + DaysP2Z + DaysP3Z'\\\n",
    "        ' + C(Sex) + BaseAgeZ + C(Season) + PC1Z + PC2Z + PC3Z + PC4Z + PC5Z'\n",
    "    for model_n in tempD2.keys():\n",
    "        print(' - '+tempD2[model_n])\n",
    "        #Fit LMM\n",
    "        model = sm.MixedLM.from_formula(formula=fe_formula, data=tempD3[model_n],\n",
    "                                        groups='public_client_id',\n",
    "                                        re_formula='DaysZ')#Random intercepts for each group are included as default\n",
    "        t_start = time.time()\n",
    "        #fit_res = model.fit(method=['bfgs', 'lbfgs', 'cg', 'powell', 'nm'])#Back up for failure of convergence\n",
    "        fit_res = model.fit(method=['powell'])#After checking convergence for all models, fix the method through the notebook\n",
    "        t_elapsed = time.time() - t_start\n",
    "        print('    - Elapsed time for fitting LMM:', round(t_elapsed//60), 'min', round(t_elapsed%60, 1), 'sec')\n",
    "        \n",
    "        tempD1[model_sex+' - '+model_n] = fit_res\n",
    "        \n",
    "        #Check the fitted models\n",
    "        display(fit_res.summary())\n",
    "        print('(Note that the estimates for fixed effects and the variances for random effects are shown in the single table.)')\n",
    "        ##Visualize b-coefs and g-coefs\n",
    "        tempDF1 = pd.DataFrame({'Coef':fit_res.params,\n",
    "                                'Coef_ci_l':fit_res.conf_int(alpha=0.05).iloc[:, 0],\n",
    "                                'Coef_ci_h':fit_res.conf_int(alpha=0.05).iloc[:, 1],\n",
    "                                'Pval':fit_res.pvalues})\n",
    "        tempDF1['Coef_ci'] = (tempDF1['Coef_ci_h'] - tempDF1['Coef_ci_l'])/2\n",
    "        tempDF1 = tempDF1.reset_index()\n",
    "        sns.set(style='ticks', font='Arial', context='notebook')\n",
    "        plt.figure(figsize=(4, 0.25*len(tempDF1)))\n",
    "        plt.errorbar(x=tempDF1['Coef'], y=tempDF1['index'], xerr=tempDF1['Coef_ci'],\n",
    "                     fmt='ok', ecolor='k', capsize=5)\n",
    "        plt.xlim(-0.1, 0.1)#Clip\n",
    "        sns.despine()\n",
    "        plt.xlabel(r'$\\beta$'+' or '+r'$\\gamma$'+' coefficient in LMM model\\n(Mean with 95% CI)')\n",
    "        plt.ylabel('')\n",
    "        plt.axvline(x=0, **{'linestyle':'--', 'color':'k'})\n",
    "        for row_i in range(len(tempDF1)):\n",
    "            if tempDF1['Pval'].iloc[row_i]<0.05:\n",
    "                plt.axhspan(ymin=row_i-0.5, ymax=row_i+0.5, facecolor=yvar_color, alpha=0.2, zorder=0)\n",
    "        plt.gca().invert_yaxis()\n",
    "        plt.margins(y=0.01, tight=True)\n",
    "        plt.show()\n",
    "        print('')\n",
    "    print('')\n",
    "\n",
    "LMM1_BS = tempD1['Sex-mixed '+yvar+' model - LMM1']\n",
    "LMM2_BS = tempD1['Sex-mixed '+yvar+' model - LMM2']\n",
    "LMM3_BS = tempD1['Sex-mixed '+yvar+' model - LMM3']\n",
    "LMM4_BS = tempD1['Sex-mixed '+yvar+' model - LMM4']\n",
    "LMM1_FM = tempD1['Sex-stratified '+yvar+' model - LMM1']\n",
    "LMM2_FM = tempD1['Sex-stratified '+yvar+' model - LMM2']\n",
    "LMM3_FM = tempD1['Sex-stratified '+yvar+' model - LMM3']\n",
    "LMM4_FM = tempD1['Sex-stratified '+yvar+' model - LMM4']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2-4. Estimated transition of biological BMI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> The days in program of the fitted values were different between individuals.  \n",
    "> –> In statsmodels, MixedLMResults.predict() method use only the fixed effects parameters for prediction.  \n",
    "> ***–> To evaluate the variability in the in-sample population, the rondom effects component should be added manually.***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD1 = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD2 = {'Sex-mixed '+yvar+' model':scaler_BS, 'Sex-stratified '+yvar+' model':scaler_FM}\n",
    "tempD3 = {'Sex-mixed '+yvar+' model':LMM1_BS, 'Sex-stratified '+yvar+' model':LMM1_FM}\n",
    "tempD4 = {'Sex-mixed '+yvar+' model':LMM2_BS, 'Sex-stratified '+yvar+' model':LMM2_FM}\n",
    "tempD5 = {'Sex-mixed '+yvar+' model':LMM3_BS, 'Sex-stratified '+yvar+' model':LMM3_FM}\n",
    "tempD6 = {'Sex-mixed '+yvar+' model':LMM4_BS, 'Sex-stratified '+yvar+' model':LMM4_FM}\n",
    "tempD4 = {'Normal':tempD4, 'Overweight':tempD5, 'Obese':tempD6}\n",
    "tempD5 = {}\n",
    "for model_sex in tempD1.keys():\n",
    "    #Prepare DF to impute days_in_program but maintain in-sample population\n",
    "    tempDF = tempD1[model_sex]\n",
    "    tempA1 = tempDF['public_client_id'].unique()\n",
    "    tempA2 = np.arange(0, 365.25/12*month_threshold + 1, 365.25/12*month_threshold/nknots)\n",
    "    tempL = ['Spring', 'Summer', 'Autumn', 'Winter']\n",
    "    tempDF = pd.DataFrame(data={'public_client_id':np.repeat(tempA1, len(tempA2)*len(tempL)),\n",
    "                                'days_in_program':np.tile(np.repeat(tempA2, len(tempL)), len(tempA1)),\n",
    "                                'Season':np.tile(tempL, len(tempA1)*len(tempA2))})\n",
    "    ##Add the covariates unique to the in-sample individuals\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempL = ['public_client_id', 'Sex', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5', 'BaseBMI_class']\n",
    "    tempDF1 = tempDF1[tempL]\n",
    "    tempDF1 = tempDF1.drop_duplicates('public_client_id', keep='first')\n",
    "    tempDF = pd.merge(tempDF, tempDF1, on='public_client_id', how='left')\n",
    "    ##Add dummy variables for the linear regression spline for time\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempL1 = []#2nd piecewise\n",
    "    tempL2 = []#3rd piecewise\n",
    "    for days in tempDF['days_in_program'].tolist():\n",
    "        if days < tempL[0]:#[0 month, 6 months)\n",
    "            tempL1.append(0.0)\n",
    "            tempL2.append(0.0)\n",
    "        elif days < tempL[1]:#[6 months, 12 months)\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(0.0)\n",
    "        elif days <= tempL[2]:#[12 months, 18 months]\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(days - tempL[1])\n",
    "        else:\n",
    "            tempL1.append('Error? Check time span setting.')\n",
    "            tempL2.append('Error? Check time span setting.')\n",
    "    tempDF['DaysP2'] = tempL1\n",
    "    tempDF['DaysP3'] = tempL2\n",
    "    ##Add a dummy index to use when merging later\n",
    "    tempDF['KeyIndex'] = tempDF['public_client_id'] + '_day' + tempDF['days_in_program'].astype(str)\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    ##Standardize numeric independent variables using the already fitted StandardScaler\n",
    "    scaler = tempD2[model_sex]\n",
    "    tempL = ['days_in_program', 'DaysP2', 'DaysP3', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5']\n",
    "    tempDF1 = pd.DataFrame(scaler.transform(tempDF[tempL]), index=tempDF.index,\n",
    "                           columns=['DaysZ', 'DaysP2Z', 'DaysP3Z', 'BaseAgeZ', 'PC1Z', 'PC2Z', 'PC3Z', 'PC4Z', 'PC5Z'])\n",
    "    tempDF = pd.merge(tempDF, tempDF1, left_index=True, right_index=True, how='inner')\n",
    "    \n",
    "    #Model 1 prediction\n",
    "    fit_res = tempD3[model_sex]\n",
    "    ##Predicted component from the fixed effects\n",
    "    tempS = fit_res.predict(tempDF)#Automatically recognize the column name in formula\n",
    "    ##Create design matrix for the random effects\n",
    "    tempDF1 = pd.DataFrame({'Intercept_re':np.repeat(1, len(tempDF)),\n",
    "                            'Slope_re':tempDF['DaysZ']})\n",
    "    tempD = fit_res.random_effects#Dictionary of key=group label and value=pd.Series\n",
    "    ##Predicted component from the random effects\n",
    "    tempL = [np.dot(tempDF1.iloc[row_i], tempD[group_n])\n",
    "             for (row_i, group_n) in enumerate(tempDF['public_client_id'].tolist())]\n",
    "    ##Total prediction for in-sample individual\n",
    "    tempS = tempS + tempL\n",
    "    ##Convert to original scale\n",
    "    tempDF[yvar+'_predicted1'] = np.e**tempS\n",
    "    \n",
    "    #Model 2-4 prediction\n",
    "    for model_i, bmi_class in enumerate(tempD4.keys(), start=2):\n",
    "        tempD = tempD4[bmi_class]\n",
    "        fit_res = tempD[model_sex]\n",
    "        ##Prepare the target class\n",
    "        tempDF1 = tempDF.loc[tempDF['BaseBMI_class']==bmi_class]\n",
    "        ##Predicted component from the fixed effects\n",
    "        tempS = fit_res.predict(tempDF1)#Automatically recognize the column name in formula\n",
    "        ##Create design matrix for the random effects\n",
    "        tempDF2 = pd.DataFrame({'Intercept_re':np.repeat(1, len(tempDF1)),\n",
    "                                'Slope_re':tempDF1['DaysZ']})\n",
    "        tempD = fit_res.random_effects#Dictionary of key=group label and value=pd.Series\n",
    "        ##Predicted component from the random effects\n",
    "        tempL = [np.dot(tempDF2.iloc[row_i], tempD[group_n])\n",
    "                 for (row_i, group_n) in enumerate(tempDF1['public_client_id'].tolist())]\n",
    "        ##Total prediction for in-sample individual\n",
    "        tempS = tempS + tempL\n",
    "        ##Convert to original scale\n",
    "        tempDF1[yvar+'_predicted'+str(model_i)] = np.e**tempS\n",
    "        ##Add the predictions to the working DF\n",
    "        tempDF = pd.merge(tempDF, tempDF1[yvar+'_predicted'+str(model_i)],\n",
    "                          left_index=True, right_index=True, how='outer')#NaN for out-of-class\n",
    "    \n",
    "    #Flatten effects of the dummy season variable\n",
    "    tempDF = tempDF.groupby(by=['public_client_id', 'days_in_program'])\n",
    "    tempDF = tempDF.agg({yvar+'_predicted1':'mean',\n",
    "                         yvar+'_predicted2':'mean',\n",
    "                         yvar+'_predicted3':'mean',\n",
    "                         yvar+'_predicted4':'mean'})\n",
    "    tempDF = tempDF.reset_index()\n",
    "    ##Recover the baseline BMI class to use later\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempL = ['public_client_id', 'BaseBMI_class']\n",
    "    tempDF1 = tempDF1[tempL].drop_duplicates('public_client_id', keep='first')\n",
    "    tempDF = pd.merge(tempDF, tempDF1, on='public_client_id', how='left')\n",
    "    tempDF['KeyIndex'] = tempDF['public_client_id'] + '_day' + tempDF['days_in_program'].astype(str)\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    \n",
    "    print(model_sex)\n",
    "    display(tempDF)\n",
    "    display(tempDF.describe(include='all'))\n",
    "    \n",
    "    tempD5[model_sex] = tempDF\n",
    "\n",
    "protLMM_BS = tempD5['Sex-mixed '+yvar+' model']\n",
    "protLMM_FM = tempD5['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check visually\n",
    "tempD1 = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD2 = {'Sex-mixed '+yvar+' model':protLMM_BS, 'Sex-stratified '+yvar+' model':protLMM_FM}\n",
    "for model_sex in tempD1.keys():\n",
    "    print(model_sex)\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempDF2 = tempD2[model_sex]\n",
    "    \n",
    "    #Plot1\n",
    "    print(' - Model 1: Overall LMM')\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(5, 3))\n",
    "    sns.lineplot(data=tempDF1, x='days_in_program', y=yvar,\n",
    "                 units='public_client_id', estimator=None, color='gray', lw=1, alpha=0.1, legend=None)\n",
    "    sns.scatterplot(data=tempDF1, x='days_in_program', y=yvar,\n",
    "                    color='gray', edgecolor='k', s=20, alpha=0.3)\n",
    "    p = sns.lineplot(data=tempDF2, x='days_in_program', y=yvar+'_predicted1',\n",
    "                     estimator='mean', ci=95, color=yvar_color)\n",
    "    p.set(xlim=(0, 365.25/12*month_threshold), xticks=np.arange(0, 365.25/12*month_threshold, 100),\n",
    "          ylim=(17.5, 42.5), yticks=np.arange(20, 40.1, 5))\n",
    "    sns.despine()\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.ylabel(yvar+' [kg/m'+r'$^2$'+']')\n",
    "    plt.show()\n",
    "    print('')\n",
    "    \n",
    "    #Plot2\n",
    "    print(' - Model 2-4: Baseline BMI class-stratified LMM')\n",
    "    tempD = {'Normal':'green', 'Overweight':'orange', 'Obese':'red'}\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(5, 3))\n",
    "    p = sns.scatterplot(data=tempDF1[tempDF1['BaseBMI_class']!='Underweight'],\n",
    "                        x='days_in_program', y=yvar,\n",
    "                        hue='BaseBMI_class', hue_order=tempD.keys(), palette=tempD,\n",
    "                        edgecolor='k', s=20, alpha=0.3)\n",
    "    for model_i, bmi_class in enumerate(tempD.keys(), start=2):\n",
    "        sns.lineplot(data=tempDF2[tempDF2['BaseBMI_class']==bmi_class],\n",
    "                     x='days_in_program', y=yvar+'_predicted'+str(model_i),\n",
    "                     color=tempD[bmi_class], estimator='mean', ci=95, ax=p)\n",
    "    p.set(xlim=(0, 365.25/12*month_threshold), xticks=np.arange(0, 365.25/12*month_threshold, 100),\n",
    "          ylim=(17.5, 42.5), yticks=np.arange(20, 40.1, 5))\n",
    "    sns.despine()\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.ylabel(yvar+' [kg/m'+r'$^2$'+']')\n",
    "    plt.legend(bbox_to_anchor=(1, 0.5), loc='center left', borderaxespad=1)\n",
    "    plt.show()\n",
    "    print('')\n",
    "    \n",
    "    print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Clinical labs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yvar = 'ChemBMI'\n",
    "yvar_color = 'g'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3-1. Prepare the time-series DF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Import time-series biological BMI\n",
    "tempD = {'Sex-mixed '+yvar+' model':'BothSex', 'Sex-stratified '+yvar+' model':'FemaleMale'}\n",
    "tempD1 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    fileDir = './ExportData/'\n",
    "    ipynbName = '220805_Multiomics-BMI-NatMed1stRevision_BMI-longitudinal-LASSO_'\n",
    "    fileName = yvar+'-'+tempD[model_sex]+'.tsv'\n",
    "    tempDF = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    tempD1[model_sex] = tempDF\n",
    "    \n",
    "    print(model_sex)\n",
    "    display(tempDF)\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    print('')\n",
    "\n",
    "tsDF_BS = tempD1['Sex-mixed '+yvar+' model']\n",
    "tsDF_FM = tempD1['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3-2. Prepare DFs for LMM with the linear regresion spline for time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check measurement distribution\n",
    "sns.set(style='ticks', font='Arial', context='notebook')\n",
    "plt.figure(figsize=(4, 3))\n",
    "sns.distplot(tsDF_BS['days_in_program'], color=yvar_color)\n",
    "sns.despine()\n",
    "for knot in [0, 6, 12, 18]:\n",
    "    plt.axvline(x=365.25/12*knot, **{'linestyle':'--', 'color':'gray'})\n",
    "plt.axvspan(xmin=365.25/12*18, xmax=plt.gca().get_xlim()[1], facecolor='gray', alpha=0.2)\n",
    "plt.ylabel('Density')\n",
    "plt.xlabel('Days in program')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD1 = {}\n",
    "tempD2 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    print(model_sex)\n",
    "    tempDF = tempD[model_sex]\n",
    "    print('Time-series DF nrows before filtering:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select measurements for the longitudinal analysis\n",
    "    month_threshold = 18\n",
    "    tempDF = tempDF.loc[tempDF['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    print('Time-series DF nrows after filtering with the taget period:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select the cohort for the longitudinal analysis\n",
    "    ##Select participants who have 2 or more measuremnts in BMI\n",
    "    fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "    ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "    fileName = 'time-series-bmiDF-without-imputation_final-cohort.tsv'\n",
    "    tempDF1 = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF1 = tempDF1.set_index('KeyIndex')\n",
    "    tempDF1 = tempDF1.loc[tempDF1['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    tempS = tempDF1['public_client_id'].value_counts()\n",
    "    tempL1 = tempS.loc[tempS>=2].index.tolist()\n",
    "    ##Select participants who have 2 or more measuremnts in omics\n",
    "    fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "    ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "    fileName = 'time-series-combiDF-without-imputation_final-cohort.tsv'\n",
    "    tempDF1 = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF1 = tempDF1.set_index('KeyIndex')\n",
    "    tempDF1 = tempDF1.loc[tempDF1['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    tempS = tempDF1['public_client_id'].value_counts()\n",
    "    tempL2 = tempS.loc[tempS>=2].index.tolist()\n",
    "    ##Select participants who have 2 or more measuremnts in both BMI and omics\n",
    "    tempL = list(set(tempL1) & set(tempL2))\n",
    "    print('Participants who have 2 or more measurements in both BMI and omics:', len(tempL))\n",
    "    tempDF = tempDF.loc[tempDF['public_client_id'].isin(tempL)]\n",
    "    print('Time-series DF nrows after filtering with the number of measurements:',\n",
    "          len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select the participants who have all covariates (just in case)\n",
    "    tempDF = tempDF.dropna()\n",
    "    print('Time-series DF nrows after filtering with covariates:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Check measurement distribution\n",
    "    sns.set(style='ticks', font='Arial', context='talk')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    sns.distplot(tempDF['days_in_program'], color=yvar_color)\n",
    "    sns.despine()\n",
    "    for knot in [0, 6, 12, 18]:\n",
    "        plt.axvline(x=365.25/12*knot, **{'linestyle':'--', 'color':'gray'})\n",
    "    plt.axvspan(xmin=365.25/12*18, xmax=plt.gca().get_xlim()[1], facecolor='gray', alpha=0.2)\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.show()\n",
    "    \n",
    "    #Check dependent variable distribution\n",
    "    print('Skewness of the dependent variable:', stats.skew(tempDF['log_'+yvar]))\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    sns.distplot(tempDF['log_'+yvar], color=yvar_color)\n",
    "    sns.despine()\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel(yvar+' (log-scale)')\n",
    "    plt.show()\n",
    "    \n",
    "    #Add dummy variables for the linear regression spline for time\n",
    "    nknots = 3\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempL1 = []#2nd piece\n",
    "    tempL2 = []#3rd piece\n",
    "    for days in tempDF['days_in_program'].tolist():\n",
    "        if days < tempL[0]:#[0 month, 6 months)\n",
    "            tempL1.append(0.0)\n",
    "            tempL2.append(0.0)\n",
    "        elif days < tempL[1]:#[6 months, 12 months)\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(0.0)\n",
    "        elif days <= tempL[2]:#[12 months, 18 months]\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(days - tempL[1])\n",
    "        else:\n",
    "            tempL1.append('Error? Check time span setting.')\n",
    "            tempL2.append('Error? Check time span setting.')\n",
    "    tempDF['DaysP2'] = tempL1\n",
    "    tempDF['DaysP3'] = tempL2\n",
    "    \n",
    "    #Standardize numeric independent variables\n",
    "    scaler = StandardScaler(copy=True, with_mean=True, with_std=True)\n",
    "    tempL = ['days_in_program', 'DaysP2', 'DaysP3', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5']\n",
    "    scaler.fit(tempDF[tempL])\n",
    "    tempDF1 = pd.DataFrame(scaler.transform(tempDF[tempL]), index=tempDF.index,\n",
    "                           columns=['DaysZ', 'DaysP2Z', 'DaysP3Z', 'BaseAgeZ', 'PC1Z', 'PC2Z', 'PC3Z', 'PC4Z', 'PC5Z'])\n",
    "    tempDF = pd.merge(tempDF, tempDF1, left_index=True, right_index=True, how='inner')\n",
    "    \n",
    "    #Confirm examples\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    for variable in ['DaysZ', 'BaseAgeZ', 'PC1Z']:\n",
    "        sns.distplot(tempDF[variable], label=variable)\n",
    "    sns.despine()\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel(r'$Z$'+'-score')\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempS = (pd.Series(tempL) - scaler.mean_[0]) / scaler.scale_[0]#Z-score transformation\n",
    "    for knot in tempS.tolist():\n",
    "        plt.axvline(x=knot, **{'linestyle':'--', 'color':'gray'})\n",
    "    plt.legend(bbox_to_anchor=(1, 0.5), loc='center left', borderaxespad=1)\n",
    "    plt.show()\n",
    "    \n",
    "    display(tempDF)\n",
    "    print('')\n",
    "    \n",
    "    tempD1[model_sex] = tempDF\n",
    "    tempD2[model_sex] = scaler#Use later again during predictions\n",
    "\n",
    "#Update\n",
    "tsDF_BS = tempD1['Sex-mixed '+yvar+' model']\n",
    "scaler_BS = tempD2['Sex-mixed '+yvar+' model']\n",
    "tsDF_FM = tempD1['Sex-stratified '+yvar+' model']\n",
    "scaler_FM = tempD2['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3-3. LMM with random intercepts and random slopes for time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Days in the program is fitted as a numeric variable using linear regression splines (cf. Zubair, N. et al. Sci. Rep. 2019).  \n",
    "> ***–> Allows differences in the trajectory of BMI/bBMI change.***\n",
    ">\n",
    "> Random intercepts and random slopes for days in the program are used as random effects (cf. Zubair, N. et al. Sci. Rep. 2019).  \n",
    "> ***–> Allows individual differences in the baseline and longitudinal change over the program.***  \n",
    "> –> Random slope is applied not for the piecewise variable of time but only for the overall variable of time; otherwise, leads to \"LinAlgError: Singular matrix\" due to small sample size per most individuals. Also, it is not so weird to assume the random slope is consistent during the program.  \n",
    ">\n",
    "> ***Eliminate \"Underweight\" class in the baseline obesity class-stratified models.***  \n",
    "> <– Otherwise, leads to \"LinAlgError: Singular matrix\" due to small sample population.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD1 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    print(model_sex)\n",
    "    tempDF = tempD[model_sex]\n",
    "    \n",
    "    tempD2 = {'LMM1':'Overall LMM',\n",
    "              'LMM2':'Normal class LMM',\n",
    "              'LMM3':'Overweight class LMM',\n",
    "              'LMM4':'Obese class LMM'}\n",
    "    tempD3 = {'LMM1':tempDF,\n",
    "              'LMM2':tempDF.loc[tempDF['BaseBMI_class']=='Normal'],\n",
    "              'LMM3':tempDF.loc[tempDF['BaseBMI_class']=='Overweight'],\n",
    "              'LMM4':tempDF.loc[tempDF['BaseBMI_class']=='Obese']}\n",
    "    fe_formula = 'log_'+yvar+' ~ DaysZ + DaysP2Z + DaysP3Z'\\\n",
    "        ' + C(Sex) + BaseAgeZ + C(Season) + PC1Z + PC2Z + PC3Z + PC4Z + PC5Z'\n",
    "    for model_n in tempD2.keys():\n",
    "        print(' - '+tempD2[model_n])\n",
    "        #Fit LMM\n",
    "        model = sm.MixedLM.from_formula(formula=fe_formula, data=tempD3[model_n],\n",
    "                                        groups='public_client_id',\n",
    "                                        re_formula='DaysZ')#Random intercepts for each group are included as default\n",
    "        t_start = time.time()\n",
    "        #fit_res = model.fit(method=['bfgs', 'lbfgs', 'cg', 'powell', 'nm'])#Back up for failure of convergence\n",
    "        fit_res = model.fit(method=['powell'])#After checking convergence for all models, fix the method through the notebook\n",
    "        t_elapsed = time.time() - t_start\n",
    "        print('    - Elapsed time for fitting LMM:', round(t_elapsed//60), 'min', round(t_elapsed%60, 1), 'sec')\n",
    "        \n",
    "        tempD1[model_sex+' - '+model_n] = fit_res\n",
    "        \n",
    "        #Check the fitted models\n",
    "        display(fit_res.summary())\n",
    "        print('(Note that the estimates for fixed effects and the variances for random effects are shown in the single table.)')\n",
    "        ##Visualize b-coefs and g-coefs\n",
    "        tempDF1 = pd.DataFrame({'Coef':fit_res.params,\n",
    "                                'Coef_ci_l':fit_res.conf_int(alpha=0.05).iloc[:, 0],\n",
    "                                'Coef_ci_h':fit_res.conf_int(alpha=0.05).iloc[:, 1],\n",
    "                                'Pval':fit_res.pvalues})\n",
    "        tempDF1['Coef_ci'] = (tempDF1['Coef_ci_h'] - tempDF1['Coef_ci_l'])/2\n",
    "        tempDF1 = tempDF1.reset_index()\n",
    "        sns.set(style='ticks', font='Arial', context='notebook')\n",
    "        plt.figure(figsize=(4, 0.25*len(tempDF1)))\n",
    "        plt.errorbar(x=tempDF1['Coef'], y=tempDF1['index'], xerr=tempDF1['Coef_ci'],\n",
    "                     fmt='ok', ecolor='k', capsize=5)\n",
    "        plt.xlim(-0.1, 0.1)#Clip\n",
    "        sns.despine()\n",
    "        plt.xlabel(r'$\\beta$'+' or '+r'$\\gamma$'+' coefficient in LMM model\\n(Mean with 95% CI)')\n",
    "        plt.ylabel('')\n",
    "        plt.axvline(x=0, **{'linestyle':'--', 'color':'k'})\n",
    "        for row_i in range(len(tempDF1)):\n",
    "            if tempDF1['Pval'].iloc[row_i]<0.05:\n",
    "                plt.axhspan(ymin=row_i-0.5, ymax=row_i+0.5, facecolor=yvar_color, alpha=0.2, zorder=0)\n",
    "        plt.gca().invert_yaxis()\n",
    "        plt.margins(y=0.01, tight=True)\n",
    "        plt.show()\n",
    "        print('')\n",
    "    print('')\n",
    "\n",
    "LMM1_BS = tempD1['Sex-mixed '+yvar+' model - LMM1']\n",
    "LMM2_BS = tempD1['Sex-mixed '+yvar+' model - LMM2']\n",
    "LMM3_BS = tempD1['Sex-mixed '+yvar+' model - LMM3']\n",
    "LMM4_BS = tempD1['Sex-mixed '+yvar+' model - LMM4']\n",
    "LMM1_FM = tempD1['Sex-stratified '+yvar+' model - LMM1']\n",
    "LMM2_FM = tempD1['Sex-stratified '+yvar+' model - LMM2']\n",
    "LMM3_FM = tempD1['Sex-stratified '+yvar+' model - LMM3']\n",
    "LMM4_FM = tempD1['Sex-stratified '+yvar+' model - LMM4']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3-4. Estimated transition of biological BMI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> The days in program of the fitted values were different between individuals.  \n",
    "> –> In statsmodels, MixedLMResults.predict() method use only the fixed effects parameters for prediction.  \n",
    "> ***–> To evaluate the variability in the in-sample population, the rondom effects component should be added manually.***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD1 = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD2 = {'Sex-mixed '+yvar+' model':scaler_BS, 'Sex-stratified '+yvar+' model':scaler_FM}\n",
    "tempD3 = {'Sex-mixed '+yvar+' model':LMM1_BS, 'Sex-stratified '+yvar+' model':LMM1_FM}\n",
    "tempD4 = {'Sex-mixed '+yvar+' model':LMM2_BS, 'Sex-stratified '+yvar+' model':LMM2_FM}\n",
    "tempD5 = {'Sex-mixed '+yvar+' model':LMM3_BS, 'Sex-stratified '+yvar+' model':LMM3_FM}\n",
    "tempD6 = {'Sex-mixed '+yvar+' model':LMM4_BS, 'Sex-stratified '+yvar+' model':LMM4_FM}\n",
    "tempD4 = {'Normal':tempD4, 'Overweight':tempD5, 'Obese':tempD6}\n",
    "tempD5 = {}\n",
    "for model_sex in tempD1.keys():\n",
    "    #Prepare DF to impute days_in_program but maintain in-sample population\n",
    "    tempDF = tempD1[model_sex]\n",
    "    tempA1 = tempDF['public_client_id'].unique()\n",
    "    tempA2 = np.arange(0, 365.25/12*month_threshold + 1, 365.25/12*month_threshold/nknots)\n",
    "    tempL = ['Spring', 'Summer', 'Autumn', 'Winter']\n",
    "    tempDF = pd.DataFrame(data={'public_client_id':np.repeat(tempA1, len(tempA2)*len(tempL)),\n",
    "                                'days_in_program':np.tile(np.repeat(tempA2, len(tempL)), len(tempA1)),\n",
    "                                'Season':np.tile(tempL, len(tempA1)*len(tempA2))})\n",
    "    ##Add the covariates unique to the in-sample individuals\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempL = ['public_client_id', 'Sex', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5', 'BaseBMI_class']\n",
    "    tempDF1 = tempDF1[tempL]\n",
    "    tempDF1 = tempDF1.drop_duplicates('public_client_id', keep='first')\n",
    "    tempDF = pd.merge(tempDF, tempDF1, on='public_client_id', how='left')\n",
    "    ##Add dummy variables for the linear regression spline for time\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempL1 = []#2nd piecewise\n",
    "    tempL2 = []#3rd piecewise\n",
    "    for days in tempDF['days_in_program'].tolist():\n",
    "        if days < tempL[0]:#[0 month, 6 months)\n",
    "            tempL1.append(0.0)\n",
    "            tempL2.append(0.0)\n",
    "        elif days < tempL[1]:#[6 months, 12 months)\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(0.0)\n",
    "        elif days <= tempL[2]:#[12 months, 18 months]\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(days - tempL[1])\n",
    "        else:\n",
    "            tempL1.append('Error? Check time span setting.')\n",
    "            tempL2.append('Error? Check time span setting.')\n",
    "    tempDF['DaysP2'] = tempL1\n",
    "    tempDF['DaysP3'] = tempL2\n",
    "    ##Add a dummy index to use when merging later\n",
    "    tempDF['KeyIndex'] = tempDF['public_client_id'] + '_day' + tempDF['days_in_program'].astype(str)\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    ##Standardize numeric independent variables using the already fitted StandardScaler\n",
    "    scaler = tempD2[model_sex]\n",
    "    tempL = ['days_in_program', 'DaysP2', 'DaysP3', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5']\n",
    "    tempDF1 = pd.DataFrame(scaler.transform(tempDF[tempL]), index=tempDF.index,\n",
    "                           columns=['DaysZ', 'DaysP2Z', 'DaysP3Z', 'BaseAgeZ', 'PC1Z', 'PC2Z', 'PC3Z', 'PC4Z', 'PC5Z'])\n",
    "    tempDF = pd.merge(tempDF, tempDF1, left_index=True, right_index=True, how='inner')\n",
    "    \n",
    "    #Model 1 prediction\n",
    "    fit_res = tempD3[model_sex]\n",
    "    ##Predicted component from the fixed effects\n",
    "    tempS = fit_res.predict(tempDF)#Automatically recognize the column name in formula\n",
    "    ##Create design matrix for the random effects\n",
    "    tempDF1 = pd.DataFrame({'Intercept_re':np.repeat(1, len(tempDF)),\n",
    "                            'Slope_re':tempDF['DaysZ']})\n",
    "    tempD = fit_res.random_effects#Dictionary of key=group label and value=pd.Series\n",
    "    ##Predicted component from the random effects\n",
    "    tempL = [np.dot(tempDF1.iloc[row_i], tempD[group_n])\n",
    "             for (row_i, group_n) in enumerate(tempDF['public_client_id'].tolist())]\n",
    "    ##Total prediction for in-sample individual\n",
    "    tempS = tempS + tempL\n",
    "    ##Convert to original scale\n",
    "    tempDF[yvar+'_predicted1'] = np.e**tempS\n",
    "    \n",
    "    #Model 2-4 prediction\n",
    "    for model_i, bmi_class in enumerate(tempD4.keys(), start=2):\n",
    "        tempD = tempD4[bmi_class]\n",
    "        fit_res = tempD[model_sex]\n",
    "        ##Prepare the target class\n",
    "        tempDF1 = tempDF.loc[tempDF['BaseBMI_class']==bmi_class]\n",
    "        ##Predicted component from the fixed effects\n",
    "        tempS = fit_res.predict(tempDF1)#Automatically recognize the column name in formula\n",
    "        ##Create design matrix for the random effects\n",
    "        tempDF2 = pd.DataFrame({'Intercept_re':np.repeat(1, len(tempDF1)),\n",
    "                                'Slope_re':tempDF1['DaysZ']})\n",
    "        tempD = fit_res.random_effects#Dictionary of key=group label and value=pd.Series\n",
    "        ##Predicted component from the random effects\n",
    "        tempL = [np.dot(tempDF2.iloc[row_i], tempD[group_n])\n",
    "                 for (row_i, group_n) in enumerate(tempDF1['public_client_id'].tolist())]\n",
    "        ##Total prediction for in-sample individual\n",
    "        tempS = tempS + tempL\n",
    "        ##Convert to original scale\n",
    "        tempDF1[yvar+'_predicted'+str(model_i)] = np.e**tempS\n",
    "        ##Add the predictions to the working DF\n",
    "        tempDF = pd.merge(tempDF, tempDF1[yvar+'_predicted'+str(model_i)],\n",
    "                          left_index=True, right_index=True, how='outer')#NaN for out-of-class\n",
    "    \n",
    "    #Flatten effects of the dummy season variable\n",
    "    tempDF = tempDF.groupby(by=['public_client_id', 'days_in_program'])\n",
    "    tempDF = tempDF.agg({yvar+'_predicted1':'mean',\n",
    "                         yvar+'_predicted2':'mean',\n",
    "                         yvar+'_predicted3':'mean',\n",
    "                         yvar+'_predicted4':'mean'})\n",
    "    tempDF = tempDF.reset_index()\n",
    "    ##Recover the baseline BMI class to use later\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempL = ['public_client_id', 'BaseBMI_class']\n",
    "    tempDF1 = tempDF1[tempL].drop_duplicates('public_client_id', keep='first')\n",
    "    tempDF = pd.merge(tempDF, tempDF1, on='public_client_id', how='left')\n",
    "    tempDF['KeyIndex'] = tempDF['public_client_id'] + '_day' + tempDF['days_in_program'].astype(str)\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    \n",
    "    print(model_sex)\n",
    "    display(tempDF)\n",
    "    display(tempDF.describe(include='all'))\n",
    "    \n",
    "    tempD5[model_sex] = tempDF\n",
    "\n",
    "chemLMM_BS = tempD5['Sex-mixed '+yvar+' model']\n",
    "chemLMM_FM = tempD5['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check visually\n",
    "tempD1 = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD2 = {'Sex-mixed '+yvar+' model':chemLMM_BS, 'Sex-stratified '+yvar+' model':chemLMM_FM}\n",
    "for model_sex in tempD1.keys():\n",
    "    print(model_sex)\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempDF2 = tempD2[model_sex]\n",
    "    \n",
    "    #Plot1\n",
    "    print(' - Model 1: Overall LMM')\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(5, 3))\n",
    "    sns.lineplot(data=tempDF1, x='days_in_program', y=yvar,\n",
    "                 units='public_client_id', estimator=None, color='gray', lw=1, alpha=0.1, legend=None)\n",
    "    sns.scatterplot(data=tempDF1, x='days_in_program', y=yvar,\n",
    "                    color='gray', edgecolor='k', s=20, alpha=0.3)\n",
    "    p = sns.lineplot(data=tempDF2, x='days_in_program', y=yvar+'_predicted1',\n",
    "                     estimator='mean', ci=95, color=yvar_color)\n",
    "    p.set(xlim=(0, 365.25/12*month_threshold), xticks=np.arange(0, 365.25/12*month_threshold, 100),\n",
    "          ylim=(17.5, 42.5), yticks=np.arange(20, 40.1, 5))\n",
    "    sns.despine()\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.ylabel(yvar+' [kg/m'+r'$^2$'+']')\n",
    "    plt.show()\n",
    "    print('')\n",
    "    \n",
    "    #Plot2\n",
    "    print(' - Model 2-4: Baseline BMI class-stratified LMM')\n",
    "    tempD = {'Normal':'green', 'Overweight':'orange', 'Obese':'red'}\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(5, 3))\n",
    "    p = sns.scatterplot(data=tempDF1[tempDF1['BaseBMI_class']!='Underweight'],\n",
    "                        x='days_in_program', y=yvar,\n",
    "                        hue='BaseBMI_class', hue_order=tempD.keys(), palette=tempD,\n",
    "                        edgecolor='k', s=20, alpha=0.3)\n",
    "    for model_i, bmi_class in enumerate(tempD.keys(), start=2):\n",
    "        sns.lineplot(data=tempDF2[tempDF2['BaseBMI_class']==bmi_class],\n",
    "                     x='days_in_program', y=yvar+'_predicted'+str(model_i),\n",
    "                     color=tempD[bmi_class], estimator='mean', ci=95, ax=p)\n",
    "    p.set(xlim=(0, 365.25/12*month_threshold), xticks=np.arange(0, 365.25/12*month_threshold, 100),\n",
    "          ylim=(17.5, 42.5), yticks=np.arange(20, 40.1, 5))\n",
    "    sns.despine()\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.ylabel(yvar+' [kg/m'+r'$^2$'+']')\n",
    "    plt.legend(bbox_to_anchor=(1, 0.5), loc='center left', borderaxespad=1)\n",
    "    plt.show()\n",
    "    print('')\n",
    "    \n",
    "    print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Measured BMI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yvar = 'BMI'\n",
    "yvar_color = 'k'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4-1. Prepare the time-series DF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Import time-series BMI\n",
    "fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "fileName = 'time-series-bmiDF-without-imputation_final-cohort.tsv'\n",
    "tempDF = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "tempDF = tempDF.set_index('KeyIndex')\n",
    "##Chage column name\n",
    "tempDF = tempDF.rename(columns={'BMI_CALC':'BMI'})\n",
    "\n",
    "#Add obesity classification (based on BMI)\n",
    "tempDF['BaseBMI'] = np.e**tempDF['log_BaseBMI']\n",
    "tempDF = tempDF.drop(columns=['log_BaseBMI'])\n",
    "tempL = []\n",
    "for row_i in range(len(tempDF)):\n",
    "    value = tempDF['BaseBMI'].iloc[row_i]\n",
    "    if np.isnan(value):\n",
    "        tempL.append('NotCalculated')\n",
    "    elif value < 18.5:\n",
    "        tempL.append('Underweight')\n",
    "    elif value < 25:\n",
    "        tempL.append('Normal')\n",
    "    elif value < 30:\n",
    "        tempL.append('Overweight')\n",
    "    elif value >= 30:\n",
    "        tempL.append('Obese')\n",
    "    else:#Just in case\n",
    "        tempL.append('Error?')\n",
    "tempDF['BaseBMI_class'] = tempL\n",
    "##Confirmation\n",
    "print('BaseBMI_class:')\n",
    "tempDF1 = tempDF.drop_duplicates('public_client_id', keep='first')\n",
    "tempS = tempDF1['BaseBMI_class'].value_counts()\n",
    "tempDF1 = pd.DataFrame({'Count':tempS, 'Percentage':tempS/len(tempDF1)*100})\n",
    "display(tempDF1)\n",
    "\n",
    "display(tempDF)\n",
    "print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "\n",
    "tsDF = tempDF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4-2. Prepare DFs for LMM with the linear regresion spline for time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check measurement distribution\n",
    "sns.set(style='ticks', font='Arial', context='notebook')\n",
    "plt.figure(figsize=(4, 3))\n",
    "sns.distplot(tsDF['days_in_program'], color=yvar_color)\n",
    "sns.despine()\n",
    "for knot in [0, 6, 12, 18]:\n",
    "    plt.axvline(x=365.25/12*knot, **{'linestyle':'--', 'color':'gray'})\n",
    "plt.axvspan(xmin=365.25/12*18, xmax=plt.gca().get_xlim()[1], facecolor='gray', alpha=0.2)\n",
    "plt.ylabel('Density')\n",
    "plt.xlabel('Days in program')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD = {'Sex-independent '+yvar+' measurement':tsDF}\n",
    "tempD1 = {}\n",
    "tempD2 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    print(model_sex)\n",
    "    tempDF = tempD[model_sex]\n",
    "    print('Time-series DF nrows before filtering:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select measurements for the longitudinal analysis\n",
    "    month_threshold = 18\n",
    "    tempDF = tempDF.loc[tempDF['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    print('Time-series DF nrows after filtering with the taget period:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select the cohort for the longitudinal analysis\n",
    "    ##Select participants who have 2 or more measuremnts in BMI\n",
    "    fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "    ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "    fileName = 'time-series-bmiDF-without-imputation_final-cohort.tsv'\n",
    "    tempDF1 = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF1 = tempDF1.set_index('KeyIndex')\n",
    "    tempDF1 = tempDF1.loc[tempDF1['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    tempS = tempDF1['public_client_id'].value_counts()\n",
    "    tempL1 = tempS.loc[tempS>=2].index.tolist()\n",
    "    ##Select participants who have 2 or more measuremnts in omics\n",
    "    fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "    ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "    fileName = 'time-series-combiDF-without-imputation_final-cohort.tsv'\n",
    "    tempDF1 = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF1 = tempDF1.set_index('KeyIndex')\n",
    "    tempDF1 = tempDF1.loc[tempDF1['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    tempS = tempDF1['public_client_id'].value_counts()\n",
    "    tempL2 = tempS.loc[tempS>=2].index.tolist()\n",
    "    ##Select participants who have 2 or more measuremnts in both BMI and omics\n",
    "    tempL = list(set(tempL1) & set(tempL2))\n",
    "    print('Participants who have 2 or more measurements in both BMI and omics:', len(tempL))\n",
    "    tempDF = tempDF.loc[tempDF['public_client_id'].isin(tempL)]\n",
    "    print('Time-series DF nrows after filtering with the number of measurements:',\n",
    "          len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select the participants who have all covariates (just in case)\n",
    "    tempDF = tempDF.dropna()\n",
    "    print('Time-series DF nrows after filtering with covariates:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Check measurement distribution\n",
    "    sns.set(style='ticks', font='Arial', context='talk')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    sns.distplot(tempDF['days_in_program'], color=yvar_color)\n",
    "    sns.despine()\n",
    "    for knot in [0, 6, 12, 18]:\n",
    "        plt.axvline(x=365.25/12*knot, **{'linestyle':'--', 'color':'gray'})\n",
    "    plt.axvspan(xmin=365.25/12*18, xmax=plt.gca().get_xlim()[1], facecolor='gray', alpha=0.2)\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.show()\n",
    "    \n",
    "    #Check dependent variable distribution\n",
    "    print('Skewness of the dependent variable:', stats.skew(tempDF['log_'+yvar]))\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    sns.distplot(tempDF['log_'+yvar], color=yvar_color)\n",
    "    sns.despine()\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel(yvar+' (log-scale)')\n",
    "    plt.show()\n",
    "    \n",
    "    #Add dummy variables for the linear regression spline for time\n",
    "    nknots = 3\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempL1 = []#2nd piece\n",
    "    tempL2 = []#3rd piece\n",
    "    for days in tempDF['days_in_program'].tolist():\n",
    "        if days < tempL[0]:#[0 month, 6 months)\n",
    "            tempL1.append(0.0)\n",
    "            tempL2.append(0.0)\n",
    "        elif days < tempL[1]:#[6 months, 12 months)\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(0.0)\n",
    "        elif days <= tempL[2]:#[12 months, 18 months]\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(days - tempL[1])\n",
    "        else:\n",
    "            tempL1.append('Error? Check time span setting.')\n",
    "            tempL2.append('Error? Check time span setting.')\n",
    "    tempDF['DaysP2'] = tempL1\n",
    "    tempDF['DaysP3'] = tempL2\n",
    "    \n",
    "    #Standardize numeric independent variables\n",
    "    scaler = StandardScaler(copy=True, with_mean=True, with_std=True)\n",
    "    tempL = ['days_in_program', 'DaysP2', 'DaysP3', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5']\n",
    "    scaler.fit(tempDF[tempL])\n",
    "    tempDF1 = pd.DataFrame(scaler.transform(tempDF[tempL]), index=tempDF.index,\n",
    "                           columns=['DaysZ', 'DaysP2Z', 'DaysP3Z', 'BaseAgeZ', 'PC1Z', 'PC2Z', 'PC3Z', 'PC4Z', 'PC5Z'])\n",
    "    tempDF = pd.merge(tempDF, tempDF1, left_index=True, right_index=True, how='inner')\n",
    "    \n",
    "    #Confirm examples\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    for variable in ['DaysZ', 'BaseAgeZ', 'PC1Z']:\n",
    "        sns.distplot(tempDF[variable], label=variable)\n",
    "    sns.despine()\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel(r'$Z$'+'-score')\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempS = (pd.Series(tempL) - scaler.mean_[0]) / scaler.scale_[0]#Z-score transformation\n",
    "    for knot in tempS.tolist():\n",
    "        plt.axvline(x=knot, **{'linestyle':'--', 'color':'gray'})\n",
    "    plt.legend(bbox_to_anchor=(1, 0.5), loc='center left', borderaxespad=1)\n",
    "    plt.show()\n",
    "    \n",
    "    display(tempDF)\n",
    "    print('')\n",
    "    \n",
    "    tempD1[model_sex] = tempDF\n",
    "    tempD2[model_sex] = scaler#Use later again during predictions\n",
    "\n",
    "#Update\n",
    "tsDF = tempD1['Sex-independent '+yvar+' measurement']\n",
    "scaler = tempD2['Sex-independent '+yvar+' measurement']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4-3. LMM with random intercepts and random slopes for time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Days in the program is fitted as a numeric variable using linear regression splines (cf. Zubair, N. et al. Sci. Rep. 2019).  \n",
    "> ***–> Allows differences in the trajectory of BMI/bBMI change.***\n",
    ">\n",
    "> Random intercepts and random slopes for days in the program are used as random effects (cf. Zubair, N. et al. Sci. Rep. 2019).  \n",
    "> ***–> Allows individual differences in the baseline and longitudinal change over the program.***  \n",
    "> –> Random slope is applied not for the piecewise variable of time but only for the overall variable of time; otherwise, leads to \"LinAlgError: Singular matrix\" due to small sample size per most individuals. Also, it is not so weird to assume the random slope is consistent during the program.  \n",
    ">\n",
    "> ***Eliminate \"Underweight\" class in the baseline obesity class-stratified models.***  \n",
    "> <– Otherwise, leads to \"LinAlgError: Singular matrix\" due to small sample population.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD = {'Sex-independent '+yvar+' measurement':tsDF}\n",
    "tempD1 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    print(model_sex)\n",
    "    tempDF = tempD[model_sex]\n",
    "    \n",
    "    tempD2 = {'LMM1':'Overall LMM',\n",
    "              'LMM2':'Normal class LMM',\n",
    "              'LMM3':'Overweight class LMM',\n",
    "              'LMM4':'Obese class LMM'}\n",
    "    tempD3 = {'LMM1':tempDF,\n",
    "              'LMM2':tempDF.loc[tempDF['BaseBMI_class']=='Normal'],\n",
    "              'LMM3':tempDF.loc[tempDF['BaseBMI_class']=='Overweight'],\n",
    "              'LMM4':tempDF.loc[tempDF['BaseBMI_class']=='Obese']}\n",
    "    fe_formula = 'log_'+yvar+' ~ DaysZ + DaysP2Z + DaysP3Z'\\\n",
    "        ' + C(Sex) + BaseAgeZ + C(Season) + PC1Z + PC2Z + PC3Z + PC4Z + PC5Z'\n",
    "    for model_n in tempD2.keys():\n",
    "        print(' - '+tempD2[model_n])\n",
    "        #Fit LMM\n",
    "        model = sm.MixedLM.from_formula(formula=fe_formula, data=tempD3[model_n],\n",
    "                                        groups='public_client_id',\n",
    "                                        re_formula='DaysZ')#Random intercepts for each group are included as default\n",
    "        t_start = time.time()\n",
    "        #fit_res = model.fit(method=['bfgs', 'lbfgs', 'cg', 'powell', 'nm'])#Back up for failure of convergence\n",
    "        fit_res = model.fit(method=['powell'])#After checking convergence for all models, fix the method through the notebook\n",
    "        t_elapsed = time.time() - t_start\n",
    "        print('    - Elapsed time for fitting LMM:', round(t_elapsed//60), 'min', round(t_elapsed%60, 1), 'sec')\n",
    "        \n",
    "        tempD1[model_sex+' - '+model_n] = fit_res\n",
    "        \n",
    "        #Check the fitted models\n",
    "        display(fit_res.summary())\n",
    "        print('(Note that the estimates for fixed effects and the variances for random effects are shown in the single table.)')\n",
    "        ##Visualize b-coefs and g-coefs\n",
    "        tempDF1 = pd.DataFrame({'Coef':fit_res.params,\n",
    "                                'Coef_ci_l':fit_res.conf_int(alpha=0.05).iloc[:, 0],\n",
    "                                'Coef_ci_h':fit_res.conf_int(alpha=0.05).iloc[:, 1],\n",
    "                                'Pval':fit_res.pvalues})\n",
    "        tempDF1['Coef_ci'] = (tempDF1['Coef_ci_h'] - tempDF1['Coef_ci_l'])/2\n",
    "        tempDF1 = tempDF1.reset_index()\n",
    "        sns.set(style='ticks', font='Arial', context='notebook')\n",
    "        plt.figure(figsize=(4, 0.25*len(tempDF1)))\n",
    "        plt.errorbar(x=tempDF1['Coef'], y=tempDF1['index'], xerr=tempDF1['Coef_ci'],\n",
    "                     fmt='ok', ecolor='k', capsize=5)\n",
    "        plt.xlim(-0.1, 0.1)#Clip\n",
    "        sns.despine()\n",
    "        plt.xlabel(r'$\\beta$'+' or '+r'$\\gamma$'+' coefficient in LMM model\\n(Mean with 95% CI)')\n",
    "        plt.ylabel('')\n",
    "        plt.axvline(x=0, **{'linestyle':'--', 'color':'k'})\n",
    "        for row_i in range(len(tempDF1)):\n",
    "            if tempDF1['Pval'].iloc[row_i]<0.05:\n",
    "                plt.axhspan(ymin=row_i-0.5, ymax=row_i+0.5, facecolor=yvar_color, alpha=0.2, zorder=0)\n",
    "        plt.gca().invert_yaxis()\n",
    "        plt.margins(y=0.01, tight=True)\n",
    "        plt.show()\n",
    "        print('')\n",
    "    print('')\n",
    "\n",
    "LMM1 = tempD1['Sex-independent '+yvar+' measurement - LMM1']\n",
    "LMM2 = tempD1['Sex-independent '+yvar+' measurement - LMM2']\n",
    "LMM3 = tempD1['Sex-independent '+yvar+' measurement - LMM3']\n",
    "LMM4 = tempD1['Sex-independent '+yvar+' measurement - LMM4']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4-4. Estimated transition of biological BMI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> The days in program of the fitted values were different between individuals.  \n",
    "> –> In statsmodels, MixedLMResults.predict() method use only the fixed effects parameters for prediction.  \n",
    "> ***–> To evaluate the variability in the in-sample population, the rondom effects component should be added manually.***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD1 = {'Sex-independent '+yvar+' measurement':tsDF}\n",
    "tempD2 = {'Sex-independent '+yvar+' measurement':scaler}\n",
    "tempD3 = {'Sex-independent '+yvar+' measurement':LMM1}\n",
    "tempD4 = {'Sex-independent '+yvar+' measurement':LMM2}\n",
    "tempD5 = {'Sex-independent '+yvar+' measurement':LMM3}\n",
    "tempD6 = {'Sex-independent '+yvar+' measurement':LMM4}\n",
    "tempD4 = {'Normal':tempD4, 'Overweight':tempD5, 'Obese':tempD6}\n",
    "tempD5 = {}\n",
    "for model_sex in tempD1.keys():\n",
    "    #Prepare DF to impute days_in_program but maintain in-sample population\n",
    "    tempDF = tempD1[model_sex]\n",
    "    tempA1 = tempDF['public_client_id'].unique()\n",
    "    tempA2 = np.arange(0, 365.25/12*month_threshold + 1, 365.25/12*month_threshold/nknots)\n",
    "    tempL = ['Spring', 'Summer', 'Autumn', 'Winter']\n",
    "    tempDF = pd.DataFrame(data={'public_client_id':np.repeat(tempA1, len(tempA2)*len(tempL)),\n",
    "                                'days_in_program':np.tile(np.repeat(tempA2, len(tempL)), len(tempA1)),\n",
    "                                'Season':np.tile(tempL, len(tempA1)*len(tempA2))})\n",
    "    ##Add the covariates unique to the in-sample individuals\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempL = ['public_client_id', 'Sex', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5', 'BaseBMI_class']\n",
    "    tempDF1 = tempDF1[tempL]\n",
    "    tempDF1 = tempDF1.drop_duplicates('public_client_id', keep='first')\n",
    "    tempDF = pd.merge(tempDF, tempDF1, on='public_client_id', how='left')\n",
    "    ##Add dummy variables for the linear regression spline for time\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempL1 = []#2nd piecewise\n",
    "    tempL2 = []#3rd piecewise\n",
    "    for days in tempDF['days_in_program'].tolist():\n",
    "        if days < tempL[0]:#[0 month, 6 months)\n",
    "            tempL1.append(0.0)\n",
    "            tempL2.append(0.0)\n",
    "        elif days < tempL[1]:#[6 months, 12 months)\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(0.0)\n",
    "        elif days <= tempL[2]:#[12 months, 18 months]\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(days - tempL[1])\n",
    "        else:\n",
    "            tempL1.append('Error? Check time span setting.')\n",
    "            tempL2.append('Error? Check time span setting.')\n",
    "    tempDF['DaysP2'] = tempL1\n",
    "    tempDF['DaysP3'] = tempL2\n",
    "    ##Add a dummy index to use when merging later\n",
    "    tempDF['KeyIndex'] = tempDF['public_client_id'] + '_day' + tempDF['days_in_program'].astype(str)\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    ##Standardize numeric independent variables using the already fitted StandardScaler\n",
    "    scaler = tempD2[model_sex]\n",
    "    tempL = ['days_in_program', 'DaysP2', 'DaysP3', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5']\n",
    "    tempDF1 = pd.DataFrame(scaler.transform(tempDF[tempL]), index=tempDF.index,\n",
    "                           columns=['DaysZ', 'DaysP2Z', 'DaysP3Z', 'BaseAgeZ', 'PC1Z', 'PC2Z', 'PC3Z', 'PC4Z', 'PC5Z'])\n",
    "    tempDF = pd.merge(tempDF, tempDF1, left_index=True, right_index=True, how='inner')\n",
    "    \n",
    "    #Model 1 prediction\n",
    "    fit_res = tempD3[model_sex]\n",
    "    ##Predicted component from the fixed effects\n",
    "    tempS = fit_res.predict(tempDF)#Automatically recognize the column name in formula\n",
    "    ##Create design matrix for the random effects\n",
    "    tempDF1 = pd.DataFrame({'Intercept_re':np.repeat(1, len(tempDF)),\n",
    "                            'Slope_re':tempDF['DaysZ']})\n",
    "    tempD = fit_res.random_effects#Dictionary of key=group label and value=pd.Series\n",
    "    ##Predicted component from the random effects\n",
    "    tempL = [np.dot(tempDF1.iloc[row_i], tempD[group_n])\n",
    "             for (row_i, group_n) in enumerate(tempDF['public_client_id'].tolist())]\n",
    "    ##Total prediction for in-sample individual\n",
    "    tempS = tempS + tempL\n",
    "    ##Convert to original scale\n",
    "    tempDF[yvar+'_predicted1'] = np.e**tempS\n",
    "    \n",
    "    #Model 2-4 prediction\n",
    "    for model_i, bmi_class in enumerate(tempD4.keys(), start=2):\n",
    "        tempD = tempD4[bmi_class]\n",
    "        fit_res = tempD[model_sex]\n",
    "        ##Prepare the target class\n",
    "        tempDF1 = tempDF.loc[tempDF['BaseBMI_class']==bmi_class]\n",
    "        ##Predicted component from the fixed effects\n",
    "        tempS = fit_res.predict(tempDF1)#Automatically recognize the column name in formula\n",
    "        ##Create design matrix for the random effects\n",
    "        tempDF2 = pd.DataFrame({'Intercept_re':np.repeat(1, len(tempDF1)),\n",
    "                                'Slope_re':tempDF1['DaysZ']})\n",
    "        tempD = fit_res.random_effects#Dictionary of key=group label and value=pd.Series\n",
    "        ##Predicted component from the random effects\n",
    "        tempL = [np.dot(tempDF2.iloc[row_i], tempD[group_n])\n",
    "                 for (row_i, group_n) in enumerate(tempDF1['public_client_id'].tolist())]\n",
    "        ##Total prediction for in-sample individual\n",
    "        tempS = tempS + tempL\n",
    "        ##Convert to original scale\n",
    "        tempDF1[yvar+'_predicted'+str(model_i)] = np.e**tempS\n",
    "        ##Add the predictions to the working DF\n",
    "        tempDF = pd.merge(tempDF, tempDF1[yvar+'_predicted'+str(model_i)],\n",
    "                          left_index=True, right_index=True, how='outer')#NaN for out-of-class\n",
    "    \n",
    "    #Flatten effects of the dummy season variable\n",
    "    tempDF = tempDF.groupby(by=['public_client_id', 'days_in_program'])\n",
    "    tempDF = tempDF.agg({yvar+'_predicted1':'mean',\n",
    "                         yvar+'_predicted2':'mean',\n",
    "                         yvar+'_predicted3':'mean',\n",
    "                         yvar+'_predicted4':'mean'})\n",
    "    tempDF = tempDF.reset_index()\n",
    "    ##Recover the baseline BMI class to use later\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempL = ['public_client_id', 'BaseBMI_class']\n",
    "    tempDF1 = tempDF1[tempL].drop_duplicates('public_client_id', keep='first')\n",
    "    tempDF = pd.merge(tempDF, tempDF1, on='public_client_id', how='left')\n",
    "    tempDF['KeyIndex'] = tempDF['public_client_id'] + '_day' + tempDF['days_in_program'].astype(str)\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    \n",
    "    print(model_sex)\n",
    "    display(tempDF)\n",
    "    display(tempDF.describe(include='all'))\n",
    "    \n",
    "    tempD5[model_sex] = tempDF\n",
    "\n",
    "measLMM = tempD5['Sex-independent '+yvar+' measurement']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check visually\n",
    "tempD1 = {'Sex-independent '+yvar+' measurement':tsDF}\n",
    "tempD2 = {'Sex-independent '+yvar+' measurement':measLMM}\n",
    "for model_sex in tempD1.keys():\n",
    "    print(model_sex)\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempDF2 = tempD2[model_sex]\n",
    "    \n",
    "    #Plot1\n",
    "    print(' - Model 1: Overall LMM')\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(5, 3))\n",
    "    sns.lineplot(data=tempDF1, x='days_in_program', y=yvar,\n",
    "                 units='public_client_id', estimator=None, color='gray', lw=1, alpha=0.1, legend=None)\n",
    "    sns.scatterplot(data=tempDF1, x='days_in_program', y=yvar,\n",
    "                    color='gray', edgecolor='k', s=20, alpha=0.3)\n",
    "    p = sns.lineplot(data=tempDF2, x='days_in_program', y=yvar+'_predicted1',\n",
    "                     estimator='mean', ci=95, color=yvar_color)\n",
    "    p.set(xlim=(0, 365.25/12*month_threshold), xticks=np.arange(0, 365.25/12*month_threshold, 100),\n",
    "          ylim=(17.5, 42.5), yticks=np.arange(20, 40.1, 5))\n",
    "    sns.despine()\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.ylabel(yvar+' [kg/m'+r'$^2$'+']')\n",
    "    plt.show()\n",
    "    print('')\n",
    "    \n",
    "    #Plot2\n",
    "    print(' - Model 2-4: Baseline BMI class-stratified LMM')\n",
    "    tempD = {'Normal':'green', 'Overweight':'orange', 'Obese':'red'}\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(5, 3))\n",
    "    p = sns.scatterplot(data=tempDF1[tempDF1['BaseBMI_class']!='Underweight'],\n",
    "                        x='days_in_program', y=yvar,\n",
    "                        hue='BaseBMI_class', hue_order=tempD.keys(), palette=tempD,\n",
    "                        edgecolor='k', s=20, alpha=0.3)\n",
    "    for model_i, bmi_class in enumerate(tempD.keys(), start=2):\n",
    "        sns.lineplot(data=tempDF2[tempDF2['BaseBMI_class']==bmi_class],\n",
    "                     x='days_in_program', y=yvar+'_predicted'+str(model_i),\n",
    "                     color=tempD[bmi_class], estimator='mean', ci=95, ax=p)\n",
    "    p.set(xlim=(0, 365.25/12*month_threshold), xticks=np.arange(0, 365.25/12*month_threshold, 100),\n",
    "          ylim=(17.5, 42.5), yticks=np.arange(20, 40.1, 5))\n",
    "    sns.despine()\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.ylabel(yvar+' [kg/m'+r'$^2$'+']')\n",
    "    plt.legend(bbox_to_anchor=(1, 0.5), loc='center left', borderaxespad=1)\n",
    "    plt.show()\n",
    "    print('')\n",
    "    \n",
    "    print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Comparison of estimates b/w BMI types"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5-1. Overall cohort and baseline BMI class stratification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#For presentation\n",
    "tempD = {'BMI':'k', 'MetBMI':'b', 'ProtBMI':'r', 'ChemBMI':'g'}\n",
    "tempL1 = [metLMM_BS, protLMM_BS, chemLMM_BS]\n",
    "tempL2 = [metLMM_FM, protLMM_FM, chemLMM_FM]\n",
    "tempD1 = {'Sex-mixed bBMI model':tempL1, 'Sex-stratified bBMI model':tempL2}\n",
    "tempD2 = {'Sex-mixed bBMI model':'BothSex', 'Sex-stratified bBMI model':'FemaleMale'}\n",
    "\n",
    "for model_sex in tempD1.keys():\n",
    "    print(model_sex)\n",
    "    \n",
    "    #Merge prediction DFs\n",
    "    tempDF = measLMM\n",
    "    tempL = tempD1[model_sex]\n",
    "    for tempDF1 in tempL:\n",
    "        tempDF1 = tempDF1.loc[:, ~tempDF1.columns.isin(tempDF.columns.tolist())]\n",
    "        tempDF = pd.merge(tempDF, tempDF1, left_index=True, right_index=True, how='outer')\n",
    "    \n",
    "    #Calculate longitudinal change\n",
    "    tempL = tempDF.loc[:, tempDF.columns.str.contains('_predicted')].columns.tolist()\n",
    "    for col_n in tempL:\n",
    "        tempDF1 = tempDF[['public_client_id', 'days_in_program', col_n]].dropna()\n",
    "        #Retrieve the estimate at days = 0 for each participant\n",
    "        tempDF2 = tempDF1.loc[tempDF1['days_in_program']==0.0]\n",
    "        tempS = tempDF2.reset_index().set_index('public_client_id')[col_n]\n",
    "        tempS.name = 'day0'\n",
    "        #Merge\n",
    "        tempDF1 = pd.merge(tempDF1, tempS, left_on='public_client_id', right_index=True, how='left')\n",
    "        #Calculate longitudinal change [%]\n",
    "        label = col_n.replace('_predicted', 'change_predicted')\n",
    "        tempDF1[label] = (tempDF1[col_n] - tempDF1['day0']) / tempDF1['day0'] * 100\n",
    "        #Add to the working DF\n",
    "        tempDF = pd.merge(tempDF, tempDF1[label], left_index=True, right_index=True, how='left')\n",
    "    \n",
    "    #Model 1\n",
    "    print(' - Model 1: Overall LMM')\n",
    "    ##Prepare DF\n",
    "    tempL = tempDF.loc[:, tempDF.columns.str.contains('change_predicted1')].columns.tolist()\n",
    "    tempDF1 = tempDF.reset_index().melt(var_name='LMM', value_name='PredictedChange', value_vars=tempL,\n",
    "                                        id_vars=['public_client_id', 'days_in_program', 'BaseBMI_class'])\n",
    "    tempDF1 = tempDF1.dropna()\n",
    "    tempDF1['LMM'] = tempDF1['LMM'].str.replace('change_predicted1', '')\n",
    "    ##Check sample size\n",
    "    tempL = []\n",
    "    for yvar in tempD.keys():\n",
    "        tempA = tempDF1['public_client_id'].loc[tempDF1['LMM']==yvar].unique()\n",
    "        tempL.append(yvar+' (n = '+str(len(tempA))+')')\n",
    "    print('  -> Confirm each sample size:', tempL)\n",
    "    ##Visualization\n",
    "    sns.set(style='ticks', font='Arial', context='talk')\n",
    "    plt.figure(figsize=(4.5, 3))\n",
    "    p = sns.lineplot(data=tempDF1, x='days_in_program', y='PredictedChange',\n",
    "                     hue='LMM', hue_order=tempD.keys(), palette=tempD, estimator='mean', ci=95)\n",
    "    p.set(xlim=(0, 365.25/12*12), xticks=np.arange(0, 365.25/12*12+1, 90),\n",
    "          ylim=(-7, 1), yticks=np.arange(-6, 0.1, 2))\n",
    "    sns.despine()\n",
    "    p.set_title('Overall cohort', {'fontsize':'large'})\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.ylabel('Estimated change [%]')\n",
    "    plt.axhline(y=0, **{'linestyle':'--', 'color':'k'}, zorder=0)\n",
    "    #for knot in range(nknots-1):\n",
    "    #    if knot%2 == 1:\n",
    "    #        plt.axvspan(xmin=365.25/12*month_threshold/nknots*knot,\n",
    "    #                    xmax=365.25/12*month_threshold/nknots*(knot+1),\n",
    "    #                    facecolor='gray', alpha=0.2, zorder=-1)\n",
    "    ###Generate standard Matplotlib legend manually\n",
    "    tempL = []\n",
    "    for yvar in tempD.keys():\n",
    "        tempL.append(mlines.Line2D([], [], color=tempD[yvar], label=yvar, linewidth=3))\n",
    "    plt.legend(handles=tempL, fontsize='medium', title='LMM estimate', title_fontsize='large',\n",
    "               bbox_to_anchor=(0.5, 0), loc='upper center', ncol=2, borderaxespad=4,\n",
    "               handlelength=1.5, handletextpad=0.5, columnspacing=1.0, frameon=True)\n",
    "    ###Save\n",
    "    fileDir = './ExportFigures/'\n",
    "    ipynbName = '220806_Multiomics-BMI-NatMed1stRevision_BMI-longitudinal-LMM-ver2_'\n",
    "    fileName = 'overall-'+tempD2[model_sex]+'.tif'\n",
    "    plt.gcf().savefig(fileDir+ipynbName+fileName, dpi=300, bbox_inches='tight', pad_inches=0.04,\n",
    "                      pil_kwargs={'compression':'tiff_lzw'})\n",
    "    plt.show()\n",
    "    \n",
    "    #Model 2-4\n",
    "    print(' - Model 2-4: Baseline BMI class-stratified LMM')\n",
    "    tempL1 = ['Normal', 'Overweight', 'Obese']\n",
    "    ##Prepare DF\n",
    "    tempL = tempDF.loc[:, tempDF.columns.str.contains('change_predicted[234]')].columns.tolist()\n",
    "    tempDF1 = tempDF.reset_index().melt(var_name='LMM', value_name='PredictedChange', value_vars=tempL,\n",
    "                                        id_vars=['public_client_id', 'days_in_program', 'BaseBMI_class'])\n",
    "    tempDF1 = tempDF1.dropna()\n",
    "    tempDF1['LMM'] = tempDF1['LMM'].str.replace('change_predicted[234]', '')\n",
    "    ##Check sample size\n",
    "    for bmi_class in tempL1:\n",
    "        print('   - '+bmi_class+' class based on the baseline BMI')\n",
    "        tempDF2 = tempDF1.loc[tempDF1['BaseBMI_class']==bmi_class]\n",
    "        tempL = []\n",
    "        for yvar in tempD.keys():\n",
    "            tempA = tempDF2['public_client_id'].loc[tempDF2['LMM']==yvar].unique()\n",
    "            tempL.append(yvar+' (n = '+str(len(tempA))+')')\n",
    "        print('     -> Confirm each sample size:', tempL)\n",
    "    ##Visualization\n",
    "    sns.set(style='ticks', font='Arial', context='talk')\n",
    "    p = sns.relplot(kind='line', data=tempDF1, x='days_in_program', y='PredictedChange',\n",
    "                    row='BaseBMI_class', row_order=tempL1,\n",
    "                    hue='LMM', hue_order=tempD.keys(), palette=tempD, estimator='mean', ci=95,\n",
    "                    height=3, aspect=4.5/3+0.125, legend=False)#Manual adjustment\n",
    "    p.set(xlim=(0, 365.25/12*12), xticks=np.arange(0, 365.25/12*12+1, 90),\n",
    "          ylim=(-7, 1), yticks=np.arange(-6, 0.1, 2))\n",
    "    tempL = [bmi_class+' class' for bmi_class in tempL1]\n",
    "    for ax, ax_title in zip(p.axes.ravel(), tempL):\n",
    "        #Draw initial line\n",
    "        ax.axhline(y=0, **{'linestyle':'--', 'color':'k'}, zorder=0)\n",
    "        #Draw span b/w knots\n",
    "        #for knot in range(nknots-1):\n",
    "        #    if knot%2 == 1:\n",
    "        #        ax.axvspan(xmin=365.25/12*month_threshold/nknots*knot,\n",
    "        #                   xmax=365.25/12*month_threshold/nknots*(knot+1),\n",
    "        #                   facecolor='gray', alpha=0.2, zorder=-1)\n",
    "        #Change facet label\n",
    "        ax.set_title(ax_title, {'fontsize':'large'})\n",
    "    ###Reset and generate common axis title\n",
    "    for row_i, ax in enumerate(p.axes.ravel()):\n",
    "        if row_i == 2:\n",
    "            ax.set_xlabel('Days in program')\n",
    "            ax.set_ylabel('')\n",
    "        elif row_i == 1:\n",
    "            ax.set_xlabel('')\n",
    "            ax.set_ylabel('Estimated change [%]')\n",
    "        else:\n",
    "            ax.set_xlabel('')\n",
    "            ax.set_ylabel('')\n",
    "    plt.tight_layout()\n",
    "    ###Generate standard Matplotlib legend manually\n",
    "    tempL = []\n",
    "    for yvar in tempD.keys():\n",
    "        tempL.append(mlines.Line2D([], [], color=tempD[yvar], label=yvar, linewidth=3))\n",
    "    plt.legend(handles=tempL, fontsize='medium', title='LMM estimate', title_fontsize='large',\n",
    "               bbox_to_anchor=(0.5, 0), loc='upper center', ncol=2, borderaxespad=4,\n",
    "               handlelength=1.5, handletextpad=0.5, columnspacing=1.0, frameon=True)\n",
    "    ###Save\n",
    "    fileDir = './ExportFigures/'\n",
    "    ipynbName = '220806_Multiomics-BMI-NatMed1stRevision_BMI-longitudinal-LMM-ver2_'\n",
    "    fileName = 'baselineBMIclass-'+tempD2[model_sex]+'.tif'\n",
    "    plt.gcf().savefig(fileDir+ipynbName+fileName, dpi=300, bbox_inches='tight', pad_inches=0.04,\n",
    "                      pil_kwargs={'compression':'tiff_lzw'})\n",
    "    plt.show()\n",
    "    \n",
    "    print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Misclassification (vs. MetBMI class)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vs_bbmi = 'MetBMI'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6-1. Prepare misclassification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Prepare misclassification\n",
    "tempD = {'Sex-mixed '+vs_bbmi+' model':'BothSex', 'Sex-stratified '+vs_bbmi+' model':'FemaleMale'}\n",
    "tempD1 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    #Import time-series biological BMI\n",
    "    fileDir = './ExportData/'\n",
    "    ipynbName = '220805_Multiomics-BMI-NatMed1stRevision_BMI-longitudinal-LASSO_'\n",
    "    fileName = vs_bbmi+'-'+tempD[model_sex]+'.tsv'\n",
    "    tempDF = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    \n",
    "    #Retrieve baseline BMI/bBMI class\n",
    "    tempL = ['public_client_id', 'BaseBMI_class', 'Base'+vs_bbmi+'_class']\n",
    "    tempDF = tempDF[tempL].drop_duplicates('public_client_id', keep='first')\n",
    "    tempDF = tempDF.set_index('public_client_id')\n",
    "    \n",
    "    #Misclassification\n",
    "    tempL = []\n",
    "    for row_i in range(len(tempDF)):\n",
    "        bmi_class = tempDF['BaseBMI_class'].iloc[row_i]\n",
    "        vs_bbmi_class = tempDF['Base'+vs_bbmi+'_class'].iloc[row_i]\n",
    "        if bmi_class==vs_bbmi_class:\n",
    "            tempL.append('Matched')\n",
    "        else:\n",
    "            tempL.append('Mismatched')\n",
    "    tempDF['vs_Base'+vs_bbmi+'_class'] = tempL\n",
    "    \n",
    "    tempD1[model_sex] = tempDF\n",
    "    \n",
    "    #Check\n",
    "    print(model_sex)\n",
    "    display(tempDF)\n",
    "    tempS1 = tempDF['BaseBMI_class'].value_counts()\n",
    "    tempS2 = tempDF['Base'+vs_bbmi+'_class'].value_counts()\n",
    "    tempDF1 = pd.merge(tempS1, tempS2, left_index=True, right_index=True, how='outer')\n",
    "    display(tempDF1)\n",
    "    tempS = tempDF.groupby(by='BaseBMI_class')['vs_Base'+vs_bbmi+'_class'].value_counts()\n",
    "    display(tempS)\n",
    "    print('')\n",
    "\n",
    "misclassDF_BS = tempD1['Sex-mixed '+vs_bbmi+' model']\n",
    "misclassDF_FM = tempD1['Sex-stratified '+vs_bbmi+' model']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6-2. Metabolomics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Repeat the above code, while adding the misclassification info.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yvar = 'MetBMI'\n",
    "yvar_color = 'b'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6-2-1. Prepare the time-series DF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD1 = {'Sex-mixed '+yvar+' model':'BothSex', 'Sex-stratified '+yvar+' model':'FemaleMale'}\n",
    "tempD2 = {'Sex-mixed '+yvar+' model':misclassDF_BS, 'Sex-stratified '+yvar+' model':misclassDF_FM}\n",
    "tempD3 = {}\n",
    "for model_sex in tempD1.keys():\n",
    "    #Import time-series biological BMI\n",
    "    fileDir = './ExportData/'\n",
    "    ipynbName = '220805_Multiomics-BMI-NatMed1stRevision_BMI-longitudinal-LASSO_'\n",
    "    fileName = yvar+'-'+tempD1[model_sex]+'.tsv'\n",
    "    tempDF = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    \n",
    "    #Add the misclassification\n",
    "    tempDF1 = tempD2[model_sex]\n",
    "    tempDF = pd.merge(tempDF, tempDF1['vs_Base'+vs_bbmi+'_class'],\n",
    "                      left_on='public_client_id', right_index=True, how='left')\n",
    "    \n",
    "    tempD3[model_sex] = tempDF\n",
    "    \n",
    "    print(model_sex)\n",
    "    display(tempDF)\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    print('')\n",
    "\n",
    "tsDF_BS = tempD3['Sex-mixed '+yvar+' model']\n",
    "tsDF_FM = tempD3['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6-2-2. Prepare DFs for LMM with the linear regresion spline for time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD1 = {}\n",
    "tempD2 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    print(model_sex)\n",
    "    tempDF = tempD[model_sex]\n",
    "    print('Time-series DF nrows before filtering:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select measurements for the longitudinal analysis\n",
    "    month_threshold = 18\n",
    "    tempDF = tempDF.loc[tempDF['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    print('Time-series DF nrows after filtering with the taget period:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select the cohort for the longitudinal analysis\n",
    "    ##Select participants who have 2 or more measuremnts in BMI\n",
    "    fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "    ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "    fileName = 'time-series-bmiDF-without-imputation_final-cohort.tsv'\n",
    "    tempDF1 = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF1 = tempDF1.set_index('KeyIndex')\n",
    "    tempDF1 = tempDF1.loc[tempDF1['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    tempS = tempDF1['public_client_id'].value_counts()\n",
    "    tempL1 = tempS.loc[tempS>=2].index.tolist()\n",
    "    ##Select participants who have 2 or more measuremnts in omics\n",
    "    fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "    ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "    fileName = 'time-series-combiDF-without-imputation_final-cohort.tsv'\n",
    "    tempDF1 = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF1 = tempDF1.set_index('KeyIndex')\n",
    "    tempDF1 = tempDF1.loc[tempDF1['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    tempS = tempDF1['public_client_id'].value_counts()\n",
    "    tempL2 = tempS.loc[tempS>=2].index.tolist()\n",
    "    ##Select participants who have 2 or more measuremnts in both BMI and omics\n",
    "    tempL = list(set(tempL1) & set(tempL2))\n",
    "    print('Participants who have 2 or more measurements in both BMI and omics:', len(tempL))\n",
    "    tempDF = tempDF.loc[tempDF['public_client_id'].isin(tempL)]\n",
    "    print('Time-series DF nrows after filtering with the number of measurements:',\n",
    "          len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select the participants who have all covariates (just in case)\n",
    "    tempDF = tempDF.dropna()\n",
    "    print('Time-series DF nrows after filtering with covariates:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Check measurement distribution\n",
    "    sns.set(style='ticks', font='Arial', context='talk')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    sns.distplot(tempDF['days_in_program'], color=yvar_color)\n",
    "    sns.despine()\n",
    "    for knot in [0, 6, 12, 18]:\n",
    "        plt.axvline(x=365.25/12*knot, **{'linestyle':'--', 'color':'gray'})\n",
    "    plt.axvspan(xmin=365.25/12*18, xmax=plt.gca().get_xlim()[1], facecolor='gray', alpha=0.2)\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.show()\n",
    "    \n",
    "    #Check dependent variable distribution\n",
    "    print('Skewness of the dependent variable:', stats.skew(tempDF['log_'+yvar]))\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    sns.distplot(tempDF['log_'+yvar], color=yvar_color)\n",
    "    sns.despine()\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel(yvar+' (log-scale)')\n",
    "    plt.show()\n",
    "    \n",
    "    #Add dummy variables for the linear regression spline for time\n",
    "    nknots = 3\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempL1 = []#2nd piece\n",
    "    tempL2 = []#3rd piece\n",
    "    for days in tempDF['days_in_program'].tolist():\n",
    "        if days < tempL[0]:#[0 month, 6 months)\n",
    "            tempL1.append(0.0)\n",
    "            tempL2.append(0.0)\n",
    "        elif days < tempL[1]:#[6 months, 12 months)\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(0.0)\n",
    "        elif days <= tempL[2]:#[12 months, 18 months]\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(days - tempL[1])\n",
    "        else:\n",
    "            tempL1.append('Error? Check time span setting.')\n",
    "            tempL2.append('Error? Check time span setting.')\n",
    "    tempDF['DaysP2'] = tempL1\n",
    "    tempDF['DaysP3'] = tempL2\n",
    "    \n",
    "    #Standardize numeric independent variables\n",
    "    scaler = StandardScaler(copy=True, with_mean=True, with_std=True)\n",
    "    tempL = ['days_in_program', 'DaysP2', 'DaysP3', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5']\n",
    "    scaler.fit(tempDF[tempL])\n",
    "    tempDF1 = pd.DataFrame(scaler.transform(tempDF[tempL]), index=tempDF.index,\n",
    "                           columns=['DaysZ', 'DaysP2Z', 'DaysP3Z', 'BaseAgeZ', 'PC1Z', 'PC2Z', 'PC3Z', 'PC4Z', 'PC5Z'])\n",
    "    tempDF = pd.merge(tempDF, tempDF1, left_index=True, right_index=True, how='inner')\n",
    "    \n",
    "    #Confirm examples\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    for variable in ['DaysZ', 'BaseAgeZ', 'PC1Z']:\n",
    "        sns.distplot(tempDF[variable], label=variable)\n",
    "    sns.despine()\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel(r'$Z$'+'-score')\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempS = (pd.Series(tempL) - scaler.mean_[0]) / scaler.scale_[0]#Z-score transformation\n",
    "    for knot in tempS.tolist():\n",
    "        plt.axvline(x=knot, **{'linestyle':'--', 'color':'gray'})\n",
    "    plt.legend(bbox_to_anchor=(1, 0.5), loc='center left', borderaxespad=1)\n",
    "    plt.show()\n",
    "    \n",
    "    #Sort to make 'Matched' the standard (bcoef = 0) in LMM (Matched < Mismatched)\n",
    "    tempDF = tempDF.sort_values(by=['vs_Base'+vs_bbmi+'_class'], ascending=True)\n",
    "    \n",
    "    display(tempDF)\n",
    "    print('')\n",
    "    \n",
    "    tempD1[model_sex] = tempDF\n",
    "    tempD2[model_sex] = scaler#Use later again during predictions\n",
    "\n",
    "#Update\n",
    "tsDF_BS = tempD1['Sex-mixed '+yvar+' model']\n",
    "scaler_BS = tempD2['Sex-mixed '+yvar+' model']\n",
    "tsDF_FM = tempD1['Sex-stratified '+yvar+' model']\n",
    "scaler_FM = tempD2['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6-2-3. LMM with random intercepts and random slopes for time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Days in the program is fitted as a numeric variable using linear regression splines (cf. Zubair, N. et al. Sci. Rep. 2019).  \n",
    "> ***–> Allows differences in the trajectory of BMI/bBMI change.***\n",
    ">\n",
    "> Random intercepts and random slopes for days in the program are used as random effects (cf. Zubair, N. et al. Sci. Rep. 2019).  \n",
    "> ***–> Allows individual differences in the baseline and longitudinal change over the program.***  \n",
    "> –> Random slope is applied not for the piecewise variable of time but only for the overall variable of time; otherwise, leads to \"LinAlgError: Singular matrix\" due to small sample size per most individuals. Also, it is not so weird to assume the random slope is consistent during the program.  \n",
    ">\n",
    "> ***Eliminate \"Underweight\" class in the baseline obesity class-stratified models.***  \n",
    "> <– Otherwise, leads to \"LinAlgError: Singular matrix\" due to small sample population.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD1 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    print(model_sex)\n",
    "    tempDF = tempD[model_sex]\n",
    "    \n",
    "    tempD2 = {'LMM5':'Normal class LMM',\n",
    "              'LMM6':'Obese class LMM'}\n",
    "    tempD3 = {'LMM5':tempDF.loc[tempDF['BaseBMI_class']=='Normal'],\n",
    "              'LMM6':tempDF.loc[tempDF['BaseBMI_class']=='Obese']}\n",
    "    fe_formula = 'log_'+yvar+' ~ DaysZ + DaysP2Z + DaysP3Z'\\\n",
    "        ' + C(Sex) + BaseAgeZ + C(Season) + PC1Z + PC2Z + PC3Z + PC4Z + PC5Z'\\\n",
    "        ' + C(vs_Base'+vs_bbmi+'_class) + DaysZ:C(vs_Base'+vs_bbmi+'_class)'\\\n",
    "        ' + DaysP2Z:C(vs_Base'+vs_bbmi+'_class) + DaysP3Z:C(vs_Base'+vs_bbmi+'_class)'\n",
    "    for model_n in tempD2.keys():\n",
    "        print(' - '+tempD2[model_n])\n",
    "        #Fit LMM\n",
    "        model = sm.MixedLM.from_formula(formula=fe_formula, data=tempD3[model_n],\n",
    "                                        groups='public_client_id',\n",
    "                                        re_formula='DaysZ')#Random intercepts for each group are included as default\n",
    "        t_start = time.time()\n",
    "        #fit_res = model.fit(method=['bfgs', 'lbfgs', 'cg', 'powell', 'nm'])#Back up for failure of convergence\n",
    "        fit_res = model.fit(method=['powell'])#After checking convergence for all models, fix the method through the notebook\n",
    "        t_elapsed = time.time() - t_start\n",
    "        print('    - Elapsed time for fitting LMM:', round(t_elapsed//60), 'min', round(t_elapsed%60, 1), 'sec')\n",
    "        \n",
    "        tempD1[model_sex+' - '+model_n] = fit_res\n",
    "        \n",
    "        #Check the fitted models\n",
    "        display(fit_res.summary())\n",
    "        print('(Note that the estimates for fixed effects and the variances for random effects are shown in the single table.)')\n",
    "        ##Visualize b-coefs and g-coefs\n",
    "        tempDF1 = pd.DataFrame({'Coef':fit_res.params,\n",
    "                                'Coef_ci_l':fit_res.conf_int(alpha=0.05).iloc[:, 0],\n",
    "                                'Coef_ci_h':fit_res.conf_int(alpha=0.05).iloc[:, 1],\n",
    "                                'Pval':fit_res.pvalues})\n",
    "        tempDF1['Coef_ci'] = (tempDF1['Coef_ci_h'] - tempDF1['Coef_ci_l'])/2\n",
    "        tempDF1 = tempDF1.reset_index()\n",
    "        sns.set(style='ticks', font='Arial', context='notebook')\n",
    "        plt.figure(figsize=(4, 0.25*len(tempDF1)))\n",
    "        plt.errorbar(x=tempDF1['Coef'], y=tempDF1['index'], xerr=tempDF1['Coef_ci'],\n",
    "                     fmt='ok', ecolor='k', capsize=5)\n",
    "        plt.xlim(-0.1, 0.1)#Clip\n",
    "        sns.despine()\n",
    "        plt.xlabel(r'$\\beta$'+' or '+r'$\\gamma$'+' coefficient in LMM model\\n(Mean with 95% CI)')\n",
    "        plt.ylabel('')\n",
    "        plt.axvline(x=0, **{'linestyle':'--', 'color':'k'})\n",
    "        for row_i in range(len(tempDF1)):\n",
    "            if tempDF1['Pval'].iloc[row_i]<0.05:\n",
    "                plt.axhspan(ymin=row_i-0.5, ymax=row_i+0.5, facecolor=yvar_color, alpha=0.2, zorder=0)\n",
    "        plt.gca().invert_yaxis()\n",
    "        plt.margins(y=0.01, tight=True)\n",
    "        plt.show()\n",
    "        print('')\n",
    "    print('')\n",
    "\n",
    "LMM5_BS = tempD1['Sex-mixed '+yvar+' model - LMM5']\n",
    "LMM6_BS = tempD1['Sex-mixed '+yvar+' model - LMM6']\n",
    "LMM5_FM = tempD1['Sex-stratified '+yvar+' model - LMM5']\n",
    "LMM6_FM = tempD1['Sex-stratified '+yvar+' model - LMM6']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6-2-4. Estimated transition of biological BMI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> The days in program of the fitted values were different between individuals.  \n",
    "> –> In statsmodels, MixedLMResults.predict() method use only the fixed effects parameters for prediction.  \n",
    "> ***–> To evaluate the variability in the in-sample population, the rondom effects component should be added manually.***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD1 = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD2 = {'Sex-mixed '+yvar+' model':scaler_BS, 'Sex-stratified '+yvar+' model':scaler_FM}\n",
    "tempD3 = {'Sex-mixed '+yvar+' model':LMM5_BS, 'Sex-stratified '+yvar+' model':LMM5_FM}\n",
    "tempD4 = {'Sex-mixed '+yvar+' model':LMM6_BS, 'Sex-stratified '+yvar+' model':LMM6_FM}\n",
    "tempD3 = {'Normal':tempD3, 'Obese':tempD4}\n",
    "tempD4 = {}\n",
    "for model_sex in tempD1.keys():\n",
    "    #Prepare DF to impute days_in_program but maintain in-sample population\n",
    "    tempDF = tempD1[model_sex]\n",
    "    tempA1 = tempDF['public_client_id'].unique()\n",
    "    tempA2 = np.arange(0, 365.25/12*month_threshold + 1, 365.25/12*month_threshold/nknots)\n",
    "    tempL = ['Spring', 'Summer', 'Autumn', 'Winter']\n",
    "    tempDF = pd.DataFrame(data={'public_client_id':np.repeat(tempA1, len(tempA2)*len(tempL)),\n",
    "                                'days_in_program':np.tile(np.repeat(tempA2, len(tempL)), len(tempA1)),\n",
    "                                'Season':np.tile(tempL, len(tempA1)*len(tempA2))})\n",
    "    ##Add the covariates unique to the in-sample individuals\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempL = ['public_client_id', 'Sex', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5',\n",
    "             'BaseBMI_class', 'vs_Base'+vs_bbmi+'_class']\n",
    "    tempDF1 = tempDF1[tempL]\n",
    "    tempDF1 = tempDF1.drop_duplicates('public_client_id', keep='first')\n",
    "    tempDF = pd.merge(tempDF, tempDF1, on='public_client_id', how='left')\n",
    "    ##Add dummy variables for the linear regression spline for time\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempL1 = []#2nd piecewise\n",
    "    tempL2 = []#3rd piecewise\n",
    "    for days in tempDF['days_in_program'].tolist():\n",
    "        if days < tempL[0]:#[0 month, 6 months)\n",
    "            tempL1.append(0.0)\n",
    "            tempL2.append(0.0)\n",
    "        elif days < tempL[1]:#[6 months, 12 months)\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(0.0)\n",
    "        elif days <= tempL[2]:#[12 months, 18 months]\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(days - tempL[1])\n",
    "        else:\n",
    "            tempL1.append('Error? Check time span setting.')\n",
    "            tempL2.append('Error? Check time span setting.')\n",
    "    tempDF['DaysP2'] = tempL1\n",
    "    tempDF['DaysP3'] = tempL2\n",
    "    ##Add a dummy index to use when merging later\n",
    "    tempDF['KeyIndex'] = tempDF['public_client_id'] + '_day' + tempDF['days_in_program'].astype(str)\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    ##Standardize numeric independent variables using the already fitted StandardScaler\n",
    "    scaler = tempD2[model_sex]\n",
    "    tempL = ['days_in_program', 'DaysP2', 'DaysP3', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5']\n",
    "    tempDF1 = pd.DataFrame(scaler.transform(tempDF[tempL]), index=tempDF.index,\n",
    "                           columns=['DaysZ', 'DaysP2Z', 'DaysP3Z', 'BaseAgeZ', 'PC1Z', 'PC2Z', 'PC3Z', 'PC4Z', 'PC5Z'])\n",
    "    tempDF = pd.merge(tempDF, tempDF1, left_index=True, right_index=True, how='inner')\n",
    "    \n",
    "    #Model 5-6 prediction\n",
    "    for model_i, bmi_class in enumerate(tempD3.keys(), start=5):\n",
    "        tempD = tempD3[bmi_class]\n",
    "        fit_res = tempD[model_sex]\n",
    "        ##Prepare the target class\n",
    "        tempDF1 = tempDF.loc[tempDF['BaseBMI_class']==bmi_class]\n",
    "        ##Predicted component from the fixed effects\n",
    "        tempS = fit_res.predict(tempDF1)#Automatically recognize the column name in formula\n",
    "        ##Create design matrix for the random effects\n",
    "        tempDF2 = pd.DataFrame({'Intercept_re':np.repeat(1, len(tempDF1)),\n",
    "                                'Slope_re':tempDF1['DaysZ']})\n",
    "        tempD = fit_res.random_effects#Dictionary of key=group label and value=pd.Series\n",
    "        ##Predicted component from the random effects\n",
    "        tempL = [np.dot(tempDF2.iloc[row_i], tempD[group_n])\n",
    "                 for (row_i, group_n) in enumerate(tempDF1['public_client_id'].tolist())]\n",
    "        ##Total prediction for in-sample individual\n",
    "        tempS = tempS + tempL\n",
    "        ##Convert to original scale\n",
    "        tempDF1[yvar+'_predicted'+str(model_i)] = np.e**tempS\n",
    "        ##Add the predictions to the working DF\n",
    "        tempDF = pd.merge(tempDF, tempDF1[yvar+'_predicted'+str(model_i)],\n",
    "                          left_index=True, right_index=True, how='outer')#NaN for out-of-class\n",
    "    \n",
    "    #Flatten effects of the dummy season variable\n",
    "    tempDF = tempDF.groupby(by=['public_client_id', 'days_in_program'])\n",
    "    tempDF = tempDF.agg({yvar+'_predicted5':'mean',\n",
    "                         yvar+'_predicted6':'mean'})\n",
    "    tempDF = tempDF.reset_index()\n",
    "    ##Recover the dropped covariates and save the prediction DF to use later\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempL = tempDF1.loc[:, tempDF1.columns.str.contains('BMI_class')].columns.tolist()\n",
    "    tempL = [col_n for sublist in [['public_client_id', 'Sex'], tempL] for col_n in sublist]\n",
    "    tempDF1 = tempDF1[tempL].drop_duplicates('public_client_id', keep='first')\n",
    "    tempDF = pd.merge(tempDF, tempDF1, on='public_client_id', how='left')\n",
    "    tempDF['KeyIndex'] = tempDF['public_client_id'] + '_day' + tempDF['days_in_program'].astype(str)\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    \n",
    "    print(model_sex)\n",
    "    display(tempDF)\n",
    "    display(tempDF.describe(include='all'))\n",
    "    \n",
    "    tempD4[model_sex] = tempDF\n",
    "\n",
    "metLMM_BS = tempD4['Sex-mixed '+yvar+' model']\n",
    "metLMM_FM = tempD4['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check visually\n",
    "tempD1 = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD2 = {'Sex-mixed '+yvar+' model':metLMM_BS, 'Sex-stratified '+yvar+' model':metLMM_FM}\n",
    "tempD3 = {'Model 5':'Normal', 'Model 6':'Obese'}\n",
    "for model_sex in tempD1.keys():\n",
    "    print(model_sex)\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempDF2 = tempD2[model_sex]\n",
    "    \n",
    "    #Plot\n",
    "    for model_i, model_n in enumerate(tempD3.keys(), start=5):\n",
    "        bmi_class = tempD3[model_n]\n",
    "        print(' - '+model_n+': Baseline '+bmi_class+' class-stratified LMM')\n",
    "        tempD = {'Matched':'tab:blue', 'Mismatched':'tab:orange'}\n",
    "        sns.set(style='ticks', font='Arial', context='notebook')\n",
    "        plt.figure(figsize=(5, 3))\n",
    "        sns.scatterplot(data=tempDF1[tempDF1['BaseBMI_class']==bmi_class],\n",
    "                        x='days_in_program', y=yvar,\n",
    "                        hue='vs_Base'+vs_bbmi+'_class', hue_order=tempD.keys(), palette=tempD,\n",
    "                        edgecolor='k', s=20, alpha=0.3)\n",
    "        p = sns.lineplot(data=tempDF2[tempDF2['BaseBMI_class']==bmi_class],\n",
    "                         x='days_in_program', y=yvar+'_predicted'+str(model_i),\n",
    "                         hue='vs_Base'+vs_bbmi+'_class', hue_order=tempD.keys(), palette=tempD,\n",
    "                         estimator='mean', ci=95)\n",
    "        p.set(xlim=(0, 365.25/12*month_threshold), xticks=np.arange(0, 365.25/12*month_threshold, 100),\n",
    "              ylim=(17.5, 42.5), yticks=np.arange(20, 40.1, 5))\n",
    "        sns.despine()\n",
    "        plt.xlabel('Days in program')\n",
    "        plt.ylabel(yvar+' [kg/m'+r'$^2$'+']')\n",
    "        plt.legend(bbox_to_anchor=(1, 0.5), loc='center left', borderaxespad=1)\n",
    "        plt.show()\n",
    "        print('')\n",
    "    \n",
    "    print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6-3. Measured BMI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Repeat the above code, while adding the misclassification info. Because the misclassification info is dependent on sex-mixed or sex-stratified model for bBMI, it is required to be devided by cases.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yvar = 'BMI'\n",
    "yvar_color = 'k'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6-3-1. Prepare the time-series DF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD1 = {'Sex-mixed '+yvar+' model':misclassDF_BS, 'Sex-stratified '+yvar+' model':misclassDF_FM}\n",
    "tempD2 = {}\n",
    "for model_sex in tempD1.keys():\n",
    "    #Import time-series BMI\n",
    "    fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "    ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "    fileName = 'time-series-bmiDF-without-imputation_final-cohort.tsv'\n",
    "    tempDF = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    ##Chage column name\n",
    "    tempDF = tempDF.rename(columns={'BMI_CALC':'BMI'})\n",
    "    \n",
    "    #Add obesity classification (based on BMI)\n",
    "    tempDF['BaseBMI'] = np.e**tempDF['log_BaseBMI']\n",
    "    tempDF = tempDF.drop(columns=['log_BaseBMI'])\n",
    "    ##-> No need to re-calculate this time! Utilize the following misclassification DF\n",
    "    \n",
    "    #Add the misclassification\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempDF = pd.merge(tempDF, tempDF1, left_on='public_client_id', right_index=True, how='left')\n",
    "    \n",
    "    tempD2[model_sex] = tempDF\n",
    "    \n",
    "    print(model_sex)\n",
    "    display(tempDF)\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    print('')\n",
    "\n",
    "tsDF_BS = tempD2['Sex-mixed '+yvar+' model']\n",
    "tsDF_FM = tempD2['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6-3-2. Prepare DFs for LMM with the linear regresion spline for time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD1 = {}\n",
    "tempD2 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    print(model_sex)\n",
    "    tempDF = tempD[model_sex]\n",
    "    print('Time-series DF nrows before filtering:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select measurements for the longitudinal analysis\n",
    "    month_threshold = 18\n",
    "    tempDF = tempDF.loc[tempDF['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    print('Time-series DF nrows after filtering with the taget period:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select the cohort for the longitudinal analysis\n",
    "    ##Select participants who have 2 or more measuremnts in BMI\n",
    "    fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "    ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "    fileName = 'time-series-bmiDF-without-imputation_final-cohort.tsv'\n",
    "    tempDF1 = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF1 = tempDF1.set_index('KeyIndex')\n",
    "    tempDF1 = tempDF1.loc[tempDF1['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    tempS = tempDF1['public_client_id'].value_counts()\n",
    "    tempL1 = tempS.loc[tempS>=2].index.tolist()\n",
    "    ##Select participants who have 2 or more measuremnts in omics\n",
    "    fileDir = '../210104_Biological-BMI-paper/ExportData/'\n",
    "    ipynbName = '210104_Biological-BMI-paper_data-cleaning-BMI-omics_'\n",
    "    fileName = 'time-series-combiDF-without-imputation_final-cohort.tsv'\n",
    "    tempDF1 = pd.read_csv(fileDir+ipynbName+fileName, sep='\\t', dtype={'public_client_id': str})\n",
    "    tempDF1 = tempDF1.set_index('KeyIndex')\n",
    "    tempDF1 = tempDF1.loc[tempDF1['days_in_program'] <= 365.25/12*month_threshold]\n",
    "    tempS = tempDF1['public_client_id'].value_counts()\n",
    "    tempL2 = tempS.loc[tempS>=2].index.tolist()\n",
    "    ##Select participants who have 2 or more measuremnts in both BMI and omics\n",
    "    tempL = list(set(tempL1) & set(tempL2))\n",
    "    print('Participants who have 2 or more measurements in both BMI and omics:', len(tempL))\n",
    "    tempDF = tempDF.loc[tempDF['public_client_id'].isin(tempL)]\n",
    "    print('Time-series DF nrows after filtering with the number of measurements:',\n",
    "          len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Select the participants who have all covariates (just in case)\n",
    "    tempDF = tempDF.dropna()\n",
    "    print('Time-series DF nrows after filtering with covariates:', len(tempDF))\n",
    "    print(' -> Unique ID:', len(tempDF['public_client_id'].unique()))\n",
    "    \n",
    "    #Check measurement distribution\n",
    "    sns.set(style='ticks', font='Arial', context='talk')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    sns.distplot(tempDF['days_in_program'], color=yvar_color)\n",
    "    sns.despine()\n",
    "    for knot in [0, 6, 12, 18]:\n",
    "        plt.axvline(x=365.25/12*knot, **{'linestyle':'--', 'color':'gray'})\n",
    "    plt.axvspan(xmin=365.25/12*18, xmax=plt.gca().get_xlim()[1], facecolor='gray', alpha=0.2)\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel('Days in program')\n",
    "    plt.show()\n",
    "    \n",
    "    #Check dependent variable distribution\n",
    "    print('Skewness of the dependent variable:', stats.skew(tempDF['log_'+yvar]))\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    sns.distplot(tempDF['log_'+yvar], color=yvar_color)\n",
    "    sns.despine()\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel(yvar+' (log-scale)')\n",
    "    plt.show()\n",
    "    \n",
    "    #Add dummy variables for the linear regression spline for time\n",
    "    nknots = 3\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempL1 = []#2nd piece\n",
    "    tempL2 = []#3rd piece\n",
    "    for days in tempDF['days_in_program'].tolist():\n",
    "        if days < tempL[0]:#[0 month, 6 months)\n",
    "            tempL1.append(0.0)\n",
    "            tempL2.append(0.0)\n",
    "        elif days < tempL[1]:#[6 months, 12 months)\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(0.0)\n",
    "        elif days <= tempL[2]:#[12 months, 18 months]\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(days - tempL[1])\n",
    "        else:\n",
    "            tempL1.append('Error? Check time span setting.')\n",
    "            tempL2.append('Error? Check time span setting.')\n",
    "    tempDF['DaysP2'] = tempL1\n",
    "    tempDF['DaysP3'] = tempL2\n",
    "    \n",
    "    #Standardize numeric independent variables\n",
    "    scaler = StandardScaler(copy=True, with_mean=True, with_std=True)\n",
    "    tempL = ['days_in_program', 'DaysP2', 'DaysP3', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5']\n",
    "    scaler.fit(tempDF[tempL])\n",
    "    tempDF1 = pd.DataFrame(scaler.transform(tempDF[tempL]), index=tempDF.index,\n",
    "                           columns=['DaysZ', 'DaysP2Z', 'DaysP3Z', 'BaseAgeZ', 'PC1Z', 'PC2Z', 'PC3Z', 'PC4Z', 'PC5Z'])\n",
    "    tempDF = pd.merge(tempDF, tempDF1, left_index=True, right_index=True, how='inner')\n",
    "    \n",
    "    #Confirm examples\n",
    "    sns.set(style='ticks', font='Arial', context='notebook')\n",
    "    plt.figure(figsize=(4, 3))\n",
    "    for variable in ['DaysZ', 'BaseAgeZ', 'PC1Z']:\n",
    "        sns.distplot(tempDF[variable], label=variable)\n",
    "    sns.despine()\n",
    "    plt.ylabel('Density')\n",
    "    plt.xlabel(r'$Z$'+'-score')\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempS = (pd.Series(tempL) - scaler.mean_[0]) / scaler.scale_[0]#Z-score transformation\n",
    "    for knot in tempS.tolist():\n",
    "        plt.axvline(x=knot, **{'linestyle':'--', 'color':'gray'})\n",
    "    plt.legend(bbox_to_anchor=(1, 0.5), loc='center left', borderaxespad=1)\n",
    "    plt.show()\n",
    "    \n",
    "    #Sort to make 'Matched' the standard (bcoef = 0) in LMM (Matched < Mismatched)\n",
    "    tempDF = tempDF.sort_values(by=['vs_Base'+vs_bbmi+'_class'], ascending=True)\n",
    "    \n",
    "    display(tempDF)\n",
    "    print('')\n",
    "    \n",
    "    tempD1[model_sex] = tempDF\n",
    "    tempD2[model_sex] = scaler#Use later again during predictions\n",
    "\n",
    "#Update\n",
    "tsDF_BS = tempD1['Sex-mixed '+yvar+' model']\n",
    "scaler_BS = tempD2['Sex-mixed '+yvar+' model']\n",
    "tsDF_FM = tempD1['Sex-stratified '+yvar+' model']\n",
    "scaler_FM = tempD2['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6-3-3. LMM with random intercepts and random slopes for time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Days in the program is fitted as a numeric variable using linear regression splines (cf. Zubair, N. et al. Sci. Rep. 2019).  \n",
    "> ***–> Allows differences in the trajectory of BMI/bBMI change.***\n",
    ">\n",
    "> Random intercepts and random slopes for days in the program are used as random effects (cf. Zubair, N. et al. Sci. Rep. 2019).  \n",
    "> ***–> Allows individual differences in the baseline and longitudinal change over the program.***  \n",
    "> –> Random slope is applied not for the piecewise variable of time but only for the overall variable of time; otherwise, leads to \"LinAlgError: Singular matrix\" due to small sample size per most individuals. Also, it is not so weird to assume the random slope is consistent during the program.  \n",
    ">\n",
    "> ***Eliminate \"Underweight\" class in the baseline obesity class-stratified models.***  \n",
    "> <– Otherwise, leads to \"LinAlgError: Singular matrix\" due to small sample population.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD1 = {}\n",
    "for model_sex in tempD.keys():\n",
    "    print(model_sex)\n",
    "    tempDF = tempD[model_sex]\n",
    "    \n",
    "    tempD2 = {'LMM5':'Normal class LMM',\n",
    "              'LMM6':'Obese class LMM'}\n",
    "    tempD3 = {'LMM5':tempDF.loc[tempDF['BaseBMI_class']=='Normal'],\n",
    "              'LMM6':tempDF.loc[tempDF['BaseBMI_class']=='Obese']}\n",
    "    fe_formula = 'log_'+yvar+' ~ DaysZ + DaysP2Z + DaysP3Z'\\\n",
    "        ' + C(Sex) + BaseAgeZ + C(Season) + PC1Z + PC2Z + PC3Z + PC4Z + PC5Z'\\\n",
    "        ' + C(vs_Base'+vs_bbmi+'_class) + DaysZ:C(vs_Base'+vs_bbmi+'_class)'\\\n",
    "        ' + DaysP2Z:C(vs_Base'+vs_bbmi+'_class) + DaysP3Z:C(vs_Base'+vs_bbmi+'_class)'\n",
    "    for model_n in tempD2.keys():\n",
    "        print(' - '+tempD2[model_n])\n",
    "        #Fit LMM\n",
    "        model = sm.MixedLM.from_formula(formula=fe_formula, data=tempD3[model_n],\n",
    "                                        groups='public_client_id',\n",
    "                                        re_formula='DaysZ')#Random intercepts for each group are included as default\n",
    "        t_start = time.time()\n",
    "        #fit_res = model.fit(method=['bfgs', 'lbfgs', 'cg', 'powell', 'nm'])#Back up for failure of convergence\n",
    "        fit_res = model.fit(method=['powell'])#After checking convergence for all models, fix the method through the notebook\n",
    "        t_elapsed = time.time() - t_start\n",
    "        print('    - Elapsed time for fitting LMM:', round(t_elapsed//60), 'min', round(t_elapsed%60, 1), 'sec')\n",
    "        \n",
    "        tempD1[model_sex+' - '+model_n] = fit_res\n",
    "        \n",
    "        #Check the fitted models\n",
    "        display(fit_res.summary())\n",
    "        print('(Note that the estimates for fixed effects and the variances for random effects are shown in the single table.)')\n",
    "        ##Visualize b-coefs and g-coefs\n",
    "        tempDF1 = pd.DataFrame({'Coef':fit_res.params,\n",
    "                                'Coef_ci_l':fit_res.conf_int(alpha=0.05).iloc[:, 0],\n",
    "                                'Coef_ci_h':fit_res.conf_int(alpha=0.05).iloc[:, 1],\n",
    "                                'Pval':fit_res.pvalues})\n",
    "        tempDF1['Coef_ci'] = (tempDF1['Coef_ci_h'] - tempDF1['Coef_ci_l'])/2\n",
    "        tempDF1 = tempDF1.reset_index()\n",
    "        sns.set(style='ticks', font='Arial', context='notebook')\n",
    "        plt.figure(figsize=(4, 0.25*len(tempDF1)))\n",
    "        plt.errorbar(x=tempDF1['Coef'], y=tempDF1['index'], xerr=tempDF1['Coef_ci'],\n",
    "                     fmt='ok', ecolor='k', capsize=5)\n",
    "        plt.xlim(-0.1, 0.1)#Clip\n",
    "        sns.despine()\n",
    "        plt.xlabel(r'$\\beta$'+' or '+r'$\\gamma$'+' coefficient in LMM model\\n(Mean with 95% CI)')\n",
    "        plt.ylabel('')\n",
    "        plt.axvline(x=0, **{'linestyle':'--', 'color':'k'})\n",
    "        for row_i in range(len(tempDF1)):\n",
    "            if tempDF1['Pval'].iloc[row_i]<0.05:\n",
    "                plt.axhspan(ymin=row_i-0.5, ymax=row_i+0.5, facecolor=yvar_color, alpha=0.2, zorder=0)\n",
    "        plt.gca().invert_yaxis()\n",
    "        plt.margins(y=0.01, tight=True)\n",
    "        plt.show()\n",
    "        print('')\n",
    "    print('')\n",
    "\n",
    "LMM5_BS = tempD1['Sex-mixed '+yvar+' model - LMM5']\n",
    "LMM6_BS = tempD1['Sex-mixed '+yvar+' model - LMM6']\n",
    "LMM5_FM = tempD1['Sex-stratified '+yvar+' model - LMM5']\n",
    "LMM6_FM = tempD1['Sex-stratified '+yvar+' model - LMM6']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6-3-4. Estimated transition of biological BMI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> The days in program of the fitted values were different between individuals.  \n",
    "> –> In statsmodels, MixedLMResults.predict() method use only the fixed effects parameters for prediction.  \n",
    "> ***–> To evaluate the variability in the in-sample population, the rondom effects component should be added manually.***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempD1 = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD2 = {'Sex-mixed '+yvar+' model':scaler_BS, 'Sex-stratified '+yvar+' model':scaler_FM}\n",
    "tempD3 = {'Sex-mixed '+yvar+' model':LMM5_BS, 'Sex-stratified '+yvar+' model':LMM5_FM}\n",
    "tempD4 = {'Sex-mixed '+yvar+' model':LMM6_BS, 'Sex-stratified '+yvar+' model':LMM6_FM}\n",
    "tempD3 = {'Normal':tempD3, 'Obese':tempD4}\n",
    "tempD4 = {}\n",
    "for model_sex in tempD1.keys():\n",
    "    #Prepare DF to impute days_in_program but maintain in-sample population\n",
    "    tempDF = tempD1[model_sex]\n",
    "    tempA1 = tempDF['public_client_id'].unique()\n",
    "    tempA2 = np.arange(0, 365.25/12*month_threshold + 1, 365.25/12*month_threshold/nknots)\n",
    "    tempL = ['Spring', 'Summer', 'Autumn', 'Winter']\n",
    "    tempDF = pd.DataFrame(data={'public_client_id':np.repeat(tempA1, len(tempA2)*len(tempL)),\n",
    "                                'days_in_program':np.tile(np.repeat(tempA2, len(tempL)), len(tempA1)),\n",
    "                                'Season':np.tile(tempL, len(tempA1)*len(tempA2))})\n",
    "    ##Add the covariates unique to the in-sample individuals\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempL = ['public_client_id', 'Sex', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5',\n",
    "             'BaseBMI_class', 'vs_Base'+vs_bbmi+'_class']\n",
    "    tempDF1 = tempDF1[tempL]\n",
    "    tempDF1 = tempDF1.drop_duplicates('public_client_id', keep='first')\n",
    "    tempDF = pd.merge(tempDF, tempDF1, on='public_client_id', how='left')\n",
    "    ##Add dummy variables for the linear regression spline for time\n",
    "    tempL = [365.25/12 * month_threshold/nknots * (knot+1) for knot in range(nknots)]#Days of knots\n",
    "    tempL1 = []#2nd piecewise\n",
    "    tempL2 = []#3rd piecewise\n",
    "    for days in tempDF['days_in_program'].tolist():\n",
    "        if days < tempL[0]:#[0 month, 6 months)\n",
    "            tempL1.append(0.0)\n",
    "            tempL2.append(0.0)\n",
    "        elif days < tempL[1]:#[6 months, 12 months)\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(0.0)\n",
    "        elif days <= tempL[2]:#[12 months, 18 months]\n",
    "            tempL1.append(days - tempL[0])\n",
    "            tempL2.append(days - tempL[1])\n",
    "        else:\n",
    "            tempL1.append('Error? Check time span setting.')\n",
    "            tempL2.append('Error? Check time span setting.')\n",
    "    tempDF['DaysP2'] = tempL1\n",
    "    tempDF['DaysP3'] = tempL2\n",
    "    ##Add a dummy index to use when merging later\n",
    "    tempDF['KeyIndex'] = tempDF['public_client_id'] + '_day' + tempDF['days_in_program'].astype(str)\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    ##Standardize numeric independent variables using the already fitted StandardScaler\n",
    "    scaler = tempD2[model_sex]\n",
    "    tempL = ['days_in_program', 'DaysP2', 'DaysP3', 'BaseAge', 'PC1', 'PC2', 'PC3', 'PC4', 'PC5']\n",
    "    tempDF1 = pd.DataFrame(scaler.transform(tempDF[tempL]), index=tempDF.index,\n",
    "                           columns=['DaysZ', 'DaysP2Z', 'DaysP3Z', 'BaseAgeZ', 'PC1Z', 'PC2Z', 'PC3Z', 'PC4Z', 'PC5Z'])\n",
    "    tempDF = pd.merge(tempDF, tempDF1, left_index=True, right_index=True, how='inner')\n",
    "    \n",
    "    #Model 5-6 prediction\n",
    "    for model_i, bmi_class in enumerate(tempD3.keys(), start=5):\n",
    "        tempD = tempD3[bmi_class]\n",
    "        fit_res = tempD[model_sex]\n",
    "        ##Prepare the target class\n",
    "        tempDF1 = tempDF.loc[tempDF['BaseBMI_class']==bmi_class]\n",
    "        ##Predicted component from the fixed effects\n",
    "        tempS = fit_res.predict(tempDF1)#Automatically recognize the column name in formula\n",
    "        ##Create design matrix for the random effects\n",
    "        tempDF2 = pd.DataFrame({'Intercept_re':np.repeat(1, len(tempDF1)),\n",
    "                                'Slope_re':tempDF1['DaysZ']})\n",
    "        tempD = fit_res.random_effects#Dictionary of key=group label and value=pd.Series\n",
    "        ##Predicted component from the random effects\n",
    "        tempL = [np.dot(tempDF2.iloc[row_i], tempD[group_n])\n",
    "                 for (row_i, group_n) in enumerate(tempDF1['public_client_id'].tolist())]\n",
    "        ##Total prediction for in-sample individual\n",
    "        tempS = tempS + tempL\n",
    "        ##Convert to original scale\n",
    "        tempDF1[yvar+'_predicted'+str(model_i)] = np.e**tempS\n",
    "        ##Add the predictions to the working DF\n",
    "        tempDF = pd.merge(tempDF, tempDF1[yvar+'_predicted'+str(model_i)],\n",
    "                          left_index=True, right_index=True, how='outer')#NaN for out-of-class\n",
    "    \n",
    "    #Flatten effects of the dummy season variable\n",
    "    tempDF = tempDF.groupby(by=['public_client_id', 'days_in_program'])\n",
    "    tempDF = tempDF.agg({yvar+'_predicted5':'mean',\n",
    "                         yvar+'_predicted6':'mean'})\n",
    "    tempDF = tempDF.reset_index()\n",
    "    ##Recover the dropped covariates and save the prediction DF to use later\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempL = tempDF1.loc[:, tempDF1.columns.str.contains('BMI_class')].columns.tolist()\n",
    "    tempL = [col_n for sublist in [['public_client_id', 'Sex'], tempL] for col_n in sublist]\n",
    "    tempDF1 = tempDF1[tempL].drop_duplicates('public_client_id', keep='first')\n",
    "    tempDF = pd.merge(tempDF, tempDF1, on='public_client_id', how='left')\n",
    "    tempDF['KeyIndex'] = tempDF['public_client_id'] + '_day' + tempDF['days_in_program'].astype(str)\n",
    "    tempDF = tempDF.set_index('KeyIndex')\n",
    "    \n",
    "    print(model_sex)\n",
    "    display(tempDF)\n",
    "    display(tempDF.describe(include='all'))\n",
    "    \n",
    "    tempD4[model_sex] = tempDF\n",
    "\n",
    "measLMM_BS = tempD4['Sex-mixed '+yvar+' model']\n",
    "measLMM_FM = tempD4['Sex-stratified '+yvar+' model']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check visually\n",
    "tempD1 = {'Sex-mixed '+yvar+' model':tsDF_BS, 'Sex-stratified '+yvar+' model':tsDF_FM}\n",
    "tempD2 = {'Sex-mixed '+yvar+' model':measLMM_BS, 'Sex-stratified '+yvar+' model':measLMM_FM}\n",
    "tempD3 = {'Model 5':'Normal', 'Model 6':'Obese'}\n",
    "for model_sex in tempD1.keys():\n",
    "    print(model_sex)\n",
    "    tempDF1 = tempD1[model_sex]\n",
    "    tempDF2 = tempD2[model_sex]\n",
    "    \n",
    "    #Plot\n",
    "    for model_i, model_n in enumerate(tempD3.keys(), start=5):\n",
    "        bmi_class = tempD3[model_n]\n",
    "        print(' - '+model_n+': Baseline '+bmi_class+' class-stratified LMM')\n",
    "        tempD = {'Matched':'tab:blue', 'Mismatched':'tab:orange'}\n",
    "        sns.set(style='ticks', font='Arial', context='notebook')\n",
    "        plt.figure(figsize=(5, 3))\n",
    "        sns.scatterplot(data=tempDF1[tempDF1['BaseBMI_class']==bmi_class],\n",
    "                        x='days_in_program', y=yvar,\n",
    "                        hue='vs_Base'+vs_bbmi+'_class', hue_order=tempD.keys(), palette=tempD,\n",
    "                        edgecolor='k', s=20, alpha=0.3)\n",
    "        p = sns.lineplot(data=tempDF2[tempDF2['BaseBMI_class']==bmi_class],\n",
    "                         x='days_in_program', y=yvar+'_predicted'+str(model_i),\n",
    "                         hue='vs_Base'+vs_bbmi+'_class', hue_order=tempD.keys(), palette=tempD,\n",
    "                         estimator='mean', ci=95)\n",
    "        p.set(xlim=(0, 365.25/12*month_threshold), xticks=np.arange(0, 365.25/12*month_threshold, 100),\n",
    "              ylim=(17.5, 42.5), yticks=np.arange(20, 40.1, 5))\n",
    "        sns.despine()\n",
    "        plt.xlabel('Days in program')\n",
    "        plt.ylabel(yvar+' [kg/m'+r'$^2$'+']')\n",
    "        plt.legend(bbox_to_anchor=(1, 0.5), loc='center left', borderaxespad=1)\n",
    "        plt.show()\n",
    "        print('')\n",
    "    \n",
    "    print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6-4. Comparison of estimates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#For presentation\n",
    "tempD1 = {'BMI':'k', 'MetBMI':'b'}\n",
    "tempD2 = {'Matched':'tab:blue', 'Mismatched':'tab:orange'}\n",
    "tempL1 = [measLMM_BS, metLMM_BS]\n",
    "tempL2 = [measLMM_FM, metLMM_FM]\n",
    "tempD3 = {'Sex-mixed bBMI model':tempL1, 'Sex-stratified bBMI model':tempL2}\n",
    "tempD4 = {'Model 5':'Normal', 'Model 6':'Obese'}\n",
    "tempD5 = {'Sex-mixed bBMI model':'BothSex', 'Sex-stratified bBMI model':'FemaleMale'}\n",
    "\n",
    "for model_sex in tempD3.keys():\n",
    "    print(model_sex)\n",
    "    \n",
    "    #Merge prediction DFs\n",
    "    tempL = tempD3[model_sex]\n",
    "    tempDF1 = tempL[0]\n",
    "    tempDF2 = tempL[1]\n",
    "    tempDF2 = tempDF2.loc[:, ~tempDF2.columns.isin(tempDF1.columns.tolist())]\n",
    "    tempDF = pd.merge(tempDF1, tempDF2, left_index=True, right_index=True, how='outer')\n",
    "    \n",
    "    #Calculate longitudinal change\n",
    "    tempL = tempDF.loc[:, tempDF.columns.str.contains('_predicted')].columns.tolist()\n",
    "    for col_n in tempL:\n",
    "        tempDF1 = tempDF[['public_client_id', 'days_in_program', col_n]].dropna()\n",
    "        #Retrieve the estimate at days = 0 for each participant\n",
    "        tempDF2 = tempDF1.loc[tempDF1['days_in_program']==0.0]\n",
    "        tempS = tempDF2.reset_index().set_index('public_client_id')[col_n]\n",
    "        tempS.name = 'day0'\n",
    "        #Merge\n",
    "        tempDF1 = pd.merge(tempDF1, tempS, left_on='public_client_id', right_index=True, how='left')\n",
    "        #Calculate longitudinal change [%]\n",
    "        label = col_n.replace('_predicted', 'change_predicted')\n",
    "        tempDF1[label] = (tempDF1[col_n] - tempDF1['day0']) / tempDF1['day0'] * 100\n",
    "        #Add to the working DF\n",
    "        tempDF = pd.merge(tempDF, tempDF1[label], left_index=True, right_index=True, how='left')\n",
    "    \n",
    "    #Model 5-6\n",
    "    for model_i, model_n in enumerate(tempD4.keys(), start=5):\n",
    "        bmi_class = tempD4[model_n]\n",
    "        print(' - '+model_n+': Baseline '+bmi_class+' class-stratified LMM')\n",
    "        #Prepare DF\n",
    "        tempL = tempDF.loc[:, tempDF.columns.str.contains('change_predicted'+str(model_i))].columns.tolist()\n",
    "        tempDF1 = tempDF.reset_index().melt(var_name='LMM', value_name='PredictedChange', value_vars=tempL,\n",
    "                                            id_vars=['public_client_id', 'days_in_program',\n",
    "                                                     'BaseBMI_class', 'vs_Base'+vs_bbmi+'_class'])\n",
    "        tempDF1 = tempDF1.dropna()\n",
    "        tempDF1['LMM'] = tempDF1['LMM'].str.replace('change_predicted'+str(model_i), '')\n",
    "        #Check sample size\n",
    "        for yvar in tempD1.keys():\n",
    "            tempDF2 = tempDF1.loc[tempDF1['LMM']==yvar]\n",
    "            print('   - '+yvar+': n =', len(tempDF2['public_client_id'].unique()))\n",
    "            tempL = []\n",
    "            for misclass in tempD2.keys():\n",
    "                tempA = tempDF2['public_client_id'].loc[tempDF2['vs_Base'+vs_bbmi+'_class']==misclass].unique()\n",
    "                tempL.append(misclass+' (n = '+str(len(tempA))+')')\n",
    "            print('     -> Confirm each sample size:', tempL)\n",
    "        #Visualization\n",
    "        sns.set(style='ticks', font='Arial', context='talk')\n",
    "        p = sns.relplot(kind='line', data=tempDF1, x='days_in_program', y='PredictedChange',\n",
    "                        row='LMM', row_order=tempD1.keys(),\n",
    "                        hue='vs_Base'+vs_bbmi+'_class', hue_order=tempD2.keys(), palette=tempD2,\n",
    "                        estimator='mean', ci=95, height=3.15, aspect=4.5/3-0.025, legend=False)#Manual adjustment\n",
    "        if bmi_class=='Normal':\n",
    "            p.set(xlim=(0, 365.25/12*12), xticks=np.arange(0, 365.25/12*12+1, 90),\n",
    "                  ylim=(-8.5, 1), yticks=np.arange(-8, 0.1, 2))\n",
    "        elif bmi_class=='Obese':\n",
    "            p.set(xlim=(0, 365.25/12*12), xticks=np.arange(0, 365.25/12*12+1, 90),\n",
    "                  ylim=(-8.5, 1), yticks=np.arange(-8, 0.1, 2))\n",
    "        tempL = ['LMM estimate: '+estimate for estimate in tempD1.keys()]\n",
    "        for ax, ax_title in zip(p.axes.ravel(), tempL):\n",
    "            #Draw initial line\n",
    "            ax.axhline(y=0, **{'linestyle':'--', 'color':'k'}, zorder=0)\n",
    "            #Draw span b/w knots\n",
    "            #for knot in range(nknots-1):\n",
    "            #    if knot%2 == 1:\n",
    "            #        ax.axvspan(xmin=365.25/12*month_threshold/nknots*knot,\n",
    "            #                   xmax=365.25/12*month_threshold/nknots*(knot+1),\n",
    "            #                   facecolor='gray', alpha=0.2, zorder=-1)\n",
    "            #Change facet label\n",
    "            ax.set_title(ax_title, {'fontsize':'large'})\n",
    "        ##Reset and generate common axis title\n",
    "        for row_i, ax in enumerate(p.axes.ravel()):\n",
    "            if row_i == 1:\n",
    "                ax.set_xlabel('Days in program')\n",
    "                ax.set_ylabel('')\n",
    "            else:\n",
    "                ax.set_xlabel('')\n",
    "                ax.set_ylabel('')\n",
    "        p.fig.text(x=0, y=0.5, s='Estimated change [%]', fontsize='medium',\n",
    "                   verticalalignment='center', horizontalalignment='left', rotation='vertical')\n",
    "        plt.tight_layout()\n",
    "        ##Generate standard Matplotlib legend manually\n",
    "        tempL = []\n",
    "        for misclass in tempD2.keys():\n",
    "            tempL.append(mlines.Line2D([], [], color=tempD2[misclass], label=misclass, linewidth=3))\n",
    "        plt.legend(handles=tempL, fontsize='medium',\n",
    "                   title='Baseline BMI class\\n(vs. '+vs_bbmi+' class)', title_fontsize='large',\n",
    "                   bbox_to_anchor=(0.5, 0), loc='upper center', ncol=2, borderaxespad=4,\n",
    "                   handlelength=1.5, handletextpad=0.5, columnspacing=1.0, frameon=True)\n",
    "        ##Save\n",
    "        fileDir = './ExportFigures/'\n",
    "        ipynbName = '220806_Multiomics-BMI-NatMed1stRevision_BMI-longitudinal-LMM-ver2_'\n",
    "        fileName = 'baseline'+bmi_class+'class(vs'+vs_bbmi+')-'+tempD5[model_sex]+'.tif'\n",
    "        plt.gcf().savefig(fileDir+ipynbName+fileName, dpi=300, bbox_inches='tight', pad_inches=0.04,\n",
    "                          pil_kwargs={'compression':'tiff_lzw'})\n",
    "        plt.show()\n",
    "        print('')\n",
    "    print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# — End of this notebook —"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "arivale-py3 - Python",
   "language": "python",
   "name": "conda-env-arivale-py3-py"
  },
  "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.7"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}