helper.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dir(dir_path):\n",
    "    if not os.path.exists(dir_path):\n",
    "        os.mkdir(dir_path)\n",
    "    \n",
    "def plot_heatmaps(maps, x_labels, y_labels, x_title, y_title, subfolder, filename, \n",
    "                  plot_title=\"\", subplot_titles=None, annotate=False, fmt='png'):\n",
    "    columns, rows = len(maps), 1\n",
    "    fig = plt.figure()\n",
    "    for i in range(1, columns * rows + 1):\n",
    "        ax = fig.add_subplot(rows, columns, i)\n",
    "        ax.set_title(subplot_titles[i-1] if subplot_titles is not None else \"\")\n",
    "        ax.set_xticks(np.arange(len(x_labels)))\n",
    "        ax.set_yticks(np.arange(len(y_labels)))\n",
    "        ax.set_xticklabels(x_labels)\n",
    "        ax.set_yticklabels(y_labels)\n",
    "        ax.set_xlabel(x_title)\n",
    "        ax.set_ylabel(y_title)\n",
    "        \n",
    "        im = ax.imshow(maps[i-1], cmap=plt.cm.Blues)\n",
    "        plt.setp(ax.xaxis.get_majorticklabels(), rotation=90) # rotate x tick labels by 90 degrees\n",
    "        \n",
    "        fig.colorbar(im, ax=ax)\n",
    "        \n",
    "        if annotate:\n",
    "            fmt_str = '{:.3f}' if maps[i-1].dtype == float else '{:.0f}'\n",
    "            thresh = (maps[i-1].max() + maps[i-1].min()) / 2.\n",
    "            for k in range(maps[i-1].shape[0]):\n",
    "                for j in range(maps[i-1].shape[1]):\n",
    "                    ax.text(j, k, fmt_str.format(maps[i-1][k, j]),\n",
    "                            ha=\"center\", va=\"center\", fontsize=8,\n",
    "                            color=\"white\" if maps[i-1][k, j] > thresh else \"black\")\n",
    "\n",
    "    fig.suptitle(plot_title)\n",
    "#     fig.subplots_adjust(right=0.8)\n",
    "#     cbar_ax = fig.add_axes([0.85, 0.15, 0.05, 0.7])\n",
    "#     fig.colorbar(im, cax=cbar_ax)\n",
    "\n",
    "    create_dir('{}'.format(subfolder))\n",
    "    plt.savefig('{}/{}.{}'.format(subfolder, filename, fmt), bbox_inches='tight')\n",
    "    plt.show()\n",
    "    plt.close()\n",
    "\n",
    "def plot_lines(lines, errors, x_labels, x_vals, plot_title, x_title, y_title, linestyles, \n",
    "               colors, labels, subfolder, filename, fmt='png', separate_legend=False):\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.set_title(plot_title)\n",
    "    ax.set_xlabel(x_title)\n",
    "    ax.set_ylabel(y_title)\n",
    "    ax.set_xticks(x_vals)\n",
    "    ax.set_xticklabels(x_labels)\n",
    "\n",
    "    for i in range(len(lines)):\n",
    "        ax.errorbar(x_vals, lines[i], yerr=errors[i], color=colors[i], \n",
    "                    label=labels[i], linestyle=linestyles[i], marker='o')\n",
    "    \n",
    "    create_dir('{}'.format(subfolder))\n",
    "    \n",
    "    if not separate_legend:\n",
    "        box = ax.get_position()\n",
    "        ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])\n",
    "        ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "    else:\n",
    "        fig_legend = plt.figure(figsize=(3, 3))\n",
    "        handles, labels = ax.get_legend_handles_labels()\n",
    "        fig_legend.legend(handles, labels, 'center', ncol=1)\n",
    "        fig_legend.savefig('{}/{}_legend.{}'.format(subfolder, filename, fmt), \n",
    "                           bbox_inches='tight')\n",
    "    \n",
    "    \n",
    "    fig.savefig('{}/{}.{}'.format(subfolder, filename, fmt), bbox_inches='tight')\n",
    "    plt.show()\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_scatter(all_x_vals, all_y_vals, colors, shapes, annotations, annotation_positions, plot_title, \n",
    "                 x_title, y_title, subfolder, filename, labels, separate_legend=False, fmt='png'):\n",
    "    if isinstance(annotations[0], int) or isinstance(annotations[0], float):\n",
    "        annotations = list(map(str, annotations))\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.set_title(plot_title)\n",
    "    ax.set_xlabel(x_title)\n",
    "    ax.set_ylabel(y_title)\n",
    "    for i, (x_vals, y_vals) in enumerate(zip(all_x_vals, all_y_vals)):\n",
    "        ax.scatter(x_vals, y_vals, s=64, c=colors[i], marker=shapes[i], label=labels[i])\n",
    "    flattened_x_vals, flattened_y_vals = [], []\n",
    "    for x_vals in all_x_vals:\n",
    "        for x in x_vals:\n",
    "            flattened_x_vals.append(x)\n",
    "    for y_vals in all_y_vals:\n",
    "        for y in y_vals:\n",
    "            flattened_y_vals.append(y)\n",
    "    arrowprops = dict(arrowstyle = \"<-\")\n",
    "    for i, txt in enumerate(annotations):\n",
    "        if txt != '':\n",
    "            ax.annotate(txt, xytext=annotation_positions[i], xy=(flattened_x_vals[i], flattened_y_vals[i]), \n",
    "                        arrowprops=arrowprops)\n",
    "    \n",
    "    create_dir('{}'.format(subfolder))\n",
    "    \n",
    "    if not separate_legend:\n",
    "        box = ax.get_position()\n",
    "        ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])\n",
    "        ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "    else:\n",
    "        fig_legend = plt.figure(figsize=(3, 3))\n",
    "        handles, labels = ax.get_legend_handles_labels()\n",
    "        fig_legend.legend(handles, labels, 'center', ncol=1)\n",
    "        fig_legend.savefig('{}/{}_legend.{}'.format(subfolder, filename, fmt), \n",
    "                           bbox_inches='tight')\n",
    "    \n",
    "    fig.savefig('{}/{}.{}'.format(subfolder, filename, fmt), bbox_inches='tight')\n",
    "    plt.show()\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}