2. Loading FE-ANN(s) EoS.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "87a5c00c",
   "metadata": {},
   "source": [
    "# 2. Loading the FE-ANN(s) EoS\n",
    "\n",
    "The FE-ANN(s) EoS is implemented using `jax` and `flax`. The trained FE-ANN(s) EoS checkpoints are stored using `orbax`. First, we need to upload the required packages:\n",
    "- **numpy**: numerical python\n",
    "- **jax**: high-performance numerical computing and machine learning\n",
    "- **nest_asyncio**: needed to load the checkpoints\n",
    "- **feanns_eos**: library with implemented FE-ANN(s) EoS and phase equilibria solvers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c20978aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from jax import numpy as jnp\n",
    "\n",
    "from feanns_eos import HelmholtzModel\n",
    "from feanns_eos import load_feanns_params\n",
    "from feanns_eos import helper_get_alpha\n",
    "\n",
    "from jax.config import config\n",
    "config.update(\"jax_enable_x64\", True)\n",
    "type_np = np.float64\n",
    "type_jax = jnp.float64\n",
    "\n",
    "# needed to read checkpoints in jupyter notebooks\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca4c4843",
   "metadata": {},
   "source": [
    "The checkpoints can be loaded with the `load_feanns_params` function. This function requires the folder (`ckpt_dir`) where the checkpoint is located and its name (`prefix`). The loaded state is a Python dictionary with some information about the FE-ANN(s) EoS, as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d59add2b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Information included in the checkpoint: dict_keys(['epoch', 'features', 'learning_rate', 'params', 'seed'])\n",
      "# epochs:  50000\n",
      "features:  [40, 40, 40] ; this equals to 3 layers of 40 neurons\n",
      "seed:  0\n",
      "learning_rate:  1e-06\n"
     ]
    }
   ],
   "source": [
    "# directory where checkpoints are saved\n",
    "ckpt_dir = \"./eos_params\"\n",
    "# name of the checkpoint\n",
    "prefix = \"feanns_\"\n",
    "\n",
    "# loading the current checkpoint\n",
    "state = load_feanns_params(ckpt_dir, prefix=prefix)\n",
    "print('Information included in the checkpoint:', state.keys())\n",
    "\n",
    "print(\"# epochs: \", state['epoch'])\n",
    "print(\"features: \", state['features'], f'; this equals to {len(state[\"features\"])} layers of {state[\"features\"][0]} neurons')\n",
    "print(\"seed: \", state['seed'])\n",
    "print(\"learning_rate: \", state['learning_rate'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad44b610",
   "metadata": {},
   "source": [
    "The `state['params']` value is a dictionary with the kernel and biases of the used in the FE-ANN(s) EoS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9f5b25a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'helmholtz_layer': {'kernel': array([[-1.91865365],\n",
       "         [-3.42008273],\n",
       "         [-2.40948245],\n",
       "         [ 2.17382488],\n",
       "         [ 1.87712497],\n",
       "         [ 1.92623504],\n",
       "         [ 1.98035204],\n",
       "         [-1.99994356],\n",
       "         [ 2.02550444],\n",
       "         [ 0.75905708],\n",
       "         [-1.8433767 ],\n",
       "         [-2.47055048],\n",
       "         [-0.52027889],\n",
       "         [ 1.31926447],\n",
       "         [ 1.3268983 ],\n",
       "         [-1.33231753],\n",
       "         [-2.69135919],\n",
       "         [ 2.77863568],\n",
       "         [ 0.42604366],\n",
       "         [-2.42654057],\n",
       "         [-8.27893431],\n",
       "         [ 1.68126863],\n",
       "         [-1.98010807],\n",
       "         [ 0.83462859],\n",
       "         [ 2.60469435],\n",
       "         [-3.31947377],\n",
       "         [-2.7173411 ],\n",
       "         [ 1.19479068],\n",
       "         [ 1.13801348],\n",
       "         [ 2.43388489],\n",
       "         [-1.27424681],\n",
       "         [-1.80180761],\n",
       "         [ 1.78896363],\n",
       "         [-2.11985688],\n",
       "         [-1.43679681],\n",
       "         [ 1.33957413],\n",
       "         [-1.06817372],\n",
       "         [ 2.47614909],\n",
       "         [ 0.91262327],\n",
       "         [-0.47107328]])},\n",
       " 'hidden_layers_0': {'bias': array([-0.14041495,  0.05864247, -0.26801362,  0.00872826, -0.35548193,\n",
       "         -0.37191372,  0.34049223, -0.60711805,  0.03188073,  0.11482699,\n",
       "          0.29429513, -0.04368035,  0.20598373,  0.2423712 ,  0.3972854 ,\n",
       "         -0.24474726,  0.36761617, -0.41083928,  0.35169503,  0.45307416,\n",
       "          0.08150771,  0.16426346,  0.28789166,  0.06790798, -0.71262576,\n",
       "          0.13066686, -0.03351084,  0.17871126, -0.21907065,  0.34691385,\n",
       "         -0.26556311, -0.37210992, -0.31785243,  0.26197533, -0.05307883,\n",
       "          0.31534835,  0.14199051,  0.27205322,  0.13874754, -0.31981226]),\n",
       "  'kernel': array([[ 1.01475238,  0.02005405,  0.0259894 ,  0.02891581, -0.38581285,\n",
       "          -0.30500901,  0.04373799, -0.25841077,  0.54248902, -0.66356478,\n",
       "          -0.05885583,  0.23904932,  0.26964657,  0.2533398 ,  0.43330858,\n",
       "          -0.51553442,  0.47749262,  0.28740781,  0.20732127, -0.40623666,\n",
       "           0.80889292, -0.14549449, -0.40167946, -0.34375152,  0.27816398,\n",
       "           0.23201532, -0.81590214, -0.44805256, -0.83600251,  0.01068955,\n",
       "           0.12699801, -0.15492695, -0.06607532,  0.24558218,  0.46162747,\n",
       "          -0.31064266, -0.03848031, -0.03607182,  0.30843885,  0.4882486 ],\n",
       "         [ 0.02183318, -0.0901233 ,  0.35296115,  0.0453627 ,  0.13723986,\n",
       "           0.53916466, -0.34191087, -0.20906375, -0.21711248,  0.11853542,\n",
       "          -0.42963561, -0.06328177, -0.42550925, -0.32817991, -0.39112069,\n",
       "           0.36267274, -0.18335842,  0.09216314, -0.53277829, -0.07192099,\n",
       "          -0.17355303,  0.18352806,  0.26976217, -0.47775584,  0.22348567,\n",
       "          -0.26744881,  0.21598909, -0.05599516,  0.43304148, -0.28653973,\n",
       "          -0.08579825,  0.54009798,  0.31148715, -0.32098969,  0.07246888,\n",
       "          -0.14132523, -0.17461726, -0.29358173, -0.33374286,  0.15404032],\n",
       "         [ 0.0520343 ,  0.09780132, -0.39096523, -0.00803742,  0.01497137,\n",
       "          -0.6045606 , -0.21599828, -2.40877504,  0.51393927, -0.36970683,\n",
       "           0.18978028, -2.53875347, -0.10990678, -0.12870743, -0.54460624,\n",
       "           0.37183532, -0.38537326,  0.12371135,  0.56744688, -0.12221545,\n",
       "           1.52860712, -0.23076871,  0.39822327,  0.31801281,  0.23532556,\n",
       "          -0.0417672 , -1.46597814, -0.17880963, -0.94411259, -0.08349777,\n",
       "           0.4341476 , -0.40342387,  0.07290207, -0.92926026, -2.10021171,\n",
       "           0.24227902,  0.08737664,  0.00689684,  0.10781924, -0.02777451]])},\n",
       " 'hidden_layers_1': {'bias': array([-0.10523192,  0.65090285,  0.07125775,  0.24536489,  0.15011906,\n",
       "          0.04090735, -0.07355307,  0.18668136,  0.1132408 ,  0.03973782,\n",
       "         -0.21500483,  0.23703741, -0.20092357, -0.19889392, -0.4401155 ,\n",
       "          0.34545362, -0.21407945,  0.02534702, -0.19728223, -0.10896258,\n",
       "          0.18060155,  0.11064651, -0.01026664,  0.18269361, -0.13326023,\n",
       "         -0.67490565, -0.19952501,  0.05365538,  0.04113639, -0.0729803 ,\n",
       "         -0.24676103, -0.20417987, -0.27150293,  0.02585729, -0.00437915,\n",
       "         -0.27463694, -0.04641005, -0.03388116, -0.55101811, -0.05673829]),\n",
       "  'kernel': array([[-0.72328042, -0.31289081,  0.27483429, ..., -0.61904931,\n",
       "          -0.70807698, -0.25796445],\n",
       "         [-0.45498142, -0.39916556,  0.14605232, ..., -0.23984086,\n",
       "           0.02427937, -0.19350098],\n",
       "         [ 0.17180435,  0.08514957, -0.23231588, ..., -0.28940572,\n",
       "          -0.32193711,  0.24388677],\n",
       "         ...,\n",
       "         [-0.62829625, -0.60321474,  0.0231589 , ...,  0.12333449,\n",
       "           0.59705802, -0.3622333 ],\n",
       "         [-0.09669691, -0.28747982, -0.127237  , ...,  0.40760873,\n",
       "           0.10791957, -0.52337793],\n",
       "         [ 0.08787958,  0.49798089,  0.62726678, ..., -0.53116352,\n",
       "          -1.71866141, -0.09253584]])},\n",
       " 'hidden_layers_2': {'bias': array([ 0.19349619, -0.08511994,  0.09771231,  0.19375971, -0.18698296,\n",
       "          0.31582219,  0.06804172,  0.07876567, -0.04852523, -0.28101584,\n",
       "         -0.06528024, -0.02099682, -0.1842872 ,  0.03906725, -0.37905565,\n",
       "         -0.0667445 , -0.14562579,  0.04925542,  0.10386684,  0.07176431,\n",
       "          0.1837394 , -0.2648383 , -0.23214424, -0.12740967,  0.0541879 ,\n",
       "          0.22681277, -0.11907119, -0.09484706, -0.00264673,  0.11956076,\n",
       "         -0.00721151,  0.08074111,  0.00687319,  0.03322777, -0.08010637,\n",
       "         -0.10564207,  0.05191083, -0.12241992,  0.02368262, -0.06859319]),\n",
       "  'kernel': array([[-0.44889281, -0.60272633, -0.25959694, ..., -0.02434016,\n",
       "          -0.27376682,  0.42550641],\n",
       "         [-0.12759051, -0.00703998,  0.3354465 , ..., -0.06008251,\n",
       "           0.07274301, -0.01226686],\n",
       "         [ 0.72677646,  0.46289062,  0.29208677, ..., -0.17373293,\n",
       "           0.0100499 ,  0.19904372],\n",
       "         ...,\n",
       "         [-0.42725198,  0.1362645 , -0.7970737 , ...,  0.17257789,\n",
       "          -0.03458357, -0.38016795],\n",
       "         [-0.2312926 ,  0.39838125,  0.33885562, ...,  0.23117707,\n",
       "          -0.72438752,  0.08262066],\n",
       "         [ 0.20250777,  0.18881076,  0.12767289, ..., -0.41790285,\n",
       "          -0.01604   ,  0.72958194]])}}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    " state['params']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "399eb548",
   "metadata": {},
   "source": [
    "To create an instance of the FE-ANN(s) EoS, you need the list of features (`state['features']`). The parameters (`state['params']`) will be needed to use the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7930c8bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# creating the model with the given features\n",
    "model = HelmholtzModel(features=state['features'])\n",
    "params = {'params': state['params']}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7029601",
   "metadata": {},
   "source": [
    "The `model.apply` function evaluates the model; this function requires the parameters and the model inputs (alpha, rhoad, and Tad). The alpha parameter is used as the descriptor of the Mie potential and can be computed with the `helper_get_alpha` function.\n",
    "\n",
    "$$ \\alpha_{\\mathrm{vdw}} = \\mathcal{C}^{\\mathrm{Mie}} \\left( \\frac{1}{\\lambda_a - 3} - \\frac{1}{\\lambda_r - 3}\\right) $$\n",
    "\n",
    "By default, the method `model.apply` computes the residual Helmholtz free energy. See the next notebook for examples of how to get other thermophysical properties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "201ede2d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1.12845457], dtype=float64)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lr = np.array([12])\n",
    "la = np.array([6])\n",
    "alpha = helper_get_alpha(lr, la)\n",
    "rhoad = np.array([0.8])\n",
    "Tad = np.array([2.4])\n",
    "\n",
    "# computing the residual Helmholtz free energy\n",
    "model.apply(params, alpha, rhoad, Tad)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e996291",
   "metadata": {},
   "source": [
    "The FE-ANN(s) EoS is already parallelized. Multiple state points can be evaluated per call. The inputs must have the same length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "84442ea6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([ 1.12845457, -0.50372673], dtype=float64)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lr = np.array([12, 12])\n",
    "la = np.array([6, 6])\n",
    "alpha = helper_get_alpha(lr, la)\n",
    "rhoad = np.array([0.8, 0.135])\n",
    "Tad = np.array([2.4, 1.4])\n",
    "\n",
    "# computing the residual Helmholtz free energy\n",
    "model.apply(params, alpha, rhoad, Tad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3d23e02",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jax",
   "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}