{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "%matplotlib inline\n",
    "from matplotlib import pyplot as plt\n",
    "from pylab import rcParams\n",
    "rcParams['figure.figsize'] = 7,7 \n",
    "import seaborn as sns\n",
    "sns.set(color_codes=True, font_scale=1.2)\n",
    "\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 197916 characters]\n",
      "text/plain": [
       "<Figure size 504x504 with 2 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 459,
       "width": 524
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Step 0 - Read the dataset, calculate column correlations and make a seaborn heatmap\n",
    "\n",
    "data = pd.read_csv('https://raw.githubusercontent.com/drazenz/heatmap/master/autos.clean.csv')\n",
    "\n",
    "corr = data.corr()\n",
    "ax = sns.heatmap(\n",
    "    corr, \n",
    "    vmin=-1, vmax=1, center=0,\n",
    "    cmap=sns.diverging_palette(20, 220, n=200),\n",
    "    square=True\n",
    ")\n",
    "ax.set_xticklabels(\n",
    "    ax.get_xticklabels(),\n",
    "    rotation=45,\n",
    "    horizontalalignment='right'\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 87348 characters]\n",
      "text/plain": [
       "<Figure size 504x504 with 2 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 487,
       "width": 569
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Step 5 - Add color legend\n",
    "import heatmap as hm\n",
    "# def heatmap(x, y, size, color):\n",
    "#     fig, ax = plt.subplots()\n",
    "    \n",
    "#     # Mapping from column names to integer coordinates\n",
    "#     x_labels = [v for v in sorted(x.unique())]\n",
    "#     y_labels = [v for v in sorted(y.unique())]\n",
    "#     x_to_num = {p[1]:p[0] for p in enumerate(x_labels)} \n",
    "#     y_to_num = {p[1]:p[0] for p in enumerate(y_labels)} \n",
    "    \n",
    "#     size_scale = 500\n",
    "    \n",
    "#     n_colors = 256 # Use 256 colors for the diverging color palette\n",
    "#     palette = sns.diverging_palette(20, 220, n=n_colors) # Create the palette\n",
    "#     color_min, color_max = [-1, 1] # Range of values that will be mapped to the palette, i.e. min and max possible correlation\n",
    "\n",
    "#     def value_to_color(val):\n",
    "#         val_position = float((val - color_min)) / (color_max - color_min) # position of value in the input range, relative to the length of the input range\n",
    "#         ind = int(val_position * (n_colors - 1)) # target index in the color palette\n",
    "#         return palette[ind]\n",
    "    \n",
    "#     plot_grid = plt.GridSpec(1, 15, hspace=0.2, wspace=0.1) # Setup a 1x15 grid\n",
    "#     ax = plt.subplot(plot_grid[:,:-1]) # Use the leftmost 14 columns of the grid for the main plot\n",
    "    \n",
    "#     ax.scatter(\n",
    "#         x=x.map(x_to_num), # Use mapping for x\n",
    "#         y=y.map(y_to_num), # Use mapping for y\n",
    "#         s=size * size_scale, # Vector of square sizes, proportional to size parameter\n",
    "#         c=color.apply(value_to_color), # Vector of square colors, mapped to color palette\n",
    "#         marker='s' # Use square as scatterplot marker\n",
    "#     )\n",
    "    \n",
    "#     # Show column labels on the axes\n",
    "#     ax.set_xticks([x_to_num[v] for v in x_labels])\n",
    "#     ax.set_xticklabels(x_labels, rotation=45, horizontalalignment='right')\n",
    "#     ax.set_yticks([y_to_num[v] for v in y_labels])\n",
    "#     ax.set_yticklabels(y_labels)\n",
    "    \n",
    "#     ax.grid(False, 'major') # Turn off major gridlines\n",
    "#     ax.grid(True, 'minor') # Turn on minor gridlines\n",
    "#     ax.set_xticks([t + 0.5 for t in ax.get_xticks()], minor=True) # Set gridlines to appear between integer coordinates\n",
    "#     ax.set_yticks([t + 0.5 for t in ax.get_yticks()], minor=True) # Do the same for y axis\n",
    "    \n",
    "#     ax.set_xlim([-0.5, max([v for v in x_to_num.values()]) + 0.5]) \n",
    "#     ax.set_ylim([-0.5, max([v for v in y_to_num.values()]) + 0.5])\n",
    "       \n",
    "#     # Add color legend on the right side of the plot\n",
    "#     ax = plt.subplot(plot_grid[:,-1]) # Use the rightmost column of the plot\n",
    "\n",
    "#     col_x = [0]*len(palette) # Fixed x coordinate for the bars\n",
    "#     y=np.linspace(color_min, color_max, n_colors) # y coordinates for each of the n_colors bars\n",
    "\n",
    "#     bar_height = y[1] - y[0]\n",
    "#     ax.barh(\n",
    "#         y=y,\n",
    "#         width=[5]*len(palette), # Make bars 5 units wide\n",
    "#         left=col_x, # Make bars start at 0\n",
    "#         height=bar_height,\n",
    "#         color=palette,\n",
    "#         linewidth=0\n",
    "#     )\n",
    "#     ax.set_xlim(1, 2) # Bars are going from 0 to 5, so lets crop the plot somewhere in the middle\n",
    "#     ax.grid(False) # Hide grid\n",
    "#     ax.set_facecolor('white') # Make background white\n",
    "#     ax.set_xticks([]) # Remove horizontal ticks\n",
    "#     ax.set_yticks(np.linspace(min(y), max(y), 3)) # Show vertical ticks for min, middle and max\n",
    "#     ax.yaxis.tick_right() # Show vertical ticks on the right \n",
    "    \n",
    "columns = ['bore', 'stroke', 'compression-ratio', 'horsepower', 'city-mpg', 'price'] \n",
    "corr = data[columns].corr()\n",
    "corr = pd.melt(corr.reset_index(), id_vars='index') # Unpivot the dataframe, so we can get pair of arrays for x and y\n",
    "corr.columns = ['x', 'y', 'value']\n",
    "hm.heatmap(\n",
    "    x=corr['x'],\n",
    "    y=corr['y'],\n",
    "    size=corr['value'].abs(),\n",
    "    color=corr['value']\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 6 - Extract all the parameters into **kwargs and make the function more robust\n",
    "\n",
    "def heatmap(x, y, **kwargs):\n",
    "    if 'color' in kwargs:\n",
    "        color = kwargs['color']\n",
    "    else:\n",
    "        color = [1]*len(x)\n",
    "        \n",
    "    if 'palette' in kwargs:\n",
    "        palette = kwargs['palette']\n",
    "        n_colors = len(palette)\n",
    "    else:\n",
    "        n_colors = 256 # Use 256 colors for the diverging color palette\n",
    "        palette = sns.color_palette(\"Blues\", n_colors) \n",
    "        \n",
    "    if 'color_range' in kwargs:\n",
    "        color_min, color_max = kwargs['color_range']\n",
    "    else:\n",
    "        color_min, color_max = min(color), max(color) # Range of values that will be mapped to the palette, i.e. min and max possible correlation\n",
    "\n",
    "    def value_to_color(val):\n",
    "        if color_min == color_max:\n",
    "            return palette[-1]\n",
    "        else:\n",
    "            val_position = float((val - color_min)) / (color_max - color_min) # position of value in the input range, relative to the length of the input range\n",
    "            val_position = min(max(val_position, 0), 1) # bound the position betwen 0 and 1\n",
    "            ind = int(val_position * (n_colors - 1)) # target index in the color palette\n",
    "            return palette[ind]\n",
    "    \n",
    "    if 'size' in kwargs:\n",
    "        size = kwargs['size']\n",
    "    else:\n",
    "        size = [1]*len(x)\n",
    "    \n",
    "    if 'size_range' in kwargs:\n",
    "        size_min, size_max = kwargs['size_range'][0], kwargs['size_range'][1]\n",
    "    else:\n",
    "        size_min, size_max = min(size), max(size)\n",
    "        \n",
    "    size_scale = kwargs.get('size_scale', 500)\n",
    "    \n",
    "    def value_to_size(val):\n",
    "        if size_min == size_max:\n",
    "            return 1 * size_scale\n",
    "        else:\n",
    "            val_position = (val - size_min) * 0.99 / (size_max - size_min) + 0.01 # position of value in the input range, relative to the length of the input range\n",
    "            val_position = min(max(val_position, 0), 1) # bound the position betwen 0 and 1\n",
    "            return val_position * size_scale\n",
    "    if 'x_order' in kwargs: \n",
    "        x_names = [t for t in kwargs['x_order']]\n",
    "    else:\n",
    "        x_names = [t for t in sorted(set([v for v in x]))]\n",
    "    x_to_num = {p[1]:p[0] for p in enumerate(x_names)}\n",
    "    \n",
    "    if 'y_order' in kwargs: \n",
    "        y_names = [t for t in kwargs['y_order']]\n",
    "    else:\n",
    "        y_names = [t for t in sorted(set([v for v in y]))]\n",
    "    y_to_num = {p[1]:p[0] for p in enumerate(y_names)}\n",
    "    \n",
    "    plot_grid = plt.GridSpec(1, 15, hspace=0.2, wspace=0.1) # Setup a 1x10 grid\n",
    "    ax = plt.subplot(plot_grid[:,:-1]) # Use the left 14/15ths of the grid for the main plot\n",
    "        \n",
    "    marker = kwargs.get('marker', 's')\n",
    "    \n",
    "    kwargs_pass_on = {k:v for k,v in kwargs.items() if k not in [\n",
    "         'color', 'palette', 'color_range', 'size', 'size_range', 'size_scale', 'marker', 'x_order', 'y_order'\n",
    "    ]}\n",
    "    \n",
    "    ax.scatter(\n",
    "        x=[x_to_num[v] for v in x],\n",
    "        y=[y_to_num[v] for v in y],\n",
    "        marker=marker,\n",
    "        s=[value_to_size(v) for v in size], \n",
    "        c=[value_to_color(v) for v in color],\n",
    "        **kwargs_pass_on\n",
    "    )\n",
    "    ax.set_xticks([v for k,v in x_to_num.items()])\n",
    "    ax.set_xticklabels([k for k in x_to_num], rotation=45, horizontalalignment='right')\n",
    "    ax.set_yticks([v for k,v in y_to_num.items()])\n",
    "    ax.set_yticklabels([k for k in y_to_num])\n",
    "    \n",
    "    ax.grid(False, 'major')\n",
    "    ax.grid(True, 'minor')\n",
    "    ax.set_xticks([t + 0.5 for t in ax.get_xticks()], minor=True)\n",
    "    ax.set_yticks([t + 0.5 for t in ax.get_yticks()], minor=True)\n",
    "\n",
    "    ax.set_xlim([-0.5, max([v for v in x_to_num.values()]) + 0.5])\n",
    "    ax.set_ylim([-0.5, max([v for v in y_to_num.values()]) + 0.5])\n",
    "    ax.set_facecolor('#F1F1F1')\n",
    "    \n",
    "    # Add color legend on the right side of the plot\n",
    "    if color_min < color_max:\n",
    "        ax = plt.subplot(plot_grid[:,-1]) # Use the rightmost column of the plot\n",
    "\n",
    "        col_x = [0]*len(palette) # Fixed x coordinate for the bars\n",
    "        y=np.linspace(color_min, color_max, n_colors) # y coordinates for each of the n_colors bars\n",
    "\n",
    "        bar_height = y[1] - y[0]\n",
    "        ax.barh(\n",
    "            y=y,\n",
    "            width=[5]*len(palette), # Make bars 5 units wide\n",
    "            left=col_x, # Make bars start at 0\n",
    "            height=bar_height,\n",
    "            color=palette,\n",
    "            linewidth=0\n",
    "        )\n",
    "        ax.set_xlim(1, 2) # Bars are going from 0 to 5, so lets crop the plot somewhere in the middle\n",
    "        ax.grid(False) # Hide grid\n",
    "        ax.set_facecolor('white') # Make background white\n",
    "        ax.set_xticks([]) # Remove horizontal ticks\n",
    "        ax.set_yticks(np.linspace(min(y), max(y), 3)) # Show vertical ticks for min, middle and max\n",
    "        ax.yaxis.tick_right() # Show vertical ticks on the right \n",
    "     \n",
    "\n",
    "def corrplot(data, size_scale=500, marker='s'):\n",
    "    corr = pd.melt(data.reset_index(), id_vars='index')\n",
    "    corr.columns = ['x', 'y', 'value']\n",
    "    heatmap(\n",
    "        corr['x'], corr['y'],\n",
    "        color=corr['value'], color_range=[-1, 1],\n",
    "        palette=sns.diverging_palette(20, 220, n=256),\n",
    "        size=corr['value'].abs(), size_range=[0,1],\n",
    "        marker=marker,\n",
    "        x_order=data.columns,\n",
    "        y_order=data.columns[::-1],\n",
    "        size_scale=size_scale\n",
    "    )\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index([' symboling', 'normalized-losses', 'make', 'fuel-type', 'aspiration',\n",
       "       'num-of-doors', 'body-style', 'drive-wheels', 'engine-location',\n",
       "       'wheel-base', 'length', 'width', 'height', 'curb-weight', 'engine-type',\n",
       "       'num-of-cylinders', 'engine-size', 'fuel-system', 'bore', 'stroke',\n",
       "       'compression-ratio', 'horsepower', 'peak-rpm', 'city-mpg',\n",
       "       'highway-mpg', 'price'],\n",
       "      dtype='object')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th>drive-wheels</th>\n",
       "      <th>4wd</th>\n",
       "      <th>fwd</th>\n",
       "      <th>rwd</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>horsepower-group</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Low (0-100)</th>\n",
       "      <td>5.0</td>\n",
       "      <td>89.0</td>\n",
       "      <td>15.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Medium (100-150)</th>\n",
       "      <td>3.0</td>\n",
       "      <td>24.0</td>\n",
       "      <td>35.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>High (150+)</th>\n",
       "      <td>NaN</td>\n",
       "      <td>5.0</td>\n",
       "      <td>25.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "drive-wheels      4wd   fwd   rwd\n",
       "horsepower-group                 \n",
       "Low (0-100)       5.0  89.0  15.0\n",
       "Medium (100-150)  3.0  24.0  35.0\n",
       "High (150+)       NaN   5.0  25.0"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 79136 characters]\n",
      "text/plain": [
       "<Figure size 504x504 with 2 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 487,
       "width": 544
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 504x504 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 46788 characters]\n",
      "text/plain": [
       "<Figure size 504x504 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 447,
       "width": 513
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 53544 characters]\n",
      "text/plain": [
       "<Figure size 504x504 with 2 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 399,
       "width": 523
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "from heatmap import heatmap, corrplot\n",
    "columns = ['bore', 'stroke', 'compression-ratio', 'horsepower', 'city-mpg', 'price'] \n",
    "plt.figure(figsize=(7, 7))\n",
    "corrplot(data[columns].corr().replace(np.nan, 0), size_scale=500)\n",
    "\n",
    "plt.figure()\n",
    "\n",
    "display(data.columns)\n",
    "data['cnt'] = np.ones(len(data))\n",
    "bin_labels = ['Low (0-100)', 'Medium (100-150)', 'High (150+)']\n",
    "\n",
    "data['horsepower-group'] = pd.cut(data['horsepower'], [0, 100, 150, data['horsepower'].max()], labels=bin_labels)\n",
    "g = data.groupby(['horsepower-group', 'drive-wheels']).count()[['cnt']].reset_index()\n",
    "plt.figure(figsize=(7, 7))\n",
    "heatmap(y=g['horsepower-group'], x=g['drive-wheels'], size=g['cnt'], size_scale=10000, x_order=['4wd', 'fwd', 'rwd'], y_order=bin_labels[::-1])\n",
    "plt.xlabel('Drivetrain layout')\n",
    "plt.ylabel('Horsepower')\n",
    "\n",
    "plt.figure()\n",
    "sns.heatmap(g.pivot('horsepower-group', 'drive-wheels', 'cnt'), square=True)\n",
    "plt.xlabel('Drivetrain layout')\n",
    "plt.ylabel('Horsepower')\n",
    "display(g.pivot('horsepower-group', 'drive-wheels', 'cnt'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 132240 characters]\n",
      "text/plain": [
       "<Figure size 216x792 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 706,
       "width": 271
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "bin_labels = ['Low (0-100)', 'Medium (100-150)', 'High (150+)']\n",
    "data['horsepower-group'] = pd.cut(\n",
    "    data['horsepower'],\n",
    "    bins=[0, 100, 150, data['horsepower'].max()],\n",
    "    labels=bin_labels\n",
    ")\n",
    "\n",
    "data['cnt'] = np.ones(len(data))\n",
    "g = data.groupby(['horsepower-group', 'make']).count()[['cnt']].reset_index().replace(np.nan, 0)\n",
    "plt.figure(figsize=(3, 11))\n",
    "heatmap(\n",
    "    x=g['horsepower-group'],\n",
    "    y=g['make'],\n",
    "    size=g['cnt'],\n",
    "    marker='h',\n",
    "    x_order=bin_labels\n",
    ")"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
