MDN.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Built upon this blopgpost: https://blog.otoro.net/2015/11/24/mixture-density-networks-with-tensorflow/\n",
    "\n",
    "import glob, os\n",
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import IPython\n",
    "from IPython.display import HTML\n",
    "from itertools import chain\n",
    "import random\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "import pickle\n",
    "import re\n",
    "from matplotlib.ticker import FormatStrFormatter\n",
    "from scipy.stats import norm\n",
    "from scipy.stats import binned_statistic\n",
    "\n",
    "import multiprocessing as mp\n",
    "import math as m\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker as mticker\n",
    "import matplotlib \n",
    "#os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "\n",
    "from joblib import Parallel, delayed, dump, load"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "pathSave = \".\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "class getData:\n",
    "    def __init__(self, _paraVecIndices, _xparaVecIndices, _tProfBounds, _trainPercent, _pathSave, _noiseLevel=0.):\n",
    "        self.paraVecIndices = _paraVecIndices\n",
    "        self.xparaVecIndices = _xparaVecIndices\n",
    "        self.tProfBounds = _tProfBounds\n",
    "        self.trainPercent = _trainPercent\n",
    "        self.pathSave = _pathSave\n",
    "        self.noiseLevel = _noiseLevel\n",
    "        \n",
    "    def visualizeData(self, picID, fontSize, Dict_baseline = {}, plotTProf=True, plotBaseline=False): \n",
    "        x_data, x_test, x_cv, y_data, y_test, y_cv, paraVec, xparaVec, rProf, startPoint, endPoint, indexer \\\n",
    "        = self.x_data, self.x_test, self.x_cv, self.y_data, self.y_test, self.y_cv, self.paraVec, self.xparaVec, \\\n",
    "            self.rProf, self.startPoint, self.endPoint, self.indexer    \n",
    "                \n",
    "        print('Added noise ' + str(self.noiseLevel))\n",
    "        print(\"X_train shape: \" + str(x_data.shape))\n",
    "        print(\"Y_train shape: \" + str(y_data.shape))\n",
    "        print(\"X_test shape: \" + str(x_test.shape))\n",
    "        print(\"Y_test shape: \" + str(y_test.shape))\n",
    "        print(\"X_cv shape: \" + str(x_cv.shape))\n",
    "        print(\"Y_cv shape: \" + str(y_cv.shape))\n",
    "        print()\n",
    "        \n",
    "        alphaVal = 0.1\n",
    "        tProfBoolean = False\n",
    "\n",
    "        if 'Tprof_4p5Gyr' in xparaVec:\n",
    "            xparaVec = xparaVec[0:-1]\n",
    "            tProfBoolean = True\n",
    "\n",
    "        for ind in range(np.size(paraVec)):\n",
    "            fig, ax = plt.subplots()\n",
    "            n, bins, patches = ax.hist(y_data[:, ind], 100, density=1)\n",
    "            ax.set_xlabel(paraVec[ind])\n",
    "            fig = plt.figure(figsize=(14,4), dpi=200)\n",
    "            plt.subplots_adjust(top=0.92, bottom=0.1, left=0.05, right=0.95, hspace=0.10,wspace=0.35)\n",
    "            nCols = np.size(xparaVec)\n",
    "            nRows = 1\n",
    "            for indX in range(nCols):\n",
    "                ax = fig.add_subplot(nRows,nCols,indX+1)\n",
    "                ax.plot(x_data[:,indX], y_data[:, ind],'ro', alpha=alphaVal)\n",
    "                if plotBaseline:\n",
    "                    ax.plot(x_data[:,indX],Dict_baseline['mu_var' + str(indX)] \\\n",
    "                                 +Dict_baseline['variance_var' + str(indX)],'k.')\n",
    "                    ax.plot(x_data[:,indX],Dict_baseline['mu_var' + str(indX)] \\\n",
    "                                 -Dict_baseline['variance_var' + str(indX)],'k.')\n",
    "                ax.yaxis.labelpad = 0.8\n",
    "                ax.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "                plt.title(\"train\")\n",
    "                plt.xlabel(xparaVec[indX])\n",
    "                plt.ylabel(paraVec[ind])\n",
    "            plt.show()\n",
    "            #fig.savefig(pathSave + \"inverseProblemPics/\" + \"Input0_fig\" + str(picID) + \".pdf\", bbox_inches='tight')\n",
    "            \n",
    "        if  tProfBoolean and plotTProf:   \n",
    "            length = int(np.size(paraVec))\n",
    "            fig = plt.figure(figsize=(length,4), dpi=200)\n",
    "            plt.subplots_adjust(top=0.9, bottom=0.1, left=0.05, right=0.95, hspace=0.10,wspace=0.3)\n",
    "            nCols = np.size(paraVec)\n",
    "            nRows = 2\n",
    "            for indY in range(nCols):\n",
    "                vecInterest = y_data[:, indY]\n",
    "                colorsVec= [plt.cm.jet(i) for i in vecInterest]\n",
    "                ax = fig.add_subplot(nRows,nCols,indY+1)\n",
    "                ax.set_prop_cycle('color',colorsVec)\n",
    "                for i, ii in enumerate(indexer):\n",
    "                    rProf = self.rProf[startPoint:endPoint] # profiles['Dict_Rprof_4p5Gyr' + str(ii)]\n",
    "                    rProf = (rProf-rProf.min())/(rProf.max()-rProf.min())\n",
    "                    tProf = x_data[i,np.size(xparaVec):]\n",
    "                    ax.plot(tProf, rProf, linewidth=0.5)\n",
    "                ax.yaxis.labelpad = 0.8\n",
    "                ax.yaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
    "                if indY > 0:\n",
    "                    ax.axes.get_yaxis().set_ticks([])\n",
    "                ax.set_title(paraVec[indY])\n",
    "                ax.set_xlabel(\"T\") #[K]\")\n",
    "                ax.set_ylabel(\"R\") #[km] \")\n",
    "                for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] +\n",
    "                     ax.get_xticklabels() + ax.get_yticklabels()):\n",
    "                    item.set_fontsize(12)\n",
    "                #ax.set_xlim([250,2250])\n",
    "                #ax.set_ylim([1700,3400])\n",
    "            plt.show()\n",
    "           # fig.savefig(pathSave + \"inverseProblemPics/\" + \"Input1_fig\" + str(picID) + \".pdf\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "class plotResults:\n",
    "    def __init__(self, data, ax, ax1, alphaVal, title, _y_actual, Dict_prob_paras ,indPara,\\\n",
    "                 Dict_baseline, numParameters, kernel, plotIndividualPDFs=False):\n",
    "        \n",
    "        if \"Test\" in title:\n",
    "            self.Dict_Test = Dict_prob_paras\n",
    "        \n",
    "        rho_m  = 3500. \n",
    "        g     = 3.7 \n",
    "        alpha_m = 2.5e-5\n",
    "        T_delta = 2000. \n",
    "        D = 1700e+3         \n",
    "        k_diffusive = 1e-6 \n",
    "        R = 8.314   \n",
    "        \n",
    "        def format_func(_val, tick_number):\n",
    "            f = mticker.ScalarFormatter(useOffset=False, useMathText=True)\n",
    "            _g = lambda x,pos : \"${}$\".format(f._formatSciNotation('%1.2e' % x))\n",
    "            fmt = mticker.FuncFormatter(_g)\n",
    "            return \"{}\".format(fmt(_val))\n",
    "\n",
    "        def dimensionalize(_val,_ind,isVariance=False):\n",
    "            _min = data.pMin[_ind]\n",
    "            _max = data.pMax[_ind]\n",
    "            _val = _val*(_max-_min) \n",
    "            if not isVariance:\n",
    "                _val = _val + _min\n",
    "            \n",
    "            if \"Ra\" in title and not isVariance:\n",
    "                _val = np.log10(rho_m * g * alpha_m * T_delta * np.power(D,3.)/(np.power(10.,_val) * k_diffusive))\n",
    "            \n",
    "            if \"ERef\" in title:\n",
    "                _val = _val*(R * T_delta)\n",
    "            \n",
    "            if \"VRef\" in title:\n",
    "                _val = _val*(R * T_delta) /(rho_m * g * D)\n",
    "            \n",
    "            if \"iniTempTop\" in title:\n",
    "                _val = _val*2000 \n",
    "                if not isVariance:\n",
    "                    _val = _val + 250\n",
    "                    \n",
    "            return _val\n",
    "\n",
    "        xx = dimensionalize(_y_actual,indPara)    \n",
    "        variance = []\n",
    "        combined_mdn_std = []\n",
    "        combined_mdn_mean = []\n",
    "        xxSorted = np.sort(xx)        \n",
    "        \n",
    "        xxSize = np.size(xx)\n",
    "        plotEvery = int(np.floor(xxSize*0.1))\n",
    "        xxSorted.shape = (xxSize,1)\n",
    "        Pr_xx = np.zeros((xxSize,xxSize))\n",
    "        Pr_baseline = np.zeros((xxSize,xxSize))\n",
    "        colors = ['r', 'g', 'm']\n",
    "        \n",
    "        def pdf(x):\n",
    "            return 1./np.sqrt(2.*np.pi) * np.exp(-x**2/2.)\n",
    "        #def cdf(x):\n",
    "        #    return (1. + erf(x/np.sqrt(2))) / 2.\n",
    "        def SSkew(x,a,e,w):\n",
    "            t = (x-e) / w\n",
    "            return 2. / w * pdf(t) * cdf(a*t)\n",
    "        def getGamma(x):\n",
    "            return m.gamma(x)\n",
    "        def tf_beta(y, alpha, beta):\n",
    "            Z = np.divide(np.multiply(getGamma(alpha),getGamma(beta)), getGamma(alpha + beta))\n",
    "            result = np.divide(np.multiply(np.power(y,(alpha - 1.)), np.power((1. - y),(beta - 1.))),Z)\n",
    "            return  result\n",
    "             \n",
    "        if plotIndividualPDFs:\n",
    "            fig3 = plt.figure(figsize=(14,18))\n",
    "           # plt.subplots_adjust(top=0.90, bottom=0.08, left=0.10, right=0.95, hspace=0.3,wspace=0.4)\n",
    "            nRows = np.ceil((xxSize/plotEvery)/3)\n",
    "            nCols = 3\n",
    "            plotCounter = 1\n",
    "        if kernel=='skewed_gaussian':\n",
    "            out_skew_test = np.asarray(Dict_prob_paras[\"skew\"+str(ind)])\n",
    "            out_sigma_test = np.asarray(Dict_prob_paras[\"sigma\"+str(ind)])\n",
    "            out_mu_test = np.asarray(Dict_prob_paras[\"mu\"+str(ind)])\n",
    "            out_pi_test = np.asarray(Dict_prob_paras[\"pi\"+str(ind)])\n",
    "            for ind,val in enumerate(xxSorted):\n",
    "                index = np.where(val == xx)[0][0]\n",
    "                muIntermediate = []\n",
    "                piIntermediate = []\n",
    "                sigmaIntermediate = []\n",
    "                skewIntermediate = []\n",
    "                for i in range(out_sigma_test.shape[1]):\n",
    "                    muIntermediate.append(out_mu_test[index,i])\n",
    "                    sigmaIntermediate.append(out_sigma_test[index,i])\n",
    "                    skewIntermediate.append(out_skew_test[index,i])\n",
    "                    piIntermediate.append(out_pi_test[index,i])\n",
    "                #mu = np.sum(np.asarray(piIntermediate) * np.asarray(muIntermediate))\n",
    "                #sigma = np.sum(np.asarray(piIntermediate) * np.asarray(sigmaIntermediate))\n",
    "                #skew = np.sum(np.asarray(piIntermediate) * np.asarray(skewIntermediate))\n",
    "                pr = piIntermediate[0] * SSkew(xxSorted,skewIntermediate[0],muIntermediate[0],sigmaIntermediate[0])\n",
    "                pr.shape = (pr.shape[0])\n",
    "                for i in range(1,out_sigma_test.shape[1]):\n",
    "                    prI = SSkew(xxSorted,skewIntermediate[i],muIntermediate[i],sigmaIntermediate[i])\n",
    "                    prI.shape = (prI.shape[0])\n",
    "                    pr += piIntermediate[i] * prI\n",
    "                Pr_xx[ind,:] = pr\n",
    "\n",
    "                if  plotIndividualPDFs and ind%plotEvery == 0:\n",
    "                    ax3 = fig3.add_subplot(nRows,nCols,plotCounter) \n",
    "                    plotCounter += 1\n",
    "                    legendStr = []\n",
    "                    for i in range(np.size(muIntermediate)):\n",
    "                        if piIntermediate[i] >= np.max(piIntermediate)*1e-15:\n",
    "                            yPDF = SSkew(xxSorted,skewIntermediate[i],muIntermediate[i],sigmaIntermediate[i])\n",
    "                            ax3.plot(xxSorted,yPDF) #/np.max(yPDF))\n",
    "                            legendStr.append(\"%.4f\" % piIntermediate[i])\n",
    "                            #print(np.trapz(yPDF[:,0], xxSorted[:,0]))\n",
    "                    #Pr_w = SSkew(xxSorted,skew,mu,sigma)\n",
    "                    ax3.plot(xxSorted, Pr_xx[ind,:],\"--\")\n",
    "                    print(np.trapz(Pr_xx[ind,:], xxSorted[:,0]))\n",
    "                    ax3.plot(val,max(Pr_xx[ind,:]), \"kx\")\n",
    "                    ax3.legend(legendStr)\n",
    "                    \n",
    "                \n",
    "        elif kernel=='gaussian':\n",
    "            out_pi_test = np.asarray(Dict_prob_paras[\"pi\"+str(indPara)])\n",
    "            out_sigma_test = np.asarray(Dict_prob_paras[\"sigma\"+str(indPara)])\n",
    "            out_mu_test = np.asarray(Dict_prob_paras[\"mu\"+str(indPara)])\n",
    "            for ind,val in enumerate(xxSorted):\n",
    "                index = np.where(val == xx)[0][0]\n",
    "                muIntermediate = []\n",
    "                piIntermediate = []\n",
    "                sigmaIntermediate = []\n",
    "                for i in range(out_sigma_test.shape[1]):\n",
    "                    muIntermediate.append(dimensionalize(out_mu_test[index,i],indPara))\n",
    "                    sigmaIntermediate.append(dimensionalize(out_sigma_test[index,i],indPara,True))\n",
    "                    piIntermediate.append(out_pi_test[index,i])\n",
    "                    \n",
    "                #mu = np.sum(np.asarray(piIntermediate) * np.asarray(muIntermediate))\n",
    "                #sigma = np.sum(np.asarray(piIntermediate) * np.asarray(sigmaIntermediate))\n",
    "                pr = piIntermediate[0] * norm.pdf(xxSorted, muIntermediate[0], sigmaIntermediate[0])\n",
    "                for i in range(1,out_sigma_test.shape[1]):\n",
    "                    pr += piIntermediate[i] * norm.pdf(xxSorted, muIntermediate[i], sigmaIntermediate[i])\n",
    "                pr.shape = (pr.shape[0],)\n",
    "                \n",
    "                mean_mdn = 0\n",
    "                for i in range(out_sigma_test.shape[1]):\n",
    "                    mean_mdn += piIntermediate[i]*muIntermediate[i]\n",
    "                var_mdn = 0\n",
    "                for i in range(out_sigma_test.shape[1]):\n",
    "                    var_mdn += piIntermediate[i]*(np.power(sigmaIntermediate[i],2) + np.power(muIntermediate[i],2) - np.power(mean_mdn,2))\n",
    "                combined_mdn_std.append(np.power(var_mdn,0.5))\n",
    "                combined_mdn_mean.append(mean_mdn)\n",
    "                \n",
    "                Pr_xx[ind,:] = pr #norm.pdf(xxSorted, mu, sigma)\n",
    "                if  plotIndividualPDFs and ind%plotEvery == 0:\n",
    "                    ax3 = fig3.add_subplot(nRows,nCols,plotCounter) \n",
    "                    plotCounter += 1\n",
    "                    legendStr = []\n",
    "                    for i in range(np.size(muIntermediate)):\n",
    "                        if piIntermediate[i] >= np.max(piIntermediate)*1e-15:\n",
    "                            yPDF = piIntermediate[i] * norm.pdf(xxSorted, muIntermediate[i], sigmaIntermediate[i])\n",
    "                            ax3.plot(xxSorted,yPDF,colors[i]+'--', linewidth=2.5)\n",
    "                            legendStr.append(\"%.4f\" % piIntermediate[i])\n",
    "                            if \"Test\" in title:\n",
    "                                self.Dict_Test[\"forFig_\" + title + \"_pdf_mixture\" + str(i) + \"_case_\" + str(ind)] = yPDF\n",
    "                    #Pr_w = norm.pdf(xxSorted, mu, sigma)\n",
    "                    ax3.plot(xxSorted, Pr_xx[ind,:],\"k-\", linewidth=2.5)\n",
    "                    ax3.plot([val, val],[0,max(Pr_xx[ind,:])], \"-\", color='grey', linewidth=2.5)\n",
    "                    \n",
    "                    if \"Test\" in title:\n",
    "                        self.Dict_Test[\"forFig_\" + title + \"_trueVal_\" + str(ind)] = val\n",
    "                    \n",
    "                    #ax3.plot([combined_mdn_mean[ind]+combined_mdn_std[ind], combined_mdn_mean[ind]+combined_mdn_std[ind]],[0,max(Pr_xx[ind,:])], \"k--\", color='grey', linewidth=2.5)\n",
    "                    #ax3.plot([combined_mdn_mean[ind]-combined_mdn_std[ind], combined_mdn_mean[ind]-combined_mdn_std[ind]],[0,max(Pr_xx[ind,:])], \"k--\", color='grey', linewidth=2.5)\n",
    "                    for item in ([ax3.title, ax3.xaxis.label, ax3.yaxis.label] +\n",
    "                     ax3.get_xticklabels() + ax3.get_yticklabels()):\n",
    "                        item.set_fontsize(20)\n",
    "                    if \"ERef\" in title:\n",
    "                        ax3.set_xticks([1e+5,3e+5,5e+5])\n",
    "                        #ax3.set_yticks([1e+5,2e+5,3e+5,4e+5,5e+5])\n",
    "                        ax3.yaxis.set_major_formatter(plt.FuncFormatter(format_func))\n",
    "                        ax3.xaxis.set_major_formatter(plt.FuncFormatter(format_func))\n",
    "                    elif \"VRef\" in title:\n",
    "                        ax3.set_xticks([4e-6, 7e-6, 10e-6])\n",
    "                        #ax3.set_yticks([4e-6, 5e-6, 6e-6, 7e-6, 8e-6, 9e-6, 10e-6])\n",
    "                        ax3.yaxis.set_major_formatter(plt.FuncFormatter(format_func))\n",
    "                        ax3.xaxis.set_major_formatter(plt.FuncFormatter(format_func))\n",
    "                    plt.tight_layout()\n",
    "\n",
    "                    #ax3.legend(legendStr)\n",
    "        elif kernel=='beta':\n",
    "            out_pi_test = np.asarray(Dict_prob_paras[\"pi\"+str(ind)])\n",
    "            out_alpha_test = np.asarray(Dict_prob_paras[\"alpha\"+str(ind)])\n",
    "            out_beta_test = np.asarray(Dict_prob_paras[\"beta\"+str(ind)])\n",
    "            for ind,val in enumerate(xxSorted):\n",
    "                index = np.where(val == xx)[0][0]\n",
    "                alphaIntermediate = []\n",
    "                piIntermediate = []\n",
    "                betaIntermediate = []\n",
    "                for i in range(out_beta_test.shape[1]):\n",
    "                    alphaIntermediate.append(out_alpha_test[index,i])\n",
    "                    betaIntermediate.append(out_beta_test[index,i])\n",
    "                    piIntermediate.append(out_pi_test[index,i])\n",
    "                alpha = np.sum(np.asarray(piIntermediate) * np.asarray(alphaIntermediate))\n",
    "                betaP = np.sum(np.asarray(piIntermediate) * np.asarray(betaIntermediate))\n",
    "                xxSorted.shape = (xxSize)\n",
    "                Pr_xx[ind,:] = tf_beta(xxSorted, alpha, betaP)\n",
    "                if  plotIndividualPDFs and ind%plotEvery == 0:\n",
    "                    ax3 = fig3.add_subplot(nRows,nCols,plotCounter) \n",
    "                    plotCounter += 1\n",
    "                    legendStr = []\n",
    "                    for i in range(np.size(piIntermediate)):\n",
    "                        if piIntermediate[i] >= np.max(piIntermediate)*1e-1:\n",
    "                            yPDF = tf_beta(xxSorted, alphaIntermediate[i], betaIntermediate[i])\n",
    "                            ax3.plot(xxSorted,yPDF)\n",
    "                            legendStr.append(\"%.4f\" % piIntermediate[i])\n",
    "                    Pr_w = tf_beta(xxSorted, alpha,betaP)\n",
    "                    #ax3.ticklabel_format(self, *, axis='both', style='sci')\n",
    "                    ax3.plot(xxSorted, Pr_w,\"--\")\n",
    "                    ax3.plot(val,max(Pr_w), \"kx\")\n",
    "        if \"Test\" in title:\n",
    "            self.Dict_Test[\"forFig_\" + title + \"_xxSorted\"] = xxSorted\n",
    "            self.Dict_Test[\"forFig_\" + title + \"_Pr_xx\"] = Pr_xx\n",
    "        \n",
    "        for ind,val in enumerate(xxSorted):\n",
    "            index = np.where(val == xx)[0][0]\n",
    "            muBase = []\n",
    "            sigmaBase = []\n",
    "            for baseInd in range(numParameters):\n",
    "                muBase.append(Dict_baseline[\"mu_var\" + str(baseInd)][index])\n",
    "                sigmaBase.append(Dict_baseline[\"variance_var\" + str(baseInd)][index])\n",
    "            muBase = dimensionalize(np.asarray(muBase),indPara)[0]\n",
    "            sigmaBase = dimensionalize(np.asarray(sigmaBase),indPara,True)[0]\n",
    "            #varMP = np.min(np.asarray(sigmaBase))\n",
    "            #muMP = muBase[np.where(varMP == sigmaBase)[0][0]]\n",
    "            prb = 0\n",
    "            #weights = 1/len(sigmaBase)\n",
    "            #for _i,_s in enumerate(sigmaBase):\n",
    "                #prb += weights*norm.pdf(xxSorted, muBase[_i], _s)  \n",
    "            prb = norm.pdf(xxSorted, muBase, sigmaBase)  \n",
    "            prb.shape = (prb.shape[0],)\n",
    "            Pr_baseline[ind,:] = prb \n",
    "            \n",
    "        if \"Test\" in title:\n",
    "            self.Dict_Test[\"forFig_\" + title + \"_averageSTD\"] = np.mean(np.asarray(combined_mdn_std))\n",
    "        \n",
    "        print(\"Average standard deviation: \" + str(np.mean(np.asarray(combined_mdn_std))))\n",
    "        \n",
    "        if \"Ra\" in title:\n",
    "            titlep = \"$\\log(\\eta_{ref})$ [Pa s]\"\n",
    "        elif \"ERef\" in title:\n",
    "            titlep = r\"$E$ [J mol$^{-1}$]\"\n",
    "        elif \"VRef\" in title:\n",
    "            titlep = r\"$V$ [m$^3$ mol$^{-1}$]\"\n",
    "        elif \"Enrichment_cr\" in title:\n",
    "            titlep = \"$\\Lambda$\"\n",
    "        elif \"iniTempTop\" in title:\n",
    "            titlep = \"$T_{ini}$ [K]\"\n",
    "        \n",
    "        x = np.zeros((np.size(xxSorted),np.size(xxSorted)))\n",
    "        y = np.zeros((np.size(xxSorted),np.size(xxSorted)))\n",
    "        for ind,val in enumerate(xxSorted):\n",
    "            x[ind,:] = val\n",
    "            y[:,ind] = xxSorted[ind]\n",
    "        ax.contourf(x,y,Pr_xx)\n",
    "        \n",
    "        if \"Test\" in title:\n",
    "            self.Dict_Test[\"forFig_\" + title + \"x\"] = x\n",
    "            self.Dict_Test[\"forFig_\" + title + \"y\"] = y\n",
    "        \n",
    "        ax.set_xlabel(\"True\")\n",
    "        ax.set_ylabel(\"Predicted\")\n",
    "        ax.set_title(titlep + \"; MDN\") #, fontname=\"Times New Roman Bold\") \n",
    "        #tex = r'$\\bar{\\sigma}$ = ' + str(np.mean(np.asarray(combined_mdn_std)))\n",
    "        \n",
    "        ax1.contourf(x,y,Pr_baseline)  \n",
    "        ax1.set_xlabel(\"True\")\n",
    "        ax1.set_ylabel(\"Predicted\")\n",
    "        ax1.set_title(titlep + \"; MP\")#, fontname=\"Times New Roman Bold) \n",
    "        print()  \n",
    "        if \"Ra\" in title:\n",
    "            ax.set_xticks([19,20,21,22])\n",
    "            ax.set_yticks([19,20,21,22])\n",
    "            #ax.text(1, 1, tex, fontsize=20, va='bottom', color='white')\n",
    "            ax1.set_xticks([19,20,21,22])\n",
    "            ax1.set_yticks([19,20,21,22])\n",
    "            \n",
    "        if \"ERef\" in title:\n",
    "            ax.set_xticks([1e+5,3e+5,5e+5])\n",
    "            ax.set_yticks([1e+5,2e+5,3e+5,4e+5,5e+5])\n",
    "            ax1.set_xticks([1e+5,3e+5,5e+5])\n",
    "            ax1.set_yticks([1e+5,2e+5,3e+5,4e+5,5e+5])\n",
    "            ax.yaxis.set_major_formatter(plt.FuncFormatter(format_func))\n",
    "            ax.xaxis.set_major_formatter(plt.FuncFormatter(format_func))\n",
    "            ax1.yaxis.set_major_formatter(plt.FuncFormatter(format_func))\n",
    "            ax1.xaxis.set_major_formatter(plt.FuncFormatter(format_func))\n",
    "        elif \"VRef\" in title:\n",
    "            #ax.set_xscale('log')\n",
    "            #ax.set_yscale('log')\n",
    "            #ax1.set_xscale('log')\n",
    "            #ax1.set_yscale('log')\n",
    "            ax.set_xticks([4e-6, 7e-6, 10e-6])\n",
    "            ax.set_yticks([4e-6, 5e-6, 6e-6, 7e-6, 8e-6, 9e-6, 10e-6])\n",
    "            ax1.set_xticks([4e-6, 7e-6, 10e-6])\n",
    "            ax1.set_yticks([4e-6, 5e-6, 6e-6, 7e-6, 8e-6, 9e-6, 10e-6])\n",
    "            ax.yaxis.set_major_formatter(plt.FuncFormatter(format_func))\n",
    "            ax.xaxis.set_major_formatter(plt.FuncFormatter(format_func))\n",
    "            ax1.yaxis.set_major_formatter(plt.FuncFormatter(format_func))\n",
    "            ax1.xaxis.set_major_formatter(plt.FuncFormatter(format_func))\n",
    "            \n",
    "        elif \"Enrichment_cr\" in title:\n",
    "            ax.set_xticks([1,10,20,30,40,50])\n",
    "            ax.set_yticks([1,10,20,30,40,50])\n",
    "            ax1.set_xticks([1,10,20,30,40,50])\n",
    "            ax1.set_yticks([1,10,20,30,40,50])\n",
    "        elif \"iniTempTop\" in title:\n",
    "            ax.set_xticks([1600,1650,1700,1750,1800])\n",
    "            ax.set_yticks([1600,1650,1700,1750,1800])\n",
    "            ax1.set_xticks([1600,1650,1700,1750,1800])\n",
    "            ax1.set_yticks([1600,1650,1700,1750,1800])\n",
    "\n",
    "        for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] +\n",
    "         ax.get_xticklabels() + ax.get_yticklabels()):\n",
    "            item.set_fontsize(20)\n",
    "            \n",
    "        for item in ([ax1.title, ax1.xaxis.label, ax1.yaxis.label] +\n",
    "         ax1.get_xticklabels() + ax1.get_yticklabels()):\n",
    "            item.set_fontsize(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MDN:\n",
    "    def __init__(self, data, x_data,y_data, x_test,y_test, x_cv,y_cv, hSize, KMIX, NEPOCH, learnRate, \\\n",
    "                 paraVec, xparaVec, Dict_baseline_train, Dict_baseline_test, picID, repeats, id_string, \\\n",
    "                 kernel='gaussian',activation='tanh', multivar=False, multivarString=\"p_ra\", trainORload='train'):\n",
    "        print(\"Hidden Layers: \" + str(hSize))\n",
    "        print(\"Number of Mixtures: \" + str(KMIX))\n",
    "        print('Kernel is ' + str(kernel))\n",
    "        print('Activation is ' + str(activation))\n",
    "        \n",
    "        print(\"X_train shape: \" + str(x_data.shape))\n",
    "        print(\"Y_train shape: \" + str(y_data.shape))\n",
    "        print(\"X_test shape: \" + str(x_test.shape))\n",
    "        print(\"Y_test shape: \" + str(y_test.shape))\n",
    "        print(\"X_cv shape: \" + str(x_cv.shape))\n",
    "        print(\"Y_cv shape: \" + str(y_cv.shape))\n",
    "        \n",
    "        if multivar:\n",
    "            print(multivarString)\n",
    "        print()\n",
    "        \n",
    "        self.pMin = data.pMin\n",
    "        self.pMax = data.pMax\n",
    "        tf.reset_default_graph()\n",
    "        yNumParameters = (y_data.shape)[1]\n",
    "        numParameters = np.size(xparaVec)\n",
    "        STDEV = 0.2\n",
    "        MEAN = 0.\n",
    "        xSize = (x_data.shape)[1] #size of each input\n",
    "        \n",
    "        # set up parameters\n",
    "        W = []\n",
    "        b = []\n",
    "        layer = []\n",
    "\n",
    "        x = tf.placeholder(dtype=tf.float64, shape=[None,xSize], name=\"x\")\n",
    "        y = tf.placeholder(dtype=tf.float64, shape=[None,yNumParameters], name=\"y\")\n",
    "        # first layer\n",
    "        W.append(tf.Variable(tf.random_normal([xSize, hSize[0]], mean=MEAN, stddev=STDEV, dtype=tf.float64)))\n",
    "        b.append(tf.Variable(tf.random_normal([1,hSize[0]], mean=MEAN, stddev=STDEV, dtype=tf.float64)))\n",
    "        # add hidden layers (variable number)\n",
    "        for i in range(1,len(hSize)):\n",
    "            W.append(tf.Variable(tf.random_normal([hSize[i-1], hSize[i]], mean=MEAN, stddev=STDEV, dtype=tf.float64)))\n",
    "            b.append(tf.Variable(tf.random_normal([1,hSize[i]], mean=MEAN, stddev=STDEV, dtype=tf.float64)))\n",
    "            \n",
    "        def pdf(x):\n",
    "            return 1./np.sqrt(2.*np.pi) * tf.exp(-x**2/2.)\n",
    "        #def cdf(x):\n",
    "        #    return (1. + tf.erf(x/np.sqrt(2))) / 2.\n",
    "        def skew(x,e,w,a):\n",
    "            t = (x-e) / w\n",
    "            return 2. / w * pdf(t) * cdf(a*t)\n",
    "        def trapz(y, x):\n",
    "            d = tf.subtract(x[1:],x[0:-1])\n",
    "            return tf.reduce_sum((y[0:-1] + y[1:]) * d / 2.)\n",
    "        def inner_function_gmc(out_piI):\n",
    "            max_piI = tf.reduce_max(out_piI, 1, keepdims=True)\n",
    "            out_piI = tf.subtract(out_piI, max_piI)\n",
    "            out_piI = tf.exp(out_piI)\n",
    "            normalize_piI = tf.reciprocal(tf.reduce_sum(out_piI, 1, keepdims=True))\n",
    "            out_piI = tf.multiply(normalize_piI, out_piI)\n",
    "            return out_piI\n",
    "        def getOutput(W,b,layer,hSize):\n",
    "            if activation=='relu':\n",
    "                layer.append(tf.nn.relu(tf.matmul(x, W[0]) + b[0]))\n",
    "                for i in range(1,len(hSize)):\n",
    "                    layer.append(tf.nn.relu(tf.matmul(layer[i-1], W[i]) + b[i]))\n",
    "            if activation=='leaky_relu':\n",
    "                layer.append(tf.nn.leaky_relu(tf.matmul(x, W[0]) + b[0]))\n",
    "                for i in range(1,len(hSize)):\n",
    "                    layer.append(tf.nn.leaky_relu(tf.matmul(layer[i-1], W[i]) + b[i]))\n",
    "            if activation=='elu':\n",
    "                layer.append(tf.nn.elu(tf.matmul(x, W[0]) + b[0]))\n",
    "                for i in range(1,len(hSize)):\n",
    "                    layer.append(tf.nn.elu(tf.matmul(layer[i-1], W[i]) + b[i]))\n",
    "            if activation=='tanh':\n",
    "                layer.append(tf.nn.tanh(tf.matmul(x, W[0]) + b[0]))\n",
    "                for i in range(1,len(hSize)):\n",
    "                    layer.append(tf.nn.tanh(tf.matmul(layer[i-1], W[i]) + b[i]))\n",
    "            return tf.matmul(layer[-1],W[-1]) + b[-1]  \n",
    "        def get_mixture_coef3(output):\n",
    "            out_pi = tf.placeholder(dtype=tf.float64, shape=[None,KMIX* yNumParameters], name=\"mixparam\")\n",
    "            out_sigma = tf.placeholder(dtype=tf.float64, shape=[None,KMIX* yNumParameters], name=\"mixparam\")\n",
    "            out_mu = tf.placeholder(dtype=tf.float64, shape=[None,KMIX* yNumParameters], name=\"mixparam\")\n",
    "            out_pi, out_sigma, out_mu = tf.split(output,3,1)\n",
    "            out_piI = inner_function_gmc((tf.split(out_pi, yNumParameters, 1))[0])\n",
    "            for ind in range(1,yNumParameters):\n",
    "                out_piI = tf.concat([out_piI, inner_function_gmc((tf.split(out_pi, yNumParameters, 1))[ind])], 1)\n",
    "            out_pi = out_piI\n",
    "            return out_pi, out_sigma, out_mu\n",
    "        \n",
    "        def get_mixture_coef3B(output):\n",
    "            out_pi = tf.placeholder(dtype=tf.float64, shape=[None,KMIX* yNumParameters], name=\"mixparam\")\n",
    "            out_alpha = tf.placeholder(dtype=tf.float64, shape=[None,KMIX* yNumParameters], name=\"mixparam\")\n",
    "            out_beta = tf.placeholder(dtype=tf.float64, shape=[None,KMIX* yNumParameters], name=\"mixparam\")\n",
    "            out_pi, out_alpha, out_beta = tf.split(output,3,1)\n",
    "            out_piI = inner_function_gmc((tf.split(out_pi, yNumParameters, 1))[0])\n",
    "            for ind in range(1,yNumParameters):\n",
    "                out_piI = tf.concat([out_piI, inner_function_gmc((tf.split(out_pi, yNumParameters, 1))[ind])], 1)\n",
    "            out_pi = out_piI\n",
    "            return out_pi, out_alpha, out_beta\n",
    "        \n",
    "        if kernel == 'skewed_gaussian':\n",
    "            ySize = KMIX * 4 # pi, mu, stdev, skew\n",
    "            # add final layer\n",
    "            W.append(tf.Variable(tf.random_normal([hSize[-1], ySize * yNumParameters], mean=MEAN, stddev=STDEV, dtype=tf.float64)))\n",
    "            b.append(tf.Variable(tf.random_normal([1,ySize * yNumParameters], mean=MEAN, stddev=STDEV, dtype=tf.float64)))\n",
    "            # define model\n",
    "            output = getOutput(W,b,layer,hSize)\n",
    "\n",
    "            def get_mixture_coef(output):\n",
    "                out_pi = tf.placeholder(dtype=tf.float64, shape=[None,KMIX* yNumParameters], name=\"mixparam\")\n",
    "                out_sigma = tf.placeholder(dtype=tf.float64, shape=[None,KMIX* yNumParameters], name=\"mixparam\")\n",
    "                out_mu = tf.placeholder(dtype=tf.float64, shape=[None,KMIX* yNumParameters], name=\"mixparam\")\n",
    "                out_skew = tf.placeholder(dtype=tf.float64, shape=[None,KMIX* yNumParameters], name=\"mixparam\")\n",
    "                out_pi, out_sigma, out_mu, out_skew = tf.split(output,4,1)\n",
    "                out_piI = inner_function_gmc((tf.split(out_pi, yNumParameters, 1))[0])\n",
    "                for ind in range(1,yNumParameters):\n",
    "                    out_piI = tf.concat([out_piI, inner_function_gmc((tf.split(out_pi, yNumParameters, 1))[ind])], 1)\n",
    "                out_pi = out_piI\n",
    "                out_sigma = tf.exp(out_sigma)\n",
    "                out_mu = out_mu\n",
    "                return out_pi, out_sigma, out_mu, out_skew\n",
    "\n",
    "            out_pi, out_sigma, out_mu, out_skew = get_mixture_coef(output)\n",
    "            oneDivSqrtTwoPI = 1. / np.sqrt(2*np.pi) # normalisation factor for gaussian\n",
    "   \n",
    "            def tf_skew(y, mu, sigma, a):\n",
    "                result = skew(y,mu,sigma,a)\n",
    "                return result \n",
    "            def get_lossfunc(out_pi, out_sigma, out_mu, out_skew, y):\n",
    "                yI = (tf.split(y,yNumParameters,1))[0]\n",
    "                out_muI = (tf.split(out_mu,yNumParameters,1))[0]\n",
    "                out_sigmaI = (tf.split(out_sigma,yNumParameters,1))[0]\n",
    "                out_skewI = (tf.split(out_skew,yNumParameters,1))[0]\n",
    "                resultI = tf_skew(yI, out_muI, out_sigmaI, out_skewI)\n",
    "                for ind in range(1,yNumParameters):\n",
    "                    yI = (tf.split(y,yNumParameters,1))[ind]\n",
    "                    out_muI = (tf.split(out_mu,yNumParameters,1))[ind]\n",
    "                    out_sigmaI = (tf.split(out_sigma,yNumParameters,1))[ind]\n",
    "                    out_skewI = (tf.split(out_skew,yNumParameters,1))[ind]\n",
    "                    resultII = tf_skew(yI, out_muI, out_sigmaI, out_skewI)\n",
    "                    resultI = tf.concat([resultI, resultII], 1)\n",
    "                result = resultI\n",
    "                result = tf.multiply(result, out_pi)\n",
    "                result = tf.reduce_sum(result, 1, keepdims=True)\n",
    "                result = -tf.log(result)\n",
    "                return tf.reduce_mean(result)\n",
    "\n",
    "            lossfunc = get_lossfunc(out_pi, out_sigma, out_mu, out_skew, y)\n",
    "            train_op = tf.train.AdamOptimizer(learning_rate=learnRate).minimize(lossfunc)\n",
    "        \n",
    "        elif kernel=='gaussian':\n",
    "            ySize = KMIX * 3 # pi, mu, stdev\n",
    "            # add final layer\n",
    "            W.append(tf.Variable(tf.random_normal([hSize[-1], ySize * yNumParameters], stddev=STDEV, dtype=tf.float64)))\n",
    "            b.append(tf.Variable(tf.random_normal([1,ySize * yNumParameters], stddev=STDEV, dtype=tf.float64)))\n",
    "            # define model\n",
    "            output = getOutput(W,b,layer,hSize)   \n",
    "           \n",
    "            out_pi, out_sigma, out_mu = get_mixture_coef3(output)\n",
    "            out_sigma = tf.exp(out_sigma)\n",
    "            oneDivSqrtTwoPI = 1. / np.sqrt(2*np.pi) # normalisation factor for gaussian\n",
    "\n",
    "            def tf_normal(y, mu, sigma):\n",
    "                result = tf.subtract(y, mu)\n",
    "                result = tf.multiply(result,tf.reciprocal(sigma))\n",
    "                result = -tf.square(result)/2.\n",
    "                result = tf.multiply(tf.multiply(tf.exp(result),tf.reciprocal(sigma)),oneDivSqrtTwoPI)\n",
    "                return result \n",
    "            def get_lossfunc(out_pi, out_sigma, out_mu, y):\n",
    "                yI = (tf.split(y,yNumParameters,1))[0]\n",
    "                out_muI = (tf.split(out_mu,yNumParameters,1))[0]\n",
    "                out_sigmaI = (tf.split(out_sigma,yNumParameters,1))[0]\n",
    "                resultI = tf_normal(yI, out_muI, out_sigmaI)\n",
    "                for ind in range(1,yNumParameters):\n",
    "                    yI = (tf.split(y,yNumParameters,1))[ind]\n",
    "                    out_muI = (tf.split(out_mu,yNumParameters,1))[ind]\n",
    "                    out_sigmaI = (tf.split(out_sigma,yNumParameters,1))[ind]\n",
    "                    resultII = tf_normal(yI, out_muI, out_sigmaI)\n",
    "                    resultI = tf.concat([resultI, resultII], 1)\n",
    "                result = resultI\n",
    "                result = tf.multiply(result, out_pi)\n",
    "                result = tf.reduce_sum(result, 1, keepdims=True)\n",
    "                result = -tf.log(result)\n",
    "                return tf.reduce_mean(result)\n",
    "            lossfunc = get_lossfunc(out_pi, out_sigma, out_mu, y) + 0.01*tf.add_n([tf.nn.l2_loss(w) for w in W])\n",
    "            train_op = tf.train.AdamOptimizer(learning_rate=learnRate).minimize(lossfunc)\n",
    "            \n",
    "        elif kernel=='beta':\n",
    "            ySize = KMIX * 3 # pi, alpha, beta\n",
    "            # add final layer\n",
    "            W.append(tf.Variable(tf.random_normal([hSize[-1], ySize * yNumParameters], mean=MEAN, \\\n",
    "                                                  stddev=STDEV, dtype=tf.float64)))\n",
    "            b.append(tf.Variable(tf.random_normal([1,ySize * yNumParameters], mean=MEAN, \\\n",
    "                                                  stddev=STDEV, dtype=tf.float64)))\n",
    "            # define model\n",
    "            output = getOutput(W,b,layer,hSize)   \n",
    "            out_pi, out_alpha, out_beta = get_mixture_coef3B(output)\n",
    "            def getGamma(z):\n",
    "                return tf.exp(tf.lgamma(tf.subtract(z,1)))\n",
    "            def tf_beta(y, alpha, beta):\n",
    "                sumIs = tf.add(alpha,beta)\n",
    "                Z = tf.divide(tf.multiply(getGamma(alpha),getGamma(beta)), getGamma(sumIs))\n",
    "                result = tf.divide(tf.multiply(tf.pow(y,(alpha - 1.)), tf.pow((1. - y),(beta - 1.))),Z)\n",
    "                return  result\n",
    "            def get_lossfunc(out_pi, out_alpha, out_beta, y):\n",
    "                yI = (tf.split(y,yNumParameters,1))[0]\n",
    "                out_alphaI = (tf.split(out_alpha,yNumParameters,1))[0]\n",
    "                out_betaI = (tf.split(out_beta,yNumParameters,1))[0]\n",
    "                resultI = tf_beta(yI, out_alphaI, out_betaI)\n",
    "                for ind in range(1,yNumParameters):\n",
    "                    yI = (tf.split(y,yNumParameters,1))[ind]\n",
    "                    out_alphaI = (tf.split(out_alpha,yNumParameters,1))[ind]\n",
    "                    out_betaI = (tf.split(out_beta,yNumParameters,1))[ind]\n",
    "                    resultII = tf_beta(yI, out_alphaI, out_betaI)\n",
    "                    resultI = tf.concat([resultI, resultII], 1)\n",
    "                result = resultI\n",
    "                result = tf.multiply(result, out_pi)\n",
    "                result = tf.abs(tf.reduce_sum(result, 1, keepdims=True))\n",
    "                result = -tf.log(result)\n",
    "                result = tf.reduce_mean(result)\n",
    "                return result \n",
    "            lossfunc = get_lossfunc(out_pi, out_alpha, out_beta, y)\n",
    "            train_op = tf.train.AdamOptimizer(learning_rate=learnRate).minimize(lossfunc)\n",
    "            \n",
    "        lossCheck = np.nan\n",
    "        trials = 0        \n",
    "        lossList = []\n",
    "        loss_cvList = []\n",
    "        trackLoss = True\n",
    "        \n",
    "        sess = tf.Session(config=tf.ConfigProto(log_device_placement=False))\n",
    "        saver = tf.train.Saver()\n",
    "        if trainORload=='train':\n",
    "            while (np.isnan(lossCheck) and trials < 20) or trials < repeats:\n",
    "                #sess = tf.InteractiveSession()\n",
    "                sess.run(tf.global_variables_initializer())\n",
    "                #print(sess.run(output, feed_dict={x: x_data, y: y_data}))\n",
    "\n",
    "                loss = [] #np.zeros(NEPOCH) # store the training progress here.\n",
    "                loss_cv = []\n",
    "                loss_testList = []\n",
    "                sess.run(train_op,feed_dict={x: x_data, y: y_data})\n",
    "                loss.append(sess.run(lossfunc, feed_dict={x: x_data, y: y_data}))\n",
    "                loss_cv.append(sess.run(lossfunc, feed_dict={x: x_cv, y: y_cv}))\n",
    "\n",
    "                sess.run(train_op,feed_dict={x: x_data, y: y_data})\n",
    "                loss.append(sess.run(lossfunc, feed_dict={x: x_data, y: y_data}))\n",
    "                loss_cv.append(sess.run(lossfunc, feed_dict={x: x_cv, y: y_cv}))\n",
    "                i = 1  \n",
    "                while loss_cv[i] <= loss_cv[i-1]:\n",
    "                    sess.run(train_op,feed_dict={x: x_data, y: y_data})\n",
    "                    loss.append(sess.run(lossfunc, feed_dict={x: x_data, y: y_data}))\n",
    "                    loss_cv.append(sess.run(lossfunc, feed_dict={x: x_cv, y: y_cv}))\n",
    "                    i += 1\n",
    "                    if trackLoss and i%10000==0:\n",
    "                        print(i,loss[-1],loss_cv[-1])\n",
    "\n",
    "                while loss_cv[i] <= loss_cv[i-int(np.floor(NEPOCH*i/100))]:\n",
    "                    sess.run(train_op,feed_dict={x: x_data, y: y_data})\n",
    "                    loss.append(sess.run(lossfunc, feed_dict={x: x_data, y: y_data}))\n",
    "                    loss_cv.append(sess.run(lossfunc, feed_dict={x: x_cv, y: y_cv}))\n",
    "                    i += 1\n",
    "                    if trackLoss and i%10000==0:\n",
    "                        print(i,loss[-1],loss_cv[-1])\n",
    "\n",
    "                lossList.append(loss[-1])\n",
    "                loss_cvList.append(loss_cv[-1])\n",
    "                lossCheck = loss[-1]\n",
    "                trials += 1\n",
    "\n",
    "                loss_testList.append(sess.run(lossfunc, feed_dict={x: x_test, y: y_test}))\n",
    "\n",
    "                print(\"Training loss dropped to \" + str(lossCheck) + \"in trial \" + str(trials))\n",
    "                print(\"CV loss dropped to \" + str(loss_cv[-1]) + \"in trial \" + str(trials))\n",
    "            \n",
    "                #saver.save(sess, pathSave + \"TrainedNetworks/\" + id_string + \"/MDN_trial_\" + str(trials))\n",
    "            print('--------------------------------------------------------------------------')\n",
    "            lossTrainStore_mean = np.mean(np.asarray(lossList))\n",
    "            lossTrainStore_std = np.std(np.asarray(lossList))\n",
    "            lossTestStore_mean = np.mean(np.asarray(loss_testList))\n",
    "            lossTestStore_std = np.std(np.asarray(loss_testList))\n",
    "            print(\"Train loss is \" + str(lossTrainStore_mean) + ' +/- ' + str(lossTrainStore_std))\n",
    "            print(\"Test loss is \" + str(lossTestStore_mean) + ' +/- ' + str(lossTestStore_std))\n",
    "            print('--------------------------------------------------------------------------')\n",
    "        else:\n",
    "            saver.restore(sess, pathSave + \"TrainedNetworks/\" + id_string + \"/MDN_trial_0\")\n",
    "                \n",
    "        Dict_Train = {}\n",
    "        Dict_Test = {}\n",
    "        Dict_CV = {}\n",
    "\n",
    "        def generateDict(x_I, yNumParameters):\n",
    "            Dict_I = {}\n",
    "            if kernel=='skewed_gaussian':\n",
    "                out_pi_test, out_sigma_test, out_mu_test, out_skew_test = sess.run(get_mixture_coef(output), feed_dict={x: x_I})   \n",
    "                for i in range(yNumParameters):\n",
    "                    Dict_I[\"sigma\" + str(i)] = (np.split(out_sigma_test, yNumParameters, axis=1))[i]\n",
    "                    Dict_I[\"mu\" + str(i)] = (np.split(out_mu_test, yNumParameters, axis=1))[i]\n",
    "                    Dict_I[\"pi\" + str(i)] = (np.split(out_pi_test, yNumParameters, axis=1))[i]\n",
    "                    Dict_I[\"skew\" + str(i)] = (np.split(out_skew_test, yNumParameters, axis=1))[i]\n",
    "            elif kernel=='gaussian':\n",
    "                out_pi_test, out_sigma_test, out_mu_test = sess.run(get_mixture_coef3(output), feed_dict={x: x_I}) \n",
    "                out_sigma_test = np.exp(out_sigma_test)\n",
    "                for i in range(yNumParameters):\n",
    "                    Dict_I[\"sigma\" + str(i)] = (np.split(out_sigma_test, yNumParameters, axis=1))[i]\n",
    "                    Dict_I[\"mu\" + str(i)] = (np.split(out_mu_test, yNumParameters, axis=1))[i]\n",
    "                    Dict_I[\"pi\" + str(i)] = (np.split(out_pi_test, yNumParameters, axis=1))[i]   \n",
    "            elif kernel=='beta':\n",
    "                out_pi_test, out_alpha_test, out_beta_test = sess.run(get_mixture_coef3(output), feed_dict={x: x_I})   \n",
    "                for i in range(yNumParameters):\n",
    "                    Dict_I[\"alpha\" + str(i)] = (np.split(out_alpha_test, yNumParameters, axis=1))[i]\n",
    "                    Dict_I[\"beta\" + str(i)] = (np.split(out_beta_test, yNumParameters, axis=1))[i]\n",
    "                    Dict_I[\"pi\" + str(i)] = (np.split(out_pi_test, yNumParameters, axis=1))[i]   \n",
    "            return Dict_I\n",
    "\n",
    "        alphaVal = 0.5\n",
    "        Dict_Train = generateDict(x_data, yNumParameters)\n",
    "        Dict_Test = generateDict(x_test, yNumParameters)\n",
    "\n",
    "        for ind in range(yNumParameters):\n",
    "            fig = plt.figure(figsize=(14,9), dpi=200)\n",
    "            plt.subplots_adjust(top=0.90, bottom=0.08, left=0.10, right=0.95, hspace=0.3,wspace=0.4)\n",
    "            nCols = 3 # 2\n",
    "            nRows = 2 #yNumParameters//nCols + 2\n",
    "            plotCounter = 1        \n",
    "            Dict_Test[\"forFig_\" + paraVec[ind] + \"_loss\"] = loss\n",
    "            Dict_Test[\"forFig_\" + paraVec[ind] + \"_loss_cv\"] = loss_cv\n",
    "            ax = fig.add_subplot(nRows,nCols,plotCounter) \n",
    "            ax.plot(loss, 'r-', label='Training')\n",
    "            ax.plot(loss_cv,\"b--\", label='Validation')\n",
    "            ax.set_xlabel('Epochs')\n",
    "            ax.set_ylabel('Negative log-likelihood')\n",
    "            #plt.title('Loss')\n",
    "            ax.legend()\n",
    "            #ax.set_xticks([0, 50000, 100000])\n",
    "            ax.locator_params(axis='x', nbins=3)\n",
    "            plotCounter += 1\n",
    "            for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] +\n",
    "            ax.get_xticklabels() + ax.get_yticklabels()):\n",
    "                item.set_fontsize(20)\n",
    "            ax = fig.add_subplot(nRows,nCols,plotCounter)\n",
    "            plotCounter += 1\n",
    "            ax1 = fig.add_subplot(nRows,nCols,plotCounter)\n",
    "            plotCounter += 2\n",
    "\n",
    "            pR = plotResults(data,ax,ax1, alphaVal,(paraVec[ind] + \", Train\"), y_data[:,ind], Dict_Train, \\\n",
    "                             ind, Dict_baseline_train, numParameters, kernel, False)\n",
    "            Dict_Train[\"loss_train_mean\"] = lossTrainStore_mean\n",
    "            Dict_Train[\"loss_train_std\"] = lossTrainStore_std\n",
    "            ax = fig.add_subplot(nRows,nCols,plotCounter)\n",
    "            plotCounter += 1\n",
    "            ax1 = fig.add_subplot(nRows,nCols,plotCounter)\n",
    "            plotCounter += 1\n",
    "            pRT = plotResults(data,ax,ax1, alphaVal,(paraVec[ind] + \", Test\"), y_test[:,ind], Dict_Test, \\\n",
    "                        ind, Dict_baseline_test, numParameters, kernel, True)\n",
    "            Dict_Test = pRT.Dict_Test\n",
    "            Dict_Test[\"loss_test_mean\"] = lossTestStore_mean\n",
    "            Dict_Test[\"loss_test_std\"] = lossTestStore_std\n",
    "            #fig.savefig(pathSave + \"inverseProblemPics/\" + \"Output_fig\" + str(picID) + \".pdf\", bbox_inches='tight')\n",
    "\n",
    "            fig.tight_layout()\n",
    "        plt.show()\n",
    "        print(\"-------------------------------------------------------------------------------------------------------\")\n",
    "        print(\"-------------------------------------------------------------------------------------------------------\")\n",
    "        print(\"-------------------------------------------------------------------------------------------------------\")\n",
    "        print()\n",
    "        sess.close()\n",
    "\n",
    "        self.Dict_Train = Dict_Train\n",
    "        self.Dict_Test = Dict_Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "class meanPredictor:\n",
    "    def __init__(self, x_data,y_data,xparaVec, plotMP=False):\n",
    "        Dict_Mean_Predictor = {}\n",
    "        if 'Tprof_4p5Gyr' in xparaVec and np.size(xparaVec)==1:\n",
    "            tProfOnly = True\n",
    "        else:\n",
    "            tProfOnly = False\n",
    "        for i in range(np.size(xparaVec)):\n",
    "            if tProfOnly or (xparaVec[i] == 'Tprof_4p5Gyr'):\n",
    "                observable = np.mean(x_data,axis=1)\n",
    "            else:\n",
    "                observable = x_data[:,i] \n",
    "            para = y_data[:,0]\n",
    "            steps = 8\n",
    "            \n",
    "            bin_means, bin_edges, binnumber = binned_statistic(observable, para, statistic='mean',bins=steps)\n",
    "            bin_variance = [np.std(para[np.where(binnumber == bIndex)[0]]) for bIndex in range(1,steps+1)]\n",
    "            removeVec = []\n",
    "            notRemoveVec = []\n",
    "            for i3 in range(np.size(bin_means)):\n",
    "                if np.isnan(bin_means[i3]) or bin_variance[i3] == 0.0:\n",
    "                    removeVec.append(i3)\n",
    "                else:\n",
    "                    notRemoveVec.append(i3)\n",
    "\n",
    "            for r in removeVec:\n",
    "                bin_variance[r] = np.average([bin_variance[notInd] for notInd in notRemoveVec])\n",
    "                bin_means[r] = np.average([bin_means[notInd] for notInd in notRemoveVec])\n",
    "\n",
    "            digitizer = np.digitize(observable,bin_edges) - 1\n",
    "            mu = np.zeros(np.size(digitizer))\n",
    "            variance = np.zeros(np.size(digitizer))\n",
    "            for ind in range(np.size(digitizer)):\n",
    "                mu[ind] = bin_means[digitizer[ind]-1]\n",
    "                variance[ind] = bin_variance[digitizer[ind]-1]\n",
    "            Dict_Mean_Predictor[\"mu_var\" + str(i)] = mu\n",
    "            Dict_Mean_Predictor[\"variance_var\" + str(i)] = variance\n",
    "            if plotMP:\n",
    "                plt.figure(figsize=(4,6))\n",
    "                plt.plot(observable,para, \"ro\", alpha = 0.4)\n",
    "                plt.plot((bin_edges[1:]+bin_edges[0:-1])/2.,bin_means, \"k-\")\n",
    "                plt.plot((bin_edges[1:]+bin_edges[0:-1])/2.,bin_means+bin_variance,\"b--\")\n",
    "                plt.plot((bin_edges[1:]+bin_edges[0:-1])/2.,bin_means-bin_variance,\"b--\")\n",
    "                plt.legend([_, \"$\\mu$\", \"$\\mu + \\sigma$\", \"$\\mu - \\sigma$\"], loc=\"lower_right\")\n",
    "        self.Dict_Mean_Predictor = Dict_Mean_Predictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bigLoop(_zipped, repeats=5, multivar=False):\n",
    "    Dict_kls = {}\n",
    "    hSize, KMIX, x_o, x_b, trainPercent, yIndex, kernel, activation, noiseLevel = _zipped\n",
    "    fontSize = 14\n",
    "    picID = 0 #yIndices  [yIndices[i5]]\n",
    "    data = getData([yIndex], x_o, x_b, trainPercent, pathSave, noiseLevel)\n",
    "    \n",
    "    id_string = str(_zipped)\n",
    "    with open(pathSave + \"/Data_Files/processedData\" + id_string + \".txt\", \"rb\") as fkl:\n",
    "        dataDict = load(fkl)\n",
    "    \n",
    "    data.x_data = dataDict['x_data']\n",
    "    data.x_test = dataDict['x_test']\n",
    "    data.x_cv = dataDict['x_cv']     \n",
    "    data.y_data = dataDict['y_data']\n",
    "    data.y_test = dataDict['y_test']\n",
    "    data.y_cv = dataDict['y_cv'] \n",
    "    data.paraVec = dataDict['paraVec']\n",
    "    data.xparaVec = dataDict['xparaVec']\n",
    "    data.rProf = dataDict['rProf']\n",
    "    data.startPoint = dataDict['startPoint']\n",
    "    data.endPoint = dataDict['endPoint']    \n",
    "    data.indexer = dataDict['indexer']    \n",
    "    data.pMax = dataDict['pMax']  \n",
    "    data.pMin = dataDict['pMin']  \n",
    "    data.oMax = dataDict['oMax']\n",
    "    data.oMin = dataDict['oMin']\n",
    "    \n",
    "    trainORload = 'train'\n",
    "    \n",
    "    Dict_MP_Test = meanPredictor(data.x_test, data.y_test, data.xparaVec, False)\n",
    "    Dict_MP_Train = meanPredictor(data.x_data, data.y_data, data.xparaVec, False)\n",
    "\n",
    "    #data.visualizeData(data, picID, fontSize, Dict_MP_Train.Dict_Mean_Predictor)\n",
    "\n",
    "    learn_rate = 0.1\n",
    "    Dict_MDN = MDN(data,data.x_data,data.y_data, data.x_test,data.y_test, data.x_cv, \\\n",
    "                       data.y_cv, hSize, KMIX, 1, learn_rate, \\\n",
    "                       data.paraVec, data.xparaVec, Dict_MP_Train.Dict_Mean_Predictor, \\\n",
    "                       Dict_MP_Test.Dict_Mean_Predictor, picID, repeats, id_string, kernel, activation, False, \n",
    "                       \"p_ra\", trainORload)\n",
    "\n",
    "    Dict_kls[id_string + \"Loss_train_mean\"] = Dict_MDN.Dict_Train[\"loss_train_mean\"]\n",
    "    Dict_kls[id_string + \"Loss_train_std\"] = Dict_MDN.Dict_Train[\"loss_train_std\"]\n",
    "    Dict_kls[id_string + \"Loss_test_mean\"] = Dict_MDN.Dict_Test[\"loss_test_mean\"]\n",
    "    Dict_kls[id_string + \"Loss_test_std\"] = Dict_MDN.Dict_Test[\"loss_test_std\"]\n",
    "\n",
    "    picID += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#Parameter indices: 0-Rayleigh number, 4-activation energy, 5-activation volume, 6-enrichment factor, 7-initial temperature\n",
    "#Obsbervables indices: 0-Q_c, 1-Q_s, 2-R_th, 3-D_e, 4-D_melt, 5-t_volc, 6-T_profs\n",
    "\n",
    "# hidden layers, mixtures, observables, bounds for temperature profile, size of training dataset, parameter, kernel, activation, Noise level\n",
    "zipped_new = [[[12, 6], 3, [0, 1, 2, 3, 4, 5, 6], [0, 100], 0.8, 0, 'gaussian', 'tanh', 0.0],\n",
    "              [[12, 6], 3, [0, 1, 2, 3, 4, 5, 6], [0, 100], 0.8, 4, 'gaussian', 'tanh', 0.0],\n",
    "              [[12, 6], 3, [0, 1, 2, 3, 4, 5, 6], [0, 100], 0.8, 5, 'gaussian', 'tanh', 0.0],\n",
    "              [[12, 6], 3, [0, 1, 2, 3, 4, 5, 6], [0, 100], 0.8, 6, 'gaussian', 'tanh', 0.0],\n",
    "              [[12, 6], 3, [0, 1, 2, 3, 4, 5, 6], [0, 100], 0.8, 7, 'gaussian', 'tanh', 0.0]\n",
    "             ]\n",
    "runParallel = False\n",
    "if runParallel:                   \n",
    "    Parallel(n_jobs=1, verbose=10, backend='loky', prefer='processes')(delayed(bigLoop)(_z, 1, True) for _z in zipped_new)\n",
    "    \n",
    "else:\n",
    "    [bigLoop(_z) for _z in zipped_new] "
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}