{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LLEPE: Liquid-Liquid Equilibrium Parameter Estimator\n",
    "\n",
    "Copyright © 2020, UChicago Argonne, LLC. All rights reserved.\n",
    "\n",
    "Released under the modified BSD license. See LICENSE for more details."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This script performs the Pitzer fit method and shows the parity plots for each of the species. This code requires that the elementz.xml file in the data folder has been copied into Cantera's data folder. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting Pitzer and linear parameters to single component data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llepe import LLEPE, lmse_perturbed_obj, ind_lmse_perturbed_obj\n",
    "import numpy as np\n",
    "import itertools\n",
    "import pandas as pd\n",
    "import matplotlib\n",
    "from scipy.optimize import curve_fit\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import json\n",
    "\n",
    "font = {'family': 'sans serif',\n",
    "        'size': 24}\n",
    "matplotlib.rc('font', **font)\n",
    "plt.rc('xtick', labelsize=18)\n",
    "plt.rc('ytick', labelsize=18)\n",
    "plt.rcParams['lines.linewidth'] = 4\n",
    "matplotlib.rcParams['lines.markersize'] = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Nd\n",
      "Pr\n",
      "Ce\n",
      "La\n",
      "Dy\n",
      "Sm\n",
      "Y\n"
     ]
    }
   ],
   "source": [
    "def linear(x, a, b):\n",
    "    return a * x + b\n",
    "lp_dict = {'species': [],\n",
    "           'slope': [],\n",
    "           'intercept': []}\n",
    "species_list = 'Nd,Pr,Ce,La,Dy,Sm,Y'.split(',')\n",
    "pitzer_params = pd.read_csv(\"../../data/csvs/may_pitzer_params.csv\")\n",
    "x_guesses = [[-5178500.0, -1459500.0],\n",
    "             [-5178342.857142857, -1460300.0],\n",
    "             [-5178342.857142857, -1459500.0],\n",
    "             [-5178342.857142857, -1458300.0],\n",
    "             [-5178185.714285715, -1459900.0],\n",
    "             [-5178185.714285715, -1459500.0],\n",
    "             [-5178185.714285715, -1459100.0],\n",
    "             [-5178185.714285715, -1458300.0],\n",
    "             [-5178028.571428572, -1459900.0],\n",
    "             [-5178028.571428572, -1459100.0],\n",
    "             [-5178028.571428572, -1458300.0],\n",
    "             [-5177557.142857143, -1459900.0],\n",
    "             [-5177400.0, -1460300.0]]\n",
    "optimizer = 'scipy_minimize'\n",
    "opt_dfs = {}\n",
    "pitzer_params_dict = {}\n",
    "for species in species_list:\n",
    "    print(species)\n",
    "    lower_species = species.lower()\n",
    "\n",
    "    labeled_data = pd.read_csv(\"../../data/csvs/labeled_zeroes_removed_{0}.csv\".format(lower_species))\n",
    "    exp_data = labeled_data.drop(labeled_data.columns[0], axis=1)\n",
    "    xml_file = '../../data/xmls/PC88A_HCL_{0}_w_pitzer.xml'.format(species)\n",
    "    opt_values = {\n",
    "        '(HA)2(org)_h0': [],\n",
    "        '{0}(H(A)2)3(org)_h0'.format(species): [],\n",
    "        'beta0': [],\n",
    "        'beta1': [],\n",
    "        'Cphi': [],\n",
    "        'obj_value': [],\n",
    "        'guess': []}\n",
    "    new_x_guesses = []\n",
    "    for x_guess in x_guesses:\n",
    "        opt_dict = {'(HA)2(org)_h0': {'upper_element_name': 'species',\n",
    "                                      'upper_attrib_name': 'name',\n",
    "                                      'upper_attrib_value': '(HA)2(org)',\n",
    "                                      'lower_element_name': 'h0',\n",
    "                                      'lower_attrib_name': None,\n",
    "                                      'lower_attrib_value': None,\n",
    "                                      'input_format': '{0}',\n",
    "                                      'input_value': x_guess[1]},\n",
    "                    '{0}(H(A)2)3(org)_h0'.format(species): {\n",
    "                        'upper_element_name': 'species',\n",
    "                        'upper_attrib_name': 'name',\n",
    "                        'upper_attrib_value': '{0}(H(A)2)3(org)'.format(\n",
    "                            species),\n",
    "                        'lower_element_name': 'h0',\n",
    "                        'lower_attrib_name': None,\n",
    "                        'lower_attrib_value': None,\n",
    "                        'input_format': '{0}',\n",
    "                        'input_value': x_guess[0]},\n",
    "                    }\n",
    "        pitzer_param_list = ['beta0', 'beta1', 'Cphi']\n",
    "        for pitzer_param in pitzer_param_list:\n",
    "            inner_dict = {'upper_element_name': 'binarySaltParameters',\n",
    "                          'upper_attrib_name': 'cation',\n",
    "                          'upper_attrib_value': '{0}+++'.format(species),\n",
    "                          'lower_element_name': pitzer_param,\n",
    "                          'lower_attrib_name': None,\n",
    "                          'lower_attrib_value': None}\n",
    "            pitzer_row = pitzer_params[\n",
    "                pitzer_params['species'] == species]\n",
    "\n",
    "            inner_dict2 = {'input_format': ' {0}, 0.0, 0.0, 0.0, 0.0 ',\n",
    "                           'input_value': pitzer_row[pitzer_param].values[\n",
    "                               0]}\n",
    "            inner_dict.update(inner_dict2)\n",
    "            opt_dict['{0}_{1}'.format(species, pitzer_param)] = inner_dict\n",
    "\n",
    "        reeps_params = {\n",
    "            'exp_data': exp_data,\n",
    "            'phases_xml_filename': xml_file,\n",
    "            'opt_dict': opt_dict,\n",
    "            'phase_names': ['HCl_electrolyte', 'PC88A_liquid'],\n",
    "            'aq_solvent_name': 'H2O(L)',\n",
    "            'extractant_name': '(HA)2(org)',\n",
    "            'diluant_name': 'dodecane',\n",
    "            'complex_names': ['{0}(H(A)2)3(org)'.format(species)],\n",
    "            'extracted_species_ion_names': ['{0}+++'.format(species)],\n",
    "            'aq_solvent_rho': 1000.0,\n",
    "            'extractant_rho': 960.0,\n",
    "            'diluant_rho': 750.0,\n",
    "            'objective_function': lmse_perturbed_obj,\n",
    "            'optimizer': optimizer\n",
    "        }\n",
    "        estimator = LLEPE(**reeps_params)\n",
    "        estimator.update_xml(reeps_params['opt_dict'])\n",
    "        exp_df = estimator.get_exp_df()\n",
    "        pred_dict = estimator.get_predicted_dict()\n",
    "        init_obj_value = lmse_perturbed_obj(pred_dict, exp_df, [species], 1e-100)\n",
    "        \n",
    "        obj_kwargs = {'species_list': [species], 'epsilon': 1e-100}\n",
    "        bounds = [(1e-1, 1e1)] * len(opt_dict)\n",
    "        optimizer_kwargs = {\"method\": 'l-bfgs-b',\n",
    "                            \"bounds\": bounds}\n",
    "        opt_dict, obj_value = estimator.fit(\n",
    "            objective_kwargs=obj_kwargs,\n",
    "            optimizer_kwargs=optimizer_kwargs)\n",
    "        keys = list(opt_dict.keys())\n",
    "        info1 = [opt_dict[key]['input_value'] for key in keys]\n",
    "        info1.append(obj_value)\n",
    "        info1.append(x_guess)\n",
    "        if obj_value > init_obj_value*0.95:\n",
    "            continue\n",
    "        \n",
    "        opt_values_keys = opt_values.keys()\n",
    "        for ind, key in enumerate(opt_values_keys):\n",
    "            opt_values[key].append(info1[ind])\n",
    "        new_x_guesses.append(x_guess)\n",
    "    x_guesses = new_x_guesses\n",
    "    opt_df = pd.DataFrame(opt_values)\n",
    "    opt_dfs[species] = opt_df\n",
    "                                               \n",
    "    p_opt, p_cov = curve_fit(linear,\n",
    "                             opt_df['(HA)2(org)_h0'].values,\n",
    "                             opt_df['{0}(H(A)2)3(org)_h0'.format(\n",
    "                                 species)].values)\n",
    "    slope, intercept = p_opt\n",
    "    lp_dict['species'].append(species)\n",
    "    lp_dict['slope'].append(slope)\n",
    "    lp_dict['intercept'].append(intercept)\n",
    "    min_h0_row = opt_df[opt_df['(HA)2(org)_h0'] == opt_df['(HA)2(org)_h0'].min()]\n",
    "    for pitzer_param in pitzer_param_list:\n",
    "        value = min_h0_row[pitzer_param].values[0]\n",
    "        inner_dict = {'upper_element_name': 'binarySaltParameters',\n",
    "                      'upper_attrib_name': 'cation',\n",
    "                      'upper_attrib_value': '{0}+++'.format(species),\n",
    "                      'lower_element_name': pitzer_param,\n",
    "                      'lower_attrib_name': None,\n",
    "                      'lower_attrib_value': None,\n",
    "                     'input_format': ' {0}, 0.0, 0.0, 0.0, 0.0 ',\n",
    "                       'input_value': value}\n",
    "        pitzer_params_dict['{0}_{1}'.format(species, pitzer_param)] = inner_dict\n",
    "lp_df = pd.DataFrame(lp_dict)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting extractant enthalpy to combined single and multicomponent dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\users\\tq220\\documents\\tits things\\anl research\\gitlab\\parameter-estimation\\llepe\\llepe.py:988: RuntimeWarning: invalid value encountered in true_divide\n",
      "  d_array = extracted_species_org_array / extracted_species_aq_array\n"
     ]
    }
   ],
   "source": [
    "def ext_to_complex(h0, custom_objects_dict, species):\n",
    "    linear_params = custom_objects_dict['lp_df']\n",
    "    row = linear_params[linear_params['species'] == species]\n",
    "    return row['slope'].values[0] * h0[0] + row['intercept'].values[0]\n",
    "\n",
    "labeled_data = pd.read_csv(\"../../data/csvs/zeroes_removed_PC88A_HCL_NdPrCeLaDySmY.csv\")\n",
    "exp_data = labeled_data.drop(labeled_data.columns[0], axis=1)\n",
    "xml_file = '../../data/xmls/PC88A_HCL_NdPrCeLaDySmY_w_pitzer.xml'\n",
    "complex_names = ['{0}(H(A)2)3(org)'.format(species)\n",
    "                 for species in species_list]\n",
    "dependant_params_dict = {}\n",
    "for species, complex_name in zip(species_list, complex_names):\n",
    "    inner_dict = {'upper_element_name': 'species',\n",
    "                  'upper_attrib_name': 'name',\n",
    "                  'upper_attrib_value': complex_name,\n",
    "                  'lower_element_name': 'h0',\n",
    "                  'lower_attrib_name': None,\n",
    "                  'lower_attrib_value': None,\n",
    "                  'input_format': '{0}',\n",
    "                  'function': ext_to_complex,\n",
    "                  'kwargs': {\"species\": species},\n",
    "                  'independent_params': '(HA)2(org)_h0'}\n",
    "    dependant_params_dict['{0}_h0'.format(complex_name)] = inner_dict\n",
    "extractant_guess = -1.37e6\n",
    "info_dict = {'(HA)2(org)_h0': {'upper_element_name': 'species',\n",
    "                              'upper_attrib_name': 'name',\n",
    "                              'upper_attrib_value': '(HA)2(org)',\n",
    "                              'lower_element_name': 'h0',\n",
    "                              'lower_attrib_name': None,\n",
    "                              'lower_attrib_value': None,\n",
    "                              'input_format': '{0}',\n",
    "                              'input_value': extractant_guess},\n",
    "            }\n",
    "reeps_params = {\n",
    "    'exp_data': exp_data,\n",
    "    'phases_xml_filename': xml_file,\n",
    "    'opt_dict': info_dict,\n",
    "    'phase_names': ['HCl_electrolyte', 'PC88A_liquid'],\n",
    "    'aq_solvent_name': 'H2O(L)',\n",
    "    'extractant_name': '(HA)2(org)',\n",
    "    'diluant_name': 'dodecane',\n",
    "    'complex_names': ['{0}(H(A)2)3(org)'.format(species)\n",
    "                      for species in species_list],\n",
    "    'extracted_species_ion_names': ['{0}+++'.format(species)\n",
    "                                    for species in species_list],\n",
    "    'aq_solvent_rho': 1000.0,\n",
    "    'extractant_rho': 960.0,\n",
    "    'diluant_rho': 750.0,\n",
    "    'objective_function': lmse_perturbed_obj,\n",
    "    'custom_objects_dict': {'lp_df': lp_df}\n",
    "}\n",
    "estimator = LLEPE(**reeps_params)\n",
    "estimator.set_dependant_params_dict(dependant_params_dict)\n",
    "estimator.update_xml(pitzer_params_dict)\n",
    "estimator.update_xml(reeps_params['opt_dict'],\n",
    "                     dependant_params_dict=dependant_params_dict)\n",
    "obj_kwargs = {'species_list': species_list,\n",
    "              'epsilon': 1e-100}\n",
    "optimizer_kwargs = {\"method\": 'l-bfgs-b',\n",
    "                    \"bounds\": [(1e-1, 1e1)] * len(info_dict)}\n",
    "opt_dict, obj_value = estimator.fit(\n",
    "    objective_kwargs=obj_kwargs,\n",
    "    optimizer_kwargs=optimizer_kwargs)\n",
    "estimator.update_xml(opt_dict,\n",
    "                     dependant_params_dict=dependant_params_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting organic phase parity plots for complete unfiltered dataset "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\users\\tq220\\documents\\tits things\\anl research\\gitlab\\parameter-estimation\\llepe\\llepe.py:988: RuntimeWarning: invalid value encountered in true_divide\n",
      "  d_array = extracted_species_org_array / extracted_species_aq_array\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 65260 characters]\n",
      "text/plain": [
       "<Figure size 743.226x572.185 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 58612 characters]\n",
      "text/plain": [
       "<Figure size 743.226x572.185 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 54300 characters]\n",
      "text/plain": [
       "<Figure size 743.226x572.185 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 51652 characters]\n",
      "text/plain": [
       "<Figure size 743.226x572.185 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 57116 characters]\n",
      "text/plain": [
       "<Figure size 743.226x572.185 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 55296 characters]\n",
      "text/plain": [
       "<Figure size 743.226x572.185 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 57936 characters]\n",
      "text/plain": [
       "<Figure size 743.226x572.185 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "feed_cols = []\n",
    "labeled_complete_data = pd.read_csv(\"../../data/csvs/PC88A_HCL_NdPrCeLaDySmY.csv\")\n",
    "complete_exp_df = labeled_data.drop(labeled_data.columns[0], axis=1)\n",
    "estimator.set_exp_df(complete_exp_df)\n",
    "exp_data = estimator.get_exp_df()\n",
    "for col in exp_data.columns:\n",
    "    if 'aq_i' in col:\n",
    "        feed_cols.append(col)\n",
    "exp_data['total_re'] = exp_data[feed_cols].sum(axis=1)\n",
    "for species in species_list:\n",
    "    fig, ax = estimator.parity_plot('{0}_org_eq'.format(species),\n",
    "                                    c_data=\n",
    "                                    exp_data['total_re'].values,\n",
    "                                    c_label='Feed total RE '\n",
    "                                            'molarity (mol/L)',\n",
    "                                    print_r_squared=True)"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
