{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# statsmodels Principal Component Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Key ideas:* Principal component analysis, world bank data, fertility\n",
    "\n",
    "In this notebook, we use principal components analysis (PCA) to analyze the time series of fertility rates in 192 countries, using data obtained from the World Bank.  The main goal is to understand how the trends in fertility over time differ from country to country.  This is a slightly atypical illustration of PCA because the data are time series.  Methods such as functional PCA have been developed for this setting, but since the fertility data are very smooth, there is no real disadvantage to using standard PCA in this case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import statsmodels.api as sm\n",
    "from statsmodels.multivariate.pca import PCA\n",
    "\n",
    "plt.rc(\"figure\", figsize=(16, 8))\n",
    "plt.rc(\"font\", size=14)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data can be obtained from the [World Bank web site](http://data.worldbank.org/indicator/SP.DYN.TFRT.IN), but here we work with a slightly cleaned-up version of the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "data = sm.datasets.fertility.load_pandas().data\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we construct a DataFrame that contains only the numerical fertility rate data and set the index to the country names.  We also drop all the countries with any missing data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "columns = list(map(str, range(1960, 2012)))\n",
    "data.set_index(\"Country Name\", inplace=True)\n",
    "dta = data[columns]\n",
    "dta = dta.dropna()\n",
    "dta.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are two ways to use PCA to analyze a rectangular matrix: we can treat the rows as the \"objects\" and the columns as the \"variables\", or vice-versa.  Here we will treat the fertility measures as \"variables\" used to measure the countries as \"objects\".  Thus the goal will be to reduce the yearly fertility rate values to a small number of fertility rate \"profiles\" or \"basis functions\" that capture most of the variation over time in the different countries."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The mean trend is removed in PCA, but its worthwhile taking a look at it.  It shows that fertility has dropped steadily over the time period covered in this dataset.  Note that the mean is calculated using a country as the unit of analysis, ignoring population size.  This is also true for the PC analysis conducted below.  A more sophisticated analysis might weight the countries, say by population in 1980."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "ax = dta.mean().plot(grid=False)\n",
    "ax.set_xlabel(\"Year\", size=17)\n",
    "ax.set_ylabel(\"Fertility rate\", size=17)\n",
    "ax.set_xlim(0, 51)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we perform the PCA:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "pca_model = PCA(dta.T, standardize=False, demean=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Based on the eigenvalues, we see that the first PC dominates, with perhaps a small amount of meaningful variation captured in the second and third PC's."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "fig = pca_model.plot_scree(log_scale=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we will plot the PC factors.  The dominant factor is monotonically increasing.  Countries with a positive score on the first factor will increase faster (or decrease slower) compared to the mean shown above.  Countries with a negative score on the first factor will decrease faster than the mean.  The second factor is U-shaped with a positive peak at around 1985.  Countries with a large positive score on the second factor will have lower than average fertilities at the beginning and end of the data range, but higher than average fertility in the middle of the range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "lines = ax.plot(pca_model.factors.iloc[:, :3], lw=4, alpha=0.6)\n",
    "ax.set_xticklabels(dta.columns.values[::10])\n",
    "ax.set_xlim(0, 51)\n",
    "ax.set_xlabel(\"Year\", size=17)\n",
    "fig.subplots_adjust(0.1, 0.1, 0.85, 0.9)\n",
    "legend = fig.legend(lines, [\"PC 1\", \"PC 2\", \"PC 3\"], loc=\"center right\")\n",
    "legend.draw_frame(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To better understand what is going on, we will plot the fertility trajectories for sets of countries with similar PC scores.  The following convenience function produces such a plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "idx = pca_model.loadings.iloc[:, 0].argsort()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we plot the five countries with the greatest scores on PC 1.  These countries have a higher rate of fertility increase than the global mean (which is decreasing)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def make_plot(labels):\n",
    "    fig, ax = plt.subplots(figsize=(9, 5))\n",
    "    ax = dta.loc[labels].T.plot(legend=False, grid=False, ax=ax)\n",
    "    dta.mean().plot(ax=ax, grid=False, label=\"Mean\")\n",
    "    ax.set_xlim(0, 51)\n",
    "    fig.subplots_adjust(0.1, 0.1, 0.75, 0.9)\n",
    "    ax.set_xlabel(\"Year\", size=17)\n",
    "    ax.set_ylabel(\"Fertility\", size=17)\n",
    "    legend = ax.legend(\n",
    "        *ax.get_legend_handles_labels(), loc=\"center left\", bbox_to_anchor=(1, 0.5)\n",
    "    )\n",
    "    legend.draw_frame(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "labels = dta.index[idx[-5:]]\n",
    "make_plot(labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are the five countries with the greatest scores on factor 2.  These are countries that reached peak fertility around 1980, later than much of the rest of the world, followed by a rapid decrease in fertility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "idx = pca_model.loadings.iloc[:, 1].argsort()\n",
    "make_plot(dta.index[idx[-5:]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we have the countries with the most negative scores on PC 2.  These are the countries where the fertility rate declined much faster than the global mean during the 1960's and 1970's, then flattened out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "make_plot(dta.index[idx[:5]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also look at a scatterplot of the first two principal component scores.  We see that the variation among countries is fairly continuous, except perhaps that the two countries with highest scores for PC 2 are somewhat separated from the other points.  These countries, Oman and Yemen, are unique in having a sharp spike in fertility around 1980.  No other country has such a spike.  In contrast, the countries with high scores on PC 1 (that have continuously increasing fertility), are part of a continuum of variation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "pca_model.loadings.plot.scatter(x=\"comp_00\", y=\"comp_01\", ax=ax)\n",
    "ax.set_xlabel(\"PC 1\", size=17)\n",
    "ax.set_ylabel(\"PC 2\", size=17)\n",
    "dta.index[pca_model.loadings.iloc[:, 1] > 0.2].values"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
