examples/how_to_compute_integrals.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f9499194-335f-46c4-9845-6e9f4b29acc0",
   "metadata": {},
   "source": [
    "# How To Compute the Integrals for the Free Energy\n",
    "\n",
    "This notebook shows how to compute the required energy and gradient functions. That will be used for the MeanFieldVarGP algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0405a8d-ff1a-47cb-973f-b18e0fe2cc08",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import sys\n",
    "\n",
    "import dill\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import sympy as sym\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from sympy.stats import E, Normal\n",
    "\n",
    "from sympy.utilities.iterables import flatten\n",
    "\n",
    "from copy import deepcopy\n",
    "\n",
    "module_path = os.path.abspath(os.path.join('../src/'))\n",
    "\n",
    "if module_path not in sys.path:\n",
    "    sys.path.append(module_path)\n",
    "# _end_if_\n",
    "\n",
    "from numerical.symbolics import LagrangePolynomial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8887ae9f-482d-47b3-a284-285f37df512e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Help make the output easier to read.\n",
    "sym.init_printing(use_unicode=True)\n",
    "\n",
    "# Objects referred to in the global dictionary\n",
    "# are recursively traced and serialized.\n",
    "dill.settings['recurse'] = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fffe1703-a2e6-46fa-920f-9237b612bde2",
   "metadata": {},
   "source": [
    "## Dynamical System equations\n",
    "\n",
    "To generate the equations we use the following settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e756644-c93d-4f05-922d-0ad9c9805cba",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# |------------------|\n",
    "# | Double Well (1D) |\n",
    "# |------------------|\n",
    "# \n",
    "# D = 1\n",
    "# \n",
    "# p = sym.Matrix([sym.symbols(\"theta\", real=True)])\n",
    "# \n",
    "# fx = sym.Matrix((4 * xt) * (p - xt**2), real=True)\n",
    "# \n",
    "\n",
    "# |-------------------------|\n",
    "# | Ornstein-Uhlenbeck (1D) |\n",
    "# |-------------------------|\n",
    "# \n",
    "# D = 1\n",
    "# \n",
    "# p = sym.Matrix([sym.symbols(\"theta mu\", real=True)])\n",
    "# \n",
    "# fx = sym.Matrix([p[0] * (p[1] - xt[0])], real=True)\n",
    "# \n",
    "\n",
    "# |---------------|\n",
    "# | Lorenz63 (3D) |\n",
    "# |---------------|\n",
    "# \n",
    "# D = 3\n",
    "# \n",
    "# p = sym.Matrix([sym.symbols(\"sigma rho beta\", real=True)])\n",
    "#\n",
    "# fx = sym.Matrix([p[0]*(xt[1] - xt[0]),\n",
    "#                 xt[0]*(p[1] - xt[2]) - xt[1],\n",
    "#                 xt[0]*xt[1] - p[2]*xt[2]], real=True);\n",
    "# \n",
    "\n",
    "# |---------------|\n",
    "# | Lorenz96 (nD) |\n",
    "# |---------------|\n",
    "# \n",
    "# D = 4\n",
    "# \n",
    "# p = sym.Matrix([sym.symbols(\"theta\", real=True)])\n",
    "# \n",
    "# fx = sym.Matrix([(xt[1] - xt[3]) * xt[2] - xt[0] + p[0]], real=True)\n",
    "# "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78283361-9825-4f20-967b-4f293326783e",
   "metadata": {},
   "source": [
    "## First define the main variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37230d60-f326-4a28-a17d-fff9f96b9336",
   "metadata": {},
   "outputs": [],
   "source": [
    "# State vector dimensions.\n",
    "D = 1\n",
    "\n",
    "# Mean and variance functions: m(t), s(t).\n",
    "mt = sym.symbols(f\"m:{D+1}(t)\", real=True)\n",
    "st = sym.symbols(f\"s:{D+1}(t)\", real=True, positive=True)\n",
    "\n",
    "# Random (state) vector: x(t).\n",
    "xt = sym.Matrix([Normal(f\"x{i}(t)\", mt[i], sym.sqrt(st[i])) for i in range(D)])\n",
    "\n",
    "# System noise (diffusion) coefficient vector: sig := sigma^2.\n",
    "sig = sym.symbols(f\"Sig:{D}\", real=True, positive=True)\n",
    "\n",
    "# Model (drift) parameters vector.\n",
    "p = sym.Matrix([sym.symbols(\"theta\", real=True)])\n",
    "\n",
    "# Drift function.\n",
    "fx = sym.Matrix((4 * xt) * (p - xt**2), real=True)\n",
    "\n",
    "# Define the limits of integration as symbolic variables.\n",
    "ti, tj = sym.symbols(\"ti, tj\", real=True, positive=True)\n",
    "\n",
    "# Number of systems dimensions.\n",
    "# NOTE: this can be different from\n",
    "#       the state vector dimensions!\n",
    "dim_D = len(fx)\n",
    "\n",
    "# Pretty printing.\n",
    "sym.pprint(fx)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a8299b9-c2c7-44d8-ad0d-44c671bef574",
   "metadata": {},
   "source": [
    "## Compute the mean/variance Lagrange functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80c0661e-df68-4350-a4a8-2450b5bfd1ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store the mean and variance functions.\n",
    "Mfun = []\n",
    "Sfun = []\n",
    "\n",
    "# Store the time derivative of the mean and variance functions.\n",
    "dMt = []\n",
    "dSt = []\n",
    "\n",
    "# Store the mean and variance points.\n",
    "Mpar = []\n",
    "Spar = []\n",
    "\n",
    "# Create one polynomial for each state vector dimension.\n",
    "for i in range(D):\n",
    "    \n",
    "    # Create the mean polynomials.\n",
    "    mt_i, t, mk_i, t_h = LagrangePolynomial(letter=f\"d{i}m\", order=3, fp=\"h\")\n",
    "    \n",
    "    # Replace time points with ti, tj.\n",
    "    mt_i = mt_i.replace(t_h[0], ti).replace(t_h[3], tj)\n",
    "    \n",
    "    # Get the time derivative functions.\n",
    "    dMt.append(mt_i.diff(t))\n",
    "    \n",
    "    # Store the functions.\n",
    "    Mfun.append(mt_i)\n",
    "    \n",
    "    # Store the Lagrange coefficients.\n",
    "    Mpar.append(mk_i)\n",
    "    \n",
    "    # Create the variance polynomials.\n",
    "    st_i, t, sk_i, t_c = LagrangePolynomial(letter=f\"d{i}s\", order=2, fp=\"c\")\n",
    "    \n",
    "    # Replace time points with ti, tj.\n",
    "    st_i = st_i.replace(t_c[0], ti).replace(t_c[2], tj)\n",
    "    \n",
    "    # Get the time derivative functions.\n",
    "    dSt.append(st_i.diff(t))\n",
    "    \n",
    "    # Store the functions.\n",
    "    Sfun.append(st_i)\n",
    "    \n",
    "    # Store the Lagrange coefficients.\n",
    "    Spar.append(sk_i)\n",
    "    \n",
    "# _end_for_\n",
    "\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "075c0ea8-b6e6-4e5b-bbb6-13654a1e4474",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "04ed8de8-ad92-4bf6-b7c2-552f07f41484",
   "metadata": {},
   "source": [
    "## Compute the expectations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f31b645-2516-4286-b737-b05bbf52f1b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# E[f(x)]\n",
    "Efx = E(fx).simplify(); Efx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e12af1a-4cd0-41c9-9ae3-80c32c9856c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# E[f(x)^2]\n",
    "Efx2 = E(fx.multiply_elementwise(fx)).simplify(); Efx2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6f440b4-c4cf-476c-91e8-dcd0b1044a50",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Jacobian: dfx/dx\n",
    "dfxdx = fx.jacobian(xt); dfxdx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b382278d-741d-4477-b8d7-b22476ffc21b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# E[dfx/dx]\n",
    "Edfxdx = E(dfxdx).diagonal().simplify(); Edfxdx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37a3fc7f-402d-4bf7-b65a-734bc1a4af46",
   "metadata": {},
   "source": [
    "## Replace m(t), s(t) with their polynomial equivalents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "471fb3fe-8cab-4eaa-9d10-83277b4ecfa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replace the mean functions.\n",
    "for j, mf_j in enumerate(Mfun):\n",
    "    Efx = Efx.replace(mt[j], mf_j)\n",
    "    \n",
    "    Efx2 = Efx2.replace(mt[j], mf_j)\n",
    "        \n",
    "    Edfxdx = Edfxdx.replace(mt[j], mf_j)\n",
    "# _end_for_\n",
    "\n",
    "# Replace the variance functions.\n",
    "for j, sf_j in enumerate(Sfun):\n",
    "    Efx = Efx.replace(st[j], sf_j)\n",
    "    \n",
    "    Efx2 = Efx2.replace(st[j], sf_j)\n",
    "        \n",
    "    Edfxdx = Edfxdx.replace(st[j], sf_j)\n",
    "# _end_for_\n",
    "\n",
    "# Finale message.\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9b6406d-cba5-47f9-8d08-27da8f5cf7e1",
   "metadata": {},
   "source": [
    "## Put all the $E_{sde}$ terms together.\n",
    "\n",
    "These is given by the expression:\n",
    "\n",
    "$\\int_a^b E_{sde}(t)dt = \\frac{1}{2} \\sum_{i=1}^{D} \\frac{1}{\\sigma_{i}^2} \\int_a^b \\bigg( E[f_{i}^2]_{q_t} -2 E[f_{i}]_{q_t} \\dot{m_i}(t) + \\dot{m_i}^2(t) + \\frac{(\\dot{s_i}(t) - \\sigma_{i}^2)^2}{4 s_{i}(t)} + (\\sigma_{i}^2 - \\dot{s_i}(t)) E\\big[\\frac{\\partial f_i}{\\partial x_i}\\big]_{q_t}\\bigg) dt$\n",
    "\n",
    "where $E[\\cdot]_{q_t}$ is the expected value with respect to the Gaussian process.\n",
    "\n",
    "In essence the following cell computes the integral of each of these terms and adds them all together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd19f1e2-271c-4d04-9bf1-f3748477113a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Stores the Integral of Esde_i.\n",
    "integral_Esde = []\n",
    "\n",
    "# Get the expression for each dimension.\n",
    "for i in range(dim_D):\n",
    "        \n",
    "    # Display info.\n",
    "    print(f\"Estimating the Integral of Esde_{i} ...\")\n",
    "    \n",
    "    # Part: 1\n",
    "    ig_01 = sym.integrate(Efx2[i], (t, ti, tj))\n",
    "    print(\" >> ig_01 is finished \")\n",
    "    \n",
    "    # Part: 2\n",
    "    ig_02 =  sym.integrate((Efx[i] * dMt[i]), (t, ti, tj))\n",
    "    print(\" >> ig_02 is finished \")\n",
    "    \n",
    "    # Part: 3\n",
    "    ig_03 = sym.integrate(dMt[i]**2, (t, ti, tj))\n",
    "    print(\" >> ig_03 is finished \")\n",
    "        \n",
    "    # Part: 4\n",
    "    ig_04_01 = sym.integrate((dSt[i]**2)/Sfun[i], (t, ti, tj))\n",
    "    print(\" >> ig_04_01 is finished \")\n",
    "    \n",
    "    ig_04_02 = sym.integrate((dSt[i]*sig[i])/Sfun[i], (t, ti, tj))\n",
    "    print(\" >> ig_04_02 is finished \")\n",
    "    \n",
    "    ig_04_03 = sym.integrate((sig[i]**2)/Sfun[i], (t, ti, tj))\n",
    "    print(\" >> ig_04_03 is finished \")\n",
    "    \n",
    "    # Part: 5\n",
    "    ig_05 = sym.integrate((sig[i] - dSt[i])*Edfxdx[i], (t, ti, tj))\n",
    "    print(\" >> ig_05 is finished \")\n",
    "\n",
    "    # Final expression of the Integral ...\n",
    "    integral_Esde.append(ig_01 - (2.0 * ig_02) + ig_03 + (0.25 * ig_04_01) - (0.5 * ig_04_02) + (0.25 * ig_04_03) + ig_05)\n",
    "    print(\"  \")\n",
    "    \n",
    "# _end_for_\n",
    "\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4eafe06-4cde-4e6c-a113-9fe0a3c740ce",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Print the integral functions.\n",
    "for n, En in enumerate(integral_Esde_dk):\n",
    "    print(n, \":\")\n",
    "    sym.pprint(En)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b962ffa-7d55-4def-8dcb-e8f04d8c191b",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Compute the gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f31bade0-20f7-410e-a611-65950d208abd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define the gradient lists.\n",
    "gradMp = []\n",
    "gradSp = []\n",
    "\n",
    "# Process each dimension separately.\n",
    "for i in range(dim_D):\n",
    "    \n",
    "    # Display info.\n",
    "    print(f\"Calculating gradients for dimension: {i} ...\")\n",
    "    \n",
    "    # Reset the lists for each dimension.\n",
    "    gradMi = []\n",
    "    gradSi = []\n",
    "    \n",
    "    # Derivatives with respect to the mean parameters.\n",
    "    for dim_k in Mpar:\n",
    "        for mk in dim_k:\n",
    "            gradMi.append(integral_Esde_dk[i].diff(mk))\n",
    "        # _end_for_\n",
    "    # _end_for_\n",
    "    \n",
    "    # Derivatives with respect to the variance parameters.\n",
    "    for dim_k in Spar:\n",
    "        for sk in dim_k:\n",
    "            gradSi.append(integral_Esde_dk[i].diff(sk))\n",
    "        # _end_for_\n",
    "    # _end_for_\n",
    "    \n",
    "    # Update the lists.\n",
    "    gradMp.append(gradMi)\n",
    "    gradSp.append(gradSi)\n",
    "    \n",
    "# _end_for_\n",
    "\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af10f8be-b9b8-4ff0-a453-6617ee0eaf69",
   "metadata": {},
   "source": [
    "## Generating lambda functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6db82990-9a47-4dc2-934e-06ec4563a5a2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Lists to store the results.\n",
    "func_ig_Esde = []\n",
    "grad_ig_Esde_mt = []\n",
    "grad_ig_Esde_st = []\n",
    "\n",
    "for i in range(dim_D):\n",
    "    \n",
    "    # Display information.\n",
    "    print(f\"Generating Lambda functions: {i} ... \")\n",
    "    \n",
    "    func_ig_Esde.append(sym.lambdify([*flatten([ti, tj, t_h[1:3], t_c[1:2], Mpar, Spar, sig, p])],\n",
    "                                     integral_Esde_dk[i], modules=[\"scipy\", \"numpy\"], cse=True))\n",
    "    \n",
    "    grad_ig_Esde_mt.append(sym.lambdify([*flatten([ti, tj, t_h[1:3], t_c[1:2], Mpar, Spar, sig, p])],\n",
    "                                        gradMp[i], modules=[\"scipy\", \"numpy\"], cse=True))\n",
    "\n",
    "    grad_ig_Esde_st.append(sym.lambdify([*flatten([ti, tj, t_h[1:3], t_c[1:2], Mpar, Spar, sig, p])],\n",
    "                                        gradSp[i], modules=[\"scipy\", \"numpy\"], cse=True))\n",
    "# _end_for_\n",
    "\n",
    "# Final message.\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "024e3665-be17-4c8e-962a-dc631d0fda33",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "87385143-7df7-482c-b268-39d24c8761be",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Save the functions to files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "081b7962-f42d-4e8a-b101-04577a06abc6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Where to save the functions (parent directory).\n",
    "# NOTE: We save the energy and the gradient (lambdified) functions in separate directories\n",
    "# to improve the readability of the code. In practise they can be anywhere.\n",
    "parent_path = Path(\"path/to/save/the/equations/of/MeanFieldVarGP/src/dynamical_systems\")\n",
    "\n",
    "for i in range(dim_D):\n",
    "    \n",
    "    with open(Path(parent_path / f\"integrals/integral_DW_Esde_{i}.sym\"), \"wb\") as sym_file:\n",
    "        \n",
    "        dill.dump(func_ig_Esde[i], sym_file)\n",
    "        \n",
    "    # _end_with_\n",
    "    \n",
    "    with open(Path(parent_path / f\"integrals/integral_dDW_Esde_dM{i}.sym\"), \"wb\") as sym_file:\n",
    "        \n",
    "        dill.dump(grad_ig_Esde_mt[i], sym_file)\n",
    "        \n",
    "    # _end_with_\n",
    "    \n",
    "    with open(Path(parent_path / f\"integrals/integral_dDW_Esde_dS{i}.sym\"), \"wb\") as sym_file:\n",
    "        \n",
    "        dill.dump(grad_ig_Esde_st[i], sym_file)\n",
    "        \n",
    "    # _end_with_\n",
    "    \n",
    "# _end_for_\n",
    "\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22eaea1a-3785-4d3f-96f4-193b1d04c1eb",
   "metadata": {},
   "source": [
    "## End of notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1d35aa3-7caf-49de-8822-c99ef1ef1d97",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tensorflow_25",
   "language": "python",
   "name": "tensorflow_25"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}