heatMAPP.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob, os\n",
    "import numpy as np\n",
    "import IPython\n",
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "fontSize =27\n",
    "plt.rcParams.update({'font.size': fontSize})\n",
    "import pickle\n",
    "import matplotlib.ticker as mticker\n",
    "from joblib import load, dump"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "pathSave = \"Data_Files/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class plotHeatMap:\n",
    "    def __init__(self,  _observables, _parameters, _klArray, _yLabel, _xLabel, _plotString, _pathSave, _title, _stdArray={}, figSize=(18,10), fontSize=fontSize):\n",
    "        self.parameters = _parameters\n",
    "        self.observables = _observables\n",
    "        self.klArray = _klArray\n",
    "        self.stdArray = _stdArray\n",
    "        self.yLabel = _yLabel\n",
    "        self.xLabel = _xLabel\n",
    "        self.plotString = _plotString\n",
    "        self.pathSave = _pathSave\n",
    "        self.title = _title\n",
    "        self.figSize = figSize\n",
    "        self.fontSize = fontSize\n",
    "        \n",
    "    def heatmap(self, data, row_labels, col_labels, ax=None,\n",
    "            cbar_kw={}, cbarlabel=\"\", **kwargs):\n",
    "        \"\"\"\n",
    "        Create a heatmap from a numpy array and two lists of labels.\n",
    "\n",
    "        Arguments:\n",
    "            data       : A 2D numpy array of shape (N,M)\n",
    "            row_labels : A list or array of length N with the labels\n",
    "                         for the rows\n",
    "            col_labels : A list or array of length M with the labels\n",
    "                         for the columns\n",
    "        Optional arguments:\n",
    "            ax         : A matplotlib.axes.Axes instance to which the heatmap\n",
    "                         is plotted. If not provided, use current axes or\n",
    "                         create a new one.\n",
    "            cbar_kw    : A dictionary with arguments to\n",
    "                         :meth:`matplotlib.Figure.colorbar`.\n",
    "            cbarlabel  : The label for the colorbar\n",
    "        All other arguments are directly passed on to the imshow call.\n",
    "        \"\"\"\n",
    "        if not ax:\n",
    "            ax = plt.gca()\n",
    "\n",
    "        # Plot the heatmap\n",
    "        im = ax.imshow(data, **kwargs)\n",
    "        divider = make_axes_locatable(ax)\n",
    "        cax = divider.append_axes(\"right\", size=\"5%\", pad=0.2)\n",
    "        # Create colorbar\n",
    "        cbar = ax.figure.colorbar(im, cax=cax, **cbar_kw)\n",
    "        cbar.ax.set_ylabel(cbarlabel, rotation=-90, va=\"bottom\")\n",
    "\n",
    "        # We want to show all ticks...\n",
    "        ax.set_xticks(np.arange(data.shape[1]))\n",
    "        ax.set_yticks(np.arange(data.shape[0]))\n",
    "        # ... and label them with the respective list entries.\n",
    "\n",
    "        ax.set_xticklabels(col_labels)\n",
    "        ax.set_yticklabels(row_labels)\n",
    "    \n",
    "        # Let the horizontal axes labeling appear on top.\n",
    "        ax.tick_params(top=True, bottom=False,\n",
    "                       labeltop=True, labelbottom=False)\n",
    "        \n",
    "        ax.set_xlabel(self.xLabel)\n",
    "        ax.set_ylabel(self.yLabel)\n",
    "        # Rotate the tick labels and set their alignment.\n",
    "        plt.setp(ax.get_xticklabels(), rotation=0, ha=\"right\", #=-30\n",
    "                 rotation_mode=\"anchor\")\n",
    "\n",
    "        # Turn spines off and create white grid.\n",
    "        for edge, spine in ax.spines.items():\n",
    "            spine.set_visible(False)\n",
    "\n",
    "        ax.set_xticks(np.arange(data.shape[1]+1)-.5, minor=True)\n",
    "        ax.set_yticks(np.arange(data.shape[0]+1)-.5, minor=True)\n",
    "        ax.grid(which=\"minor\", color=\"w\", linestyle='-', linewidth=3)\n",
    "        ax.tick_params(which=\"minor\", bottom=False, left=False)\n",
    "\n",
    "        return im, cbar\n",
    "\n",
    "\n",
    "    def annotate_heatmap(self, im, stdBool=False, data=None, valfmt=\"{x:.2f}\",\n",
    "                         textcolors=[\"black\", \"white\"],\n",
    "                         threshold=None, **textkw):\n",
    "        \"\"\"\n",
    "        A function to annotate a heatmap.\n",
    "\n",
    "        Arguments:\n",
    "            im         : The AxesImage to be labeled.\n",
    "        Optional arguments:\n",
    "            data       : Data used to annotate. If None, the image's data is used.\n",
    "            valfmt     : The format of the annotations inside the heatmap.\n",
    "                         This should either use the string format method, e.g.\n",
    "                         \"$ {x:.2f}\", or be a :class:`matplotlib.ticker.Formatter`.\n",
    "            textcolors : A list or array of two color specifications. The first is\n",
    "                         used for values below a threshold, the second for those\n",
    "                         above.\n",
    "            threshold  : Value in data units according to which the colors from\n",
    "                         textcolors are applied. If None (the default) uses the\n",
    "                         middle of the colormap as separation.\n",
    "\n",
    "        Further arguments are passed on to the created text labels.\n",
    "        \"\"\"\n",
    "\n",
    "        if not isinstance(data, (list, np.ndarray)):\n",
    "            data = im.get_array()\n",
    "\n",
    "        # Normalize the threshold to the images color range.\n",
    "        if threshold is not None:\n",
    "            threshold = im.norm(threshold)\n",
    "        else:\n",
    "            threshold = im.norm(data.max())/2.3\n",
    "\n",
    "        # Set default alignment to center, but allow it to be\n",
    "        # overwritten by textkw.\n",
    "        kw = dict(horizontalalignment=\"center\",\n",
    "                  verticalalignment=\"bottom\")\n",
    "        kw.update(textkw)\n",
    "        \n",
    "        kw1 = dict(horizontalalignment=\"center\",\n",
    "                  verticalalignment=\"top\")\n",
    "        kw1.update(textkw)\n",
    "\n",
    "        # Get the formatter in case a string is supplied\n",
    "        if isinstance(valfmt, str):\n",
    "            valfmt = matplotlib.ticker.StrMethodFormatter(valfmt)\n",
    "\n",
    "        # Loop over the data and create a `Text` for each \"pixel\".\n",
    "        # Change the text's color depending on the data.\n",
    "        texts = []\n",
    "        for i in range(data.shape[0]):\n",
    "            for j in range(data.shape[1]):\n",
    "                if not stdBool:\n",
    "                    kw.update(color=textcolors[im.norm(data[i, j]) > threshold])\n",
    "                    text = im.axes.text(j, i, valfmt(data[i, j], None), **kw)\n",
    "                else:\n",
    "                    kw1.update(color=textcolors[im.norm(data[i, j]) > threshold])\n",
    "                    text = im.axes.text(j, i, r'$\\pm$' + valfmt(self.stdArray[i,j], None), **kw1)\n",
    "                texts.append(text)\n",
    "\n",
    "        return texts\n",
    "    \n",
    "    def plotMap(self):\n",
    "        plt.rcParams.update({'font.size': self.fontSize})\n",
    "        fig, ax = plt.subplots(figsize=self.figSize,dpi=480)\n",
    "        im, cbar = self.heatmap(self.klArray,  self.parameters, self.observables, ax=ax, cmap=\"YlGn\", cbarlabel=self.title) #\"$KL(MP||Ref)-KL(NN||Ref)$\")\n",
    "        texts = self.annotate_heatmap(im, False, valfmt=\"{x:.2f}\", size=self.fontSize/1.1)\n",
    "        texts = self.annotate_heatmap(im, True, valfmt=\"{x:.2f}\", size=self.fontSize/1.5)\n",
    "        fig.tight_layout()\n",
    "        #fig.savefig(self.pathSave + \"/Pics/\" + self.plotString + '.pdf', bbox_inches='tight')\n",
    "        plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with open(pathSave + \"Parameters_Observables.txt\", \"rb\") as fkl:\n",
    "    Dict_Parameters_Observables = load(fkl)\n",
    "\n",
    "klArray = Dict_Parameters_Observables['log-likelihood_mean']\n",
    "stdArray = Dict_Parameters_Observables['log-likelihood_std']\n",
    "x = Dict_Parameters_Observables['x']\n",
    "y = Dict_Parameters_Observables['y']\n",
    "xLabel = Dict_Parameters_Observables['xLabel'] \n",
    "yLabel = Dict_Parameters_Observables['yLabel']\n",
    "\n",
    "counter = 0\n",
    "for rInd in [[0,43],[43,85],[85,127]]:\n",
    "    _y = y[rInd[0]:rInd[1]]\n",
    "    _klArray = klArray[rInd[0]:rInd[1],:]\n",
    "    _klArray.shape = (rInd[1]-rInd[0],5)\n",
    "\n",
    "    plotString = \"Parameters_Observables_\" + str(counter)\n",
    "    HM = plotHeatMap(x, _y, _klArray, yLabel, xLabel, plotString, pathSave \\\n",
    "                     ,\"Log likelihood \" + r'($\\mu \\pm \\sigma$)', stdArray, (24,24), 15)\n",
    "    HM.plotMap()\n",
    "    counter += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(pathSave + \"Parameters_Observables_Select.txt\", \"rb\") as fkl:\n",
    "    Dict_Parameters_Observables = load(fkl)\n",
    "\n",
    "klArray = Dict_Parameters_Observables['log-likelihood_mean']\n",
    "stdArray = Dict_Parameters_Observables['log-likelihood_std']\n",
    "x = Dict_Parameters_Observables['x']\n",
    "y = Dict_Parameters_Observables['y']\n",
    "xLabel = Dict_Parameters_Observables['xLabel'] \n",
    "yLabel = Dict_Parameters_Observables['yLabel']\n",
    "\n",
    "counter = 0\n",
    "for rInd in [[0,len(y)]]:\n",
    "    _y = y[rInd[0]:rInd[1]]\n",
    "    _klArray = klArray[rInd[0]:rInd[1],:]\n",
    "    _klArray.shape = (rInd[1]-rInd[0],5)\n",
    "\n",
    "    plotString = \"Parameters_Observables_Select\"\n",
    "    HM = plotHeatMap(x, _y, _klArray, yLabel, xLabel, plotString, pathSave \\\n",
    "                     ,\"Log likelihood \" + r'($\\mu \\pm \\sigma$)', stdArray, (24,24), 35)\n",
    "    HM.plotMap()\n",
    "    counter += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plotString = \"Parameters_Observables_Tprof\"\n",
    "with open(pathSave + plotString + \".txt\", \"rb\") as fkl:\n",
    "    Dict_Parameters_Observables = load(fkl)\n",
    "\n",
    "klArray = Dict_Parameters_Observables['log-likelihood_mean']\n",
    "stdArray = Dict_Parameters_Observables['log-likelihood_std']\n",
    "x = Dict_Parameters_Observables['x']\n",
    "y = Dict_Parameters_Observables['y']\n",
    "xLabel = Dict_Parameters_Observables['xLabel'] \n",
    "yLabel = Dict_Parameters_Observables['yLabel']\n",
    "\n",
    "HM = plotHeatMap(x, y, klArray, yLabel, xLabel, plotString, pathSave \\\n",
    "                 ,\"Log likelihood \" + r'($\\mu \\pm \\sigma$)', stdArray)\n",
    "HM.plotMap()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plotString = \"Parameters_Observables_Tprof_inverse\"\n",
    "with open(pathSave + plotString + \".txt\", \"rb\") as fkl:\n",
    "    Dict_Parameters_Observables = load(fkl)\n",
    "\n",
    "klArray = Dict_Parameters_Observables['log-likelihood_mean']\n",
    "stdArray = Dict_Parameters_Observables['log-likelihood_std']\n",
    "x = Dict_Parameters_Observables['x']\n",
    "y = Dict_Parameters_Observables['y']\n",
    "xLabel = Dict_Parameters_Observables['xLabel'] \n",
    "yLabel = Dict_Parameters_Observables['yLabel']\n",
    "\n",
    "HM = plotHeatMap(x, y, klArray, yLabel, xLabel, plotString, pathSave \\\n",
    "                 ,\"Log likelihood \" + r'($\\mu \\pm \\sigma$)', stdArray)\n",
    "HM.plotMap()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plotString = \"Parameters_Noise\"\n",
    "with open(pathSave + plotString + \".txt\", \"rb\") as fkl:\n",
    "    Dict_Parameters_Observables = load(fkl)\n",
    "\n",
    "klArray = Dict_Parameters_Observables['log-likelihood_mean']\n",
    "stdArray = Dict_Parameters_Observables['log-likelihood_std']\n",
    "x = Dict_Parameters_Observables['x']\n",
    "y = Dict_Parameters_Observables['y']\n",
    "xLabel = Dict_Parameters_Observables['xLabel'] \n",
    "yLabel = Dict_Parameters_Observables['yLabel']\n",
    "\n",
    "HM = plotHeatMap(x, y, klArray, yLabel, xLabel, plotString, pathSave \\\n",
    "                 ,\"Log likelihood \" + r'($\\mu \\pm \\sigma$)', stdArray)\n",
    "HM.plotMap()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for p in [0,4,5,6,7]:\n",
    "    plotString = \"Simulations_Observables\" + str(p)\n",
    "    with open(pathSave + plotString + \".txt\", \"rb\") as fkl:\n",
    "        Dict_Parameters_Observables = load(fkl)\n",
    "\n",
    "    klArray = Dict_Parameters_Observables['log-likelihood_mean']\n",
    "    stdArray = Dict_Parameters_Observables['log-likelihood_std']\n",
    "    x = Dict_Parameters_Observables['x']\n",
    "    y = Dict_Parameters_Observables['y']\n",
    "    xLabel = Dict_Parameters_Observables['xLabel'] \n",
    "    yLabel = Dict_Parameters_Observables['yLabel']\n",
    "\n",
    "    HM = plotHeatMap(x, y, klArray, yLabel, \"Number of simulations\", plotString, pathSave \\\n",
    "                     ,\"Log likelihood \" + r'($\\mu \\pm \\sigma$)', stdArray)\n",
    "    HM.plotMap()"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}