paper/kinn_datagen_reg.ipynb

{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "Created on Wed Jul+ 2020\n",
    "\n",
    "@author:gusmaogabriels@gmail.com // gusmaogabriels@gatech.edu\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jaxlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.1.73'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jaxlib.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.2.25'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "-*- coding: utf-8 -*-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ggusmao3/.pyenv/versions/ML/lib/python3.7/site-packages/jax/experimental/optimizers.py:30: FutureWarning: jax.experimental.optimizers is deprecated, import jax.example_libraries.optimizers instead\n",
      "  FutureWarning)\n"
     ]
    }
   ],
   "source": [
    "from kinn.basis.trainer import Trainer, TrainerSeparate, TrainerCV\n",
    "from kinn.basis.model import model\n",
    "from kinn.basis.nnx import nn\n",
    "from kinn.basis import plt, jit, optimizers, onp, np\n",
    "from kinn.graphics.generators import gen_gif, gen_sample_gif\n",
    "from kinn.basis import np, random, itertools, clear_output, grad, jacfwd, hessian, vmap\n",
    "from kinn.basis import plt, MaxNLocator, FormatStrFormatter, HTML, animation, cm,\\\n",
    "                    left, right, bottom, top, wspace, hspace, np, make_axes_locatable, griddata\n",
    "from scipy.integrate import solve_ivp\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Graphics Aesthetics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "### Paper Examples\n",
    "\n",
    "**0. Kinetics Framework**  \n",
    "**1. Overall Kinetics Definitions** (all types, stoichiometry matrix, parameters)   \n",
    "\n",
    "    i.   Simple Kinetics Type g (Homogeneous) \n",
    "    ii.  Latent Kinetics Type gda (Heterogeneous) \n",
    "    iii. Latent Kinetics Type gdac (Heterogeneous)\n",
    "    iv.  Latent Kinetics Type gdacs** (Heterogeneous)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Kinetics Framework"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MF MKM's consist of a set of elementary or irreducible reaction, also referred to as those comprising rare events. Such events involve the interaction of at most two individual entities whose probability of occurrence is proportional to their concentrations, which gives rise to the power law kinetics. In a nut shell, power-law kinetics can be summarized as a linear combination of rare events to occur whose sampling frequency scales linearly with the concentration of each participant, and whose probability of success is given by the Arrhenius law. MKM models comprise the scenario in which one or more of such elementary rare events occur. The mathematical representation of such intertwined set of events embodies a stack of ODE's, which assigns to the rate of change of each state a linear combination of power-law kinetics according to the system's stoichiometries. Let $\\mathbf{c}$ represent an array of concentrations or fractions of the system constituents at some given time $t$, and $\\mathbf{\\dot{c}}$ the rate of change of $\\mathbf{c}$, the MKM can be conveyed as in Eq. \\ref{eq:kin_ode}.\n",
    "\n",
    "\\begin{equation}\n",
    "    \\mathbf{\\dot{c}}=\\mathbf{M}\\left(\\mathbf{k}(\\theta)\\odot f(\\mathbf{c})\\right)\\label{eq:kin_ode}\\tag{1}\n",
    "\\end{equation}\n",
    "\n",
    "Where $\\mathbf{M}\\in\\mathbb{Z}^{n\\times m}$ is the corresponding stochiometry matrix and $~{f(\\cdot):\\mathbb{R}^n_+\\to\\mathbb{R}^m_+}$ maps $~{\\mathbf{c}:=\\{\\mathbf{c}\\,|\\,\\mathbf{c}\\in\\mathbb{R}^n_+\\}}$ concentrations into the concentration-related parcel of the power-law kinetics and $~{\\mathbf{k}:=\\{\\mathbf{k}(\\theta)~\\in~\\mathbb{R}^m_+,\\,\\theta\\in\\mathbb{R}_{+}\\}}$ is the temperature- and binding-energies-dependent Arrhenius-like rate constant term. $\\mathbf{c}$ encompasses both unbounded, e.g. gas, and bounded (adsorbates or intermediates) species concentration-related state variables, i.e. partial pressures, concentrations, coverage fractions. For further reference, let $\\mathbf{c_g} := \\{{c}_i\\,|\\, i\\in\\mathcal{C}_g\\}$ be the subset of $\\mathbf{c}$ corresponding to gas-phase species and $\\mathbf{c_a}:=\\{c_i\\,|\\,i\\in\\mathcal{C}_a\\}$ is the subset of $\\mathbf{c}$ related to bounded gas-counterpart molecules coverage fractions, and $\\mathbf{c_s}:=\\{c_i\\,|\\,i\\in\\mathcal{C}_s\\}$ be the subset of $\\mathbf{c}$ related to intermediates/radicals species on the catalyst surface, such that $\\cup_{i\\in\\{\\mathbf{g},\\mathbf{a},\\mathbf{s}\\}}\\mathcal{C}_i=\\{1,2,...,m\\,|\\,m\\in\\mathbb{N}\\}$ and $\\cap_{i\\in\\{\\mathbf{g},\\mathbf{a},\\mathbf{s}\\}}\\mathcal{C}_i=\\emptyset$.\n",
    "\n",
    "The state-dependent reaction rate vector, denoted as $\\mathbf{r}(\\mathbf{c},\\theta):=\\{\\mathbf{r}\\in\\mathbb{R}^m\\,|\\,\\mathbf{r}={k(\\theta)}\\odot f(\\mathbf{c})\\}$, comprises a the rate of reaction (frequency) associated with different kinds of elementary reaction: (i) reaction in homogeneous phase, $\\mathbf{r_g} := \\{{r}_i\\,|\\, i\\in\\mathcal{G}_g\\}$, (ii) those involving adsorption-desorption, $\\mathbf{r_d} := \\{{r}_i\\,|\\, i\\in\\mathcal{G}_d\\}$, , (iii) reactions between adsorbed molecules, $\\mathbf{r_a} := \\{{r}_i\\,|\\, i\\in\\mathcal{G}_a\\}$, (iv) reactions involving adsorbed molecules and radicals/intermediates on the surface, $\\mathbf{r_c} := \\{{r}_i\\,|\\, i\\in\\mathcal{G}_c\\}$, and (v) reactions between radicals on the surface, $\\mathbf{r_s} := \\{{r}_i\\,|\\, i\\in\\mathcal{G}_s\\}$, such that $\\cup_{i\\in\\{\\mathbf{g},\\mathbf{d},\\mathbf{a},\\mathbf{c},\\mathbf{s}\\}}\\mathcal{G}_i=\\{1,2,...,n\\,|\\,n\\in\\mathbb{N}\\}$ and $\\cap_{i\\in\\{\\mathbf{g},\\mathbf{d},\\mathbf{a},\\mathbf{c},\\mathbf{s}\\}}\\mathcal{G}_i=\\emptyset$. The stoichiometry matrix in turn consits of a composition of reaction-types and submatrices constituents of the full stoichiometry matrix, as follows in Eq. \\ref{eq:stoich_full}.\n",
    "\n",
    "\\begin{align}\n",
    "    \\begin{bmatrix} \\mathbf{\\dot{c}_g}\\\\\\mathbf{\\dot{c}_a}\\\\\\mathbf{\\dot{c}_s}\\end{bmatrix}=\\begin{bmatrix} \\mathbf{M_{gg}}&\\mathbf{{M_{gd}}}&\\mathbf{0}&\\mathbf{0}&\\mathbf{0}\\\\\\mathbf{0}&\\mathbf{{M_{ad}}}&\\mathbf{{M_{aa}}}&\\mathbf{{M_{ac}}}&\\mathbf{0}\\\\\\mathbf{0}&\\mathbf{0}&\\mathbf{0}&\\mathbf{M_{sc}}&\\mathbf{M_{ss}}\\end{bmatrix}\\mathbf{r}(\\mathbf{c},\\theta)\n",
    "    \\label{eq:stoich_full}\\tag{2}\n",
    "\\end{align}\n",
    "\n",
    "The full stoichiometry matrix in Eq. \\ref{eq:stoich_full} generally conveys the different types of chemical kinetics: (i) purely homogeneous reactions, only $\\mathbf{M_{:g}}\\ne\\mathbf{0}$, (ii) pure adsorption/desorption reactions, $\\mathbf{M_{:\\{gd\\}}}\\ne\\mathbf{0}$, (iii) single intermediate surface reactions, $\\mathbf{M_{:\\{gdac\\}}}\\ne\\mathbf{0}$, and (iv) reaction containing elementary reactions involving surface intermediates, $\\mathbf{M_{:\\{gdacs\\}}}\\ne\\mathbf{0}$, where $:$ denotes all rows and $\\{\\mathbf{i}\\}\\subset\\{\\mathbf{g},\\mathbf{d},\\mathbf{a},\\mathbf{c},\\mathbf{s}\\}$ are the corresponding reaction types $\\mathbf{i}$ as columns of $\\mathbf{M}$. This framework allows for further classification of types of kinetics solely based on the chemical reaction stoichiometry matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### BC-Constrained Neural Net\n",
    "\n",
    "What we do here is to use an operator that enforces $\\mathbf{x}(\\omega_{sm},t_0)=\\mathbf{x}_0~\\forall~\\omega_{sm}$. A natural choice for this purpose is $c(t)=tanh(t)\\in \\mathcal{C}^1$, since $tanh(0)=0$ and $\\partial_t\\tanh(0)=1$. Let $\\mathbf{x}(\\omega_{sm},t)$ denote a surrogate model, in this case a neural network with parameters $\\omega_{sm}$ and independent variable $t$, the boundary condition operator is defined as $\\mathbf{C}[\\mathbf{x}(\\omega_{sm},t),\\mathbf{x}_0]=\\mathbf{x}(\\omega_{sm},t)c(t-t_0)+\\mathbf{x}_0$. \n",
    "\n",
    "From simple inspection, we have that \n",
    "$$\\mathbf{x}(\\omega_{sm},t)c(t-t_0)|_{t=t_0}+\\mathbf{x}_0=\\mathbf{x}(\\omega_{sm},t_0)c(0)+\\mathbf{x}_0=\\mathbf{x}_0$$\n",
    "$$\\partial_t~\\mathbf{x}(\\omega_{sm},t)c(t-t_0)|_{t=t_0}+\\mathbf{x}_0=\\mathbf{x}_t(\\omega_{sm},t_0)c(0)+\\mathbf{x}(\\omega_{sm},t_0)c_t(0)=\\mathbf{x}(\\omega_{sm},t_0)$$\n",
    "\n",
    "Such that the neural net vanishes at $t=t_0$ with continuous time-derivative equals the neural net output. Therefore, any *Dirichlet* boundary condition can be automatically satisfied in this form for UDEs. The same approach can be extend to *Neumann* boundary conditions with proper choice of $c$ function.\n",
    "\n",
    "$$\\mathbf{\\dot{x}}(\\omega_{sm},t)=Mf(\\mathbf{x}(\\omega_{sm},t))$$\n",
    "$$\\mathbf{\\dot{x}}(\\omega_{sm},t_0)=Mf(\\mathbf{x}(\\omega_{sm},t_0))$$\n",
    "$$\\mathbf{x}(\\omega_{sm},t)=\\mathbf{C}[\\mathbf{x'}(\\omega_{sm},t)]=\\mathbf{x'}(\\omega_{sm},t)c(t)$$\n",
    "$$\\partial_t[{\\mathbf{x'}(\\omega_{sm},t)c(t-t_0)}+\\mathbf{x}_0]=Mf(\\mathbf{x'}(\\omega_{sm},t)c(t-t_0)+\\mathbf{x}_0)$$\n",
    "$$\\mathbf{\\dot{x}'}(\\omega_{sm},t)c(t-t_0)+\\mathbf{x'}(\\omega_{sm},t)\\dot{c}(t-t_0)=Mf(\\mathbf{x'}(\\omega_{sm},t)c(t-t_0)+\\mathbf{x}_0)$$\n",
    "$$\\mathbf{x'}(\\omega_{sm},t_0)=Mf(\\mathbf{x}_0)$$\n",
    "\n",
    "Below `nn_c` is recasted from the parent `nn` function by adding $\\mathcal{C}[x]$ as a constraint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hypersphere Contrained Subclassed NN for Latent Variables\n",
    "\n",
    "#### General Definition\n",
    "\n",
    "Kinetic models under mean-field approximation deal with states as the evolution of descriptive statistics over the scale under considerations. In particular, MKM treat concentration of intermediate species in terms of fractions of the total number of active sites available. As a result, it is desirable to develop surrogate models that structurally enforce normalization by eliminating the additional degree of freedom. To this end, we propose the use of the projection od radius-one hypersphere onto the natural basis. \n",
    "\n",
    "Let $\\mathbf{x_s}(\\omega_{sm},t)\\in\\mathbb{R}^p$ be the output values of the constrained surrogate related to the neural network $x'(\\omega_{sm},t)\\in\\mathbb{R}^{p-1}$, such that\n",
    "\n",
    "\\begin{align}\n",
    "    {x_s}_i(\\omega_{sm},t)&=\\left(1-\\sin^{2}\\left({x'}_i(\\omega_{sm},t)\\right)\\right)\\prod_{j<i}\\sin^{2}\\left({x'}_j(\\omega_{sm},t)\\right)\\;\\forall\\;i<p;\\;{i,j}\\in\\mathbb{N}\\\\\n",
    "    {x_s}_p(\\omega_{sm},t)&=\\prod_{j<p}\\sin^{2}\\left({x'}_j(\\omega_{sm},t)\\right)\n",
    "\\end{align}\n",
    "\n",
    "Such trigonometric transformation enforces that $0\\le\\mathbf{x_s}\\le1$ and $\\sum\\mathbf{x_s}=1\\;\\forall\\;\\mathbf{x'}\\in\\mathbb{R}^{p-1}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Hypersphere Constrained with Boundary Conditions \n",
    "*Dirichlet*\n",
    "\n",
    "Since the hypersphere constraint maps from $\\mathbb{R}^{p-1}$ to $\\mathbb{R}^p$, applying boundary conditions operators on the surrogate model output would lead to increasing stiffness in the underlying neural net training (**hand-wavy: may need additional math elaboration**). The workaround is to transform the output layer boundary conditions to the output of the inner neural network, i.e. convert the boundary contidition to equivalent hypersphere angles. The latter can be undertaken by properly finding the bijection from the hypersphere transformation, as follows.\n",
    "\n",
    "\\begin{align}\n",
    "    {x'_0}_i=\\arcsin{\\left(\\sqrt{1-\\frac{{x_0}_i}{\\prod_{j<i}\\sin^2\\left({x_0}_j\\right)}}\\right)};\\forall\\;i<p;\\;i,j\\in\\mathbb{N}\n",
    "\\end{align}\n",
    "\n",
    "Where $\\mathbf{x'_0}$ is the angle-transformed boundary conditions and $\\mathbf{x_0}$ is the natural boundary condition for the surrogate model output. For the degenerate case where one of the outputs is the unit, the corresponding angle should be set to $0$ and all other angles to $\\frac{\\pi}{2}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Overall Kinetics Definitions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### i. Simple Kinetics Type gg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The simple initial model represents the following fully-reversible chemical reaction:\n",
    "\n",
    "$$A+B\\underset{k_2}{\\stackrel{k_1}{\\rightleftharpoons}} C\\notag$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ii. Latent Kinetics Type gda"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Latent Kinetics type gad involves ad/desorption steps and a surface reaction between adsorbed molecules.\n",
    "\n",
    "$$A\\underset{k_2}{\\stackrel{k_1}{\\rightleftharpoons}} A*\\notag\\\\\n",
    "B\\underset{k_4}{\\stackrel{k_3}{\\rightleftharpoons}} B*\\notag\\\\\n",
    "C\\underset{k_6}{\\stackrel{k_5}{\\rightleftharpoons}} C*\\notag\\\\\n",
    "A*+B*\\underset{k_8}{\\stackrel{k_7}{\\rightleftharpoons}} C*\\notag$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### iii. Latent Kinetics Type gdac"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An intermediate species (radicals, $D*$) that do not have a corresponding gas phase species is part of the reaction.\n",
    "\n",
    "$$A\\underset{k_2}{\\stackrel{k_1}{\\rightleftharpoons}} A*\\notag\\\\\n",
    "B\\underset{k_4}{\\stackrel{k_3}{\\rightleftharpoons}} B*\\notag\\\\\n",
    "C\\underset{k_6}{\\stackrel{k_5}{\\rightleftharpoons}} C*\\notag\\\\\n",
    "B*+*\\underset{k_8}{\\stackrel{k_7}{\\rightleftharpoons}} 2D*\\notag\\\\\n",
    "A*\\:+\\:D*\\underset{k_{10}}{\\stackrel{k_9}{\\rightleftharpoons}} C*\\:+\\:*\\notag$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### iv. Latent Kinetics Type gdacs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reaction between radicals $D*$, $E*$ and $F*$ adds further complexity to surface reaction.\n",
    "\n",
    "$$A\\underset{k_2}{\\stackrel{k_1}{\\rightleftharpoons}} A*\\notag\\\\\n",
    "B\\underset{k_4}{\\stackrel{k_3}{\\rightleftharpoons}} B*\\notag\\\\\n",
    "C\\underset{k_6}{\\stackrel{k_5}{\\rightleftharpoons}} C*\\notag\\\\\n",
    "A*+*\\underset{k_8}{\\stackrel{k_7}{\\rightleftharpoons}} 2D*\\notag\\\\\n",
    "B*+*\\underset{k_10}{\\stackrel{k_9}{\\rightleftharpoons}} 2E*\\notag\\\\\n",
    "D*\\:+\\:E*\\underset{k_{12}}{\\stackrel{k_{11}}{\\rightleftharpoons}} F*\\:+\\:*\\notag\\\\\n",
    "F*\\:+\\:E*\\underset{k_{14}}{\\stackrel{k_{13}}{\\rightleftharpoons}} C*\\:+\\:*\\notag$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Concentrations of chemical species $A$, $B$ and $C$ are observable variables, i.e. for topics [1], [2] and [3], there are not any latent (hidden) variables."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A dictionary of stoichiometry matrices and with kinetic parameters (rate constants).\n",
    "\n",
    "All reactions are considered reversible"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Parameter `dict` for trainer object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trainer_source import pars"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model parameters (rate constants): `model_params0`  \n",
    "Boundary conditions (concentrations at $t=0$): `bc0`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model function, `x` are state variables, `t` the independet variable (e.g. time)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trainer_source import nn_c, nn_cn, nn_cn_bc, normtrig, nn_combo \n",
    "from trainer_source import trainer as trainer_fwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'type': 'heterogeneous',\n",
       " 'sps': ['A', 'B', 'C', 'A*', 'B*', 'C*', 'D*', '*'],\n",
       " 'stoich': DeviceArray([[-1,  1,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "              [ 0,  0, -1,  1,  0,  0,  0,  0,  0,  0],\n",
       "              [ 0,  0,  0,  0, -1,  1,  0,  0,  0,  0],\n",
       "              [ 1, -1,  0,  0,  0,  0,  0,  0, -1,  1],\n",
       "              [ 0,  0,  1, -1,  0,  0, -1,  1,  0,  0],\n",
       "              [ 0,  0,  0,  0,  1, -1,  0,  0,  1, -1],\n",
       "              [ 0,  0,  0,  0,  0,  0,  2, -2, -1,  1],\n",
       "              [-1,  1, -1,  1, -1,  1, -1,  1,  1, -1]], dtype=int64),\n",
       " 'kinpars': DeviceArray([  20.,    8.,   16.,    4.,   12.,    8., 1200.,  400.,\n",
       "              2000., 1600.], dtype=float64),\n",
       " 'bc': [[DeviceArray([0.], dtype=float64),\n",
       "   DeviceArray([0.6, 0.4, 0. , 0. , 0. , 0. , 0. , 1. ], dtype=float64)],\n",
       "  [DeviceArray([0.], dtype=float64),\n",
       "   DeviceArray([0.2, 0.3, 0.5, 0. , 0. , 0. , 0. , 1. ], dtype=float64)]],\n",
       " 'nncpars': [{'forward': {'trig': False,\n",
       "    'damp': True,\n",
       "    'nobs': 3,\n",
       "    'bc': [DeviceArray([0.], dtype=float64),\n",
       "     DeviceArray([0.6, 0.4, 0. , 0. , 0. , 0. , 0. , 1. ], dtype=float64)]},\n",
       "   'inverse': {'trig': True,\n",
       "    'damp': False,\n",
       "    'nobs': 3,\n",
       "    'bc': [DeviceArray([0.], dtype=float64),\n",
       "     DeviceArray([0.6, 0.4, 0. , 0. , 0. , 0. , 0. , 1. ], dtype=float64)]}},\n",
       "  {'forward': {'trig': False,\n",
       "    'damp': True,\n",
       "    'nobs': 3,\n",
       "    'bc': [DeviceArray([0.], dtype=float64),\n",
       "     DeviceArray([0.2, 0.3, 0.5, 0. , 0. , 0. , 0. , 1. ], dtype=float64)]},\n",
       "   'inverse': {'trig': True,\n",
       "    'damp': False,\n",
       "    'nobs': 3,\n",
       "    'bc': [DeviceArray([0.], dtype=float64),\n",
       "     DeviceArray([0.2, 0.3, 0.5, 0. , 0. , 0. , 0. , 1. ], dtype=float64)]}}],\n",
       " 'nnmpars': {'forward': {'layers_sizes': [[1, 16, 16, 16, 8]],\n",
       "   'act_fun': [<CompiledFunction of <function <lambda> at 0x7f2d541e2d90>>,\n",
       "    <CompiledFunction of <function <lambda> at 0x7f2c95c7fe18>>,\n",
       "    <CompiledFunction of <function <lambda> at 0x7f2d541e2d90>>],\n",
       "   'nn_scale': 0.001},\n",
       "  'inverse': {'layers_sizes': [[1, 14, 14, 14, 7]],\n",
       "   'act_fun': [<CompiledFunction of <function <lambda> at 0x7f2d541e2d90>>,\n",
       "    <CompiledFunction of <function <lambda> at 0x7f2c95c7fe18>>,\n",
       "    <CompiledFunction of <function <lambda> at 0x7f2d541e2d90>>],\n",
       "   'nn_scale': 0.001}},\n",
       " 'nntpars': {'layers_sizes': [[1, 6, 1]],\n",
       "  'act_fun': [<CompiledFunction of <function <lambda> at 0x7f2c95c7fe18>>,\n",
       "   <CompiledFunction of <function <lambda> at 0x7f2d541e2d90>>,\n",
       "   <CompiledFunction of <function <lambda> at 0x7f2c95c7fe18>>,\n",
       "   <CompiledFunction of <function <lambda> at 0x7f2d541e2d90>>],\n",
       "  'nn_scale': 0.001},\n",
       " 'kinpars_tags': ['k_{1}^d',\n",
       "  'k_{\\text{-}1}^d',\n",
       "  'k_{2}^d',\n",
       "  'k_{\\text{-}2}^d',\n",
       "  'k_{3}^d',\n",
       "  'k_{\\text{-}3}^d',\n",
       "  'k_{1}^c',\n",
       "  'k_{\\text{-}1}^c',\n",
       "  'k_{2}^c',\n",
       "  'k_{\\text{-}2}^c']}"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pars[2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate All Case Studios"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 7999, Loss Batch: 2.24595e+00, Loss Data: 2.24595e+00 Loss CV: 2.24595e+00, Fit MODEL: 4.08571e-02, Fit DATA: 2.20509e-04\n",
      "2.2459515817087516\n",
      "['diff_obs', 'state_obs', 'diff_latent', 'state_latent']\n",
      "['diff_obs', 'state_obs', 'diff_latent', 'state_latent']\n"
     ]
    }
   ],
   "source": [
    "__reload__ = False\n",
    "__errors__ = {}\n",
    "__results__ = {}\n",
    "\n",
    "addon = '_alpha13'\n",
    "#_alphas = np.logspace(-2,4,5)\n",
    "_alphas = np.logspace(-2,4,6)\n",
    "for no in range(2,3):\n",
    "    \n",
    "    for tag, mode in zip(['fwd','inv','invsc','invvwn'],['forward']+['inverse']*3):\n",
    "    \n",
    "        if tag.startswith('inv'):\n",
    "            alphas = _alphas\n",
    "        else:\n",
    "            alphas = [1]\n",
    "\n",
    "        model_scale = 1e-2 # model scale (assume low so as not to be bias)\n",
    "        model_nn = model(pars[no]['stoich'],model_scale=model_scale)\n",
    "        if tag == 'fwd':\n",
    "            model_nn.params = [pars[no]['kinpars']]\n",
    "\n",
    "        model_ = model(pars[no]['stoich'])\n",
    "        model_.params = [pars[no]['kinpars']]\n",
    "\n",
    "        def gen_data(n_points,i):\n",
    "            @jit\n",
    "            def ode(t,C):\n",
    "                return model_.single_eval([pars[no]['kinpars']],[t,C]).flatten()\n",
    "\n",
    "            tmax = 20 # max time to evaluate\n",
    "            t_eval = (np.logspace(0,np.log10(tmax+1),n_points)-1.)/tmax\n",
    "            print((pars[no]['bc'][i][0], tmax), pars[no]['bc'][i][1])\n",
    "            sol = solve_ivp(ode, (pars[no]['bc'][i][0], tmax), pars[no]['bc'][i][1], t_eval = t_eval, method='LSODA',atol=1e-20,rtol=1e-20)\n",
    "\n",
    "            return t_eval.reshape(-1,1), sol.y.T\n",
    "\n",
    "        nnms = []\n",
    "        nnts = []\n",
    "        nncs = [] # one per dataset\n",
    "        for i in range(len(pars[no]['bc'])):\n",
    "            nnm = nn(**pars[no]['nnmpars'][mode])\n",
    "            nnms += [nnm]\n",
    "            nnt = nn(**pars[no]['nntpars'])\n",
    "            nnts += [nnt]\n",
    "            if mode == 'forward':\n",
    "                nncs += [nn_combo([nnm,nnt], mode=mode,**pars[no]['nncpars'][i][mode])]\n",
    "            elif mode == 'inverse': \n",
    "                nncs += [nn_combo([nnm], mode=mode,**pars[no]['nncpars'][i][mode])]\n",
    "            else:\n",
    "                raise Exception('mode not implemented ({})'.format(mode))\n",
    "\n",
    "        num_epochs = 1000\n",
    "        num_iter = 100\n",
    "\n",
    "        if tag.endswith('sc') or tag.endswith('wn'):\n",
    "            scale = True\n",
    "        else:\n",
    "            scale = False\n",
    "\n",
    "        print('scale',scale)\n",
    "        trainerf = trainer_fwd(nncs, model_nn, num_iter=num_iter, num_epochs=num_epochs, batch_size=1.,\\\n",
    "                               split=1., verbose=True, mode=mode, scale=scale, historian=True, tol=1e-10, nobs=pars[no]['nncpars'][i][mode]['nobs'],iter_data=[])\n",
    "        trainerf.__setattr__('err_tags',['MODEL','DATA'])\n",
    "\n",
    "        if __reload__:# and not (tag.endswith('sc') or tag.endswith('wn')):\n",
    "            try:\n",
    "                trainerf.load('trainer_{}_{}{}.npz'.format(tag,no,addon))\n",
    "            except:\n",
    "                print('Load failed: {} {}'.format(tag,no))\n",
    "                __errors__.update({'LOAD_FAILED':(tag,no)})\n",
    "\n",
    "        datas0 = []\n",
    "        for i in range(len(pars[no]['bc'])):\n",
    "            t, x = gen_data(100,i)\n",
    "            data0 = [(t,x.copy())]\n",
    "            datas0 += data0\n",
    "        dstack = np.vstack([_[1][:,pars[no]['nncpars'][i][mode]['nobs']:] for _ in datas0])\n",
    "        s = np.std(dstack,axis=0)   \n",
    "\n",
    "        datas  = []\n",
    "\n",
    "        for i in range(len(pars[no]['bc'])):\n",
    "\n",
    "            t, x = datas0[i]\n",
    "\n",
    "            if tag == 'fwd': \n",
    "                data = [(t,[])]\n",
    "            elif tag == 'inv':\n",
    "                xinv  = x[:,:pars[no]['nncpars'][i][mode]['nobs']]\n",
    "                data = [(t,xinv)]\n",
    "            elif tag == 'invsc':\n",
    "                d = x[:,pars[no]['nncpars'][i][mode]['nobs']:]\n",
    "                d = d/s\n",
    "                x_sc = np.hstack((x[:,:pars[no]['nncpars'][i][mode]['nobs']],d))\n",
    "                data = [(t,x_sc)]\n",
    "            elif tag == 'invvwn':\n",
    "                d = x[:,pars[no]['nncpars'][i][mode]['nobs']:]\n",
    "                d = d/s\n",
    "                x_sc = np.hstack((x[:,:pars[no]['nncpars'][i][mode]['nobs']],d))\n",
    "                x_scwn = x_sc + random.normal(random.PRNGKey(0),x_sc.shape)*0.025\n",
    "                data = [(t,x_scwn)]\n",
    "            datas  += data\n",
    "\n",
    "        for tol, h in zip([1e-12],[1e-3]): \n",
    "            trainerf.tol = tol  \n",
    "            adam_kwargs  = {'step_size':h,'b1':0.9,'b2':0.9999,'eps':1e-12}\n",
    "            trainerf.set_optimizer(optimizers.adam,adam_kwargs)\n",
    "            trainerf.rejit()\n",
    "            for alpha in alphas:\n",
    "                trainerf.reinit()\n",
    "                print(no,alpha)\n",
    "                trainerf.train(datas,alpha=alpha)\n",
    "        for tol, h in zip([1e-12]*2,[1e-4,1e-5]): \n",
    "            trainerf.tol = tol  \n",
    "            adam_kwargs  = {'step_size':h,'b1':0.9,'b2':0.9999,'eps':1e-12}\n",
    "            trainerf.set_optimizer(optimizers.adam,adam_kwargs)\n",
    "            trainerf.rejit()\n",
    "            trainerf.reinit()\n",
    "            print(no,alpha)\n",
    "            trainerf.train(datas,alpha=alpha)\n",
    "\n",
    "        __results__['trainer_{}_{}{}'.format(tag,no,addon)] = trainerf.get_error_metrics(datas0)\n",
    "\n",
    "        trainerf.dump('trainer_{}_{}{}'.format(tag,no,addon))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate Pareto Sensitivities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 36599, Loss Batch: 1.98822e-05, Loss Data: 1.98822e-05 Loss CV: 1.98822e-05, Fit MODEL: 1.98774e-05, Fit DATA: 4.86432e-03\n",
      "1.988223701804666e-05\n",
      "['diff_obs', 'state_obs', 'diff_latent', 'state_latent']\n",
      "['diff_obs', 'state_obs', 'diff_latent', 'state_latent']\n"
     ]
    }
   ],
   "source": [
    "__reload__ = True\n",
    "__errors__ = {}\n",
    "__results__ = {}\n",
    "\n",
    "addon = '_alpha17sens'\n",
    "_alphas = np.concatenate((np.logspace(-6,6,21),np.logspace(-6,6,41)[-2::-1]))\n",
    "\n",
    "for no in [3]:\n",
    "    \n",
    "    for tag, mode in zip(['invvwn','invsc'],['inverse']*2):\n",
    "    \n",
    "        if tag.startswith('inv'):\n",
    "            alphas = _alphas\n",
    "        else:\n",
    "            alphas = [1]\n",
    "\n",
    "        model_scale = 1e-1 # model scale (assume low so as not to be bias)\n",
    "        model_nn = model(pars[no]['stoich'],model_scale=model_scale)\n",
    "        if tag == 'fwd':\n",
    "            model_nn.params = [pars[no]['kinpars']]\n",
    "\n",
    "        model_ = model(pars[no]['stoich'])\n",
    "        model_.params = [pars[no]['kinpars']]\n",
    "\n",
    "        def gen_data(n_points,i):\n",
    "            @jit\n",
    "            def ode(t,C):\n",
    "                return model_.single_eval([pars[no]['kinpars']],[t,C]).flatten()\n",
    "\n",
    "            tmax = 20 # max time to evaluate\n",
    "            t_eval = (np.logspace(0,np.log10(tmax+1),n_points)-1.)/tmax\n",
    "            print((pars[no]['bc'][i][0], tmax), pars[no]['bc'][i][1])\n",
    "            sol = solve_ivp(ode, (pars[no]['bc'][i][0], tmax), pars[no]['bc'][i][1], t_eval = t_eval, method='LSODA',atol=1e-20,rtol=1e-20)\n",
    "            return t_eval.reshape(-1,1), sol.y.T\n",
    "\n",
    "        nnms = []\n",
    "        nnts = []\n",
    "        nncs = [] # one per dataset\n",
    "        for i in range(len(pars[no]['bc'])):\n",
    "            nnm = nn(**pars[no]['nnmpars'][mode])\n",
    "            nnms += [nnm]\n",
    "            nnt = nn(**pars[no]['nntpars'])\n",
    "            nnts += [nnt]\n",
    "            if mode == 'forward':\n",
    "                nncs += [nn_combo([nnm,nnt], mode=mode,**pars[no]['nncpars'][i][mode])]\n",
    "            elif mode == 'inverse': \n",
    "                nncs += [nn_combo([nnm], mode=mode,**pars[no]['nncpars'][i][mode])]\n",
    "            else:\n",
    "                raise Exception('mode not implemented ({})'.format(mode))\n",
    "\n",
    "        num_epochs = 300\n",
    "        num_iter   = 150\n",
    "\n",
    "        if tag.endswith('sc') or tag.endswith('wn'):\n",
    "            scale = True\n",
    "        else:\n",
    "            scale = False\n",
    "\n",
    "        print('scale',scale)\n",
    "        trainerf = trainer_fwd(nncs, model_nn, num_iter=num_iter, num_epochs=num_epochs, batch_size=1.,\\\n",
    "                               split=1., verbose=True, mode=mode, scale=scale, historian=True, tol=1e-10, nobs=pars[no]['nncpars'][i][mode]['nobs'],iter_data=[])\n",
    "        trainerf.__setattr__('err_tags',['MODEL','DATA'])\n",
    "\n",
    "        if __reload__:# and not (tag.endswith('sc') or tag.endswith('wn')):\n",
    "            try:\n",
    "                trainerf.load('trainer_{}_{}{}.npz'.format(tag,no,addon))\n",
    "            except:\n",
    "                print('Load failed: {} {}'.format(tag,no))\n",
    "                __errors__.update({'LOAD_FAILED':(tag,no)})\n",
    "\n",
    "\n",
    "        datas0 = []\n",
    "        for i in range(len(pars[no]['bc'])):\n",
    "            t, x = gen_data(100,i)\n",
    "            data0 = [(t,x.copy())]\n",
    "            datas0 += data0\n",
    "        dstack = np.vstack([_[1][:,pars[no]['nncpars'][i][mode]['nobs']:] for _ in datas0])\n",
    "        s = np.std(dstack,axis=0)\n",
    "\n",
    "        datas  = []\n",
    "\n",
    "        for i in range(len(pars[no]['bc'])):\n",
    "\n",
    "            t, x = datas0[i]\n",
    "\n",
    "            if tag == 'fwd': \n",
    "                data = [(t,[])]\n",
    "            elif tag == 'inv':\n",
    "                xinv  = x[:,:pars[no]['nncpars'][i][mode]['nobs']]\n",
    "                data = [(t,xinv)]\n",
    "            elif tag == 'invsc':\n",
    "                d = x[:,pars[no]['nncpars'][i][mode]['nobs']:]\n",
    "                d = d/s\n",
    "                x_sc = np.hstack((x[:,:pars[no]['nncpars'][i][mode]['nobs']],d))\n",
    "                data = [(t,x_sc)]\n",
    "            elif tag == 'invvwn':\n",
    "                d = x[:,pars[no]['nncpars'][i][mode]['nobs']:]\n",
    "                d = d/s\n",
    "                x_sc = np.hstack((x[:,:pars[no]['nncpars'][i][mode]['nobs']],d))\n",
    "                x_scwn = x_sc + random.normal(random.PRNGKey(0),x_sc.shape)*0.025\n",
    "                data = [(t,x_scwn)]\n",
    "            datas  += data\n",
    "\n",
    "        trainerf.tol = 1e-12  \n",
    "        adam_kwargs  = {'step_size':1e-4,'b1':0.9,'b2':0.9999,'eps':1e-12}\n",
    "        trainerf.set_optimizer(optimizers.adam,adam_kwargs)\n",
    "        trainerf.rejit()\n",
    "            \n",
    "        for alpha in alphas:\n",
    "            trainerf.alpha=alpha\n",
    "            trainerf.reinit()\n",
    "            print(no,alpha)\n",
    "            trainerf.train(datas,alpha=alpha)\n",
    "\n",
    "        __results__['trainer_{}_{}{}'.format(tag,no,addon)] = trainerf.get_error_metrics(datas0)\n",
    "\n",
    "        trainerf.dump('trainer_{}_{}{}'.format(tag,no,addon))"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "encoding": "# -*- coding: utf-8 -*-",
   "executable": "/usr/bin/env python3",
   "formats": "py:light,ipynb",
   "notebook_metadata_filter": "-all"
  },
  "kernelspec": {
   "display_name": "ML",
   "language": "python",
   "name": "ml"
  },
  "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"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}