02-figure2-objective-correlation-matrix.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2e372c41-beb6-49c3-a00e-8bd938fa354a",
   "metadata": {},
   "source": [
    "# How to select an objective function using information theory"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fbc99a9-9a6c-48ff-bfcf-4c5c30b3a04d",
   "metadata": {},
   "source": [
    "Reproduce figure 2 of the paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ef72a63-4417-4cb0-beea-260bb1b80c42",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# compute likelihood\n",
    "import numpy as np\n",
    "from scipy.stats import pearsonr\n",
    "import scipy.stats\n",
    "\n",
    "\n",
    "def normal_ll(y, y_hat, transform=None, gradient=1):\n",
    "    '''Log likelihood for the normal distribution with change of variable\n",
    "    \n",
    "    The normal distribution is the formal likelihood for the mean squared error (MSE).\n",
    "    \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    y : array_like\n",
    "        Observations.\n",
    "    y_hat : array_like\n",
    "        Predictions.\n",
    "    transform : function\n",
    "        Change of variable transformation.\n",
    "    gradient : function\n",
    "        Gradient of the transform function.\n",
    "        \n",
    "    Proof\n",
    "    -----\n",
    "    https://www.statlect.com/probability-distributions/normal-distribution\n",
    "    '''\n",
    "    if transform is not None:\n",
    "        y = transform(y)\n",
    "        y_hat = transform(y_hat)\n",
    "        \n",
    "    e = y - y_hat\n",
    "    n = len(e)\n",
    "    sigma = e.std()\n",
    "    log_gradient = np.sum(np.log(np.abs(gradient)))\n",
    "    ll = -n * np.log(sigma) - n/2*np.log(2*np.pi) - 1/(2*sigma**2) * (e**2).sum() + log_gradient\n",
    "    return ll\n",
    "\n",
    "\n",
    "def laplace_ll(y, y_hat, transform=None, gradient=1):\n",
    "    '''Log likelihood for Laplace distribution with change of variable\n",
    "    \n",
    "    The laplace distribution is the formal likelihood for the mean absolute\n",
    "    error (MAE).\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    y : array_like\n",
    "        Observations.\n",
    "    y_hat : array_like\n",
    "        Predictions.\n",
    "    transform : function\n",
    "        Change of variable transformation.\n",
    "    gradient : function\n",
    "        Gradient of the transform function.\n",
    "    '''\n",
    "    if transform is not None:\n",
    "        y = transform(y)\n",
    "        y_hat = transform(y_hat)\n",
    "        \n",
    "    e = (y - y_hat).abs()\n",
    "    n = len(e)\n",
    "    b = e.mean()\n",
    "    log_gradient = np.sum(np.log(np.abs(gradient)))\n",
    "    ll = -n * np.log(2*b) - 1/b * e.sum() + log_gradient\n",
    "    return ll.sum()\n",
    "                                   \n",
    "\n",
    "def msre_ll(y, y_hat):\n",
    "    '''Log likelihood for mean squared square-root error\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    y : array_like\n",
    "    y_hat : array_like\n",
    "    '''\n",
    "    return normal_ll(y, y_hat, transform=lambda x: np.sqrt(x), gradient=-1/(2*np.sqrt(y)))\n",
    "\n",
    "\n",
    "def mare_ll(y, y_hat):\n",
    "    '''Log likelihood for mean absolute square-root error\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    y : array_like\n",
    "    y_hat : array_like\n",
    "    '''\n",
    "    return laplace_ll(y, y_hat, transform=lambda x: np.sqrt(x), gradient=-1/(2*np.sqrt(y)))\n",
    "\n",
    "\n",
    "def lognormal_ll(y, y_hat):\n",
    "    '''Lognormal log likelihood\n",
    "    \n",
    "    The lognormal distribution is the formal likelihood for the mean squared\n",
    "    log error (MSLE).\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    y : array_like\n",
    "    y_hat : array_like\n",
    "    '''\n",
    "    return normal_ll(y, y_hat, transform=lambda x: np.log(x), gradient=1/y)\n",
    "\n",
    "\n",
    "def mspe_ll(y, y_hat):\n",
    "    '''Log likelhood for mean squared percentage error\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    y : array_like\n",
    "    y_hat : array_like\n",
    "    \n",
    "    '''\n",
    "    return normal_ll(y, y_hat, transform=lambda x: x/y, gradient=-1/(y**2)) \n",
    "\n",
    "\n",
    "def nse_ll(y, y_hat, group='gage_id'):\n",
    "    '''Log likelihood for normalized squared error (NSE)\n",
    "    \n",
    "    NSE is equivalent to the Nash–Sutcliffe model efficiency coefficient.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    y : array_like\n",
    "    y_hat : array_like\n",
    "    '''\n",
    "    sigma_o = y.groupby('gage_id').transform(lambda x: x.std())\n",
    "    return normal_ll(y, y_hat, transform=lambda x: x/sigma_o, gradient=1/sigma_o)\n",
    "\n",
    "\n",
    "def loglaplace_ll(y, y_hat):\n",
    "    '''Log likelihood for log Laplace distribution\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    y : array_like\n",
    "    y_hat : array_like\n",
    "    '''\n",
    "    return laplace_ll(y, y_hat, transform=lambda x: np.log(x), gradient=1/y)\n",
    "\n",
    "\n",
    "def uniform_ll(y, y_hat):\n",
    "    '''Log likelihood for uniform distribution.\n",
    "    \n",
    "    The uniform log likelihood minimizes the maximum error.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    y : array_like\n",
    "    y_hat : array_like\n",
    "    '''\n",
    "    e = np.abs(y - y_hat)\n",
    "    n = len(e)\n",
    "    #ll = -n * np.log(e.max()-e.min()) # standard formulation\n",
    "    ll = -n * np.log(e.max() - 0)\n",
    "    return ll\n",
    "\n",
    "\n",
    "def bernoulli_ll(y, y_hat, groupby=None):\n",
    "    '''TODO and use within zi_ll\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    y : array_like\n",
    "    y_hat : array_like\n",
    "    '''\n",
    "    pass\n",
    "\n",
    "\n",
    "\n",
    "def zi_ll(y, y_hat, ll=normal_ll, threshold=0.01, groupby=None):\n",
    "    ''' Zero-inflated log likelihood.\n",
    "    \n",
    "     Parameters\n",
    "    ----------\n",
    "    y : array_like\n",
    "    y_hat : array_like\n",
    "    ll : function\n",
    "        Zero-inflated log likelihood \n",
    "    threshold : float\n",
    "        Value below which is treated as zero\n",
    "    groupby : string\n",
    "        Optional groupby term (testing)\n",
    "    '''\n",
    "    y_o = y <= threshold\n",
    "    y_hat_o = y_hat <= threshold\n",
    "    \n",
    "    if groupby is None:\n",
    "        n1 = (y_o & y_hat_o).sum() # correct zero-flow prediction\n",
    "        n2 = (y_o ^ y_hat_o).sum() # incorrect zero-flow prediction \n",
    "    else:\n",
    "        n1 = (y_o & y_hat_o).groupby(groupby).sum() # correct zero-flow prediction\n",
    "        n2 = (y_o ^ y_hat_o).groupby(groupby).sum() # incorrect zero-flow prediction\n",
    "\n",
    "    n3 = (~y_o & ~y_hat_o) # correct flow predictions\n",
    "    \n",
    "    # fraction of correctly predicted zero flows\n",
    "    rho = np.where( (n1+n2) == 0, 0, n1 / (n1 + n2))\n",
    "    n_rho = 1-rho\n",
    "    \n",
    "    # n1 * np.log(rho) + n2 * np.log(1-rho)\n",
    "    ll_zero = n1[rho!=0] * np.log(rho[rho!=0]) + n2[n_rho!=0]* np.log(n_rho[n_rho!=0])\n",
    "    \n",
    "    return ll_zero.sum() + ll(y[n3], y_hat[n3])\n",
    "\n",
    "\n",
    "def zilognormal_ll(y, y_hat):\n",
    "    '''Log likelihood for zero-inflated lognormal.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    y : array_like\n",
    "    y_hat : array_like\n",
    "    '''\n",
    "       \n",
    "    return zi_ll(y, y_hat, ll=lognormal_ll, threshold=0.01)\n",
    "\n",
    "\n",
    "def ziloglaplace_ll(y, y_hat):\n",
    "    '''Log likelihood for zero-inflated laplace.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    y : array_like\n",
    "    y_hat : array_like\n",
    "    '''\n",
    "    return zi_ll(y, y_hat, ll=loglaplace_ll, threshold=0.01)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1a1cdb3-b7c1-4fd1-8317-e0312047d90b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "# read local copy\n",
    "df = pd.read_parquet('gages2_nndar.parquet')\n",
    "df[df < 0.01] = 0.01\n",
    "\n",
    "sigma_o = df['obs'].groupby('gage_id').transform(lambda x: x.std())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b31809c-6c5d-43ca-9058-270b85abf28b",
   "metadata": {},
   "outputs": [],
   "source": [
    "count = df.groupby('gage_id').count()['obs']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c27dbc9f-6724-4123-90c3-60ff62c5e877",
   "metadata": {},
   "outputs": [],
   "source": [
    "# step 1: create a table of objective functions\n",
    "objectives = {\n",
    "    'U' : {'name':'uniformly distributed error', 'f':uniform_ll},\n",
    "    'MSE' : {'name':'mean squared error', 'f':normal_ll},\n",
    "    'NSE' : {'name':'normalized squared error', 'f':nse_ll},\n",
    "    'MAE' : {'name': 'mean absolute error', 'f':laplace_ll},\n",
    "    'MSPE' : {'name': 'mean squared percent error', 'f':mspe_ll},\n",
    "    'MSLE' : {'name':'mean squared log error*', 'f':lognormal_ll},\n",
    "    'MALE' : {'name':'mean absolute log error*', 'f':loglaplace_ll},\n",
    "    'ZMSLE' : {'name':'zero-inflated MSLE', 'f':zilognormal_ll},\n",
    "    'ZMALE' : {'name':'zero-inflated MALE', 'f':ziloglaplace_ll},\n",
    "    'MARE' : {'name':'mean absolute square root error', 'f':mare_ll},\n",
    "}\n",
    "\n",
    "obj_df = pd.DataFrame.from_dict(objectives, orient='index')\n",
    "\n",
    "# step 2: compute the information in each objective function\n",
    "results = []\n",
    "for index, row in obj_df.iterrows():\n",
    "    # nats is the negative log likelihood or the info in the error\n",
    "    #results.append(- row.f(df.obs, df.NNDAR)/count/np.log(2) )\n",
    "    results.append(- df.groupby('gage_id').apply(lambda x: row.f(x['obs'], x['NNDAR']))/count/np.log(2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f6e7d7d-9605-48e6-aed1-6683f428f3a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = ['MSPE','U','MSE','NSE','MAE', 'MSLE','MARE','ZMSLE','MALE','ZMALE']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4ec7002-8306-45e3-a9b8-55fb03bc186b",
   "metadata": {},
   "outputs": [],
   "source": [
    "results_df = pd.DataFrame({c: i for i,c in zip(results,columns)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ce97ab8-8cc6-46fb-9744-0de58cc10910",
   "metadata": {},
   "outputs": [],
   "source": [
    "results_df[results_df > 200] = np.nan # nan any crazy values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1d1d1ad-3de5-4fb5-a72b-f469a364b898",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "#sns.set_theme(style=\"white\")\n",
    "plt.rc('font', family='serif')\n",
    "plt.rc('xtick', labelsize='x-small')\n",
    "plt.rc('ytick', labelsize='x-small')\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.rcParams['pdf.fonttype'] = 42\n",
    "matplotlib.rcParams['ps.fonttype'] = 42"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "726b6925-331e-45f3-81c1-a04cf1b6fd4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from string import ascii_letters\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Compute the correlation matrix\n",
    "corr = results_df.corr()\n",
    "\n",
    "# Generate a mask for the upper triangle\n",
    "mask = np.triu(np.ones_like(corr, dtype=bool)) - np.eye(corr.shape[0])\n",
    "\n",
    "# Set up the matplotlib figure\n",
    "#f, ax = plt.subplots(figsize=(5.512, 5))\n",
    "f, ax = plt.subplots(figsize=(3.54, 3))\n",
    "# Generate a custom diverging colormap\n",
    "#cmap = sns.diverging_palette(230, 20, as_cmap=True)\n",
    "\n",
    "# Draw the heatmap with the mask and correct aspect ratio\n",
    "sns.heatmap(corr, mask=mask, vmax=1, vmin=0.7, cmap='flare',\n",
    "            square=True, linewidths=0.5, cbar_kws={\"shrink\": 1})\n",
    "\n",
    "f.tight_layout()\n",
    "f.savefig('figure2.pdf')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}